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

feat(goose-acp): enable parallel sessions with isolated agent state (#6392)

Browse Source 
Signed-off-by: Adrian Cole <adrian@tetrate.io>
This commit is contained in:
Adrian Cole
2026-01-15 06:19:17 +08:00
committed by GitHub
parent fb0eca2c36
commit 7d4a6bd8ff
86 changed files with 2594 additions and 1938 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/Cargo.toml
+8 -11
View File
@@ -16,36 +16,34 @@ path = "src/main.rs"
[dependencies]
goose = { path = "../goose" }
goose-acp = { path = "../goose-acp" }
goose-bench = { path = "../goose-bench" }
goose-mcp = { path = "../goose-mcp" }
rmcp = { workspace = true }
sacp = { workspace = true }
clap = { version = "4.4", features = ["derive"] }
cliclack = "0.3.5"
console = "0.16.1"
uuid = { version = "1.11", features = ["v4"] }
dotenvy = "0.15.7"
bat = "0.25.0"
anyhow = "1.0"
serde_json = "1.0"
tokio = { version = "1.43", features = ["full"] }
futures = "0.3"
anyhow = { workspace = true }
serde_json = { workspace = true }
tokio = { workspace = true }
futures = { workspace = true }
serde = { version = "1.0", features = ["derive"] } # For serialization
serde_yaml = "0.9"
tempfile = "3"
etcetera = { workspace = true }
rand = "0.8.5"
rustyline = "15.0.0"
tracing = "0.1"
tracing = { workspace = true }
chrono = "0.4"
tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt", "json", "time"] }
tracing-appender = "0.2"
once_cell = "1.20.2"
shlex = "1.3.0"
async-trait = "0.1.89"
base64 = "0.22.1"
regex = "1.11.1"
nix = { version = "0.30.1", features = ["process", "signal"] }
tar = "0.4"
# Web server dependencies
axum = { version = "0.8.1", features = ["ws", "macros"] }
@@ -55,7 +53,6 @@ webbrowser = {workspace = true}
indicatif = "0.18.1"
tokio-util = { version = "0.7.15", features = ["compat", "rt"] }
anstream = "0.6.18"
url = "2.5.7"
open = "5.3.2"
urlencoding = "2.1"
clap_complete = "4.5.62"
@@ -71,5 +68,5 @@ disable-update = []
tempfile = "3"
temp-env = { version = "0.3.6", features = ["async_closure"] }
test-case = "3.3"
tokio = { version = "1.43", features = ["rt", "macros"] }
serial_test = "3.2.0"
tokio = { workspace = true }
serial_test = { workspace = true }
crates/goose-cli/src/cli.rs
+38 -26
View File
@@ -8,7 +8,6 @@ use goose_mcp::{
AutoVisualiserRouter, ComputerControllerServer, DeveloperServer, MemoryServer, TutorialServer,
};
use crate::commands::acp::run_acp_agent;
use crate::commands::bench::agent_generator;
use crate::commands::configure::{configure_telemetry_consent_dialog, handle_configure};
use crate::commands::info::handle_info;
@@ -320,21 +319,24 @@ async fn get_or_create_session_id(
return Ok(None);
}
let session_manager = SessionManager::instance();
let Some(id) = identifier else {
return if resume {
let sessions = SessionManager::list_sessions().await?;
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))
} else {
let session = SessionManager::create_session(
std::env::current_dir()?,
"CLI Session".to_string(),
SessionType::User,
)
.await?;
let session = session_manager
.create_session(
std::env::current_dir()?,
"CLI Session".to_string(),
SessionType::User,
)
.await?;
Ok(Some(session.id))
};
};
@@ -343,7 +345,7 @@ async fn get_or_create_session_id(
Ok(Some(session_id))
} else if let Some(name) = id.name {
if resume {
let sessions = SessionManager::list_sessions().await?;
let sessions = session_manager.list_sessions().await?;
let session_id = sessions
.into_iter()
.find(|s| s.name == name || s.id == name)
@@ -351,14 +353,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(),
SessionType::User,
)
.await?;
let session = session_manager
.create_session(std::env::current_dir()?, name.clone(), SessionType::User)
.await?;
SessionManager::update_session(&session.id)
session_manager
.update(&session.id)
.user_provided_name(name)
.apply()
.await?;
@@ -373,12 +373,13 @@ 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(),
SessionType::User,
)
.await?;
let session = session_manager
.create_session(
std::env::current_dir()?,
"CLI Session".to_string(),
SessionType::User,
)
.await?;
Ok(Some(session.id))
}
}
@@ -387,7 +388,8 @@ async fn lookup_session_id(identifier: Identifier) -> Result<String> {
if let Some(session_id) = identifier.session_id {
Ok(session_id)
} else if let Some(name) = identifier.name {
let sessions = SessionManager::list_sessions().await?;
let session_manager = SessionManager::instance();
let sessions = session_manager.list_sessions().await?;
sessions
.into_iter()
.find(|s| s.name == name || s.id == name)
@@ -1008,10 +1010,15 @@ async fn handle_session_subcommand(command: SessionCommand) -> Result<()> {
output,
format,
} => {
let session_manager = SessionManager::instance();
let session_identifier = if let Some(id) = identifier {
lookup_session_id(id).await?
} else {
match crate::commands::session::prompt_interactive_session_selection().await {
match crate::commands::session::prompt_interactive_session_selection(
&session_manager,
)
.await
{
Ok(id) => id,
Err(e) => {
eprintln!("Error: {}", e);
@@ -1023,10 +1030,15 @@ async fn handle_session_subcommand(command: SessionCommand) -> Result<()> {
.await?;
}
SessionCommand::Diagnostics { identifier, output } => {
let session_manager = SessionManager::instance();
let session_id = if let Some(id) = identifier {
lookup_session_id(id).await?
} else {
match crate::commands::session::prompt_interactive_session_selection().await {
match crate::commands::session::prompt_interactive_session_selection(
&session_manager,
)
.await
{
Ok(id) => id,
Err(e) => {
eprintln!("Error: {}", e);
@@ -1469,7 +1481,7 @@ pub async fn cli() -> anyhow::Result<()> {
Some(Command::Configure {}) => handle_configure().await,
Some(Command::Info { verbose }) => handle_info(verbose),
Some(Command::Mcp { server }) => handle_mcp_command(server).await,
Some(Command::Acp { builtins }) => run_acp_agent(builtins).await,
Some(Command::Acp { builtins }) => goose_acp::server::run(builtins).await,
Some(Command::Session {
command: Some(cmd), ..
}) => handle_session_subcommand(cmd).await,
crates/goose-cli/src/commands/acp.rs
-1233
View File
File diff suppressed because it is too large Load Diff
crates/goose-cli/src/commands/bench.rs
+3 -2
View File
@@ -3,6 +3,7 @@ use crate::session::SessionBuilderConfig;
use crate::{logging, CliSession};
use async_trait::async_trait;
use goose::conversation::Conversation;
use goose::session::session_manager::Session;
use goose_bench::bench_session::{BenchAgent, BenchBaseSession};
use goose_bench::eval_suites::ExtensionRequirements;
use std::sync::Arc;
@@ -25,8 +26,8 @@ impl BenchBaseSession for CliSession {
})
}
fn get_session_id(&self) -> anyhow::Result<String> {
Ok(self.session_id().to_string())
async fn get_session(&self) -> anyhow::Result<Session> {
self.get_session().await
}
}
pub async fn agent_generator(
crates/goose-cli/src/commands/configure.rs
+15 -12
View File
@@ -22,7 +22,7 @@ use goose::model::ModelConfig;
use goose::posthog::{get_telemetry_choice, TELEMETRY_ENABLED_KEY};
use goose::providers::provider_test::test_provider_configuration;
use goose::providers::{create, providers, retry_operation, RetryConfig};
use goose::session::{SessionManager, SessionType};
use goose::session::SessionType;
use serde_json::Value;
use std::collections::HashMap;
@@ -1081,8 +1081,7 @@ pub fn remove_extension_dialog() -> anyhow::Result<()> {
for name in selected {
remove_extension(&name_to_key(name));
let mut permission_manager = PermissionManager::default();
permission_manager.remove_extension(&name_to_key(name));
PermissionManager::instance().remove_extension(&name_to_key(name));
cliclack::outro(format!("Removed {} extension", style(name).green()))?;
}
@@ -1396,15 +1395,19 @@ pub async fn configure_tool_permissions_dialog() -> anyhow::Result<()> {
.expect("No model configured. Please set model first");
let model_config = ModelConfig::new(&model)?;
let session = SessionManager::create_session(
std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from(".")),
"Tool Permission Configuration".to_string(),
SessionType::Hidden,
)
.await?;
let agent = Agent::new();
let new_provider = create(&provider_name, model_config).await?;
let session = agent
.config
.session_manager
.create_session(
std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from(".")),
"Tool Permission Configuration".to_string(),
SessionType::Hidden,
)
.await?;
agent.update_provider(new_provider, &session.id).await?;
if let Some(config) = get_extension_by_name(&selected_extension_name) {
agent
@@ -1426,9 +1429,9 @@ pub async fn configure_tool_permissions_dialog() -> anyhow::Result<()> {
return Ok(());
}
let mut permission_manager = PermissionManager::default();
let permission_manager = PermissionManager::instance();
let selected_tools = agent
.list_tools(Some(selected_extension_name.clone()))
.list_tools(&session.id, Some(selected_extension_name.clone()))
.await
.into_iter()
.map(|tool| {
crates/goose-cli/src/commands/mod.rs
-1
View File
@@ -1,4 +1,3 @@
pub mod acp;
pub mod bench;
pub mod configure;
pub mod info;
crates/goose-cli/src/commands/schedule.rs
+31 -13
View File
@@ -3,17 +3,10 @@ use goose::scheduler::{
get_default_scheduled_recipes_dir, get_default_scheduler_storage_path, ScheduledJob, Scheduler,
SchedulerError,
};
use goose::session::SessionManager;
use std::path::Path;
use std::sync::Arc;
async fn create_scheduler() -> Result<Arc<Scheduler>> {
let storage_path =
get_default_scheduler_storage_path().context("Failed to get scheduler storage path")?;
Scheduler::new(storage_path)
.await
.context("Failed to initialize scheduler")
}
fn validate_cron_expression(cron: &str) -> Result<()> {
// Basic validation and helpful suggestions
if cron.trim().is_empty() {
@@ -97,7 +90,12 @@ pub async fn handle_schedule_add(
process_start_time: None,
};
let scheduler = create_scheduler().await?;
let scheduler_storage_path =
get_default_scheduler_storage_path().context("Failed to get scheduler storage path")?;
let session_manager = Arc::new(SessionManager::instance());
let scheduler = Scheduler::new(scheduler_storage_path, session_manager)
.await
.context("Failed to initialize scheduler")?;
match scheduler.add_scheduled_job(job, true).await {
Ok(_) => {
@@ -139,7 +137,12 @@ pub async fn handle_schedule_add(
}
pub async fn handle_schedule_list() -> Result<()> {
let scheduler = create_scheduler().await?;
let scheduler_storage_path =
get_default_scheduler_storage_path().context("Failed to get scheduler storage path")?;
let session_manager = Arc::new(SessionManager::instance());
let scheduler = Scheduler::new(scheduler_storage_path, session_manager)
.await
.context("Failed to initialize scheduler")?;
let jobs = scheduler.list_scheduled_jobs().await;
if jobs.is_empty() {
@@ -170,7 +173,12 @@ pub async fn handle_schedule_list() -> Result<()> {
}
pub async fn handle_schedule_remove(schedule_id: String) -> Result<()> {
let scheduler = create_scheduler().await?;
let scheduler_storage_path =
get_default_scheduler_storage_path().context("Failed to get scheduler storage path")?;
let session_manager = Arc::new(SessionManager::instance());
let scheduler = Scheduler::new(scheduler_storage_path, session_manager)
.await
.context("Failed to initialize scheduler")?;
match scheduler.remove_scheduled_job(&schedule_id, true).await {
Ok(_) => {
@@ -193,7 +201,12 @@ pub async fn handle_schedule_remove(schedule_id: String) -> Result<()> {
}
pub async fn handle_schedule_sessions(schedule_id: String, limit: Option<usize>) -> Result<()> {
let scheduler = create_scheduler().await?;
let scheduler_storage_path =
get_default_scheduler_storage_path().context("Failed to get scheduler storage path")?;
let session_manager = Arc::new(SessionManager::instance());
let scheduler = Scheduler::new(scheduler_storage_path, session_manager)
.await
.context("Failed to initialize scheduler")?;
match scheduler.sessions(&schedule_id, limit.unwrap_or(50)).await {
Ok(sessions) => {
@@ -225,7 +238,12 @@ pub async fn handle_schedule_sessions(schedule_id: String, limit: Option<usize>)
}
pub async fn handle_schedule_run_now(schedule_id: String) -> Result<()> {
let scheduler = create_scheduler().await?;
let scheduler_storage_path =
get_default_scheduler_storage_path().context("Failed to get scheduler storage path")?;
let session_manager = Arc::new(SessionManager::instance());
let scheduler = Scheduler::new(scheduler_storage_path, session_manager)
.await
.context("Failed to initialize scheduler")?;
match scheduler.run_now(&schedule_id).await {
Ok(session_id) => {
crates/goose-cli/src/commands/session.rs
+22 -14
View File
@@ -11,7 +11,7 @@ use std::path::PathBuf;
const TRUNCATED_DESC_LENGTH: usize = 60;
pub async fn remove_sessions(sessions: Vec<Session>) -> Result<()> {
async fn remove_sessions(session_manager: &SessionManager, sessions: Vec<Session>) -> Result<()> {
println!("The following sessions will be removed:");
for session in &sessions {
println!("- {} {}", session.id, session.name);
@@ -23,7 +23,7 @@ pub async fn remove_sessions(sessions: Vec<Session>) -> Result<()> {
if should_delete {
for session in sessions {
SessionManager::delete_session(&session.id).await?;
session_manager.delete_session(&session.id).await?;
println!("Session `{}` removed.", session.id);
}
} else {
@@ -76,7 +76,8 @@ pub async fn handle_session_remove(
name: Option<String>,
regex_string: Option<String>,
) -> Result<()> {
let all_sessions = match SessionManager::list_sessions().await {
let session_manager = SessionManager::instance();
let all_sessions = match session_manager.list_sessions().await {
Ok(sessions) => sessions,
Err(e) => {
tracing::error!("Failed to retrieve sessions: {:?}", e);
@@ -125,7 +126,7 @@ pub async fn handle_session_remove(
return Ok(());
}
remove_sessions(matched_sessions).await
remove_sessions(&session_manager, matched_sessions).await
}
pub async fn handle_session_list(
@@ -134,7 +135,8 @@ pub async fn handle_session_list(
working_dir: Option<PathBuf>,
limit: Option<usize>,
) -> Result<()> {
let mut sessions = SessionManager::list_sessions().await?;
let session_manager = SessionManager::instance();
let mut sessions = session_manager.list_sessions().await?;
if let Some(ref pat) = working_dir {
let pat_lower = pat.to_string_lossy().to_lowercase();
@@ -181,7 +183,8 @@ pub async fn handle_session_export(
output_path: Option<PathBuf>,
format: String,
) -> Result<()> {
let session = match SessionManager::get_session(&session_id, true).await {
let session_manager = SessionManager::instance();
let session = match session_manager.get_session(&session_id, true).await {
Ok(session) => session,
Err(e) => {
return Err(anyhow::anyhow!(
@@ -222,12 +225,15 @@ pub async fn handle_diagnostics(session_id: &str, output_path: Option<PathBuf>)
session_id
);
let diagnostics_data = generate_diagnostics(session_id).await.with_context(|| {
format!(
"Failed to write to generate diagnostics bundle for session '{}'",
session_id
)
})?;
let session_manager = SessionManager::instance();
let diagnostics_data = generate_diagnostics(&session_manager, session_id)
.await
.with_context(|| {
format!(
"Failed to write to generate diagnostics bundle for session '{}'",
session_id
)
})?;
let output_file = if let Some(path) = output_path {
path.clone()
@@ -319,8 +325,10 @@ fn export_session_to_markdown(
/// Prompt the user to interactively select a session
///
/// Shows a list of available sessions and lets the user select one
pub async fn prompt_interactive_session_selection() -> Result<String> {
let sessions = SessionManager::list_sessions().await?;
pub async fn prompt_interactive_session_selection(
session_manager: &SessionManager,
) -> Result<String> {
let sessions = session_manager.list_sessions().await?;
if sessions.is_empty() {
return Err(anyhow::anyhow!("No sessions found"));
crates/goose-cli/src/commands/term.rs
+25 -15
View File
@@ -1,8 +1,7 @@
use anyhow::{anyhow, Result};
use chrono;
use goose::conversation::message::{Message, MessageContent, MessageMetadata};
use goose::session::SessionManager;
use goose::session::SessionType;
use goose::session::{SessionManager, SessionType};
use rmcp::model::Role;
use crate::session::{build_session, SessionBuilderConfig};
@@ -119,10 +118,13 @@ pub async fn handle_term_init(
with_command_not_found: bool,
) -> Result<()> {
let config = shell.config();
let session_manager = SessionManager::instance();
let working_dir = std::env::current_dir()?;
let named_session = if let Some(ref name) = name {
let sessions = SessionManager::list_sessions_by_types(&[SessionType::Terminal]).await?;
let sessions = session_manager
.list_sessions_by_types(&[SessionType::Terminal])
.await?;
sessions.into_iter().find(|s| s.name == *name)
} else {
None
@@ -131,15 +133,17 @@ pub async fn handle_term_init(
let session = match named_session {
Some(s) => s,
None => {
let session = SessionManager::create_session(
working_dir,
"Goose Term Session".to_string(),
SessionType::Terminal,
)
.await?;
let session = session_manager
.create_session(
working_dir,
"Goose Term Session".to_string(),
SessionType::Terminal,
)
.await?;
if let Some(name) = name {
SessionManager::update_session(&session.id)
session_manager
.update(&session.id)
.user_provided_name(name)
.apply()
.await?;
@@ -184,7 +188,8 @@ pub async fn handle_term_log(command: String) -> Result<()> {
)
.with_metadata(MessageMetadata::user_only());
SessionManager::add_message(&session_id, &message).await?;
let session_manager = SessionManager::instance();
session_manager.add_message(&session_id, &message).await?;
Ok(())
}
@@ -201,13 +206,15 @@ pub async fn handle_term_run(prompt: Vec<String>) -> Result<()> {
})?;
let working_dir = std::env::current_dir()?;
let session_manager = SessionManager::instance();
SessionManager::update_session(&session_id)
session_manager
.update(&session_id)
.working_dir(working_dir)
.apply()
.await?;
let session = SessionManager::get_session(&session_id, true).await?;
let session = session_manager.get_session(&session_id, true).await?;
let user_messages_after_last_assistant: Vec<&Message> =
if let Some(conv) = &session.conversation {
conv.messages()
@@ -220,7 +227,9 @@ pub async fn handle_term_run(prompt: Vec<String>) -> Result<()> {
};
if let Some(oldest_user) = user_messages_after_last_assistant.last() {
SessionManager::truncate_conversation(&session_id, oldest_user.created).await?;
session_manager
.truncate_conversation(&session_id, oldest_user.created)
.await?;
}
let prompt_with_context = if user_messages_after_last_assistant.is_empty() {
@@ -260,7 +269,8 @@ pub async fn handle_term_info() -> Result<()> {
Err(_) => return Ok(()),
};
let session = SessionManager::get_session(&session_id, false).await.ok();
let session_manager = SessionManager::instance();
let session = session_manager.get_session(&session_id, false).await.ok();
let total_tokens = session.as_ref().and_then(|s| s.total_tokens).unwrap_or(0) as usize;
let config = goose::config::Config::global();
crates/goose-cli/src/commands/web.rs
+40 -20
View File
@@ -16,7 +16,6 @@ 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;
use std::{net::ToSocketAddrs, sync::Arc};
@@ -171,14 +170,17 @@ fn get_provider_and_model() -> (String, String) {
async fn create_agent(provider_name: &str, model: &str) -> Result<Agent> {
let model_config = goose::model::ModelConfig::new(model)?;
let init_session = SessionManager::create_session(
std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from(".")),
"Web Agent Initialization".to_string(),
SessionType::Hidden,
)
.await?;
let agent = Agent::new();
let session_manager = agent.config.session_manager.clone();
let init_session = session_manager
.create_session(
std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from(".")),
"Web Agent Initialization".to_string(),
SessionType::Hidden,
)
.await?;
let provider = goose::providers::create(provider_name, model_config).await?;
agent.update_provider(provider, &init_session.id).await?;
@@ -284,14 +286,21 @@ pub async fn handle_web(
Ok(())
}
async fn serve_index(uri: Uri) -> 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()))?;
async fn serve_index(
State(state): State<AppState>,
uri: Uri,
) -> Result<Redirect, (http::StatusCode, String)> {
let session = state
.agent
.config
.session_manager
.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()))?;
let redirect_url = if let Some(query) = uri.query() {
format!("/session/{}?{}", session.id, query)
@@ -351,8 +360,8 @@ async fn health_check() -> Json<serde_json::Value> {
}))
}
async fn list_sessions() -> Json<serde_json::Value> {
match SessionManager::list_sessions().await {
async fn list_sessions(State(state): State<AppState>) -> Json<serde_json::Value> {
match state.agent.config.session_manager.list_sessions().await {
Ok(sessions) => {
let mut session_info = Vec::new();
@@ -375,9 +384,16 @@ async fn list_sessions() -> Json<serde_json::Value> {
}
}
async fn get_session(
State(state): State<AppState>,
axum::extract::Path(session_id): axum::extract::Path<String>,
) -> Json<serde_json::Value> {
match SessionManager::get_session(&session_id, true).await {
match state
.agent
.config
.session_manager
.get_session(&session_id, true)
.await
{
Ok(session) => Json(serde_json::json!({
"metadata": session,
"messages": session.conversation.unwrap_or_default().messages()
@@ -544,7 +560,11 @@ async fn process_message_streaming(
return Ok(());
}
let session = SessionManager::get_session(&session_id, true).await?;
let session = agent
.config
.session_manager
.get_session(&session_id, true)
.await?;
let mut messages = session.conversation.unwrap_or_default();
messages.push(user_message.clone());
crates/goose-cli/src/scenario_tests/mock_client.rs
+2 -1
View File
@@ -1,7 +1,7 @@
//! MockClient is a mock implementation of the McpClientTrait for testing purposes.
//! add a tool you want to have around and then add the client to the extension router
use goose::agents::mcp_client::{Error, McpClientTrait};
use goose::agents::mcp_client::{Error, McpClientTrait, McpMeta};
use rmcp::{
model::{
CallToolResult, Content, ErrorData, GetPromptResult, ListPromptsResult,
@@ -94,6 +94,7 @@ impl McpClientTrait for MockClient {
&self,
name: &str,
arguments: Option<serde_json::Map<String, Value>>,
_meta: McpMeta,
_cancel_token: CancellationToken,
) -> Result<CallToolResult, Error> {
if let Some(handler) = self.handlers.get(name) {
crates/goose-cli/src/scenario_tests/scenario_runner.rs
+18 -8
View File
@@ -6,7 +6,9 @@ use crate::scenario_tests::mock_client::weather_client;
use crate::scenario_tests::provider_configs::{get_provider_configs, ProviderConfig};
use crate::session::CliSession;
use anyhow::Result;
use goose::agents::Agent;
use goose::agents::{Agent, AgentConfig};
use goose::config::permission::PermissionManager;
use goose::config::GooseMode;
use goose::model::ModelConfig;
use goose::providers::{create, testprovider::TestProvider};
use goose::session::session_manager::SessionType;
@@ -14,6 +16,7 @@ use goose::session::SessionManager;
use std::collections::{HashMap, HashSet};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tempfile::TempDir;
use tokio_util::sync::CancellationToken;
pub const SCENARIO_TESTS_DIR: &str = "src/scenario_tests";
@@ -198,7 +201,11 @@ where
let mock_client = weather_client();
let agent = Agent::new();
let temp_dir = TempDir::new()?;
let session_manager = Arc::new(SessionManager::new(temp_dir.path().to_path_buf()));
let permission_manager = Arc::new(PermissionManager::new(temp_dir.path().to_path_buf()));
let agent_config = AgentConfig::new(session_manager, permission_manager, None, GooseMode::Auto); // no scheduler needed for scenario tests
let agent = Agent::with_config(agent_config);
agent
.extension_manager
.add_client(
@@ -217,12 +224,15 @@ where
)
.await;
let session = SessionManager::create_session(
PathBuf::default(),
"scenario-runner".to_string(),
SessionType::Hidden,
)
.await?;
let session = agent
.config
.session_manager
.create_session(
PathBuf::default(),
"scenario-runner".to_string(),
SessionType::Hidden,
)
.await?;
agent
.update_provider(
crates/goose-cli/src/session/builder.rs
+28 -30
View File
@@ -10,9 +10,7 @@ use goose::config::{
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;
use std::collections::HashSet;
@@ -147,12 +145,15 @@ async fn offer_extension_debugging_help(
// Create a minimal agent for debugging
let debug_agent = Agent::new();
let session = SessionManager::create_session(
std::env::current_dir()?,
"CLI Session".to_string(),
SessionType::Hidden,
)
.await?;
let session = debug_agent
.config
.session_manager
.create_session(
std::env::current_dir()?,
"CLI Session".to_string(),
SessionType::Hidden,
)
.await?;
debug_agent.update_provider(provider, &session.id).await?;
@@ -252,10 +253,12 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession {
goose::posthog::set_session_context("cli", session_config.resume);
let config = Config::global();
let agent: Agent = Agent::new();
let session_manager = agent.config.session_manager.clone();
let (saved_provider, saved_model_config) = if session_config.resume {
if let Some(ref session_id) = session_config.session_id {
match SessionManager::get_session(session_id, false).await {
match session_manager.get_session(session_id, false).await {
Ok(session_data) => (session_data.provider_name, session_data.model_config),
Err(_) => (None, None),
}
@@ -310,8 +313,6 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession {
.with_temperature(temperature)
};
let agent: Agent = Agent::new();
agent
.apply_recipe_components(
session_config.sub_recipes,
@@ -348,17 +349,14 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession {
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");
let session = session_manager
.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 {
match session_manager.get_session(&session_id, false).await {
Ok(_) => session_id,
Err(_) => {
output::render_error(&format!(
@@ -369,7 +367,7 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession {
}
}
} else {
match SessionManager::list_sessions().await {
match session_manager.list_sessions().await {
Ok(sessions) if !sessions.is_empty() => sessions[0].id.clone(),
_ => {
output::render_error("Cannot resume - no previous sessions found");
@@ -389,16 +387,11 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession {
process::exit(1);
});
agent
.extension_manager
.set_context(PlatformExtensionContext {
session_id: Some(session_id.clone()),
extension_manager: Some(Arc::downgrade(&agent.extension_manager)),
})
.await;
if session_config.resume {
let session = SessionManager::get_session(&session_id, false)
let session = agent
.config
.session_manager
.get_session(&session_id, false)
.await
.unwrap_or_else(|e| {
output::render_error(&format!("Failed to read session metadata: {}", e));
@@ -451,7 +444,12 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession {
let extensions_to_run: Vec<_> = if let Some(extensions) = session_config.extensions_override {
extensions.into_iter().collect()
} else if session_config.resume {
match SessionManager::get_session(&session_id, false).await {
match agent
.config
.session_manager
.get_session(&session_id, false)
.await
{
Ok(session_data) => {
if let Some(saved_state) =
EnabledExtensionsState::from_extension_data(&session_data.extension_data)
crates/goose-cli/src/session/mod.rs
+32 -8
View File
@@ -33,7 +33,6 @@ use goose::agents::extension::{Envs, ExtensionConfig, PLATFORM_EXTENSIONS};
use goose::agents::types::RetryConfig;
use goose::agents::{Agent, SessionConfig, COMPACT_TRIGGERS};
use goose::config::{Config, GooseMode};
use goose::session::SessionManager;
use input::InputResult;
use rmcp::model::PromptMessage;
use rmcp::model::ServerNotification;
@@ -229,7 +228,10 @@ impl CliSession {
retry_config: Option<RetryConfig>,
output_format: String,
) -> Self {
let messages = SessionManager::get_session(&session_id, true)
let messages = agent
.config
.session_manager
.get_session(&session_id, true)
.await
.map(|session| session.conversation.unwrap_or_default())
.unwrap();
@@ -693,14 +695,22 @@ impl CliSession {
}
async fn handle_clear(&mut self) -> Result<()> {
if let Err(e) =
SessionManager::replace_conversation(&self.session_id, &Conversation::default()).await
if let Err(e) = self
.agent
.config
.session_manager
.replace_conversation(&self.session_id, &Conversation::default())
.await
{
output::render_error(&format!("Failed to clear session: {}", e));
return Ok(());
}
if let Err(e) = SessionManager::update_session(&self.session_id)
if let Err(e) = self
.agent
.config
.session_manager
.update(&self.session_id)
.total_tokens(Some(0))
.input_tokens(Some(0))
.output_tokens(Some(0))
@@ -1021,7 +1031,13 @@ impl CliSession {
}
if is_json_mode {
let metadata = match SessionManager::get_session(&self.session_id, false).await {
let metadata = match self
.agent
.config
.session_manager
.get_session(&self.session_id, false)
.await
{
Ok(session) => JsonMetadata {
total_tokens: session.total_tokens,
status: "completed".to_string(),
@@ -1037,7 +1053,11 @@ impl CliSession {
};
println!("{}", serde_json::to_string_pretty(&json_output)?);
} else if is_stream_json_mode {
let total_tokens = SessionManager::get_session(&self.session_id, false)
let total_tokens = self
.agent
.config
.session_manager
.get_session(&self.session_id, false)
.await
.ok()
.and_then(|s| s.total_tokens);
@@ -1207,7 +1227,11 @@ impl CliSession {
}
pub async fn get_session(&self) -> Result<goose::session::Session> {
SessionManager::get_session(&self.session_id, false).await
self.agent
.config
.session_manager
.get_session(&self.session_id, false)
.await
}
// Get the session's total token usage