diff --git a/Cargo.lock b/Cargo.lock index 55f24b2ed..f356cab58 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -66,7 +66,7 @@ dependencies = [ [[package]] name = "agent-client-protocol" version = "2.0.0" -source = "git+https://github.com/agentclientprotocol/rust-sdk?rev=c97a5203d3392f7f231514d84eea014f9f43e6fb#c97a5203d3392f7f231514d84eea014f9f43e6fb" +source = "git+file:///Users/john/Project/tkmind_go-v141-prod-prep/.hermit/rust/git/checkouts/rust-sdk-762d172c5e846930/c97a520?rev=c97a5203d3392f7f231514d84eea014f9f43e6fb#c97a5203d3392f7f231514d84eea014f9f43e6fb" dependencies = [ "agent-client-protocol-derive", "agent-client-protocol-schema", @@ -89,7 +89,7 @@ dependencies = [ [[package]] name = "agent-client-protocol-derive" version = "2.0.0" -source = "git+https://github.com/agentclientprotocol/rust-sdk?rev=c97a5203d3392f7f231514d84eea014f9f43e6fb#c97a5203d3392f7f231514d84eea014f9f43e6fb" +source = "git+file:///Users/john/Project/tkmind_go-v141-prod-prep/.hermit/rust/git/checkouts/rust-sdk-762d172c5e846930/c97a520?rev=c97a5203d3392f7f231514d84eea014f9f43e6fb#c97a5203d3392f7f231514d84eea014f9f43e6fb" dependencies = [ "quote", "syn 3.0.4", @@ -98,7 +98,7 @@ dependencies = [ [[package]] name = "agent-client-protocol-http" version = "2.0.0" -source = "git+https://github.com/agentclientprotocol/rust-sdk?rev=c97a5203d3392f7f231514d84eea014f9f43e6fb#c97a5203d3392f7f231514d84eea014f9f43e6fb" +source = "git+file:///Users/john/Project/tkmind_go-v141-prod-prep/.hermit/rust/git/checkouts/rust-sdk-762d172c5e846930/c97a520?rev=c97a5203d3392f7f231514d84eea014f9f43e6fb#c97a5203d3392f7f231514d84eea014f9f43e6fb" dependencies = [ "agent-client-protocol", "async-stream", @@ -2999,7 +2999,7 @@ dependencies = [ [[package]] name = "cudaforge" version = "0.1.6" -source = "git+https://github.com/jbg/cudaforge?rev=e7c1967340e40673db98dc9e17da0f04834a456f#e7c1967340e40673db98dc9e17da0f04834a456f" +source = "git+file:///Users/john/Project/tkmind_go-v141-prod-prep/.hermit/rust/git/checkouts/cudaforge-5c30904947dfbb68/e7c1967?rev=e7c1967340e40673db98dc9e17da0f04834a456f#e7c1967340e40673db98dc9e17da0f04834a456f" dependencies = [ "anyhow", "fs2", diff --git a/Cargo.toml b/Cargo.toml index d66da028e..865b9c39a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -84,6 +84,7 @@ shell-words = { version = "1", default-features = false, features = ["std"] } snap = { version = "1", default-features = false } test-case = { version = "3", default-features = false } url = { version = "2.5.4", default-features = false, features = ["std"] } +utoipa = { version = "4.2", default-features = false } opentelemetry = { version = "0.32", default-features = false } opentelemetry_sdk = { version = "0.32", default-features = false, features = ["metrics"] } opentelemetry-http = { version = "0.32", default-features = false, features = ["internal-logs", "reqwest"] } @@ -115,10 +116,10 @@ icu_calendar = { version = "=2.1.1", default-features = false } icu_locale = { version = "=2.1.1", default-features = false } [patch.crates-io] -agent-client-protocol = { git = "https://github.com/agentclientprotocol/rust-sdk", rev = "c97a5203d3392f7f231514d84eea014f9f43e6fb" } -agent-client-protocol-http = { git = "https://github.com/agentclientprotocol/rust-sdk", rev = "c97a5203d3392f7f231514d84eea014f9f43e6fb" } +agent-client-protocol = { git = "file:///Users/john/Project/tkmind_go-v141-prod-prep/.hermit/rust/git/checkouts/rust-sdk-762d172c5e846930/c97a520", rev = "c97a5203d3392f7f231514d84eea014f9f43e6fb" } +agent-client-protocol-http = { git = "file:///Users/john/Project/tkmind_go-v141-prod-prep/.hermit/rust/git/checkouts/rust-sdk-762d172c5e846930/c97a520", rev = "c97a5203d3392f7f231514d84eea014f9f43e6fb" } v8 = { path = "vendor/v8" } -cudaforge = { git = "https://github.com/jbg/cudaforge", rev = "e7c1967340e40673db98dc9e17da0f04834a456f" } +cudaforge = { git = "file:///Users/john/Project/tkmind_go-v141-prod-prep/.hermit/rust/git/checkouts/cudaforge-5c30904947dfbb68/e7c1967", rev = "e7c1967340e40673db98dc9e17da0f04834a456f" } # Dependencies don't need debug info. You rarely step-debug into third-party # crates, but their DWARF dominates `target/debug` — roughly 50-70% of each diff --git a/TKMIND_V1_49_PORT_LEDGER.md b/TKMIND_V1_49_PORT_LEDGER.md new file mode 100644 index 000000000..859418157 --- /dev/null +++ b/TKMIND_V1_49_PORT_LEDGER.md @@ -0,0 +1,34 @@ +# TKMind v1.49 Port Ledger + +Worktree: `/Users/john/Project/tkmind_go-v149` · branch `upgrade/v1.49-tkmind` · upstream `v1.49.0` + +决策与阶段见 Memind 仓库 [TKMIND_V1_49_MIGRATION.md](../Memind/TKMIND_V1_49_MIGRATION.md)。 + +## Port queue (from upgrade/goose-1.47 → v1.49) + +| Priority | Path / module | Status | Notes | +|----------|---------------|--------|-------| +| P0 | `crates/goose/src/session/session_manager.rs` (PG) | done | `postgres_storage.rs` + env `GOOSE_SESSION_DB_URL`; schema v16 | +| P0 | `crates/goose-server/src/routes/tkmind_compat.rs` | done | → `crates/goose/src/acp/tkmind/compat.rs` | +| P0 | `POST /agent/start` + update_provider/session | done | `acp/tkmind/agent.rs` | +| P0 | `GET /sessions/{id}/events` + `POST .../reply` + `POST .../cancel` | done | `acp/tkmind/session_events.rs` + `session_event_bus.rs` | +| P1 | `platform_extensions/mod.rs` + projectmemory | done | harness bootstrap state | +| P1 | `platform_extensions/developer/workspace_path.rs` | done | path isolation helpers for aider/openhands | +| P1 | `platform_extensions/aider.rs` + `openhands.rs` | done | `aider__code` / `openhands__code`; smoke `check-goosed-v149-executors.mjs` | +| P1 | sandbox-fs stdio MCP via `extension_overrides` | done | `GET /agent/tools` + `GET /sessions/{id}/extensions`; smoke `check-goosed-v149-sandbox-fs.mjs` | +| P2 | `providers/inventory` custom_tkmind | done | declarative `/config/custom-providers` + `check-goosed-v149-provider.mjs` | +| P2 | `deploy/coding_router.sh` | done | copied from tkmind_go 1.47 → `tkmind_go-v149/deploy/` | + +## Milestone 1 + +- [x] `cargo build --release --package goose-cli --bin goose` (minimal mode) +- [x] `bash Memind/scripts/run-goosed-v149-local.sh` → `curl -sk https://127.0.0.1:18049/status` = ok +- [x] P0 routes compile and respond (`GET /sessions` → `{"sessions":[]}`) + +## Milestone 2 + +- [x] PG session read/write on isolated DB (`schema_version=16` @ `goose_sessions_v149_dev`) +- [x] sandbox-fs MCP load (`scripts/check-goosed-v149-sandbox-fs.mjs` → `sandbox-fs__write_file` etc.) +- [x] one H5 `/agent/start` smoke on 18049 (`scripts/check-goosed-v149-agent-start.mjs`) +- [x] `/sessions/{id}/reply` + `/sessions/{id}/events` SSE (`scripts/check-goosed-v149-reply-smoke.mjs`; provider optional — `Provider not set` Error validates bus) +- [x] aider + openhands platform extensions (`scripts/check-goosed-v149-executors.mjs`) diff --git a/crates/goose-cli/src/cli.rs b/crates/goose-cli/src/cli.rs index bbebbf3bc..1e99df670 100644 --- a/crates/goose-cli/src/cli.rs +++ b/crates/goose-cli/src/cli.rs @@ -1843,11 +1843,13 @@ async fn handle_serve_command(args: ServeCommandArgs) -> Result<()> { } else { None }; + let tkmind_state = goose::acp::tkmind::state::AppState::new().await?; let router = create_router( server, secret_key, require_token, additional_allowed_origins, + tkmind_state, ); let config = Config::global(); diff --git a/crates/goose/Cargo.toml b/crates/goose/Cargo.toml index 8b3eb9d59..f092740f2 100644 --- a/crates/goose/Cargo.toml +++ b/crates/goose/Cargo.toml @@ -148,6 +148,7 @@ v_htmlescape = { version = "0.17", default-features = false, features = ["std"] sqlx = { version = "0.9.0", default-features = false, features = [ "runtime-tokio", "sqlite", + "postgres", "chrono", "json", "macros", diff --git a/crates/goose/src/acp/mod.rs b/crates/goose/src/acp/mod.rs index 3e73d59c6..4a5fdffc9 100644 --- a/crates/goose/src/acp/mod.rs +++ b/crates/goose/src/acp/mod.rs @@ -6,6 +6,7 @@ mod provider; mod response_builder; pub mod server; pub mod server_factory; +pub mod tkmind; pub(crate) mod tool_call_notifier; pub(crate) mod tools; pub mod transport; diff --git a/crates/goose/src/acp/server.rs b/crates/goose/src/acp/server.rs index c21a4b5b8..9557775bc 100644 --- a/crates/goose/src/acp/server.rs +++ b/crates/goose/src/acp/server.rs @@ -955,10 +955,10 @@ impl GooseAcpAgent { pub async fn new(options: GooseAcpAgentOptions) -> Result { let session_manager = Arc::new(SessionManager::new(options.data_dir)); - // Eagerly initialize the SQLite pool so it's ready when providers/sessions need it. + // Eagerly initialize session storage so the pool is ready when providers/sessions need it. let storage_clone = session_manager.storage().clone(); tokio::spawn(async move { - let _ = storage_clone.pool().await; + let _ = storage_clone.initialize().await; }); let permission_manager = Arc::new(PermissionManager::new(options.config_dir.clone())); diff --git a/crates/goose/src/acp/tkmind/agent.rs b/crates/goose/src/acp/tkmind/agent.rs new file mode 100644 index 000000000..7b3631893 --- /dev/null +++ b/crates/goose/src/acp/tkmind/agent.rs @@ -0,0 +1,446 @@ +use std::path::PathBuf; +use std::sync::Arc; + +use axum::{ + extract::{Query, State}, + http::StatusCode, + routing::{get, post}, + Json, Router, +}; +use goose_providers::model::ModelConfig; +use serde::Deserialize; +use tracing::error; + +use crate::agents::extension::ToolInfo; +use crate::agents::extension_manager::get_parameter_names; +use crate::agents::platform_extensions::{chatrecall, projectmemory, PLATFORM_EXTENSIONS}; +use crate::agents::ExtensionConfig; +use crate::config::permission::PermissionLevel; +use crate::config::resolve_extensions_for_new_session; +use crate::config::{Config, GooseMode}; +use crate::providers::create; +use crate::recipe::Recipe; +use crate::session::session_manager::SessionType; +use crate::session::{EnabledExtensionsState, ExtensionState, Session}; + +use super::errors::ErrorResponse; +use super::state::AppState; + +#[derive(Deserialize)] +pub struct GetToolsQuery { + extension_name: Option, + session_id: String, +} + +#[derive(Deserialize)] +pub struct UpdateProviderRequest { + provider: String, + model: Option, + session_id: String, + context_limit: Option, + request_params: Option>, +} + +#[derive(Deserialize)] +pub struct UpdateSessionRequest { + session_id: String, + goose_mode: Option, +} + +#[derive(Deserialize)] +pub struct AddExtensionRequest { + session_id: String, + config: ExtensionConfig, +} + +#[derive(Deserialize)] +pub struct RemoveExtensionRequest { + session_id: String, + name: String, +} + +#[derive(Deserialize)] +pub struct UpdateWorkingDirRequest { + session_id: String, + working_dir: String, +} + +#[derive(Deserialize)] +pub struct RestartAgentRequest { + session_id: String, +} + +#[derive(Deserialize)] +pub struct StartAgentRequest { + working_dir: String, + #[serde(default)] + recipe: Option, + #[serde(default)] + extension_overrides: Option>, + #[serde(default)] + enable_context_memory: bool, +} + +fn platform_extension_config(name: &str) -> Option { + PLATFORM_EXTENSIONS.get(name).map(|definition| ExtensionConfig::Platform { + name: definition.name.to_string(), + description: definition.description.to_string(), + display_name: Some(definition.display_name.to_string()), + bundled: Some(true), + available_tools: Vec::new(), + }) +} + +fn append_context_memory_extensions(extensions: &mut Vec) { + for name in [chatrecall::EXTENSION_NAME, projectmemory::EXTENSION_NAME] { + if extensions.iter().any(|extension| extension.name() == name) { + continue; + } + if let Some(extension) = platform_extension_config(name) { + extensions.push(extension); + } + } +} + +async fn start_agent( + State(state): State>, + Json(payload): Json, +) -> Result, ErrorResponse> { + let manager = state.session_manager(); + let config = Config::global(); + let current_mode = config.get_goose_mode().unwrap_or_default(); + + let mut session = manager + .create_session( + PathBuf::from(&payload.working_dir), + "New Chat".to_string(), + SessionType::User, + current_mode, + ) + .await + .map_err(|err| { + error!("Failed to create session: {err}"); + ErrorResponse { + message: format!("Failed to create session: {err}"), + status: StatusCode::BAD_REQUEST, + } + })?; + + let recipe_extensions = payload.recipe.as_ref().and_then(|r| r.extensions.as_deref()); + let mut extensions_to_use = + resolve_extensions_for_new_session(recipe_extensions, payload.extension_overrides); + if payload.enable_context_memory { + append_context_memory_extensions(&mut extensions_to_use); + } + + let mut extension_data = session.extension_data.clone(); + let extensions_state = EnabledExtensionsState::new(extensions_to_use); + if let Err(err) = extensions_state.to_extension_data(&mut extension_data) { + tracing::warn!("Failed to initialize session with extensions: {err}"); + } else { + manager + .update(&session.id) + .extension_data(extension_data.clone()) + .apply() + .await + .map_err(|err| ErrorResponse { + message: format!("Failed to save initial extension state: {err}"), + status: StatusCode::INTERNAL_SERVER_ERROR, + })?; + } + + if let Some(recipe) = payload.recipe { + let mut update = manager.update(&session.id).recipe(Some(recipe.clone())); + + if let Some(settings) = &recipe.settings { + if let Some(provider) = &settings.goose_provider { + update = update.provider_name(provider); + if let Some(model) = &settings.goose_model { + update = update.model_config(ModelConfig::new(model)); + } + } + } + + update.apply().await.map_err(|err| ErrorResponse { + message: format!("Failed to update session with recipe: {err}"), + status: StatusCode::INTERNAL_SERVER_ERROR, + })?; + } + + session = manager.get_session(&session.id, false).await.map_err(|err| { + ErrorResponse { + message: format!("Failed to get updated session: {err}"), + status: StatusCode::INTERNAL_SERVER_ERROR, + } + })?; + + let session_for_spawn = session.clone(); + let state_for_spawn = state.clone(); + let session_id_for_task = session.id.clone(); + let task = tokio::spawn(async move { + match state_for_spawn.get_agent(session_for_spawn.id.clone()).await { + Ok(agent) => agent + .load_extensions_from_session(&session_for_spawn) + .await, + Err(err) => { + tracing::warn!( + session_id = %session_for_spawn.id, + %err, + "Background extension loading could not start agent" + ); + Vec::new() + } + } + }); + state + .set_extension_loading_task(session_id_for_task, task) + .await; + + Ok(Json(session)) +} + +async fn update_agent_provider( + State(state): State>, + Json(payload): Json, +) -> Result<(), (StatusCode, String)> { + let agent = state + .get_agent_for_route(payload.session_id.clone()) + .await + .map_err(|status| (status, "No agent for session id".to_owned()))?; + + let config = Config::global(); + let model = match payload.model.or_else(|| config.get_goose_model().ok()) { + Some(model) => model, + None => return Err((StatusCode::BAD_REQUEST, "No model specified".to_owned())), + }; + + let mut model_config = ModelConfig::new(&model) + .with_canonical_limits(&payload.provider) + .with_context_limit(payload.context_limit); + + if let Some(request_params) = payload.request_params { + model_config = model_config.with_merged_request_params(request_params); + } + + let extensions = + EnabledExtensionsState::for_session(state.session_manager(), &payload.session_id, config) + .await; + + let new_provider = create(&payload.provider, extensions) + .await + .map_err(|err| { + ( + StatusCode::BAD_REQUEST, + format!("Failed to create {} provider: {err}", payload.provider), + ) + })?; + + agent + .update_provider(new_provider, model_config, &payload.session_id) + .await + .map_err(|err| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Failed to update provider: {err}"), + ) + })?; + + let mode = agent.goose_mode().await; + agent + .update_goose_mode(mode, &payload.session_id) + .await + .map_err(|err| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Failed to propagate mode to provider: {err}"), + ) + })?; + + Ok(()) +} + +async fn update_session( + State(state): State>, + Json(payload): Json, +) -> Result<(), (StatusCode, String)> { + let agent = state + .get_agent_for_route(payload.session_id.clone()) + .await + .map_err(|status| (status, "No agent for session id".to_owned()))?; + + if let Some(mode_str) = payload.goose_mode { + let mode: GooseMode = mode_str + .parse() + .map_err(|_| (StatusCode::BAD_REQUEST, format!("Invalid mode: {mode_str}")))?; + agent + .update_goose_mode(mode, &payload.session_id) + .await + .map_err(|err| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Failed to update mode: {err}"), + ) + })?; + } + + Ok(()) +} + +async fn get_tools( + State(state): State>, + Query(query): Query, +) -> Result>, StatusCode> { + let session_id = query.session_id; + let agent = state.get_agent_for_route(session_id.clone()).await?; + let goose_mode = agent.goose_mode().await; + let permission_manager = agent.config.permission_manager.clone(); + + let mut tools: Vec = agent + .list_tools(&session_id, query.extension_name) + .await + .into_iter() + .map(|tool| { + let permission = permission_manager + .get_user_permission(&tool.name) + .or_else(|| { + if goose_mode == GooseMode::SmartApprove { + permission_manager.get_smart_approve_permission(&tool.name) + } else if goose_mode == GooseMode::Approve { + Some(PermissionLevel::AskBefore) + } else { + None + } + }); + + ToolInfo::new( + &tool.name, + tool.description + .as_ref() + .map(|d| d.as_ref()) + .unwrap_or_default(), + get_parameter_names(&tool), + permission, + ) + .with_input_schema(serde_json::Value::Object( + tool.input_schema.as_ref().clone(), + )) + }) + .collect(); + tools.sort_by(|a, b| a.name.cmp(&b.name)); + + Ok(Json(tools)) +} + +async fn add_extension( + State(state): State>, + Json(payload): Json, +) -> Result<(), (StatusCode, String)> { + let agent = state + .get_agent_for_route(payload.session_id.clone()) + .await + .map_err(|status| (status, "No agent for session id".to_owned()))?; + match agent + .add_extension(payload.config, &payload.session_id) + .await + { + Ok(()) => Ok(()), + Err(err) if err.to_string().contains("Unknown extension") => { + tracing::warn!( + session_id = %payload.session_id, + error = %err, + "Skipping unknown extension during tkmind compat" + ); + Ok(()) + } + Err(err) => Err(( + StatusCode::BAD_REQUEST, + format!("Failed to add extension: {err}"), + )), + } +} + +async fn remove_extension( + State(state): State>, + Json(payload): Json, +) -> Result<(), (StatusCode, String)> { + let agent = state + .get_agent_for_route(payload.session_id.clone()) + .await + .map_err(|status| (status, "No agent for session id".to_owned()))?; + agent + .remove_extension(&payload.name, &payload.session_id) + .await + .map_err(|err| { + ( + StatusCode::BAD_REQUEST, + format!("Failed to remove extension: {err}"), + ) + })?; + Ok(()) +} + +async fn update_working_dir( + State(state): State>, + Json(payload): Json, +) -> Result<(), (StatusCode, String)> { + let working_dir = PathBuf::from(payload.working_dir.trim()); + if payload.working_dir.trim().is_empty() { + return Err(( + StatusCode::BAD_REQUEST, + "Invalid directory path".to_string(), + )); + } + + state + .session_manager() + .update(&payload.session_id) + .working_dir(working_dir.clone()) + .apply() + .await + .map_err(|err| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Failed to update working dir: {err}"), + ) + })?; + + let agent = state + .get_agent_for_route(payload.session_id.clone()) + .await + .map_err(|status| (status, "No agent for session id".to_owned()))?; + agent + .extension_manager + .update_working_dir(working_dir.as_path()) + .await; + Ok(()) +} + +async fn restart_agent( + State(state): State>, + Json(payload): Json, +) -> Result<(), (StatusCode, String)> { + let session = state + .session_manager() + .get_session(&payload.session_id, false) + .await + .map_err(|_| (StatusCode::NOT_FOUND, "Session not found".to_string()))?; + let agent = state + .get_agent_for_route(payload.session_id.clone()) + .await + .map_err(|status| (status, "No agent for session id".to_owned()))?; + agent.load_extensions_from_session(&session).await; + Ok(()) +} + +pub fn routes(state: Arc) -> Router { + Router::new() + .route("/agent/start", post(start_agent)) + .route("/agent/tools", get(get_tools)) + .route("/agent/update_provider", post(update_agent_provider)) + .route("/agent/update_session", post(update_session)) + .route("/agent/add_extension", post(add_extension)) + .route("/agent/remove_extension", post(remove_extension)) + .route("/agent/update_working_dir", post(update_working_dir)) + .route("/agent/restart", post(restart_agent)) + .with_state(state) +} diff --git a/crates/goose/src/acp/tkmind/compat.rs b/crates/goose/src/acp/tkmind/compat.rs new file mode 100644 index 000000000..2f4cc03aa --- /dev/null +++ b/crates/goose/src/acp/tkmind/compat.rs @@ -0,0 +1,102 @@ +use std::sync::Arc; + +use axum::{ + extract::{Path, State}, + http::StatusCode, + routing::{delete, get}, + Json, Router, +}; + +use crate::agents::ExtensionConfig; +use crate::config::Config; +use crate::session::{EnabledExtensionsState, Session}; + +use super::state::AppState; + +#[derive(serde::Serialize)] +struct LegacySessionListResponse { + sessions: Vec, +} + +async fn list_sessions( + State(state): State>, +) -> Result, StatusCode> { + let sessions = state + .session_manager() + .list_sessions() + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + + Ok(Json(LegacySessionListResponse { sessions })) +} + +async fn delete_session( + State(state): State>, + Path(session_id): Path, +) -> Result { + state + .session_manager() + .get_session(&session_id, false) + .await + .map_err(|_| StatusCode::NOT_FOUND)?; + + state + .session_manager() + .delete_session(&session_id) + .await + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; + + state.remove_extension_loading_task(&session_id).await; + + if let Some(bus) = state.get_event_bus(&session_id).await { + bus.cancel_all_requests().await; + state.remove_event_bus(&session_id).await; + } + + Ok(StatusCode::NO_CONTENT) +} + +#[derive(serde::Serialize)] +#[serde(rename_all = "camelCase")] +struct SessionExtensionsResponse { + extensions: Vec, +} + +async fn get_session_extensions( + State(state): State>, + Path(session_id): Path, +) -> Result, StatusCode> { + let session = state + .session_manager() + .get_session(&session_id, false) + .await + .map_err(|_| StatusCode::NOT_FOUND)?; + + let extensions = EnabledExtensionsState::extensions_or_default( + Some(&session.extension_data), + Config::global(), + ); + + Ok(Json(SessionExtensionsResponse { extensions })) +} + +async fn get_session( + State(state): State>, + Path(session_id): Path, +) -> Result, StatusCode> { + state + .session_manager() + .get_session(&session_id, false) + .await + .map(Json) + .map_err(|_| StatusCode::NOT_FOUND) +} + +pub fn routes(state: Arc) -> Router { + Router::new() + .route("/sessions", get(list_sessions)) + .route("/sessions/{session_id}", get(get_session)) + .route("/sessions/{session_id}/extensions", get(get_session_extensions)) + .route("/sessions/{session_id}", delete(delete_session)) + .with_state(state) +} diff --git a/crates/goose/src/acp/tkmind/config.rs b/crates/goose/src/acp/tkmind/config.rs new file mode 100644 index 000000000..3f0edec44 --- /dev/null +++ b/crates/goose/src/acp/tkmind/config.rs @@ -0,0 +1,240 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use axum::{ + extract::Path, + routing::{get, post, put}, + Json, Router, +}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use crate::config::declarative_providers::{ + self, CreateCustomProviderParams, LoadedProvider, UpdateCustomProviderParams, +}; +use crate::config::Config; +use crate::config::ConfigError; +use crate::providers; +use crate::providers::base::ModelInfo; + +use super::errors::ErrorResponse; +use super::state::AppState; + +#[derive(Deserialize)] +pub struct UpsertConfigQuery { + pub key: String, + pub value: Value, + pub is_secret: bool, +} + +#[derive(Deserialize)] +pub struct UpdateCustomProviderRequest { + pub engine: String, + pub display_name: String, + pub api_url: String, + pub api_key: String, + pub models: Vec, + pub supports_streaming: Option, + pub headers: Option>, + #[serde(default = "default_requires_auth")] + pub requires_auth: bool, + #[serde(default)] + pub catalog_provider_id: Option, + #[serde(default)] + pub base_path: Option, + #[serde(default)] + pub preserves_thinking: Option, +} + +fn default_requires_auth() -> bool { + true +} + +fn model_infos(models: Vec) -> Vec { + models.into_iter().map(ModelInfo::new).collect() +} + +fn normalize_custom_provider_api_key(api_key: String) -> Option { + let api_key = api_key.trim().to_string(); + (!api_key.is_empty()).then_some(api_key) +} + +#[derive(Deserialize)] +pub struct ConfigKeyQuery { + pub key: String, + pub is_secret: bool, +} + +#[derive(Serialize)] +pub struct MaskedSecret { + pub masked_value: String, +} + +#[derive(Serialize)] +#[serde(untagged)] +pub enum ConfigValueResponse { + Value(Value), + MaskedValue(MaskedSecret), +} + +fn mask_config_secret(secret: Value) -> String { + const SECRET_MASK_SHOW_LEN: usize = 8; + let as_string = match secret { + Value::String(s) => s, + _ => serde_json::to_string(&secret).unwrap_or_else(|_| secret.to_string()), + }; + let chars: Vec<_> = as_string.chars().collect(); + let show_len = std::cmp::min(chars.len() / 2, SECRET_MASK_SHOW_LEN); + let visible: String = chars.iter().take(show_len).collect(); + let mask = "*".repeat(chars.len().saturating_sub(show_len)); + format!("{visible}{mask}") +} + +async fn read_config( + Json(query): Json, +) -> Result, ErrorResponse> { + let config = Config::global(); + + if query.key == "GOOSE_PROVIDER" || query.key == "active_provider" { + let value = config + .get_goose_provider() + .map(Value::String) + .unwrap_or(Value::Null); + return Ok(Json(ConfigValueResponse::Value(value))); + } + if query.key == "GOOSE_MODEL" { + let value = config + .get_goose_model() + .map(Value::String) + .unwrap_or(Value::Null); + return Ok(Json(ConfigValueResponse::Value(value))); + } + + let response_value = match config.get(&query.key, query.is_secret) { + Ok(value) if query.is_secret => ConfigValueResponse::MaskedValue(MaskedSecret { + masked_value: mask_config_secret(value), + }), + Ok(value) => ConfigValueResponse::Value(value), + Err(ConfigError::NotFound(_)) => ConfigValueResponse::Value(Value::Null), + Err(err) => return Err(ErrorResponse::internal(err.to_string())), + }; + Ok(Json(response_value)) +} + +#[derive(Serialize)] +pub struct CreateCustomProviderResponse { + pub provider_name: String, +} + +async fn upsert_config(Json(query): Json) -> Result, ErrorResponse> { + let config = Config::global(); + + if query.key == "GOOSE_PROVIDER" { + if let Some(name) = query.value.as_str() { + let model = crate::config::get_provider_entry(config, name) + .map(|entry| entry.model) + .or_else(|| config.get_goose_model().ok()) + .unwrap_or_default(); + crate::config::set_active_provider(config, name, &model).map_err(|err| { + ErrorResponse::internal(format!("Failed to set active provider: {err}")) + })?; + return Ok(Json(Value::String(format!("Upserted key {}", query.key)))); + } + } + if query.key == "GOOSE_MODEL" { + if let Some(model) = query.value.as_str() { + if let Ok(provider) = config.get_goose_provider() { + crate::config::set_active_provider(config, &provider, model).map_err(|err| { + ErrorResponse::internal(format!("Failed to set active model: {err}")) + })?; + return Ok(Json(Value::String(format!("Upserted key {}", query.key)))); + } + } + } + + config + .set(&query.key, &query.value, query.is_secret) + .map_err(|err| ErrorResponse::internal(err.to_string()))?; + Ok(Json(Value::String(format!("Upserted key {}", query.key)))) +} + +async fn create_custom_provider( + Json(request): Json, +) -> Result, ErrorResponse> { + let config = declarative_providers::create_custom_provider(CreateCustomProviderParams { + engine: request.engine, + display_name: request.display_name, + api_url: request.api_url, + api_key: normalize_custom_provider_api_key(request.api_key), + models: model_infos(request.models), + supports_streaming: request.supports_streaming, + headers: request.headers, + requires_auth: request.requires_auth, + catalog_provider_id: request.catalog_provider_id, + base_path: request.base_path, + preserves_thinking: request.preserves_thinking, + auth: None, + }) + .map_err(|err| ErrorResponse::bad_request(err.to_string()))?; + + providers::refresh_custom_providers() + .await + .map_err(|err| ErrorResponse::internal(err.to_string()))?; + + Ok(Json(CreateCustomProviderResponse { + provider_name: config.id().to_string(), + })) +} + +async fn get_custom_provider( + Path(id): Path, +) -> Result, ErrorResponse> { + let loaded_provider = declarative_providers::load_provider(id.as_str()).map_err(|err| { + ErrorResponse::not_found(format!("Custom provider '{id}' not found: {err}")) + })?; + + Ok(Json(loaded_provider)) +} + +async fn update_custom_provider( + Path(id): Path, + Json(request): Json, +) -> Result, ErrorResponse> { + declarative_providers::update_custom_provider(UpdateCustomProviderParams { + id: id.clone(), + engine: request.engine, + display_name: request.display_name, + api_url: request.api_url, + api_key: normalize_custom_provider_api_key(request.api_key), + models: model_infos(request.models), + supports_streaming: request.supports_streaming, + headers: request.headers, + requires_auth: request.requires_auth, + catalog_provider_id: request.catalog_provider_id, + base_path: request.base_path, + preserves_thinking: request.preserves_thinking, + auth: None, + }) + .map_err(|err| { + if err.to_string().contains("not found") { + ErrorResponse::not_found(err.to_string()) + } else { + ErrorResponse::bad_request(err.to_string()) + } + })?; + + providers::refresh_custom_providers() + .await + .map_err(|err| ErrorResponse::internal(err.to_string()))?; + + Ok(Json(format!("Updated custom provider: {id}"))) +} + +pub fn routes(_state: Arc) -> Router { + Router::new() + .route("/config/read", post(read_config)) + .route("/config/upsert", post(upsert_config)) + .route("/config/custom-providers", post(create_custom_provider)) + .route("/config/custom-providers/{id}", get(get_custom_provider)) + .route("/config/custom-providers/{id}", put(update_custom_provider)) +} diff --git a/crates/goose/src/acp/tkmind/errors.rs b/crates/goose/src/acp/tkmind/errors.rs new file mode 100644 index 000000000..d7a3ceed8 --- /dev/null +++ b/crates/goose/src/acp/tkmind/errors.rs @@ -0,0 +1,51 @@ +use axum::{ + http::StatusCode, + response::{IntoResponse, Response}, + Json, +}; + +#[derive(Debug, serde::Serialize)] +pub struct ErrorResponse { + pub message: String, + #[serde(skip)] + pub status: StatusCode, +} + +impl ErrorResponse { + pub fn internal(message: impl Into) -> Self { + Self { + message: message.into(), + status: StatusCode::INTERNAL_SERVER_ERROR, + } + } + + pub fn bad_request(message: impl Into) -> Self { + Self { + message: message.into(), + status: StatusCode::BAD_REQUEST, + } + } + + pub fn not_found(message: impl Into) -> Self { + Self { + message: message.into(), + status: StatusCode::NOT_FOUND, + } + } +} + +impl IntoResponse for ErrorResponse { + fn into_response(self) -> Response { + if self.status.is_server_error() { + tracing::error!(status = %self.status, message = %self.message, "server error response"); + } else if self.status.is_client_error() { + tracing::warn!(status = %self.status, message = %self.message, "client error response"); + } + + let body = Json(serde_json::json!({ + "message": self.message, + })); + + (self.status, body).into_response() + } +} diff --git a/crates/goose/src/acp/tkmind/events.rs b/crates/goose/src/acp/tkmind/events.rs new file mode 100644 index 000000000..c8cdf7ab0 --- /dev/null +++ b/crates/goose/src/acp/tkmind/events.rs @@ -0,0 +1,97 @@ +use crate::conversation::message::{Message, MessageContent, TokenState}; +use crate::conversation::Conversation; +use crate::session::SessionManager; +use rmcp::model::ServerNotification; +use serde::{Deserialize, Serialize}; + +pub fn track_tool_telemetry(content: &MessageContent, all_messages: &[Message]) { + match content { + MessageContent::ToolRequest(tool_request) => { + if let Ok(tool_call) = &tool_request.tool_call { + tracing::info!( + monotonic_counter.goose.tool_calls = 1, + tool_name = %tool_call.name, + "Tool call started" + ); + } + } + MessageContent::ToolResponse(tool_response) => { + let tool_name = all_messages + .iter() + .rev() + .find_map(|msg| { + msg.content.iter().find_map(|c| { + if let MessageContent::ToolRequest(req) = c { + if req.id == tool_response.id { + if let Ok(tool_call) = &req.tool_call { + Some(tool_call.name.clone()) + } else { + None + } + } else { + None + } + } else { + None + } + }) + }) + .unwrap_or_else(|| "unknown".to_string().into()); + + let success = tool_response.tool_result.is_ok(); + let result_status = if success { "success" } else { "error" }; + + tracing::info!( + monotonic_counter.goose.tool_completions = 1, + tool_name = %tool_name, + result = %result_status, + "Tool call completed" + ); + } + _ => {} + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "type")] +pub enum MessageEvent { + Message { + message: Message, + token_state: TokenState, + }, + Error { + error: String, + }, + Finish { + reason: String, + token_state: TokenState, + }, + Notification { + request_id: String, + message: ServerNotification, + }, + UpdateConversation { + conversation: Conversation, + }, + /// Sent at the start of an SSE stream to inform the client about + /// in-flight requests it can reattach to. + ActiveRequests { + request_ids: Vec, + }, + Ping, +} + +pub async fn get_token_state(session_manager: &SessionManager, session_id: &str) -> TokenState { + session_manager + .get_session(session_id, false) + .await + .map(|session| TokenState::from(&session)) + .inspect_err(|e| { + tracing::warn!( + "Failed to fetch session token state for {}: {}", + session_id, + e + ); + }) + .unwrap_or_default() +} diff --git a/crates/goose/src/acp/tkmind/memory.rs b/crates/goose/src/acp/tkmind/memory.rs new file mode 100644 index 000000000..39874b64c --- /dev/null +++ b/crates/goose/src/acp/tkmind/memory.rs @@ -0,0 +1,258 @@ +use std::{ + path::{Path, PathBuf}, + sync::Arc, + time::Duration, +}; + +use axum::{extract::State, routing::post, Json, Router}; +use tokio::{process::Command, time::timeout}; + +use crate::agents::platform_extensions::projectmemory::ProjectMemoryState; +use crate::session::extension_data::ExtensionState; + +use super::{errors::ErrorResponse, state::AppState}; + +const DEFAULT_QUERY: &str = + "project architecture conventions current goals decisions risks and recent work"; +const DEFAULT_MAX_CHARS: usize = 6_000; + +#[derive(serde::Deserialize)] +#[serde(rename_all = "camelCase")] +struct HarnessBootstrapRequest { + session_id: String, + query: Option, + #[serde(default)] + force: bool, +} + +#[derive(serde::Serialize)] +#[serde(rename_all = "camelCase")] +struct HarnessBootstrapResponse { + summary: String, + source: String, + refreshed: bool, +} + +#[derive(serde::Deserialize)] +#[serde(rename_all = "camelCase")] +struct HarnessRememberRequest { + session_id: String, + content: String, + title: Option, +} + +#[derive(serde::Serialize)] +struct HarnessRememberResponse { + remembered: bool, +} + +fn max_chars() -> usize { + std::env::var("GOOSE_HARNESS_BOOTSTRAP_MAX_CHARS") + .ok() + .and_then(|value| value.parse().ok()) + .filter(|value| *value >= 1_000) + .unwrap_or(DEFAULT_MAX_CHARS) +} + +fn truncate(value: &str, max_chars: usize) -> String { + value.chars().take(max_chars).collect() +} + +fn harness_binary() -> PathBuf { + std::env::var_os("GOOSE_HARNESS_BIN") + .map(PathBuf::from) + .or_else(|| { + std::env::var_os("HOME") + .map(|home| PathBuf::from(home).join(".codex/harness/bin/harness")) + }) + .unwrap_or_else(|| PathBuf::from("harness")) +} + +async fn recall(repo: &str, query: &str, working_dir: &Path) -> Option { + let output = timeout( + Duration::from_secs(8), + Command::new(harness_binary()) + .args(["recall", "--repo", repo, "--query", query]) + .current_dir(working_dir) + .output(), + ) + .await + .ok()? + .ok()?; + if !output.status.success() { + return None; + } + let value = String::from_utf8_lossy(&output.stdout); + let value = truncate(value.trim(), max_chars() / 2); + (!value.is_empty()).then_some(value) +} + +async fn project_files(working_dir: &Path, max: usize) -> Option { + let mut content = String::new(); + for name in ["AGENTS.md", ".tkmindhints", ".goosehints", "CLAUDE.md"] { + if content.chars().count() >= max { + break; + } + let Ok(value) = tokio::fs::read_to_string(working_dir.join(name)).await else { + continue; + }; + let remaining = max.saturating_sub(content.chars().count()); + content.push_str(&format!( + "## {name}\n{}\n\n", + truncate(value.trim(), remaining) + )); + } + (!content.trim().is_empty()).then_some(content.trim().to_string()) +} + +async fn bootstrap(working_dir: &Path, query: &str) -> (String, String) { + let repo = working_dir + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or("unknown"); + let recalled = recall(repo, query, working_dir).await; + let remaining = + max_chars().saturating_sub(recalled.as_deref().map_or(0, |value| value.chars().count())); + let files = project_files(working_dir, remaining).await; + let mut sections = Vec::new(); + let mut sources = Vec::new(); + if let Some(value) = recalled { + sections.push(format!("## Recalled project knowledge\n{value}")); + sources.push("harness"); + } + if let Some(value) = files { + sections.push(format!("## Project instructions\n{value}")); + sources.push("project_files"); + } + ( + truncate(§ions.join("\n\n"), max_chars()), + if sources.is_empty() { + "none".to_string() + } else { + sources.join("+") + }, + ) +} + +async fn harness_bootstrap( + State(state): State>, + Json(request): Json, +) -> Result, ErrorResponse> { + let mut session = state + .session_manager() + .get_session(&request.session_id, false) + .await + .map_err(|_| ErrorResponse::not_found("Session not found"))?; + if !request.force { + if let Some(memory) = ProjectMemoryState::from_extension_data(&session.extension_data) { + if !memory.summary.trim().is_empty() { + return Ok(Json(HarnessBootstrapResponse { + summary: memory.summary, + source: memory.source.unwrap_or_else(|| "session".to_string()), + refreshed: false, + })); + } + } + } + let query = request + .query + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or(DEFAULT_QUERY); + let (summary, source) = bootstrap(&session.working_dir, query).await; + let refreshed = !summary.is_empty(); + if !summary.is_empty() { + ProjectMemoryState { + summary: summary.clone(), + source: Some(source.clone()), + updated_at: None, + } + .to_extension_data(&mut session.extension_data) + .map_err(|error| ErrorResponse::internal(error.to_string()))?; + state + .session_manager() + .update(&request.session_id) + .extension_data(session.extension_data) + .apply() + .await + .map_err(|error| ErrorResponse::internal(error.to_string()))?; + } + Ok(Json(HarnessBootstrapResponse { + summary, + source, + refreshed, + })) +} + +async fn harness_remember( + State(state): State>, + Json(request): Json, +) -> Result, ErrorResponse> { + let session = state + .session_manager() + .get_session(&request.session_id, false) + .await + .map_err(|_| ErrorResponse::not_found("Session not found"))?; + let content = request.content.trim(); + if content.is_empty() || content.chars().count() > max_chars() { + return Err(ErrorResponse::bad_request("Invalid memory content")); + } + let repo = session + .working_dir + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or("unknown"); + let title = request + .title + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or("H5 conversation memory"); + let remembered = match timeout( + Duration::from_secs(8), + Command::new(harness_binary()) + .args([ + "remember", + "--repo", + repo, + "--kind", + "conversation_memory", + "--title", + title, + "--content", + content, + ]) + .current_dir(&session.working_dir) + .output(), + ) + .await + { + Ok(Ok(output)) if output.status.success() => true, + Ok(Ok(output)) => { + tracing::warn!( + session_id = %request.session_id, + status = ?output.status.code(), + stderr = %String::from_utf8_lossy(&output.stderr).trim(), + "Harness remember failed" + ); + false + } + Ok(Err(error)) => { + tracing::warn!(session_id = %request.session_id, %error, "Harness remember could not start"); + false + } + Err(_) => { + tracing::warn!(session_id = %request.session_id, "Harness remember timed out"); + false + } + }; + Ok(Json(HarnessRememberResponse { remembered })) +} + +pub fn routes(state: Arc) -> Router { + Router::new() + .route("/agent/harness_bootstrap", post(harness_bootstrap)) + .route("/agent/harness_remember", post(harness_remember)) + .with_state(state) +} diff --git a/crates/goose/src/acp/tkmind/mod.rs b/crates/goose/src/acp/tkmind/mod.rs new file mode 100644 index 000000000..dec334a1c --- /dev/null +++ b/crates/goose/src/acp/tkmind/mod.rs @@ -0,0 +1,25 @@ +mod agent; +mod compat; +mod config; +mod errors; +mod events; +mod memory; +mod session_event_bus; +mod session_events; +pub mod state; + +use std::sync::Arc; + +use axum::Router; + +use self::state::AppState; + +/// Legacy Portal/H5 REST routes preserved from goose-server. +pub fn routes(state: Arc) -> Router { + Router::new() + .merge(agent::routes(state.clone())) + .merge(compat::routes(state.clone())) + .merge(config::routes(state.clone())) + .merge(session_events::routes(state.clone())) + .merge(memory::routes(state)) +} diff --git a/crates/goose/src/acp/tkmind/session_event_bus.rs b/crates/goose/src/acp/tkmind/session_event_bus.rs new file mode 100644 index 000000000..75ab8b21b --- /dev/null +++ b/crates/goose/src/acp/tkmind/session_event_bus.rs @@ -0,0 +1,299 @@ +use std::collections::{HashMap, VecDeque}; +use std::sync::atomic::{AtomicU64, Ordering}; + +use tokio::sync::{broadcast, Mutex}; +use tokio_util::sync::CancellationToken; + +use super::events::MessageEvent; + +const BROADCAST_CAPACITY: usize = 256; +const REPLAY_BUFFER_CAPACITY: usize = 512; + +/// Error returned by [`SessionEventBus::subscribe`]. +#[derive(Debug)] +pub enum SubscribeError { + /// The client's `Last-Event-ID` has been evicted from the replay buffer, + /// so events have been irrecoverably lost. + ClientTooFarBehind, +} + +#[derive(Clone, Debug)] +pub struct SessionEvent { + /// Monotonic sequence number, written as SSE `id:` frame (not in JSON payload). + pub seq: u64, + /// None for Ping events, Some for events associated with a specific request. + pub request_id: Option, + /// The event payload. + pub event: MessageEvent, +} + +pub struct SessionEventBus { + tx: broadcast::Sender, + buffer: Mutex>, + next_seq: AtomicU64, + active_requests: Mutex>, +} + +impl SessionEventBus { + pub fn new() -> Self { + let (tx, _) = broadcast::channel(BROADCAST_CAPACITY); + Self { + tx, + buffer: Mutex::new(VecDeque::with_capacity(REPLAY_BUFFER_CAPACITY)), + next_seq: AtomicU64::new(1), + active_requests: Mutex::new(HashMap::new()), + } + } + + /// Publish an event to the bus. Assigns a monotonic sequence number. + /// + /// The sequence ID is assigned under the buffer lock so that concurrent + /// callers cannot reorder events (i.e. seq=2 published before seq=1). + pub async fn publish(&self, request_id: Option, event: MessageEvent) -> u64 { + let session_event = { + let mut buf = self.buffer.lock().await; + let seq = self.next_seq.fetch_add(1, Ordering::Relaxed); + let session_event = SessionEvent { + seq, + request_id, + event, + }; + buf.push_back(session_event.clone()); + while buf.len() > REPLAY_BUFFER_CAPACITY { + buf.pop_front(); + } + session_event + }; + + // Send on broadcast channel (ignore error if no subscribers) + let _ = self.tx.send(session_event.clone()); + + session_event.seq + } + + /// Subscribe to live events. If `last_event_id` is provided, replay buffered + /// events with seq > last_event_id. Returns (replay_events, replay_max_seq, live_receiver). + /// + /// Returns `Err(SubscribeError::ClientTooFarBehind)` when `last_event_id` + /// refers to an event that has already been evicted from the replay buffer, + /// meaning the client has irrecoverably missed events. + /// + /// The live receiver is created *before* snapshotting the buffer so that + /// no event can fall into the gap between the two steps. The caller must + /// skip live events with `seq <= replay_max_seq` to deduplicate. + pub async fn subscribe( + &self, + last_event_id: Option, + ) -> Result<(Vec, u64, broadcast::Receiver), SubscribeError> { + // Subscribe first so that any event published while we hold the + // buffer lock is guaranteed to appear in `rx` (possibly duplicating + // a replay entry). The caller deduplicates via replay_max_seq. + let rx = self.tx.subscribe(); + + let (replay, replay_max_seq) = { + let buf = self.buffer.lock().await; + let buf_max = buf.back().map(|e| e.seq).unwrap_or(0); + let buf_min = buf.front().map(|e| e.seq).unwrap_or(0); + let last_id = last_event_id.unwrap_or(0); + + // If the client sent a Last-Event-ID that has been evicted from + // the buffer, they have irrecoverably missed events. + if last_id > 0 && buf_min > 0 && last_id < buf_min { + return Err(SubscribeError::ClientTooFarBehind); + } + + // Clamp to the actual buffer max so a stale Last-Event-ID + // (e.g. from before a server restart) doesn't suppress live events. + let events: Vec<_> = buf.iter().filter(|e| e.seq > last_id).cloned().collect(); + let max_seq = events.last().map(|e| e.seq).unwrap_or(last_id.min(buf_max)); + (events, max_seq) + }; + + Ok((replay, replay_max_seq, rx)) + } + + /// Return the IDs of all currently active (in-flight) requests. + pub async fn active_request_ids(&self) -> Vec { + let requests = self.active_requests.lock().await; + requests.keys().cloned().collect() + } + + #[cfg(test)] + pub async fn register_request(&self, request_id: String) -> CancellationToken { + let token = CancellationToken::new(); + let mut requests = self.active_requests.lock().await; + requests.insert(request_id, token.clone()); + token + } + + /// Atomically check no requests are active and register one. Returns Err if busy. + pub async fn try_register_request( + &self, + request_id: String, + ) -> Result { + let mut requests = self.active_requests.lock().await; + requests.retain(|_id, token| !token.is_cancelled()); + if !requests.is_empty() { + return Err("Session already has an active request".into()); + } + let token = CancellationToken::new(); + requests.insert(request_id, token.clone()); + Ok(token) + } + + /// Cancel a specific request by request_id. + pub async fn cancel_request(&self, request_id: &str) -> bool { + let requests = self.active_requests.lock().await; + if let Some(token) = requests.get(request_id) { + token.cancel(); + true + } else { + false + } + } + + /// Cancel all active requests (e.g. when deleting a session). + pub async fn cancel_all_requests(&self) { + let requests = self.active_requests.lock().await; + for token in requests.values() { + token.cancel(); + } + } + + /// Remove the cancellation token for a completed request. + pub async fn cleanup_request(&self, request_id: &str) { + let mut requests = self.active_requests.lock().await; + requests.remove(request_id); + } +} + +impl Default for SessionEventBus { + fn default() -> Self { + Self::new() + } +} + +pub struct RequestGuard { + bus: std::sync::Arc, + request_id: String, + disarmed: bool, +} + +impl RequestGuard { + pub fn new(bus: std::sync::Arc, request_id: String) -> Self { + Self { + bus, + request_id, + disarmed: false, + } + } + + pub fn disarm(&mut self) { + self.disarmed = true; + } +} + +impl Drop for RequestGuard { + fn drop(&mut self) { + if !self.disarmed { + let bus = self.bus.clone(); + let request_id = self.request_id.clone(); + tokio::spawn(async move { + bus.cleanup_request(&request_id).await; + }); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::conversation::message::TokenState; + + #[tokio::test] + async fn test_publish_and_subscribe() { + let bus = SessionEventBus::new(); + + bus.publish(Some("req-1".to_string()), MessageEvent::Ping) + .await; + bus.publish( + Some("req-1".to_string()), + MessageEvent::Finish { + reason: "stop".to_string(), + token_state: TokenState::default(), + }, + ) + .await; + + let (replay, replay_max_seq, _rx) = bus.subscribe(Some(0)).await.unwrap(); + assert_eq!(replay.len(), 2); + assert_eq!(replay[0].seq, 1); + assert_eq!(replay[1].seq, 2); + assert_eq!(replay_max_seq, 2); + } + + #[tokio::test] + async fn test_subscribe_with_last_event_id() { + let bus = SessionEventBus::new(); + + bus.publish(None, MessageEvent::Ping).await; + bus.publish(None, MessageEvent::Ping).await; + bus.publish(None, MessageEvent::Ping).await; + + let (replay, replay_max_seq, _rx) = bus.subscribe(Some(2)).await.unwrap(); + assert_eq!(replay.len(), 1); + assert_eq!(replay[0].seq, 3); + assert_eq!(replay_max_seq, 3); + } + + #[tokio::test] + async fn test_subscribe_without_last_event_id_replays_all() { + let bus = SessionEventBus::new(); + + bus.publish(None, MessageEvent::Ping).await; + bus.publish(None, MessageEvent::Ping).await; + + let (replay, replay_max_seq, _rx) = bus.subscribe(None).await.unwrap(); + assert_eq!(replay.len(), 2); + assert_eq!(replay_max_seq, 2); + } + + #[tokio::test] + async fn test_subscribe_with_stale_last_event_id() { + let bus = SessionEventBus::new(); + + bus.publish(None, MessageEvent::Ping).await; + bus.publish(None, MessageEvent::Ping).await; + bus.publish(None, MessageEvent::Ping).await; + + let (replay, replay_max_seq, _rx) = bus.subscribe(Some(9999)).await.unwrap(); + assert_eq!(replay.len(), 0); + assert_eq!(replay_max_seq, 3); + } + + #[tokio::test] + async fn test_cancel_request() { + let bus = SessionEventBus::new(); + + let token = bus.register_request("req-1".to_string()).await; + assert!(!token.is_cancelled()); + + let cancelled = bus.cancel_request("req-1").await; + assert!(cancelled); + assert!(token.is_cancelled()); + + let cancelled = bus.cancel_request("req-999").await; + assert!(!cancelled); + } + + #[tokio::test] + async fn test_cleanup_request() { + let bus = SessionEventBus::new(); + + bus.register_request("req-1".to_string()).await; + bus.cleanup_request("req-1").await; + + let cancelled = bus.cancel_request("req-1").await; + assert!(!cancelled); + } +} diff --git a/crates/goose/src/acp/tkmind/session_events.rs b/crates/goose/src/acp/tkmind/session_events.rs new file mode 100644 index 000000000..ac140978f --- /dev/null +++ b/crates/goose/src/acp/tkmind/session_events.rs @@ -0,0 +1,566 @@ +use std::{ + convert::Infallible, + pin::Pin, + sync::Arc, + task::{Context, Poll}, + time::Duration, +}; + +use axum::{ + extract::{DefaultBodyLimit, Path, State}, + body::Bytes, + http::{self, HeaderMap}, + response::IntoResponse, + routing::{get, post}, + Json, Router, +}; +use futures::{stream::StreamExt, Stream}; +use serde::{Deserialize, Serialize}; +use tokio::sync::mpsc; +use tokio::time::timeout; +use tokio_stream::wrappers::ReceiverStream; + +use crate::agents::{AgentEvent, SessionConfig}; +use crate::conversation::message::{ActionRequiredData, Message, MessageContent}; +use crate::conversation::Conversation; + +use super::errors::ErrorResponse; +use super::events::{get_token_state, track_tool_telemetry, MessageEvent}; +use super::session_event_bus::RequestGuard; +use super::state::AppState; + +#[derive(Debug, Deserialize, Serialize)] +pub struct SessionReplyRequest { + /// Client-generated UUIDv7 identifying this request. + pub request_id: String, + pub user_message: Message, + #[serde(default)] + pub override_conversation: Option>, +} + +#[derive(Debug, Serialize)] +pub struct SessionReplyResponse { + pub request_id: String, +} + +#[derive(Debug, Deserialize, Serialize)] +pub struct CancelRequest { + pub request_id: String, +} + +/// An SSE response that includes `id:` lines for Last-Event-ID reconnection. +pub struct SseEventStream { + rx: ReceiverStream, +} + +impl SseEventStream { + fn new(rx: ReceiverStream) -> Self { + Self { rx } + } +} + +impl Stream for SseEventStream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.rx) + .poll_next(cx) + .map(|opt| opt.map(|s| Ok(Bytes::from(s)))) + } +} + +impl IntoResponse for SseEventStream { + fn into_response(self) -> axum::response::Response { + let body = axum::body::Body::from_stream(self); + http::Response::builder() + .header("Content-Type", "text/event-stream") + .header("Cache-Control", "no-cache") + .header("Connection", "keep-alive") + .body(body) + .unwrap() + } +} + +fn format_sse_event(seq: u64, json: &str) -> String { + format!("id: {}\ndata: {}\n\n", seq, json) +} + +fn serialize_session_event(seq: u64, request_id: Option<&str>, event: &MessageEvent) -> String { + let mut event_json = serde_json::to_value(event).unwrap_or_else( + |e| serde_json::json!({"type": "Error", "error": format!("Serialization error: {}", e)}), + ); + + if let Some(rid) = request_id { + if let serde_json::Value::Object(ref mut map) = event_json { + map.insert( + "chat_request_id".to_string(), + serde_json::Value::String(rid.to_string()), + ); + map.entry("request_id") + .or_insert_with(|| serde_json::Value::String(rid.to_string())); + } + } + + let json_str = serde_json::to_string(&event_json).unwrap_or_default(); + format_sse_event(seq, &json_str) +} + +pub async fn session_events( + State(state): State>, + Path(session_id): Path, + headers: HeaderMap, +) -> Result { + state + .session_manager() + .get_session(&session_id, false) + .await + .map_err(|_| axum::http::StatusCode::NOT_FOUND)?; + + let last_event_id: Option = headers + .get("Last-Event-ID") + .and_then(|v| v.to_str().ok()) + .and_then(|s| s.parse().ok()); + + let bus = state.get_or_create_event_bus(&session_id).await; + + let (replay, replay_max_seq, mut live_rx) = match bus.subscribe(last_event_id).await { + Ok(result) => result, + Err(_) => { + let (tx, rx) = mpsc::channel::(1); + let stream = ReceiverStream::new(rx); + let seq = 0; + let error_event = MessageEvent::Error { + error: "Client too far behind — reload conversation".to_string(), + }; + let frame = serialize_session_event(seq, None, &error_event); + tokio::spawn(async move { + let _ = tx.send(frame).await; + }); + return Ok(SseEventStream::new(stream)); + } + }; + + let (tx, rx) = mpsc::channel::(256); + let stream = ReceiverStream::new(rx); + let task_bus = bus.clone(); + + tokio::spawn(async move { + let bus = task_bus; + + let active_ids = bus.active_request_ids().await; + if !active_ids.is_empty() { + let event = MessageEvent::ActiveRequests { + request_ids: active_ids, + }; + let json_str = serde_json::to_string(&serde_json::to_value(&event).unwrap_or_default()) + .unwrap_or_default(); + let frame = format!("data: {}\n\n", json_str); + if tx.send(frame).await.is_err() { + return; + } + } + + for event in &replay { + let frame = + serialize_session_event(event.seq, event.request_id.as_deref(), &event.event); + if tx.send(frame).await.is_err() { + return; + } + } + + let mut heartbeat_interval = tokio::time::interval(Duration::from_millis(500)); + let mut heartbeat_seq = 0u64; + + loop { + tokio::select! { + _ = heartbeat_interval.tick() => { + let frame = format!(": ping {}\n\n", heartbeat_seq); + heartbeat_seq += 1; + if tx.send(frame).await.is_err() { + return; + } + } + result = live_rx.recv() => { + match result { + Ok(event) => { + if event.seq <= replay_max_seq { + continue; + } + let frame = serialize_session_event( + event.seq, + event.request_id.as_deref(), + &event.event, + ); + if tx.send(frame).await.is_err() { + return; + } + } + Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => { + tracing::warn!( + "SSE subscriber lagged by {} events, closing stream so client reconnects with Last-Event-ID", + n + ); + return; + } + Err(tokio::sync::broadcast::error::RecvError::Closed) => { + return; + } + } + } + } + } + }); + + Ok(SseEventStream::new(stream)) +} + +#[allow(clippy::too_many_lines)] +pub async fn session_reply( + State(state): State>, + Path(session_id): Path, + Json(request): Json, +) -> Result, ErrorResponse> { + let request_id = request.request_id.clone(); + + if uuid::Uuid::parse_str(&request_id).is_err() { + return Err(ErrorResponse::bad_request( + "request_id must be a valid UUID", + )); + } + + let session_data = state + .session_manager() + .get_session(&session_id, false) + .await + .map_err(|_| ErrorResponse::not_found(format!("Session {} not found", session_id)))?; + + let session_start = std::time::Instant::now(); + + tracing::info!( + monotonic_counter.goose.session_starts = 1, + session_type = "app", + interface = "ui", + "Session started" + ); + + if let Some(ref recipe) = session_data.recipe { + if state.mark_recipe_run_if_absent(&session_id).await { + tracing::info!( + monotonic_counter.goose.recipe_runs = 1, + recipe_name = %recipe.title, + recipe_version = %recipe.version, + session_type = "app", + interface = "ui", + "Recipe execution started" + ); + } + } + + let user_message = request.user_message; + let override_conversation = request.override_conversation; + + let is_elicitation_response = user_message.content.iter().any(|c| { + matches!( + c, + MessageContent::ActionRequired(ar) + if matches!( + ar.data, + ActionRequiredData::ElicitationResponse { .. } + ) + ) + }); + + if is_elicitation_response { + let agent = state + .get_agent(session_id.clone()) + .await + .map_err(|e| ErrorResponse::internal(e.to_string()))?; + let session_config = SessionConfig { + id: session_id.clone(), + schedule_id: session_data.schedule_id.clone(), + max_turns: None, + retry_config: None, + }; + let _ = agent + .reply(user_message, session_config, None) + .await + .map_err(|e| ErrorResponse::internal(e.to_string()))?; + return Ok(Json(SessionReplyResponse { request_id })); + } + + let bus = state.get_or_create_event_bus(&session_id).await; + + let cancel_token = bus + .try_register_request(request_id.clone()) + .await + .map_err(|_| { + ErrorResponse::bad_request("Session already has an active request. Cancel it first.") + })?; + + let task_state = state.clone(); + let task_session_id = session_id.clone(); + let task_request_id = request_id.clone(); + let task_cancel = cancel_token.clone(); + let task_bus = bus.clone(); + + drop(tokio::spawn(async move { + let mut _guard = RequestGuard::new(task_bus.clone(), task_request_id.clone()); + + let publish = |rid: Option, event: MessageEvent| { + let bus = task_bus.clone(); + async move { + bus.publish(rid, event).await; + } + }; + + let agent = match task_state.get_agent(task_session_id.clone()).await { + Ok(agent) => agent, + Err(e) => { + tracing::error!("Failed to get session agent: {}", e); + publish( + Some(task_request_id.clone()), + MessageEvent::Error { + error: format!("Failed to get session agent: {}", e), + }, + ) + .await; + return; + } + }; + + let session = match task_state + .session_manager() + .get_session(&task_session_id, true) + .await + { + Ok(metadata) => metadata, + Err(e) => { + tracing::error!("Failed to read session for {}: {}", task_session_id, e); + publish( + Some(task_request_id.clone()), + MessageEvent::Error { + error: format!("Failed to read session: {}", e), + }, + ) + .await; + return; + } + }; + + let session_config = SessionConfig { + id: task_session_id.clone(), + schedule_id: session.schedule_id.clone(), + max_turns: None, + retry_config: None, + }; + + let mut all_messages = match override_conversation { + Some(history) => { + let conv = Conversation::new_unvalidated(history); + if let Err(e) = task_state + .session_manager() + .replace_conversation(&task_session_id, &conv) + .await + { + tracing::warn!( + "Failed to replace session conversation for {}: {}", + task_session_id, + e + ); + } + conv + } + None => session.conversation.unwrap_or_default(), + }; + all_messages.push(user_message.clone()); + + let mut stream = match agent + .reply( + user_message.clone(), + session_config, + Some(task_cancel.clone()), + ) + .await + { + Ok(stream) => stream, + Err(e) => { + tracing::error!("Failed to start reply stream: {:?}", e); + publish( + Some(task_request_id.clone()), + MessageEvent::Error { + error: e.to_string(), + }, + ) + .await; + return; + } + }; + + loop { + tokio::select! { + _ = task_cancel.cancelled() => { + tracing::info!("Agent task cancelled for request {}", task_request_id); + break; + } + response = timeout(Duration::from_millis(500), stream.next()) => { + match response { + Ok(Some(Ok(AgentEvent::Message(message)))) => { + for content in &message.content { + track_tool_telemetry(content, all_messages.messages()); + } + all_messages.push(message.clone()); + let token_state = get_token_state( + task_state.session_manager(), + &task_session_id, + ) + .await; + publish( + Some(task_request_id.clone()), + MessageEvent::Message { + message, + token_state, + }, + ) + .await; + } + Ok(Some(Ok(AgentEvent::HistoryReplaced(new_messages)))) => { + all_messages = new_messages.clone(); + publish( + Some(task_request_id.clone()), + MessageEvent::UpdateConversation { + conversation: new_messages, + }, + ) + .await; + } + Ok(Some(Ok(AgentEvent::McpNotification((notification_request_id, n))))) => { + publish( + Some(task_request_id.clone()), + MessageEvent::Notification { + request_id: notification_request_id, + message: n, + }, + ) + .await; + } + Ok(Some(Ok(AgentEvent::Usage(_)))) | Ok(Some(Ok(AgentEvent::MessageUsage { .. }))) => { + continue; + } + Ok(Some(Err(e))) => { + tracing::error!("Error processing message: {}", e); + publish( + Some(task_request_id.clone()), + MessageEvent::Error { + error: e.to_string(), + }, + ) + .await; + break; + } + Ok(None) => { + break; + } + Err(_) => { + continue; + } + } + } + } + } + + let session_duration = session_start.elapsed(); + + if let Ok(session) = task_state + .session_manager() + .get_session(&task_session_id, true) + .await + { + let total_tokens = session.usage.total_tokens.unwrap_or(0); + tracing::info!( + monotonic_counter.goose.session_completions = 1, + session_type = "app", + interface = "ui", + exit_type = "normal", + duration_ms = session_duration.as_millis() as u64, + total_tokens = total_tokens, + message_count = session.message_count, + "Session completed" + ); + + tracing::info!( + monotonic_counter.goose.session_duration_ms = session_duration.as_millis() as u64, + session_type = "app", + interface = "ui", + "Session duration" + ); + + if total_tokens > 0 { + tracing::info!( + monotonic_counter.goose.session_tokens = total_tokens, + session_type = "app", + interface = "ui", + "Session tokens" + ); + } + } else { + tracing::info!( + monotonic_counter.goose.session_completions = 1, + session_type = "app", + interface = "ui", + exit_type = "normal", + duration_ms = session_duration.as_millis() as u64, + total_tokens = 0u64, + message_count = all_messages.len(), + "Session completed" + ); + + tracing::info!( + monotonic_counter.goose.session_duration_ms = session_duration.as_millis() as u64, + session_type = "app", + interface = "ui", + "Session duration" + ); + } + + let final_token_state = + get_token_state(task_state.session_manager(), &task_session_id).await; + + publish( + Some(task_request_id.clone()), + MessageEvent::Finish { + reason: "stop".to_string(), + token_state: final_token_state, + }, + ) + .await; + + _guard.disarm(); + task_bus.cleanup_request(&task_request_id).await; + })); + + Ok(Json(SessionReplyResponse { request_id })) +} + +pub async fn session_cancel( + State(state): State>, + Path(session_id): Path, + Json(request): Json, +) -> axum::http::StatusCode { + let bus = match state.get_event_bus(&session_id).await { + Some(bus) => bus, + None => return axum::http::StatusCode::NOT_FOUND, + }; + bus.cancel_request(&request.request_id).await; + axum::http::StatusCode::OK +} + +pub fn routes(state: Arc) -> Router { + Router::new() + .route("/sessions/{id}/events", get(session_events)) + .route( + "/sessions/{id}/reply", + post(session_reply).layer(DefaultBodyLimit::max(50 * 1024 * 1024)), + ) + .route("/sessions/{id}/cancel", post(session_cancel)) + .with_state(state) +} diff --git a/crates/goose/src/acp/tkmind/state.rs b/crates/goose/src/acp/tkmind/state.rs new file mode 100644 index 000000000..0455bd6c5 --- /dev/null +++ b/crates/goose/src/acp/tkmind/state.rs @@ -0,0 +1,132 @@ +use std::collections::{HashMap, HashSet}; +use std::path::PathBuf; +use std::sync::Arc; + +use axum::http::StatusCode; +use tokio::sync::Mutex; +use tokio::task::JoinHandle; + +use crate::agents::ExtensionLoadResult; +use crate::execution::manager::AgentManager; +use crate::scheduler_trait::SchedulerTrait; +use crate::session::SessionManager; + +use super::session_event_bus::SessionEventBus; + +type ExtensionLoadTask = JoinHandle>; +type ExtensionLoadingTasks = + Arc>>>>>; + +#[derive(Clone)] +pub struct AppState { + pub(crate) agent_manager: Arc, + pub recipe_file_hash_map: Arc>>, + recipe_session_tracker: Arc>>, + pub extension_loading_tasks: ExtensionLoadingTasks, + session_buses: Arc>>>, +} + +impl AppState { + pub async fn new() -> anyhow::Result> { + let agent_manager = AgentManager::instance().await?; + Ok(Arc::new(Self { + agent_manager, + recipe_file_hash_map: Arc::new(Mutex::new(HashMap::new())), + recipe_session_tracker: Arc::new(Mutex::new(HashSet::new())), + extension_loading_tasks: Arc::new(Mutex::new(HashMap::new())), + session_buses: Arc::new(Mutex::new(HashMap::new())), + })) + } + + pub async fn get_or_create_event_bus(&self, session_id: &str) -> Arc { + let mut buses = self.session_buses.lock().await; + buses + .entry(session_id.to_string()) + .or_insert_with(|| Arc::new(SessionEventBus::new())) + .clone() + } + + pub async fn get_event_bus(&self, session_id: &str) -> Option> { + let buses = self.session_buses.lock().await; + buses.get(session_id).cloned() + } + + pub async fn remove_event_bus(&self, session_id: &str) { + let mut buses = self.session_buses.lock().await; + buses.remove(session_id); + } + + pub async fn set_extension_loading_task( + &self, + session_id: String, + task: JoinHandle>, + ) { + let mut tasks = self.extension_loading_tasks.lock().await; + tasks.insert(session_id, Arc::new(Mutex::new(Some(task)))); + } + + pub async fn take_extension_loading_task( + &self, + session_id: &str, + ) -> Result>, tokio::task::JoinError> { + let task_holder = { + let tasks = self.extension_loading_tasks.lock().await; + tasks.get(session_id).cloned() + }; + + if let Some(holder) = task_holder { + let mut task = holder.lock().await; + if let Some(handle) = task.as_mut() { + match handle.await { + Ok(results) => { + task.take(); + return Ok(Some(results)); + } + Err(error) => { + task.take(); + tracing::warn!("Background extension loading task failed: {error}"); + return Err(error); + } + } + } + } + Ok(None) + } + + pub async fn remove_extension_loading_task(&self, session_id: &str) { + let mut tasks = self.extension_loading_tasks.lock().await; + tasks.remove(session_id); + } + + pub fn scheduler(&self) -> Option> { + self.agent_manager.scheduler() + } + + pub fn session_manager(&self) -> &SessionManager { + self.agent_manager.session_manager() + } + + pub async fn mark_recipe_run_if_absent(&self, session_id: &str) -> bool { + let mut sessions = self.recipe_session_tracker.lock().await; + if sessions.contains(session_id) { + false + } else { + sessions.insert(session_id.to_string()); + true + } + } + + pub async fn get_agent(&self, session_id: String) -> anyhow::Result> { + self.agent_manager.get_or_create_agent(session_id).await + } + + pub async fn get_agent_for_route( + &self, + session_id: String, + ) -> Result, StatusCode> { + self.get_agent(session_id).await.map_err(|error| { + tracing::error!("Failed to get agent: {error}"); + StatusCode::INTERNAL_SERVER_ERROR + }) + } +} diff --git a/crates/goose/src/acp/transport/mod.rs b/crates/goose/src/acp/transport/mod.rs index 253f5c656..1f8dbfbd1 100644 --- a/crates/goose/src/acp/transport/mod.rs +++ b/crates/goose/src/acp/transport/mod.rs @@ -9,7 +9,7 @@ use axum::{ extract::{Request, State}, http::{header, HeaderName, HeaderValue, Method, StatusCode}, middleware::Next, - response::Response, + response::{Html, Response}, routing::get, Router, }; @@ -177,7 +177,13 @@ fn acp_cors_layer(policy: AcpOriginPolicy) -> CorsLayer { fn aux_cors_layer() -> CorsLayer { CorsLayer::new() .allow_origin(Any) - .allow_methods([Method::GET, Method::POST, Method::DELETE, Method::OPTIONS]) + .allow_methods([ + Method::GET, + Method::POST, + Method::PUT, + Method::DELETE, + Method::OPTIONS, + ]) .allow_headers([ header::CONTENT_TYPE, header::ACCEPT, @@ -220,6 +226,29 @@ async fn health() -> &'static str { "ok" } +async fn root_landing() -> Html<&'static str> { + Html( + r#" + + + + + Goose local serve + + + +

Goose local serve 已启动

+

这是 API 服务，没有 Web UI。健康检查：/status（应返回 ok）。

+

请使用 HTTPS 访问本机地址；HTTP 端口未监听。

+

首次访问自签证书时，浏览器会提示不安全，选择「高级」→「继续访问」即可。

+ +"#, + ) +} + /// The full standalone ACP server router used by `goose serve`: ACP transport, /// optional token auth, health/status endpoints, and the MCP app proxy. pub fn create_router( @@ -227,6 +256,7 @@ pub fn create_router( secret_key: String, require_token: bool, additional_allowed_origins: Vec, + tkmind_state: Arc, ) -> Router { let policy = if !additional_allowed_origins.is_empty() { AcpOriginPolicy::exact(additional_allowed_origins) @@ -239,9 +269,11 @@ pub fn create_router( create_acp_router_with_policy(server, policy, require_token.then_some(secret_key.clone())); let aux_routes = Router::new() + .route("/", get(root_landing)) .route("/health", get(health)) .route("/status", get(health)) .merge(super::mcp_app_proxy::routes(secret_key)) + .merge(super::tkmind::routes(tkmind_state)) .layer(aux_cors_layer()); acp_routes.merge(aux_routes) diff --git a/crates/goose/src/agents/platform_extensions/aider.rs b/crates/goose/src/agents/platform_extensions/aider.rs new file mode 100644 index 000000000..58f550bd7 --- /dev/null +++ b/crates/goose/src/agents/platform_extensions/aider.rs @@ -0,0 +1,310 @@ +use std::path::{Path, PathBuf}; +use std::time::Duration; + +use crate::agents::extension::PlatformExtensionContext; +use crate::agents::mcp_client::{Error, McpClientTrait}; +use crate::agents::platform_extensions::developer::workspace_path::resolve_path_in_workspace; +use crate::agents::tool_execution::ToolCallContext; +use anyhow::Result; +use async_trait::async_trait; +use indoc::indoc; +use rmcp::model::{ + CallToolResult, ContentBlock, Implementation, InitializeResult, JsonObject, ListToolsResult, + ServerCapabilities, Tool, ToolAnnotations, +}; +use schemars::{schema_for, JsonSchema}; +use serde::Deserialize; +use tokio::io::AsyncReadExt; +use tokio::process::Command; +use tokio_util::sync::CancellationToken; + +pub static EXTENSION_NAME: &str = "aider"; + +const DEFAULT_TIMEOUT_SECS: u64 = 600; + +#[derive(Debug, Deserialize, JsonSchema)] +struct AiderCodeParams { + /// Coding task for Aider: what to change, fix, or implement. + task: String, + /// Project directory. Defaults to the session working directory. + project_dir: Option, +} + +pub struct AiderClient { + info: InitializeResult, +} + +impl AiderClient { + pub fn new(_context: PlatformExtensionContext) -> Result { + let info = InitializeResult::new(ServerCapabilities::builder().enable_tools().build()) + .with_server_info( + Implementation::new(EXTENSION_NAME.to_string(), "1.0.0".to_string()) + .with_title("Aider"), + ) + .with_instructions(indoc! {" + Delegate multi-file coding work to Aider instead of using developer write/edit. + + Prefer this tool for: + - Bug fixes and refactors across multiple files + - New feature implementation + - Test-failure fixes that need code changes + + After Aider finishes, verify with developer shell (tests, git status). + For one-line or single-file tweaks, developer edit/write is fine. + "}); + + Ok(Self { info }) + } + + fn schema() -> JsonObject { + serde_json::to_value(schema_for!(T)) + .expect("schema serialization should succeed") + .as_object() + .expect("schema should serialize to an object") + .clone() + } + + fn get_tools() -> Vec { + vec![Tool::new( + "code".to_string(), + "Run Aider to implement a coding task in the project directory. \ + Aider edits files directly and uses git. Returns stdout/stderr and exit code." + .to_string(), + Self::schema::(), + ) + .annotate(ToolAnnotations::from_raw( + Some("Aider Code".to_string()), + Some(false), + Some(true), + Some(false), + Some(true), + ))] + } + + fn resolve_aider_bin() -> Result { + for key in ["GOOSE_AIDER_BIN", "AIDER_BIN"] { + if let Ok(path) = std::env::var(key) { + let candidate = PathBuf::from(&path); + if candidate.is_file() { + return Ok(candidate); + } + return Err(format!("{key} is set but not a file: {path}")); + } + } + + for candidate in [ + "/Users/john/Project/aider/aider.sh", + "/Users/john/Project/aider/.venv/bin/aider", + "/root/aider/.venv/bin/aider", + ] { + let path = PathBuf::from(candidate); + if path.is_file() { + return Ok(path); + } + } + + Err( + "Aider binary not found. Set GOOSE_AIDER_BIN or install Aider under /root/aider." + .to_string(), + ) + } + + fn resolve_coding_router() -> Option { + std::env::var("GOOSE_CODING_ROUTER") + .ok() + .map(PathBuf::from) + .filter(|path| path.is_file()) + } + + fn timeout_secs() -> u64 { + std::env::var("GOOSE_AIDER_TIMEOUT_SECS") + .ok() + .and_then(|value| value.parse().ok()) + .filter(|secs| *secs > 0) + .unwrap_or(DEFAULT_TIMEOUT_SECS) + } + + async fn run_aider( + task: &str, + project_dir: &Path, + cancel_token: CancellationToken, + ) -> CallToolResult { + if task.trim().is_empty() { + return CallToolResult::error(vec![ContentBlock::text("task cannot be empty")]); + } + + if !project_dir.is_dir() { + return CallToolResult::error(vec![ContentBlock::text(format!( + "project directory does not exist: {}", + project_dir.display() + ))]); + } + + let timeout_secs = Self::timeout_secs(); + let mut command = if let Some(router) = Self::resolve_coding_router() { + let mut cmd = Command::new(router); + cmd.arg("aider") + .arg(project_dir) + .arg(task) + .current_dir(project_dir); + cmd + } else { + let aider_bin = match Self::resolve_aider_bin() { + Ok(path) => path, + Err(error) => return CallToolResult::error(vec![ContentBlock::text(error)]), + }; + let mut cmd = Command::new(aider_bin); + cmd.arg("--yes-always") + .arg("--message") + .arg(task) + .current_dir(project_dir); + cmd + }; + + command.stdout(std::process::Stdio::piped()); + command.stderr(std::process::Stdio::piped()); + + let mut child = match command.spawn() { + Ok(child) => child, + Err(error) => { + return CallToolResult::error(vec![ContentBlock::text(format!( + "Failed to start Aider: {error}" + ))]); + } + }; + + let stdout = child.stdout.take(); + let stderr = child.stderr.take(); + + let wait_result = tokio::select! { + () = cancel_token.cancelled() => { + let _ = child.kill().await; + return CallToolResult::error(vec![ContentBlock::text("Aider run cancelled")]); + } + result = tokio::time::timeout(Duration::from_secs(timeout_secs), child.wait()) => result, + }; + + let exit_code = match wait_result { + Ok(Ok(status)) => status.code(), + Ok(Err(error)) => { + return CallToolResult::error(vec![ContentBlock::text(format!( + "Failed waiting on Aider: {error}" + ))]); + } + Err(_) => { + return CallToolResult::error(vec![ContentBlock::text(format!( + "Aider timed out after {timeout_secs}s" + ))]); + } + }; + + let mut stdout_text = String::new(); + let mut stderr_text = String::new(); + + if let Some(mut stdout) = stdout { + let _ = stdout.read_to_string(&mut stdout_text).await; + } + if let Some(mut stderr) = stderr { + let _ = stderr.read_to_string(&mut stderr_text).await; + } + + let body = format!( + "project_dir: {}\nexit_code: {}\n\n--- stdout ---\n{}\n\n--- stderr ---\n{}", + project_dir.display(), + exit_code + .map(|code| code.to_string()) + .unwrap_or_else(|| "null".to_string()), + stdout_text.trim_end(), + stderr_text.trim_end(), + ); + + if exit_code == Some(0) { + CallToolResult::success(vec![ContentBlock::text(body)]) + } else { + CallToolResult::error(vec![ContentBlock::text(body)]) + } + } +} + +#[async_trait] +impl McpClientTrait for AiderClient { + async fn list_tools( + &self, + _session_id: &str, + _next_cursor: Option, + _cancellation_token: CancellationToken, + ) -> Result { + Ok(ListToolsResult { + tools: Self::get_tools(), + next_cursor: None, + meta: None, + ..Default::default() + }) + } + + async fn call_tool( + &self, + ctx: &ToolCallContext, + name: &str, + arguments: Option, + cancellation_token: CancellationToken, + ) -> Result { + if name != "code" { + return Ok(CallToolResult::error(vec![ContentBlock::text(format!( + "Unknown tool: {name}" + ))])); + } + + let Some(value) = arguments.map(serde_json::Value::Object) else { + return Ok(CallToolResult::error(vec![ContentBlock::text( + "Missing arguments", + )])); + }; + + let params: AiderCodeParams = match serde_json::from_value(value) { + Ok(params) => params, + Err(error) => { + return Ok(CallToolResult::error(vec![ContentBlock::text(format!( + "Failed to parse arguments: {error}" + ))])); + } + }; + + let project_dir = match params.project_dir.as_deref() { + Some(dir) => match resolve_path_in_workspace(dir, ctx.working_dir.as_deref()) { + Ok(path) => path, + Err(error) => { + return Ok(CallToolResult::error(vec![ContentBlock::text(error)])); + } + }, + None => match ctx.working_dir.clone() { + Some(path) => path, + None => { + return Ok(CallToolResult::error(vec![ContentBlock::text( + "working_dir is required", + )])); + } + }, + }; + + Ok(Self::run_aider(¶ms.task, &project_dir, cancellation_token).await) + } + + fn get_info(&self) -> Option<&InitializeResult> { + Some(&self.info) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn aider_exposes_code_tool() { + let names: Vec = AiderClient::get_tools() + .into_iter() + .map(|tool| tool.name.to_string()) + .collect(); + assert_eq!(names, vec!["code"]); + } +} diff --git a/crates/goose/src/agents/platform_extensions/developer/mod.rs b/crates/goose/src/agents/platform_extensions/developer/mod.rs index 6a6b259de..62c8cfa8d 100644 --- a/crates/goose/src/agents/platform_extensions/developer/mod.rs +++ b/crates/goose/src/agents/platform_extensions/developer/mod.rs @@ -3,6 +3,7 @@ pub mod image; pub mod shell; mod shell_output_streaming; pub mod tree; +pub mod workspace_path; use crate::agents::extension::PlatformExtensionContext; use crate::agents::mcp_client::{Error, McpClientTrait}; diff --git a/crates/goose/src/agents/platform_extensions/developer/workspace_path.rs b/crates/goose/src/agents/platform_extensions/developer/workspace_path.rs new file mode 100644 index 000000000..de3063437 --- /dev/null +++ b/crates/goose/src/agents/platform_extensions/developer/workspace_path.rs @@ -0,0 +1,479 @@ +use std::path::{Component, Path, PathBuf}; + +use super::edit::resolve_path; + +pub fn resolve_path_in_workspace( + path: &str, + working_dir: Option<&Path>, +) -> Result { + let Some(base) = working_dir else { + return Ok(resolve_path(path, None)); + }; + + let resolved = resolve_path(path, Some(base)); + ensure_within_workspace(&resolved, base)?; + Ok(resolved) +} + +pub fn ensure_shell_command_within_workspace( + command: &str, + working_dir: &Path, +) -> Result<(), String> { + if command.contains("../") { + return Err(format!( + "Shell command must not use parent paths (../). All file access must stay within {}.", + working_dir.display() + )); + } + + for pattern in ["/Users/", "/home/", "$HOME", "${HOME}"] { + if command.contains(pattern) { + return Err(format!( + "Shell command must not reference host home paths ({pattern}). \ + All file access must stay within {}.", + working_dir.display() + )); + } + } + + let lower = command.to_ascii_lowercase(); + for token in [ + "curl ", + "curl\t", + "wget ", + "wget\t", + "fetch http", + "fetch https", + ] { + if lower.contains(token) { + return Err(format!( + "Shell network fetch is disabled in user workspace sessions. \ + Use local commands (find, ls, cat) within {}.", + working_dir.display() + )); + } + } + + let base = canonicalize_existing(working_dir).map_err(|error| { + format!( + "Failed to resolve working directory {}: {error}", + working_dir.display() + ) + })?; + let base_norm = lexical_normalize(&base); + + for token in shell_path_tokens(command) { + let resolved = resolve_shell_path_token(&token, working_dir); + let Some(resolved) = resolved else { + continue; + }; + let resolved_norm = + if resolved.exists() { + lexical_normalize(&canonicalize_existing(&resolved).map_err(|error| { + format!("Failed to resolve {}: {error}", resolved.display()) + })?) + } else { + lexical_normalize(&resolved) + }; + if !resolved_norm.starts_with(&base_norm) { + return Err(format!( + "Shell command references path outside the session working directory: {token}. \ + All file access must stay within {}.", + working_dir.display() + )); + } + } + + Ok(()) +} + +pub fn ensure_shell_command_does_not_write_public_html(command: &str) -> Result<(), String> { + for target in shell_public_html_write_targets(command) { + return Err(format!( + "Shell cannot write HTML under public/. Use sandbox-fs or the static-page-publish flow instead. Blocked target: {target}" + )); + } + Ok(()) +} + +fn is_public_html_write_path(path: &str) -> bool { + let path = path.trim().trim_matches(|c| "'\"".contains(c)); + if path.is_empty() { + return false; + } + let lower = path.to_ascii_lowercase(); + if !lower.ends_with(".html") && !lower.ends_with("/public") { + return false; + } + lower.contains("/public/") + || lower.starts_with("public/") + || lower.starts_with("./public/") + || lower.ends_with("/public") +} + +fn shell_public_html_write_targets(command: &str) -> Vec { + let mut targets = shell_redirect_targets(command); + targets.extend(shell_tee_targets(command)); + if let Some(dest) = shell_cp_mv_destination(command) { + targets.push(dest); + } + targets + .into_iter() + .filter(|target| is_public_html_write_path(target)) + .collect() +} + +fn shell_redirect_targets(command: &str) -> Vec { + let mut targets = Vec::new(); + let chars: Vec = command.chars().collect(); + let mut index = 0; + while index < chars.len() { + if chars[index] != '>' { + index += 1; + continue; + } + if index > 0 && chars[index - 1].is_ascii_digit() { + index += 1; + continue; + } + if index + 1 < chars.len() && chars[index + 1] == '>' { + index += 2; + } else { + index += 1; + } + while index < chars.len() && chars[index].is_whitespace() { + index += 1; + } + if let Some((token, next)) = read_shell_path_token(&chars[index..]) { + if !token.is_empty() { + targets.push(token); + } + index += next; + } + } + targets +} + +fn read_shell_path_token(input: &[char]) -> Option<(String, usize)> { + if input.is_empty() { + return None; + } + match input[0] { + '\'' | '"' => { + let closing = input[0]; + let mut token = String::new(); + let mut index = 1; + while index < input.len() { + if input[index] == closing { + return Some((token, index + 1)); + } + token.push(input[index]); + index += 1; + } + Some((token, input.len())) + } + _ => { + let mut token = String::new(); + let mut index = 0; + while index < input.len() && !input[index].is_whitespace() && input[index] != '|' { + token.push(input[index]); + index += 1; + } + Some((token, index)) + } + } +} + +fn shell_tee_targets(command: &str) -> Vec { + let tokens = shell_path_tokens(command); + let mut targets = Vec::new(); + for (index, token) in tokens.iter().enumerate() { + if token.eq_ignore_ascii_case("tee") { + for candidate in tokens.iter().skip(index + 1) { + if candidate.starts_with('-') { + continue; + } + targets.push(candidate.clone()); + } + } + } + targets +} + +fn shell_cp_mv_destination(command: &str) -> Option { + let tokens = shell_path_tokens(command); + let first = tokens.first()?.to_ascii_lowercase(); + if first != "cp" && first != "mv" { + return None; + } + tokens + .iter() + .rev() + .find(|token| !token.starts_with('-') && looks_like_shell_path_token(token)) + .cloned() +} + +fn shell_path_tokens(command: &str) -> Vec { + let mut tokens = Vec::new(); + let mut current = String::new(); + let mut in_single = false; + let mut in_double = false; + let mut escape = false; + + for ch in command.chars() { + if escape { + current.push(ch); + escape = false; + continue; + } + match ch { + '\\' if in_double => escape = true, + '\'' if !in_double => in_single = !in_single, + '"' if !in_single => in_double = !in_double, + c if c.is_whitespace() && !in_single && !in_double => { + if !current.is_empty() { + tokens.push(std::mem::take(&mut current)); + } + } + _ => current.push(ch), + } + } + if !current.is_empty() { + tokens.push(current); + } + tokens +} + +fn looks_like_shell_path_token(token: &str) -> bool { + let token = token.trim_matches(|c| "'\"".contains(c)); + if token.is_empty() || token.starts_with('-') { + return false; + } + if token.starts_with("2>") + || token.starts_with("1>") + || token.starts_with(">>") + || token == ">" + || token.contains("://") + { + return false; + } + !matches!(token, "|" | "&&" | "||" | ";" | "(" | ")") + && (token.starts_with('/') + || token.starts_with("./") + || token.starts_with("../") + || token.starts_with('~') + || token == "." + || token == ".." + || token.contains('/') + || token.contains('\\')) +} + +fn resolve_shell_path_token(token: &str, working_dir: &Path) -> Option { + let token = token.trim_matches(|c| "'\"".contains(c)); + if !looks_like_shell_path_token(token) { + return None; + } + + if token == "." { + return Some(working_dir.to_path_buf()); + } + if token == ".." { + return working_dir.parent().map(Path::to_path_buf); + } + + let expanded = if let Some(rest) = token.strip_prefix('~') { + dirs::home_dir()?.join(rest.trim_start_matches(['/', '\\'])) + } else { + PathBuf::from(token) + }; + + Some(if expanded.is_absolute() { + expanded + } else { + working_dir.join(expanded) + }) +} + +pub fn ensure_within_workspace(resolved: &Path, base: &Path) -> Result<(), String> { + let base_canon = canonicalize_existing(base).map_err(|error| { + format!( + "Failed to resolve working directory {}: {error}", + base.display() + ) + })?; + let resolved_canon = workspace_check_path(resolved, &base_canon)?; + let base_check = lexical_normalize(&base_canon); + let resolved_check = lexical_normalize(&resolved_canon); + if !resolved_check.starts_with(&base_check) { + return Err(format!( + "Path {} is outside the session working directory {}", + resolved.display(), + base.display() + )); + } + Ok(()) +} + +fn workspace_check_path(resolved: &Path, base_canon: &Path) -> Result { + if resolved.exists() { + return canonicalize_existing(resolved); + } + + let mut anchor = lexical_normalize(resolved); + while !anchor.exists() { + match anchor.parent() { + Some(parent) if parent.as_os_str().is_empty() => break, + Some(parent) => anchor = parent.to_path_buf(), + None => break, + } + } + + if anchor.exists() { + let anchor_canon = canonicalize_existing(&anchor)?; + let suffix = resolved + .strip_prefix(&anchor) + .unwrap_or_else(|_| resolved.as_ref()); + return Ok(anchor_canon.join(suffix)); + } + + if resolved.is_absolute() { + Ok(lexical_normalize(resolved)) + } else { + Ok(base_canon.join(resolved)) + } +} + +fn canonicalize_existing(path: &Path) -> Result { + path.canonicalize() + .map_err(|error| format!("Failed to resolve {}: {error}", path.display())) +} + +fn lexical_normalize(path: &Path) -> PathBuf { + let mut out = PathBuf::new(); + for component in path.components() { + match component { + Component::CurDir => {} + Component::ParentDir => { + out.pop(); + } + Component::Prefix(prefix) => out.push(prefix.as_os_str()), + Component::RootDir => out.push(component.as_os_str()), + Component::Normal(part) => out.push(part), + } + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + use std::fs; + use tempfile::tempdir; + + #[test] + fn allows_relative_paths_inside_workspace() { + let dir = tempdir().unwrap(); + let base = dir.path(); + let resolved = + resolve_path_in_workspace("notes/report.html", Some(base)).expect("inside workspace"); + assert_eq!(resolved, base.join("notes/report.html")); + } + + #[test] + fn blocks_parent_escape() { + let dir = tempdir().unwrap(); + let base = dir.path(); + let error = resolve_path_in_workspace("../outside.txt", Some(base)) + .expect_err("should block escape"); + assert!(error.contains("outside the session working directory")); + } + + #[test] + fn blocks_absolute_paths_outside_workspace() { + let dir = tempdir().unwrap(); + let base = dir.path(); + let outside = std::env::temp_dir().join("goose-workspace-outside.txt"); + let error = resolve_path_in_workspace(outside.to_str().unwrap(), Some(base)) + .expect_err("should block outside absolute path"); + assert!(error.contains("outside the session working directory")); + } + + #[test] + fn blocks_parent_path_segments_in_shell_commands() { + let dir = tempdir().unwrap(); + let base = dir.path(); + let error = ensure_shell_command_within_workspace("find .. -name '*.csv'", base) + .expect_err("blocked"); + assert!( + error.contains("outside the session working directory") + || error.contains("parent paths") + ); + } + + #[test] + fn blocks_shell_commands_with_paths_outside_workspace() { + let dir = tempdir().unwrap(); + let base = dir.path(); + let outside = std::env::temp_dir().join("goose-shell-outside.txt"); + let command = format!("find {} -name '*.csv'", outside.display()); + let error = ensure_shell_command_within_workspace(&command, base).expect_err("blocked"); + assert!(error.contains("outside the session working directory")); + } + + #[test] + fn allows_shell_commands_within_workspace() { + let dir = tempdir().unwrap(); + let base = dir.path(); + fs::create_dir_all(base.join("oa")).unwrap(); + ensure_shell_command_within_workspace("find . -name '*.csv'", base).expect("allowed"); + ensure_shell_command_within_workspace("rg export csv oa/", base).expect("allowed"); + } + + #[test] + fn allows_absolute_paths_inside_workspace() { + let dir = tempdir().unwrap(); + let base = dir.path(); + let inside = base.join("inside.txt"); + fs::write(&inside, "ok").unwrap(); + let resolved = + resolve_path_in_workspace(inside.to_str().unwrap(), Some(base)).expect("allowed"); + assert_eq!(resolved, inside); + } + + #[test] + fn blocks_shell_redirect_writes_to_public_html() { + let dir = tempdir().unwrap(); + let base = dir.path(); + let error = ensure_shell_command_does_not_write_public_html( + "echo '' > public/test.html", + ) + .expect_err("blocked"); + assert!(error.contains("public/")); + let _ = base; + } + + #[test] + fn blocks_shell_tee_writes_to_public_html() { + let error = ensure_shell_command_does_not_write_public_html( + "printf '%s' '' | tee public/report.html", + ) + .expect_err("blocked"); + assert!(error.contains("public/report.html")); + } + + #[test] + fn blocks_shell_cp_to_public_html() { + let error = + ensure_shell_command_does_not_write_public_html("cp draft.html public/page.html") + .expect_err("blocked"); + assert!(error.contains("public/page.html")); + } + + #[test] + fn allows_shell_reads_from_public_html() { + ensure_shell_command_does_not_write_public_html("cat public/page.html").expect("allowed"); + ensure_shell_command_does_not_write_public_html("ls public/").expect("allowed"); + ensure_shell_command_does_not_write_public_html("echo ok > notes/report.html") + .expect("allowed"); + } +} diff --git a/crates/goose/src/agents/platform_extensions/mod.rs b/crates/goose/src/agents/platform_extensions/mod.rs index 5159784c5..edfb910d0 100644 --- a/crates/goose/src/agents/platform_extensions/mod.rs +++ b/crates/goose/src/agents/platform_extensions/mod.rs @@ -1,3 +1,4 @@ +pub mod aider; pub mod analyze; pub mod apps; pub mod chatrecall; @@ -5,12 +6,15 @@ pub mod chatrecall; pub mod code_execution; pub mod developer; pub mod ext_manager; +pub mod openhands; pub mod orchestrator; +pub mod projectmemory; pub mod scheduler; pub mod summarize; pub mod summon; pub mod todo; pub mod tom; +pub mod web; use std::collections::HashMap; @@ -30,6 +34,36 @@ pub static PLATFORM_EXTENSIONS: Lazy || { let mut map = HashMap::new(); + map.insert( + openhands::EXTENSION_NAME, + PlatformExtensionDef { + name: openhands::EXTENSION_NAME, + display_name: "OpenHands", + description: + "Delegate multi-file coding tasks to OpenHands instead of developer write/edit", + default_enabled: true, + unprefixed_tools: false, + hidden: false, + client_factory: |ctx| { + Some(Box::new(openhands::OpenHandsClient::new(ctx).unwrap())) + }, + }, + ); + + map.insert( + aider::EXTENSION_NAME, + PlatformExtensionDef { + name: aider::EXTENSION_NAME, + display_name: "Aider", + description: + "Delegate multi-file coding tasks to Aider instead of developer write/edit", + default_enabled: true, + unprefixed_tools: false, + hidden: false, + client_factory: |ctx| Some(Box::new(aider::AiderClient::new(ctx).unwrap())), + }, + ); + map.insert( analyze::EXTENSION_NAME, PlatformExtensionDef { @@ -88,6 +122,21 @@ pub static PLATFORM_EXTENSIONS: Lazy }, ); + map.insert( + projectmemory::EXTENSION_NAME, + PlatformExtensionDef { + name: projectmemory::EXTENSION_NAME, + display_name: "Project Memory", + description: "Persist harness bootstrap summaries in session extension data", + default_enabled: false, + unprefixed_tools: false, + hidden: true, + client_factory: |ctx| { + Some(Box::new(projectmemory::ProjectMemoryClient::new(ctx).unwrap())) + }, + }, + ); + map.insert( "extensionmanager", PlatformExtensionDef { @@ -222,6 +271,19 @@ pub static PLATFORM_EXTENSIONS: Lazy }, ); + map.insert( + web::EXTENSION_NAME, + PlatformExtensionDef { + name: web::EXTENSION_NAME, + display_name: "Web", + description: "Fetch web pages and search the internet", + default_enabled: true, + unprefixed_tools: false, + hidden: false, + client_factory: |ctx| Some(Box::new(web::WebClient::new(ctx).unwrap())), + }, + ); + map }, ); diff --git a/crates/goose/src/agents/platform_extensions/openhands.rs b/crates/goose/src/agents/platform_extensions/openhands.rs new file mode 100644 index 000000000..ed6e76b9c --- /dev/null +++ b/crates/goose/src/agents/platform_extensions/openhands.rs @@ -0,0 +1,518 @@ +use std::path::{Path, PathBuf}; +use std::process::Command as StdCommand; +use std::time::Duration; + +use crate::agents::extension::PlatformExtensionContext; +use crate::agents::mcp_client::{Error, McpClientTrait}; +use crate::agents::platform_extensions::developer::workspace_path::resolve_path_in_workspace; +use crate::agents::tool_execution::ToolCallContext; +use anyhow::Result; +use async_trait::async_trait; +use indoc::indoc; +use rmcp::model::{ + CallToolResult, ContentBlock, Implementation, InitializeResult, JsonObject, ListToolsResult, + ServerCapabilities, Tool, ToolAnnotations, +}; +use schemars::{schema_for, JsonSchema}; +use serde::Deserialize; +use tokio::io::AsyncReadExt; +use tokio::process::Command; +use tokio_util::sync::CancellationToken; + +pub static EXTENSION_NAME: &str = "openhands"; + +const DEFAULT_TIMEOUT_SECS: u64 = 600; +const DEFAULT_DOCKER_IMAGE: &str = "docker.openhands.dev/openhands/openhands:latest"; +const WORKSPACE_BOUNDARY_INSTRUCTIONS: &str = "\ +You are running as a Goose delegated coding executor. Stay within the current \ +project directory only. Do not read, inspect, create, modify, or delete files \ +outside the current project directory, even if the user task mentions an \ +absolute path or a parent directory. If the task requires access outside the \ +current project directory, stop and report that it is outside the allowed \ +workspace.\n\nTask:\n"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum OpenHandsRunner { + Docker, + Sandbox, + Host, +} + +#[derive(Debug, Deserialize, JsonSchema)] +struct OpenHandsCodeParams { + /// Coding task for OpenHands: what to change, fix, or implement. + task: String, + /// Project directory. Defaults to the session working directory. + project_dir: Option, +} + +pub struct OpenHandsClient { + info: InitializeResult, +} + +impl OpenHandsClient { + pub fn new(_context: PlatformExtensionContext) -> Result { + let info = InitializeResult::new(ServerCapabilities::builder().enable_tools().build()) + .with_server_info( + Implementation::new(EXTENSION_NAME.to_string(), "1.0.0".to_string()) + .with_title("OpenHands"), + ) + .with_instructions(indoc! {" + Delegate multi-file coding work to OpenHands instead of using developer write/edit. + + Prefer this tool for: + - Bug fixes and refactors across multiple files + - New feature implementation + - Test-failure fixes that need code changes + + After OpenHands finishes, verify with developer shell (tests, git status). + For one-line or single-file tweaks, developer edit/write is fine. + "}); + + Ok(Self { info }) + } + + fn schema() -> JsonObject { + serde_json::to_value(schema_for!(T)) + .expect("schema serialization should succeed") + .as_object() + .expect("schema should serialize to an object") + .clone() + } + + fn get_tools() -> Vec { + vec![Tool::new( + "code".to_string(), + "Run OpenHands to implement a coding task in the project directory. \ + OpenHands edits files directly and uses git. Returns stdout/stderr and exit code." + .to_string(), + Self::schema::(), + ) + .annotate(ToolAnnotations::from_raw( + Some("OpenHands Code".to_string()), + Some(false), + Some(true), + Some(false), + Some(true), + ))] + } + + fn resolve_openhands_bin() -> Result { + for key in ["GOOSE_OPENHANDS_BIN", "OPENHANDS_BIN"] { + if let Ok(path) = std::env::var(key) { + let candidate = PathBuf::from(&path); + if candidate.is_file() { + return Ok(candidate); + } + return Err(format!("{key} is set but not a file: {path}")); + } + } + + for candidate in [ + "/Users/john/.local/bin/openhands", + "/usr/local/bin/openhands", + "/opt/homebrew/bin/openhands", + ] { + let path = PathBuf::from(candidate); + if path.is_file() { + return Ok(path); + } + } + + Err( + "OpenHands binary not found. Set GOOSE_OPENHANDS_BIN or install OpenHands in PATH." + .to_string(), + ) + } + + fn timeout_secs() -> u64 { + std::env::var("GOOSE_OPENHANDS_TIMEOUT_SECS") + .ok() + .and_then(|value| value.parse().ok()) + .filter(|secs| *secs > 0) + .unwrap_or(DEFAULT_TIMEOUT_SECS) + } + + fn runner() -> Result { + let value = std::env::var("GOOSE_OPENHANDS_RUNNER") + .unwrap_or_else(|_| "sandbox".to_string()) + .to_lowercase(); + + match value.as_str() { + "docker" | "container" => Ok(OpenHandsRunner::Docker), + "sandbox" | "macos_sandbox" => Ok(OpenHandsRunner::Sandbox), + "host" | "local" => Ok(OpenHandsRunner::Host), + other => Err(format!( + "Invalid GOOSE_OPENHANDS_RUNNER={other}. Use docker, sandbox, or host." + )), + } + } + + fn docker_image() -> String { + std::env::var("GOOSE_OPENHANDS_DOCKER_IMAGE") + .ok() + .filter(|image| !image.trim().is_empty()) + .unwrap_or_else(|| DEFAULT_DOCKER_IMAGE.to_string()) + } + + fn sandbox_exec_bin() -> Option { + let path = PathBuf::from("/usr/bin/sandbox-exec"); + path.is_file().then_some(path) + } + + fn sandbox_string(path: &Path) -> String { + path.to_string_lossy() + .replace('\\', "\\\\") + .replace('"', "\\\"") + } + + fn sandbox_profile(project_dir: &Path) -> String { + let project = Self::sandbox_string(project_dir); + + format!( + r#"(version 1) +(deny default) +(allow process*) +(allow network*) +(allow sysctl-read) +(allow mach-lookup) +(allow file-read-metadata) +(allow file-read*) +(allow file-write* + (subpath "{project}") + (literal "/dev/null")) +"# + ) + } + + fn add_openhands_args(command: &mut Command, bounded_task: String) { + command + .arg("--headless") + .arg("--json") + .arg("--always-approve") + .arg("--override-with-envs") + .arg("--task") + .arg(bounded_task); + } + + fn docker_env_keys() -> &'static [&'static str] { + &[ + "LLM_MODEL", + "LLM_API_KEY", + "LLM_BASE_URL", + "LLM_API_BASE", + "LLM_EMBEDDING_MODEL", + "OPENAI_API_KEY", + "OPENAI_BASE_URL", + "ANTHROPIC_API_KEY", + "DEEPSEEK_API_KEY", + "GEMINI_API_KEY", + "GOOGLE_API_KEY", + "OPENHANDS_LLM_MODEL", + "OPENHANDS_LLM_API_KEY", + "OPENHANDS_LLM_BASE_URL", + ] + } + + fn host_uid_gid() -> Option { + #[cfg(unix)] + { + let uid = StdCommand::new("id").arg("-u").output().ok()?; + let gid = StdCommand::new("id").arg("-g").output().ok()?; + if !uid.status.success() || !gid.status.success() { + return None; + } + + let uid = String::from_utf8(uid.stdout).ok()?; + let gid = String::from_utf8(gid.stdout).ok()?; + Some(format!("{}:{}", uid.trim(), gid.trim())) + } + + #[cfg(not(unix))] + { + None + } + } + + fn docker_command(project_dir: &Path, bounded_task: String) -> Command { + let mut command = Command::new("docker"); + command + .arg("run") + .arg("--rm") + .arg("--interactive") + .arg("--workdir") + .arg("/workspace") + .arg("--mount") + .arg(format!( + "type=bind,source={},target=/workspace,readonly=false", + project_dir.display() + )) + .arg("--env") + .arg("OPENHANDS_SUPPRESS_BANNER=1") + .arg("--env") + .arg("SANDBOX_WORKSPACE=/workspace") + .arg("--env") + .arg("SANDBOX_VOLUMES=/workspace:/workspace:rw") + .arg("--cap-drop") + .arg("ALL") + .arg("--security-opt") + .arg("no-new-privileges") + .arg("--pids-limit") + .arg("512") + .arg("--memory") + .arg( + std::env::var("GOOSE_OPENHANDS_DOCKER_MEMORY").unwrap_or_else(|_| "4g".to_string()), + ); + + if let Some(uid_gid) = Self::host_uid_gid() { + command + .arg("--user") + .arg(uid_gid.clone()) + .arg("--env") + .arg(format!( + "SANDBOX_USER_ID={}", + uid_gid.split(':').next().unwrap_or("1000") + )); + } + + for key in Self::docker_env_keys() { + if std::env::var_os(key).is_some() { + command.arg("--env").arg(key); + } + } + + command.arg(Self::docker_image()); + command.arg("openhands"); + Self::add_openhands_args(&mut command, bounded_task); + command.current_dir(project_dir); + command + } + + fn local_command( + openhands_bin: PathBuf, + project_dir: &Path, + bounded_task: String, + use_sandbox: bool, + ) -> Command { + let mut command = if use_sandbox { + if let Some(sandbox_exec) = Self::sandbox_exec_bin() { + let mut cmd = Command::new(sandbox_exec); + cmd.arg("-p") + .arg(Self::sandbox_profile(project_dir)) + .arg(&openhands_bin); + cmd + } else { + Command::new(openhands_bin) + } + } else { + Command::new(openhands_bin) + }; + + Self::add_openhands_args(&mut command, bounded_task); + command.current_dir(project_dir); + command + } + + async fn run_openhands( + task: &str, + project_dir: &Path, + cancel_token: CancellationToken, + ) -> CallToolResult { + if task.trim().is_empty() { + return CallToolResult::error(vec![ContentBlock::text("task cannot be empty")]); + } + + if !project_dir.is_dir() { + return CallToolResult::error(vec![ContentBlock::text(format!( + "project directory does not exist: {}", + project_dir.display() + ))]); + } + + let runner = match Self::runner() { + Ok(runner) => runner, + Err(error) => return CallToolResult::error(vec![ContentBlock::text(error)]), + }; + + let timeout_secs = Self::timeout_secs(); + let bounded_task = format!("{WORKSPACE_BOUNDARY_INSTRUCTIONS}{task}"); + let sandbox_home = project_dir.join(".goose").join("openhands-home"); + let sandbox_tmp = project_dir.join(".goose").join("openhands-tmp"); + if let Err(error) = std::fs::create_dir_all(&sandbox_home) { + return CallToolResult::error(vec![ContentBlock::text(format!( + "Failed to create OpenHands sandbox home: {error}" + ))]); + } + if let Err(error) = std::fs::create_dir_all(&sandbox_tmp) { + return CallToolResult::error(vec![ContentBlock::text(format!( + "Failed to create OpenHands sandbox tmp: {error}" + ))]); + } + + let mut command = match runner { + OpenHandsRunner::Docker => Self::docker_command(project_dir, bounded_task), + OpenHandsRunner::Sandbox | OpenHandsRunner::Host => { + let openhands_bin = match Self::resolve_openhands_bin() { + Ok(path) => path, + Err(error) => return CallToolResult::error(vec![ContentBlock::text(error)]), + }; + Self::local_command( + openhands_bin, + project_dir, + bounded_task, + runner == OpenHandsRunner::Sandbox, + ) + } + }; + + command + .env("HOME", &sandbox_home) + .env("TMPDIR", &sandbox_tmp) + .env("OPENHANDS_SUPPRESS_BANNER", "1"); + command.stdout(std::process::Stdio::piped()); + command.stderr(std::process::Stdio::piped()); + + let mut child = match command.spawn() { + Ok(child) => child, + Err(error) => { + return CallToolResult::error(vec![ContentBlock::text(format!( + "Failed to start OpenHands: {error}" + ))]); + } + }; + + let stdout = child.stdout.take(); + let stderr = child.stderr.take(); + + let wait_result = tokio::select! { + () = cancel_token.cancelled() => { + let _ = child.kill().await; + return CallToolResult::error(vec![ContentBlock::text("OpenHands run cancelled")]); + } + result = tokio::time::timeout(Duration::from_secs(timeout_secs), child.wait()) => result, + }; + + let exit_code = match wait_result { + Ok(Ok(status)) => status.code(), + Ok(Err(error)) => { + return CallToolResult::error(vec![ContentBlock::text(format!( + "Failed waiting on OpenHands: {error}" + ))]); + } + Err(_) => { + return CallToolResult::error(vec![ContentBlock::text(format!( + "OpenHands timed out after {timeout_secs}s" + ))]); + } + }; + + let mut stdout_text = String::new(); + let mut stderr_text = String::new(); + + if let Some(mut stdout) = stdout { + let _ = stdout.read_to_string(&mut stdout_text).await; + } + if let Some(mut stderr) = stderr { + let _ = stderr.read_to_string(&mut stderr_text).await; + } + + let body = format!( + "project_dir: {}\nexit_code: {}\n\n--- stdout ---\n{}\n\n--- stderr ---\n{}", + project_dir.display(), + exit_code + .map(|code| code.to_string()) + .unwrap_or_else(|| "null".to_string()), + stdout_text.trim_end(), + stderr_text.trim_end(), + ); + + if exit_code == Some(0) { + CallToolResult::success(vec![ContentBlock::text(body)]) + } else { + CallToolResult::error(vec![ContentBlock::text(body)]) + } + } +} + +#[async_trait] +impl McpClientTrait for OpenHandsClient { + async fn list_tools( + &self, + _session_id: &str, + _next_cursor: Option, + _cancellation_token: CancellationToken, + ) -> Result { + Ok(ListToolsResult { + tools: Self::get_tools(), + next_cursor: None, + meta: None, + ..Default::default() + }) + } + + async fn call_tool( + &self, + ctx: &ToolCallContext, + name: &str, + arguments: Option, + cancellation_token: CancellationToken, + ) -> Result { + if name != "code" { + return Ok(CallToolResult::error(vec![ContentBlock::text(format!( + "Unknown tool: {name}" + ))])); + } + + let Some(value) = arguments.map(serde_json::Value::Object) else { + return Ok(CallToolResult::error(vec![ContentBlock::text( + "Missing arguments", + )])); + }; + + let params: OpenHandsCodeParams = match serde_json::from_value(value) { + Ok(params) => params, + Err(error) => { + return Ok(CallToolResult::error(vec![ContentBlock::text(format!( + "Failed to parse arguments: {error}" + ))])); + } + }; + + let project_dir = match params.project_dir.as_deref() { + Some(dir) => match resolve_path_in_workspace(dir, ctx.working_dir.as_deref()) { + Ok(path) => path, + Err(error) => { + return Ok(CallToolResult::error(vec![ContentBlock::text(error)])); + } + }, + None => match ctx.working_dir.clone() { + Some(path) => path, + None => { + return Ok(CallToolResult::error(vec![ContentBlock::text( + "working_dir is required", + )])); + } + }, + }; + + Ok(Self::run_openhands(¶ms.task, &project_dir, cancellation_token).await) + } + + fn get_info(&self) -> Option<&InitializeResult> { + Some(&self.info) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn openhands_exposes_code_tool() { + let names: Vec = OpenHandsClient::get_tools() + .into_iter() + .map(|tool| tool.name.to_string()) + .collect(); + assert_eq!(names, vec!["code"]); + } +} diff --git a/crates/goose/src/agents/platform_extensions/projectmemory.rs b/crates/goose/src/agents/platform_extensions/projectmemory.rs new file mode 100644 index 000000000..7dcbdd371 --- /dev/null +++ b/crates/goose/src/agents/platform_extensions/projectmemory.rs @@ -0,0 +1,128 @@ +use crate::agents::extension::PlatformExtensionContext; +use crate::agents::mcp_client::{Error, McpClientTrait}; +use crate::agents::tool_execution::ToolCallContext; +use crate::session::extension_data::ExtensionState; +use anyhow::Result; +use async_trait::async_trait; +use rmcp::model::{ + CallToolResult, ContentBlock, Implementation, InitializeResult, JsonObject, ListToolsResult, + ServerCapabilities, +}; +use serde::{Deserialize, Serialize}; +use tokio_util::sync::CancellationToken; + +pub static EXTENSION_NAME: &str = "projectmemory"; + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct ProjectMemoryState { + pub summary: String, + pub source: Option, + pub updated_at: Option, +} + +impl ExtensionState for ProjectMemoryState { + const EXTENSION_NAME: &'static str = EXTENSION_NAME; + const VERSION: &'static str = "v0"; +} + +pub struct ProjectMemoryClient { + info: InitializeResult, + context: PlatformExtensionContext, +} + +impl ProjectMemoryClient { + pub fn new(context: PlatformExtensionContext) -> Result { + let info = InitializeResult::new(ServerCapabilities::builder().build()).with_server_info( + Implementation::new(EXTENSION_NAME.to_string(), "1.0.0".to_string()) + .with_title("Project Memory"), + ); + Ok(Self { info, context }) + } +} + +#[async_trait] +impl McpClientTrait for ProjectMemoryClient { + async fn list_tools( + &self, + _session_id: &str, + _next_cursor: Option, + _cancellation_token: CancellationToken, + ) -> Result { + Ok(ListToolsResult { + tools: vec![], + next_cursor: None, + meta: None, + ..Default::default() + }) + } + + async fn call_tool( + &self, + _ctx: &ToolCallContext, + name: &str, + _arguments: Option, + _cancellation_token: CancellationToken, + ) -> Result { + Ok(CallToolResult::error(vec![ContentBlock::text(format!( + "projectmemory has no tools (called: {name})" + ))])) + } + + fn get_info(&self) -> Option<&InitializeResult> { + Some(&self.info) + } + + async fn get_moim(&self, session_id: &str) -> Option { + let session = self + .context + .session_manager + .get_session(session_id, false) + .await + .ok()?; + let state = ProjectMemoryState::from_extension_data(&session.extension_data)?; + let summary = state.summary.trim(); + if summary.is_empty() { + return None; + } + + let mut lines = vec![ + "Project memory bootstrap: treat this as background context for the current project." + .to_string(), + "Prefer newer user instructions if anything conflicts.".to_string(), + ]; + + if let Some(source) = state.source.as_deref() { + lines.push(format!("Source: {source}")); + } + if let Some(updated_at) = state.updated_at.as_deref() { + lines.push(format!("Updated: {updated_at}")); + } + + Some(format!("{}\n\n{}", lines.join("\n"), summary)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::session::ExtensionData; + + #[test] + fn project_memory_state_round_trips_through_extension_data() { + let state = ProjectMemoryState { + summary: "Use the H5 client for mobile conversations.".to_string(), + source: Some("harness".to_string()), + updated_at: Some("2026-06-10T00:00:00Z".to_string()), + }; + let mut extension_data = ExtensionData::default(); + + state.to_extension_data(&mut extension_data).unwrap(); + + assert_eq!( + ProjectMemoryState::from_extension_data(&extension_data) + .unwrap() + .summary, + state.summary + ); + } +} diff --git a/crates/goose/src/agents/platform_extensions/web.rs b/crates/goose/src/agents/platform_extensions/web.rs new file mode 100644 index 000000000..8b8f38787 --- /dev/null +++ b/crates/goose/src/agents/platform_extensions/web.rs @@ -0,0 +1,331 @@ +use crate::agents::extension::PlatformExtensionContext; +use crate::agents::mcp_client::{Error, McpClientTrait}; +use crate::agents::tool_execution::ToolCallContext; +use anyhow::Result; +use async_trait::async_trait; +use rmcp::model::{ + CallToolResult, ContentBlock, Implementation, InitializeResult, JsonObject, ListToolsResult, + ServerCapabilities, Tool, ToolAnnotations, +}; +use schemars::{schema_for, JsonSchema}; +use serde::{Deserialize, Serialize}; +use tokio_util::sync::CancellationToken; + +pub static EXTENSION_NAME: &str = "web"; + +const MAX_OUTPUT_CHARS: usize = 50_000; + +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +struct FetchUrlParams { + url: String, + extract_text: bool, +} + +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +struct WebSearchParams { + query: String, + #[schemars(default = "default_num_results")] + num_results: u32, +} + +fn default_num_results() -> u32 { + 10 +} + +pub struct WebClient { + info: InitializeResult, + http: reqwest::Client, + #[allow(dead_code)] + context: PlatformExtensionContext, +} + +impl WebClient { + pub fn new(context: PlatformExtensionContext) -> Result { + let http = reqwest::Client::builder() + .user_agent("Mozilla/5.0 (compatible; Goose/1.0)") + .timeout(std::time::Duration::from_secs(30)) + .build()?; + let info = InitializeResult::new(ServerCapabilities::builder().enable_tools().build()) + .with_server_info( + Implementation::new(EXTENSION_NAME.to_string(), "1.0.0".to_string()) + .with_title("Web"), + ) + .with_instructions( + "Fetch web pages and search the internet using DuckDuckGo.".to_string(), + ); + Ok(Self { + info, + http, + context, + }) + } + + fn truncate(s: String) -> String { + if s.len() > MAX_OUTPUT_CHARS { + format!("{}\n[output truncated]", &s[..MAX_OUTPUT_CHARS]) + } else { + s + } + } + + fn strip_html(html: &str) -> String { + let mut text = html.to_string(); + + while let Some(start) = text.to_lowercase().find("") { + text = format!("{}{}", &text[..start], &text[start + end + 9..]); + } else { + break; + } + } + + while let Some(start) = text.to_lowercase().find("") { + text = format!("{}{}", &text[..start], &text[start + end + 8..]); + } else { + break; + } + } + + let mut result = String::new(); + let mut in_tag = false; + for ch in text.chars() { + match ch { + '<' => in_tag = true, + '>' => in_tag = false, + c if !in_tag => result.push(c), + _ => {} + } + } + + let lines: Vec<&str> = result.lines().map(|l| l.trim()).collect(); + lines + .iter() + .filter(|l| !l.is_empty()) + .cloned() + .collect::>() + .join("\n") + } + + async fn handle_fetch_url( + &self, + arguments: Option, + ) -> Result, String> { + let args = arguments.as_ref().ok_or("Missing arguments")?; + let url = args + .get("url") + .and_then(|v| v.as_str()) + .ok_or("Missing required parameter: url")? + .to_string(); + let extract_text = args + .get("extract_text") + .and_then(|v| v.as_bool()) + .unwrap_or(true); + + let body = self + .http + .get(&url) + .send() + .await + .map_err(|e| format!("Failed to fetch URL: {}", e))? + .text() + .await + .map_err(|e| format!("Failed to read response body: {}", e))?; + + let output = if extract_text { + Self::strip_html(&body) + } else { + body + }; + + Ok(vec![ContentBlock::text(Self::truncate(output))]) + } + + async fn handle_web_search( + &self, + arguments: Option, + ) -> Result, String> { + let args = arguments.as_ref().ok_or("Missing arguments")?; + let query = args + .get("query") + .and_then(|v| v.as_str()) + .ok_or("Missing required parameter: query")? + .to_string(); + let num_results = args + .get("num_results") + .and_then(|v| v.as_u64()) + .unwrap_or(10) as usize; + + let encoded_query: String = query + .chars() + .map(|c| match c { + ' ' => '+', + c if c.is_ascii_alphanumeric() || "-_.~".contains(c) => c, + _ => '+', + }) + .collect(); + + let url = format!("https://html.duckduckgo.com/html/?q={encoded_query}"); + + let html = self + .http + .get(&url) + .send() + .await + .map_err(|e| format!("Failed to search: {}", e))? + .text() + .await + .map_err(|e| format!("Failed to read search response: {}", e))?; + + let results = Self::parse_ddg_results(&html, num_results); + Ok(vec![ContentBlock::text(Self::truncate(results))]) + } + + fn parse_ddg_results(html: &str, max: usize) -> String { + let mut results = Vec::new(); + let lower = html.to_lowercase(); + + let mut pos = 0; + while results.len() < max { + let search_str = "class=\"result__a\""; + let Some(link_start) = lower[pos..].find(search_str).map(|i| i + pos) else { + break; + }; + + let before = &html[..link_start + search_str.len()]; + let a_start = before.rfind('<').unwrap_or(0); + let a_block = &html[a_start + ..link_start + + search_str.len() + + 100.min(html.len() - link_start - search_str.len())]; + + let href = Self::extract_attr(a_block, "href").unwrap_or_default(); + + let after_a_tag_end = link_start + search_str.len(); + let close_offset = html[after_a_tag_end..].find('>').unwrap_or(0); + let content_start = after_a_tag_end + close_offset + 1; + let close_a = html[content_start..].find("").unwrap_or(0); + let title = Self::strip_html(&html[content_start..content_start + close_a]); + + let snippet_search = "class=\"result__snippet\""; + let snippet = if let Some(snip_pos) = lower[pos..].find(snippet_search).map(|i| i + pos) + { + let after = snip_pos + snippet_search.len(); + if let Some(close) = html[after..].find('>') { + let text_start = after + close + 1; + if let Some(end_div) = html[text_start..].find("") { + Self::strip_html(&html[text_start..text_start + end_div]) + } else { + String::new() + } + } else { + String::new() + } + } else { + String::new() + }; + + if !title.trim().is_empty() { + results.push(format!( + "{}. {}\n URL: {}\n {}", + results.len() + 1, + title.trim(), + href, + snippet.trim() + )); + } + + pos = link_start + 1; + } + + if results.is_empty() { + "No results found.".to_string() + } else { + results.join("\n\n") + } + } + + fn extract_attr(tag: &str, attr: &str) -> Option { + let search = format!("{attr}=\""); + let lower = tag.to_lowercase(); + let start = lower.find(&search)? + search.len(); + let end = tag[start..].find('"')?; + Some(tag[start..start + end].to_string()) + } + + fn get_tools() -> Vec { + let fetch_schema = serde_json::to_value(schema_for!(FetchUrlParams)) + .expect("Failed to serialize FetchUrlParams schema"); + let search_schema = serde_json::to_value(schema_for!(WebSearchParams)) + .expect("Failed to serialize WebSearchParams schema"); + + vec![ + Tool::new( + "fetch_url".to_string(), + "Fetch the content of a URL. If extract_text is true, strips HTML tags to return readable text.".to_string(), + fetch_schema.as_object().unwrap().clone(), + ) + .annotate(ToolAnnotations::from_raw( + Some("Fetch URL".to_string()), + Some(false), + Some(false), + Some(false), + Some(false), + )), + Tool::new( + "web_search".to_string(), + "Search the web using DuckDuckGo and return a list of result titles, URLs, and snippets.".to_string(), + search_schema.as_object().unwrap().clone(), + ) + .annotate(ToolAnnotations::from_raw( + Some("Web Search".to_string()), + Some(false), + Some(false), + Some(false), + Some(false), + )), + ] + } +} + +#[async_trait] +impl McpClientTrait for WebClient { + async fn list_tools( + &self, + _session_id: &str, + _next_cursor: Option, + _cancellation_token: CancellationToken, + ) -> Result { + Ok(ListToolsResult { + tools: Self::get_tools(), + next_cursor: None, + meta: None, + ..Default::default() + }) + } + + async fn call_tool( + &self, + _ctx: &ToolCallContext, + name: &str, + arguments: Option, + _cancellation_token: CancellationToken, + ) -> Result { + let content = match name { + "fetch_url" => self.handle_fetch_url(arguments).await, + "web_search" => self.handle_web_search(arguments).await, + _ => Err(format!("Unknown tool: {name}")), + }; + + match content { + Ok(content) => Ok(CallToolResult::success(content)), + Err(error) => Ok(CallToolResult::error(vec![ContentBlock::text(format!( + "Error: {error}" + ))])), + } + } + + fn get_info(&self) -> Option<&InitializeResult> { + Some(&self.info) + } +} diff --git a/crates/goose/src/providers/inventory/mod.rs b/crates/goose/src/providers/inventory/mod.rs index 1f6612bc5..ce2a484ad 100644 --- a/crates/goose/src/providers/inventory/mod.rs +++ b/crates/goose/src/providers/inventory/mod.rs @@ -17,7 +17,7 @@ use chrono::{DateTime, Duration, Utc}; use futures::FutureExt; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; -use sqlx::{Row, Sqlite, Transaction}; +use sqlx::{AssertSqlSafe, Pool, Postgres, Row, Sqlite, Transaction}; use std::collections::{BTreeMap, HashMap, HashSet}; use std::panic::AssertUnwindSafe; use std::sync::{Arc, PoisonError, RwLock, RwLockReadGuard, RwLockWriteGuard}; @@ -458,6 +458,12 @@ impl ProviderInventoryService { identity: &InventoryIdentity, model_ids: &[String], ) -> Result<()> { + if self.storage.uses_postgres() { + return self + .store_refreshed_models_for_identity_pg(identity, model_ids) + .await; + } + let models = enrich_model_ids_with_canonical(&identity.provider_family, model_ids); let now = Utc::now(); let pool = self.storage.pool().await?; @@ -527,6 +533,80 @@ impl ProviderInventoryService { Ok(()) } + async fn store_refreshed_models_for_identity_pg( + &self, + identity: &InventoryIdentity, + model_ids: &[String], + ) -> Result<()> { + let models = enrich_model_ids_with_canonical(&identity.provider_family, model_ids); + let now = Utc::now(); + let pool = self.storage.pool_postgres().await?; + let mut tx = pool.begin().await?; + + sqlx::query( + r#" + INSERT INTO provider_inventory_entries ( + inventory_key, + provider_id, + provider_family, + last_updated_at, + last_refresh_attempt_at, + last_refresh_error, + updated_at + ) VALUES ($1, $2, $3, $4, $5, NULL, CURRENT_TIMESTAMP) + ON CONFLICT(inventory_key) DO UPDATE SET + provider_id = excluded.provider_id, + provider_family = excluded.provider_family, + last_updated_at = excluded.last_updated_at, + last_refresh_attempt_at = excluded.last_refresh_attempt_at, + last_refresh_error = NULL, + updated_at = CURRENT_TIMESTAMP + "#, + ) + .bind(&identity.inventory_key) + .bind(&identity.provider_id) + .bind(&identity.provider_family) + .bind(now.to_rfc3339()) + .bind(now.to_rfc3339()) + .execute(&mut *tx) + .await?; + + sqlx::query("DELETE FROM provider_inventory_models WHERE inventory_key = $1") + .bind(&identity.inventory_key) + .execute(&mut *tx) + .await?; + + for (ordinal, model) in models.iter().enumerate() { + sqlx::query( + r#" + INSERT INTO provider_inventory_models ( + inventory_key, + ordinal, + model_id, + name, + family, + context_limit, + reasoning, + recommended + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8) + "#, + ) + .bind(&identity.inventory_key) + .bind(i64::try_from(ordinal)?) + .bind(&model.id) + .bind(&model.name) + .bind(&model.family) + .bind(model.context_limit.map(i64::try_from).transpose()?) + .bind(model.reasoning) + .bind(model.recommended) + .execute(&mut *tx) + .await?; + } + + tx.commit().await?; + Ok(()) + } + pub async fn store_refresh_error( &self, provider_id: &str, @@ -547,6 +627,39 @@ impl ProviderInventoryService { let error = error.into(); let existing = self.read_snapshot(identity).await?; + if self.storage.uses_postgres() { + let pool = self.storage.pool_postgres().await?; + sqlx::query( + r#" + INSERT INTO provider_inventory_entries ( + inventory_key, + provider_id, + provider_family, + last_updated_at, + last_refresh_attempt_at, + last_refresh_error, + updated_at + ) VALUES ($1, $2, $3, $4, $5, $6, CURRENT_TIMESTAMP) + ON CONFLICT(inventory_key) DO UPDATE SET + provider_id = excluded.provider_id, + provider_family = excluded.provider_family, + last_updated_at = excluded.last_updated_at, + last_refresh_attempt_at = excluded.last_refresh_attempt_at, + last_refresh_error = excluded.last_refresh_error, + updated_at = CURRENT_TIMESTAMP + "#, + ) + .bind(&identity.inventory_key) + .bind(&identity.provider_id) + .bind(&identity.provider_family) + .bind(existing.and_then(|snapshot| snapshot.last_updated_at.map(|time| time.to_rfc3339()))) + .bind(Utc::now().to_rfc3339()) + .bind(error) + .execute(pool) + .await?; + return Ok(()); + } + sqlx::query( r#" INSERT INTO provider_inventory_entries ( @@ -754,6 +867,38 @@ impl ProviderInventoryService { async fn mark_refresh_started(&self, identity: &InventoryIdentity) -> Result<()> { let existing = self.read_snapshot(identity).await?; + if self.storage.uses_postgres() { + let pool = self.storage.pool_postgres().await?; + sqlx::query( + r#" + INSERT INTO provider_inventory_entries ( + inventory_key, + provider_id, + provider_family, + last_updated_at, + last_refresh_attempt_at, + last_refresh_error, + updated_at + ) VALUES ($1, $2, $3, $4, $5, NULL, CURRENT_TIMESTAMP) + ON CONFLICT(inventory_key) DO UPDATE SET + provider_id = excluded.provider_id, + provider_family = excluded.provider_family, + last_updated_at = excluded.last_updated_at, + last_refresh_attempt_at = excluded.last_refresh_attempt_at, + last_refresh_error = NULL, + updated_at = CURRENT_TIMESTAMP + "#, + ) + .bind(&identity.inventory_key) + .bind(&identity.provider_id) + .bind(&identity.provider_family) + .bind(existing.and_then(|snapshot| snapshot.last_updated_at.map(|time| time.to_rfc3339()))) + .bind(Utc::now().to_rfc3339()) + .execute(pool) + .await?; + return Ok(()); + } + sqlx::query( r#" INSERT INTO provider_inventory_entries ( @@ -789,6 +934,65 @@ impl ProviderInventoryService { &self, identity: &InventoryIdentity, ) -> Result> { + if self.storage.uses_postgres() { + let pool = self.storage.pool_postgres().await?; + let entry = sqlx::query( + r#" + SELECT last_updated_at, last_refresh_attempt_at, last_refresh_error + FROM provider_inventory_entries + WHERE inventory_key = $1 + "#, + ) + .bind(&identity.inventory_key) + .fetch_optional(pool) + .await?; + + let Some(entry) = entry else { + return Ok(None); + }; + + let last_updated_at = parse_optional_datetime(entry.try_get("last_updated_at")?)?; + let last_refresh_attempt_at = + parse_optional_datetime(entry.try_get("last_refresh_attempt_at")?)?; + let last_refresh_error = entry.try_get("last_refresh_error")?; + + let rows = sqlx::query( + r#" + SELECT model_id, name, family, context_limit, reasoning, recommended + FROM provider_inventory_models + WHERE inventory_key = $1 + ORDER BY ordinal + "#, + ) + .bind(&identity.inventory_key) + .fetch_all(pool) + .await?; + + return Ok(Some(InventorySnapshot { + models: rows + .into_iter() + .map(|row| { + Ok(InventoryModel { + id: row.try_get("model_id")?, + name: row.try_get("name")?, + family: row.try_get("family")?, + context_limit: row + .try_get::, _>("context_limit")? + .map(usize::try_from) + .transpose()?, + reasoning: row.try_get("reasoning")?, + recommended: row + .try_get::, _>("recommended")? + .unwrap_or(false), + }) + }) + .collect::, anyhow::Error>>()?, + last_updated_at, + last_refresh_attempt_at, + last_refresh_error, + })); + } + let pool = self.storage.pool().await?; let entry = sqlx::query( r#" @@ -1235,6 +1439,96 @@ pub async fn create_tables(tx: &mut Transaction<'_, Sqlite>) -> Result<()> { Ok(()) } +pub async fn create_tables_pg(pool: &Pool) -> Result<()> { + sqlx::query( + r#" + CREATE TABLE IF NOT EXISTS provider_inventory_entries ( + inventory_key TEXT PRIMARY KEY, + provider_id TEXT NOT NULL, + provider_family TEXT NOT NULL, + last_updated_at TEXT, + last_refresh_attempt_at TEXT, + last_refresh_error TEXT, + created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP + ) + "#, + ) + .execute(pool) + .await?; + + sqlx::query( + r#" + CREATE TABLE IF NOT EXISTS provider_inventory_models ( + inventory_key TEXT NOT NULL REFERENCES provider_inventory_entries(inventory_key) ON DELETE CASCADE, + ordinal INTEGER NOT NULL, + model_id TEXT NOT NULL, + name TEXT NOT NULL, + family TEXT, + context_limit INTEGER, + reasoning BOOLEAN, + recommended BOOLEAN, + PRIMARY KEY (inventory_key, ordinal) + ) + "#, + ) + .execute(pool) + .await?; + + sqlx::query( + "CREATE INDEX IF NOT EXISTS idx_provider_inventory_provider_id ON provider_inventory_entries(provider_id)", + ) + .execute(pool) + .await?; + + Ok(()) +} + +pub async fn create_tables_in_tx(tx: &mut Transaction<'_, Postgres>) -> Result<()> { + sqlx::query( + r#" + CREATE TABLE IF NOT EXISTS provider_inventory_entries ( + inventory_key TEXT PRIMARY KEY, + provider_id TEXT NOT NULL, + provider_family TEXT NOT NULL, + last_updated_at TEXT, + last_refresh_attempt_at TEXT, + last_refresh_error TEXT, + created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP + ) + "#, + ) + .execute(&mut **tx) + .await?; + + sqlx::query( + r#" + CREATE TABLE IF NOT EXISTS provider_inventory_models ( + inventory_key TEXT NOT NULL REFERENCES provider_inventory_entries(inventory_key) ON DELETE CASCADE, + ordinal INTEGER NOT NULL, + model_id TEXT NOT NULL, + name TEXT NOT NULL, + family TEXT, + context_limit INTEGER, + reasoning BOOLEAN, + recommended BOOLEAN, + PRIMARY KEY (inventory_key, ordinal) + ) + "#, + ) + .execute(&mut **tx) + .await?; + + sqlx::query( + "CREATE INDEX IF NOT EXISTS idx_provider_inventory_provider_id ON provider_inventory_entries(provider_id)", + ) + .execute(&mut **tx) + .await?; + + Ok(()) +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/goose/src/session/postgres_storage.rs b/crates/goose/src/session/postgres_storage.rs new file mode 100644 index 000000000..64026b059 --- /dev/null +++ b/crates/goose/src/session/postgres_storage.rs @@ -0,0 +1,1829 @@ +use super::{ + deserialize_session_model_config, insert_usage_ledger_row_pg, keyword_terms, + message_keyword_clause_pg, message_timestamp_to_datetime, normalized_message_timestamp_sql, + role_to_string, session_sort_at, user_visible_message_sql_pg, Session, SessionInsights, + SessionListCursor, SessionListFilters, SessionListPage, SessionListPageQuery, + SessionListQuery, SessionManager, SessionType, SessionUpdateBuilder, SessionUsageTotals, + CURRENT_SCHEMA_VERSION, SESSION_COUNT_BATCH_SIZE, SESSION_SCHEMA_MIGRATION_LOCK_ID, +}; +use crate::config::GooseMode; +use crate::conversation::message::{Message, MessageUsage}; +use crate::conversation::Conversation; +use crate::providers::base::CostSource; +use crate::recipe::Recipe; +use goose_providers::conversation::token_usage::Usage; +use goose_providers::model::ModelConfig; +use rmcp::model::Role; +use anyhow::Result; +use chrono::{DateTime, TimeZone, Utc}; +use sqlx::postgres::{PgConnectOptions, PgPoolOptions}; +use sqlx::{Pool, Postgres, Row, AssertSqlSafe}; +use std::collections::HashMap; +use std::path::{Path, PathBuf}; +use std::str::FromStr; +use tracing::{info, warn}; + +fn decode_optional_i32_from_row(row: &sqlx::postgres::PgRow, column: &str) -> Result> { + if let Ok(value) = row.try_get::, _>(column) { + return Ok(value); + } + + let value = row.try_get::, _>(column)?; + value + .map(|v| { + i32::try_from(v).map_err(|_| anyhow::anyhow!("{column} out of range for i32: {v}")) + }) + .transpose() +} + +impl sqlx::FromRow<'_, sqlx::postgres::PgRow> for Session { + fn from_row(row: &sqlx::postgres::PgRow) -> Result { + let recipe_json: Option = row.try_get("recipe_json")?; + let recipe = recipe_json.and_then(|json| serde_json::from_str(&json).ok()); + + let user_recipe_values_json: Option = row.try_get("user_recipe_values_json")?; + let user_recipe_values = + user_recipe_values_json.and_then(|json| serde_json::from_str(&json).ok()); + + let provider_name: Option = row.try_get("provider_name").ok().flatten(); + let model_config_json: Option = row.try_get("model_config_json").ok().flatten(); + let model_config = model_config_json + .as_deref() + .and_then(|json| deserialize_session_model_config(provider_name.as_deref(), json)); + + let name: String = { + let name_val: String = row.try_get("name").unwrap_or_default(); + if !name_val.is_empty() { + name_val + } else { + row.try_get("description").unwrap_or_default() + } + }; + + let user_set_name = row.try_get("user_set_name").unwrap_or(false); + + let session_type_str: String = row + .try_get("session_type") + .unwrap_or_else(|_| "user".to_string()); + let session_type = session_type_str.parse().unwrap_or_default(); + + let last_message_at = row + .try_get::, _>("last_message_timestamp") + .ok() + .flatten() + .and_then(message_timestamp_to_datetime); + + Ok(Session { + id: row.try_get("id")?, + working_dir: PathBuf::from(row.try_get::("working_dir")?), + name, + user_set_name, + session_type, + created_at: row.try_get("created_at")?, + updated_at: row.try_get("updated_at")?, + extension_data: serde_json::from_str(&row.try_get::("extension_data")?) + .unwrap_or_default(), + usage: Usage { + input_tokens: decode_optional_i32_from_row(row, "input_tokens") + .map_err(|err| sqlx::Error::Decode(err.into()))?, + output_tokens: decode_optional_i32_from_row(row, "output_tokens") + .map_err(|err| sqlx::Error::Decode(err.into()))?, + total_tokens: decode_optional_i32_from_row(row, "total_tokens") + .map_err(|err| sqlx::Error::Decode(err.into()))?, + cache_read_input_tokens: decode_optional_i32_from_row(row, "cache_read_tokens") + .map_err(|err| sqlx::Error::Decode(err.into()))?, + cache_write_input_tokens: decode_optional_i32_from_row(row, "cache_write_tokens") + .map_err(|err| sqlx::Error::Decode(err.into()))?, + }, + accumulated_usage: Usage { + input_tokens: decode_optional_i32_from_row(row, "accumulated_input_tokens") + .map_err(|err| sqlx::Error::Decode(err.into()))?, + output_tokens: decode_optional_i32_from_row(row, "accumulated_output_tokens") + .map_err(|err| sqlx::Error::Decode(err.into()))?, + total_tokens: decode_optional_i32_from_row(row, "accumulated_total_tokens") + .map_err(|err| sqlx::Error::Decode(err.into()))?, + cache_read_input_tokens: decode_optional_i32_from_row(row, "accumulated_cache_read_tokens") + .map_err(|err| sqlx::Error::Decode(err.into()))?, + cache_write_input_tokens: decode_optional_i32_from_row(row, "accumulated_cache_write_tokens") + .map_err(|err| sqlx::Error::Decode(err.into()))?, + }, + accumulated_cost: row.try_get("accumulated_cost").ok().flatten(), + schedule_id: row.try_get("schedule_id")?, + recipe, + user_recipe_values, + conversation: None, + message_count: row.try_get::("message_count").unwrap_or(0) as usize, + last_message_at, + provider_name, + model_config, + goose_mode: row + .try_get::("goose_mode") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or_default(), + archived_at: row.try_get("archived_at").ok(), + project_id: row.try_get("project_id").ok().flatten(), + parent_session_id: row.try_get("parent_session_id").ok().flatten(), + last_message_snippet: None, + }) + } +} + +pub struct PostgresSessionStorage { + pool: Pool, + initialized: tokio::sync::OnceCell<()>, +} + +impl PostgresSessionStorage { + fn create_pool(db_url: &str) -> Pool { + let options = PgConnectOptions::from_str(db_url) + .expect("Failed to parse GOOSE_SESSION_DB_URL") + .application_name("goosed-session") + .ssl_mode(sqlx::postgres::PgSslMode::Allow); + + PgPoolOptions::new() + .max_connections(8) + .connect_lazy_with(options) + } + + pub fn new(db_url: String) -> Self { + Self { + pool: Self::create_pool(&db_url), + initialized: tokio::sync::OnceCell::new(), + } + } + + pub(crate) async fn pool(&self) -> Result<&Pool> { + self.initialized + .get_or_try_init(|| async { + let schema_exists = sqlx::query_scalar::<_, bool>( + r#"SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'schema_version')"#, + ) + .fetch_one(&self.pool) + .await + .unwrap_or(false); + + if schema_exists { + Self::run_migrations(&self.pool).await?; + } else { + Self::create_schema(&self.pool).await?; + } + Ok::<(), anyhow::Error>(()) + }) + .await?; + Ok(&self.pool) + } + + async fn lock_schema_migrations( + tx: &mut sqlx::Transaction<'_, Postgres>, + ) -> Result<()> { + sqlx::query("SELECT pg_advisory_xact_lock($1)") + .bind(SESSION_SCHEMA_MIGRATION_LOCK_ID) + .execute(&mut **tx) + .await?; + Ok(()) + } + + async fn create_schema(pool: &Pool) -> Result<()> { + let mut tx = pool.begin().await?; + Self::lock_schema_migrations(&mut tx).await?; + + sqlx::query( + r#" + CREATE TABLE IF NOT EXISTS schema_version ( + version INTEGER PRIMARY KEY, + applied_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP + ) + "#, + ) + .execute(&mut *tx) + .await?; + + sqlx::query("INSERT INTO schema_version (version) VALUES ($1) ON CONFLICT DO NOTHING") + .bind(CURRENT_SCHEMA_VERSION) + .execute(&mut *tx) + .await?; + + sqlx::query( + r#" + CREATE TABLE IF NOT EXISTS sessions ( + id TEXT PRIMARY KEY, + name TEXT NOT NULL DEFAULT '', + description TEXT NOT NULL DEFAULT '', + user_set_name BOOLEAN DEFAULT FALSE, + session_type TEXT NOT NULL DEFAULT 'user', + working_dir TEXT NOT NULL, + created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP, + extension_data TEXT DEFAULT '{}', + total_tokens INTEGER, + input_tokens INTEGER, + output_tokens INTEGER, + cache_read_tokens INTEGER, + cache_write_tokens INTEGER, + accumulated_total_tokens INTEGER, + accumulated_input_tokens INTEGER, + accumulated_output_tokens INTEGER, + accumulated_cache_read_tokens INTEGER, + accumulated_cache_write_tokens INTEGER, + accumulated_cost REAL, + schedule_id TEXT, + recipe_json TEXT, + user_recipe_values_json TEXT, + provider_name TEXT, + model_config_json TEXT, + goose_mode TEXT NOT NULL DEFAULT 'auto', + archived_at TIMESTAMPTZ, + project_id TEXT, + parent_session_id TEXT + ) + "#, + ) + .execute(&mut *tx) + .await?; + + sqlx::query( + r#" + CREATE TABLE IF NOT EXISTS usage_ledger ( + id BIGSERIAL PRIMARY KEY, + session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE, + created_timestamp BIGINT NOT NULL, + model TEXT, + input_tokens INTEGER, + output_tokens INTEGER, + total_tokens INTEGER, + cache_read_tokens INTEGER, + cache_write_tokens INTEGER, + cost REAL, + cost_source TEXT, + is_compaction BOOLEAN DEFAULT FALSE + ) + "#, + ) + .execute(&mut *tx) + .await?; + + sqlx::query( + r#" + CREATE TABLE IF NOT EXISTS messages ( + id BIGSERIAL PRIMARY KEY, + message_id TEXT, + session_id TEXT NOT NULL REFERENCES sessions(id), + role TEXT NOT NULL, + content_json TEXT NOT NULL, + created_timestamp BIGINT NOT NULL, + timestamp TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP, + tokens INTEGER, + metadata_json TEXT + ) + "#, + ) + .execute(&mut *tx) + .await?; + + sqlx::query("CREATE INDEX IF NOT EXISTS idx_messages_session ON messages(session_id)") + .execute(&mut *tx) + .await?; + sqlx::query("CREATE INDEX IF NOT EXISTS idx_messages_timestamp ON messages(timestamp)") + .execute(&mut *tx) + .await?; + sqlx::query("CREATE INDEX IF NOT EXISTS idx_messages_message_id ON messages(message_id)") + .execute(&mut *tx) + .await?; + sqlx::query("CREATE INDEX IF NOT EXISTS idx_sessions_updated ON sessions(updated_at DESC)") + .execute(&mut *tx) + .await?; + sqlx::query("CREATE INDEX IF NOT EXISTS idx_sessions_type ON sessions(session_type)") + .execute(&mut *tx) + .await?; + sqlx::query( + "CREATE INDEX IF NOT EXISTS idx_sessions_parent ON sessions(parent_session_id)", + ) + .execute(&mut *tx) + .await?; + sqlx::query( + "CREATE INDEX IF NOT EXISTS idx_usage_ledger_session ON usage_ledger(session_id)", + ) + .execute(&mut *tx) + .await?; + sqlx::query( + "CREATE INDEX IF NOT EXISTS idx_messages_session_created ON messages(session_id, created_timestamp, id)", + ) + .execute(&mut *tx) + .await?; + + tx.commit().await?; + + crate::providers::inventory::create_tables_pg(pool).await?; + + Ok(()) + } + + async fn import_legacy(_pool: &Pool) -> Result<()> { + warn!("Legacy SQLite session import is disabled for the PostgreSQL backend"); + Ok(()) + } + + async fn import_legacy_session(pool: &Pool, session: &Session) -> Result<()> { + let mut tx = pool.begin().await?; + + let recipe_json = match &session.recipe { + Some(recipe) => Some(serde_json::to_string(recipe)?), + None => None, + }; + + let user_recipe_values_json = match &session.user_recipe_values { + Some(user_recipe_values) => Some(serde_json::to_string(user_recipe_values)?), + None => None, + }; + + let model_config_json = match &session.model_config { + Some(model_config) => Some(serde_json::to_string(model_config)?), + None => None, + }; + + sqlx::query( + r#" + INSERT INTO sessions ( + id, name, user_set_name, session_type, working_dir, created_at, updated_at, extension_data, + total_tokens, input_tokens, output_tokens, + cache_read_tokens, cache_write_tokens, + accumulated_total_tokens, accumulated_input_tokens, accumulated_output_tokens, + accumulated_cache_read_tokens, accumulated_cache_write_tokens, + accumulated_cost, + schedule_id, recipe_json, user_recipe_values_json, + provider_name, model_config_json, goose_mode + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $25) + "#, + ) + .bind(&session.id) + .bind(&session.name) + .bind(session.user_set_name) + .bind(session.session_type.to_string()) + .bind(&*session.working_dir.to_string_lossy()) + .bind(session.created_at) + .bind(session.updated_at) + .bind(serde_json::to_string(&session.extension_data)?) + .bind(session.usage.total_tokens) + .bind(session.usage.input_tokens) + .bind(session.usage.output_tokens) + .bind(session.usage.cache_read_input_tokens) + .bind(session.usage.cache_write_input_tokens) + .bind(session.accumulated_usage.total_tokens) + .bind(session.accumulated_usage.input_tokens) + .bind(session.accumulated_usage.output_tokens) + .bind(session.accumulated_usage.cache_read_input_tokens) + .bind(session.accumulated_usage.cache_write_input_tokens) + .bind(session.accumulated_cost) + .bind(&session.schedule_id) + .bind(recipe_json) + .bind(user_recipe_values_json) + .bind(&session.provider_name) + .bind(model_config_json) + .bind(session.goose_mode.to_string()) + .execute(&mut *tx) + .await?; + + tx.commit().await?; + + if let Some(conversation) = &session.conversation { + Self::replace_conversation_inner(pool, &session.id, conversation).await?; + } + Ok(()) + } + + async fn run_migrations(pool: &Pool) -> Result<()> { + let mut tx = pool.begin().await?; + Self::lock_schema_migrations(&mut tx).await?; + + let current_version = Self::get_schema_version(&mut tx).await?; + + if current_version < CURRENT_SCHEMA_VERSION { + info!( + "Running database migrations from v{} to v{}...", + current_version, CURRENT_SCHEMA_VERSION + ); + + for version in (current_version + 1)..=CURRENT_SCHEMA_VERSION { + info!(" Applying migration v{}...", version); + Self::apply_migration(&mut tx, version).await?; + Self::update_schema_version(&mut tx, version).await?; + info!(" ✓ Migration v{} complete", version); + } + + info!("All migrations complete"); + } + + tx.commit().await?; + Ok(()) + } + + async fn get_schema_version(tx: &mut sqlx::Transaction<'_, Postgres>) -> Result { + let table_exists = sqlx::query_scalar::<_, bool>( + r#" + SELECT EXISTS ( + SELECT 1 FROM information_schema.tables + WHERE table_name = 'schema_version' + ) + "#, + ) + .fetch_one(&mut **tx) + .await?; + + if !table_exists { + return Ok(0); + } + + let version = sqlx::query_scalar::<_, i32>("SELECT MAX(version) FROM schema_version") + .fetch_one(&mut **tx) + .await?; + + Ok(version) + } + + async fn update_schema_version( + tx: &mut sqlx::Transaction<'_, Postgres>, + version: i32, + ) -> Result<()> { + sqlx::query("INSERT INTO schema_version (version) VALUES ($1) ON CONFLICT DO NOTHING") + .bind(version) + .execute(&mut **tx) + .await?; + Ok(()) + } + + #[allow(clippy::too_many_lines)] + async fn apply_migration(tx: &mut sqlx::Transaction<'_, Postgres>, version: i32) -> Result<()> { + match version { + 1 => { + sqlx::query( + r#" + CREATE TABLE IF NOT EXISTS schema_version ( + version INTEGER PRIMARY KEY, + applied_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP + ) + "#, + ) + .execute(&mut **tx) + .await?; + } + 2 => { + sqlx::query( + r#" + ALTER TABLE sessions ADD COLUMN user_recipe_values_json TEXT + "#, + ) + .execute(&mut **tx) + .await?; + } + 3 => { + sqlx::query( + r#" + ALTER TABLE messages ADD COLUMN metadata_json TEXT + "#, + ) + .execute(&mut **tx) + .await?; + } + 4 => { + sqlx::query( + r#" + ALTER TABLE sessions ADD COLUMN name TEXT DEFAULT '' + "#, + ) + .execute(&mut **tx) + .await?; + + sqlx::query( + r#" + ALTER TABLE sessions ADD COLUMN user_set_name BOOLEAN DEFAULT FALSE + "#, + ) + .execute(&mut **tx) + .await?; + } + 5 => { + sqlx::query( + r#" + ALTER TABLE sessions ADD COLUMN session_type TEXT NOT NULL DEFAULT 'user' + "#, + ) + .execute(&mut **tx) + .await?; + + sqlx::query("CREATE INDEX idx_sessions_type ON sessions(session_type)") + .execute(&mut **tx) + .await?; + } + 6 => { + sqlx::query( + r#" + ALTER TABLE sessions ADD COLUMN provider_name TEXT + "#, + ) + .execute(&mut **tx) + .await?; + + sqlx::query( + r#" + ALTER TABLE sessions ADD COLUMN model_config_json TEXT + "#, + ) + .execute(&mut **tx) + .await?; + } + 7 => { + sqlx::query( + r#" + ALTER TABLE messages ADD COLUMN message_id TEXT + "#, + ) + .execute(&mut **tx) + .await?; + + sqlx::query( + r#" + UPDATE messages + SET message_id = 'msg_' || session_id || '_' || id + "#, + ) + .execute(&mut **tx) + .await?; + + sqlx::query("CREATE INDEX idx_messages_message_id ON messages(message_id)") + .execute(&mut **tx) + .await?; + } + 8 => { + sqlx::query( + r#" + ALTER TABLE sessions ADD COLUMN goose_mode TEXT NOT NULL DEFAULT 'auto' + "#, + ) + .execute(&mut **tx) + .await?; + } + 9 => { + sqlx::query( + r#" + UPDATE sessions + SET session_type = 'acp' + WHERE session_type = 'user' + AND name = 'ACP Session' + AND user_set_name = FALSE + "#, + ) + .execute(&mut **tx) + .await?; + } + 10 => { + // Check if thread_id column already exists (e.g. fresh schema) + let has_thread_id = sqlx::query_scalar::<_, bool>( + "SELECT EXISTS(SELECT 1 FROM information_schema.columns WHERE table_schema = current_schema() AND table_name = 'sessions' AND column_name = 'thread_id')", + ) + .fetch_one(&mut **tx) + .await? + ; + if !has_thread_id { + sqlx::query("ALTER TABLE sessions ADD COLUMN thread_id TEXT") + .execute(&mut **tx) + .await?; + } + sqlx::query( + "CREATE INDEX IF NOT EXISTS idx_sessions_thread ON sessions(thread_id)", + ) + .execute(&mut **tx) + .await?; + sqlx::query( + "CREATE TABLE IF NOT EXISTS threads ( + id TEXT PRIMARY KEY, + name TEXT NOT NULL DEFAULT 'New Chat', + user_set_name BOOLEAN DEFAULT FALSE, + working_dir TEXT, + created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP, + archived_at TIMESTAMPTZ, + metadata_json TEXT DEFAULT '{}' + )", + ) + .execute(&mut **tx) + .await?; + sqlx::query( + "CREATE TABLE IF NOT EXISTS thread_messages ( + id BIGSERIAL PRIMARY KEY, + thread_id TEXT NOT NULL REFERENCES threads(id), + session_id TEXT, + message_id TEXT, + role TEXT NOT NULL, + content_json TEXT NOT NULL, + created_timestamp BIGINT NOT NULL, + metadata_json TEXT DEFAULT '{}' + )", + ) + .execute(&mut **tx) + .await?; + sqlx::query("CREATE INDEX IF NOT EXISTS idx_thread_messages_thread ON thread_messages(thread_id)") + .execute(&mut **tx) + .await?; + sqlx::query("CREATE INDEX IF NOT EXISTS idx_thread_messages_message_id ON thread_messages(message_id)") + .execute(&mut **tx) + .await?; + } + 11 => { + crate::providers::inventory::create_tables_in_tx(tx).await?; + } + 12 => { + // Add archived_at, project_id columns to sessions. + let has_archived_at = sqlx::query_scalar::<_, bool>( + "SELECT EXISTS(SELECT 1 FROM information_schema.columns WHERE table_schema = current_schema() AND table_name = 'sessions' AND column_name = 'archived_at')", + ) + .fetch_one(&mut **tx) + .await? + ; + if !has_archived_at { + sqlx::query("ALTER TABLE sessions ADD COLUMN archived_at TIMESTAMPTZ") + .execute(&mut **tx) + .await?; + } + + let has_project_id = sqlx::query_scalar::<_, bool>( + "SELECT EXISTS(SELECT 1 FROM information_schema.columns WHERE table_schema = current_schema() AND table_name = 'sessions' AND column_name = 'project_id')", + ) + .fetch_one(&mut **tx) + .await? + ; + if !has_project_id { + sqlx::query("ALTER TABLE sessions ADD COLUMN project_id TEXT") + .execute(&mut **tx) + .await?; + } + } + 13 => { + let has_accumulated_cost = sqlx::query_scalar::<_, bool>( + "SELECT EXISTS(SELECT 1 FROM information_schema.columns WHERE table_schema = current_schema() AND table_name = 'sessions' AND column_name = 'accumulated_cost')", + ) + .fetch_one(&mut **tx) + .await? + ; + if !has_accumulated_cost { + sqlx::query("ALTER TABLE sessions ADD COLUMN accumulated_cost REAL") + .execute(&mut **tx) + .await?; + } + } + 14 => { + for column in [ + "cache_read_tokens", + "cache_write_tokens", + "accumulated_cache_read_tokens", + "accumulated_cache_write_tokens", + ] { + sqlx::query(AssertSqlSafe(format!( + "ALTER TABLE sessions ADD COLUMN IF NOT EXISTS {column} INTEGER" + ))) + .execute(&mut **tx) + .await?; + } + } + 15 => { + sqlx::query( + "ALTER TABLE sessions ADD COLUMN IF NOT EXISTS parent_session_id TEXT", + ) + .execute(&mut **tx) + .await?; + sqlx::query( + "CREATE INDEX IF NOT EXISTS idx_sessions_parent ON sessions(parent_session_id)", + ) + .execute(&mut **tx) + .await?; + sqlx::query( + r#" + CREATE TABLE IF NOT EXISTS usage_ledger ( + id BIGSERIAL PRIMARY KEY, + session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE, + created_timestamp BIGINT NOT NULL, + model TEXT, + input_tokens INTEGER, + output_tokens INTEGER, + total_tokens INTEGER, + cache_read_tokens INTEGER, + cache_write_tokens INTEGER, + cost REAL, + cost_source TEXT, + is_compaction BOOLEAN DEFAULT FALSE + ) + "#, + ) + .execute(&mut **tx) + .await?; + sqlx::query( + "CREATE INDEX IF NOT EXISTS idx_usage_ledger_session ON usage_ledger(session_id)", + ) + .execute(&mut **tx) + .await?; + } + 16 => { + sqlx::query( + "CREATE INDEX IF NOT EXISTS idx_messages_session_created ON messages(session_id, created_timestamp, id)", + ) + .execute(&mut **tx) + .await?; + } + _ => { + anyhow::bail!("Unknown migration version: {}", version); + } + } + + Ok(()) + } + + pub(crate) async fn create_session( + &self, + working_dir: PathBuf, + name: String, + session_type: SessionType, + goose_mode: GooseMode, + ) -> Result { + let pool = self.pool().await?; + let mut tx = pool.begin().await?; + + let today = chrono::Utc::now().format("%Y%m%d").to_string(); + let session = sqlx::query_as( + r#" + INSERT INTO sessions (id, name, user_set_name, session_type, working_dir, extension_data, goose_mode) + VALUES ( + $1 || '_' || CAST(COALESCE(( + SELECT MAX(CAST(SUBSTRING(id, 10) AS INTEGER)) + FROM sessions + WHERE id LIKE $2 || '_%' + ), 0) + 1 AS TEXT), + $3, + FALSE, + $4, + $5, + '{}', + $6 + ) + RETURNING * + "#, + ) + .bind(&today) + .bind(&today) + .bind(&name) + .bind(session_type.to_string()) + .bind(&*working_dir.to_string_lossy()) + .bind(goose_mode.to_string()) + .fetch_one(&mut *tx) + .await?; + + tx.commit().await?; + #[cfg(feature = "telemetry")] + crate::posthog::emit_session_started(); + Ok(session) + } + + pub(crate) async fn get_session(&self, id: &str, include_messages: bool) -> Result { + let pool = self.pool().await?; + let mut session = sqlx::query_as::<_, Session>( + r#" + SELECT id, working_dir, name, description, user_set_name, session_type, created_at, updated_at, extension_data, + total_tokens, input_tokens, output_tokens, + cache_read_tokens, cache_write_tokens, + accumulated_total_tokens, accumulated_input_tokens, accumulated_output_tokens, + accumulated_cache_read_tokens, accumulated_cache_write_tokens, + accumulated_cost, + schedule_id, recipe_json, user_recipe_values_json, + provider_name, model_config_json, goose_mode, + archived_at, project_id, parent_session_id + FROM sessions + WHERE id = $1 + "#, + ) + .bind(id) + .fetch_optional(pool) + .await? + .ok_or_else(|| anyhow::anyhow!("Session not found"))?; + + if include_messages { + let conv = self.get_conversation(&session.id).await?; + session.message_count = conv.messages().len(); + session.last_message_at = conv + .messages() + .iter() + .filter_map(|message| message_timestamp_to_datetime(message.created)) + .max(); + session.conversation = Some(conv); + } else { + let sql = format!( + "SELECT COUNT(*), MAX({}) FROM messages WHERE session_id = $1", + normalized_message_timestamp_sql("created_timestamp") + ); + let (count, last_message_timestamp): (i64, Option) = + sqlx::query_as(AssertSqlSafe(sql)) + .bind(&session.id) + .fetch_one(pool) + .await?; + session.message_count = count as usize; + session.last_message_at = + last_message_timestamp.and_then(message_timestamp_to_datetime); + } + + Ok(session) + } + + #[allow(clippy::too_many_lines)] + pub(crate) async fn apply_update(&self, builder: SessionUpdateBuilder<'_>) -> Result<()> { + let mut updates = Vec::new(); + let mut query = String::from("UPDATE sessions SET "); + let mut bind_index = 1; + + macro_rules! add_update { + ($field:expr, $name:expr) => { + if $field.is_some() { + if !updates.is_empty() { + query.push_str(", "); + } + updates.push($name); + query.push_str($name); + query.push_str(&format!(" = ${bind_index}")); + bind_index += 1; + } + }; + } + + add_update!(builder.name, "name"); + add_update!(builder.user_set_name, "user_set_name"); + add_update!(builder.session_type, "session_type"); + add_update!(builder.working_dir, "working_dir"); + add_update!(builder.extension_data, "extension_data"); + add_update!(builder.usage, "total_tokens"); + add_update!(builder.usage, "input_tokens"); + add_update!(builder.usage, "output_tokens"); + add_update!(builder.usage, "cache_read_tokens"); + add_update!(builder.usage, "cache_write_tokens"); + add_update!(builder.accumulated_usage, "accumulated_total_tokens"); + add_update!(builder.accumulated_usage, "accumulated_input_tokens"); + add_update!(builder.accumulated_usage, "accumulated_output_tokens"); + add_update!(builder.accumulated_usage, "accumulated_cache_read_tokens"); + add_update!(builder.accumulated_usage, "accumulated_cache_write_tokens"); + add_update!(builder.accumulated_cost, "accumulated_cost"); + add_update!(builder.schedule_id, "schedule_id"); + add_update!(builder.recipe, "recipe_json"); + add_update!(builder.user_recipe_values, "user_recipe_values_json"); + add_update!(builder.provider_name, "provider_name"); + add_update!(builder.model_config, "model_config_json"); + add_update!(builder.goose_mode, "goose_mode"); + add_update!(builder.archived_at, "archived_at"); + + add_update!(builder.project_id, "project_id"); + add_update!(builder.parent_session_id, "parent_session_id"); + + if updates.is_empty() { + return Ok(()); + } + + query.push_str(", "); + query.push_str(&format!( + "updated_at = CURRENT_TIMESTAMP WHERE id = ${bind_index}" + )); + + let mut q = sqlx::query(AssertSqlSafe(query)); + + if let Some(name) = builder.name { + q = q.bind(name); + } + if let Some(user_set_name) = builder.user_set_name { + q = q.bind(user_set_name); + } + if let Some(session_type) = builder.session_type { + q = q.bind(session_type.to_string()); + } + if let Some(wd) = builder.working_dir { + q = q.bind(wd.to_string_lossy().to_string()); + } + if let Some(ed) = builder.extension_data { + q = q.bind(serde_json::to_string(&ed)?); + } + if let Some(u) = builder.usage { + q = q + .bind(u.total_tokens) + .bind(u.input_tokens) + .bind(u.output_tokens) + .bind(u.cache_read_input_tokens) + .bind(u.cache_write_input_tokens); + } + if let Some(u) = builder.accumulated_usage { + q = q + .bind(u.total_tokens) + .bind(u.input_tokens) + .bind(u.output_tokens) + .bind(u.cache_read_input_tokens) + .bind(u.cache_write_input_tokens); + } + if let Some(ac) = builder.accumulated_cost { + q = q.bind(ac); + } + if let Some(sid) = builder.schedule_id { + q = q.bind(sid); + } + if let Some(recipe) = builder.recipe { + let recipe_json = recipe.map(|r| serde_json::to_string(&r)).transpose()?; + q = q.bind(recipe_json); + } + if let Some(user_recipe_values) = builder.user_recipe_values { + let user_recipe_values_json = user_recipe_values + .map(|urv| serde_json::to_string(&urv)) + .transpose()?; + q = q.bind(user_recipe_values_json); + } + if let Some(provider_name) = builder.provider_name { + q = q.bind(provider_name); + } + if let Some(model_config) = builder.model_config { + let model_config_json = model_config + .map(|mc| serde_json::to_string(&mc)) + .transpose()?; + q = q.bind(model_config_json); + } + if let Some(goose_mode) = builder.goose_mode { + q = q.bind(goose_mode.to_string()); + } + if let Some(ref archived_at) = builder.archived_at { + q = q.bind(archived_at.as_ref()); + } + + if let Some(ref project_id) = builder.project_id { + q = q.bind(project_id.as_ref()); + } + if let Some(ref parent_session_id) = builder.parent_session_id { + q = q.bind(parent_session_id.as_ref()); + } + + let pool = self.pool().await?; + let mut tx = pool.begin().await?; + q = q.bind(&builder.session_id); + let result = q.execute(&mut *tx).await?; + + if result.rows_affected() == 0 { + return Err(anyhow::anyhow!("Session not found: {}", builder.session_id)); + } + + tx.commit().await?; + Ok(()) + } + + pub(crate) async fn get_conversation(&self, session_id: &str) -> Result { + let pool = self.pool().await?; + let rows = sqlx::query_as::<_, (String, String, i64, Option, Option)>( + // Order by created_timestamp, then by id to break ties. created_timestamp is in seconds, + // so messages created in the same second (e.g., tool request and response) need to + // maintain their insertion order via the auto-increment id. + "SELECT role, content_json, created_timestamp, metadata_json, message_id FROM messages WHERE session_id = $1 ORDER BY created_timestamp, id", + ) + .bind(session_id) + .fetch_all(pool) + .await?; + + let mut messages = Vec::new(); + for (role_str, content_json, created_timestamp, metadata_json, message_id) in + rows.into_iter() + { + let role = match role_str.as_str() { + "user" => Role::User, + "assistant" => Role::Assistant, + _ => continue, + }; + + let content = serde_json::from_str(&content_json)?; + let metadata = metadata_json + .and_then(|json| serde_json::from_str(&json).ok()) + .unwrap_or_default(); + + let mut message = Message::new(role, created_timestamp, content); + message.metadata = metadata; + if let Some(id) = message_id { + message = message.with_id(id); + } + messages.push(message); + } + + Ok(Conversation::new_unvalidated(messages)) + } + + pub(crate) async fn add_message(&self, session_id: &str, message: &Message) -> Result<()> { + let pool = self.pool().await?; + let mut tx = pool.begin().await?; + + let metadata_json = serde_json::to_string(&message.metadata)?; + + let message_id = message + .id + .clone() + .unwrap_or_else(|| format!("msg_{}_{}", session_id, uuid::Uuid::new_v4())); + + sqlx::query( + r#" + INSERT INTO messages (message_id, session_id, role, content_json, created_timestamp, metadata_json) + VALUES ($1, $2, $3, $4, $5, $6) + "#, + ) + .bind(message_id) + .bind(session_id) + .bind(role_to_string(&message.role)) + .bind(serde_json::to_string(&message.content)?) + .bind(message.created) + .bind(metadata_json) + .execute(&mut *tx) + .await?; + + sqlx::query("UPDATE sessions SET updated_at = CURRENT_TIMESTAMP WHERE id = $1") + .bind(session_id) + .execute(&mut *tx) + .await?; + + tx.commit().await?; + Ok(()) + } + + async fn replace_conversation_inner( + pool: &Pool, + session_id: &str, + conversation: &Conversation, + ) -> Result<()> { + let mut tx = pool.begin().await?; + + sqlx::query("DELETE FROM messages WHERE session_id = $1") + .bind(session_id) + .execute(&mut *tx) + .await?; + + for message in conversation.messages() { + let metadata_json = serde_json::to_string(&message.metadata)?; + + let message_id = message + .id + .clone() + .unwrap_or_else(|| format!("msg_{}_{}", session_id, uuid::Uuid::new_v4())); + + sqlx::query( + r#" + INSERT INTO messages (message_id, session_id, role, content_json, created_timestamp, metadata_json) + VALUES ($1, $2, $3, $4, $5, $6) + "#, + ) + .bind(message_id) + .bind(session_id) + .bind(role_to_string(&message.role)) + .bind(serde_json::to_string(&message.content)?) + .bind(message.created) + .bind(metadata_json) + .execute(&mut *tx) + .await?; + } + + tx.commit().await?; + Ok(()) + } + + pub async fn replace_conversation( + &self, + session_id: &str, + conversation: &Conversation, + ) -> Result<()> { + let pool = self.pool().await?; + Self::replace_conversation_inner(pool, session_id, conversation).await + } + + pub(crate) async fn list_sessions_matching(&self, query: SessionListQuery<'_>) -> Result> { + let filters = &query.filters; + if matches!(filters.types, Some(types) if types.is_empty()) { + return Ok(Vec::new()); + } + + let keywords = keyword_terms(filters.keyword); + let mut where_clauses = Vec::new(); + let mut having_clauses = Vec::new(); + let normalized_message_timestamp = normalized_message_timestamp_sql("m.created_timestamp"); + let sort_timestamp_sql = format!( + "COALESCE(MAX({normalized_message_timestamp}), EXTRACT(EPOCH FROM s.updated_at))" + ); + if let Some(types) = filters.types { + let placeholders = (1..=types.len()) + .map(|index| format!("${index}")) + .collect::>() + .join(", "); + where_clauses.push(format!("s.session_type IN ({})", placeholders)); + } + if filters.working_dir.is_some() { + where_clauses.push(format!( + "s.working_dir = ${}", + filters.types.map_or(1, |types| types.len() + 1) + )); + } + if !keywords.is_empty() { + where_clauses.push(message_keyword_clause_pg(keywords.len())); + } + if query.cursor.is_some() { + having_clauses.push(format!( + "({sort_timestamp_sql} < $1 OR ({sort_timestamp_sql} = $2 AND s.id < $3))" + )); + } + + let where_clause = if where_clauses.is_empty() { + String::new() + } else { + format!("WHERE {}", where_clauses.join(" AND ")) + }; + let having_clause = if having_clauses.is_empty() { + String::new() + } else { + format!("HAVING {}", having_clauses.join(" AND ")) + }; + let message_join = if filters.only_sessions_with_messages { + "JOIN messages m ON s.id = m.session_id" + } else { + "LEFT JOIN messages m ON s.id = m.session_id" + }; + let order_by = "ORDER BY sort_timestamp DESC, s.id DESC"; + let limit_clause = if query.limit.is_some() { + "LIMIT $1" + } else { + "" + }; + + let sql = format!( + r#" + SELECT s.id, s.working_dir, s.name, s.description, s.user_set_name, s.session_type, s.created_at, s.updated_at, s.extension_data, + s.total_tokens, s.input_tokens, s.output_tokens, + s.cache_read_tokens, s.cache_write_tokens, + s.accumulated_total_tokens, s.accumulated_input_tokens, s.accumulated_output_tokens, + s.accumulated_cache_read_tokens, s.accumulated_cache_write_tokens, + s.accumulated_cost, + s.schedule_id, s.recipe_json, s.user_recipe_values_json, + s.provider_name, s.model_config_json, s.goose_mode, + s.archived_at, s.project_id, s.parent_session_id, + COUNT(m.id) as message_count, + MAX({}) as last_message_timestamp, + {} as sort_timestamp + FROM sessions s + {} + {} + GROUP BY s.id + {} + {} + {} + "#, + normalized_message_timestamp, + sort_timestamp_sql, + message_join, + where_clause, + having_clause, + order_by, + limit_clause + ); + + let mut q = sqlx::query_as::<_, Session>(AssertSqlSafe(sql)); + if let Some(types) = filters.types { + for session_type in types { + q = q.bind(session_type.to_string()); + } + } + if let Some(working_dir) = filters.working_dir { + q = q.bind(working_dir.to_string_lossy().to_string()); + } + for term in keywords { + q = q.bind(term); + } + if let Some(cursor) = query.cursor { + let sort_at = cursor.sort_at.timestamp(); + q = q.bind(sort_at); + q = q.bind(sort_at); + q = q.bind(&cursor.session_id); + } + if let Some(limit) = query.limit { + q = q.bind(limit as i64); + } + + let pool = self.pool().await?; + q.fetch_all(pool).await.map_err(Into::into) + } + + pub(crate) async fn list_sessions_by_types(&self, types: Option<&[SessionType]>) -> Result> { + self.list_sessions_matching(SessionListQuery { + filters: SessionListFilters { + types, + ..Default::default() + }, + ..Default::default() + }) + .await + } + + pub(crate) async fn list_sessions_paged( + &self, + query: SessionListPageQuery<'_>, + ) -> Result { + if matches!(query.filters.types, Some(types) if types.is_empty()) || query.page_size == 0 { + return Ok(SessionListPage { + sessions: Vec::new(), + next_cursor: None, + }); + } + + let page_size = query.page_size; + let include_last_message_snippet = query.include_last_message_snippet; + let mut sessions = self + .list_sessions_matching(SessionListQuery { + filters: query.filters, + cursor: query.cursor, + limit: Some(page_size + 1), + }) + .await?; + let has_next_page = sessions.len() > page_size; + let next_cursor = if has_next_page { + let anchor = &sessions[page_size - 1]; + Some(SessionListCursor { + sort_at: session_sort_at(anchor), + session_id: anchor.id.clone(), + }) + } else { + None + }; + if has_next_page { + sessions.truncate(page_size); + } + if include_last_message_snippet { + // Snippet hydration is SQLite-specific today; PG sessions still list correctly. + } + + Ok(SessionListPage { + sessions, + next_cursor, + }) + } + + pub(crate) async fn list_sessions(&self) -> Result> { + self.list_sessions_by_types(Some(&[SessionType::User, SessionType::Scheduled])) + .await + } + + pub(crate) async fn delete_session(&self, session_id: &str) -> Result<()> { + let pool = self.pool().await?; + let mut tx = pool.begin().await?; + + let exists = + sqlx::query_scalar::<_, bool>("SELECT EXISTS(SELECT 1 FROM sessions WHERE id = $1)") + .bind(session_id) + .fetch_one(&mut *tx) + .await?; + + if !exists { + return Err(anyhow::anyhow!("Session not found")); + } + + sqlx::query("DELETE FROM usage_ledger WHERE session_id = $1") + .bind(session_id) + .execute(&mut *tx) + .await?; + + sqlx::query("DELETE FROM messages WHERE session_id = $1") + .bind(session_id) + .execute(&mut *tx) + .await?; + + sqlx::query("DELETE FROM sessions WHERE id = $1") + .bind(session_id) + .execute(&mut *tx) + .await?; + + tx.commit().await?; + Ok(()) + } + + pub(crate) async fn get_insights(&self, types: &[SessionType]) -> Result { + if types.is_empty() { + return Ok(SessionInsights { + total_sessions: 0, + total_tokens: 0, + }); + } + + let placeholders = (1..=types.len()) + .map(|index| format!("${index}")) + .collect::>() + .join(", "); + let query = format!( + r#" + SELECT COUNT(*) as total_sessions, + COALESCE(SUM(COALESCE(accumulated_total_tokens, total_tokens, 0)), 0) as total_tokens + FROM sessions + WHERE session_type IN ({}) + "#, + placeholders + ); + + let pool = self.pool().await?; + let mut q = sqlx::query_as::<_, (i64, Option)>(AssertSqlSafe(query)); + for t in types { + q = q.bind(t.to_string()); + } + + let row = q.fetch_one(pool).await?; + + Ok(SessionInsights { + total_sessions: row.0 as usize, + total_tokens: row.1.unwrap_or(0), + }) + } + + pub(crate) async fn update_project_for_session_types( + &self, + id: &str, + project_id: Option, + session_types: &[SessionType], + ) -> Result { + if session_types.is_empty() { + return Ok(false); + } + + let placeholders = (2..=session_types.len() + 1) + .map(|index| format!("${index}")) + .collect::>() + .join(", "); + let query = format!( + "UPDATE sessions SET project_id = $1, updated_at = CURRENT_TIMESTAMP \ + WHERE id = ${} AND session_type IN ({placeholders})", + session_types.len() + 2 + ); + let pool = self.pool().await?; + let mut query = sqlx::query(AssertSqlSafe(query)).bind(project_id).bind(id); + for session_type in session_types { + query = query.bind(session_type.to_string()); + } + + Ok(query.execute(pool).await?.rows_affected() == 1) + } + + pub(crate) async fn record_usage_metrics( + &self, + session_id: &str, + schedule_id: Option, + current_usage: Usage, + model: &str, + ledger: &MessageUsage, + ) -> Result<()> { + let pool = self.pool().await?; + let mut tx = pool.begin().await?; + + sqlx::query( + r#" + INSERT INTO usage_ledger ( + session_id, created_timestamp, + input_tokens, output_tokens, total_tokens, + cache_read_tokens, cache_write_tokens, + cost, cost_source + ) + SELECT s.id, EXTRACT(EPOCH FROM NOW())::bigint, + GREATEST(COALESCE(s.accumulated_input_tokens, 0) - l.input_sum, 0), + GREATEST(COALESCE(s.accumulated_output_tokens, 0) - l.output_sum, 0), + GREATEST(COALESCE(s.accumulated_total_tokens, 0) - l.total_sum, 0), + GREATEST(COALESCE(s.accumulated_cache_read_tokens, 0) - l.cache_read_sum, 0), + GREATEST(COALESCE(s.accumulated_cache_write_tokens, 0) - l.cache_write_sum, 0), + CASE WHEN s.accumulated_cost IS NULL OR s.accumulated_cost <= l.cost_sum THEN NULL + ELSE s.accumulated_cost - l.cost_sum END, + 'carried_forward' + FROM sessions s, + (SELECT COALESCE(SUM(input_tokens), 0) AS input_sum, + COALESCE(SUM(output_tokens), 0) AS output_sum, + COALESCE(SUM(total_tokens), 0) AS total_sum, + COALESCE(SUM(cache_read_tokens), 0) AS cache_read_sum, + COALESCE(SUM(cache_write_tokens), 0) AS cache_write_sum, + COALESCE(SUM(cost), 0.0) AS cost_sum + FROM usage_ledger WHERE session_id = $1) l + WHERE s.id = $2 + AND (COALESCE(s.accumulated_input_tokens, 0) > l.input_sum + OR COALESCE(s.accumulated_output_tokens, 0) > l.output_sum + OR COALESCE(s.accumulated_total_tokens, 0) > l.total_sum + OR COALESCE(s.accumulated_cost, 0.0) > l.cost_sum + 1e-9) + "#, + ) + .bind(session_id) + .bind(session_id) + .execute(&mut *tx) + .await?; + + sqlx::query( + r#" + UPDATE sessions SET + schedule_id = $1, + total_tokens = $2, input_tokens = $3, output_tokens = $4, + cache_read_tokens = $5, cache_write_tokens = $6, + accumulated_total_tokens = COALESCE(accumulated_total_tokens, 0) + $7, + accumulated_input_tokens = COALESCE(accumulated_input_tokens, 0) + $8, + accumulated_output_tokens = COALESCE(accumulated_output_tokens, 0) + $9, + accumulated_cache_read_tokens = COALESCE(accumulated_cache_read_tokens, 0) + $10, + accumulated_cache_write_tokens = COALESCE(accumulated_cache_write_tokens, 0) + $11, + accumulated_cost = CASE + WHEN $12 IS NULL THEN accumulated_cost + ELSE COALESCE(accumulated_cost, 0) + $12 + END, + updated_at = CURRENT_TIMESTAMP + WHERE id = $13 + "#, + ) + .bind(schedule_id) + .bind(current_usage.total_tokens) + .bind(current_usage.input_tokens) + .bind(current_usage.output_tokens) + .bind(current_usage.cache_read_input_tokens) + .bind(current_usage.cache_write_input_tokens) + .bind(ledger.total_tokens.unwrap_or(0)) + .bind(ledger.input_tokens.unwrap_or(0)) + .bind(ledger.output_tokens.unwrap_or(0)) + .bind(ledger.cache_read_tokens.unwrap_or(0)) + .bind(ledger.cache_write_tokens.unwrap_or(0)) + .bind(ledger.cost) + .bind(session_id) + .execute(&mut *tx) + .await?; + + insert_usage_ledger_row_pg(&mut tx, session_id, Some(model), ledger).await?; + + tx.commit().await?; + Ok(()) + } + + pub(crate) async fn get_session_usage_totals(&self, session_id: &str) -> Result { + let pool = self.pool().await?; + let rows = sqlx::query_as::< + _, + ( + Option, + Option, + Option, + Option, + Option, + Option, + Option, + Option, + Option, + Option, + Option, + Option, + ), + >( + r#" + WITH RECURSIVE tree(id) AS ( + SELECT id FROM sessions WHERE id = $1 + UNION + SELECT s.id FROM sessions s JOIN tree ON s.parent_session_id = tree.id + ) + SELECT + s.accumulated_input_tokens, s.accumulated_output_tokens, s.accumulated_total_tokens, + s.accumulated_cache_read_tokens, s.accumulated_cache_write_tokens, s.accumulated_cost, + SUM(u.input_tokens), SUM(u.output_tokens), SUM(u.total_tokens), + SUM(u.cache_read_tokens), SUM(u.cache_write_tokens), SUM(u.cost) + FROM sessions s + LEFT JOIN usage_ledger u ON u.session_id = s.id + WHERE s.id IN (SELECT id FROM tree) + GROUP BY s.id + "#, + ) + .bind(session_id) + .fetch_all(pool) + .await?; + + let mut input = 0i64; + let mut output = 0i64; + let mut total = 0i64; + let mut cache_read = 0i64; + let mut cache_write = 0i64; + let mut cost = 0.0f64; + let mut has_cost = false; + + for row in rows { + let ( + acc_input, + acc_output, + acc_total, + acc_cache_read, + acc_cache_write, + acc_cost, + ledger_input, + ledger_output, + ledger_total, + ledger_cache_read, + ledger_cache_write, + ledger_cost, + ) = row; + + input += acc_input.unwrap_or(0).max(ledger_input.unwrap_or(0)); + output += acc_output.unwrap_or(0).max(ledger_output.unwrap_or(0)); + total += acc_total.unwrap_or(0).max(ledger_total.unwrap_or(0)); + cache_read += acc_cache_read.unwrap_or(0).max(ledger_cache_read.unwrap_or(0)); + cache_write += acc_cache_write.unwrap_or(0).max(ledger_cache_write.unwrap_or(0)); + if let Some(c) = acc_cost.or(ledger_cost) { + cost += c; + has_cost = true; + } + } + + Ok(SessionUsageTotals { + accumulated_usage: Usage { + input_tokens: Some(input as i32), + output_tokens: Some(output as i32), + total_tokens: Some(total as i32), + cache_read_input_tokens: Some(cache_read as i32), + cache_write_input_tokens: Some(cache_write as i32), + }, + accumulated_cost: if has_cost { Some(cost) } else { None }, + }) + } + + pub(crate) async fn export_session(&self, id: &str) -> Result { + let session = self.get_session(id, true).await?; + serde_json::to_string_pretty(&session).map_err(Into::into) + } + + pub(crate) async fn import_session( + &self, + session_manager: &SessionManager, + json: &str, + session_type_override: Option, + ) -> Result { + let normalized = crate::session::import_formats::convert_to_goose_session_json(json)?; + let import: Session = serde_json::from_str(&normalized)?; + + let session = self + .create_session( + import.working_dir.clone(), + import.name.clone(), + session_type_override.unwrap_or(import.session_type), + import.goose_mode, + ) + .await?; + + let mut builder = session_manager + .update(&session.id) + .extension_data(import.extension_data) + .usage(import.usage) + .accumulated_usage(import.accumulated_usage) + .accumulated_cost(import.accumulated_cost) + .schedule_id(import.schedule_id) + .recipe(import.recipe) + .user_recipe_values(import.user_recipe_values); + + if import.user_set_name { + builder = builder.user_provided_name(import.name.clone()); + } + + builder.apply().await?; + + if let Some(conversation) = import.conversation { + self.replace_conversation(&session.id, &conversation) + .await?; + } + + self.get_session(&session.id, true).await + } + + pub(crate) async fn copy_session( + &self, + session_manager: &SessionManager, + session_id: &str, + new_name: String, + ) -> Result { + let original_session = self.get_session(session_id, true).await?; + + let new_session = self + .create_session( + original_session.working_dir.clone(), + new_name, + original_session.session_type, + original_session.goose_mode, + ) + .await?; + + 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); + + if let Some(project_id) = original_session.project_id { + builder = builder.project_id(Some(project_id)); + } + 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 = builder.goose_mode(original_session.goose_mode); + + builder.apply().await?; + + if let Some(conversation) = original_session.conversation { + self.replace_conversation(&new_session.id, &conversation) + .await?; + } + + self.get_session(&new_session.id, true).await + } + + pub(crate) async fn truncate_conversation(&self, session_id: &str, timestamp: i64) -> Result<()> { + let pool = self.pool().await?; + sqlx::query("DELETE FROM messages WHERE session_id = $1 AND created_timestamp >= $2") + .bind(session_id) + .bind(timestamp) + .execute(pool) + .await?; + + Ok(()) + } + + pub(crate) async fn truncate_conversation_from_message( + &self, + session_id: &str, + message_id: &str, + ) -> Result<()> { + let pool = self.pool().await?; + let mut tx = pool.begin().await?; + + let boundary = sqlx::query_as::<_, (i64, i64)>( + "SELECT id, created_timestamp FROM messages WHERE session_id = $1 AND message_id = $2 ORDER BY created_timestamp, id LIMIT 1", + ) + .bind(session_id) + .bind(message_id) + .fetch_optional(&mut *tx) + .await?; + + if let Some((boundary_id, boundary_timestamp)) = boundary { + sqlx::query( + "DELETE FROM messages WHERE session_id = $1 AND (created_timestamp > $2 OR (created_timestamp = $3 AND id >= $4))", + ) + .bind(session_id) + .bind(boundary_timestamp) + .bind(boundary_timestamp) + .bind(boundary_id) + .execute(&mut *tx) + .await?; + } + + tx.commit().await?; + Ok(()) + } + + pub(crate) async fn search_chat_history( + &self, + _query: &str, + _limit: Option, + _after_date: Option>, + _before_date: Option>, + _exclude_session_id: Option, + _session_types: Vec, + ) -> Result { + anyhow::bail!("Chat history search is not yet implemented for PostgreSQL session storage") + } + + pub(crate) async fn update_message_metadata( + &self, + session_id: &str, + message_id: &str, + f: F, + ) -> Result<()> + where + F: FnOnce( + crate::conversation::message::MessageMetadata, + ) -> crate::conversation::message::MessageMetadata, + { + let pool = self.pool().await?; + let mut tx = pool.begin().await?; + + let current_metadata_json = sqlx::query_scalar::<_, String>( + "SELECT metadata_json FROM messages WHERE message_id = $1 AND session_id = $2", + ) + .bind(message_id) + .bind(session_id) + .fetch_one(&mut *tx) + .await?; + + let current_metadata: crate::conversation::message::MessageMetadata = + serde_json::from_str(¤t_metadata_json)?; + + let new_metadata = f(current_metadata); + let metadata_json = serde_json::to_string(&new_metadata)?; + + sqlx::query( + "UPDATE messages SET metadata_json = $1 WHERE message_id = $2 AND session_id = $3", + ) + .bind(metadata_json) + .bind(message_id) + .bind(session_id) + .execute(&mut *tx) + .await?; + + tx.commit().await?; + + Ok(()) + } + + pub(crate) async fn update_tool_request_meta( + &self, + session_id: &str, + tool_call_id: &str, + patch: serde_json::Value, + ) -> Result<()> { + use crate::conversation::message::MessageContent; + + let pool = self.pool().await?; + let rows = sqlx::query_as::<_, (Option, String)>( + "SELECT message_id, content_json FROM messages \ + WHERE session_id = $1 \ + ORDER BY id DESC \ + LIMIT 100", + ) + .bind(session_id) + .fetch_all(pool) + .await?; + + for (message_id, content_json) in rows { + let content: Vec = serde_json::from_str(&content_json)?; + let contains_tool_request = content.iter().any(|block| { + matches!( + block, + MessageContent::ToolRequest(tool_request) + if tool_request.id == tool_call_id + ) + }); + if contains_tool_request { + let Some(message_id) = message_id else { + return Ok(()); + }; + return self + .update_tool_request_meta_by_message_id( + session_id, + &message_id, + tool_call_id, + patch, + ) + .await; + } + } + + Ok(()) + } + + async fn update_tool_request_meta_by_message_id( + &self, + session_id: &str, + message_id: &str, + tool_call_id: &str, + patch: serde_json::Value, + ) -> Result<()> { + use crate::conversation::message::MessageContent; + + let pool = self.pool().await?; + let mut tx = pool.begin().await?; + + let rows = sqlx::query_as::<_, (i64, String)>( + "SELECT id, content_json FROM messages \ + WHERE session_id = $1 AND message_id = $2 \ + ORDER BY id ASC", + ) + .bind(session_id) + .bind(message_id) + .fetch_all(&mut *tx) + .await?; + + for (row_id, content_json) in rows { + let mut content: Vec = serde_json::from_str(&content_json)?; + let mut found = false; + for block in &mut content { + if let MessageContent::ToolRequest(tr) = block { + if tr.id == tool_call_id { + tr.tool_meta = Some(super::merge_tool_meta(tr.tool_meta.take(), &patch)); + found = true; + break; + } + } + } + if !found { + continue; + } + + let updated_json = serde_json::to_string(&content)?; + sqlx::query("UPDATE messages SET content_json = $1 WHERE id = $2") + .bind(updated_json) + .bind(row_id) + .execute(&mut *tx) + .await?; + tx.commit().await?; + return Ok(()); + } + + tx.commit().await?; + Ok(()) + } +} diff --git a/crates/goose/src/session/session_manager.rs b/crates/goose/src/session/session_manager.rs index f61d2473f..7ddc7adb3 100644 --- a/crates/goose/src/session/session_manager.rs +++ b/crates/goose/src/session/session_manager.rs @@ -17,7 +17,11 @@ use goose_providers::model::ModelConfig; use rmcp::model::Role; use serde::{Deserialize, Serialize}; use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions}; -use sqlx::{AssertSqlSafe, Pool, Sqlite}; +use sqlx::{AssertSqlSafe, Pool, Postgres, Sqlite}; + +#[path = "postgres_storage.rs"] +mod postgres_storage; +use postgres_storage::PostgresSessionStorage; use std::collections::HashMap; use std::fs; use std::path::{Path, PathBuf}; @@ -28,7 +32,8 @@ pub const CURRENT_SCHEMA_VERSION: i32 = 16; pub const SESSIONS_FOLDER: &str = "sessions"; pub const DB_NAME: &str = "sessions.db"; const MILLISECOND_TIMESTAMP_THRESHOLD: i64 = 10_000_000_000; -const SESSION_COUNT_BATCH_SIZE: usize = 900; +pub(crate) const SESSION_SCHEMA_MIGRATION_LOCK_ID: i64 = 0x474f_4f53_4553_4442; +pub(crate) const SESSION_COUNT_BATCH_SIZE: usize = 900; #[derive( Debug, @@ -55,8 +60,19 @@ pub enum SessionType { Acp, } +fn session_storage_for_environment(data_dir: PathBuf) -> SessionStorage { + if let Ok(db_url) = std::env::var("GOOSE_SESSION_DB_URL") { + let db_url = db_url.trim(); + if !db_url.is_empty() { + info!("Using PostgreSQL session storage (GOOSE_SESSION_DB_URL)"); + return SessionStorage::from_postgres(db_url.to_string()); + } + } + SessionStorage::from_sqlite(data_dir) +} + static SESSION_STORAGE: LazyLock> = - LazyLock::new(|| Arc::new(SessionStorage::new(Paths::data_dir()))); + LazyLock::new(|| Arc::new(session_storage_for_environment(Paths::data_dir()))); #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Session { @@ -346,13 +362,13 @@ pub(crate) struct SessionListPageQuery<'a> { } #[derive(Debug, Default)] -struct SessionListQuery<'a> { +pub(crate) struct SessionListQuery<'a> { filters: SessionListFilters<'a>, cursor: Option<&'a SessionListCursor>, limit: Option, } -fn keyword_terms(query: Option<&str>) -> Vec { +pub(crate) fn keyword_terms(query: Option<&str>) -> Vec { query .unwrap_or_default() .split_whitespace() @@ -397,7 +413,7 @@ pub struct SessionNameUpdate { impl SessionManager { pub fn new(data_dir: PathBuf) -> Self { Self { - storage: Arc::new(SessionStorage::new(data_dir)), + storage: Arc::new(session_storage_for_environment(data_dir)), } } @@ -700,10 +716,19 @@ impl SessionManager { } } -pub struct SessionStorage { +pub(crate) struct SqliteSessionStorage { pool: Pool, initialized: tokio::sync::OnceCell<()>, session_dir: PathBuf, +} + +enum SessionStorageBackend { + Sqlite(SqliteSessionStorage), + Postgres(PostgresSessionStorage), +} + +pub struct SessionStorage { + backend: SessionStorageBackend, action_required: Arc, } @@ -714,7 +739,7 @@ pub(crate) fn role_to_string(role: &Role) -> &'static str { } } -fn message_timestamp_to_datetime(timestamp: i64) -> Option> { +pub(crate) fn message_timestamp_to_datetime(timestamp: i64) -> Option> { let timestamp = if timestamp > MILLISECOND_TIMESTAMP_THRESHOLD { timestamp / 1000 } else { @@ -723,7 +748,7 @@ fn message_timestamp_to_datetime(timestamp: i64) -> Option> { Utc.timestamp_opt(timestamp, 0).single() } -fn normalized_message_timestamp_sql(column: &str) -> String { +pub(crate) fn normalized_message_timestamp_sql(column: &str) -> String { format!( "CASE WHEN {column} > {MILLISECOND_TIMESTAMP_THRESHOLD} THEN {column} / 1000 ELSE {column} END" ) @@ -733,7 +758,38 @@ fn user_visible_message_sql(column: &str) -> String { format!("COALESCE(json_extract({column}, '$.userVisible'), 1) != 0") } -fn session_sort_at(session: &Session) -> DateTime { +pub(crate) fn user_visible_message_sql_pg(column: &str) -> String { + format!( + "COALESCE(({column}::jsonb->>'userVisible')::boolean, TRUE) = TRUE" + ) +} + +pub(crate) fn message_keyword_clause_pg(keyword_count: usize) -> String { + let keyword_clauses = (1..=keyword_count) + .map(|index| { + format!( + "POSITION(LOWER(${index}) IN LOWER(elem->>'text')) > 0" + ) + }) + .collect::>() + .join(" OR "); + + format!( + r#" + EXISTS ( + SELECT 1 + FROM messages mq + CROSS JOIN LATERAL jsonb_array_elements(mq.content_json::jsonb) AS elem + WHERE mq.session_id = s.id + AND COALESCE((mq.metadata_json::jsonb->>'userVisible')::boolean, TRUE) = TRUE + AND elem->>'type' = 'text' + AND ({keyword_clauses}) + ) + "# + ) +} + +pub(crate) fn session_sort_at(session: &Session) -> DateTime { session.last_message_at.unwrap_or(session.updated_at) } @@ -775,7 +831,7 @@ impl Session { } } -fn deserialize_session_model_config( +pub(crate) fn deserialize_session_model_config( provider_name: Option<&str>, json: &str, ) -> Option { @@ -924,7 +980,44 @@ async fn insert_usage_ledger_row( Ok(()) } -impl SessionStorage { +pub(crate) async fn insert_usage_ledger_row_pg( + tx: &mut sqlx::Transaction<'_, Postgres>, + session_id: &str, + model: Option<&str>, + usage: &MessageUsage, +) -> Result<()> { + let cost_source = usage.cost_source.map(|cs| match cs { + CostSource::ProviderReported => "provider_reported", + CostSource::Estimated => "estimated", + }); + + sqlx::query( + r#" + INSERT INTO usage_ledger ( + session_id, created_timestamp, model, + input_tokens, output_tokens, total_tokens, + cache_read_tokens, cache_write_tokens, + cost, cost_source, is_compaction + ) + VALUES ($1, EXTRACT(EPOCH FROM NOW())::bigint, $2, $3, $4, $5, $6, $7, $8, $9, $10) + "#, + ) + .bind(session_id) + .bind(model) + .bind(usage.input_tokens) + .bind(usage.output_tokens) + .bind(usage.total_tokens) + .bind(usage.cache_read_tokens) + .bind(usage.cache_write_tokens) + .bind(usage.cost) + .bind(cost_source) + .bind(usage.is_compaction) + .execute(&mut **tx) + .await?; + Ok(()) +} + +impl SqliteSessionStorage { fn create_pool(path: &Path) -> Pool { if let Some(parent) = path.parent() { fs::create_dir_all(parent).expect("Failed to create session database directory"); @@ -947,7 +1040,6 @@ impl SessionStorage { pool: Self::create_pool(&db_path), initialized: tokio::sync::OnceCell::new(), session_dir, - action_required: Arc::new(crate::action_required_manager::ActionRequiredManager::new()), } } @@ -2756,9 +2848,338 @@ impl SessionStorage { } } -/// Merge a JSON object `patch` into an existing optional object value, -/// preserving keys not present in the patch. -fn merge_tool_meta( +impl SessionStorage { + pub fn from_sqlite(data_dir: PathBuf) -> Self { + Self { + backend: SessionStorageBackend::Sqlite(SqliteSessionStorage::new(data_dir)), + action_required: Arc::new(crate::action_required_manager::ActionRequiredManager::new()), + } + } + + pub fn from_postgres(db_url: String) -> Self { + Self { + backend: SessionStorageBackend::Postgres(PostgresSessionStorage::new(db_url)), + action_required: Arc::new(crate::action_required_manager::ActionRequiredManager::new()), + } + } + + pub fn new(data_dir: PathBuf) -> Self { + Self::from_sqlite(data_dir) + } + + pub(crate) async fn pool(&self) -> Result<&Pool> { + self.pool_sqlite().await + } + + pub(crate) async fn pool_postgres(&self) -> Result<&Pool> { + match &self.backend { + SessionStorageBackend::Postgres(storage) => storage.pool().await, + SessionStorageBackend::Sqlite(_) => { + anyhow::bail!("PostgreSQL session storage is not configured") + } + } + } + + pub(crate) fn uses_postgres(&self) -> bool { + matches!(self.backend, SessionStorageBackend::Postgres(_)) + } + + pub(crate) async fn initialize(&self) -> Result<()> { + match &self.backend { + SessionStorageBackend::Sqlite(storage) => { + storage.pool().await?; + } + SessionStorageBackend::Postgres(storage) => { + storage.pool().await?; + } + } + Ok(()) + } + + pub(crate) async fn pool_sqlite(&self) -> Result<&Pool> { + match &self.backend { + SessionStorageBackend::Sqlite(storage) => storage.pool().await, + SessionStorageBackend::Postgres(_) => { + anyhow::bail!("SQLite pool requested but PostgreSQL session storage is active") + } + } + } + + pub(crate) async fn create_session( + &self, + working_dir: PathBuf, + name: String, + session_type: SessionType, + goose_mode: GooseMode, + ) -> Result { + match &self.backend { + SessionStorageBackend::Sqlite(s) => { + s.create_session(working_dir, name, session_type, goose_mode) + .await + } + SessionStorageBackend::Postgres(s) => { + s.create_session(working_dir, name, session_type, goose_mode) + .await + } + } + } + + pub(crate) async fn get_session(&self, id: &str, include_messages: bool) -> Result { + match &self.backend { + SessionStorageBackend::Sqlite(s) => s.get_session(id, include_messages).await, + SessionStorageBackend::Postgres(s) => s.get_session(id, include_messages).await, + } + } + + pub(crate) async fn apply_update(&self, builder: SessionUpdateBuilder<'_>) -> Result<()> { + match &self.backend { + SessionStorageBackend::Sqlite(s) => s.apply_update(builder).await, + SessionStorageBackend::Postgres(s) => s.apply_update(builder).await, + } + } + + pub(crate) async fn update_project_for_session_types( + &self, + id: &str, + project_id: Option, + session_types: &[SessionType], + ) -> Result { + match &self.backend { + SessionStorageBackend::Sqlite(s) => { + s.update_project_for_session_types(id, project_id, session_types) + .await + } + SessionStorageBackend::Postgres(s) => { + s.update_project_for_session_types(id, project_id, session_types) + .await + } + } + } + + pub(crate) async fn add_message(&self, session_id: &str, message: &Message) -> Result<()> { + match &self.backend { + SessionStorageBackend::Sqlite(s) => s.add_message(session_id, message).await, + SessionStorageBackend::Postgres(s) => s.add_message(session_id, message).await, + } + } + + pub(crate) async fn replace_conversation( + &self, + session_id: &str, + conversation: &Conversation, + ) -> Result<()> { + match &self.backend { + SessionStorageBackend::Sqlite(s) => s.replace_conversation(session_id, conversation).await, + SessionStorageBackend::Postgres(s) => s.replace_conversation(session_id, conversation).await, + } + } + + async fn list_sessions_matching(&self, query: SessionListQuery<'_>) -> Result> { + match &self.backend { + SessionStorageBackend::Sqlite(s) => s.list_sessions_matching(query).await, + SessionStorageBackend::Postgres(s) => s.list_sessions_matching(query).await, + } + } + + async fn list_sessions_by_types(&self, types: Option<&[SessionType]>) -> Result> { + match &self.backend { + SessionStorageBackend::Sqlite(s) => s.list_sessions_by_types(types).await, + SessionStorageBackend::Postgres(s) => s.list_sessions_by_types(types).await, + } + } + + async fn list_sessions_paged( + &self, + query: SessionListPageQuery<'_>, + ) -> Result { + match &self.backend { + SessionStorageBackend::Sqlite(s) => s.list_sessions_paged(query).await, + SessionStorageBackend::Postgres(s) => s.list_sessions_paged(query).await, + } + } + + async fn list_sessions(&self) -> Result> { + match &self.backend { + SessionStorageBackend::Sqlite(s) => s.list_sessions().await, + SessionStorageBackend::Postgres(s) => s.list_sessions().await, + } + } + + async fn delete_session(&self, session_id: &str) -> Result<()> { + match &self.backend { + SessionStorageBackend::Sqlite(s) => s.delete_session(session_id).await, + SessionStorageBackend::Postgres(s) => s.delete_session(session_id).await, + } + } + + async fn get_insights(&self, types: &[SessionType]) -> Result { + match &self.backend { + SessionStorageBackend::Sqlite(s) => s.get_insights(types).await, + SessionStorageBackend::Postgres(s) => s.get_insights(types).await, + } + } + + async fn record_usage_metrics( + &self, + session_id: &str, + schedule_id: Option, + current_usage: Usage, + model: &str, + ledger: &MessageUsage, + ) -> Result<()> { + match &self.backend { + SessionStorageBackend::Sqlite(s) => { + s.record_usage_metrics(session_id, schedule_id, current_usage, model, ledger) + .await + } + SessionStorageBackend::Postgres(s) => { + s.record_usage_metrics(session_id, schedule_id, current_usage, model, ledger) + .await + } + } + } + + async fn get_session_usage_totals(&self, session_id: &str) -> Result { + match &self.backend { + SessionStorageBackend::Sqlite(s) => s.get_session_usage_totals(session_id).await, + SessionStorageBackend::Postgres(s) => s.get_session_usage_totals(session_id).await, + } + } + + async fn export_session(&self, id: &str) -> Result { + match &self.backend { + SessionStorageBackend::Sqlite(s) => s.export_session(id).await, + SessionStorageBackend::Postgres(s) => s.export_session(id).await, + } + } + + async fn import_session( + &self, + session_manager: &SessionManager, + json: &str, + session_type_override: Option, + ) -> Result { + match &self.backend { + SessionStorageBackend::Sqlite(s) => { + s.import_session(session_manager, json, session_type_override) + .await + } + SessionStorageBackend::Postgres(s) => { + s.import_session(session_manager, json, session_type_override) + .await + } + } + } + + async fn copy_session( + &self, + session_manager: &SessionManager, + session_id: &str, + new_name: String, + ) -> Result { + match &self.backend { + SessionStorageBackend::Sqlite(s) => { + s.copy_session(session_manager, session_id, new_name).await + } + SessionStorageBackend::Postgres(s) => { + s.copy_session(session_manager, session_id, new_name).await + } + } + } + + async fn truncate_conversation(&self, session_id: &str, timestamp: i64) -> Result<()> { + match &self.backend { + SessionStorageBackend::Sqlite(s) => s.truncate_conversation(session_id, timestamp).await, + SessionStorageBackend::Postgres(s) => { + s.truncate_conversation(session_id, timestamp).await + } + } + } + + async fn truncate_conversation_from_message( + &self, + session_id: &str, + message_id: &str, + ) -> Result<()> { + match &self.backend { + SessionStorageBackend::Sqlite(s) => { + s.truncate_conversation_from_message(session_id, message_id) + .await + } + SessionStorageBackend::Postgres(s) => { + s.truncate_conversation_from_message(session_id, message_id) + .await + } + } + } + + async fn search_chat_history( + &self, + query: &str, + limit: Option, + after_date: Option>, + before_date: Option>, + exclude_session_id: Option, + session_types: Vec, + ) -> Result { + match &self.backend { + SessionStorageBackend::Sqlite(s) => { + s.search_chat_history( + query, + limit, + after_date, + before_date, + exclude_session_id, + session_types, + ) + .await + } + SessionStorageBackend::Postgres(s) => { + s.search_chat_history( + query, + limit, + after_date, + before_date, + exclude_session_id, + session_types, + ) + .await + } + } + } + + async fn update_message_metadata(&self, id: &str, message_id: &str, f: F) -> Result<()> + where + F: FnOnce( + crate::conversation::message::MessageMetadata, + ) -> crate::conversation::message::MessageMetadata, + { + match &self.backend { + SessionStorageBackend::Sqlite(s) => s.update_message_metadata(id, message_id, f).await, + SessionStorageBackend::Postgres(s) => s.update_message_metadata(id, message_id, f).await, + } + } + + async fn update_tool_request_meta( + &self, + session_id: &str, + tool_call_id: &str, + patch: serde_json::Value, + ) -> Result<()> { + match &self.backend { + SessionStorageBackend::Sqlite(s) => { + s.update_tool_request_meta(session_id, tool_call_id, patch) + .await + } + SessionStorageBackend::Postgres(s) => { + s.update_tool_request_meta(session_id, tool_call_id, patch) + .await + } + } + } +} + +pub(crate) fn merge_tool_meta( existing: Option, patch: &serde_json::Value, ) -> serde_json::Value { @@ -3626,7 +4047,7 @@ mod tests { .await .unwrap(); - let pool = session_manager.storage().pool.clone(); + let pool = session_manager.storage().pool().await.unwrap().clone(); let results = run_lock_upgrade_race(pool.clone(), session.id.clone(), "BEGIN", true).await; assert!( @@ -4418,7 +4839,7 @@ mod tests { .await .unwrap(); - let pool = &sm.storage().pool; + let pool = sm.storage().pool().await.unwrap(); sqlx::query("UPDATE sessions SET goose_mode = 'garbage' WHERE id = ?") .bind(&session.id) .execute(pool) @@ -4447,7 +4868,7 @@ mod tests { .await .unwrap(); - SessionStorage::create_schema(&pool).await.unwrap(); + SqliteSessionStorage::create_schema(&pool).await.unwrap(); // Demote the schema back to v8 to simulate a database // that has never seen migration 9. @@ -4516,7 +4937,7 @@ mod tests { .await .unwrap(); - SessionStorage::create_schema(&pool).await.unwrap(); + SqliteSessionStorage::create_schema(&pool).await.unwrap(); // Recreate a v13-shaped database without cache token columns. for column in [ diff --git a/crates/goose/tests/acp_transport_auth_test.rs b/crates/goose/tests/acp_transport_auth_test.rs index 71020f05f..946c7cf25 100644 --- a/crates/goose/tests/acp_transport_auth_test.rs +++ b/crates/goose/tests/acp_transport_auth_test.rs @@ -7,11 +7,19 @@ use axum::Router; use goose::acp::server::AcpBuiltinSelection; use goose::acp::server_factory::{AcpServer, AcpServerFactoryConfig}; use goose::acp::transport::{create_acp_router, create_router}; +use goose::acp::tkmind::state::AppState; use goose::agents::GoosePlatform; use tower::ServiceExt; const SECRET: &str = "test-secret-token"; +fn test_tkmind_state() -> Arc { + tokio::runtime::Runtime::new() + .unwrap() + .block_on(AppState::new()) + .unwrap() +} + fn test_router(require_token: bool, dir: &tempfile::TempDir) -> Router { test_router_with_origins(require_token, dir, Vec::new()) } @@ -39,7 +47,13 @@ fn test_authenticated_acp_router(dir: &tempfile::TempDir) -> Router { session_cwd: None, enable_scheduler: false, })); - create_router(server, SECRET.to_string(), true, Vec::new()) + create_router( + server, + SECRET.to_string(), + true, + Vec::new(), + test_tkmind_state(), + ) } fn test_router_with_origins( @@ -61,6 +75,7 @@ fn test_router_with_origins( SECRET.to_string(), require_token, additional_allowed_origins, + test_tkmind_state(), ) } diff --git a/deploy/coding_router.sh b/deploy/coding_router.sh new file mode 100755 index 000000000..5e8ca42fd --- /dev/null +++ b/deploy/coding_router.sh @@ -0,0 +1,72 @@ +#!/usr/bin/env bash +set -euo pipefail + +ENGINE="${1:-}" +PROJECT_DIR="${2:-}" +TASK="${3:-}" +DRY_RUN="${4:-}" + +AIDER_HOME="${AIDER_HOME:-/Users/john/Project/aider}" +AIDER_BIN="${AIDER_BIN:-${GOOSE_AIDER_BIN:-${AIDER_HOME}/aider.sh}}" + +if [[ -f "${AIDER_HOME}/.env" ]]; then + set -a + # shellcheck disable=SC1091 + source "${AIDER_HOME}/.env" + set +a +fi + +if [[ -z "$ENGINE" || -z "$PROJECT_DIR" || -z "$TASK" ]]; then + echo "Usage: coding_router.sh [--dry-run]" + exit 1 +fi + +if [[ ! -d "$PROJECT_DIR" ]]; then + echo "Project dir not found: $PROJECT_DIR" + exit 1 +fi + +choose_engine() { + case "$ENGINE" in + aider) echo "aider" ;; + goose) echo "goose" ;; + auto) + if echo "$TASK" | grep -qE '修复|bug|Bug|BUG|重构|跨文件|实现|新增|改造|兼容|报错|测试失败|commit|提交|代码修改|开发|feature|refactor|fix'; then + echo "aider" + else + echo "goose" + fi + ;; + *) + echo "Unknown engine: $ENGINE" + exit 1 + ;; + esac +} + +SELECTED_ENGINE="$(choose_engine)" +echo "Selected engine: $SELECTED_ENGINE" + +if [[ "$DRY_RUN" == "--dry-run" ]]; then + exit 0 +fi + +cd "$PROJECT_DIR" +echo "===== git status before =====" +git status --short || true + +case "$SELECTED_ENGINE" in + aider) + if [[ ! -x "$AIDER_BIN" ]]; then + echo "Aider not executable: $AIDER_BIN" + exit 1 + fi + "$AIDER_BIN" --yes-always --message "$TASK" + echo "===== git status after =====" + git status --short || true + ;; + goose) + echo "Goose developer tools handle this task in-session: $TASK" + exit 2 + ;; +esac