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

Sessions required (#5548)

Browse Source 
Co-authored-by: Douwe Osinga <douwe@squareup.com>
This commit is contained in:
Douwe Osinga
2025-11-03 21:04:44 -05:00
committed by GitHub
parent 86c3e42e43
commit 5b93ee587f
31 changed files with 956 additions and 2678 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/cli.rs
+21 -10
View File
@@ -19,6 +19,7 @@ use crate::commands::session::{handle_session_list, handle_session_remove};
use crate::recipes::extract_from_cli::extract_recipe_info_from_cli;
use crate::recipes::recipe::{explain_recipe, render_recipe_as_yaml};
use crate::session::{build_session, SessionBuilderConfig, SessionSettings};
use goose::session::session_manager::SessionType;
use goose::session::SessionManager;
use goose_bench::bench_config::BenchRunConfig;
use goose_bench::runners::bench_runner::BenchRunner;
@@ -86,9 +87,12 @@ async fn get_or_create_session_id(
.ok_or_else(|| anyhow::anyhow!("No session found to resume"))?;
Ok(Some(session_id))
} else {
let session =
SessionManager::create_session(std::env::current_dir()?, "CLI Session".to_string())
.await?;
let session = SessionManager::create_session(
std::env::current_dir()?,
"CLI Session".to_string(),
SessionType::User,
)
.await?;
Ok(Some(session.id))
};
};
@@ -105,8 +109,12 @@ async fn get_or_create_session_id(
.ok_or_else(|| anyhow::anyhow!("No session found with name '{}'", name))?;
Ok(Some(session_id))
} else {
let session =
SessionManager::create_session(std::env::current_dir()?, name.clone()).await?;
let session = SessionManager::create_session(
std::env::current_dir()?,
name.clone(),
SessionType::User,
)
.await?;
SessionManager::update_session(&session.id)
.user_provided_name(name)
@@ -123,9 +131,12 @@ async fn get_or_create_session_id(
.ok_or_else(|| anyhow::anyhow!("Could not extract session ID from path: {:?}", path))?;
Ok(Some(session_id))
} else {
let session =
SessionManager::create_session(std::env::current_dir()?, "CLI Session".to_string())
.await?;
let session = SessionManager::create_session(
std::env::current_dir()?,
"CLI Session".to_string(),
SessionType::User,
)
.await?;
Ok(Some(session.id))
}
}
@@ -977,7 +988,7 @@ pub async fn cli() -> anyhow::Result<()> {
let exit_type = if result.is_ok() { "normal" } else { "error" };
let (total_tokens, message_count) = session
.get_metadata()
.get_session()
.await
.map(|m| (m.total_tokens.unwrap_or(0), m.message_count))
.unwrap_or((0, 0));
@@ -1198,7 +1209,7 @@ pub async fn cli() -> anyhow::Result<()> {
let exit_type = if result.is_ok() { "normal" } else { "error" };
let (total_tokens, message_count) = session
.get_metadata()
.get_session()
.await
.map(|m| (m.total_tokens.unwrap_or(0), m.message_count))
.unwrap_or((0, 0));
crates/goose-cli/src/commands/acp.rs
+22 -27
View File
@@ -3,11 +3,13 @@ use agent_client_protocol::{
ToolCallContent,
};
use anyhow::Result;
use goose::agents::Agent;
use goose::agents::{Agent, SessionConfig};
use goose::config::{get_all_extensions, Config};
use goose::conversation::message::{Message, MessageContent};
use goose::conversation::Conversation;
use goose::providers::create;
use goose::session::session_manager::SessionType;
use goose::session::SessionManager;
use rmcp::model::{RawContent, ResourceContents};
use std::collections::{HashMap, HashSet};
use std::fs;
@@ -19,17 +21,15 @@ use tokio_util::sync::CancellationToken;
use tracing::{error, info, warn};
use url::Url;
/// Represents a single goose session for ACP
struct GooseSession {
struct GooseAcpSession {
messages: Conversation,
tool_call_ids: HashMap<String, String>, // Maps internal tool IDs to ACP tool call IDs
cancel_token: Option<CancellationToken>, // Active cancellation token for prompt processing
}
/// goose ACP Agent implementation that connects to real goose agents
struct GooseAcpAgent {
session_update_tx: mpsc::UnboundedSender<(acp::SessionNotification, oneshot::Sender<()>)>,
sessions: Arc<Mutex<HashMap<String, GooseSession>>>,
session_update_tx: mpsc::UnboundedSender<(SessionNotification, oneshot::Sender<()>)>,
sessions: Arc<Mutex<HashMap<String, GooseAcpSession>>>,
agent: Agent, // Shared agent instance
}
@@ -97,7 +97,6 @@ impl GooseAcpAgent {
async fn new(
session_update_tx: mpsc::UnboundedSender<(acp::SessionNotification, oneshot::Sender<()>)>,
) -> Result<Self> {
// Load config and create provider
let config = Config::global();
let provider_name: String = config
@@ -217,7 +216,7 @@ impl GooseAcpAgent {
&self,
content_item: &MessageContent,
session_id: &acp::SessionId,
session: &mut GooseSession,
session: &mut GooseAcpSession,
) -> Result<(), acp::Error> {
match content_item {
MessageContent::Text(text) => {
@@ -273,7 +272,7 @@ impl GooseAcpAgent {
&self,
tool_request: &goose::conversation::message::ToolRequest,
session_id: &acp::SessionId,
session: &mut GooseSession,
session: &mut GooseAcpSession,
) -> Result<(), acp::Error> {
// Generate ACP tool call ID and track mapping
let acp_tool_id = format!("tool_{}", uuid::Uuid::new_v4());
@@ -341,7 +340,7 @@ impl GooseAcpAgent {
&self,
tool_response: &goose::conversation::message::ToolResponse,
session_id: &acp::SessionId,
session: &mut GooseSession,
session: &mut GooseAcpSession,
) -> Result<(), acp::Error> {
// Look up the ACP tool call ID
if let Some(acp_tool_id) = session.tool_call_ids.get(&tool_response.id) {
@@ -496,7 +495,7 @@ impl acp::Agent for GooseAcpAgent {
// Generate a unique session ID
let session_id = uuid::Uuid::new_v4().to_string();
let session = GooseSession {
let session = GooseAcpSession {
messages: Conversation::new_unvalidated(Vec::new()),
tool_call_ids: HashMap::new(),
cancel_token: None,
@@ -544,30 +543,26 @@ impl acp::Agent for GooseAcpAgent {
// Create and store cancellation token for this prompt
let cancel_token = CancellationToken::new();
// Convert ACP prompt to Goose message
let user_message = self.convert_acp_prompt_to_message(args.prompt);
// Prepare for agent reply
let messages = {
let mut sessions = self.sessions.lock().await;
let session = sessions
.get_mut(&session_id)
.ok_or_else(acp::Error::invalid_params)?;
let session = SessionManager::create_session(
std::env::current_dir().unwrap_or_default(),
"ACP Session".to_string(),
SessionType::Hidden,
)
.await?;
// Add message to conversation
session.messages.push(user_message);
// Store cancellation token
session.cancel_token = Some(cancel_token.clone());
// Clone what we need for the reply call
session.messages.clone()
let session_config = SessionConfig {
id: session.id.clone(),
schedule_id: None,
max_turns: None,
retry_config: None,
};
// Get agent's reply through the Goose agent
let mut stream = self
.agent
.reply(messages, None, Some(cancel_token.clone()))
.reply(user_message, session_config, Some(cancel_token.clone()))
.await
.map_err(|e| {
error!("Error getting agent reply: {}", e);
crates/goose-cli/src/commands/bench.rs
+1 -3
View File
@@ -26,9 +26,7 @@ impl BenchBaseSession for CliSession {
}
fn get_session_id(&self) -> anyhow::Result<String> {
self.session_id()
.cloned()
.ok_or_else(|| anyhow::anyhow!("No session ID available"))
Ok(self.session_id().to_string())
}
}
pub async fn agent_generator(
crates/goose-cli/src/commands/schedule.rs
-1
View File
@@ -98,7 +98,6 @@ pub async fn handle_schedule_add(
paused: false,
current_session_id: None,
process_start_time: None,
execution_mode: Some("background".to_string()), // Default to background for CLI
};
let scheduler_storage_path =
crates/goose-cli/src/commands/web.rs
+4 -7
View File
@@ -15,6 +15,7 @@ use base64::Engine;
use futures::{sink::SinkExt, stream::StreamExt};
use goose::agents::{Agent, AgentEvent};
use goose::conversation::message::Message as GooseMessage;
use goose::session::session_manager::SessionType;
use goose::session::SessionManager;
use serde::{Deserialize, Serialize};
use serde_json::Value;
@@ -226,6 +227,7 @@ async fn serve_index() -> Result<Redirect, (http::StatusCode, String)> {
let session = SessionManager::create_session(
std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from(".")),
"Web session".to_string(),
SessionType::User,
)
.await
.map_err(|err| (http::StatusCode::INTERNAL_SERVER_ERROR, err.to_string()))?;
@@ -467,21 +469,16 @@ async fn process_message_streaming(
let session = SessionManager::get_session(&session_id, true).await?;
let mut messages = session.conversation.unwrap_or_default();
messages.push(user_message);
messages.push(user_message.clone());
let session_config = SessionConfig {
id: session.id.clone(),
working_dir: session.working_dir,
schedule_id: None,
execution_mode: None,
max_turns: None,
retry_config: None,
};
match agent
.reply(messages.clone(), Some(session_config), None)
.await
{
match agent.reply(user_message, session_config, None).await {
Ok(mut stream) => {
while let Some(result) = stream.next().await {
match result {
crates/goose-cli/src/scenario_tests/scenario_runner.rs
+13 -6
View File
@@ -9,8 +9,10 @@ use anyhow::Result;
use goose::agents::Agent;
use goose::model::ModelConfig;
use goose::providers::{create, testprovider::TestProvider};
use goose::session::session_manager::SessionType;
use goose::session::SessionManager;
use std::collections::{HashMap, HashSet};
use std::path::Path;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio_util::sync::CancellationToken;
@@ -190,7 +192,6 @@ where
)
};
// Generate messages using the provider
let messages = vec![message_generator(&*provider_arc)];
let mock_client = weather_client();
@@ -218,11 +219,17 @@ where
.update_provider(provider_arc as Arc<dyn goose::providers::base::Provider>)
.await?;
let mut session = CliSession::new(agent, None, false, None, None, None, None).await;
let session = SessionManager::create_session(
PathBuf::default(),
"scenario-runner".to_string(),
SessionType::Hidden,
)
.await?;
let mut cli_session = CliSession::new(agent, session.id, false, None, None, None, None).await;
let mut error = None;
for message in &messages {
if let Err(e) = session
if let Err(e) = cli_session
.process_message(message.clone(), CancellationToken::default())
.await
{
@@ -230,7 +237,7 @@ where
break;
}
}
let updated_messages = session.message_history();
let updated_messages = cli_session.message_history();
if let Some(ref err_msg) = error {
if err_msg.contains("No recorded response found") {
@@ -249,7 +256,7 @@ where
validator(&result)?;
drop(session);
drop(cli_session);
if let Some(provider) = provider_for_saving {
if result.error.is_none() {
crates/goose-cli/src/session/builder.rs
+68 -63
View File
@@ -11,6 +11,7 @@ use goose::providers::create;
use goose::recipe::{Response, SubRecipe};
use goose::agents::extension::PlatformExtensionContext;
use goose::session::session_manager::SessionType;
use goose::session::SessionManager;
use goose::session::{EnabledExtensionsState, ExtensionState};
use rustyline::EditMode;
@@ -25,7 +26,7 @@ use tokio::task::JoinSet;
/// including session identification, extension configuration, and debug settings.
#[derive(Default, Clone, Debug)]
pub struct SessionBuilderConfig {
/// Optional session ID for resuming or identifying an existing session
/// Session id, optional need to deduce from context
pub session_id: Option<String>,
/// Whether to resume an existing session
pub resume: bool,
@@ -132,8 +133,14 @@ async fn offer_extension_debugging_help(
}
}
// Create the debugging session
let mut debug_session = CliSession::new(debug_agent, None, false, None, None, None, None).await;
let session = SessionManager::create_session(
std::env::current_dir()?,
"CLI Session".to_string(),
SessionType::Hidden,
)
.await?;
let mut debug_session =
CliSession::new(debug_agent, session.id, false, None, None, None, None).await;
// Process the debugging request
println!("{}", style("Analyzing the extension failure...").yellow());
@@ -278,12 +285,20 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession {
process::exit(1);
});
let session_id: Option<String> = if session_config.no_session {
None
let session_id: String = if session_config.no_session {
let working_dir = std::env::current_dir().expect("Could not get working directory");
let session = SessionManager::create_session(
working_dir,
"CLI Session".to_string(),
SessionType::Hidden,
)
.await
.expect("Could not create session");
session.id
} else if session_config.resume {
if let Some(session_id) = session_config.session_id {
match SessionManager::get_session(&session_id, false).await {
Ok(_) => Some(session_id),
Ok(_) => session_id,
Err(_) => {
output::render_error(&format!(
"Cannot resume session {} - no such session exists",
@@ -294,7 +309,7 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession {
}
} else {
match SessionManager::list_sessions().await {
Ok(sessions) if !sessions.is_empty() => Some(sessions[0].id.clone()),
Ok(sessions) if !sessions.is_empty() => sessions[0].id.clone(),
_ => {
output::render_error("Cannot resume - no previous sessions found");
process::exit(1);
@@ -302,46 +317,44 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession {
}
}
} else {
session_config.session_id
session_config.session_id.unwrap()
};
agent
.extension_manager
.set_context(PlatformExtensionContext {
session_id: session_id.clone(),
session_id: Some(session_id.clone()),
extension_manager: Some(Arc::downgrade(&agent.extension_manager)),
tool_route_manager: Some(Arc::downgrade(&agent.tool_route_manager)),
})
.await;
if session_config.resume {
if let Some(session_id) = session_id.as_ref() {
let metadata = SessionManager::get_session(session_id, false)
.await
.unwrap_or_else(|e| {
output::render_error(&format!("Failed to read session metadata: {}", e));
process::exit(1);
});
let session = SessionManager::get_session(&session_id, false)
.await
.unwrap_or_else(|e| {
output::render_error(&format!("Failed to read session metadata: {}", e));
process::exit(1);
});
let current_workdir =
std::env::current_dir().expect("Failed to get current working directory");
if current_workdir != metadata.working_dir {
let change_workdir = cliclack::confirm(format!("{} The original working directory of this session was set to {}. Your current directory is {}. Do you want to switch back to the original working directory?", style("WARNING:").yellow(), style(metadata.working_dir.display()).cyan(), style(current_workdir.display()).cyan()))
let current_workdir =
std::env::current_dir().expect("Failed to get current working directory");
if current_workdir != session.working_dir {
let change_workdir = cliclack::confirm(format!("{} The original working directory of this session was set to {}. Your current directory is {}. Do you want to switch back to the original working directory?", style("WARNING:").yellow(), style(session.working_dir.display()).cyan(), style(current_workdir.display()).cyan()))
.initial_value(true)
.interact().expect("Failed to get user input");
if change_workdir {
if !metadata.working_dir.exists() {
output::render_error(&format!(
"Cannot switch to original working directory - {} no longer exists",
style(metadata.working_dir.display()).cyan()
));
} else if let Err(e) = std::env::set_current_dir(&metadata.working_dir) {
output::render_error(&format!(
"Failed to switch to original working directory: {}",
e
));
}
if change_workdir {
if !session.working_dir.exists() {
output::render_error(&format!(
"Cannot switch to original working directory - {} no longer exists",
style(session.working_dir.display()).cyan()
));
} else if let Err(e) = std::env::set_current_dir(&session.working_dir) {
output::render_error(&format!(
"Failed to switch to original working directory: {}",
e
));
}
}
}
@@ -354,22 +367,18 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession {
agent.disable_router_for_recipe().await;
extensions.into_iter().collect()
} else if session_config.resume {
if let Some(session_id) = session_id.as_ref() {
match SessionManager::get_session(session_id, false).await {
Ok(session_data) => {
if let Some(saved_state) =
EnabledExtensionsState::from_extension_data(&session_data.extension_data)
{
check_missing_extensions_or_exit(&saved_state.extensions);
saved_state.extensions
} else {
get_enabled_extensions()
}
match SessionManager::get_session(&session_id, false).await {
Ok(session_data) => {
if let Some(saved_state) =
EnabledExtensionsState::from_extension_data(&session_data.extension_data)
{
check_missing_extensions_or_exit(&saved_state.extensions);
saved_state.extensions
} else {
get_enabled_extensions()
}
_ => get_enabled_extensions(),
}
} else {
get_enabled_extensions()
_ => get_enabled_extensions(),
}
} else {
get_enabled_extensions()
@@ -560,23 +569,19 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession {
}
}
if let Some(session_id) = session_id.as_ref() {
let session_config_for_save = SessionConfig {
id: session_id.clone(),
working_dir: std::env::current_dir().unwrap_or_default(),
schedule_id: None,
execution_mode: None,
max_turns: None,
retry_config: None,
};
let session_config_for_save = SessionConfig {
id: session_id.clone(),
schedule_id: None,
max_turns: None,
retry_config: None,
};
if let Err(e) = session
.agent
.save_extension_state(&session_config_for_save)
.await
{
tracing::warn!("Failed to save initial extension state: {}", e);
}
if let Err(e) = session
.agent
.save_extension_state(&session_config_for_save)
.await
{
tracing::warn!("Failed to save initial extension state: {}", e);
}
// Add CLI-specific system prompt extension
@@ -603,7 +608,7 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession {
session_config.resume,
&provider_name,
&model_name,
&session_id,
&Some(session_id),
Some(&provider_for_display),
);
}
crates/goose-cli/src/session/mod.rs
+33 -42
View File
@@ -58,7 +58,7 @@ pub enum RunMode {
pub struct CliSession {
agent: Agent,
messages: Conversation,
session_id: Option<String>,
session_id: String,
completion_cache: Arc<std::sync::RwLock<CompletionCache>>,
debug: bool,
run_mode: RunMode,
@@ -122,21 +122,17 @@ pub async fn classify_planner_response(
impl CliSession {
pub async fn new(
agent: Agent,
session_id: Option<String>,
session_id: String,
debug: bool,
scheduled_job_id: Option<String>,
max_turns: Option<u32>,
edit_mode: Option<EditMode>,
retry_config: Option<RetryConfig>,
) -> Self {
let messages = if let Some(session_id) = &session_id {
SessionManager::get_session(session_id, true)
.await
.map(|session| session.conversation.unwrap_or_default())
.unwrap()
} else {
Conversation::new_unvalidated(Vec::new())
};
let messages = SessionManager::get_session(&session_id, true)
.await
.map(|session| session.conversation.unwrap_or_default())
.unwrap();
CliSession {
agent,
@@ -152,8 +148,8 @@ impl CliSession {
}
}
pub fn session_id(&self) -> Option<&String> {
self.session_id.as_ref()
pub fn session_id(&self) -> &String {
&self.session_id
}
/// Add a stdio extension to the session
@@ -359,9 +355,6 @@ impl CliSession {
cancel_token: CancellationToken,
) -> Result<()> {
let cancel_token = cancel_token.clone();
// TODO(Douwe): Make sure we generate the description here still:
self.push_message(message);
self.process_agent_response(false, cancel_token).await?;
Ok(())
@@ -443,7 +436,7 @@ impl CliSession {
// Track the current directory and last instruction in projects.json
if let Err(e) = crate::project_tracker::update_project_tracker(
Some(&content),
self.session_id.as_deref(),
Some(&self.session_id),
) {
eprintln!("Warning: Failed to update project tracker with instruction: {}", e);
}
@@ -583,16 +576,14 @@ impl CliSession {
input::InputResult::Clear => {
save_history(&mut editor);
if let Some(session_id) = &self.session_id {
if let Err(e) = SessionManager::replace_conversation(
session_id,
&Conversation::default(),
)
.await
{
output::render_error(&format!("Failed to clear session: {}", e));
continue;
}
if let Err(e) = SessionManager::replace_conversation(
&self.session_id,
&Conversation::default(),
)
.await
{
output::render_error(&format!("Failed to clear session: {}", e));
continue;
}
self.messages.clear();
@@ -671,9 +662,10 @@ impl CliSession {
}
}
if let Some(id) = &self.session_id {
println!("Closing session. Session ID: {}", console::style(id).cyan());
}
println!(
"Closing session. Session ID: {}",
console::style(&self.session_id).cyan()
);
Ok(())
}
@@ -768,18 +760,20 @@ impl CliSession {
) -> Result<()> {
let cancel_token_clone = cancel_token.clone();
let session_config = self.session_id.as_ref().map(|session_id| SessionConfig {
id: session_id.clone(),
working_dir: std::env::current_dir().unwrap_or_default(),
let session_config = SessionConfig {
id: self.session_id.clone(),
schedule_id: self.scheduled_job_id.clone(),
execution_mode: None,
max_turns: self.max_turns,
retry_config: self.retry_config.clone(),
});
};
let user_message = self
.messages
.last()
.ok_or_else(|| anyhow::anyhow!("No user message"))?;
let mut stream = self
.agent
.reply(
self.messages.clone(),
user_message.clone(),
session_config.clone(),
Some(cancel_token.clone()),
)
@@ -1224,16 +1218,13 @@ impl CliSession {
);
}
pub async fn get_metadata(&self) -> Result<goose::session::Session> {
match &self.session_id {
Some(id) => SessionManager::get_session(id, false).await,
None => Err(anyhow::anyhow!("No session available")),
}
pub async fn get_session(&self) -> Result<goose::session::Session> {
SessionManager::get_session(&self.session_id, false).await
}
// Get the session's total token usage
pub async fn get_total_token_usage(&self) -> Result<Option<i32>> {
let metadata = self.get_metadata().await?;
let metadata = self.get_session().await?;
Ok(metadata.total_tokens)
}
@@ -1265,7 +1256,7 @@ impl CliSession {
}
}
match self.get_metadata().await {
match self.get_session().await {
Ok(metadata) => {
let total_tokens = metadata.total_tokens.unwrap_or(0) as usize;