tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

Add session forking capability (#5882)

Browse Source 
Co-authored-by: Zane Staggs <zane@squareup.com>
This commit is contained in:
Will Pfleger
2026-01-22 12:49:04 -05:00
committed by GitHub
parent f1a36b2bc3
commit 13bdff4bb5
13 changed files with 378 additions and 275 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/cli.rs
+84 -45
View File
@@ -322,15 +322,37 @@ async fn get_or_create_session_id(
let session_manager = SessionManager::instance();
let Some(id) = identifier else {
return if resume {
let resolved_id = if resume {
let Some(id) = identifier else {
let sessions = session_manager.list_sessions().await?;
let session_id = sessions
.first()
.map(|s| s.id.clone())
.ok_or_else(|| anyhow::anyhow!("No session found to resume"))?;
Ok(Some(session_id))
return Ok(Some(session_id));
};
if let Some(session_id) = id.session_id {
session_id
} else if let Some(name) = id.name {
let sessions = session_manager.list_sessions().await?;
sessions
.into_iter()
.find(|s| s.name == name || s.id == name)
.map(|s| s.id)
.ok_or_else(|| anyhow::anyhow!("No session found with name '{}'", name))?
} else if let Some(path) = id.path {
path.file_stem()
.and_then(|s| s.to_str())
.map(|s| s.to_string())
.ok_or_else(|| {
anyhow::anyhow!("Could not extract session ID from path: {:?}", path)
})?
} else {
return Err(anyhow::anyhow!("Invalid identifier"));
}
} else {
let Some(id) = identifier else {
let session = session_manager
.create_session(
std::env::current_dir()?,
@@ -338,58 +360,39 @@ async fn get_or_create_session_id(
SessionType::User,
)
.await?;
Ok(Some(session.id))
return Ok(Some(session.id));
};
};
if let Some(session_id) = id.session_id {
Ok(Some(session_id))
} else if let Some(name) = id.name {
if resume {
let sessions = session_manager.list_sessions().await?;
let session_id = sessions
.into_iter()
.find(|s| s.name == name || s.id == name)
.map(|s| s.id)
.ok_or_else(|| anyhow::anyhow!("No session found with name '{}'", name))?;
Ok(Some(session_id))
} else {
let session = session_manager
.create_session(std::env::current_dir()?, name.clone(), SessionType::User)
.await?;
if id.session_id.is_some() {
return Err(anyhow::anyhow!("Cannot use --session-id without --resume"));
}
let has_user_provided_name = id.name.is_some();
let name = id.name.unwrap_or_else(|| "CLI Session".to_string());
let session = session_manager
.create_session(std::env::current_dir()?, name.clone(), SessionType::User)
.await?;
if has_user_provided_name {
session_manager
.update(&session.id)
.user_provided_name(name)
.apply()
.await?;
Ok(Some(session.id))
}
} else if let Some(path) = id.path {
let session_id = path
.file_stem()
.and_then(|s| s.to_str())
.map(|s| s.to_string())
.ok_or_else(|| anyhow::anyhow!("Could not extract session ID from path: {:?}", path))?;
Ok(Some(session_id))
} else {
let session = session_manager
.create_session(
std::env::current_dir()?,
"CLI Session".to_string(),
SessionType::User,
)
.await?;
Ok(Some(session.id))
}
return Ok(Some(session.id));
};
Ok(Some(resolved_id))
}
async fn lookup_session_id(identifier: Identifier) -> Result<String> {
let session_manager = SessionManager::instance();
if let Some(session_id) = identifier.session_id {
Ok(session_id)
} else if let Some(name) = identifier.name {
let session_manager = SessionManager::instance();
let sessions = session_manager.list_sessions().await?;
sessions
.into_iter()
@@ -722,6 +725,15 @@ enum Command {
)]
resume: bool,
/// Fork a previous session (creates new session with copied history)
#[arg(
long,
requires = "resume",
help = "Fork a previous session (creates new session with copied history)",
long_help = "Create a new session by copying all messages from a previous session. Must be used with --resume. If --name or --session-id is provided, forks that specific session. Otherwise forks the most recently used session."
)]
fork: bool,
/// Show message history when resuming
#[arg(
long,
@@ -1047,6 +1059,7 @@ async fn handle_session_subcommand(command: SessionCommand) -> Result<()> {
async fn handle_interactive_session(
identifier: Option<Identifier>,
resume: bool,
fork: bool,
history: bool,
session_opts: SessionOptions,
extension_opts: ExtensionOptions,
@@ -1056,7 +1069,13 @@ async fn handle_interactive_session(
}
let session_start = std::time::Instant::now();
let session_type = if resume { "resumed" } else { "new" };
let session_type = if fork {
"forked"
} else if resume {
"resumed"
} else {
"new"
};
tracing::info!(
counter.goose.session_starts = 1,
@@ -1076,11 +1095,21 @@ async fn handle_interactive_session(
}
}
let session_id = get_or_create_session_id(identifier, resume, false).await?;
let mut session_id = get_or_create_session_id(identifier, resume, false).await?;
if fork {
if let Some(id) = session_id {
let session_manager = SessionManager::instance();
let original = session_manager.get_session(&id, false).await?;
let copied = session_manager.copy_session(&id, original.name).await?;
session_id = Some(copied.id);
}
}
let mut session: crate::CliSession = build_session(SessionBuilderConfig {
session_id,
resume,
fork,
no_session: false,
extensions: extension_opts.extensions,
streamable_http_extensions: extension_opts.streamable_http_extensions,
@@ -1099,7 +1128,7 @@ async fn handle_interactive_session(
})
.await;
if resume && history {
if (resume || fork) && history {
session.render_message_history();
}
@@ -1283,6 +1312,7 @@ async fn handle_run_command(
let mut session = build_session(SessionBuilderConfig {
session_id,
resume: run_behavior.resume,
fork: false,
no_session: run_behavior.no_session,
extensions: extension_opts.extensions,
streamable_http_extensions: extension_opts.streamable_http_extensions,
@@ -1407,6 +1437,7 @@ async fn handle_default_session() -> Result<()> {
let mut session = build_session(SessionBuilderConfig {
session_id,
resume: false,
fork: false,
no_session: false,
extensions: Vec::new(),
streamable_http_extensions: Vec::new(),
@@ -1458,12 +1489,20 @@ pub async fn cli() -> anyhow::Result<()> {
command: None,
identifier,
resume,
fork,
history,
session_opts,
extension_opts,
}) => {
handle_interactive_session(identifier, resume, history, session_opts, extension_opts)
.await
handle_interactive_session(
identifier,
resume,
fork,
history,
session_opts,
extension_opts,
)
.await
}
Some(Command::Project {}) => {
handle_project_default()?;
crates/goose-cli/src/commands/bench.rs
+1
View File
@@ -37,6 +37,7 @@ pub async fn agent_generator(
let base_session = build_session(SessionBuilderConfig {
session_id: Some(session_id),
resume: false,
fork: false,
no_session: false,
extensions: requirements.external,
streamable_http_extensions: requirements.streamable_http,
crates/goose-cli/src/session/builder.rs
+5
View File
@@ -82,6 +82,8 @@ pub struct SessionBuilderConfig {
pub session_id: Option<String>,
/// Whether to resume an existing session
pub resume: bool,
/// Whether to fork an existing session (creates a copy of the original/existing session then resumes the copy)
pub fork: bool,
/// Whether to run without a session file
pub no_session: bool,
/// List of stdio extension commands to add
@@ -121,6 +123,7 @@ impl Default for SessionBuilderConfig {
SessionBuilderConfig {
session_id: None,
resume: false,
fork: false,
no_session: false,
extensions: Vec::new(),
streamable_http_extensions: Vec::new(),
@@ -659,6 +662,7 @@ mod tests {
let config = SessionBuilderConfig {
session_id: None,
resume: false,
fork: false,
no_session: false,
extensions: vec!["echo test".to_string()],
streamable_http_extensions: vec!["http://localhost:8080/mcp".to_string()],
@@ -705,6 +709,7 @@ mod tests {
assert!(config.scheduled_job_id.is_none());
assert!(!config.interactive);
assert!(!config.quiet);
assert!(!config.fork);
}
#[tokio::test]
crates/goose-server/src/openapi.rs
+3 -4
View File
@@ -380,7 +380,7 @@ derive_utoipa!(Icon as IconSchema);
super::routes::session::export_session,
super::routes::session::import_session,
super::routes::session::update_session_user_recipe_values,
super::routes::session::edit_message,
super::routes::session::fork_session,
super::routes::session::get_session_extensions,
super::routes::schedule::create_schedule,
super::routes::schedule::list_schedules,
@@ -442,9 +442,8 @@ derive_utoipa!(Icon as IconSchema);
super::routes::session::UpdateSessionNameRequest,
super::routes::session::UpdateSessionUserRecipeValuesRequest,
super::routes::session::UpdateSessionUserRecipeValuesResponse,
super::routes::session::EditType,
super::routes::session::EditMessageRequest,
super::routes::session::EditMessageResponse,
super::routes::session::ForkRequest,
super::routes::session::ForkResponse,
super::routes::session::SessionExtensionsResponse,
Message,
MessageContent,
crates/goose-server/src/routes/session.rs
+77 -64
View File
@@ -51,28 +51,17 @@ pub struct ImportSessionRequest {
json: String,
}
#[derive(Debug, Deserialize, ToSchema)]
#[serde(rename_all = "lowercase")]
pub enum EditType {
Fork,
Edit,
}
#[derive(Deserialize, ToSchema)]
#[serde(rename_all = "camelCase")]
pub struct EditMessageRequest {
timestamp: i64,
#[serde(default = "default_edit_type")]
edit_type: EditType,
}
fn default_edit_type() -> EditType {
EditType::Fork
pub struct ForkRequest {
timestamp: Option<i64>,
truncate: bool,
copy: bool,
}
#[derive(Serialize, ToSchema)]
#[serde(rename_all = "camelCase")]
pub struct EditMessageResponse {
pub struct ForkResponse {
session_id: String,
}
@@ -369,16 +358,16 @@ async fn import_session(
#[utoipa::path(
post,
path = "/sessions/{session_id}/edit_message",
request_body = EditMessageRequest,
path = "/sessions/{session_id}/fork",
request_body = ForkRequest,
params(
("session_id" = String, Path, description = "Unique identifier for the session")
),
responses(
(status = 200, description = "Session prepared for editing - frontend should submit the edited message", body = EditMessageResponse),
(status = 400, description = "Bad request - Invalid message timestamp"),
(status = 200, description = "Session forked successfully", body = ForkResponse),
(status = 400, description = "Bad request - truncate=true requires timestamp"),
(status = 401, description = "Unauthorized - Invalid or missing API key"),
(status = 404, description = "Session or message not found"),
(status = 404, description = "Session not found"),
(status = 500, description = "Internal server error")
),
security(
@@ -386,51 +375,75 @@ async fn import_session(
),
tag = "Session Management"
)]
async fn edit_message(
async fn fork_session(
State(state): State<Arc<AppState>>,
Path(session_id): Path<String>,
Json(request): Json<EditMessageRequest>,
) -> Result<Json<EditMessageResponse>, StatusCode> {
let manager = state.session_manager();
match request.edit_type {
EditType::Fork => {
let new_session = manager
.copy_session(&session_id, "(edited)".to_string())
.await
.map_err(|e| {
tracing::error!("Failed to copy session: {}", e);
goose::posthog::emit_error("session_copy_failed", &e.to_string());
StatusCode::INTERNAL_SERVER_ERROR
})?;
manager
.truncate_conversation(&new_session.id, request.timestamp)
.await
.map_err(|e| {
tracing::error!("Failed to truncate conversation: {}", e);
goose::posthog::emit_error("session_truncate_failed", &e.to_string());
StatusCode::INTERNAL_SERVER_ERROR
})?;
Ok(Json(EditMessageResponse {
session_id: new_session.id,
}))
}
EditType::Edit => {
manager
.truncate_conversation(&session_id, request.timestamp)
.await
.map_err(|e| {
tracing::error!("Failed to truncate conversation: {}", e);
goose::posthog::emit_error("session_truncate_failed", &e.to_string());
StatusCode::INTERNAL_SERVER_ERROR
})?;
Ok(Json(EditMessageResponse {
session_id: session_id.clone(),
}))
}
Json(request): Json<ForkRequest>,
) -> Result<Json<ForkResponse>, ErrorResponse> {
if request.truncate && request.timestamp.is_none() {
return Err(ErrorResponse {
message: "truncate=true requires a timestamp".to_string(),
status: StatusCode::BAD_REQUEST,
});
}
let session_manager = state.session_manager();
let target_session_id = if request.copy {
let original = session_manager
.get_session(&session_id, false)
.await
.map_err(|e| {
tracing::error!("Failed to get session: {}", e);
goose::posthog::emit_error("session_get_failed", &e.to_string());
ErrorResponse {
message: if e.to_string().contains("not found") {
format!("Session {} not found", session_id)
} else {
format!("Failed to get session: {}", e)
},
status: if e.to_string().contains("not found") {
StatusCode::NOT_FOUND
} else {
StatusCode::INTERNAL_SERVER_ERROR
},
}
})?;
let copied = session_manager
.copy_session(&session_id, original.name)
.await
.map_err(|e| {
tracing::error!("Failed to copy session: {}", e);
goose::posthog::emit_error("session_copy_failed", &e.to_string());
ErrorResponse {
message: format!("Failed to copy session: {}", e),
status: StatusCode::INTERNAL_SERVER_ERROR,
}
})?;
copied.id
} else {
session_id.clone()
};
if request.truncate {
session_manager
.truncate_conversation(&target_session_id, request.timestamp.unwrap_or(0))
.await
.map_err(|e| {
tracing::error!("Failed to truncate conversation: {}", e);
goose::posthog::emit_error("session_truncate_failed", &e.to_string());
ErrorResponse {
message: format!("Failed to truncate conversation: {}", e),
status: StatusCode::INTERNAL_SERVER_ERROR,
}
})?;
}
Ok(Json(ForkResponse {
session_id: target_session_id,
}))
}
#[derive(Serialize, ToSchema)]
@@ -487,7 +500,7 @@ pub fn routes(state: Arc<AppState>) -> Router {
"/sessions/{session_id}/user_recipe_values",
put(update_session_user_recipe_values),
)
.route("/sessions/{session_id}/edit_message", post(edit_message))
.route("/sessions/{session_id}/fork", post(fork_session))
.route(
"/sessions/{session_id}/extensions",
get(get_session_extensions),
crates/goose/src/session/session_manager.rs
+13 -4
View File
@@ -1285,14 +1285,23 @@ impl SessionStorage {
)
.await?;
session_manager
let mut builder = session_manager
.update(&new_session.id)
.extension_data(original_session.extension_data)
.schedule_id(original_session.schedule_id)
.recipe(original_session.recipe)
.user_recipe_values(original_session.user_recipe_values)
.apply()
.await?;
.user_recipe_values(original_session.user_recipe_values);
// Preserve provider and model config from original session
if let Some(provider_name) = original_session.provider_name {
builder = builder.provider_name(provider_name);
}
if let Some(model_config) = original_session.model_config {
builder = builder.model_config(model_config);
}
builder.apply().await?;
if let Some(conversation) = original_session.conversation {
self.replace_conversation(&new_session.id, &conversation)