Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion crates/goose-cli/src/session/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -120,7 +120,7 @@ async fn offer_extension_debugging_help(
std::env::temp_dir().join(format!("goose_debug_extension_{}.jsonl", extension_name));

// 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
println!("{}", style("Analyzing the extension failure...").yellow());
Expand Down Expand Up @@ -366,6 +366,7 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> Session {
session_file.clone(),
session_config.debug,
session_config.scheduled_job_id.clone(),
!session_config.no_session, // save_session is the inverse of no_session
);

// Add extensions if provided
Expand Down
28 changes: 21 additions & 7 deletions crates/goose-cli/src/session/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@ pub struct Session {
debug: bool, // New field for debug mode
run_mode: RunMode,
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
Expand Down Expand Up @@ -113,13 +114,19 @@ impl Session {
session_file: PathBuf,
debug: bool,
scheduled_job_id: Option<String>,
save_session: bool,
) -> Self {
let messages = match session::read_messages(&session_file) {
Ok(msgs) => msgs,
Err(e) => {
eprintln!("Warning: Failed to load message history: {}", e);
Vec::new()
let messages = if save_session {
match session::read_messages(&session_file) {
Ok(msgs) => msgs,
Err(e) => {
eprintln!("Warning: Failed to load message history: {}", e);
Vec::new()
}
}
} else {
// Don't try to read messages if we're not saving sessions
Vec::new()
};

Session {
Expand All @@ -130,6 +137,7 @@ impl Session {
debug,
run_mode: RunMode::Normal,
scheduled_job_id,
save_session,
}
}

Expand Down Expand Up @@ -319,6 +327,7 @@ impl Session {
&self.messages,
Some(provider),
self.scheduled_job_id.clone(),
self.save_session,
)
.await?;

Expand Down Expand Up @@ -431,6 +440,7 @@ impl Session {
&self.messages,
Some(provider),
self.scheduled_job_id.clone(),
self.save_session,
)
.await?;

Expand Down Expand Up @@ -619,6 +629,7 @@ impl Session {
&self.messages,
Some(provider),
self.scheduled_job_id.clone(),
self.save_session,
)
.await?;

Expand Down Expand Up @@ -792,7 +803,7 @@ impl Session {
Err(ToolError::ExecutionError("Tool call cancelled by user".to_string()))
));
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);
break;
Expand Down Expand Up @@ -889,7 +900,7 @@ impl Session {
self.messages.push(message.clone());

// 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()};
let _ = progress_bars.hide();
Expand Down Expand Up @@ -1093,6 +1104,7 @@ impl Session {
&self.messages,
None,
self.scheduled_job_id.clone(),
self.save_session,
)
.await?;

Expand All @@ -1108,6 +1120,7 @@ impl Session {
&self.messages,
None,
self.scheduled_job_id.clone(),
self.save_session,
)
.await?;

Expand All @@ -1128,6 +1141,7 @@ impl Session {
&self.messages,
None,
self.scheduled_job_id.clone(),
self.save_session,
)
.await?;

Expand Down
11 changes: 8 additions & 3 deletions crates/goose/src/scheduler.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1248,6 +1248,7 @@ async fn run_scheduled_job_internal(
&session_file_path,
&updated_metadata,
&all_session_messages,
true,
) {
tracing::error!(
"[Job {}] Failed to persist final messages: {}",
Expand Down Expand Up @@ -1278,6 +1279,7 @@ async fn run_scheduled_job_internal(
&session_file_path,
&fallback_metadata,
&all_session_messages,
true,
) {
tracing::error!("[Job {}] Failed to persist final messages with fallback metadata: {}", job.id, e_fb);
}
Expand All @@ -1304,9 +1306,12 @@ async fn run_scheduled_job_internal(
message_count: 0,
..Default::default()
};
if let Err(e) =
crate::session::storage::save_messages_with_metadata(&session_file_path, &metadata, &[])
{
if let Err(e) = crate::session::storage::save_messages_with_metadata(
&session_file_path,
&metadata,
&[],
true,
) {
tracing::error!(
"[Job {}] Failed to persist metadata for empty job: {}",
job.id,
Expand Down
122 changes: 113 additions & 9 deletions crates/goose/src/session/storage.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1054,7 +1054,7 @@ pub async fn persist_messages(
messages: &[Message],
provider: Option<Arc<dyn Provider>>,
) -> 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
Expand All @@ -1071,7 +1071,13 @@ pub async fn persist_messages_with_schedule_id(
messages: &[Message],
provider: Option<Arc<dyn Provider>>,
schedule_id: Option<String>,
save_session: bool,
) -> Result<()> {
if !save_session {
tracing::debug!("Skipping session persistence (save_session=false)");
return Ok(());
}

// Validate the session file path for security
let secure_path = get_path(Identifier::Path(session_file.to_path_buf()))?;

Expand All @@ -1091,8 +1097,14 @@ pub async fn persist_messages_with_schedule_id(
match provider {
Some(provider) if user_message_count < 4 => {
//generate_description is responsible for writing the messages
generate_description_with_schedule_id(&secure_path, messages, provider, schedule_id)
.await
generate_description_with_schedule_id(
&secure_path,
messages,
provider,
schedule_id,
save_session,
)
.await
}
_ => {
// Read existing metadata
Expand All @@ -1102,7 +1114,7 @@ pub async fn persist_messages_with_schedule_id(
metadata.schedule_id = schedule_id;
}
// 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)
}
}
}
Expand All @@ -1124,7 +1136,13 @@ pub fn save_messages_with_metadata(
session_file: &Path,
metadata: &SessionMetadata,
messages: &[Message],
save_session: bool,
) -> Result<()> {
if !save_session {
tracing::debug!("Skipping session file write (save_session=false)");
return Ok(());
}

use fs2::FileExt;

// Validate the path for security
Expand Down Expand Up @@ -1239,7 +1257,7 @@ pub async fn generate_description(
messages: &[Message],
provider: Arc<dyn Provider>,
) -> 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
Expand All @@ -1256,7 +1274,13 @@ pub async fn generate_description_with_schedule_id(
messages: &[Message],
provider: Arc<dyn Provider>,
schedule_id: Option<String>,
save_session: bool,
) -> Result<()> {
if !save_session {
tracing::debug!("Skipping description generation (save_session=false)");
return Ok(());
}

// Validate the path for security
let secure_path = get_path(Identifier::Path(session_file.to_path_buf()))?;

Expand Down Expand Up @@ -1331,7 +1355,7 @@ pub async fn generate_description_with_schedule_id(
}

// 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
Expand All @@ -1347,7 +1371,7 @@ pub async fn update_metadata(session_file: &Path, metadata: &SessionMetadata) ->
let messages = read_messages(&secure_path)?;

// 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)]
Expand Down Expand Up @@ -1662,7 +1686,7 @@ mod tests {
let messages = vec![Message::user().with_text("test")];

// Write with special metadata
save_messages_with_metadata(&file_path, &metadata, &messages)?;
save_messages_with_metadata(&file_path, &metadata, &messages, true)?;

// Read back metadata
let read_metadata = read_metadata(&file_path)?;
Expand All @@ -1686,7 +1710,7 @@ mod tests {

// Test deserialization of invalid directory
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
let contents = fs::read_to_string(&file_path)?;
Expand Down Expand Up @@ -1755,4 +1779,84 @@ mod tests {
let normalized_existing = normalize_path_for_comparison(&test_path);
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(())
}
}
2 changes: 1 addition & 1 deletion crates/goose/tests/private_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -788,7 +788,7 @@ async fn test_schedule_tool_session_content_action_with_real_session() {
];

// 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();

// Test the session_content action
Expand Down
Loading