pipe the argument to storage (#3184)

This commit is contained in:
Max Novich
2025-06-30 16:55:31 -07:00
committed by GitHub
parent 00eddcb880
commit 6ecf3582ad
5 changed files with 145 additions and 21 deletions

View File

@@ -120,7 +120,7 @@ async fn offer_extension_debugging_help(
std::env::temp_dir().join(format!("goose_debug_extension_{}.jsonl", extension_name)); std::env::temp_dir().join(format!("goose_debug_extension_{}.jsonl", extension_name));
// Create the debugging session // Create the debugging session
let mut debug_session = Session::new(debug_agent, temp_session_file.clone(), false, None); let mut debug_session = Session::new(debug_agent, temp_session_file.clone(), false, None, true);
// Process the debugging request // Process the debugging request
println!("{}", style("Analyzing the extension failure...").yellow()); println!("{}", style("Analyzing the extension failure...").yellow());
@@ -366,6 +366,7 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> Session {
session_file.clone(), session_file.clone(),
session_config.debug, session_config.debug,
session_config.scheduled_job_id.clone(), session_config.scheduled_job_id.clone(),
!session_config.no_session, // save_session is the inverse of no_session
); );
// Add extensions if provided // Add extensions if provided

View File

@@ -52,6 +52,7 @@ pub struct Session {
debug: bool, // New field for debug mode debug: bool, // New field for debug mode
run_mode: RunMode, run_mode: RunMode,
scheduled_job_id: Option<String>, // ID of the scheduled job that triggered this session scheduled_job_id: Option<String>, // ID of the scheduled job that triggered this session
save_session: bool, // Whether to save session to file
} }
// Cache structure for completion data // Cache structure for completion data
@@ -113,13 +114,19 @@ impl Session {
session_file: PathBuf, session_file: PathBuf,
debug: bool, debug: bool,
scheduled_job_id: Option<String>, scheduled_job_id: Option<String>,
save_session: bool,
) -> Self { ) -> Self {
let messages = match session::read_messages(&session_file) { let messages = if save_session {
match session::read_messages(&session_file) {
Ok(msgs) => msgs, Ok(msgs) => msgs,
Err(e) => { Err(e) => {
eprintln!("Warning: Failed to load message history: {}", e); eprintln!("Warning: Failed to load message history: {}", e);
Vec::new() Vec::new()
} }
}
} else {
// Don't try to read messages if we're not saving sessions
Vec::new()
}; };
Session { Session {
@@ -130,6 +137,7 @@ impl Session {
debug, debug,
run_mode: RunMode::Normal, run_mode: RunMode::Normal,
scheduled_job_id, scheduled_job_id,
save_session,
} }
} }
@@ -319,6 +327,7 @@ impl Session {
&self.messages, &self.messages,
Some(provider), Some(provider),
self.scheduled_job_id.clone(), self.scheduled_job_id.clone(),
self.save_session,
) )
.await?; .await?;
@@ -431,6 +440,7 @@ impl Session {
&self.messages, &self.messages,
Some(provider), Some(provider),
self.scheduled_job_id.clone(), self.scheduled_job_id.clone(),
self.save_session,
) )
.await?; .await?;
@@ -619,6 +629,7 @@ impl Session {
&self.messages, &self.messages,
Some(provider), Some(provider),
self.scheduled_job_id.clone(), self.scheduled_job_id.clone(),
self.save_session,
) )
.await?; .await?;
@@ -792,7 +803,7 @@ impl Session {
Err(ToolError::ExecutionError("Tool call cancelled by user".to_string())) Err(ToolError::ExecutionError("Tool call cancelled by user".to_string()))
)); ));
self.messages.push(response_message); self.messages.push(response_message);
session::persist_messages_with_schedule_id(&self.session_file, &self.messages, None, self.scheduled_job_id.clone()).await?; session::persist_messages_with_schedule_id(&self.session_file, &self.messages, None, self.scheduled_job_id.clone(), self.save_session).await?;
drop(stream); drop(stream);
break; break;
@@ -889,7 +900,7 @@ impl Session {
self.messages.push(message.clone()); self.messages.push(message.clone());
// No need to update description on assistant messages // No need to update description on assistant messages
session::persist_messages_with_schedule_id(&self.session_file, &self.messages, None, self.scheduled_job_id.clone()).await?; session::persist_messages_with_schedule_id(&self.session_file, &self.messages, None, self.scheduled_job_id.clone(), self.save_session).await?;
if interactive {output::hide_thinking()}; if interactive {output::hide_thinking()};
let _ = progress_bars.hide(); let _ = progress_bars.hide();
@@ -1093,6 +1104,7 @@ impl Session {
&self.messages, &self.messages,
None, None,
self.scheduled_job_id.clone(), self.scheduled_job_id.clone(),
self.save_session,
) )
.await?; .await?;
@@ -1108,6 +1120,7 @@ impl Session {
&self.messages, &self.messages,
None, None,
self.scheduled_job_id.clone(), self.scheduled_job_id.clone(),
self.save_session,
) )
.await?; .await?;
@@ -1128,6 +1141,7 @@ impl Session {
&self.messages, &self.messages,
None, None,
self.scheduled_job_id.clone(), self.scheduled_job_id.clone(),
self.save_session,
) )
.await?; .await?;

