| 1 | //! Minimal Agent Client Protocol stdio adapter. |
| 2 | //! |
| 3 | //! This intentionally starts with the ACP baseline: initialize, new session, |
| 4 | //! prompt, and cancel. It keeps stdout protocol-clean for editor clients and |
| 5 | //! routes prompts through the same configured DeepSeek client as one-shot CLI |
| 6 | //! mode. |
| 7 | //! |
| 8 | //! `session/prompt` streams the provider response: each text delta is emitted |
| 9 | //! as a `session/update` agent_message_chunk as it arrives, instead of buffering |
| 10 | //! the whole turn and sending one chunk at the end. The stream is consumed |
| 11 | //! concurrently with the input reader so that a `session/cancel` for the same |
| 12 | //! session can interrupt the turn mid-stream (returning `stopReason: "cancelled"`) |
| 13 | //! instead of being queued behind it. A single writer task is preserved so |
| 14 | //! stdout stays protocol-clean. |
| 15 | |
| 16 | use std::collections::HashMap; |
| 17 | use std::path::PathBuf; |
| 18 | |
| 19 | use anyhow::{Result, anyhow}; |
| 20 | use futures_util::StreamExt; |
| 21 | use serde_json::{Value, json}; |
| 22 | use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncWrite, AsyncWriteExt, BufReader, Lines}; |
| 23 | |
| 24 | use crate::client::DeepSeekClient; |
| 25 | use crate::config::{ApiProvider, Config}; |
| 26 | use crate::llm_client::{LlmClient, StreamEventBox}; |
| 27 | use crate::models::{ |
| 28 | ContentBlock, ContentBlockStart, Delta, Message, MessageRequest, StreamEvent, SystemPrompt, |
| 29 | }; |
| 30 | |
| 31 | const ACP_PROTOCOL_VERSION: u64 = 1; |
| 32 | |
| 33 | pub async fn run_acp_server(config: Config, model: String, default_cwd: PathBuf) -> Result<()> { |
| 34 | let stdin = tokio::io::stdin(); |
| 35 | let stdout = tokio::io::stdout(); |
| 36 | let mut reader = BufReader::new(stdin).lines(); |
| 37 | let mut writer = tokio::io::BufWriter::new(stdout); |
| 38 | let mut server = AcpServer::new(config, model, default_cwd); |
| 39 | |
| 40 | while let Some(line) = reader.next_line().await? { |
| 41 | if line.trim().is_empty() { |
| 42 | continue; |
| 43 | } |
| 44 | |
| 45 | let message: Value = match serde_json::from_str(&line) { |
| 46 | Ok(value) => value, |
| 47 | Err(err) => { |
| 48 | write_jsonrpc_error(&mut writer, None, -32700, format!("invalid json: {err}")) |
| 49 | .await?; |
| 50 | continue; |
| 51 | } |
| 52 | }; |
| 53 | |
| 54 | if message.get("jsonrpc").and_then(Value::as_str) != Some("2.0") { |
| 55 | write_jsonrpc_error( |
| 56 | &mut writer, |
| 57 | message |
| 58 | .get("id") |
| 59 | .cloned() |
| 60 | .map(|id| server.response_id_policy.response_id(id)), |
| 61 | -32600, |
| 62 | "jsonrpc version must be 2.0", |
| 63 | ) |
| 64 | .await?; |
| 65 | continue; |
| 66 | } |
| 67 | |
| 68 | let id = message.get("id").cloned(); |
| 69 | let method = match message.get("method").and_then(Value::as_str) { |
| 70 | Some(method) => method, |
| 71 | None => { |
| 72 | write_jsonrpc_error( |
| 73 | &mut writer, |
| 74 | id.map(|id| server.response_id_policy.response_id(id)), |
| 75 | -32600, |
| 76 | "missing method", |
| 77 | ) |
| 78 | .await?; |
| 79 | continue; |
| 80 | } |
| 81 | }; |
| 82 | let params = message.get("params").cloned().unwrap_or_else(|| json!({})); |
| 83 | |
| 84 | // `session/prompt` is driven concurrently with the reader so a |
| 85 | // `session/cancel` can interrupt the in-flight provider call. Every |
| 86 | // other method is request/response and handled synchronously below. |
| 87 | if method == "session/prompt" { |
| 88 | match server.begin_prompt(params) { |
| 89 | Ok(prepared) => { |
| 90 | let PreparedPrompt { |
| 91 | session_id, |
| 92 | messages, |
| 93 | cwd, |
| 94 | } = prepared; |
| 95 | // Opening the stream borrows `&server` only briefly; the |
| 96 | // returned `StreamEventBox` is `'static`, so it can be raced |
| 97 | // against the reader without holding a borrow on the server, |
| 98 | // and the main task keeps exclusive ownership of stdout. |
| 99 | match server.open_prompt_stream(&messages, &cwd).await { |
| 100 | Ok(stream) => { |
| 101 | let response_id_policy = server.response_id_policy; |
| 102 | let outcome = drive_prompt_stream( |
| 103 | stream, |
| 104 | &session_id, |
| 105 | response_id_policy, |
| 106 | &mut reader, |
| 107 | &mut writer, |
| 108 | ) |
| 109 | .await; |
| 110 | match outcome { |
| 111 | Ok(PromptOutcome::Completed(output)) => { |
| 112 | // Chunks were already streamed; record the full |
| 113 | // assistant turn in history for the next prompt. |
| 114 | server.finish_prompt(&session_id, &output); |
| 115 | if let Some(id) = id { |
| 116 | let id = response_id_policy.response_id(id); |
| 117 | write_jsonrpc_result( |
| 118 | &mut writer, |
| 119 | id, |
| 120 | json!({ "stopReason": "end_turn" }), |
| 121 | ) |
| 122 | .await?; |
| 123 | } |
| 124 | } |
| 125 | Ok(PromptOutcome::Cancelled) => { |
| 126 | if let Some(id) = id { |
| 127 | let id = response_id_policy.response_id(id); |
| 128 | write_jsonrpc_result( |
| 129 | &mut writer, |
| 130 | id, |
| 131 | json!({ "stopReason": "cancelled" }), |
| 132 | ) |
| 133 | .await?; |
| 134 | } |
| 135 | } |
| 136 | Err(err) => { |
| 137 | let id = id.map(|id| response_id_policy.response_id(id)); |
| 138 | write_jsonrpc_error(&mut writer, id, -32603, err.to_string()) |
| 139 | .await?; |
| 140 | } |
| 141 | } |
| 142 | } |
| 143 | Err(err) => { |
| 144 | let id = id.map(|id| server.response_id_policy.response_id(id)); |
| 145 | write_jsonrpc_error(&mut writer, id, -32603, err.to_string()).await?; |
| 146 | } |
| 147 | } |
| 148 | } |
| 149 | Err(err) => { |
| 150 | let id = id.map(|id| server.response_id_policy.response_id(id)); |
| 151 | write_jsonrpc_error(&mut writer, id, err.code, err.message).await?; |
| 152 | } |
| 153 | } |
| 154 | continue; |
| 155 | } |
| 156 | |
| 157 | match server.handle_request(method, params).await { |
| 158 | Ok(AcpDispatch::Response(result)) => { |
| 159 | if let Some(id) = id { |
| 160 | let id = server.response_id_policy.response_id(id); |
| 161 | write_jsonrpc_result(&mut writer, id, result).await?; |
| 162 | } |
| 163 | } |
| 164 | Ok(AcpDispatch::Shutdown) => { |
| 165 | if let Some(id) = id { |
| 166 | let id = server.response_id_policy.response_id(id); |
| 167 | write_jsonrpc_result(&mut writer, id, json!(null)).await?; |
| 168 | } |
| 169 | break; |
| 170 | } |
| 171 | Err(err) => { |
| 172 | let id = id.map(|id| server.response_id_policy.response_id(id)); |
| 173 | write_jsonrpc_error(&mut writer, id, err.code, err.message).await?; |
| 174 | } |
| 175 | } |
| 176 | } |
| 177 | |
| 178 | Ok(()) |
| 179 | } |
| 180 | |
| 181 | /// Outcome of a `session/prompt` turn driven against the input stream. |
| 182 | #[derive(Debug, PartialEq, Eq)] |
| 183 | enum PromptOutcome { |
| 184 | /// The provider call finished first; carries the assistant text. |
| 185 | Completed(String), |
| 186 | /// A matching `session/cancel` arrived before the call finished. |
| 187 | Cancelled, |
| 188 | } |
| 189 | |
| 190 | /// The text payload an ACP client should see for a given stream event, if any. |
| 191 | /// ACP baseline is text-only, so thinking/tool/control events carry no chunk. |
| 192 | fn stream_text_chunk(event: &StreamEvent) -> Option<&str> { |
| 193 | match event { |
| 194 | StreamEvent::ContentBlockDelta { |
| 195 | delta: Delta::TextDelta { text }, |
| 196 | .. |
| 197 | } => Some(text), |
| 198 | StreamEvent::ContentBlockStart { |
| 199 | content_block: ContentBlockStart::Text { text }, |
| 200 | .. |
| 201 | } => Some(text), |
| 202 | _ => None, |
| 203 | } |
| 204 | } |
| 205 | |
| 206 | /// Consume a provider response `stream`, emitting each text delta as a |
| 207 | /// `session/update` chunk, while concurrently watching `reader` for a |
| 208 | /// `session/cancel` targeting `session_id`. |
| 209 | /// |
| 210 | /// This is the streaming + cancellation control point. It is generic over the |
| 211 | /// reader/writer and takes the boxed stream, so it is unit-tested with canned |
| 212 | /// in-memory streams and readers — no real provider call required. The caller |
| 213 | /// keeps the only writer, so streamed chunks and acknowledgements all stay on |
| 214 | /// the single protocol-clean stdout stream. |
| 215 | /// |
| 216 | /// Returns [`PromptOutcome::Completed`] with the full accumulated text once the |
| 217 | /// stream ends (or emits `message_stop`), so the caller can record the turn in |
| 218 | /// history. A matching `session/cancel` (request or notification form) ends it |
| 219 | /// early with [`PromptOutcome::Cancelled`] — dropping the stream aborts the |
| 220 | /// underlying provider connection. The turn is single-flight: a cancel for a |
| 221 | /// different session is acknowledged and ignored; any other concurrent *request* |
| 222 | /// is rejected with a clear error so the client is not left waiting; |
| 223 | /// notifications without an id are ignored. |
| 224 | async fn drive_prompt_stream<R, W>( |
| 225 | mut stream: StreamEventBox, |
| 226 | session_id: &str, |
| 227 | response_id_policy: JsonRpcResponseIdPolicy, |
| 228 | reader: &mut Lines<R>, |
| 229 | writer: &mut W, |
| 230 | ) -> Result<PromptOutcome> |
| 231 | where |
| 232 | R: AsyncBufRead + Unpin, |
| 233 | W: AsyncWrite + Unpin, |
| 234 | { |
| 235 | let mut accumulated = String::new(); |
| 236 | // Once input closes mid-turn we stop selecting on the reader and just drain |
| 237 | // the stream to completion, rather than spinning on repeated EOFs. |
| 238 | let mut reader_open = true; |
| 239 | loop { |
| 240 | tokio::select! { |
| 241 | event = stream.next() => { |
| 242 | match event { |
| 243 | // Stream exhausted without an explicit stop: turn is done. |
| 244 | None => return Ok(PromptOutcome::Completed(accumulated)), |
| 245 | Some(Ok(event)) => { |
| 246 | if let Some(text) = stream_text_chunk(&event) |
| 247 | && !text.is_empty() { |
| 248 | accumulated.push_str(text); |
| 249 | write_session_update(writer, session_id, text.to_string()).await?; |
| 250 | } |
| 251 | match event { |
| 252 | StreamEvent::MessageStop => { |
| 253 | return Ok(PromptOutcome::Completed(accumulated)); |
| 254 | } |
| 255 | StreamEvent::Error { error } => { |
| 256 | return Err(anyhow!("provider stream error: {error}")); |
| 257 | } |
| 258 | _ => {} |
| 259 | } |
| 260 | } |
| 261 | Some(Err(err)) => return Err(err), |
| 262 | } |
| 263 | } |
| 264 | line = reader.next_line(), if reader_open => { |
| 265 | let line = match line? { |
| 266 | Some(line) => line, |
| 267 | // Input closed mid-turn: stop watching it, keep draining. |
| 268 | None => { |
| 269 | reader_open = false; |
| 270 | continue; |
| 271 | } |
| 272 | }; |
| 273 | if line.trim().is_empty() { |
| 274 | continue; |
| 275 | } |
| 276 | let message: Value = match serde_json::from_str(&line) { |
| 277 | Ok(value) => value, |
| 278 | Err(err) => { |
| 279 | write_jsonrpc_error(writer, None, -32700, format!("invalid json: {err}")) |
| 280 | .await?; |
| 281 | continue; |
| 282 | } |
| 283 | }; |
| 284 | let id = message.get("id").cloned(); |
| 285 | match message.get("method").and_then(Value::as_str) { |
| 286 | Some("session/cancel") => { |
| 287 | let target = message.pointer("/params/sessionId").and_then(Value::as_str); |
| 288 | // A cancel with no sessionId is treated as targeting the |
| 289 | // single in-flight turn. |
| 290 | if target.is_none() || target == Some(session_id) { |
| 291 | if let Some(id) = id { |
| 292 | let id = response_id_policy.response_id(id); |
| 293 | write_jsonrpc_result(writer, id, json!(null)).await?; |
| 294 | } |
| 295 | // Dropping `stream` on return aborts the provider call. |
| 296 | return Ok(PromptOutcome::Cancelled); |
| 297 | } |
| 298 | // Cancel for some other session: acknowledge, keep going. |
| 299 | if let Some(id) = id { |
| 300 | let id = response_id_policy.response_id(id); |
| 301 | write_jsonrpc_result(writer, id, json!(null)).await?; |
| 302 | } |
| 303 | } |
| 304 | _ => { |
| 305 | // The turn is single-flight; do not silently drop a |
| 306 | // request the client expects a response to. |
| 307 | if let Some(id) = id { |
| 308 | let id = response_id_policy.response_id(id); |
| 309 | write_jsonrpc_error( |
| 310 | writer, |
| 311 | Some(id), |
| 312 | -32603, |
| 313 | "a session/prompt turn is already in progress", |
| 314 | ) |
| 315 | .await?; |
| 316 | } |
| 317 | } |
| 318 | } |
| 319 | } |
| 320 | } |
| 321 | } |
| 322 | } |
| 323 | |
| 324 | struct AcpServer { |
| 325 | config: Config, |
| 326 | model: String, |
| 327 | default_cwd: PathBuf, |
| 328 | sessions: HashMap<String, AcpSession>, |
| 329 | response_id_policy: JsonRpcResponseIdPolicy, |
| 330 | } |
| 331 | |
| 332 | struct AcpSession { |
| 333 | cwd: PathBuf, |
| 334 | messages: Vec<Message>, |
| 335 | } |
| 336 | |
| 337 | /// The `&mut self` result of validating a `session/prompt`: the user turn is |
| 338 | /// already recorded, and the cloned conversation + cwd are ready for the |
| 339 | /// borrow-free provider call that the prompt driver races against cancellation. |
| 340 | struct PreparedPrompt { |
| 341 | session_id: String, |
| 342 | messages: Vec<Message>, |
| 343 | cwd: PathBuf, |
| 344 | } |
| 345 | |
| 346 | enum AcpDispatch { |
| 347 | Response(Value), |
| 348 | Shutdown, |
| 349 | } |
| 350 | |
| 351 | #[derive(Debug)] |
| 352 | struct AcpError { |
| 353 | code: i32, |
| 354 | message: String, |
| 355 | } |
| 356 | |
| 357 | impl AcpServer { |
| 358 | fn new(config: Config, model: String, default_cwd: PathBuf) -> Self { |
| 359 | Self { |
| 360 | config, |
| 361 | model, |
| 362 | default_cwd, |
| 363 | sessions: HashMap::new(), |
| 364 | response_id_policy: JsonRpcResponseIdPolicy::Preserve, |
| 365 | } |
| 366 | } |
| 367 | |
| 368 | // `session/prompt` is handled in the main loop (it needs to run concurrently |
| 369 | // with the reader for cancellation); every other method is request/response. |
| 370 | async fn handle_request( |
| 371 | &mut self, |
| 372 | method: &str, |
| 373 | params: Value, |
| 374 | ) -> std::result::Result<AcpDispatch, AcpError> { |
| 375 | match method { |
| 376 | "initialize" => { |
| 377 | self.response_id_policy = JsonRpcResponseIdPolicy::from_initialize_params(¶ms); |
| 378 | Ok(AcpDispatch::Response(initialize_result( |
| 379 | params.get("protocolVersion").and_then(Value::as_u64), |
| 380 | &self.config, |
| 381 | ))) |
| 382 | } |
| 383 | "session/new" => Ok(AcpDispatch::Response(self.new_session(params)?)), |
| 384 | "session/listProviders" => Ok(AcpDispatch::Response(self.list_providers())), |
| 385 | "session/currentModel" => Ok(AcpDispatch::Response(self.current_model())), |
| 386 | "session/selectModel" => Ok(AcpDispatch::Response(self.select_model(params)?)), |
| 387 | // A cancel that arrives with no prompt in flight is an idempotent |
| 388 | // no-op (the in-flight case is handled by the prompt driver). |
| 389 | "session/cancel" => Ok(AcpDispatch::Response(json!(null))), |
| 390 | "shutdown" => Ok(AcpDispatch::Shutdown), |
| 391 | _ => Err(AcpError::method_not_found(method)), |
| 392 | } |
| 393 | } |
| 394 | |
| 395 | fn new_session(&mut self, params: Value) -> std::result::Result<Value, AcpError> { |
| 396 | let cwd = params |
| 397 | .get("cwd") |
| 398 | .and_then(Value::as_str) |
| 399 | .map(PathBuf::from) |
| 400 | .unwrap_or_else(|| self.default_cwd.clone()); |
| 401 | let session_id = format!("codewhale-{}", uuid::Uuid::new_v4()); |
| 402 | self.sessions.insert( |
| 403 | session_id.clone(), |
| 404 | AcpSession { |
| 405 | cwd, |
| 406 | messages: Vec::new(), |
| 407 | }, |
| 408 | ); |
| 409 | Ok(json!({ "sessionId": session_id })) |
| 410 | } |
| 411 | |
| 412 | fn list_providers(&self) -> Value { |
| 413 | let mut providers = ApiProvider::sorted_for_display() |
| 414 | .into_iter() |
| 415 | .map(|provider| { |
| 416 | json!({ |
| 417 | "id": provider.as_str(), |
| 418 | "displayName": provider.display_name(), |
| 419 | "defaultModel": provider.metadata().map(|metadata| metadata.default_model()) |
| 420 | }) |
| 421 | }) |
| 422 | .collect::<Vec<_>>(); |
| 423 | |
| 424 | // Include user-defined `[providers.<name>]` custom entries so ACP |
| 425 | // clients can discover and round-trip the provider names that |
| 426 | // `session/selectModel` now accepts (#1519). |
| 427 | if let Some(custom) = self.config.providers.as_ref().map(|p| &p.custom) { |
| 428 | let mut names = custom.keys().collect::<Vec<_>>(); |
| 429 | names.sort(); |
| 430 | for name in names { |
| 431 | providers.push(json!({ |
| 432 | "id": name, |
| 433 | "displayName": name, |
| 434 | "defaultModel": custom.get(name).and_then(|cfg| cfg.model.clone()) |
| 435 | })); |
| 436 | } |
| 437 | } |
| 438 | |
| 439 | json!({ "providers": providers }) |
| 440 | } |
| 441 | |
| 442 | fn current_model(&self) -> Value { |
| 443 | // Prefer the raw configured provider key so a custom `[providers.<name>]` |
| 444 | // entry round-trips through ACP instead of canonicalizing to "custom". |
| 445 | let provider = match self.config.provider.as_deref() { |
| 446 | Some(name) if !name.trim().is_empty() => name.to_string(), |
| 447 | _ => self.config.api_provider().as_str().to_string(), |
| 448 | }; |
| 449 | json!({ |
| 450 | "provider": provider, |
| 451 | "model": self.model.as_str() |
| 452 | }) |
| 453 | } |
| 454 | |
| 455 | fn select_model(&mut self, params: Value) -> std::result::Result<Value, AcpError> { |
| 456 | let model = params |
| 457 | .get("model") |
| 458 | .and_then(Value::as_str) |
| 459 | .ok_or_else(|| AcpError::invalid_params("model is required"))? |
| 460 | .to_string(); |
| 461 | |
| 462 | if let Some(provider_value) = params.get("provider") { |
| 463 | let provider_name = provider_value |
| 464 | .as_str() |
| 465 | .ok_or_else(|| AcpError::invalid_params("provider must be a string"))?; |
| 466 | // Accept either a built-in provider id/alias or a user-defined |
| 467 | // custom provider name that has a `[providers.<name>]` table. For |
| 468 | // custom providers, preserve the raw key so routing can still find |
| 469 | // the configured base URL / auth / model (#1519); canonicalizing to |
| 470 | // "custom" would lose that table key. |
| 471 | let is_custom = self |
| 472 | .config |
| 473 | .providers |
| 474 | .as_ref() |
| 475 | .and_then(|providers| providers.custom_provider_config(provider_name)) |
| 476 | .is_some(); |
| 477 | if !is_custom && ApiProvider::parse(provider_name).is_none() { |
| 478 | return Err(AcpError::invalid_params(format!( |
| 479 | "unknown provider: {provider_name}" |
| 480 | ))); |
| 481 | } |
| 482 | self.config.provider = Some(provider_name.to_string()); |
| 483 | } |
| 484 | |
| 485 | self.model = model; |
| 486 | Ok(self.current_model()) |
| 487 | } |
| 488 | |
| 489 | /// Validate a `session/prompt` request and append the user turn to history, |
| 490 | /// returning the cloned conversation for the (borrow-free) provider call. |
| 491 | /// |
| 492 | /// This is the `&mut self` half of a prompt turn; the streaming provider |
| 493 | /// call lives in [`AcpServer::open_prompt_stream`] (which borrows `&self` |
| 494 | /// only and returns a `'static` stream) so it can be raced against the |
| 495 | /// reader for cancellation. |
| 496 | fn begin_prompt(&mut self, params: Value) -> std::result::Result<PreparedPrompt, AcpError> { |
| 497 | let session_id = params |
| 498 | .get("sessionId") |
| 499 | .and_then(Value::as_str) |
| 500 | .ok_or_else(|| AcpError::invalid_params("sessionId is required"))? |
| 501 | .to_string(); |
| 502 | let prompt = extract_prompt_text(params.get("prompt")) |
| 503 | .filter(|text| !text.trim().is_empty()) |
| 504 | .ok_or_else(|| AcpError::invalid_params("prompt must include text content"))?; |
| 505 | |
| 506 | let (messages, cwd) = { |
| 507 | let session = self |
| 508 | .sessions |
| 509 | .get_mut(&session_id) |
| 510 | .ok_or_else(|| AcpError::invalid_params("unknown sessionId"))?; |
| 511 | session.messages.push(Message { |
| 512 | role: "user".to_string(), |
| 513 | content: vec![ContentBlock::Text { |
| 514 | text: prompt, |
| 515 | cache_control: None, |
| 516 | }], |
| 517 | }); |
| 518 | (session.messages.clone(), session.cwd.clone()) |
| 519 | }; |
| 520 | |
| 521 | Ok(PreparedPrompt { |
| 522 | session_id, |
| 523 | messages, |
| 524 | cwd, |
| 525 | }) |
| 526 | } |
| 527 | |
| 528 | /// Append a completed assistant turn to session history. A cancelled turn |
| 529 | /// never calls this, so cancelled output does not pollute the transcript. |
| 530 | fn finish_prompt(&mut self, session_id: &str, output: &str) { |
| 531 | if output.is_empty() { |
| 532 | return; |
| 533 | } |
| 534 | if let Some(session) = self.sessions.get_mut(session_id) { |
| 535 | session.messages.push(Message { |
| 536 | role: "assistant".to_string(), |
| 537 | content: vec![ContentBlock::Text { |
| 538 | text: output.to_string(), |
| 539 | cache_control: None, |
| 540 | }], |
| 541 | }); |
| 542 | } |
| 543 | } |
| 544 | |
| 545 | /// Resolve the route, build the streaming request, and open the provider |
| 546 | /// response stream. Borrows `&self` only to read config/model; the returned |
| 547 | /// [`StreamEventBox`] is `'static`, so the caller can race it against the |
| 548 | /// reader without holding any borrow on the server. The cwd guard only needs |
| 549 | /// to cover route resolution and client construction, not stream |
| 550 | /// consumption (ACP exposes no file/shell tools), so it is dropped here. |
| 551 | async fn open_prompt_stream( |
| 552 | &self, |
| 553 | messages: &[Message], |
| 554 | cwd: &PathBuf, |
| 555 | ) -> Result<StreamEventBox> { |
| 556 | let _cwd_guard = ScopedCurrentDir::new(cwd)?; |
| 557 | let last_user_text = messages |
| 558 | .iter() |
| 559 | .rev() |
| 560 | .find_map(|m| { |
| 561 | if m.role == "user" { |
| 562 | m.content.iter().find_map(|b| match b { |
| 563 | ContentBlock::Text { text, .. } => Some(text.as_str()), |
| 564 | _ => None, |
| 565 | }) |
| 566 | } else { |
| 567 | None |
| 568 | } |
| 569 | }) |
| 570 | .unwrap_or(""); |
| 571 | let route = |
| 572 | crate::resolve_cli_auto_route(&self.config, &self.model, last_user_text).await?; |
| 573 | let execution_config = crate::config_for_cli_route(&self.config, &route); |
| 574 | let client = DeepSeekClient::new(&execution_config)?; |
| 575 | let reasoning_effort = route |
| 576 | .reasoning_effort |
| 577 | .and_then(|effort| { |
| 578 | effort.api_value_for_route( |
| 579 | execution_config.api_provider(), |
| 580 | &execution_config.deepseek_base_url(), |
| 581 | &route.model, |
| 582 | ) |
| 583 | }) |
| 584 | .map(str::to_string); |
| 585 | |
| 586 | let request = MessageRequest { |
| 587 | model: route.model, |
| 588 | messages: messages.to_vec(), |
| 589 | max_tokens: 4096, |
| 590 | system: Some(SystemPrompt::Text( |
| 591 | "You are a coding assistant inside an ACP-compatible editor. Give concise, actionable responses.".to_string(), |
| 592 | )), |
| 593 | tools: None, |
| 594 | tool_choice: None, |
| 595 | metadata: None, |
| 596 | thinking: None, |
| 597 | reasoning_effort, |
| 598 | stream: Some(true), |
| 599 | temperature: Some(0.2), |
| 600 | top_p: Some(0.9), |
| 601 | }; |
| 602 | |
| 603 | client.create_message_stream(request).await |
| 604 | } |
| 605 | } |
| 606 | |
| 607 | struct ScopedCurrentDir { |
| 608 | prior: PathBuf, |
| 609 | } |
| 610 | |
| 611 | impl ScopedCurrentDir { |
| 612 | fn new(cwd: &PathBuf) -> Result<Self> { |
| 613 | let prior = std::env::current_dir()?; |
| 614 | if cwd.as_os_str().is_empty() { |
| 615 | return Ok(Self { prior }); |
| 616 | } |
| 617 | std::env::set_current_dir(cwd) |
| 618 | .map_err(|err| anyhow!("failed to enter ACP session cwd {}: {err}", cwd.display()))?; |
| 619 | Ok(Self { prior }) |
| 620 | } |
| 621 | } |
| 622 | |
| 623 | impl Drop for ScopedCurrentDir { |
| 624 | fn drop(&mut self) { |
| 625 | let _ = std::env::set_current_dir(&self.prior); |
| 626 | } |
| 627 | } |
| 628 | |
| 629 | impl AcpError { |
| 630 | fn invalid_params(message: impl Into<String>) -> Self { |
| 631 | Self { |
| 632 | code: -32602, |
| 633 | message: message.into(), |
| 634 | } |
| 635 | } |
| 636 | |
| 637 | fn method_not_found(method: &str) -> Self { |
| 638 | Self { |
| 639 | code: -32601, |
| 640 | message: format!("method not found: {method}"), |
| 641 | } |
| 642 | } |
| 643 | } |
| 644 | |
| 645 | fn initialize_result(client_protocol_version: Option<u64>, config: &Config) -> Value { |
| 646 | json!({ |
| 647 | "protocolVersion": client_protocol_version |
| 648 | .map(|version| version.min(ACP_PROTOCOL_VERSION)) |
| 649 | .unwrap_or(ACP_PROTOCOL_VERSION), |
| 650 | "agentCapabilities": { |
| 651 | "loadSession": false, |
| 652 | "modelSelection": true, |
| 653 | "promptCapabilities": { |
| 654 | "image": false, |
| 655 | "audio": false, |
| 656 | "embeddedContext": true |
| 657 | }, |
| 658 | "mcpCapabilities": { |
| 659 | "http": false, |
| 660 | "sse": false |
| 661 | }, |
| 662 | "sessionCapabilities": {} |
| 663 | }, |
| 664 | "agentInfo": { |
| 665 | "name": "codewhale", |
| 666 | "title": "codewhale", |
| 667 | "version": env!("CARGO_PKG_VERSION") |
| 668 | }, |
| 669 | "authMethods": acp_auth_methods(config) |
| 670 | }) |
| 671 | } |
| 672 | |
| 673 | fn acp_auth_methods(config: &Config) -> Value { |
| 674 | let provider = config.api_provider().as_str(); |
| 675 | json!([ |
| 676 | { |
| 677 | "id": "codewhale-terminal-auth", |
| 678 | "name": "Set Codewhale API key", |
| 679 | "description": format!("Run Codewhale's terminal credential setup for the {provider} provider."), |
| 680 | "type": "terminal", |
| 681 | "args": ["auth", "set", "--provider", provider], |
| 682 | "env": {} |
| 683 | } |
| 684 | ]) |
| 685 | } |
| 686 | |
| 687 | fn extract_prompt_text(prompt: Option<&Value>) -> Option<String> { |
| 688 | match prompt? { |
| 689 | Value::String(text) => Some(text.clone()), |
| 690 | Value::Array(blocks) => { |
| 691 | let parts = blocks |
| 692 | .iter() |
| 693 | .filter_map(content_block_text) |
| 694 | .collect::<Vec<_>>(); |
| 695 | (!parts.is_empty()).then(|| parts.join("\n\n")) |
| 696 | } |
| 697 | _ => None, |
| 698 | } |
| 699 | } |
| 700 | |
| 701 | fn content_block_text(block: &Value) -> Option<String> { |
| 702 | match block.get("type").and_then(Value::as_str)? { |
| 703 | "text" => block |
| 704 | .get("text") |
| 705 | .and_then(Value::as_str) |
| 706 | .map(str::to_string), |
| 707 | "resource" => resource_text(block), |
| 708 | "resource_link" | "resourceLink" => resource_link_text(block), |
| 709 | _ => None, |
| 710 | } |
| 711 | } |
| 712 | |
| 713 | fn resource_text(block: &Value) -> Option<String> { |
| 714 | let resource = block.get("resource").unwrap_or(block); |
| 715 | if let Some(text) = resource.get("text").and_then(Value::as_str) { |
| 716 | return Some(text.to_string()); |
| 717 | } |
| 718 | resource_link_text(resource) |
| 719 | } |
| 720 | |
| 721 | fn resource_link_text(block: &Value) -> Option<String> { |
| 722 | let uri = block |
| 723 | .get("uri") |
| 724 | .or_else(|| block.pointer("/resource/uri")) |
| 725 | .and_then(Value::as_str)?; |
| 726 | Some(format!("@{uri}")) |
| 727 | } |
| 728 | |
| 729 | async fn write_session_update<W>(writer: &mut W, session_id: &str, text: String) -> Result<()> |
| 730 | where |
| 731 | W: AsyncWrite + Unpin, |
| 732 | { |
| 733 | let notification = json!({ |
| 734 | "jsonrpc": "2.0", |
| 735 | "method": "session/update", |
| 736 | "params": { |
| 737 | "sessionId": session_id, |
| 738 | "update": { |
| 739 | "sessionUpdate": "agent_message_chunk", |
| 740 | "content": { |
| 741 | "type": "text", |
| 742 | "text": text |
| 743 | } |
| 744 | } |
| 745 | } |
| 746 | }); |
| 747 | write_json_line(writer, notification).await |
| 748 | } |
| 749 | |
| 750 | async fn write_jsonrpc_result<W>(writer: &mut W, id: Value, result: Value) -> Result<()> |
| 751 | where |
| 752 | W: AsyncWrite + Unpin, |
| 753 | { |
| 754 | write_json_line( |
| 755 | writer, |
| 756 | json!({ |
| 757 | "jsonrpc": "2.0", |
| 758 | "id": id, |
| 759 | "result": result |
| 760 | }), |
| 761 | ) |
| 762 | .await |
| 763 | } |
| 764 | |
| 765 | async fn write_jsonrpc_error<W>( |
| 766 | writer: &mut W, |
| 767 | id: Option<Value>, |
| 768 | code: i32, |
| 769 | message: impl Into<String>, |
| 770 | ) -> Result<()> |
| 771 | where |
| 772 | W: AsyncWrite + Unpin, |
| 773 | { |
| 774 | write_json_line( |
| 775 | writer, |
| 776 | json!({ |
| 777 | "jsonrpc": "2.0", |
| 778 | "id": id, |
| 779 | "error": { |
| 780 | "code": code, |
| 781 | "message": message.into() |
| 782 | } |
| 783 | }), |
| 784 | ) |
| 785 | .await |
| 786 | } |
| 787 | |
| 788 | async fn write_json_line<W>(writer: &mut W, value: Value) -> Result<()> |
| 789 | where |
| 790 | W: AsyncWrite + Unpin, |
| 791 | { |
| 792 | writer.write_all(value.to_string().as_bytes()).await?; |
| 793 | writer.write_all(b"\n").await?; |
| 794 | writer.flush().await?; |
| 795 | Ok(()) |
| 796 | } |
| 797 | |
| 798 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] |
| 799 | enum JsonRpcResponseIdPolicy { |
| 800 | /// JSON-RPC's normal contract: echo the request id without changing type. |
| 801 | Preserve, |
| 802 | /// Zed's ACP client currently decodes response ids as strings even when it |
| 803 | /// sent a number. Keep this narrow compatibility mode client-identified. |
| 804 | StringifyNumeric, |
| 805 | } |
| 806 | |
| 807 | impl JsonRpcResponseIdPolicy { |
| 808 | fn from_initialize_params(params: &Value) -> Self { |
| 809 | let client_name = params |
| 810 | .pointer("/clientInfo/name") |
| 811 | .and_then(Value::as_str) |
| 812 | .unwrap_or_default(); |
| 813 | if client_name.eq_ignore_ascii_case("zed") { |
| 814 | Self::StringifyNumeric |
| 815 | } else { |
| 816 | Self::Preserve |
| 817 | } |
| 818 | } |
| 819 | |
| 820 | fn response_id(self, id: Value) -> Value { |
| 821 | match (self, id) { |
| 822 | (Self::StringifyNumeric, Value::Number(number)) => Value::String(number.to_string()), |
| 823 | (_, id) => id, |
| 824 | } |
| 825 | } |
| 826 | } |
| 827 | |
| 828 | #[cfg(test)] |
| 829 | mod tests { |
| 830 | use super::*; |
| 831 | |
| 832 | #[test] |
| 833 | fn initialize_advertises_baseline_acp_agent() { |
| 834 | let result = initialize_result(Some(1), &Config::default()); |
| 835 | |
| 836 | assert_eq!(result["protocolVersion"], 1); |
| 837 | assert_eq!(result["agentInfo"]["name"], "codewhale"); |
| 838 | assert_eq!(result["agentCapabilities"]["loadSession"], false); |
| 839 | assert_eq!( |
| 840 | result["agentCapabilities"]["promptCapabilities"]["embeddedContext"], |
| 841 | true |
| 842 | ); |
| 843 | assert_eq!(result["authMethods"][0]["type"], "terminal"); |
| 844 | assert_eq!( |
| 845 | result["authMethods"][0]["args"], |
| 846 | json!(["auth", "set", "--provider", "deepseek"]) |
| 847 | ); |
| 848 | } |
| 849 | |
| 850 | #[test] |
| 851 | fn initialize_advertises_model_selection_capability() { |
| 852 | let result = initialize_result(Some(1), &Config::default()); |
| 853 | |
| 854 | assert_eq!(result["agentCapabilities"]["modelSelection"], true); |
| 855 | } |
| 856 | |
| 857 | #[test] |
| 858 | fn list_providers_returns_provider_set() { |
| 859 | let server = AcpServer::new( |
| 860 | Config::default(), |
| 861 | "deepseek-chat".into(), |
| 862 | PathBuf::from("/tmp"), |
| 863 | ); |
| 864 | let result = server.list_providers(); |
| 865 | let providers = result["providers"].as_array().expect("providers array"); |
| 866 | |
| 867 | assert!(!providers.is_empty()); |
| 868 | assert!( |
| 869 | providers |
| 870 | .iter() |
| 871 | .any(|provider| provider["id"] == "deepseek") |
| 872 | ); |
| 873 | } |
| 874 | |
| 875 | #[test] |
| 876 | fn current_model_reflects_constructor_default() { |
| 877 | let config = Config::default(); |
| 878 | let expected_provider = config.api_provider().as_str(); |
| 879 | let server = AcpServer::new(config, "deepseek-reasoner".into(), PathBuf::from("/tmp")); |
| 880 | let result = server.current_model(); |
| 881 | |
| 882 | assert_eq!(result["provider"], expected_provider); |
| 883 | assert_eq!(result["model"], "deepseek-reasoner"); |
| 884 | } |
| 885 | |
| 886 | #[test] |
| 887 | fn select_model_updates_active_selection() { |
| 888 | let mut server = AcpServer::new( |
| 889 | Config::default(), |
| 890 | "deepseek-chat".into(), |
| 891 | PathBuf::from("/tmp"), |
| 892 | ); |
| 893 | |
| 894 | let result = server |
| 895 | .select_model(json!({ "provider": "openai", "model": "gpt-4o" })) |
| 896 | .expect("select model"); |
| 897 | |
| 898 | assert_eq!(result["provider"], "openai"); |
| 899 | assert_eq!(result["model"], "gpt-4o"); |
| 900 | assert_eq!(server.current_model()["provider"], "openai"); |
| 901 | assert_eq!(server.current_model()["model"], "gpt-4o"); |
| 902 | } |
| 903 | |
| 904 | #[test] |
| 905 | fn select_model_rejects_unknown_provider() { |
| 906 | let mut server = AcpServer::new( |
| 907 | Config::default(), |
| 908 | "deepseek-chat".into(), |
| 909 | PathBuf::from("/tmp"), |
| 910 | ); |
| 911 | let before = server.current_model(); |
| 912 | |
| 913 | let err = server |
| 914 | .select_model(json!({ "provider": "unknown-provider", "model": "gpt-4o" })) |
| 915 | .expect_err("unknown provider rejected"); |
| 916 | |
| 917 | assert_eq!(err.code, -32602); |
| 918 | assert_eq!(server.current_model(), before); |
| 919 | } |
| 920 | |
| 921 | #[test] |
| 922 | fn select_model_rejects_missing_model() { |
| 923 | let mut server = AcpServer::new( |
| 924 | Config::default(), |
| 925 | "deepseek-chat".into(), |
| 926 | PathBuf::from("/tmp"), |
| 927 | ); |
| 928 | |
| 929 | let err = server |
| 930 | .select_model(json!({ "provider": "openai" })) |
| 931 | .expect_err("missing model rejected"); |
| 932 | |
| 933 | assert_eq!(err.code, -32602); |
| 934 | } |
| 935 | |
| 936 | #[test] |
| 937 | fn extract_prompt_text_accepts_text_and_resource_blocks() { |
| 938 | let prompt = json!([ |
| 939 | { "type": "text", "text": "Review this file" }, |
| 940 | { |
| 941 | "type": "resource", |
| 942 | "resource": { |
| 943 | "uri": "file:///tmp/app.rs", |
| 944 | "mimeType": "text/rust", |
| 945 | "text": "fn main() {}" |
| 946 | } |
| 947 | }, |
| 948 | { "type": "resource_link", "uri": "file:///tmp/lib.rs" } |
| 949 | ]); |
| 950 | |
| 951 | let text = extract_prompt_text(Some(&prompt)).expect("prompt text"); |
| 952 | |
| 953 | assert!(text.contains("Review this file")); |
| 954 | assert!(text.contains("fn main() {}")); |
| 955 | assert!(text.contains("@file:///tmp/lib.rs")); |
| 956 | } |
| 957 | |
| 958 | #[tokio::test] |
| 959 | async fn session_update_is_protocol_clean_single_line_json() { |
| 960 | let mut out = Vec::new(); |
| 961 | |
| 962 | write_session_update(&mut out, "sess_1", "hello\nworld".to_string()) |
| 963 | .await |
| 964 | .expect("write update"); |
| 965 | |
| 966 | let line = String::from_utf8(out).expect("utf8"); |
| 967 | assert_eq!(line.lines().count(), 1); |
| 968 | let value: Value = serde_json::from_str(line.trim()).expect("json"); |
| 969 | assert_eq!(value["method"], "session/update"); |
| 970 | assert_eq!(value["params"]["sessionId"], "sess_1"); |
| 971 | assert_eq!(value["params"]["update"]["content"]["text"], "hello\nworld"); |
| 972 | } |
| 973 | |
| 974 | #[tokio::test] |
| 975 | async fn jsonrpc_result_preserves_numeric_ids_for_avante_acp() { |
| 976 | let mut out = Vec::new(); |
| 977 | |
| 978 | let params = json!({ |
| 979 | "protocolVersion": 1, |
| 980 | "clientCapabilities": {} |
| 981 | }); |
| 982 | let id = JsonRpcResponseIdPolicy::from_initialize_params(¶ms).response_id(json!(1)); |
| 983 | write_jsonrpc_result(&mut out, id, json!({"ok": true})) |
| 984 | .await |
| 985 | .expect("write result"); |
| 986 | |
| 987 | let line = String::from_utf8(out).expect("utf8"); |
| 988 | let value: Value = serde_json::from_str(line.trim()).expect("json"); |
| 989 | // Numeric ID must stay numeric — avante.nvim's Lua client uses |
| 990 | // strict table keys (callbacks[1] ≠ callbacks["1"]). |
| 991 | assert!( |
| 992 | value["id"].is_number(), |
| 993 | "numeric id must stay numeric, got {:?}", |
| 994 | value["id"] |
| 995 | ); |
| 996 | assert_eq!(value["result"], json!({"ok": true})); |
| 997 | } |
| 998 | |
| 999 | #[tokio::test] |
| 1000 | async fn jsonrpc_result_stringifies_numeric_ids_for_zed_acp() { |
| 1001 | let mut out = Vec::new(); |
| 1002 | |
| 1003 | let params = json!({ |
| 1004 | "protocolVersion": 1, |
| 1005 | "clientCapabilities": {}, |
| 1006 | "clientInfo": { |
| 1007 | "name": "zed", |
| 1008 | "version": "1.2.6" |
| 1009 | } |
| 1010 | }); |
| 1011 | let id = JsonRpcResponseIdPolicy::from_initialize_params(¶ms).response_id(json!(1)); |
| 1012 | write_jsonrpc_result(&mut out, id, json!({"ok": true})) |
| 1013 | .await |
| 1014 | .expect("write result"); |
| 1015 | |
| 1016 | let line = String::from_utf8(out).expect("utf8"); |
| 1017 | let value: Value = serde_json::from_str(line.trim()).expect("json"); |
| 1018 | assert_eq!(value["id"], "1"); |
| 1019 | assert_eq!(value["result"], json!({"ok": true})); |
| 1020 | } |
| 1021 | |
| 1022 | #[tokio::test] |
| 1023 | async fn jsonrpc_error_keeps_absent_id_null() { |
| 1024 | let mut out = Vec::new(); |
| 1025 | |
| 1026 | write_jsonrpc_error(&mut out, None, -32700, "invalid json") |
| 1027 | .await |
| 1028 | .expect("write error"); |
| 1029 | |
| 1030 | let line = String::from_utf8(out).expect("utf8"); |
| 1031 | let value: Value = serde_json::from_str(line.trim()).expect("json"); |
| 1032 | assert_eq!(value["id"], Value::Null); |
| 1033 | assert_eq!(value["error"]["code"], -32700); |
| 1034 | } |
| 1035 | |
| 1036 | #[test] |
| 1037 | fn new_session_starts_with_empty_messages() { |
| 1038 | let mut server = AcpServer::new( |
| 1039 | Config::default(), |
| 1040 | "test-model".to_string(), |
| 1041 | PathBuf::from("/tmp"), |
| 1042 | ); |
| 1043 | let result = server |
| 1044 | .new_session(json!({ "cwd": "/tmp" })) |
| 1045 | .expect("new session"); |
| 1046 | let session_id = result["sessionId"].as_str().expect("session id"); |
| 1047 | let session = server.sessions.get(session_id).expect("session exists"); |
| 1048 | assert!(session.messages.is_empty()); |
| 1049 | } |
| 1050 | |
| 1051 | #[test] |
| 1052 | fn prompt_appends_user_and_assistant_messages_to_history() { |
| 1053 | let mut server = AcpServer::new( |
| 1054 | Config::default(), |
| 1055 | "test-model".to_string(), |
| 1056 | PathBuf::from("/tmp"), |
| 1057 | ); |
| 1058 | let result = server |
| 1059 | .new_session(json!({ "cwd": "/tmp" })) |
| 1060 | .expect("new session"); |
| 1061 | let session_id = result["sessionId"].as_str().unwrap().to_string(); |
| 1062 | |
| 1063 | // Simulate adding a user message (same logic as prompt() but without LLM call) |
| 1064 | { |
| 1065 | let session = server.sessions.get_mut(&session_id).unwrap(); |
| 1066 | session.messages.push(Message { |
| 1067 | role: "user".to_string(), |
| 1068 | content: vec![ContentBlock::Text { |
| 1069 | text: "1+1".to_string(), |
| 1070 | cache_control: None, |
| 1071 | }], |
| 1072 | }); |
| 1073 | } |
| 1074 | |
| 1075 | // Simulate assistant response |
| 1076 | { |
| 1077 | let session = server.sessions.get_mut(&session_id).unwrap(); |
| 1078 | session.messages.push(Message { |
| 1079 | role: "assistant".to_string(), |
| 1080 | content: vec![ContentBlock::Text { |
| 1081 | text: "2".to_string(), |
| 1082 | cache_control: None, |
| 1083 | }], |
| 1084 | }); |
| 1085 | } |
| 1086 | |
| 1087 | // Second user message |
| 1088 | { |
| 1089 | let session = server.sessions.get_mut(&session_id).unwrap(); |
| 1090 | session.messages.push(Message { |
| 1091 | role: "user".to_string(), |
| 1092 | content: vec![ContentBlock::Text { |
| 1093 | text: "add one more".to_string(), |
| 1094 | cache_control: None, |
| 1095 | }], |
| 1096 | }); |
| 1097 | } |
| 1098 | |
| 1099 | // Verify full conversation history |
| 1100 | let session = server.sessions.get(&session_id).unwrap(); |
| 1101 | assert_eq!(session.messages.len(), 3); |
| 1102 | assert_eq!(session.messages[0].role, "user"); |
| 1103 | assert_eq!(session.messages[1].role, "assistant"); |
| 1104 | assert_eq!(session.messages[2].role, "user"); |
| 1105 | |
| 1106 | // Verify text content |
| 1107 | assert_eq!( |
| 1108 | match &session.messages[0].content[0] { |
| 1109 | ContentBlock::Text { text, .. } => text.clone(), |
| 1110 | _ => String::new(), |
| 1111 | }, |
| 1112 | "1+1" |
| 1113 | ); |
| 1114 | assert_eq!( |
| 1115 | match &session.messages[1].content[0] { |
| 1116 | ContentBlock::Text { text, .. } => text.clone(), |
| 1117 | _ => String::new(), |
| 1118 | }, |
| 1119 | "2" |
| 1120 | ); |
| 1121 | assert_eq!( |
| 1122 | match &session.messages[2].content[0] { |
| 1123 | ContentBlock::Text { text, .. } => text.clone(), |
| 1124 | _ => String::new(), |
| 1125 | }, |
| 1126 | "add one more" |
| 1127 | ); |
| 1128 | } |
| 1129 | |
| 1130 | fn lines_from(input: &'static str) -> Lines<BufReader<&'static [u8]>> { |
| 1131 | BufReader::new(input.as_bytes()).lines() |
| 1132 | } |
| 1133 | |
| 1134 | fn text_delta(text: &str) -> StreamEvent { |
| 1135 | StreamEvent::ContentBlockDelta { |
| 1136 | index: 0, |
| 1137 | delta: Delta::TextDelta { |
| 1138 | text: text.to_string(), |
| 1139 | }, |
| 1140 | } |
| 1141 | } |
| 1142 | |
| 1143 | /// A stream that yields the given events immediately, then ends. |
| 1144 | fn ready_stream(events: Vec<StreamEvent>) -> StreamEventBox { |
| 1145 | Box::pin(futures_util::stream::iter( |
| 1146 | events.into_iter().map(Ok::<_, anyhow::Error>), |
| 1147 | )) |
| 1148 | } |
| 1149 | |
| 1150 | /// A stream that never yields, so a concurrent cancel always wins. |
| 1151 | fn pending_stream() -> StreamEventBox { |
| 1152 | Box::pin(futures_util::stream::pending::<Result<StreamEvent>>()) |
| 1153 | } |
| 1154 | |
| 1155 | /// A stream that yields `events` immediately, then emits `message_stop` |
| 1156 | /// after a short delay — long enough that an already-buffered reader line is |
| 1157 | /// processed first, making the ordering deterministic in tests. |
| 1158 | fn events_then_delayed_stop(events: Vec<StreamEvent>) -> StreamEventBox { |
| 1159 | let head = futures_util::stream::iter(events.into_iter().map(Ok::<_, anyhow::Error>)); |
| 1160 | let tail = futures_util::stream::once(async { |
| 1161 | tokio::time::sleep(std::time::Duration::from_millis(20)).await; |
| 1162 | Ok(StreamEvent::MessageStop) |
| 1163 | }); |
| 1164 | Box::pin(head.chain(tail)) |
| 1165 | } |
| 1166 | |
| 1167 | fn parse_lines(out: Vec<u8>) -> Vec<Value> { |
| 1168 | String::from_utf8(out) |
| 1169 | .expect("utf8") |
| 1170 | .lines() |
| 1171 | .filter(|line| !line.trim().is_empty()) |
| 1172 | .map(|line| serde_json::from_str(line).expect("json")) |
| 1173 | .collect() |
| 1174 | } |
| 1175 | |
| 1176 | #[tokio::test] |
| 1177 | async fn drive_prompt_streams_each_delta_as_a_chunk_then_completes() { |
| 1178 | let stream = ready_stream(vec![ |
| 1179 | text_delta("hello"), |
| 1180 | text_delta(" world"), |
| 1181 | StreamEvent::MessageStop, |
| 1182 | ]); |
| 1183 | let mut reader = lines_from(""); |
| 1184 | let mut out = Vec::new(); |
| 1185 | |
| 1186 | let outcome = drive_prompt_stream( |
| 1187 | stream, |
| 1188 | "sess_1", |
| 1189 | JsonRpcResponseIdPolicy::Preserve, |
| 1190 | &mut reader, |
| 1191 | &mut out, |
| 1192 | ) |
| 1193 | .await |
| 1194 | .expect("driver ok"); |
| 1195 | |
| 1196 | // Full text is accumulated for history... |
| 1197 | assert_eq!(outcome, PromptOutcome::Completed("hello world".to_string())); |
| 1198 | // ...and each delta was emitted as its own session/update chunk. |
| 1199 | let updates = parse_lines(out); |
| 1200 | assert_eq!(updates.len(), 2); |
| 1201 | assert!(updates.iter().all(|u| u["method"] == "session/update")); |
| 1202 | assert_eq!(updates[0]["params"]["update"]["content"]["text"], "hello"); |
| 1203 | assert_eq!(updates[1]["params"]["update"]["content"]["text"], " world"); |
| 1204 | } |
| 1205 | |
| 1206 | #[tokio::test] |
| 1207 | async fn drive_prompt_cancels_when_matching_cancel_arrives() { |
| 1208 | // A provider stream that never finishes within the test. |
| 1209 | let stream = pending_stream(); |
| 1210 | let mut reader = lines_from( |
| 1211 | r#"{"jsonrpc":"2.0","method":"session/cancel","params":{"sessionId":"sess_1"}}"#, |
| 1212 | ); |
| 1213 | let mut out = Vec::new(); |
| 1214 | |
| 1215 | let outcome = drive_prompt_stream( |
| 1216 | stream, |
| 1217 | "sess_1", |
| 1218 | JsonRpcResponseIdPolicy::Preserve, |
| 1219 | &mut reader, |
| 1220 | &mut out, |
| 1221 | ) |
| 1222 | .await |
| 1223 | .expect("driver ok"); |
| 1224 | |
| 1225 | assert_eq!(outcome, PromptOutcome::Cancelled); |
| 1226 | // Notification-form cancel (no id) is acknowledged by acting, not writing. |
| 1227 | assert!(out.is_empty()); |
| 1228 | } |
| 1229 | |
| 1230 | #[tokio::test] |
| 1231 | async fn drive_prompt_ignores_cancel_for_a_different_session() { |
| 1232 | // The unrelated cancel line is buffered and ready; the delayed stop makes |
| 1233 | // it process first, proving it does not abort the turn. |
| 1234 | let stream = events_then_delayed_stop(vec![text_delta("kept")]); |
| 1235 | let mut reader = lines_from( |
| 1236 | r#"{"jsonrpc":"2.0","id":7,"method":"session/cancel","params":{"sessionId":"other"}}"#, |
| 1237 | ); |
| 1238 | let mut out = Vec::new(); |
| 1239 | |
| 1240 | let outcome = drive_prompt_stream( |
| 1241 | stream, |
| 1242 | "sess_1", |
| 1243 | JsonRpcResponseIdPolicy::StringifyNumeric, |
| 1244 | &mut reader, |
| 1245 | &mut out, |
| 1246 | ) |
| 1247 | .await |
| 1248 | .expect("driver ok"); |
| 1249 | |
| 1250 | assert_eq!(outcome, PromptOutcome::Completed("kept".to_string())); |
| 1251 | // The other-session cancel carried an id, so it was acknowledged with null. |
| 1252 | let lines = parse_lines(out); |
| 1253 | assert!( |
| 1254 | lines |
| 1255 | .iter() |
| 1256 | .any(|v| v["id"] == "7" && v["result"] == Value::Null), |
| 1257 | "expected a null ack for the other-session cancel, got {lines:?}" |
| 1258 | ); |
| 1259 | } |
| 1260 | |
| 1261 | #[tokio::test] |
| 1262 | async fn drive_prompt_rejects_a_concurrent_request_but_keeps_running() { |
| 1263 | let stream = events_then_delayed_stop(vec![text_delta("done")]); |
| 1264 | // A non-cancel request arrives mid-turn. |
| 1265 | let mut reader = |
| 1266 | lines_from(r#"{"jsonrpc":"2.0","id":9,"method":"session/new","params":{}}"#); |
| 1267 | let mut out = Vec::new(); |
| 1268 | |
| 1269 | let outcome = drive_prompt_stream( |
| 1270 | stream, |
| 1271 | "sess_1", |
| 1272 | JsonRpcResponseIdPolicy::StringifyNumeric, |
| 1273 | &mut reader, |
| 1274 | &mut out, |
| 1275 | ) |
| 1276 | .await |
| 1277 | .expect("driver ok"); |
| 1278 | |
| 1279 | assert_eq!(outcome, PromptOutcome::Completed("done".to_string())); |
| 1280 | let lines = parse_lines(out); |
| 1281 | assert!( |
| 1282 | lines |
| 1283 | .iter() |
| 1284 | .any(|v| v["id"] == "9" && v["error"]["code"] == -32603), |
| 1285 | "expected a prompt-in-progress error for the concurrent request, got {lines:?}" |
| 1286 | ); |
| 1287 | } |
| 1288 | |
| 1289 | #[test] |
| 1290 | fn different_sessions_have_independent_history() { |
| 1291 | let mut server = AcpServer::new( |
| 1292 | Config::default(), |
| 1293 | "test-model".to_string(), |
| 1294 | PathBuf::from("/tmp"), |
| 1295 | ); |
| 1296 | let result1 = server |
| 1297 | .new_session(json!({ "cwd": "/tmp" })) |
| 1298 | .expect("session 1"); |
| 1299 | let result2 = server |
| 1300 | .new_session(json!({ "cwd": "/tmp" })) |
| 1301 | .expect("session 2"); |
| 1302 | let sid1 = result1["sessionId"].as_str().unwrap().to_string(); |
| 1303 | let sid2 = result2["sessionId"].as_str().unwrap().to_string(); |
| 1304 | |
| 1305 | // Add messages to session 1 |
| 1306 | { |
| 1307 | let session = server.sessions.get_mut(&sid1).unwrap(); |
| 1308 | session.messages.push(Message { |
| 1309 | role: "user".to_string(), |
| 1310 | content: vec![ContentBlock::Text { |
| 1311 | text: "hello".to_string(), |
| 1312 | cache_control: None, |
| 1313 | }], |
| 1314 | }); |
| 1315 | } |
| 1316 | |
| 1317 | // Session 2 should remain empty |
| 1318 | let session2 = server.sessions.get(&sid2).unwrap(); |
| 1319 | assert!(session2.messages.is_empty()); |
| 1320 | |
| 1321 | // Session 1 should have the message |
| 1322 | let session1 = server.sessions.get(&sid1).unwrap(); |
| 1323 | assert_eq!(session1.messages.len(), 1); |
| 1324 | } |
| 1325 | } |
| 1326 |