| 1 | //! [`AcpServer`]: the ACP v1 agent behind every Onde product. Owns sessions (in memory, backed |
| 2 | //! by [`SessionStore`]), the turn loop, auth, and the model and auto-approve config options. |
| 3 | //! Everything product-specific comes from the [`Profile`]. |
| 4 | |
| 5 | use std::collections::{HashMap, HashSet}; |
| 6 | use std::path::{Path, PathBuf}; |
| 7 | use std::sync::atomic::{AtomicBool, Ordering}; |
| 8 | use std::sync::{Arc, Mutex}; |
| 9 | use std::time::{Duration, SystemTime, UNIX_EPOCH}; |
| 10 | |
| 11 | use agent_client_protocol::schema::ProtocolVersion; |
| 12 | use agent_client_protocol::schema::v1::{ |
| 13 | AuthMethod, AuthMethodTerminal, AuthenticateRequest, AuthenticateResponse, |
| 14 | AvailableCommandsUpdate, ClientCapabilities, CloseSessionRequest, CloseSessionResponse, |
| 15 | ConfigOptionUpdate, ContentBlock, ContentChunk, DeleteSessionRequest, DeleteSessionResponse, |
| 16 | ListSessionsRequest, ListSessionsResponse, LoadSessionRequest, LoadSessionResponse, |
| 17 | LogoutResponse, McpServer, MessageId, NewSessionRequest, NewSessionResponse, PromptRequest, |
| 18 | PromptResponse, RequestId, ResumeSessionRequest, ResumeSessionResponse, SessionConfigOption, |
| 19 | SessionConfigOptionCategory, SessionConfigSelectOption, SessionConfigSelectOptions, SessionId, |
| 20 | SessionInfo, SessionInfoUpdate, SessionNotification, SessionUpdate, |
| 21 | SetSessionConfigOptionRequest, SetSessionConfigOptionResponse, StopReason, ToolCall, |
| 22 | ToolCallContent, ToolCallLocation, ToolCallStatus, ToolCallUpdateFields, UsageUpdate, |
| 23 | }; |
| 24 | use agent_client_protocol::{Client, ConnectionTo, Responder}; |
| 25 | use serde_json::{Value, json}; |
| 26 | use tokio_util::sync::CancellationToken; |
| 27 | |
| 28 | use ed_mcp::{ClientInfo, McpToolset, ToolsetConfig}; |
| 29 | |
| 30 | use crate::auth::{AuthCtx, AuthProvider}; |
| 31 | use crate::content as prompt; |
| 32 | use crate::llm::{Delta, LlmClient, LlmConfig, LlmEnv}; |
| 33 | use crate::store::{SessionStore, StoredSession}; |
| 34 | use crate::tools::{self, ToolCtx, ToolOutcome, Toolset, absolutize, normalize}; |
| 35 | use crate::{CommandCtx, Profile, PromptCtx, SessionCtx}; |
| 36 | |
| 37 | const MAX_TURNS: usize = 50; |
| 38 | /// Id of the session config option that selects the model. |
| 39 | pub const MODEL_CONFIG_ID: &str = "model"; |
| 40 | /// Id of the boolean session config option that skips permission prompts. |
| 41 | pub const AUTO_APPROVE_CONFIG_ID: &str = "auto_approve"; |
| 42 | /// Context window reported in `usage_update` when `<PREFIX>_CONTEXT_WINDOW` is unset. |
| 43 | const DEFAULT_CONTEXT_WINDOW: u64 = 128_000; |
| 44 | /// Id of the terminal auth method advertised on initialize. |
| 45 | pub const AUTH_METHOD_ID: &str = "terminal-setup"; |
| 46 | /// Cap on the sessions returned by `session/list`. |
| 47 | const MAX_LISTED: usize = 200; |
| 48 | /// How long `run_command` may run when `<PREFIX>_COMMAND_TIMEOUT_SECS` is unset. |
| 49 | const DEFAULT_COMMAND_TIMEOUT: Duration = Duration::from_secs(120); |
| 50 | |
| 51 | /// The terminal auth method: clients run `<agent> --setup`. |
| 52 | pub fn auth_methods() -> Vec<AuthMethod> { |
| 53 | vec![AuthMethod::Terminal( |
| 54 | AuthMethodTerminal::new(AUTH_METHOD_ID, "Run in terminal") |
| 55 | .description("Interactive setup: choose a provider and store an API key") |
| 56 | .args(vec!["--setup".into()]), |
| 57 | )] |
| 58 | } |
| 59 | |
| 60 | /// Whether the client can run terminal auth. ACP v1 has the `auth.terminal` capability; clients |
| 61 | /// that predate it (including the ACP registry's validator, which probes every listed agent) |
| 62 | /// say so with `_meta["terminal-auth"]: true` instead. Offering no method to such a client |
| 63 | /// fails registry validation, so either signal counts. |
| 64 | pub fn supports_terminal_auth(caps: &ClientCapabilities) -> bool { |
| 65 | caps.auth.terminal |
| 66 | || caps |
| 67 | .meta |
| 68 | .as_ref() |
| 69 | .and_then(|m| m.get("terminal-auth")) |
| 70 | .and_then(Value::as_bool) |
| 71 | .unwrap_or(false) |
| 72 | } |
| 73 | |
| 74 | pub fn auth_required_error(agent: &str) -> agent_client_protocol::Error { |
| 75 | agent_client_protocol::Error::auth_required().data(format!( |
| 76 | "No API key configured. Authenticate with the terminal method (`{agent} --setup`)." |
| 77 | )) |
| 78 | } |
| 79 | |
| 80 | /// The only protocol version this agent speaks is v1. A client proposing a newer version gets |
| 81 | /// v1 back (it can then decide whether to continue); an older (v0) client is told v1 too, as |
| 82 | /// the spec requires the latest supported version when the requested one isn't supported. |
| 83 | pub fn negotiate_version(client: ProtocolVersion) -> ProtocolVersion { |
| 84 | let _ = client; |
| 85 | ProtocolVersion::V1 |
| 86 | } |
| 87 | |
| 88 | struct Session { |
| 89 | cwd: PathBuf, |
| 90 | /// Additional workspace roots (from the request's `additional_directories`). |
| 91 | roots: Vec<PathBuf>, |
| 92 | model: String, |
| 93 | /// Conversation history without the system prompt (rebuilt every turn). |
| 94 | messages: Vec<Value>, |
| 95 | cancel: CancellationToken, |
| 96 | always_allowed: Arc<Mutex<HashSet<String>>>, |
| 97 | always_rejected: Arc<Mutex<HashSet<String>>>, |
| 98 | /// Approve tool calls without asking; the `auto_approve` config option. Not persisted: |
| 99 | /// a reopened session starts from the process default (`--yolo`). |
| 100 | auto_approve: Arc<AtomicBool>, |
| 101 | title: Option<String>, |
| 102 | /// Last activity, seconds since the Unix epoch. |
| 103 | updated_at: u64, |
| 104 | mcp: Arc<McpToolset>, |
| 105 | /// A turn is running; its history is checked out of `messages`. |
| 106 | busy: bool, |
| 107 | } |
| 108 | |
| 109 | impl Session { |
| 110 | fn fresh( |
| 111 | cwd: PathBuf, |
| 112 | roots: Vec<PathBuf>, |
| 113 | model: String, |
| 114 | yolo: bool, |
| 115 | mcp: Arc<McpToolset>, |
| 116 | ) -> Self { |
| 117 | Self::restored( |
| 118 | StoredSession { |
| 119 | cwd, |
| 120 | roots, |
| 121 | model, |
| 122 | title: None, |
| 123 | updated_at: now_secs(), |
| 124 | messages: Vec::new(), |
| 125 | }, |
| 126 | yolo, |
| 127 | mcp, |
| 128 | ) |
| 129 | } |
| 130 | |
| 131 | /// A session brought back from the store. Permission grants start fresh; MCP servers are |
| 132 | /// connected by the request that reopens it. |
| 133 | fn restored(stored: StoredSession, yolo: bool, mcp: Arc<McpToolset>) -> Self { |
| 134 | Self { |
| 135 | cwd: stored.cwd, |
| 136 | roots: stored.roots, |
| 137 | model: stored.model, |
| 138 | messages: stored.messages, |
| 139 | cancel: CancellationToken::new(), |
| 140 | always_allowed: Arc::default(), |
| 141 | always_rejected: Arc::default(), |
| 142 | auto_approve: Arc::new(AtomicBool::new(yolo)), |
| 143 | title: stored.title, |
| 144 | updated_at: stored.updated_at, |
| 145 | mcp, |
| 146 | busy: false, |
| 147 | } |
| 148 | } |
| 149 | |
| 150 | fn stored(&self) -> StoredSession { |
| 151 | StoredSession { |
| 152 | cwd: self.cwd.clone(), |
| 153 | roots: self.roots.clone(), |
| 154 | model: self.model.clone(), |
| 155 | title: self.title.clone(), |
| 156 | updated_at: self.updated_at, |
| 157 | messages: self.messages.clone(), |
| 158 | } |
| 159 | } |
| 160 | } |
| 161 | |
| 162 | fn new_message_id() -> MessageId { |
| 163 | MessageId::new(uuid::Uuid::new_v4().to_string()) |
| 164 | } |
| 165 | |
| 166 | /// ACP v1 requires absolute `cwd` and `additionalDirectories`. |
| 167 | fn check_absolute(cwd: &Path, roots: &[PathBuf]) -> agent_client_protocol::Result<()> { |
| 168 | if !cwd.is_absolute() { |
| 169 | return Err(invalid("cwd must be an absolute path")); |
| 170 | } |
| 171 | if roots.iter().any(|r| !r.is_absolute()) { |
| 172 | return Err(invalid( |
| 173 | "additionalDirectories entries must be absolute paths", |
| 174 | )); |
| 175 | } |
| 176 | Ok(()) |
| 177 | } |
| 178 | |
| 179 | pub fn now_secs() -> u64 { |
| 180 | SystemTime::now() |
| 181 | .duration_since(UNIX_EPOCH) |
| 182 | .map(|d| d.as_secs()) |
| 183 | .unwrap_or(0) |
| 184 | } |
| 185 | |
| 186 | /// ISO 8601 UTC timestamp without sub-second precision (no chrono dependency). |
| 187 | pub fn iso8601(secs: u64) -> String { |
| 188 | let days = secs / 86_400; |
| 189 | let secs_of_day = secs % 86_400; |
| 190 | let (h, m, s) = ( |
| 191 | secs_of_day / 3600, |
| 192 | secs_of_day % 3600 / 60, |
| 193 | secs_of_day % 60, |
| 194 | ); |
| 195 | // Civil-from-days algorithm (Howard Hinnant). |
| 196 | let z = days as i64 + 719_468; |
| 197 | let era = z.div_euclid(146_097); |
| 198 | let doe = z.rem_euclid(146_097); |
| 199 | let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365; |
| 200 | let y = yoe + era * 400; |
| 201 | let doy = doe - (365 * yoe + yoe / 4 - yoe / 100); |
| 202 | let mp = (5 * doy + 2) / 153; |
| 203 | let d = doy - (153 * mp + 2) / 5 + 1; |
| 204 | let mo = if mp < 10 { mp + 3 } else { mp - 9 }; |
| 205 | let y = if mo <= 2 { y + 1 } else { y }; |
| 206 | format!("{y:04}-{mo:02}-{d:02}T{h:02}:{m:02}:{s:02}Z") |
| 207 | } |
| 208 | |
| 209 | /// Derive a short session title from the first user prompt. |
| 210 | pub fn session_title(prompt: &str) -> String { |
| 211 | const MAX: usize = 80; |
| 212 | let first = prompt.lines().next().unwrap_or("").trim(); |
| 213 | if first.chars().count() <= MAX { |
| 214 | first.to_string() |
| 215 | } else { |
| 216 | let cut: String = first.chars().take(MAX - 1).collect(); |
| 217 | format!("{}…", cut.trim_end()) |
| 218 | } |
| 219 | } |
| 220 | |
| 221 | fn invalid(msg: impl Into<String>) -> agent_client_protocol::Error { |
| 222 | agent_client_protocol::Error::invalid_params().data(msg.into()) |
| 223 | } |
| 224 | |
| 225 | fn unknown_session(id: &SessionId) -> agent_client_protocol::Error { |
| 226 | invalid(format!("unknown session {id}")) |
| 227 | } |
| 228 | |
| 229 | /// A tool message; `failed` is kept in a private field that never reaches the model. |
| 230 | fn tool_message(id: &str, content: &str, failed: bool) -> Value { |
| 231 | let mut m = json!({ "role": "tool", "tool_call_id": id, "content": content }); |
| 232 | if failed { |
| 233 | m["x_failed"] = json!(true); |
| 234 | } |
| 235 | m |
| 236 | } |
| 237 | |
| 238 | fn wire_message(m: &Value) -> Value { |
| 239 | match m.as_object() { |
| 240 | Some(o) if o.contains_key("x_failed") => { |
| 241 | let mut o = o.clone(); |
| 242 | o.remove("x_failed"); |
| 243 | Value::Object(o) |
| 244 | } |
| 245 | _ => m.clone(), |
| 246 | } |
| 247 | } |
| 248 | |
| 249 | fn locations_of(paths: &[PathBuf]) -> Vec<ToolCallLocation> { |
| 250 | paths |
| 251 | .iter() |
| 252 | .map(|p| ToolCallLocation::new(absolutize(p))) |
| 253 | .collect() |
| 254 | } |
| 255 | |
| 256 | /// Text of a stored user message (`content` is a string, or an array of parts). |
| 257 | fn user_text(msg: &Value) -> String { |
| 258 | match msg.get("content") { |
| 259 | Some(Value::String(s)) => s.clone(), |
| 260 | Some(Value::Array(parts)) => parts |
| 261 | .iter() |
| 262 | .filter_map(|p| match p.get("type").and_then(Value::as_str) { |
| 263 | Some("text") => p.get("text").and_then(Value::as_str).map(String::from), |
| 264 | Some("image_url") => Some("[image]".into()), |
| 265 | _ => None, |
| 266 | }) |
| 267 | .collect::<Vec<_>>() |
| 268 | .join("\n\n"), |
| 269 | _ => String::new(), |
| 270 | } |
| 271 | } |
| 272 | |
| 273 | #[derive(Clone)] |
| 274 | pub struct AcpServer { |
| 275 | profile: Arc<dyn Profile>, |
| 276 | /// The profile's own sign-in, when it has one (see [`Profile::auth`]). |
| 277 | auth: Option<Arc<dyn AuthProvider>>, |
| 278 | /// The profile's toolsets, built once. |
| 279 | toolsets: Arc<Vec<Arc<dyn Toolset>>>, |
| 280 | env: LlmEnv, |
| 281 | /// Replaced when credentials change: `authenticate`, `logout`, or a key stored by |
| 282 | /// `--setup` in another process. |
| 283 | llm: Arc<Mutex<LlmClient>>, |
| 284 | /// Process default for each session's `auto_approve` (`--yolo` / `<PREFIX>_YOLO`). |
| 285 | yolo: bool, |
| 286 | /// How the agent was launched: `tui` (interactive terminal UI) or `acp` (editor). |
| 287 | surface: &'static str, |
| 288 | client_caps: Arc<Mutex<ClientCapabilities>>, |
| 289 | sessions: Arc<Mutex<HashMap<SessionId, Session>>>, |
| 290 | /// Models offered for selection, fetched from `GET /models` on first use and dropped |
| 291 | /// when credentials change. |
| 292 | models: Arc<tokio::sync::Mutex<Option<Vec<String>>>>, |
| 293 | store: SessionStore, |
| 294 | data_dir: PathBuf, |
| 295 | } |
| 296 | |
| 297 | impl AcpServer { |
| 298 | pub fn new(profile: Arc<dyn Profile>, yolo: bool, surface: &'static str) -> Self { |
| 299 | let env = profile.info().env; |
| 300 | let data_dir = env.data_dir(); |
| 301 | let auth = profile.auth(); |
| 302 | Self { |
| 303 | toolsets: Arc::new(profile.toolsets()), |
| 304 | llm: Arc::new(Mutex::new(LlmClient::new(Self::config_for( |
| 305 | auth.as_deref(), |
| 306 | &env, |
| 307 | )))), |
| 308 | profile, |
| 309 | auth, |
| 310 | env, |
| 311 | yolo, |
| 312 | surface, |
| 313 | client_caps: Arc::default(), |
| 314 | sessions: Arc::default(), |
| 315 | models: Arc::default(), |
| 316 | store: SessionStore::new(&data_dir), |
| 317 | data_dir, |
| 318 | } |
| 319 | } |
| 320 | |
| 321 | /// The provider's stored session when the profile has one, else the environment. A |
| 322 | /// signed-out provider yields the environment's endpoint with no key, so a stray |
| 323 | /// `ONDE_API_KEY` can't stand in for the product's own sign-in. |
| 324 | fn config_for(auth: Option<&dyn AuthProvider>, env: &LlmEnv) -> LlmConfig { |
| 325 | match auth { |
| 326 | None => LlmConfig::from_env(env), |
| 327 | Some(auth) => auth.llm_config().unwrap_or_else(|| LlmConfig { |
| 328 | api_key: None, |
| 329 | ..LlmConfig::from_env(env) |
| 330 | }), |
| 331 | } |
| 332 | } |
| 333 | |
| 334 | /// The auth methods for `initialize`. A provider's method is offered whatever the client |
| 335 | /// can do; terminal auth only to a client that can run it. |
| 336 | pub fn auth_methods(&self, caps: &ClientCapabilities) -> Vec<AuthMethod> { |
| 337 | match &self.auth { |
| 338 | Some(auth) => vec![auth.method()], |
| 339 | None if supports_terminal_auth(caps) => auth_methods(), |
| 340 | None => Vec::new(), |
| 341 | } |
| 342 | } |
| 343 | |
| 344 | pub fn set_client_caps(&self, caps: ClientCapabilities) { |
| 345 | *self.client_caps.lock().unwrap() = caps; |
| 346 | } |
| 347 | |
| 348 | fn llm(&self) -> LlmClient { |
| 349 | self.llm.lock().unwrap().clone() |
| 350 | } |
| 351 | |
| 352 | pub fn has_api_key(&self) -> bool { |
| 353 | self.llm().has_api_key() |
| 354 | } |
| 355 | |
| 356 | pub fn model(&self) -> String { |
| 357 | self.llm().model().to_string() |
| 358 | } |
| 359 | |
| 360 | /// Re-read credentials (environment + stored `env` file) and report whether a key is |
| 361 | /// configured. Editors run terminal auth (`<agent> --setup`) in a separate process |
| 362 | /// and then retry `session/new` or `session/load` on this one without calling |
| 363 | /// `authenticate`, so the stored key has to be picked up here. The model list is only |
| 364 | /// dropped when the configuration actually changed. |
| 365 | pub async fn refresh_credentials(&self) -> bool { |
| 366 | let fresh = Self::config_for(self.auth.as_deref(), &self.env); |
| 367 | let changed = fresh != *self.llm().config(); |
| 368 | if changed { |
| 369 | tracing::info!( |
| 370 | "credentials changed; reloading provider {}", |
| 371 | fresh.provider.name() |
| 372 | ); |
| 373 | *self.llm.lock().unwrap() = LlmClient::new(fresh); |
| 374 | *self.models.lock().await = None; |
| 375 | } |
| 376 | self.has_api_key() |
| 377 | } |
| 378 | |
| 379 | /// `authenticate`: the client has run `<agent> --setup`; pick up the stored key. |
| 380 | pub async fn authenticate( |
| 381 | &self, |
| 382 | req: AuthenticateRequest, |
| 383 | connection: &ConnectionTo<Client>, |
| 384 | request_id: RequestId, |
| 385 | ) -> agent_client_protocol::Result<AuthenticateResponse> { |
| 386 | if let Some(auth) = &self.auth { |
| 387 | if req.method_id != *auth.method().id() { |
| 388 | return Err(invalid(format!("unknown auth method {}", req.method_id.0))); |
| 389 | } |
| 390 | let caps = self.client_caps.lock().unwrap().clone(); |
| 391 | auth.authenticate(AuthCtx { |
| 392 | connection, |
| 393 | request_id, |
| 394 | client_caps: &caps, |
| 395 | }) |
| 396 | .await?; |
| 397 | // A new session means a new bearer, and possibly a different model list. |
| 398 | *self.models.lock().await = None; |
| 399 | *self.llm.lock().unwrap() = |
| 400 | LlmClient::new(Self::config_for(self.auth.as_deref(), &self.env)); |
| 401 | return if self.has_api_key() { |
| 402 | Ok(AuthenticateResponse::new()) |
| 403 | } else { |
| 404 | Err(self.auth_required()) |
| 405 | }; |
| 406 | } |
| 407 | if &*req.method_id.0 != AUTH_METHOD_ID { |
| 408 | return Err(invalid(format!("unknown auth method {}", req.method_id.0))); |
| 409 | } |
| 410 | if self.refresh_credentials().await { |
| 411 | Ok(AuthenticateResponse::new()) |
| 412 | } else { |
| 413 | Err(self.auth_required()) |
| 414 | } |
| 415 | } |
| 416 | |
| 417 | /// `logout`: remove the stored key and forget it. Every file the key may be loaded from |
| 418 | /// goes, or the agent would still be signed in afterwards; one failed removal doesn't stop |
| 419 | /// the others. A key exported in the agent's own environment can't be removed from here |
| 420 | /// and keeps the agent signed in. |
| 421 | pub async fn logout(&self) -> agent_client_protocol::Result<LogoutResponse> { |
| 422 | if let Some(auth) = &self.auth { |
| 423 | return match auth.logout().await { |
| 424 | Ok(()) => { |
| 425 | *self.models.lock().await = None; |
| 426 | *self.llm.lock().unwrap() = |
| 427 | LlmClient::new(Self::config_for(self.auth.as_deref(), &self.env)); |
| 428 | Ok(LogoutResponse::new()) |
| 429 | } |
| 430 | Err(e) => { |
| 431 | Err(agent_client_protocol::Error::internal_error().data(format!("{e:#}"))) |
| 432 | } |
| 433 | }; |
| 434 | } |
| 435 | let failures: Vec<String> = self |
| 436 | .env |
| 437 | .config_file_candidates() |
| 438 | .into_iter() |
| 439 | .filter_map(|path| match std::fs::remove_file(&path) { |
| 440 | Err(e) if e.kind() != std::io::ErrorKind::NotFound => { |
| 441 | Some(format!("removing {}: {e}", path.display())) |
| 442 | } |
| 443 | _ => None, |
| 444 | }) |
| 445 | .collect(); |
| 446 | self.refresh_credentials().await; |
| 447 | if !failures.is_empty() { |
| 448 | return Err(agent_client_protocol::Error::internal_error().data(failures.join("; "))); |
| 449 | } |
| 450 | Ok(LogoutResponse::new()) |
| 451 | } |
| 452 | |
| 453 | fn scratch_dir(&self, id: &SessionId) -> Option<PathBuf> { |
| 454 | let id = crate::store::sanitize_id(&id.0).ok()?; |
| 455 | Some(self.data_dir.join("scratch").join(id)) |
| 456 | } |
| 457 | |
| 458 | fn remove_scratch(&self, id: &SessionId) { |
| 459 | if let Some(dir) = self.scratch_dir(id) { |
| 460 | let _ = std::fs::remove_dir_all(dir); |
| 461 | } |
| 462 | } |
| 463 | |
| 464 | /// `AUTH_REQUIRED`, naming this agent's setup command. |
| 465 | pub fn auth_required(&self) -> agent_client_protocol::Error { |
| 466 | if let Some(auth) = &self.auth { |
| 467 | return agent_client_protocol::Error::auth_required().data(format!( |
| 468 | "Not signed in. Authenticate with the {} method.", |
| 469 | auth.method().name() |
| 470 | )); |
| 471 | } |
| 472 | auth_required_error(self.profile.info().name) |
| 473 | } |
| 474 | |
| 475 | /// How a session's MCP servers are connected. |
| 476 | fn mcp_config(&self) -> ToolsetConfig { |
| 477 | let info = self.profile.info(); |
| 478 | ToolsetConfig::new(ClientInfo::new(info.name, info.version)) |
| 479 | } |
| 480 | |
| 481 | async fn connect_mcp( |
| 482 | &self, |
| 483 | servers: &[agent_client_protocol::schema::v1::McpServer], |
| 484 | ) -> Arc<McpToolset> { |
| 485 | Arc::new( |
| 486 | McpToolset::connect( |
| 487 | self.mcp_config(), |
| 488 | servers, |
| 489 | &self.profile.builtin_mcp_servers(), |
| 490 | ) |
| 491 | .await, |
| 492 | ) |
| 493 | } |
| 494 | |
| 495 | fn system_prompt(&self, cwd: &Path, roots: &[PathBuf], mcp: &McpToolset) -> String { |
| 496 | self.profile.system_prompt(&PromptCtx { |
| 497 | surface: self.surface, |
| 498 | cwd, |
| 499 | roots, |
| 500 | mcp, |
| 501 | }) |
| 502 | } |
| 503 | |
| 504 | /// The selectable models. Falls back to just the configured model if the endpoint |
| 505 | /// can't list them; the configured model is always offered. |
| 506 | async fn models(&self) -> Vec<String> { |
| 507 | let mut cache = self.models.lock().await; |
| 508 | if let Some(ids) = &*cache { |
| 509 | return ids.clone(); |
| 510 | } |
| 511 | let llm = self.llm(); |
| 512 | let default = llm.model().to_string(); |
| 513 | // An explicit list wins; editors show this list as the model picker. |
| 514 | let configured: Vec<String> = std::env::var(self.env.var("MODELS")) |
| 515 | .unwrap_or_default() |
| 516 | .split(',') |
| 517 | .map(|m| m.trim().to_string()) |
| 518 | .filter(|m| !m.is_empty()) |
| 519 | .collect(); |
| 520 | let mut ids = if !configured.is_empty() { |
| 521 | configured |
| 522 | } else { |
| 523 | match llm.models().await { |
| 524 | Ok(models) => models.into_iter().map(|m| m.id).collect(), |
| 525 | Err(e) => { |
| 526 | tracing::warn!("listing models failed, offering only {default}: {e:#}"); |
| 527 | Vec::new() |
| 528 | } |
| 529 | } |
| 530 | }; |
| 531 | if !ids.contains(&default) { |
| 532 | ids.insert(0, default); |
| 533 | } |
| 534 | *cache = Some(ids.clone()); |
| 535 | ids |
| 536 | } |
| 537 | |
| 538 | /// Whether the client advertised `session.configOptions.boolean`, which ACP requires |
| 539 | /// before an agent may offer `type: "boolean"` options. |
| 540 | fn boolean_options_supported(&self) -> bool { |
| 541 | self.client_caps |
| 542 | .lock() |
| 543 | .unwrap() |
| 544 | .session |
| 545 | .as_ref() |
| 546 | .and_then(|s| s.config_options.as_ref()) |
| 547 | .is_some_and(|c| c.boolean.is_some()) |
| 548 | } |
| 549 | |
| 550 | async fn config_options(&self, model: &str, auto_approve: bool) -> Vec<SessionConfigOption> { |
| 551 | let options = self |
| 552 | .models() |
| 553 | .await |
| 554 | .into_iter() |
| 555 | .map(|id| SessionConfigSelectOption::new(id.clone(), id)) |
| 556 | .collect(); |
| 557 | let mut config = vec![ |
| 558 | SessionConfigOption::select( |
| 559 | MODEL_CONFIG_ID, |
| 560 | "Model", |
| 561 | model.to_string(), |
| 562 | SessionConfigSelectOptions::Ungrouped(options), |
| 563 | ) |
| 564 | .category(SessionConfigOptionCategory::Model), |
| 565 | ]; |
| 566 | if self.boolean_options_supported() { |
| 567 | config.push( |
| 568 | SessionConfigOption::boolean( |
| 569 | AUTO_APPROVE_CONFIG_ID, |
| 570 | "Auto-approve actions", |
| 571 | auto_approve, |
| 572 | ) |
| 573 | .description("Run tools that change things without asking for permission"), |
| 574 | ); |
| 575 | } |
| 576 | config |
| 577 | } |
| 578 | |
| 579 | /// The model and auto-approve setting of a live session. |
| 580 | fn session_settings(&self, id: &SessionId) -> Option<(String, bool)> { |
| 581 | self.sessions |
| 582 | .lock() |
| 583 | .unwrap() |
| 584 | .get(id) |
| 585 | .map(|s| (s.model.clone(), s.auto_approve.load(Ordering::Relaxed))) |
| 586 | } |
| 587 | |
| 588 | async fn session_config_options( |
| 589 | &self, |
| 590 | id: &SessionId, |
| 591 | ) -> agent_client_protocol::Result<Vec<SessionConfigOption>> { |
| 592 | let (model, auto_approve) = self |
| 593 | .session_settings(id) |
| 594 | .ok_or_else(|| unknown_session(id))?; |
| 595 | Ok(self.config_options(&model, auto_approve).await) |
| 596 | } |
| 597 | |
| 598 | pub async fn set_config_option( |
| 599 | &self, |
| 600 | req: SetSessionConfigOptionRequest, |
| 601 | connection: &ConnectionTo<Client>, |
| 602 | ) -> agent_client_protocol::Result<SetSessionConfigOptionResponse> { |
| 603 | let id = req.session_id.clone(); |
| 604 | self.restore(&id).await?; |
| 605 | match &*req.config_id.0 { |
| 606 | MODEL_CONFIG_ID => { |
| 607 | let Some(model) = req.value.as_value_id().map(|v| v.0.to_string()) else { |
| 608 | return Err(invalid("model must be a value id")); |
| 609 | }; |
| 610 | if !self.models().await.contains(&model) { |
| 611 | return Err(invalid(format!("unknown model {model}"))); |
| 612 | } |
| 613 | match self.sessions.lock().unwrap().get_mut(&id) { |
| 614 | Some(s) => s.model = model.clone(), |
| 615 | None => return Err(unknown_session(&id)), |
| 616 | } |
| 617 | self.persist(&id).await; |
| 618 | tracing::info!("session {id} now uses model {model}"); |
| 619 | } |
| 620 | AUTO_APPROVE_CONFIG_ID if self.boolean_options_supported() => { |
| 621 | let Some(on) = req.value.as_bool() else { |
| 622 | return Err(invalid("auto_approve must be a boolean")); |
| 623 | }; |
| 624 | match self.sessions.lock().unwrap().get(&id) { |
| 625 | Some(s) => s.auto_approve.store(on, Ordering::Relaxed), |
| 626 | None => return Err(unknown_session(&id)), |
| 627 | } |
| 628 | tracing::info!("session {id} auto-approve {on}"); |
| 629 | } |
| 630 | other => return Err(invalid(format!("unknown config option {other}"))), |
| 631 | } |
| 632 | let options = self.session_config_options(&id).await?; |
| 633 | // Tell every view of the session about the change, not just the requester. |
| 634 | connection.send_notification(SessionNotification::new( |
| 635 | id, |
| 636 | SessionUpdate::ConfigOptionUpdate(ConfigOptionUpdate::new(options.clone())), |
| 637 | ))?; |
| 638 | Ok(SetSessionConfigOptionResponse::new(options)) |
| 639 | } |
| 640 | |
| 641 | /// Save a session that is not running a turn (a turn checks the history out; the |
| 642 | /// end-of-turn save covers it). |
| 643 | async fn persist(&self, id: &SessionId) { |
| 644 | let stored = self |
| 645 | .sessions |
| 646 | .lock() |
| 647 | .unwrap() |
| 648 | .get(id) |
| 649 | .filter(|s| !s.busy) |
| 650 | .map(Session::stored); |
| 651 | let Some(stored) = stored else { return }; |
| 652 | let store = self.store.clone(); |
| 653 | let sid = id.0.to_string(); |
| 654 | match tokio::task::spawn_blocking(move || store.save(&sid, &stored)).await { |
| 655 | Ok(Ok(())) => {} |
| 656 | Ok(Err(e)) => tracing::warn!("saving session {id} failed: {e:#}"), |
| 657 | Err(e) => tracing::warn!("saving session {id} failed: {e}"), |
| 658 | } |
| 659 | } |
| 660 | |
| 661 | /// Bring a stored session into memory without rebinding it (a prompt or config change |
| 662 | /// for a thread the client reopened without `session/load`). Errors if it is unknown. |
| 663 | async fn restore(&self, id: &SessionId) -> agent_client_protocol::Result<()> { |
| 664 | if self.sessions.lock().unwrap().contains_key(id) { |
| 665 | return Ok(()); |
| 666 | } |
| 667 | let stored = self |
| 668 | .load_stored(id) |
| 669 | .await? |
| 670 | .ok_or_else(|| unknown_session(id))?; |
| 671 | tracing::info!("restored session {id} from disk"); |
| 672 | self.sessions |
| 673 | .lock() |
| 674 | .unwrap() |
| 675 | .entry(id.clone()) |
| 676 | .or_insert_with(|| Session::restored(stored, self.yolo, self.no_mcp())); |
| 677 | Ok(()) |
| 678 | } |
| 679 | |
| 680 | async fn load_stored( |
| 681 | &self, |
| 682 | id: &SessionId, |
| 683 | ) -> agent_client_protocol::Result<Option<StoredSession>> { |
| 684 | let store = self.store.clone(); |
| 685 | let sid = id.0.to_string(); |
| 686 | tokio::task::spawn_blocking(move || store.load(&sid)) |
| 687 | .await |
| 688 | .map_err(|e| agent_client_protocol::Error::internal_error().data(e.to_string()))? |
| 689 | .map_err(|e| agent_client_protocol::Error::internal_error().data(format!("{e:#}"))) |
| 690 | } |
| 691 | |
| 692 | /// A session placeholder until its MCP servers are connected. |
| 693 | fn no_mcp(&self) -> Arc<McpToolset> { |
| 694 | Arc::new(McpToolset::empty(self.mcp_config())) |
| 695 | } |
| 696 | |
| 697 | fn session_mcp(&self, id: &SessionId) -> Arc<McpToolset> { |
| 698 | self.sessions |
| 699 | .lock() |
| 700 | .unwrap() |
| 701 | .get(id) |
| 702 | .map(|s| s.mcp.clone()) |
| 703 | .unwrap_or_else(|| self.no_mcp()) |
| 704 | } |
| 705 | |
| 706 | /// `AvailableCommandsUpdate` for a session, sent after the response that created it. |
| 707 | pub fn send_commands( |
| 708 | &self, |
| 709 | connection: &ConnectionTo<Client>, |
| 710 | id: &SessionId, |
| 711 | ) -> agent_client_protocol::Result<()> { |
| 712 | let mcp = self.session_mcp(id); |
| 713 | let commands = self.profile.slash_commands(&SessionCtx { mcp: &mcp }); |
| 714 | connection.send_notification(SessionNotification::new( |
| 715 | id.clone(), |
| 716 | SessionUpdate::AvailableCommandsUpdate(AvailableCommandsUpdate::new(commands)), |
| 717 | )) |
| 718 | } |
| 719 | |
| 720 | pub async fn new_session( |
| 721 | &self, |
| 722 | req: NewSessionRequest, |
| 723 | ) -> agent_client_protocol::Result<NewSessionResponse> { |
| 724 | check_absolute(&req.cwd, &req.additional_directories)?; |
| 725 | // Lexical normalization only: the client sees `cwd` again in `session/list`, so |
| 726 | // symlinks in it are not rewritten. |
| 727 | let cwd = normalize(&req.cwd); |
| 728 | let roots: Vec<PathBuf> = req |
| 729 | .additional_directories |
| 730 | .iter() |
| 731 | .map(|r| normalize(r)) |
| 732 | .collect(); |
| 733 | let id = SessionId::new(uuid::Uuid::new_v4().to_string()); |
| 734 | let model = self.model(); |
| 735 | let mcp = self.connect_mcp(&req.mcp_servers).await; |
| 736 | let session = Session::fresh(cwd, roots, model.clone(), self.yolo, mcp); |
| 737 | self.sessions.lock().unwrap().insert(id.clone(), session); |
| 738 | // Saved right away so the thread survives a restart before its first prompt. |
| 739 | self.persist(&id).await; |
| 740 | let options = self.config_options(&model, self.yolo).await; |
| 741 | Ok(NewSessionResponse::new(id).config_options(options)) |
| 742 | } |
| 743 | |
| 744 | /// Live and stored sessions, optionally filtered by working directory, newest first. |
| 745 | pub fn list_sessions( |
| 746 | &self, |
| 747 | req: ListSessionsRequest, |
| 748 | ) -> agent_client_protocol::Result<ListSessionsResponse> { |
| 749 | // Every session fits on one page, so no cursor is ever handed out; any cursor is stale. |
| 750 | if let Some(cursor) = &req.cursor { |
| 751 | return Err(invalid(format!("invalid cursor {cursor}"))); |
| 752 | } |
| 753 | struct Row { |
| 754 | id: SessionId, |
| 755 | cwd: PathBuf, |
| 756 | roots: Vec<PathBuf>, |
| 757 | title: Option<String>, |
| 758 | updated_at: u64, |
| 759 | } |
| 760 | let mut rows: HashMap<String, Row> = self |
| 761 | .store |
| 762 | .list() |
| 763 | .into_iter() |
| 764 | .map(|(id, s)| { |
| 765 | ( |
| 766 | id.clone(), |
| 767 | Row { |
| 768 | id: SessionId::new(id), |
| 769 | cwd: s.cwd, |
| 770 | roots: s.roots, |
| 771 | title: s.title, |
| 772 | updated_at: s.updated_at, |
| 773 | }, |
| 774 | ) |
| 775 | }) |
| 776 | .collect(); |
| 777 | // Live state wins over the (possibly older) stored copy. |
| 778 | for (id, s) in self.sessions.lock().unwrap().iter() { |
| 779 | rows.insert( |
| 780 | id.0.to_string(), |
| 781 | Row { |
| 782 | id: id.clone(), |
| 783 | cwd: s.cwd.clone(), |
| 784 | roots: s.roots.clone(), |
| 785 | title: s.title.clone(), |
| 786 | updated_at: s.updated_at, |
| 787 | }, |
| 788 | ); |
| 789 | } |
| 790 | // Stored cwds are normalized, so normalize the filter the same way. |
| 791 | let wanted = req.cwd.as_deref().map(normalize); |
| 792 | let mut rows: Vec<Row> = rows |
| 793 | .into_values() |
| 794 | .filter(|r| wanted.as_ref().is_none_or(|cwd| *cwd == normalize(&r.cwd))) |
| 795 | .collect(); |
| 796 | rows.sort_by(|a, b| { |
| 797 | b.updated_at |
| 798 | .cmp(&a.updated_at) |
| 799 | .then_with(|| a.id.0.cmp(&b.id.0)) |
| 800 | }); |
| 801 | rows.truncate(MAX_LISTED); |
| 802 | Ok(ListSessionsResponse::new( |
| 803 | rows.into_iter() |
| 804 | .map(|r| { |
| 805 | SessionInfo::new(r.id, r.cwd) |
| 806 | .additional_directories(r.roots) |
| 807 | .title(r.title) |
| 808 | .updated_at(iso8601(r.updated_at)) |
| 809 | }) |
| 810 | .collect(), |
| 811 | )) |
| 812 | } |
| 813 | |
| 814 | pub fn cancel(&self, session_id: &SessionId) { |
| 815 | if let Some(s) = self.sessions.lock().unwrap().get(session_id) { |
| 816 | s.cancel.cancel(); |
| 817 | } |
| 818 | } |
| 819 | |
| 820 | /// Make `id` live for `session/load` and `session/resume`: from memory, else from the |
| 821 | /// store; rebind cwd and roots, and reconnect MCP servers from the request. A well-formed |
| 822 | /// id with no saved history (a thread from a process that never saved it) reopens empty |
| 823 | /// under the same id, so the editor's thread opens instead of failing to launch. Returns |
| 824 | /// whether the history was lost. |
| 825 | async fn reopen( |
| 826 | &self, |
| 827 | id: &SessionId, |
| 828 | cwd: &Path, |
| 829 | roots: &[PathBuf], |
| 830 | mcp_servers: &[McpServer], |
| 831 | ) -> agent_client_protocol::Result<bool> { |
| 832 | check_absolute(cwd, roots)?; |
| 833 | if crate::store::sanitize_id(&id.0).is_err() { |
| 834 | return Err(unknown_session(id)); |
| 835 | } |
| 836 | let live = self.sessions.lock().unwrap().contains_key(id); |
| 837 | let stored = if live { |
| 838 | None |
| 839 | } else { |
| 840 | self.load_stored(id).await? |
| 841 | }; |
| 842 | let history_lost = !live && stored.is_none(); |
| 843 | let cwd = normalize(cwd); |
| 844 | let roots: Vec<PathBuf> = roots.iter().map(|r| normalize(r)).collect(); |
| 845 | let mcp = self.connect_mcp(mcp_servers).await; |
| 846 | |
| 847 | let old_mcp = { |
| 848 | let mut sessions = self.sessions.lock().unwrap(); |
| 849 | if !sessions.contains_key(id) { |
| 850 | let session = match stored { |
| 851 | Some(stored) => Session::restored(stored, self.yolo, self.no_mcp()), |
| 852 | None => { |
| 853 | tracing::warn!("session {id} has no saved history; reopening it empty"); |
| 854 | Session::fresh( |
| 855 | cwd.clone(), |
| 856 | roots.clone(), |
| 857 | self.model(), |
| 858 | self.yolo, |
| 859 | self.no_mcp(), |
| 860 | ) |
| 861 | } |
| 862 | }; |
| 863 | sessions.insert(id.clone(), session); |
| 864 | } |
| 865 | let s = sessions.get_mut(id).ok_or_else(|| unknown_session(id))?; |
| 866 | s.cwd = cwd; |
| 867 | if !roots.is_empty() { |
| 868 | s.roots = roots; |
| 869 | } |
| 870 | s.updated_at = now_secs(); |
| 871 | std::mem::replace(&mut s.mcp, mcp) |
| 872 | }; |
| 873 | if let Some(old) = Arc::into_inner(old_mcp) { |
| 874 | tokio::spawn(old.shutdown()); |
| 875 | } |
| 876 | self.persist(id).await; |
| 877 | Ok(history_lost) |
| 878 | } |
| 879 | |
| 880 | /// Tell the user a reopened thread starts without its earlier history. Shown in the |
| 881 | /// thread only; the model doesn't see it. |
| 882 | fn notify_history_lost( |
| 883 | &self, |
| 884 | connection: &ConnectionTo<Client>, |
| 885 | id: &SessionId, |
| 886 | ) -> agent_client_protocol::Result<()> { |
| 887 | connection.send_notification(SessionNotification::new( |
| 888 | id.clone(), |
| 889 | SessionUpdate::AgentMessageChunk( |
| 890 | ContentChunk::new( |
| 891 | format!( |
| 892 | "{} no longer has this conversation's history, so it starts fresh \ |
| 893 | from here.", |
| 894 | self.profile.info().display_name |
| 895 | ) |
| 896 | .into(), |
| 897 | ) |
| 898 | .message_id(new_message_id()), |
| 899 | ), |
| 900 | )) |
| 901 | } |
| 902 | |
| 903 | /// Replay a session's conversation as `session/update` notifications (text, tool calls |
| 904 | /// with their status and output), then respond with config options. `session/load`. |
| 905 | pub async fn load_session( |
| 906 | &self, |
| 907 | req: LoadSessionRequest, |
| 908 | connection: &ConnectionTo<Client>, |
| 909 | ) -> agent_client_protocol::Result<LoadSessionResponse> { |
| 910 | let id = req.session_id.clone(); |
| 911 | let history_lost = self |
| 912 | .reopen(&id, &req.cwd, &req.additional_directories, &req.mcp_servers) |
| 913 | .await?; |
| 914 | let (messages, cwd, roots) = { |
| 915 | let sessions = self.sessions.lock().unwrap(); |
| 916 | let s = sessions.get(&id).ok_or_else(|| unknown_session(&id))?; |
| 917 | (s.messages.clone(), s.cwd.clone(), s.roots.clone()) |
| 918 | }; |
| 919 | let notify = |update: SessionUpdate| { |
| 920 | connection.send_notification(SessionNotification::new(id.clone(), update)) |
| 921 | }; |
| 922 | // Tool output by call id, from the stored `tool` messages. |
| 923 | let results: HashMap<&str, (&str, bool)> = messages |
| 924 | .iter() |
| 925 | .filter(|m| m.get("role").and_then(Value::as_str) == Some("tool")) |
| 926 | .filter_map(|m| { |
| 927 | Some(( |
| 928 | m.get("tool_call_id")?.as_str()?, |
| 929 | ( |
| 930 | m.get("content").and_then(Value::as_str).unwrap_or(""), |
| 931 | m.get("x_failed").and_then(Value::as_bool).unwrap_or(false), |
| 932 | ), |
| 933 | )) |
| 934 | }) |
| 935 | .collect(); |
| 936 | for msg in &messages { |
| 937 | match msg.get("role").and_then(Value::as_str).unwrap_or("") { |
| 938 | "user" => { |
| 939 | let text = user_text(msg); |
| 940 | if !text.is_empty() { |
| 941 | notify(SessionUpdate::UserMessageChunk( |
| 942 | ContentChunk::new(text.into()).message_id(new_message_id()), |
| 943 | ))?; |
| 944 | } |
| 945 | } |
| 946 | "assistant" => { |
| 947 | let text = msg.get("content").and_then(Value::as_str).unwrap_or(""); |
| 948 | if !text.is_empty() { |
| 949 | notify(SessionUpdate::AgentMessageChunk( |
| 950 | ContentChunk::new(text.to_string().into()).message_id(new_message_id()), |
| 951 | ))?; |
| 952 | } |
| 953 | let calls = msg.get("tool_calls").and_then(Value::as_array); |
| 954 | for tc in calls.into_iter().flatten() { |
| 955 | let call_id = tc.get("id").and_then(Value::as_str).unwrap_or("unknown"); |
| 956 | let name = tc |
| 957 | .pointer("/function/name") |
| 958 | .and_then(Value::as_str) |
| 959 | .unwrap_or("unknown"); |
| 960 | let args: Value = tc |
| 961 | .pointer("/function/arguments") |
| 962 | .and_then(Value::as_str) |
| 963 | .and_then(|a| serde_json::from_str(a).ok()) |
| 964 | .unwrap_or_else(|| json!({})); |
| 965 | let (title, kind, locations) = |
| 966 | tools::describe(&self.toolsets, &cwd, &roots, name, &args); |
| 967 | let mut call = ToolCall::new(call_id.to_string(), title) |
| 968 | .name(name.to_string()) |
| 969 | .kind(kind) |
| 970 | .locations(locations) |
| 971 | .raw_input(args); |
| 972 | match results.get(call_id) { |
| 973 | Some((output, failed)) => { |
| 974 | call = call |
| 975 | .status(if *failed { |
| 976 | ToolCallStatus::Failed |
| 977 | } else { |
| 978 | ToolCallStatus::Completed |
| 979 | }) |
| 980 | .content(vec![ToolCallContent::from(output.to_string())]) |
| 981 | .raw_output(json!({ "output": output, "failed": failed })); |
| 982 | } |
| 983 | // The turn was interrupted before this call produced a result. |
| 984 | None => call = call.status(ToolCallStatus::Failed), |
| 985 | } |
| 986 | notify(SessionUpdate::ToolCall(call))?; |
| 987 | } |
| 988 | } |
| 989 | // system and tool messages are not replayed on their own. |
| 990 | _ => {} |
| 991 | } |
| 992 | } |
| 993 | if history_lost { |
| 994 | self.notify_history_lost(connection, &id)?; |
| 995 | } |
| 996 | Ok(LoadSessionResponse::new().config_options(self.session_config_options(&id).await?)) |
| 997 | } |
| 998 | |
| 999 | /// Rebind a session to a new cwd / additional directories without replaying it. |
| 1000 | pub async fn resume_session( |
| 1001 | &self, |
| 1002 | req: ResumeSessionRequest, |
| 1003 | connection: &ConnectionTo<Client>, |
| 1004 | ) -> agent_client_protocol::Result<ResumeSessionResponse> { |
| 1005 | let id = req.session_id.clone(); |
| 1006 | let history_lost = self |
| 1007 | .reopen(&id, &req.cwd, &req.additional_directories, &req.mcp_servers) |
| 1008 | .await?; |
| 1009 | if history_lost { |
| 1010 | self.notify_history_lost(connection, &id)?; |
| 1011 | } |
| 1012 | Ok(ResumeSessionResponse::new().config_options(self.session_config_options(&id).await?)) |
| 1013 | } |
| 1014 | |
| 1015 | /// Cancel any in-progress work, release MCP servers and scratch audio. The stored |
| 1016 | /// conversation stays; `session/delete` removes it. |
| 1017 | pub fn close_session( |
| 1018 | &self, |
| 1019 | req: CloseSessionRequest, |
| 1020 | ) -> agent_client_protocol::Result<CloseSessionResponse> { |
| 1021 | let id = req.session_id.clone(); |
| 1022 | let removed = self.sessions.lock().unwrap().remove(&id); |
| 1023 | match removed { |
| 1024 | Some(s) => self.release(&id, s), |
| 1025 | // A session from an earlier process that was never reopened has nothing to close. |
| 1026 | None if self.store.load(&id.0).ok().flatten().is_some() => {} |
| 1027 | None => return Err(unknown_session(&id)), |
| 1028 | } |
| 1029 | Ok(CloseSessionResponse::new()) |
| 1030 | } |
| 1031 | |
| 1032 | /// Forget a session, in memory and on disk, with its scratch audio. Idempotent: deleting |
| 1033 | /// an unknown session succeeds. |
| 1034 | pub fn delete_session( |
| 1035 | &self, |
| 1036 | req: DeleteSessionRequest, |
| 1037 | ) -> agent_client_protocol::Result<DeleteSessionResponse> { |
| 1038 | let id = req.session_id.clone(); |
| 1039 | if let Some(s) = self.sessions.lock().unwrap().remove(&id) { |
| 1040 | self.release(&id, s); |
| 1041 | } |
| 1042 | // A malformed id can't name a stored session; nothing to remove. |
| 1043 | if crate::store::sanitize_id(&id.0).is_ok() { |
| 1044 | self.store.delete(&id.0).map_err(|e| { |
| 1045 | agent_client_protocol::Error::internal_error().data(format!("{e:#}")) |
| 1046 | })?; |
| 1047 | } |
| 1048 | self.remove_scratch(&id); |
| 1049 | Ok(DeleteSessionResponse::new()) |
| 1050 | } |
| 1051 | |
| 1052 | fn release(&self, id: &SessionId, s: Session) { |
| 1053 | s.cancel.cancel(); |
| 1054 | self.remove_scratch(id); |
| 1055 | if let Some(mcp) = Arc::into_inner(s.mcp) { |
| 1056 | tokio::spawn(mcp.shutdown()); |
| 1057 | } |
| 1058 | } |
| 1059 | |
| 1060 | /// Stop everything: used when the connection ends. |
| 1061 | pub async fn shutdown(&self) { |
| 1062 | let sessions: Vec<Session> = self |
| 1063 | .sessions |
| 1064 | .lock() |
| 1065 | .unwrap() |
| 1066 | .drain() |
| 1067 | .map(|(_, s)| s) |
| 1068 | .collect(); |
| 1069 | for s in sessions { |
| 1070 | s.cancel.cancel(); |
| 1071 | if let Some(mcp) = Arc::into_inner(s.mcp) { |
| 1072 | mcp.shutdown().await; |
| 1073 | } |
| 1074 | } |
| 1075 | for server in self.profile.builtin_mcp_servers() { |
| 1076 | server.shutdown().await; |
| 1077 | } |
| 1078 | } |
| 1079 | |
| 1080 | pub async fn prompt( |
| 1081 | &self, |
| 1082 | request: PromptRequest, |
| 1083 | responder: Responder<PromptResponse>, |
| 1084 | connection: ConnectionTo<Client>, |
| 1085 | ) -> agent_client_protocol::Result<()> { |
| 1086 | // A client may prompt a reopened thread without loading or resuming it first. |
| 1087 | if let Err(e) = self.restore(&request.session_id).await { |
| 1088 | return responder.respond_with_error(e); |
| 1089 | } |
| 1090 | if let Some(reply) = self.answer_slash(&request).await { |
| 1091 | connection.send_notification(SessionNotification::new( |
| 1092 | request.session_id.clone(), |
| 1093 | SessionUpdate::AgentMessageChunk( |
| 1094 | ContentChunk::new(reply.into()).message_id(new_message_id()), |
| 1095 | ), |
| 1096 | ))?; |
| 1097 | return responder.respond(PromptResponse::new(StopReason::EndTurn)); |
| 1098 | } |
| 1099 | match self.run_turn(request, connection).await { |
| 1100 | Ok(stop) => responder.respond(PromptResponse::new(stop)), |
| 1101 | Err(e) if e.downcast_ref::<UnknownSession>().is_some() => { |
| 1102 | responder.respond_with_error(invalid(e.to_string())) |
| 1103 | } |
| 1104 | Err(e) => responder.respond_with_error( |
| 1105 | agent_client_protocol::Error::internal_error().data(format!("{e:#}")), |
| 1106 | ), |
| 1107 | } |
| 1108 | } |
| 1109 | |
| 1110 | /// The profile's reply to a slash command it answers without the model, for a prompt that |
| 1111 | /// is a single text block. The reply is not added to the conversation. |
| 1112 | async fn answer_slash(&self, request: &PromptRequest) -> Option<String> { |
| 1113 | let [ContentBlock::Text(text)] = request.prompt.as_slice() else { |
| 1114 | return None; |
| 1115 | }; |
| 1116 | if !text.text.trim_start().starts_with('/') { |
| 1117 | return None; |
| 1118 | } |
| 1119 | let (model, _) = self.session_settings(&request.session_id)?; |
| 1120 | let models = self.models().await; |
| 1121 | let mcp = self.session_mcp(&request.session_id); |
| 1122 | self.profile.answer_slash( |
| 1123 | &text.text, |
| 1124 | &CommandCtx { |
| 1125 | mcp: &mcp, |
| 1126 | model: &model, |
| 1127 | models: &models, |
| 1128 | }, |
| 1129 | ) |
| 1130 | } |
| 1131 | |
| 1132 | async fn run_turn( |
| 1133 | &self, |
| 1134 | request: PromptRequest, |
| 1135 | connection: ConnectionTo<Client>, |
| 1136 | ) -> anyhow::Result<StopReason> { |
| 1137 | let session_id = request.session_id.clone(); |
| 1138 | let caps = self.client_caps.lock().unwrap().clone(); |
| 1139 | let blocks = |
| 1140 | prompt::resolve_resource_links(&request.prompt, &caps, &connection, &session_id).await; |
| 1141 | let prompt_text = prompt::prompt_to_text(&blocks, &self.profile.audio_hints()); |
| 1142 | // Check the session out for the turn. |
| 1143 | let ( |
| 1144 | cwd, |
| 1145 | roots, |
| 1146 | model, |
| 1147 | history, |
| 1148 | cancel, |
| 1149 | always_allowed, |
| 1150 | always_rejected, |
| 1151 | auto_approve, |
| 1152 | mcp, |
| 1153 | new_title, |
| 1154 | ) = { |
| 1155 | let mut sessions = self.sessions.lock().unwrap(); |
| 1156 | let s = sessions |
| 1157 | .get_mut(&session_id) |
| 1158 | .ok_or_else(|| UnknownSession(session_id.to_string()))?; |
| 1159 | if s.busy { |
| 1160 | anyhow::bail!("a prompt is already running for session {session_id}"); |
| 1161 | } |
| 1162 | s.busy = true; |
| 1163 | s.cancel = CancellationToken::new(); |
| 1164 | s.updated_at = now_secs(); |
| 1165 | let new_title = if s.title.is_none() && !prompt_text.trim().is_empty() { |
| 1166 | s.title = Some(session_title(&prompt_text)); |
| 1167 | s.title.clone() |
| 1168 | } else { |
| 1169 | None |
| 1170 | }; |
| 1171 | ( |
| 1172 | s.cwd.clone(), |
| 1173 | s.roots.clone(), |
| 1174 | s.model.clone(), |
| 1175 | std::mem::take(&mut s.messages), |
| 1176 | s.cancel.clone(), |
| 1177 | s.always_allowed.clone(), |
| 1178 | s.always_rejected.clone(), |
| 1179 | s.auto_approve.clone(), |
| 1180 | s.mcp.clone(), |
| 1181 | new_title, |
| 1182 | ) |
| 1183 | }; |
| 1184 | if let Some(title) = new_title { |
| 1185 | let update = SessionInfoUpdate::new() |
| 1186 | .title(title) |
| 1187 | .updated_at(iso8601(now_secs())); |
| 1188 | let _ = connection.send_notification(SessionNotification::new( |
| 1189 | session_id.clone(), |
| 1190 | SessionUpdate::SessionInfoUpdate(update), |
| 1191 | )); |
| 1192 | } |
| 1193 | |
| 1194 | let mut messages = vec![json!({ |
| 1195 | "role": "system", |
| 1196 | "content": self.system_prompt(&cwd, &roots, &mcp), |
| 1197 | })]; |
| 1198 | messages.extend(history); |
| 1199 | |
| 1200 | let ctx = ToolCtx { |
| 1201 | connection: connection.clone(), |
| 1202 | session_id: session_id.clone(), |
| 1203 | cwd, |
| 1204 | roots, |
| 1205 | caps, |
| 1206 | cancel, |
| 1207 | auto_approve, |
| 1208 | always_allowed, |
| 1209 | always_rejected, |
| 1210 | mcp, |
| 1211 | command_timeout: std::env::var(self.env.var("COMMAND_TIMEOUT_SECS")) |
| 1212 | .ok() |
| 1213 | .and_then(|v| v.parse().ok()) |
| 1214 | .map(Duration::from_secs) |
| 1215 | .unwrap_or(DEFAULT_COMMAND_TIMEOUT), |
| 1216 | }; |
| 1217 | let result = self.turn(&ctx, &model, &mut messages, &blocks).await; |
| 1218 | |
| 1219 | // Check the history back in, even when the turn failed. |
| 1220 | let history = messages.split_off(1); |
| 1221 | let stored = { |
| 1222 | let mut sessions = self.sessions.lock().unwrap(); |
| 1223 | sessions.get_mut(&session_id).map(|s| { |
| 1224 | s.messages = history; |
| 1225 | s.busy = false; |
| 1226 | s.updated_at = now_secs(); |
| 1227 | s.stored() |
| 1228 | }) |
| 1229 | }; |
| 1230 | if let Some(stored) = stored { |
| 1231 | let updated_at = stored.updated_at; |
| 1232 | let store = self.store.clone(); |
| 1233 | let id = session_id.0.to_string(); |
| 1234 | let saved = tokio::task::spawn_blocking(move || store.save(&id, &stored)).await; |
| 1235 | if let Ok(Err(e)) | Err(e) = saved.map_err(anyhow::Error::from) { |
| 1236 | tracing::warn!("saving session {session_id} failed: {e:#}"); |
| 1237 | } |
| 1238 | let _ = connection.send_notification(SessionNotification::new( |
| 1239 | session_id.clone(), |
| 1240 | SessionUpdate::SessionInfoUpdate( |
| 1241 | SessionInfoUpdate::new().updated_at(iso8601(updated_at)), |
| 1242 | ), |
| 1243 | )); |
| 1244 | } |
| 1245 | result |
| 1246 | } |
| 1247 | |
| 1248 | /// Everything inside a turn after the session is checked out. |
| 1249 | async fn turn( |
| 1250 | &self, |
| 1251 | ctx: &ToolCtx, |
| 1252 | model: &str, |
| 1253 | messages: &mut Vec<Value>, |
| 1254 | blocks: &[ContentBlock], |
| 1255 | ) -> anyhow::Result<StopReason> { |
| 1256 | let audio_paths = match self.scratch_dir(&ctx.session_id) { |
| 1257 | Some(dir) => prompt::save_audio_blocks(blocks, &dir).await?, |
| 1258 | None => Vec::new(), |
| 1259 | }; |
| 1260 | let hints = self.profile.audio_hints(); |
| 1261 | let slash = self.profile.expand_slash( |
| 1262 | &prompt::prompt_to_text(blocks, &hints), |
| 1263 | &SessionCtx { mcp: &ctx.mcp }, |
| 1264 | ); |
| 1265 | let content = prompt::prompt_to_content(blocks, &audio_paths, slash.as_deref(), &hints); |
| 1266 | messages.push(json!({ "role": "user", "content": content })); |
| 1267 | self.agent_loop(ctx, model, messages).await |
| 1268 | } |
| 1269 | |
| 1270 | async fn agent_loop( |
| 1271 | &self, |
| 1272 | ctx: &ToolCtx, |
| 1273 | model: &str, |
| 1274 | messages: &mut Vec<Value>, |
| 1275 | ) -> anyhow::Result<StopReason> { |
| 1276 | let llm = self.llm(); |
| 1277 | let tool_defs = tools::definitions(&self.toolsets, &ctx.mcp); |
| 1278 | let notify = |update: SessionUpdate| { |
| 1279 | ctx.connection |
| 1280 | .send_notification(SessionNotification::new(ctx.session_id.clone(), update)) |
| 1281 | }; |
| 1282 | // `size` is the model's context window, which OpenAI-style endpoints don't report. |
| 1283 | let context_window = std::env::var(self.env.var("CONTEXT_WINDOW")) |
| 1284 | .ok() |
| 1285 | .and_then(|v| v.parse().ok()) |
| 1286 | .unwrap_or(DEFAULT_CONTEXT_WINDOW); |
| 1287 | |
| 1288 | for _ in 0..MAX_TURNS { |
| 1289 | // One message id per assistant message (and per thought), so clients can group |
| 1290 | // the streamed chunks. |
| 1291 | let message_id = new_message_id(); |
| 1292 | let thought_id = new_message_id(); |
| 1293 | let wire: Vec<Value> = messages.iter().map(wire_message).collect(); |
| 1294 | let completion = tokio::select! { |
| 1295 | r = llm.complete(model, &wire, &tool_defs, |delta| { |
| 1296 | let update = match delta { |
| 1297 | Delta::Text(t) => SessionUpdate::AgentMessageChunk( |
| 1298 | ContentChunk::new(t.to_string().into()).message_id(message_id.clone()), |
| 1299 | ), |
| 1300 | Delta::Reasoning(t) => SessionUpdate::AgentThoughtChunk( |
| 1301 | ContentChunk::new(t.to_string().into()).message_id(thought_id.clone()), |
| 1302 | ), |
| 1303 | }; |
| 1304 | if let Err(e) = notify(update) { |
| 1305 | tracing::warn!("failed to send update: {e}"); |
| 1306 | } |
| 1307 | }) => r?, |
| 1308 | () = ctx.cancel.cancelled() => return Ok(StopReason::Cancelled), |
| 1309 | }; |
| 1310 | messages.push(completion.to_message()); |
| 1311 | if let Some(usage) = completion.usage { |
| 1312 | let used = usage.used(); |
| 1313 | notify(SessionUpdate::UsageUpdate(UsageUpdate::new( |
| 1314 | used, |
| 1315 | context_window.max(used), |
| 1316 | )))?; |
| 1317 | } |
| 1318 | |
| 1319 | if completion.tool_calls.is_empty() { |
| 1320 | return Ok(match completion.finish_reason.as_deref() { |
| 1321 | Some("length") => StopReason::MaxTokens, |
| 1322 | Some("content_filter") => StopReason::Refusal, |
| 1323 | _ => StopReason::EndTurn, |
| 1324 | }); |
| 1325 | } |
| 1326 | |
| 1327 | for tc in &completion.tool_calls { |
| 1328 | // Every tool call needs a matching tool message, even after cancellation. |
| 1329 | if ctx.cancel.is_cancelled() { |
| 1330 | messages.push(tool_message(&tc.id, "Cancelled by user.", true)); |
| 1331 | continue; |
| 1332 | } |
| 1333 | let parsed: Result<Value, _> = serde_json::from_str(&tc.arguments); |
| 1334 | let args = parsed.as_ref().cloned().unwrap_or_else(|_| json!({})); |
| 1335 | let (title, kind, locations) = |
| 1336 | tools::describe(&self.toolsets, &ctx.cwd, &ctx.roots, &tc.name, &args); |
| 1337 | notify(SessionUpdate::ToolCall( |
| 1338 | ToolCall::new(tc.id.clone(), title) |
| 1339 | .name(tc.name.clone()) |
| 1340 | .kind(kind) |
| 1341 | .status(ToolCallStatus::InProgress) |
| 1342 | .locations(locations) |
| 1343 | .raw_input(args.clone()), |
| 1344 | ))?; |
| 1345 | |
| 1346 | let outcome: ToolOutcome = if parsed.is_err() { |
| 1347 | ToolOutcome::err(format!("Invalid JSON arguments: {}", tc.arguments)) |
| 1348 | } else { |
| 1349 | tools::execute( |
| 1350 | ctx, |
| 1351 | &self.toolsets, |
| 1352 | self.profile.as_ref(), |
| 1353 | &tc.id, |
| 1354 | &tc.name, |
| 1355 | args, |
| 1356 | ) |
| 1357 | .await |
| 1358 | }; |
| 1359 | |
| 1360 | let mut fields = ToolCallUpdateFields::new() |
| 1361 | .status(if outcome.failed { |
| 1362 | ToolCallStatus::Failed |
| 1363 | } else { |
| 1364 | ToolCallStatus::Completed |
| 1365 | }) |
| 1366 | .content(outcome.content) |
| 1367 | .raw_output(json!({ "output": outcome.text, "failed": outcome.failed })); |
| 1368 | if !outcome.locations.is_empty() { |
| 1369 | fields = fields.locations(locations_of(&outcome.locations)); |
| 1370 | } |
| 1371 | ctx.update(&tc.id, fields)?; |
| 1372 | messages.push(tool_message(&tc.id, &outcome.text, outcome.failed)); |
| 1373 | } |
| 1374 | if ctx.cancel.is_cancelled() { |
| 1375 | return Ok(StopReason::Cancelled); |
| 1376 | } |
| 1377 | } |
| 1378 | Ok(StopReason::MaxTurnRequests) |
| 1379 | } |
| 1380 | } |
| 1381 | |
| 1382 | /// Marker error: the prompt named a session that doesn't exist. |
| 1383 | #[derive(Debug)] |
| 1384 | struct UnknownSession(String); |
| 1385 | |
| 1386 | impl std::fmt::Display for UnknownSession { |
| 1387 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { |
| 1388 | write!(f, "unknown session {}", self.0) |
| 1389 | } |
| 1390 | } |
| 1391 | |
| 1392 | impl std::error::Error for UnknownSession {} |
| 1393 | |
| 1394 | #[cfg(test)] |
| 1395 | mod tests { |
| 1396 | use super::*; |
| 1397 | |
| 1398 | #[test] |
| 1399 | fn iso8601_known_dates() { |
| 1400 | assert_eq!(iso8601(0), "1970-01-01T00:00:00Z"); |
| 1401 | assert_eq!(iso8601(1_700_000_000), "2023-11-14T22:13:20Z"); |
| 1402 | assert_eq!(iso8601(951_782_400), "2000-02-29T00:00:00Z"); |
| 1403 | } |
| 1404 | |
| 1405 | #[test] |
| 1406 | fn titles_are_trimmed_and_single_line() { |
| 1407 | assert_eq!(session_title(" hello\nworld"), "hello"); |
| 1408 | let long = "x".repeat(200); |
| 1409 | let t = session_title(&long); |
| 1410 | assert_eq!(t.chars().count(), 80); |
| 1411 | assert!(t.ends_with('…')); |
| 1412 | } |
| 1413 | |
| 1414 | #[test] |
| 1415 | fn terminal_auth_support_from_capability_or_legacy_meta() { |
| 1416 | let caps = |v: Value| serde_json::from_value::<ClientCapabilities>(v).unwrap(); |
| 1417 | assert!(!supports_terminal_auth(&caps(json!({})))); |
| 1418 | assert!(supports_terminal_auth(&caps( |
| 1419 | json!({"auth": {"terminal": true}}) |
| 1420 | ))); |
| 1421 | // The ACP registry validator's payload. |
| 1422 | assert!(supports_terminal_auth(&caps( |
| 1423 | json!({"terminal": true, "_meta": {"terminal_output": true, "terminal-auth": true}}) |
| 1424 | ))); |
| 1425 | assert!(!supports_terminal_auth(&caps( |
| 1426 | json!({"_meta": {"terminal-auth": false}}) |
| 1427 | ))); |
| 1428 | } |
| 1429 | |
| 1430 | #[test] |
| 1431 | fn version_negotiation_always_v1() { |
| 1432 | assert_eq!(negotiate_version(ProtocolVersion::V1), ProtocolVersion::V1); |
| 1433 | assert_eq!( |
| 1434 | negotiate_version(ProtocolVersion::from(7)), |
| 1435 | ProtocolVersion::V1 |
| 1436 | ); |
| 1437 | assert_eq!( |
| 1438 | negotiate_version(ProtocolVersion::from(0)), |
| 1439 | ProtocolVersion::V1 |
| 1440 | ); |
| 1441 | } |
| 1442 | |
| 1443 | #[test] |
| 1444 | fn private_failure_flag_never_reaches_the_wire() { |
| 1445 | let m = tool_message("c1", "boom", true); |
| 1446 | assert_eq!(m["x_failed"], true); |
| 1447 | let w = wire_message(&m); |
| 1448 | assert!(w.get("x_failed").is_none()); |
| 1449 | assert_eq!(w["content"], "boom"); |
| 1450 | assert!(tool_message("c1", "ok", false).get("x_failed").is_none()); |
| 1451 | } |
| 1452 | |
| 1453 | #[test] |
| 1454 | fn user_text_handles_parts() { |
| 1455 | let m = json!({"role": "user", "content": [ |
| 1456 | {"type": "text", "text": "hi"}, |
| 1457 | {"type": "image_url", "image_url": {"url": "x"}} |
| 1458 | ]}); |
| 1459 | assert_eq!(user_text(&m), "hi\n\n[image]"); |
| 1460 | assert_eq!(user_text(&json!({"content": "plain"})), "plain"); |
| 1461 | } |
| 1462 | } |