View File

@@ -1248,6 +1248,7 @@ async fn run_scheduled_job_internal(
&session_file_path, &session_file_path,
&updated_metadata, &updated_metadata,
&all_session_messages, &all_session_messages,
true,
) { ) {
tracing::error!( tracing::error!(
"[Job {}] Failed to persist final messages: {}", "[Job {}] Failed to persist final messages: {}",
@@ -1278,6 +1279,7 @@ async fn run_scheduled_job_internal(
&session_file_path, &session_file_path,
&fallback_metadata, &fallback_metadata,
&all_session_messages, &all_session_messages,
true,
) { ) {
tracing::error!("[Job {}] Failed to persist final messages with fallback metadata: {}", job.id, e_fb); tracing::error!("[Job {}] Failed to persist final messages with fallback metadata: {}", job.id, e_fb);
} }
@@ -1304,9 +1306,12 @@ async fn run_scheduled_job_internal(
message_count: 0, message_count: 0,
..Default::default() ..Default::default()
}; };
if let Err(e) = if let Err(e) = crate::session::storage::save_messages_with_metadata(
crate::session::storage::save_messages_with_metadata(&session_file_path, &metadata, &[]) &session_file_path,
{ &metadata,
&[],
true,
) {
tracing::error!( tracing::error!(
"[Job {}] Failed to persist metadata for empty job: {}", "[Job {}] Failed to persist metadata for empty job: {}",
job.id, job.id,

View File

@@ -1054,7 +1054,7 @@ pub async fn persist_messages(
messages: &[Message], messages: &[Message],
provider: Option<Arc<dyn Provider>>, provider: Option<Arc<dyn Provider>>,
) -> Result<()> { ) -> Result<()> {
persist_messages_with_schedule_id(session_file, messages, provider, None).await persist_messages_with_schedule_id(session_file, messages, provider, None, true).await
} }
/// Write messages to a session file with metadata, including an optional scheduled job ID /// Write messages to a session file with metadata, including an optional scheduled job ID
@@ -1071,7 +1071,13 @@ pub async fn persist_messages_with_schedule_id(
messages: &[Message], messages: &[Message],
provider: Option<Arc<dyn Provider>>, provider: Option<Arc<dyn Provider>>,
schedule_id: Option<String>, schedule_id: Option<String>,
save_session: bool,
) -> Result<()> { ) -> Result<()> {
if !save_session {
tracing::debug!("Skipping session persistence (save_session=false)");
return Ok(());
}
// Validate the session file path for security // Validate the session file path for security
let secure_path = get_path(Identifier::Path(session_file.to_path_buf()))?; let secure_path = get_path(Identifier::Path(session_file.to_path_buf()))?;
@@ -1091,7 +1097,13 @@ pub async fn persist_messages_with_schedule_id(
match provider { match provider {
Some(provider) if user_message_count < 4 => { Some(provider) if user_message_count < 4 => {
//generate_description is responsible for writing the messages //generate_description is responsible for writing the messages
generate_description_with_schedule_id(&secure_path, messages, provider, schedule_id) generate_description_with_schedule_id(
&secure_path,
messages,
provider,
schedule_id,
save_session,
)
.await .await
} }
_ => { _ => {
@@ -1102,7 +1114,7 @@ pub async fn persist_messages_with_schedule_id(
metadata.schedule_id = schedule_id; metadata.schedule_id = schedule_id;
} }
// Write the file with metadata and messages // Write the file with metadata and messages
save_messages_with_metadata(&secure_path, &metadata, messages) save_messages_with_metadata(&secure_path, &metadata, messages, save_session)
} }
} }
} }
@@ -1124,7 +1136,13 @@ pub fn save_messages_with_metadata(
session_file: &Path, session_file: &Path,
metadata: &SessionMetadata, metadata: &SessionMetadata,
messages: &[Message], messages: &[Message],
save_session: bool,
) -> Result<()> { ) -> Result<()> {
if !save_session {
tracing::debug!("Skipping session file write (save_session=false)");
return Ok(());
}
use fs2::FileExt; use fs2::FileExt;
// Validate the path for security // Validate the path for security
@@ -1239,7 +1257,7 @@ pub async fn generate_description(
messages: &[Message], messages: &[Message],
provider: Arc<dyn Provider>, provider: Arc<dyn Provider>,
) -> Result<()> { ) -> Result<()> {
generate_description_with_schedule_id(session_file, messages, provider, None).await generate_description_with_schedule_id(session_file, messages, provider, None, true).await
} }
/// Generate a description for the session using the provider, including an optional scheduled job ID /// Generate a description for the session using the provider, including an optional scheduled job ID
@@ -1256,7 +1274,13 @@ pub async fn generate_description_with_schedule_id(
messages: &[Message], messages: &[Message],
provider: Arc<dyn Provider>, provider: Arc<dyn Provider>,
schedule_id: Option<String>, schedule_id: Option<String>,
save_session: bool,
) -> Result<()> { ) -> Result<()> {
if !save_session {
tracing::debug!("Skipping description generation (save_session=false)");
return Ok(());
}
// Validate the path for security // Validate the path for security
let secure_path = get_path(Identifier::Path(session_file.to_path_buf()))?; let secure_path = get_path(Identifier::Path(session_file.to_path_buf()))?;
@@ -1331,7 +1355,7 @@ pub async fn generate_description_with_schedule_id(
} }
// Update the file with the new metadata and existing messages // Update the file with the new metadata and existing messages
save_messages_with_metadata(&secure_path, &metadata, messages) save_messages_with_metadata(&secure_path, &metadata, messages, save_session)
} }
/// Update only the metadata in a session file, preserving all messages /// Update only the metadata in a session file, preserving all messages
@@ -1347,7 +1371,7 @@ pub async fn update_metadata(session_file: &Path, metadata: &SessionMetadata) ->
let messages = read_messages(&secure_path)?; let messages = read_messages(&secure_path)?;
// Rewrite the file with the new metadata and existing messages // Rewrite the file with the new metadata and existing messages
save_messages_with_metadata(&secure_path, metadata, &messages) save_messages_with_metadata(&secure_path, metadata, &messages, true)
} }
#[cfg(test)] #[cfg(test)]
@@ -1662,7 +1686,7 @@ mod tests {
let messages = vec![Message::user().with_text("test")]; let messages = vec![Message::user().with_text("test")];
// Write with special metadata // Write with special metadata
save_messages_with_metadata(&file_path, &metadata, &messages)?; save_messages_with_metadata(&file_path, &metadata, &messages, true)?;
// Read back metadata // Read back metadata
let read_metadata = read_metadata(&file_path)?; let read_metadata = read_metadata(&file_path)?;
@@ -1686,7 +1710,7 @@ mod tests {
// Test deserialization of invalid directory // Test deserialization of invalid directory
let messages = vec![Message::user().with_text("test")]; let messages = vec![Message::user().with_text("test")];
save_messages_with_metadata(&file_path, &metadata, &messages)?; save_messages_with_metadata(&file_path, &metadata, &messages, true)?;
// Modify the file to include invalid directory // Modify the file to include invalid directory
let contents = fs::read_to_string(&file_path)?; let contents = fs::read_to_string(&file_path)?;
@@ -1755,4 +1779,84 @@ mod tests {
let normalized_existing = normalize_path_for_comparison(&test_path); let normalized_existing = normalize_path_for_comparison(&test_path);
assert!(!normalized_existing.as_os_str().is_empty()); assert!(!normalized_existing.as_os_str().is_empty());
} }
#[tokio::test]
async fn test_save_session_parameter() -> Result<()> {
let dir = tempdir()?;
let file_path = dir.path().join("test_save_session.jsonl");
let messages = vec![
Message::user().with_text("Hello"),
Message::assistant().with_text("Hi there"),
];
let metadata = SessionMetadata::default();
// Test with save_session = false - should not create file
save_messages_with_metadata(&file_path, &metadata, &messages, false)?;
assert!(
!file_path.exists(),
"File should not be created when save_session=false"
);
// Test with save_session = true - should create file
save_messages_with_metadata(&file_path, &metadata, &messages, true)?;
assert!(
file_path.exists(),
"File should be created when save_session=true"
);
// Verify content is correct
let read_messages = read_messages(&file_path)?;
assert_eq!(messages.len(), read_messages.len());
Ok(())
}
#[tokio::test]
async fn test_persist_messages_with_save_session_false() -> Result<()> {
let dir = tempdir()?;
let file_path = dir.path().join("test_persist_no_save.jsonl");
let messages = vec![
Message::user().with_text("Test message"),
Message::assistant().with_text("Test response"),
];
// Test persist_messages_with_schedule_id with save_session = false
persist_messages_with_schedule_id(
&file_path,
&messages,
None,
Some("test_schedule".to_string()),
false,
)
.await?;
assert!(
!file_path.exists(),
"File should not be created when save_session=false"
);
// Test persist_messages_with_schedule_id with save_session = true
persist_messages_with_schedule_id(
&file_path,
&messages,
None,
Some("test_schedule".to_string()),
true,
)
.await?;
assert!(
file_path.exists(),
"File should be created when save_session=true"
);
// Verify the schedule_id was set correctly
let metadata = read_metadata(&file_path)?;
assert_eq!(metadata.schedule_id, Some("test_schedule".to_string()));
Ok(())
}
} }

View File

@@ -788,7 +788,7 @@ async fn test_schedule_tool_session_content_action_with_real_session() {
]; ];
// Save the session file // Save the session file
goose::session::storage::save_messages_with_metadata(&session_path, &metadata, &messages) goose::session::storage::save_messages_with_metadata(&session_path, &metadata, &messages, true)
.unwrap(); .unwrap();
// Test the session_content action // Test the session_content action