Skip to content

Commit 85f284b

Browse files
zanesqzanesq
authored
Session file security updates (aaif-goose#3071)
1 parent 8a32128 commit 85f284b

10 files changed

Lines changed: 436 additions & 168 deletions

File tree

‎crates/goose-cli/src/commands/session.rs‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -175,7 +175,12 @@ pub fn handle_session_list(verbose: bool, format: String, ascending: bool) -> Re
175175
/// without creating an Agent or prompting about working directories.
176176
pub fn handle_session_export(identifier: Identifier, output_path: Option<PathBuf>) -> Result<()> {
177177
// Get the session file path
178-
let session_file_path = goose::session::get_path(identifier.clone());
178+
let session_file_path = match goose::session::get_path(identifier.clone()) {
179+
Ok(path) => path,
180+
Err(e) => {
181+
return Err(anyhow::anyhow!("Invalid session identifier: {}", e));
182+
}
183+
};
179184

180185
if !session_file_path.exists() {
181186
return Err(anyhow::anyhow!(

‎crates/goose-cli/src/commands/web.rs‎

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -250,7 +250,14 @@ async fn list_sessions() -> Json<serde_json::Value> {
250250
async fn get_session(
251251
axum::extract::Path(session_id): axum::extract::Path<String>,
252252
) -> Json<serde_json::Value> {
253-
let session_file = session::get_path(session::Identifier::Name(session_id));
253+
let session_file = match session::get_path(session::Identifier::Name(session_id)) {
254+
Ok(path) => path,
255+
Err(e) => {
256+
return Json(serde_json::json!({
257+
"error": format!("Invalid session ID: {}", e)
258+
}));
259+
}
260+
};
254261

255262
match session::read_messages(&session_file) {
256263
Ok(messages) => {
@@ -288,8 +295,15 @@ async fn handle_socket(socket: WebSocket, state: AppState) {
288295
..
289296
}) => {
290297
// Get session file path from session_id
291-
let session_file =
292-
session::get_path(session::Identifier::Name(session_id.clone()));
298+
let session_file = match session::get_path(session::Identifier::Name(
299+
session_id.clone(),
300+
)) {
301+
Ok(path) => path,
302+
Err(e) => {
303+
tracing::error!("Failed to get session path: {}", e);
304+
continue;
305+
}
306+
};
293307

294308
// Get or create session in memory (for fast access during processing)
295309
let session_messages = {

‎crates/goose-cli/src/session/builder.rs‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -234,7 +234,13 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> Session {
234234
}
235235
} else if session_config.resume {
236236
if let Some(identifier) = session_config.identifier {
237-
let session_file = session::get_path(identifier);
237+
let session_file = match session::get_path(identifier) {
238+
Ok(path) => path,
239+
Err(e) => {
240+
output::render_error(&format!("Invalid session identifier: {}", e));
241+
process::exit(1);
242+
}
243+
};
238244
if !session_file.exists() {
239245
output::render_error(&format!(
240246
"Cannot resume session {} - no such session exists",
@@ -262,7 +268,13 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> Session {
262268
};
263269

264270
// Just get the path - file will be created when needed
265-
session::get_path(id)
271+
match session::get_path(id) {
272+
Ok(path) => path,
273+
Err(e) => {
274+
output::render_error(&format!("Failed to create session path: {}", e));
275+
process::exit(1);
276+
}
277+
}
266278
};
267279

268280
if session_config.resume && !session_config.no_session {

‎crates/goose-server/src/routes/reply.rs‎

Lines changed: 25 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -210,7 +210,20 @@ async fn handler(
210210
};
211211

212212
let mut all_messages = messages.clone();
213-
let session_path = session::get_path(session::Identifier::Name(session_id.clone()));
213+
let session_path = match session::get_path(session::Identifier::Name(session_id.clone())) {
214+
Ok(path) => path,
215+
Err(e) => {
216+
tracing::error!("Failed to get session path: {}", e);
217+
let _ = stream_event(
218+
MessageEvent::Error {
219+
error: format!("Failed to get session path: {}", e),
220+
},
221+
&tx,
222+
)
223+
.await;
224+
return;
225+
}
226+
};
214227

215228
loop {
216229
tokio::select! {
@@ -390,13 +403,21 @@ async fn ask_handler(
390403
all_messages.push(response_message);
391404
}
392405

393-
let session_path = session::get_path(session::Identifier::Name(session_id.clone()));
406+
let session_path = match session::get_path(session::Identifier::Name(session_id.clone())) {
407+
Ok(path) => path,
408+
Err(e) => {
409+
tracing::error!("Failed to get session path: {}", e);
410+
return Err(StatusCode::INTERNAL_SERVER_ERROR);
411+
}
412+
};
394413

395-
let session_path = session_path.clone();
414+
let session_path_clone = session_path.clone();
396415
let messages = all_messages.clone();
397416
let provider = Arc::clone(provider.as_ref().unwrap());
398417
tokio::spawn(async move {
399-
if let Err(e) = session::persist_messages(&session_path, &messages, Some(provider)).await {
418+
if let Err(e) =
419+
session::persist_messages(&session_path_clone, &messages, Some(provider)).await
420+
{
400421
tracing::error!("Failed to store session history: {:?}", e);
401422
}
402423
});

‎crates/goose-server/src/routes/session.rs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,10 @@ async fn get_session_history(
8484
) -> Result<Json<SessionHistoryResponse>, StatusCode> {
8585
verify_secret_key(&headers, &state)?;
8686

87-
let session_path = session::get_path(session::Identifier::Name(session_id.clone()));
87+
let session_path = match session::get_path(session::Identifier::Name(session_id.clone())) {
88+
Ok(path) => path,
89+
Err(_) => return Err(StatusCode::BAD_REQUEST),
90+
};
8891

8992
// Read metadata
9093
let metadata = session::read_metadata(&session_path).map_err(|_| StatusCode::NOT_FOUND)?;

‎crates/goose/src/agents/reply_parts.rs‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -222,7 +222,12 @@ impl Agent {
222222
usage: &crate::providers::base::ProviderUsage,
223223
messages_length: usize,
224224
) -> Result<()> {
225-
let session_file_path = session::storage::get_path(session_config.id.clone());
225+
let session_file_path = match session::storage::get_path(session_config.id.clone()) {
226+
Ok(path) => path,
227+
Err(e) => {
228+
return Err(anyhow::anyhow!("Failed to get session file path: {}", e));
229+
}
230+
};
226231
let mut metadata = session::storage::read_metadata(&session_file_path)?;
227232

228233
metadata.schedule_id = session_config.schedule_id.clone();

‎crates/goose/src/agents/schedule_tool.rs‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -372,9 +372,17 @@ impl Agent {
372372
})?;
373373

374374
// Get the session file path
375-
let session_path = crate::session::storage::get_path(
375+
let session_path = match crate::session::storage::get_path(
376376
crate::session::storage::Identifier::Name(session_id.to_string()),
377-
);
377+
) {
378+
Ok(path) => path,
379+
Err(e) => {
380+
return Err(ToolError::ExecutionError(format!(
381+
"Invalid session ID '{}': {}",
382+
session_id, e
383+
)));
384+
}
385+
};
378386

379387
// Check if session file exists
380388
if !session_path.exists() {

‎crates/goose/src/scheduler.rs‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1138,9 +1138,17 @@ async fn run_scheduled_job_internal(
11381138
}
11391139
}
11401140

1141-
let session_file_path = crate::session::storage::get_path(
1141+
let session_file_path = match crate::session::storage::get_path(
11421142
crate::session::storage::Identifier::Name(session_id_for_return.clone()),
1143-
);
1143+
) {
1144+
Ok(path) => path,
1145+
Err(e) => {
1146+
return Err(JobExecutionError {
1147+
job_id: job.id.clone(),
1148+
error: format!("Failed to get session file path: {}", e),
1149+
});
1150+
}
1151+
};
11441152

11451153
if let Some(prompt_text) = recipe.prompt {
11461154
let mut all_session_messages: Vec<Message> =

0 commit comments

Comments
 (0)