diff --git a/Cargo.toml b/Cargo.toml index 9fbb98a6..a61e9a34 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -43,7 +43,7 @@ json5 = "0.4" schemars = "1.0" anyhow = "1.0" clap = { version = "4.5", features = ["derive"] } -agent-client-protocol = { version = "=2.0.0", features = ["unstable_session_fork"] } +agent-client-protocol = { version = "=2.0.0", features = ["unstable_elicitation", "unstable_session_fork"] } clap_complete = "4.5" ignore = "0.4" copypasta = "0.10" diff --git a/_docs/acp.mdx b/_docs/acp.mdx index c427019d..a46b5c62 100644 --- a/_docs/acp.mdx +++ b/_docs/acp.mdx @@ -40,31 +40,42 @@ Configure a provider in the Crabcode TUI with `/connect` before starting an ACP ## Capability Matrix -| Area | Supported behavior | Current limitations | Planned follow-up | +All Crabcode-side capabilities in this matrix are implemented. Conditions in the last column are editor, model, provider, or protocol requirements rather than incomplete server behavior. + +| Area | Full behavior | Status | Runtime or protocol requirements | | --- | --- | --- | --- | -| Transport | JSON-RPC over stdio through `crabcode acp`; clean stdin EOF shutdown. | stdout must remain protocol-only. | Add protocol-level subprocess integration coverage. | -| Sessions | Create, list, load, resume, close, and fork persisted root sessions. | Replay message IDs are deterministic per load, not durable message IDs. | Persist stable message IDs across streaming snapshots and reloads. | -| Prompts | Text, embedded text resources, and PNG, JPEG, GIF, or WebP image attachments; assistant text and reasoning stream back to the editor. | Images require an image-capable selected model; audio prompt blocks are unsupported. | Store ACP attachments persistently and add audio input support. | -| Modes and models | Visible primary Crabcode agents, selectable `/models` catalog entries, and model-supported reasoning-effort selectors are available as ACP options; changes are session-local. | Reasoning options depend on the selected model catalog metadata. | Add richer provider-specific reasoning configuration. | -| Tools | Tool calls and completed or failed results stream with ACP tool kinds, titles, raw input, and preview output. | No normalized locations, diffs, full outputs, or result images yet. | Preserve structured tool results, locations, patches, and image content in runtime events. | -| Permissions | Existing Crabcode permission prompts are forwarded with allow once, always allow, and reject choices. | Permission requests use a generated ACP call ID because the current internal prompt lacks the originating tool-call ID. | Carry the real tool-call ID and edit patch metadata through permission preflight. | -| Cancellation | `session/cancel` cancels the active turn and keeps the session reusable. | Provider stop reasons are currently reduced to normal completion, cancellation, or a safe failure. | Preserve output-limit and refusal stop reasons from the model runtime. | -| Commands and skills | Session updates publish available commands: project custom slash commands, workspace skills, plus built-in `/skills` and `/mcp`. Leading `/…` prompts expand through the same command and skill templates before the model turn. | Built-in TUI commands such as `/compact` are not ACP-available commands; unknown `/…` lines pass through as plain text. | Add more built-in commands (for example `/compact`) and richer command input schemas. | -| MCP | Project MCP from Crabcode config runs as usual. Editors may also pass MCP servers on `session/new`; those servers are merged into the session config (stdio, HTTP, and SSE). | HTTP and SSE client MCP are advertised; stdio client MCP is accepted and merged even though it is not a separate advertised capability flag. | Surface richer MCP connection status and OAuth for remote client servers. | -| Terminals | — | ACP terminal embedding is not implemented. | Add a client-terminal adapter for long-running shell sessions. | -| Questions | — | Interactive question prompts from the agent are not forwarded over ACP; the runtime skips them rather than blocking the editor. | Map Crabcode questions to ACP permission-style or dedicated question requests. | -| Usage | Estimated context-window usage is emitted when the selected model exposes a context limit. | Provider-authoritative token and cost accounting is not complete; usage is omitted when no context limit is known. | Retain provider input, output, cache, context, and cost data for authoritative usage updates. | +| Transport | JSON-RPC over stdio through `crabcode acp`, protocol-only stdout, clean stdin EOF shutdown, and subprocess initialize/response coverage. | Full | The subprocess wrapper must not write banners or logs to stdout. | +| Sessions | Create, cursor-list, load, resume, close, delete, and fork all persisted sessions, including child sessions. Lists include Crabcode parent/root IDs in ACP `_meta`; forks preserve the source title, regenerate message IDs, copy attachments independently, and publish commands for the new session. Delete removes persisted history and managed attachments. | Full | ACP has no standard nested-session tree field, so hierarchy is exposed through the `crabcode` metadata extension while the standard list remains flat. | +| Prompts | Text, embedded resources, PNG/JPEG/GIF/WebP images, and WAV/MP3 audio; assistant text and reasoning stream back to the editor. Attachments use private per-session storage, survive load/resume, copy independently on fork, delete with persisted sessions, and readable legacy paths migrate automatically on load. | Full | The selected model route must advertise the matching input modality. Audio uses verified OpenAI-compatible Chat Completions `input_audio`; unsupported provider transports return a clear error instead of dropping media. | +| Modes and models | Visible primary agents, selectable model catalog entries, and supported reasoning-effort values are session-local ACP configuration options. | Full | Available reasoning values follow the selected model's catalog capability. | +| Tools | Pending and completed/failed tool calls include ACP kinds, titles, raw input/output, full text plus bounded previews, normalized locations, native editor images/audio/resources, annotations, metadata, and full-file diffs. The model receives the structured textual/raw representation and supported image results. Unknown future MCP blocks are preserved in raw output and rendered as readable JSON text instead of being dropped. | Full | A future content type can only be native when ACP defines a matching content block; the lossless text/raw fallback remains available otherwise. | +| Permissions | Permission requests carry the originating tool-call ID, raw input, normalized locations, and preflight full-file diffs for `edit`, `write`, `write_files`, and multi-file `apply_patch`, with allow once, always allow, and reject choices. Patch previews use the same hunk matching without mutating disk. | Full | If an invalid patch cannot be simulated, the request still shows its raw patch and target locations and remains blocked until the user decides. | +| Cancellation | `session/cancel` interrupts model turns, questions, compaction, and terminal creation/execution while keeping the session reusable. Crabcode maps completion, output limit, configured turn limit, refusal/content filtering, and cancellation to ACP `end_turn`, `max_tokens`, `max_turn_requests`, `refusal`, and `cancelled`. | Full | Provider failures that are not normal stop conditions remain JSON-RPC/tool errors, as required by ACP's stop-reason model. | +| Commands and skills | Session updates publish global/workspace skills, project custom commands, and `/skills`, `/mcp`, and `/compact`. Custom command agent/model overrides apply to that turn. `/skills` and `/mcp` return local results without spending or persisting a model turn; `/mcp` reports live connection/auth/failure status. `/compact` rewrites persisted context. Unknown slash commands return an explicit error. | Full | Editor-native session/model/mode operations replace TUI-only navigation dialogs and pickers rather than duplicating their terminal UI commands. | +| MCP | Project MCP and client-supplied stdio, HTTP, and SSE servers merge into the session. Static headers, structured results, live status, resources, annotations, images, audio, and metadata are preserved. Project-configured remote MCP continues to use Crabcode's OAuth credential flow. | Full | ACP currently advertises only the HTTP/SSE transport flags; its client-server schema has no stdio flag or remote OAuth fields. Client-supplied remote auth can still be provided through headers. | +| Terminals | `terminal_session` and terminal-mode `bash` use the editor terminal host through create, embed, wait, output, kill, and release. Output is bounded for the model, and cancellation also covers terminal creation. | Full | The editor must advertise terminal hosting. User input and resize happen directly in the embedded editor terminal because ACP has no agent-issued stdin/resize requests. | +| Questions | Agent questions use capability-gated ACP form elicitation with validated non-empty prompts/options, unique labels, ordered single/multi-select answers, custom text, cardinality checks, deduplication, length bounds, cancellation, and safe skip behavior. | Full | Form elicitation is an unstable ACP capability and is only sent to editors that advertise it. | +| Usage | Provider input/output/cache-read/cache-write usage is aggregated across multi-step turns and persisted. ACP always receives context occupancy and cumulative USD cost updates; detailed token/cache values and whether the context size is known are included in `crabcode` `_meta`. Catalog pricing is cache-aware. | Full | Providers that omit usage or models without pricing cannot supply authoritative token or cost data; Crabcode still emits estimated context occupancy and marks unknown context size in metadata. | + +## Runtime requirements + +- Image and audio prompts require a selected model route with the corresponding input modality. +- ACP terminal embedding and question forms require the editor to advertise those client capabilities during initialization. +- Client-supplied remote MCP OAuth parameters are not part of the current ACP server descriptor. Use static headers from the editor, or configure the MCP server in Crabcode to use Crabcode's OAuth flow. +- Fields under `_meta.crabcode` are backwards-compatible Crabcode extensions for session hierarchy and detailed usage accounting. ## Session behavior -`session/close` detaches the editor and cancels any active turn. It does not delete Crabcode session history. Load replays the stored transcript; resume restores the session configuration without replaying prior content. Fork creates a new persisted session with a copied transcript. +`session/close` detaches the editor and cancels any active turn. It does not delete Crabcode session history. `session/delete` removes persisted history and managed attachments. List cursors page through the complete non-archived result set. Load replays the stored transcript; resume restores the session configuration without replaying prior content. Fork creates a new persisted session with a copied transcript and independently managed attachments. Child sessions are listed as normal entries with hierarchy metadata under `_meta.crabcode`. ## Safety notes Crabcode applies the same configured permission rules in ACP as it does in the TUI. When a tool needs approval, the editor receives an ACP permission request. If the editor cannot respond or disconnects, Crabcode denies the request rather than continuing unattended. +Question forms are only sent to editors that advertise ACP form elicitation support. Declining, cancelling, disconnecting, or using an editor without that capability returns empty answers to the agent so the session can continue without waiting indefinitely. + Client-supplied MCP servers run with the same trust as project-configured MCP: stdio servers can execute local processes, and remote servers can send the headers and credentials the editor provides. Only attach MCP servers you trust for that workspace. -Image attachments are decoded under a size limit and written to temporary files under the system temp directory (`…/crabcode/acp-images/`) for the model turn. Prefer cleaning those files after long ACP sessions until automatic cleanup lands. +Image and audio attachments are decoded under a 20 MiB-per-file limit and written to private session-managed storage under Crabcode's state directory (`…/crabcode/attachments//`). Audio input accepts WAV and MP3 only. Closing an editor session keeps those files because history remains loadable; deleting the persisted session removes its managed attachment directory. Forks receive independent copies so deleting either session does not break the other. Readable external or temporary attachment paths from older sessions are copied into managed storage the next time Crabcode loads the session. -The capability matrix matches what the ACP server implements today. Crabcode only advertises protocol capability flags it handles (`loadSession`, image and embedded-context prompts, HTTP/SSE MCP, and list/resume/fork/close session ops). +The capability matrix matches what the ACP server implements today. Crabcode only advertises protocol capability flags it handles (`loadSession`, image, audio, and embedded-context prompts, HTTP/SSE MCP, and list/resume/fork/close/delete session ops). Client-side form elicitation and terminal hosting are capability-gated during initialization before Crabcode sends those requests. Stdio MCP is accepted even though the current ACP capability object has no separate stdio flag. diff --git a/src/acp/server.rs b/src/acp/server.rs index 5e87fde9..49c2cfc0 100644 --- a/src/acp/server.rs +++ b/src/acp/server.rs @@ -1,11 +1,12 @@ use agent_client_protocol::schema::v1::{ AgentCapabilities, CancelNotification, CloseSessionRequest, CloseSessionResponse, - ForkSessionRequest, ForkSessionResponse, Implementation, InitializeRequest, InitializeResponse, - ListSessionsRequest, LoadSessionRequest, McpCapabilities, NewSessionRequest, - PromptCapabilities, PromptRequest, ResumeSessionRequest, SessionCapabilities, - SessionCloseCapabilities, SessionForkCapabilities, SessionListCapabilities, - SessionNotification, SessionResumeCapabilities, SessionUpdate, SetSessionConfigOptionRequest, - SetSessionModeRequest, SetSessionModeResponse, + DeleteSessionRequest, DeleteSessionResponse, ForkSessionRequest, ForkSessionResponse, + Implementation, InitializeRequest, InitializeResponse, ListSessionsRequest, LoadSessionRequest, + McpCapabilities, NewSessionRequest, PromptCapabilities, PromptRequest, ResumeSessionRequest, + SessionCapabilities, SessionCloseCapabilities, SessionDeleteCapabilities, + SessionForkCapabilities, SessionListCapabilities, SessionNotification, + SessionResumeCapabilities, SessionUpdate, SetSessionConfigOptionRequest, SetSessionModeRequest, + SetSessionModeResponse, }; use agent_client_protocol::{Agent, Stdio}; use anyhow::{Context, Result}; @@ -30,12 +31,14 @@ pub async fn run(cwd: Option) -> Result<()> { })?; let service = crate::acp::service::AcpService::new(&workspace) .map_err(|_| anyhow::anyhow!("failed to initialize ACP session storage"))?; + let initialize_service = service.clone(); Agent .builder() .name("crabcode-acp") .on_receive_request( async move |request: InitializeRequest, responder, _connection| { + initialize_service.set_client_capabilities(request.client_capabilities.clone()); let response = InitializeResponse::new(request.protocol_version) .agent_capabilities(capabilities()) .agent_info(Implementation::new("crabcode", env!("CARGO_PKG_VERSION"))); @@ -46,16 +49,39 @@ pub async fn run(cwd: Option) -> Result<()> { .on_receive_request( { let service = service.clone(); - async move |request: ForkSessionRequest, responder, _connection| { - let result = service - .fork_session(request.session_id.to_string(), request.cwd) - .await - .map(|response| { - ForkSessionResponse::new(response.session_id) + async move |request: DeleteSessionRequest, responder, _connection| { + responder.respond_with_result( + service + .delete_session(&request.session_id.to_string()) + .await + .map(|_| DeleteSessionResponse::new()), + ) + } + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + { + let service = service.clone(); + async move |request: ForkSessionRequest, responder, connection| { + let service = service.clone(); + let task_connection = connection.clone(); + connection.spawn(async move { + let response = service + .fork_session(request.session_id.to_string(), request.cwd) + .await?; + let session_id = response.session_id.clone(); + let commands = service.available_commands(&session_id.to_string()).await?; + responder.respond( + ForkSessionResponse::new(session_id.clone()) .modes(response.modes) - .config_options(response.config_options) - }); - responder.respond_with_result(result) + .config_options(response.config_options), + )?; + task_connection.send_notification(SessionNotification::new( + session_id, + SessionUpdate::AvailableCommandsUpdate(commands), + )) + }) } }, agent_client_protocol::on_receive_request!(), @@ -188,7 +214,9 @@ pub async fn run(cwd: Option) -> Result<()> { { let service = service.clone(); async move |request: ListSessionsRequest, responder, _connection| { - responder.respond_with_result(service.list_sessions(request.cwd).await) + responder.respond_with_result( + service.list_sessions(request.cwd, request.cursor).await, + ) } }, agent_client_protocol::on_receive_request!(), @@ -240,13 +268,33 @@ pub async fn run(cwd: Option) -> Result<()> { fn capabilities() -> AgentCapabilities { AgentCapabilities::new() .load_session(true) - .prompt_capabilities(PromptCapabilities::new().embedded_context(true).image(true)) + .prompt_capabilities( + PromptCapabilities::new() + .embedded_context(true) + .image(true) + .audio(true), + ) .mcp_capabilities(McpCapabilities::new().http(true).sse(true)) .session_capabilities( SessionCapabilities::new() .list(SessionListCapabilities::new()) .resume(SessionResumeCapabilities::new()) .fork(SessionForkCapabilities::new()) - .close(SessionCloseCapabilities::new()), + .close(SessionCloseCapabilities::new()) + .delete(SessionDeleteCapabilities::new()), ) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn advertises_audio_prompt_support() { + let capabilities = capabilities(); + let prompt = capabilities.prompt_capabilities; + assert!(prompt.audio); + assert!(prompt.image); + assert!(capabilities.session_capabilities.delete.is_some()); + } +} diff --git a/src/acp/service.rs b/src/acp/service.rs index 1277e7e2..acbb5039 100644 --- a/src/acp/service.rs +++ b/src/acp/service.rs @@ -2,19 +2,23 @@ use crate::config::configuration::LoadedConfig; use crate::session::manager::SessionManager; use agent_client_protocol::schema::v1::{ AvailableCommand, AvailableCommandInput, AvailableCommandsUpdate, ContentBlock, ContentChunk, - EmbeddedResourceResource, ListSessionsResponse, LoadSessionResponse, McpServer, - NewSessionResponse, PermissionOption, PermissionOptionKind, PromptResponse, + Cost as AcpCost, CreateElicitationRequest, CreateTerminalRequest, ElicitationAction, + ElicitationContentValue, ElicitationFormMode, ElicitationSchema, ElicitationSessionScope, + EmbeddedResourceResource, EnumOption, KillTerminalRequest, ListSessionsResponse, + LoadSessionResponse, McpServer, MultiSelectPropertySchema, NewSessionResponse, + PermissionOption, PermissionOptionKind, PromptResponse, ReleaseTerminalRequest, RequestPermissionOutcome, RequestPermissionRequest, ResumeSessionResponse, SessionConfigOption, SessionConfigOptionCategory, SessionConfigSelectGroup, SessionConfigSelectOption, SessionInfo, SessionMode, SessionModeState, SessionNotification, SessionUpdate, - SetSessionConfigOptionResponse, StopReason, ToolCall, ToolCallContent, ToolCallStatus, + SetSessionConfigOptionResponse, StopReason, StringPropertySchema, Terminal, + TerminalOutputRequest, ToolCall, ToolCallContent, ToolCallLocation, ToolCallStatus, ToolCallUpdate, ToolCallUpdateFields, ToolKind, UnstructuredCommandInput, UsageUpdate, + WaitForTerminalExitRequest, }; use agent_client_protocol::{Client, ConnectionTo, Error}; use base64::Engine as _; use std::collections::HashMap; use std::path::{Path, PathBuf}; -use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; use tokio::sync::Mutex as AsyncMutex; use tokio_util::sync::CancellationToken; @@ -23,6 +27,449 @@ use tokio_util::sync::CancellationToken; pub struct AcpService { sessions: Arc>>, session_manager: Arc>, + client_capabilities: Arc>, +} + +fn command_text(parts: &[ContentBlock]) -> String { + let mut text = String::new(); + for part in parts { + match part { + ContentBlock::Text(content) => text.push_str(&content.text), + ContentBlock::ResourceLink(link) => text.push_str(&format!("[{}]", link.uri)), + ContentBlock::Resource(resource) => match &resource.resource { + EmbeddedResourceResource::TextResourceContents(resource) => { + text.push_str(&format!("[{}]\n{}", resource.uri, resource.text)); + } + EmbeddedResourceResource::BlobResourceContents(resource) => { + text.push_str(&format!("[{}]", resource.uri)); + } + _ => {} + }, + _ => {} + } + } + text +} + +fn acp_session_info( + session: crate::session::manager::SessionInfo, + root_id: Option, +) -> SessionInfo { + let mut meta = serde_json::Map::new(); + meta.insert( + "crabcode".to_string(), + serde_json::json!({ + "parentSessionId": session.parent_id, + "rootSessionId": root_id, + }), + ); + SessionInfo::new(session.id, session.workspace_path) + .title(session.title) + .updated_at(system_time_to_iso8601(session.updated_at)) + .meta(meta) +} + +fn session_page( + cursor: Option<&str>, + total: usize, +) -> Result<(usize, usize, Option), Error> { + let offset = cursor + .unwrap_or("0") + .parse::() + .map_err(|_| Error::invalid_params().data("invalid session list cursor"))?; + let end = offset.saturating_add(100).min(total); + let next_cursor = (end < total).then(|| end.to_string()); + Ok((offset.min(total), end, next_cursor)) +} + +fn mcp_native_tool_content(metadata: &serde_json::Value) -> Vec { + metadata + .get("mcp_result") + .and_then(|result| result.get("content")) + .and_then(serde_json::Value::as_array) + .into_iter() + .flatten() + .filter(|block| { + !matches!( + block.get("type").and_then(serde_json::Value::as_str), + Some("text") + ) + }) + .map(|block| { + serde_json::from_value::(block.clone()) + .map(ToolCallContent::from) + .unwrap_or_else(|_| { + let kind = block + .get("type") + .and_then(serde_json::Value::as_str) + .unwrap_or("unknown"); + ToolCallContent::from(format!( + "[MCP {kind} content]\n{}", + serde_json::to_string_pretty(block).unwrap_or_else(|_| block.to_string()) + )) + }) + }) + .collect() +} + +fn metadata_has_mcp_images(metadata: Option<&serde_json::Value>) -> bool { + metadata + .and_then(|metadata| metadata.get("mcp_result")) + .and_then(|result| result.get("content")) + .and_then(serde_json::Value::as_array) + .is_some_and(|content| { + content + .iter() + .any(|block| block.get("type").and_then(serde_json::Value::as_str) == Some("image")) + }) +} + +fn permission_tool_content( + prompt: &crate::tools::PermissionPrompt, + cwd: &Path, +) -> Vec { + match prompt.tool_id.as_str() { + "write" => prompt + .raw_input + .get("file_path") + .or_else(|| prompt.raw_input.get("filePath")) + .and_then(serde_json::Value::as_str) + .zip( + prompt + .raw_input + .get("content") + .and_then(serde_json::Value::as_str), + ) + .map(|(path, new_text)| vec![preflight_diff(path, new_text, cwd)]) + .unwrap_or_default(), + "write_files" => prompt + .raw_input + .get("files") + .and_then(serde_json::Value::as_array) + .into_iter() + .flatten() + .filter_map(|file| { + let path = file.get("file_path")?.as_str()?; + let content = file.get("content")?.as_str()?; + Some(preflight_diff(path, content, cwd)) + }) + .collect(), + "edit" => permission_edit_diff(&prompt.raw_input, cwd) + .into_iter() + .collect(), + "apply_patch" => crate::tools::patch::preview_patch(&prompt.raw_input, cwd) + .unwrap_or_default() + .into_iter() + .map(|change| { + agent_client_protocol::schema::v1::Diff::new(change.path, change.new_text) + .old_text(change.old_text) + .into() + }) + .collect(), + _ => Vec::new(), + } +} + +fn preflight_diff(path: &str, new_text: &str, cwd: &Path) -> ToolCallContent { + let path = absolute_tool_path(path, cwd); + let old_text = std::fs::read_to_string(&path).ok(); + agent_client_protocol::schema::v1::Diff::new(path, new_text.to_string()) + .old_text(old_text) + .into() +} + +fn permission_edit_diff(input: &serde_json::Value, cwd: &Path) -> Option { + let path = input + .get("file_path") + .or_else(|| input.get("filePath"))? + .as_str()?; + let old_string = input.get("old_string")?.as_str()?; + let new_string = input.get("new_string")?.as_str()?; + let absolute = absolute_tool_path(path, cwd); + let old_text = std::fs::read_to_string(&absolute).ok()?; + let new_text = if input + .get("replace_all") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false) + { + old_text.replace(old_string, new_string) + } else { + old_text.replacen(old_string, new_string, 1) + }; + Some( + agent_client_protocol::schema::v1::Diff::new(absolute, new_text) + .old_text(old_text) + .into(), + ) +} + +fn write_prompt_audio( + session_id: &str, + audio: &agent_client_protocol::schema::v1::AudioContent, +) -> Result { + const MAX_AUDIO_BYTES: usize = 20 * 1024 * 1024; + let media_type = audio.mime_type.trim().to_ascii_lowercase(); + let extension = match media_type.as_str() { + "audio/wav" | "audio/x-wav" | "audio/wave" => "wav", + "audio/mpeg" | "audio/mp3" => "mp3", + _ => { + return Err(Error::invalid_params() + .data(format!("unsupported audio MIME type: {}", audio.mime_type))); + } + }; + let data = base64::engine::general_purpose::STANDARD + .decode(audio.data.trim()) + .map_err(|error| Error::invalid_params().data(format!("invalid audio data: {error}")))?; + if data.len() > MAX_AUDIO_BYTES { + return Err(Error::invalid_params().data("audio exceeds the 20 MiB size limit")); + } + let path = crate::persistence::attachments::write(session_id, extension, &data) + .map_err(|_| internal_error())?; + Ok(path.to_string_lossy().into_owned()) +} + +fn model_supports_audio(config: &LoadedConfig, provider: &str, model: &str) -> bool { + crate::model::discovery::Discovery::new_with_custom(Some( + config.merged_config.custom_providers.clone(), + )) + .ok() + .is_some_and(|discovery| discovery.model_supports_input_modality(provider, model, "audio")) +} + +struct ManagedAttachmentGuard { + paths: Vec, + committed: bool, +} + +impl ManagedAttachmentGuard { + fn new(paths: Vec) -> Self { + Self { + paths, + committed: false, + } + } + + fn commit(&mut self) { + self.committed = true; + } +} + +impl Drop for ManagedAttachmentGuard { + fn drop(&mut self) { + if !self.committed { + for path in &self.paths { + crate::persistence::attachments::remove_file(Path::new(path)); + } + } + } +} + +struct AcpQuestionField { + selection: String, + custom: String, + labels: HashMap, + multiple: bool, +} + +struct AcpQuestionForm { + request: CreateElicitationRequest, + fields: Vec, +} + +fn skipped_question_answers(questions: &serde_json::Value) -> serde_json::Value { + let count = questions.as_array().map_or(1, Vec::len); + serde_json::Value::Array( + (0..count) + .map(|_| serde_json::Value::Array(Vec::new())) + .collect(), + ) +} + +fn question_text(question: &serde_json::Value, key: &str, fallback: &str) -> String { + question + .get(key) + .and_then(serde_json::Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or(fallback) + .to_string() +} + +fn acp_question_form( + session_id: &str, + tool_call_id: Option<&str>, + questions: &serde_json::Value, +) -> AcpQuestionForm { + let question_items = questions + .as_array() + .cloned() + .unwrap_or_else(|| vec![questions.clone()]); + let mut schema = ElicitationSchema::new() + .title("Agent questions") + .description("Answer any fields you want; blank fields are treated as skipped."); + let mut fields = Vec::with_capacity(question_items.len()); + + for (question_index, question) in question_items.iter().enumerate() { + let selection = format!("question_{question_index}"); + let custom = format!("question_{question_index}_custom"); + let prompt = question_text(question, "question", "Question"); + let header = question_text( + question, + "header", + &format!("Question {}", question_index + 1), + ); + let mut labels = HashMap::new(); + let options = question + .get("options") + .and_then(serde_json::Value::as_array) + .into_iter() + .flatten() + .enumerate() + .filter_map(|(option_index, option)| { + let label = option + .get("label") + .and_then(serde_json::Value::as_str) + .or_else(|| option.as_str())? + .trim(); + if label.is_empty() { + return None; + } + + let value = format!("q{question_index}_option_{option_index}"); + labels.insert(value.clone(), label.to_string()); + let description = option + .get("description") + .and_then(serde_json::Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()); + let mut option = EnumOption::new(value, label); + if let Some(description) = description { + option = option.description(description); + } + Some(option) + }) + .collect::>(); + let multiple = question + .get("multiple") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false); + + if multiple { + schema = schema.property( + selection.clone(), + MultiSelectPropertySchema::titled(options) + .title(header.clone()) + .description(prompt.clone()), + false, + ); + } else { + schema = schema.property( + selection.clone(), + StringPropertySchema::new() + .title(header.clone()) + .description(prompt.clone()) + .one_of(options), + false, + ); + } + schema = schema.property( + custom.clone(), + StringPropertySchema::new() + .title(format!("{header}: custom answer")) + .description("Optional free-form answer."), + false, + ); + fields.push(AcpQuestionField { + selection, + custom, + labels, + multiple, + }); + } + + let scope = ElicitationSessionScope::new(session_id.to_string()) + .tool_call_id(tool_call_id.map(agent_client_protocol::schema::v1::ToolCallId::new)); + let request = CreateElicitationRequest::new( + ElicitationFormMode::new(scope, schema), + "The agent needs additional input to continue.", + ); + AcpQuestionForm { request, fields } +} + +fn acp_question_answers( + fields: &[AcpQuestionField], + action: ElicitationAction, +) -> serde_json::Value { + let ElicitationAction::Accept(accepted) = action else { + return serde_json::Value::Array( + fields + .iter() + .map(|_| serde_json::Value::Array(Vec::new())) + .collect(), + ); + }; + let content = accepted.content.unwrap_or_default(); + serde_json::Value::Array( + fields + .iter() + .map(|field| { + let mut answers = Vec::new(); + match content.get(&field.selection) { + Some(ElicitationContentValue::String(value)) if !field.multiple => { + if let Some(label) = field.labels.get(value) { + answers.push(serde_json::Value::String(label.clone())); + } + } + Some(ElicitationContentValue::StringArray(values)) if field.multiple => { + for value in values { + if let Some(label) = field.labels.get(value) { + if !answers.iter().any(|answer| answer.as_str() == Some(label)) { + answers.push(serde_json::Value::String(label.clone())); + } + } + } + } + _ => {} + } + if let Some(ElicitationContentValue::String(custom)) = content.get(&field.custom) { + let custom = custom.trim(); + if !custom.is_empty() { + if !field.multiple { + answers.clear(); + } + answers.push(serde_json::Value::String( + custom.chars().take(8_192).collect(), + )); + } + } + serde_json::Value::Array(answers) + }) + .collect(), + ) +} + +fn compact_command(prompt: &str) -> Result { + let trimmed = prompt.trim(); + let Some(command_line) = trimmed.strip_prefix('/') else { + return Ok(false); + }; + let (name, args) = command_line + .split_once(char::is_whitespace) + .map(|(name, args)| (name, args.trim())) + .unwrap_or((command_line, "")); + if name != "compact" { + return Ok(false); + } + if !args.is_empty() { + return Err(Error::invalid_params().data("Usage: /compact")); + } + Ok(true) +} + +fn permission_tool_call_id(tool_call_id: Option<&str>) -> String { + tool_call_id + .map(str::to_string) + .unwrap_or_else(|| format!("permission:{}", cuid2::create_id())) } fn resolved_reasoning( @@ -56,32 +503,36 @@ fn available_commands(session: &AcpSession) -> Vec { .merged_config .commands .iter() + .filter(|command| !matches!(command.name.as_str(), "compact" | "skills" | "mcp")) .map(|command| { let description = command .description .clone() .unwrap_or_else(|| format!("Run /{}", command.name)); - let mut available = AvailableCommand::new(command.name.clone(), description); - if command.template.contains("$ARGUMENTS") { - available = available.input(AvailableCommandInput::Unstructured( - UnstructuredCommandInput::new("Arguments"), - )); - } - available + AvailableCommand::new(command.name.clone(), description).input( + AvailableCommandInput::Unstructured(UnstructuredCommandInput::new("Arguments")), + ) }) .collect(); - commands.extend(session.skills.all().into_iter().map(|skill| { - AvailableCommand::new( - skill.name.clone(), - skill - .description - .clone() - .unwrap_or_else(|| format!("Use the {} skill", skill.name)), - ) - .input(AvailableCommandInput::Unstructured( - UnstructuredCommandInput::new("Task or context for this skill"), - )) - })); + commands.extend( + session + .skills + .all() + .into_iter() + .filter(|skill| !matches!(skill.name.as_str(), "compact" | "skills" | "mcp")) + .map(|skill| { + AvailableCommand::new( + skill.name.clone(), + skill + .description + .clone() + .unwrap_or_else(|| format!("Use the {} skill", skill.name)), + ) + .input(AvailableCommandInput::Unstructured( + UnstructuredCommandInput::new("Task or context for this skill"), + )) + }), + ); commands.push(AvailableCommand::new( "skills", "List skills available in this workspace", @@ -90,41 +541,41 @@ fn available_commands(session: &AcpSession) -> Vec { "mcp", "List configured MCP servers and their status", )); + commands.push(AvailableCommand::new( + "compact", + "Summarize this session to reduce context", + )); commands.sort_by(|left, right| left.name.cmp(&right.name)); commands.dedup_by(|left, right| left.name == right.name); commands } -async fn expand_slash_command(session: &AcpSession, prompt: &str) -> Result { +#[derive(Debug, PartialEq, Eq)] +enum SlashExpansion { + Prompt { + prompt: String, + agent: Option, + model: Option, + }, + LocalResult(String), +} + +async fn expand_slash_command(session: &AcpSession, prompt: &str) -> Result { let Some(command_line) = prompt.strip_prefix('/') else { - return Ok(prompt.to_string()); + return Ok(SlashExpansion::Prompt { + prompt: prompt.to_string(), + agent: None, + model: None, + }); }; let (name, args) = command_line .split_once(char::is_whitespace) .map(|(name, args)| (name, args.trim_start())) .unwrap_or((command_line, "")); - if let Some(command) = session - .config - .merged_config - .commands - .iter() - .find(|command| command.name == name) - { - return command - .render(args) - .await - .map(|rendered| rendered.prompt) - .map_err(|_| internal_error()); - } - if let Some(skill) = session.skills.get(name) { - let mut expanded = skill.content.clone(); + if name == "skills" { if !args.is_empty() { - expanded.push_str("\n\nUser task/context:\n"); - expanded.push_str(args); + return Err(Error::invalid_params().data("Usage: /skills")); } - return Ok(expanded); - } - if name == "skills" { let skills = session .skills .all() @@ -137,38 +588,72 @@ async fn expand_slash_command(session: &AcpSession, prompt: &str) -> Result>(); - return Ok(if skills.is_empty() { + return Ok(SlashExpansion::LocalResult(if skills.is_empty() { "No skills are available in this workspace.".to_string() } else { - format!("Available workspace skills:\n{}", skills.join("\n")) - }); + format!("Available skills:\n{}", skills.join("\n")) + })); } if name == "mcp" { - let servers = session - .config - .merged_config - .mcp - .iter() - .map(|(name, server)| { + if !args.is_empty() { + return Err(Error::invalid_params().data("Usage: /mcp")); + } + let manager = crate::mcp::McpManager::ensure( + session.config.merged_config.mcp.clone(), + session.cwd.clone(), + ); + let servers = manager + .lock() + .await + .views() + .into_iter() + .map(|server| { + let detail = server + .detail + .map(|detail| format!(": {detail}")) + .unwrap_or_default(); format!( - "- {} ({}, {})", - name, - server.kind(), - if server.enabled() { - "enabled" - } else { - "disabled" - } + "- {} ({}, {}){}", + server.name, server.kind, server.status, detail ) }) .collect::>(); - return Ok(if servers.is_empty() { + return Ok(SlashExpansion::LocalResult(if servers.is_empty() { "No MCP servers are configured for this workspace.".to_string() } else { - format!("Configured MCP servers:\n{}", servers.join("\n")) + format!("MCP servers:\n{}", servers.join("\n")) + })); + } + if let Some(command) = session + .config + .merged_config + .commands + .iter() + .find(|command| command.name == name) + { + return command + .render(args) + .await + .map_err(|_| internal_error()) + .map(|rendered| SlashExpansion::Prompt { + prompt: rendered.prompt, + agent: rendered.agent, + model: rendered.model, + }); + } + if let Some(skill) = session.skills.get(name) { + let mut expanded = skill.content.clone(); + if !args.is_empty() { + expanded.push_str("\n\nUser task/context:\n"); + expanded.push_str(args); + } + return Ok(SlashExpansion::Prompt { + prompt: expanded, + agent: None, + model: None, }); } - Ok(prompt.to_string()) + Err(Error::invalid_params().data(format!("Unknown ACP command: /{name}"))) } fn merge_acp_mcp_servers(config: &mut LoadedConfig, servers: Vec) { @@ -262,9 +747,35 @@ impl AcpService { Ok(Self { sessions: Arc::new(AsyncMutex::new(HashMap::new())), session_manager: Arc::new(Mutex::new(session_manager)), + client_capabilities: Arc::new(Mutex::new(Default::default())), }) } + pub fn set_client_capabilities( + &self, + capabilities: agent_client_protocol::schema::v1::ClientCapabilities, + ) { + if let Ok(mut current) = self.client_capabilities.lock() { + *current = capabilities; + } + } + + fn supports_form_elicitation(&self) -> bool { + self.client_capabilities + .lock() + .ok() + .and_then(|capabilities| capabilities.elicitation.clone()) + .and_then(|elicitation| elicitation.form) + .is_some() + } + + fn supports_terminals(&self) -> bool { + self.client_capabilities + .lock() + .ok() + .is_some_and(|capabilities| capabilities.terminal) + } + pub async fn available_commands( &self, session_id: &str, @@ -344,13 +855,17 @@ impl AcpService { .config_options(session_config_options(&session))) } - pub async fn list_sessions(&self, cwd: Option) -> Result { + pub async fn list_sessions( + &self, + cwd: Option, + cursor: Option, + ) -> Result { let cwd = cwd.as_deref().map(workspace_path).transpose()?; let manager = self.session_manager.lock().map_err(|_| internal_error())?; let mut sessions = manager .list_sessions() .into_iter() - .filter(|session| session.parent_id.is_none() && session.archived_at.is_none()) + .filter(|session| session.archived_at.is_none()) .filter(|session| { cwd.as_ref() .is_none_or(|cwd| session.workspace_path == cwd.to_string_lossy()) @@ -358,17 +873,19 @@ impl AcpService { .collect::>(); sessions.sort_by(|left, right| right.updated_at.cmp(&left.updated_at)); + let (offset, end, next_cursor) = session_page(cursor.as_deref(), sessions.len())?; Ok(ListSessionsResponse::new( sessions .into_iter() - .take(100) + .skip(offset) + .take(end.saturating_sub(offset)) .map(|session| { - SessionInfo::new(session.id, session.workspace_path) - .title(session.title) - .updated_at(system_time_to_iso8601(session.updated_at)) + let root_id = manager.root_session_id_for(&session.id); + acp_session_info(session, root_id) }) .collect(), - )) + ) + .next_cursor(next_cursor)) } pub async fn load_session( @@ -377,11 +894,11 @@ impl AcpService { cwd: PathBuf, connection: ConnectionTo, ) -> Result { - let (_session, messages) = self.attach_persisted_session(&session_id, cwd).await?; - replay_messages(&connection, &session_id, &messages)?; + let (session, messages) = self.attach_persisted_session(&session_id, cwd).await?; + replay_messages(&connection, &session_id, &messages, &session.cwd)?; Ok(LoadSessionResponse::new() - .modes(session_modes(&_session)) - .config_options(session_config_options(&_session))) + .modes(session_modes(&session)) + .config_options(session_config_options(&session))) } pub async fn resume_session( @@ -403,13 +920,29 @@ impl AcpService { let (source, messages) = self.attach_persisted_session(&session_id, cwd).await?; let fork_id = { let mut manager = self.session_manager.lock().map_err(|_| internal_error())?; + let source_title = manager + .get_session(&session_id) + .map(|session| session.title.clone()) + .unwrap_or_else(|| session_id.clone()); manager .switch_current_workspace_path(&source.cwd.to_string_lossy()) .map_err(|_| internal_error())?; - let fork_id = manager.create_session(Some(format!("{} (fork)", session_id))); - manager + let fork_id = manager.create_session(Some(format!("{source_title} (fork)"))); + let messages = + match crate::persistence::attachments::clone_messages(&messages, &fork_id) { + Ok(messages) => messages, + Err(_) => { + manager.delete_session(&fork_id); + return Err(internal_error()); + } + }; + if manager .replace_session_messages(&fork_id, messages) - .map_err(|_| internal_error())?; + .is_err() + { + manager.delete_session(&fork_id); + return Err(internal_error()); + } fork_id }; self.sessions @@ -429,6 +962,16 @@ impl AcpService { } } + pub async fn delete_session(&self, session_id: &str) -> Result<(), Error> { + self.close_session(session_id).await; + let mut manager = self.session_manager.lock().map_err(|_| internal_error())?; + match manager.try_delete_session(session_id) { + Ok(true) => Ok(()), + Ok(false) => Err(Error::invalid_params().data("unknown session")), + Err(_) => Err(internal_error()), + } + } + pub async fn cancel_session(&self, session_id: &str) { if let Some(session) = self.sessions.lock().await.get(session_id) { if let Some(cancellation) = &session.cancellation { @@ -579,20 +1122,93 @@ impl AcpService { prompt: Vec, connection: ConnectionTo, ) -> Result { - let session = self + let mut session = self .sessions .lock() .await .get(&session_id) .cloned() .ok_or_else(|| Error::invalid_params().data("unknown session"))?; + let compact_text = command_text(&prompt); + if compact_command(&compact_text)? { + if prompt + .iter() + .any(|part| !matches!(part, ContentBlock::Text(_))) + { + return Err(Error::invalid_params().data("/compact does not accept attachments")); + } + return self.compact_session(&session_id, session, connection).await; + } + let expansion = if compact_text.trim_start().starts_with('/') { + Some(expand_slash_command(&session, &compact_text).await?) + } else { + None + }; + let expanded_prompt = match expansion { + Some(SlashExpansion::LocalResult(text)) => { + if prompt + .iter() + .any(|part| !matches!(part, ContentBlock::Text(_))) + { + return Err(Error::invalid_params() + .data("local ACP commands do not accept attachments")); + } + let message_id = cuid2::create_id(); + send_replay_text(&connection, &session_id, &message_id, &text, false, false)?; + return Ok(PromptResponse::new(StopReason::EndTurn)); + } + Some(SlashExpansion::Prompt { + prompt, + agent, + model, + }) => { + if let Some(agent) = agent { + if session + .config + .merged_config + .agent_registry + .get(&agent) + .is_none() + { + return Err(Error::invalid_params() + .data(format!("custom command references unknown agent: {agent}"))); + } + session.agent = agent; + } + if let Some(model_ref) = model { + let (provider, model) = crate::app::parse_model_ref(&model_ref); + let canonical = format!("{provider}/{model}"); + let model = find_selectable_model(&session.models, &canonical)?; + session.provider.clone_from(&model.provider_id); + session.model.clone_from(&model.id); + session.reasoning = resolved_reasoning(&session, session.reasoning_selection); + session.context_window = + model_context_window(&session.config, &session.provider, &session.model); + } + Some(prompt) + } + None => None, + }; let supports_images = session .models .iter() .find(|model| model.provider_id == session.provider && model.id == session.model) .is_some_and(|model| model.attachment); - let (prompt, local_image_paths) = prompt_content(prompt, supports_images, &session)?; - let prompt = expand_slash_command(&session, &prompt).await?; + let supports_audio = + model_supports_audio(&session.config, &session.provider, &session.model); + let (mut prompt, local_image_paths, local_audio_paths) = prompt_content( + prompt, + supports_images, + supports_audio, + &session_id, + &session, + )?; + if let Some(expanded_prompt) = expanded_prompt { + prompt = expanded_prompt; + } + let mut managed_paths = local_image_paths.clone(); + managed_paths.extend(local_audio_paths.clone()); + let mut attachment_guard = ManagedAttachmentGuard::new(managed_paths); if prompt.trim().is_empty() { return Err(Error::invalid_params().data("prompt must include text content")); } @@ -617,6 +1233,7 @@ impl AcpService { }; let mut user_message = crate::session::types::Message::user(&prompt); user_message.local_image_paths = local_image_paths; + user_message.local_audio_paths = local_audio_paths; user_message.provider = Some(session.provider.clone()); user_message.model = Some(session.model.clone()); user_message.agent_mode = Some(session.agent.clone()); @@ -625,6 +1242,7 @@ impl AcpService { manager .add_message_to_session(&session_id, &user_message) .map_err(|_| internal_error())?; + attachment_guard.commit(); manager .set_session_status( &session_id, @@ -635,9 +1253,10 @@ impl AcpService { } messages.push(user_message); + let (sender, mut receiver) = tokio::sync::mpsc::unbounded_channel(); let process_registry = std::sync::Arc::new(crate::tools::ProcessRegistry::new()); let prompt_registry = crate::tools::initialize_tool_registry_with_dynamic_config( - None, + Some(sender.clone()), tool_permissions(&session), session.config.merged_config.agent_registry.clone(), cancellation.clone(), @@ -663,9 +1282,19 @@ impl AcpService { .await; messages.insert(0, crate::session::types::Message::system(system_prompt)); let base_context_tokens = crate::session::compaction::total_context_tokens(&messages); - send_usage(&connection, &session_id, &session, base_context_tokens)?; + let base_cost = messages + .iter() + .filter_map(|message| message.cost) + .sum::(); + send_usage( + &connection, + &session_id, + &session, + base_context_tokens, + (base_cost > 0.0).then_some(base_cost), + None, + )?; - let (sender, mut receiver) = tokio::sync::mpsc::unbounded_channel(); let stream_session_id = session_id.clone(); let stream_cancellation = cancellation.clone(); let stream_sender = sender.clone(); @@ -703,13 +1332,14 @@ impl AcpService { let _ = stream_sender.send(crate::llm::ChunkMessage::End); }); - let message_id = cuid2::create_id(); let mut assistant = crate::session::types::Message::incomplete(""); + let message_id = assistant.id.clone(); assistant.provider = Some(session.provider.clone()); assistant.model = Some(session.model.clone()); assistant.agent_mode = Some(session.agent.clone()); let mut failed = None; let mut cancelled = false; + let mut turn_stop_reason = None; while let Some(chunk) = receiver.recv().await { match chunk { @@ -723,7 +1353,7 @@ impl AcpService { } crate::llm::ChunkMessage::ToolCalls(tool_calls) => { for tool_call in tool_calls { - send_tool_call(&connection, &session_id, tool_call)?; + send_tool_call(&connection, &session_id, tool_call, &session.cwd)?; } } crate::llm::ChunkMessage::ToolResult(result) => { @@ -732,35 +1362,70 @@ impl AcpService { "name": result.name, "content": result.content, })); - send_tool_result(&connection, &session_id, result)?; + send_tool_result(&connection, &session_id, result, &session.cwd)?; } crate::llm::ChunkMessage::Metrics { token_count, duration_ms, + usage, + cost, } => { assistant.token_count = Some(token_count); assistant.duration_ms = Some(duration_ms); + if let Some(usage) = usage { + assistant.apply_usage(usage, cost); + } send_usage( &connection, &session_id, &session, base_context_tokens.saturating_add(token_count), + cost.map(|turn_cost| base_cost + turn_cost), + usage, )?; } crate::llm::ChunkMessage::Cancelled => cancelled = true, crate::llm::ChunkMessage::Failed(error) => failed = Some(error), + crate::llm::ChunkMessage::TurnStopReason(reason) => turn_stop_reason = Some(reason), crate::llm::ChunkMessage::PermissionRequest(prompt) => { let response = request_permission(&connection, &session_id, &prompt).await; let _ = prompt.response_tx.send(response); } - crate::llm::ChunkMessage::QuestionRequest { response_tx, .. } => { - let _ = response_tx.send(serde_json::json!({"skipped": true})); + crate::llm::ChunkMessage::QuestionRequest { + tool_call_id, + questions, + response_tx, + } => { + let response = if self.supports_form_elicitation() { + request_questions( + &connection, + &session_id, + tool_call_id.as_deref(), + &questions, + &cancellation, + ) + .await + } else { + skipped_question_answers(&questions) + }; + let _ = response_tx.send(response); } crate::llm::ChunkMessage::TerminalSessionRequest(request) => { - let _ = request - .control_tx - .send(crate::tools::TerminalSessionControl::Stop); - } + if self.supports_terminals() { + bridge_terminal_session( + &connection, + &session_id, + &session.cwd, + request, + &cancellation, + ) + .await; + } else { + let _ = request + .control_tx + .send(crate::tools::TerminalSessionControl::Stop); + } + } crate::llm::ChunkMessage::End => break, _ => {} } @@ -794,7 +1459,204 @@ impl AcpService { if let Some(error) = failed { return Err(internal_error_with(&error)); } - Ok(PromptResponse::new(StopReason::EndTurn)) + Ok(PromptResponse::new(acp_stop_reason(turn_stop_reason))) + } + + async fn compact_session( + &self, + session_id: &str, + session: AcpSession, + connection: ConnectionTo, + ) -> Result { + let cancellation = CancellationToken::new(); + { + let mut sessions = self.sessions.lock().await; + let current = sessions + .get_mut(session_id) + .ok_or_else(|| Error::invalid_params().data("unknown session"))?; + if current.cancellation.is_some() { + return Err(Error::invalid_params().data("session already has an active prompt")); + } + current.cancellation = Some(cancellation.clone()); + } + let status_result = self + .session_manager + .lock() + .map_err(|_| internal_error())? + .set_session_status( + session_id, + crate::session::types::SessionStatus::Streaming, + None, + ); + if status_result.is_err() { + if let Some(current) = self.sessions.lock().await.get_mut(session_id) { + current.cancellation = None; + } + return Err(internal_error()); + } + + let result = self + .run_compaction(session_id, &session, cancellation.clone()) + .await; + if let Some(current) = self.sessions.lock().await.get_mut(session_id) { + current.cancellation = None; + } + self.session_manager + .lock() + .map_err(|_| internal_error())? + .set_session_status(session_id, crate::session::types::SessionStatus::Idle, None) + .map_err(|_| internal_error())?; + + match result { + Ok(stats) => { + let feedback = format!( + "Context compacted ({})", + crate::session::compaction::format_compaction_stats(stats) + ); + send_text( + &connection, + session_id, + &cuid2::create_id(), + feedback, + false, + )?; + Ok(PromptResponse::new(StopReason::EndTurn)) + } + Err(_error) if cancellation.is_cancelled() => { + Ok(PromptResponse::new(StopReason::Cancelled)) + } + Err(error) => Err(error), + } + } + + async fn run_compaction( + &self, + session_id: &str, + session: &AcpSession, + cancellation: CancellationToken, + ) -> Result { + let messages = { + let manager = self.session_manager.lock().map_err(|_| internal_error())?; + manager + .get_session_ref(session_id) + .map(|stored| stored.messages.clone()) + .ok_or_else(|| Error::invalid_params().data("unknown session"))? + }; + let selection = crate::session::compaction::select_messages_for_compaction_with_min( + &messages, + crate::session::compaction::DEFAULT_TAIL_TURNS, + 0, + ) + .ok_or_else(|| Error::invalid_params().data("Nothing to compact"))?; + let before_tokens = crate::session::compaction::total_context_tokens(&messages); + let before_messages = + crate::session::compaction::filter_messages_for_context(&messages).len(); + let prompt = crate::session::compaction::build_prompt(&selection.messages_to_summarize); + let summary = crate::llm::client::summarize_for_compaction( + session.provider.clone(), + session.model.clone(), + compaction_reasoning(session), + prompt, + cancellation.clone(), + ) + .await + .map_err(|error| internal_error_with(&error.to_string()))?; + if cancellation.is_cancelled() { + return Err(internal_error_with("Compaction cancelled by user")); + } + let (compacted, stats) = compacted_messages( + &messages, + &selection, + &summary, + session, + before_tokens, + before_messages, + )?; + let mut manager = self.session_manager.lock().map_err(|_| internal_error())?; + manager + .replace_session_messages(session_id, compacted) + .map_err(|_| internal_error())?; + Ok(stats) + } +} + +fn compaction_reasoning(session: &AcpSession) -> Option { + use crate::model::reasoning::ReasoningEffort; + let capability = model_reasoning_capability( + &session.config, + &session.models, + &session.provider, + &session.model, + )?; + [ + ReasoningEffort::None, + ReasoningEffort::Minimal, + ReasoningEffort::Low, + ] + .into_iter() + .find(|effort| capability.values().contains(effort)) + .or(session.reasoning) + .filter(|effort| *effort != ReasoningEffort::None) +} + +fn compacted_messages( + messages: &[crate::session::types::Message], + selection: &crate::session::compaction::CompactionSelection, + summary: &str, + session: &AcpSession, + before_tokens: usize, + before_messages: usize, +) -> Result< + ( + Vec, + crate::session::types::CompactionStats, + ), + Error, +> { + let mut compacted = crate::session::compaction::apply_soft_compaction( + messages, + selection, + summary, + Some(session.model.clone()), + Some(session.provider.clone()), + Some(session.agent.clone()), + crate::session::types::CompactionStats { + before_tokens, + after_tokens: 0, + before_messages, + after_messages: 0, + }, + ); + let after_tokens = crate::session::compaction::total_context_tokens(&compacted); + let after_messages = crate::session::compaction::filter_messages_for_context(&compacted).len(); + let stats = crate::session::types::CompactionStats { + before_tokens, + after_tokens, + before_messages, + after_messages, + }; + if after_tokens >= before_tokens { + return Err(Error::invalid_params().data(format!( + "Compaction did not reduce context ({})", + crate::session::compaction::format_compaction_stats(stats) + ))); + } + if let Some(marker) = compacted + .iter_mut() + .rev() + .find(|message| crate::session::compaction::is_compaction_marker(message)) + { + marker.compaction_stats = Some(stats); + } + Ok((compacted, stats)) +} + +fn acp_stop_reason(reason: Option) -> StopReason { + match reason { + Some(crate::llm::TurnStopReason::MaxTokens) => StopReason::MaxTokens, + Some(crate::llm::TurnStopReason::MaxTurnRequests) => StopReason::MaxTurnRequests, + Some(crate::llm::TurnStopReason::Refusal) => StopReason::Refusal, + None => StopReason::EndTurn, } } @@ -1004,15 +1866,16 @@ fn workspace_path(path: &Path) -> Result { Ok(path) } -static ACP_IMAGE_SEQUENCE: AtomicU64 = AtomicU64::new(0); - fn prompt_content( parts: Vec, supports_images: bool, + supports_audio: bool, + session_id: &str, session: &AcpSession, -) -> Result<(String, Vec), Error> { +) -> Result<(String, Vec, Vec), Error> { let mut text = String::new(); let mut local_image_paths = Vec::new(); + let mut local_audio_paths = Vec::new(); for part in parts { match part { ContentBlock::Text(content) => text.push_str(&content.text), @@ -1035,18 +1898,47 @@ fn prompt_content( session.provider, session.model ))); } - local_image_paths.push(write_prompt_image(&image)?); + match write_prompt_image(session_id, &image) { + Ok(path) => local_image_paths.push(path), + Err(error) => { + for path in &local_image_paths { + crate::persistence::attachments::remove_file(Path::new(path)); + } + return Err(error); + } + } } - ContentBlock::Audio(_) => { - return Err(Error::invalid_params().data("audio ACP prompts are not supported yet")); + ContentBlock::Audio(audio) => { + if !supports_audio { + for path in local_image_paths.iter().chain(local_audio_paths.iter()) { + crate::persistence::attachments::remove_file(Path::new(path)); + } + return Err(Error::invalid_params().data(format!( + "model {}/{} does not support audio input", + session.provider, session.model + ))); + } + match write_prompt_audio(session_id, &audio) { + Ok(path) => local_audio_paths.push(path), + Err(error) => { + for path in local_image_paths.iter().chain(local_audio_paths.iter()) { + crate::persistence::attachments::remove_file(Path::new(path)); + } + return Err(error); + } + } } _ => {} } } - if text.is_empty() && !local_image_paths.is_empty() { - text.push_str("[Image attached]"); + if text.is_empty() { + if !local_image_paths.is_empty() { + text.push_str("[Image attached]"); + } else if !local_audio_paths.is_empty() { + text.push_str("[Audio attached]"); + } } - Ok((text, local_image_paths)) + Ok((text, local_image_paths, local_audio_paths)) } fn prompt_text(parts: Vec) -> Result { @@ -1076,6 +1968,7 @@ fn prompt_text(parts: Vec) -> Result { } fn write_prompt_image( + session_id: &str, image: &agent_client_protocol::schema::v1::ImageContent, ) -> Result { const MAX_IMAGE_BYTES: usize = 20 * 1024 * 1024; @@ -1099,11 +1992,8 @@ fn write_prompt_image( return Err(Error::invalid_params().data("image exceeds the 20 MiB size limit")); } - let directory = std::env::temp_dir().join("crabcode").join("acp-images"); - std::fs::create_dir_all(&directory).map_err(|_| internal_error())?; - let sequence = ACP_IMAGE_SEQUENCE.fetch_add(1, Ordering::Relaxed); - let path = directory.join(format!("{}-{sequence}.{extension}", std::process::id())); - std::fs::write(&path, data).map_err(|_| internal_error())?; + let path = crate::persistence::attachments::write(session_id, extension, &data) + .map_err(|_| internal_error())?; Ok(path.to_string_lossy().into_owned()) } @@ -1154,10 +2044,142 @@ fn send_text( .map_err(|_| internal_error()) } +fn tool_result_text(payload: &serde_json::Value) -> String { + payload + .get("output") + .or_else(|| payload.get("output_preview")) + .or_else(|| payload.get("error")) + .and_then(serde_json::Value::as_str) + .unwrap_or_default() + .to_string() +} + +fn absolute_tool_path(path: &str, cwd: &Path) -> PathBuf { + let path = PathBuf::from(path); + if path.is_absolute() { + path + } else { + cwd.join(path) + } +} + +fn tool_locations(tool_name: &str, input: &serde_json::Value, cwd: &Path) -> Vec { + let paths = if tool_name == "write_files" { + input + .get("files") + .and_then(serde_json::Value::as_array) + .into_iter() + .flatten() + .filter_map(|file| file.get("file_path").and_then(serde_json::Value::as_str)) + .map(str::to_string) + .collect() + } else if tool_name == "apply_patch" { + crate::tools::patch::patch_paths_from_params(input) + } else { + input + .get("file_path") + .or_else(|| input.get("filePath")) + .or_else(|| input.get("filepath")) + .or_else(|| input.get("path")) + .and_then(serde_json::Value::as_str) + .map(|path| vec![path.to_string()]) + .unwrap_or_default() + }; + + paths + .into_iter() + .map(|path| ToolCallLocation::new(absolute_tool_path(&path, cwd))) + .collect() +} + +fn tool_result_locations(payload: &serde_json::Value, cwd: &Path) -> Vec { + let Some(metadata) = payload.get("metadata") else { + return Vec::new(); + }; + let line = metadata + .get("line_number") + .and_then(serde_json::Value::as_u64) + .and_then(|line| u32::try_from(line).ok()); + + if let Some(changes) = metadata + .get("changes") + .and_then(serde_json::Value::as_array) + { + return changes + .iter() + .filter_map(|change| change.get("path").and_then(serde_json::Value::as_str)) + .map(|path| ToolCallLocation::new(absolute_tool_path(path, cwd))) + .collect(); + } + + metadata + .get("path") + .and_then(serde_json::Value::as_str) + .map(|path| ToolCallLocation::new(absolute_tool_path(path, cwd)).line(line)) + .into_iter() + .collect() +} + +fn tool_result_content(payload: &serde_json::Value, cwd: &Path) -> Vec { + let mut content = Vec::new(); + let text = tool_result_text(payload); + if !text.is_empty() { + content.push(ToolCallContent::from(text)); + } + + if let Some(metadata) = payload.get("metadata") { + if let Some(changes) = metadata + .get("changes") + .and_then(serde_json::Value::as_array) + { + content.extend(changes.iter().filter_map(|change| tool_diff(change, cwd))); + } else if let Some(diff) = tool_diff(metadata, cwd) { + content.push(diff); + } + content.extend(mcp_native_tool_content(metadata)); + } + + if !metadata_has_mcp_images(payload.get("metadata")) { + if let Some(images) = payload.get("images").and_then(serde_json::Value::as_array) { + content.extend(images.iter().filter_map(tool_result_image)); + } + } + + content +} + +fn tool_diff(change: &serde_json::Value, cwd: &Path) -> Option { + let path = change.get("path")?.as_str()?; + let new_text = change.get("new_text")?.as_str()?; + let old_text = change.get("old_text").and_then(serde_json::Value::as_str); + Some( + agent_client_protocol::schema::v1::Diff::new( + absolute_tool_path(path, cwd), + new_text.to_string(), + ) + .old_text(old_text.map(str::to_string)) + .into(), + ) +} + +fn tool_result_image(image: &serde_json::Value) -> Option { + let data = image.get("data_url")?.as_str()?; + let media_type = image.get("media_type")?.as_str()?; + let encoded = data + .strip_prefix("data:") + .and_then(|value| value.split_once(',')) + .map(|(_, encoded)| encoded) + .unwrap_or(data); + Some(ToolCallContent::from(ContentBlock::Image( + agent_client_protocol::schema::v1::ImageContent::new(encoded, media_type), + ))) +} + fn send_tool_call( connection: &ConnectionTo, session_id: &str, tool_call: crate::llm::ToolCall, + cwd: &Path, ) -> Result<(), Error> { let raw_input = serde_json::from_str(&tool_call.function.arguments) .unwrap_or_else(|_| serde_json::json!({ "arguments": tool_call.function.arguments })); @@ -1166,6 +2188,7 @@ fn send_tool_call( ToolCall::new(tool_call.id, title) .kind(tool_kind(&tool_call.function.name)) .status(ToolCallStatus::Pending) + .locations(tool_locations(&tool_call.function.name, &raw_input, cwd)) .raw_input(raw_input), ); connection @@ -1178,7 +2201,7 @@ async fn request_permission( session_id: &str, prompt: &crate::tools::PermissionPrompt, ) -> crate::tools::PermissionResponse { - let tool_call_id = format!("permission:{}", cuid2::create_id()); + let tool_call_id = permission_tool_call_id(prompt.tool_call_id.as_deref()); let input = serde_json::json!({ "tool": prompt.tool_id, "permission": prompt.permission, @@ -1187,15 +2210,22 @@ async fn request_permission( "command": prompt.command, "workdir": prompt.workdir, "reason": prompt.reason, + "input": prompt.raw_input, }); - let tool_call = ToolCallUpdate::new( - tool_call_id, - ToolCallUpdateFields::new() + let cwd = PathBuf::from(&prompt.workspace); + let content = permission_tool_content(prompt, &cwd); + let tool_call = ToolCallUpdate::new(tool_call_id, { + let mut fields = ToolCallUpdateFields::new() .title(permission_title(prompt)) .kind(tool_kind(&prompt.tool_id)) .status(ToolCallStatus::Pending) - .raw_input(input), - ); + .locations(tool_locations(&prompt.tool_id, &prompt.raw_input, &cwd)) + .raw_input(input); + if !content.is_empty() { + fields = fields.content(content); + } + fields + }); let request = RequestPermissionRequest::new( session_id.to_string(), tool_call, @@ -1226,6 +2256,157 @@ async fn request_permission( } } +async fn request_questions( + connection: &ConnectionTo, + session_id: &str, + tool_call_id: Option<&str>, + questions: &serde_json::Value, + cancellation: &CancellationToken, +) -> serde_json::Value { + let form = acp_question_form(session_id, tool_call_id, questions); + let request = connection.send_request(form.request).block_task(); + tokio::pin!(request); + let response = tokio::select! { + _ = cancellation.cancelled() => return skipped_question_answers(questions), + response = &mut request => response, + }; + let Ok(response) = response else { + return skipped_question_answers(questions); + }; + acp_question_answers(&form.fields, response.action) +} + +async fn bridge_terminal_session( + connection: &ConnectionTo, + session_id: &str, + session_cwd: &Path, + request: crate::tools::TerminalSessionRequest, + cancellation: &CancellationToken, +) { + let start = request.start; + let control_tx = request.control_tx; + let cwd = start + .workdir + .as_deref() + .map(PathBuf::from) + .map(|path| absolute_tool_path(&path.to_string_lossy(), session_cwd)) + .unwrap_or_else(|| session_cwd.to_path_buf()); + let create = CreateTerminalRequest::new(session_id.to_string(), "bash") + .args(vec!["-c".to_string(), start.command.clone()]) + .cwd(cwd) + .output_byte_limit(crate::tools::terminal_session::MAX_TRANSCRIPT_BYTES as u64); + let create_request = connection.send_request(create).block_task(); + tokio::pin!(create_request); + let terminal_id = match tokio::select! { + _ = cancellation.cancelled() => { + let _ = control_tx.send(crate::tools::TerminalSessionControl::ExternalResult( + crate::tools::terminal_session::external_terminal_result( + &start, "", false, None, true, + ), + )); + return; + } + response = &mut create_request => response, + } { + Ok(response) => response.terminal_id, + Err(error) => { + let _ = control_tx.send(crate::tools::TerminalSessionControl::ExternalError( + format!("ACP client could not create terminal: {error}"), + )); + return; + } + }; + + let terminal_update = SessionUpdate::ToolCallUpdate(ToolCallUpdate::new( + start.tool_call_id.clone(), + ToolCallUpdateFields::new().content(vec![ToolCallContent::Terminal(Terminal::new( + terminal_id.clone(), + ))]), + )); + if connection + .send_notification(SessionNotification::new( + session_id.to_string(), + terminal_update, + )) + .is_err() + { + let _ = connection + .send_request(ReleaseTerminalRequest::new( + session_id.to_string(), + terminal_id, + )) + .block_task() + .await; + let _ = control_tx.send(crate::tools::TerminalSessionControl::ExternalError( + "ACP client could not embed terminal".to_string(), + )); + return; + } + + let wait = connection + .send_request(WaitForTerminalExitRequest::new( + session_id.to_string(), + terminal_id.clone(), + )) + .block_task(); + tokio::pin!(wait); + let (stopped_by_user, exit_code, wait_error) = tokio::select! { + _ = cancellation.cancelled() => { + let _ = connection + .send_request(KillTerminalRequest::new(session_id.to_string(), terminal_id.clone())) + .block_task() + .await; + (true, None, None) + } + response = &mut wait => match response { + Ok(response) => ( + false, + response.exit_status.exit_code.and_then(|code| i32::try_from(code).ok()), + None, + ), + Err(error) => (false, None, Some(error.to_string())), + } + }; + + let output = connection + .send_request(TerminalOutputRequest::new( + session_id.to_string(), + terminal_id.clone(), + )) + .block_task() + .await; + let _ = connection + .send_request(ReleaseTerminalRequest::new( + session_id.to_string(), + terminal_id, + )) + .block_task() + .await; + + match (wait_error, output) { + (None, Ok(output)) => { + let result = crate::tools::terminal_session::external_terminal_result( + &start, + &output.output, + output.truncated, + exit_code, + stopped_by_user, + ); + let _ = control_tx.send(crate::tools::TerminalSessionControl::ExternalResult(result)); + } + (Some(error), _) => { + let _ = control_tx.send(crate::tools::TerminalSessionControl::ExternalError( + format!("ACP terminal wait failed: {error}"), + )); + } + (None, Err(error)) => { + let _ = control_tx.send(crate::tools::TerminalSessionControl::ExternalError( + format!("ACP terminal output failed: {error}"), + )); + } + } +} + fn permission_title(prompt: &crate::tools::PermissionPrompt) -> String { prompt .command @@ -1239,6 +2420,7 @@ fn send_tool_result( connection: &ConnectionTo, session_id: &str, result: crate::llm::ToolCallResult, + cwd: &Path, ) -> Result<(), Error> { let payload = serde_json::from_str::(&result.content).unwrap_or_else( |_| serde_json::json!({ "status": "error", "output_preview": result.content }), @@ -1247,16 +2429,15 @@ fn send_tool_result( Some("ok") => ToolCallStatus::Completed, _ => ToolCallStatus::Failed, }; - let text = payload - .get("output_preview") - .or_else(|| payload.get("error")) - .and_then(serde_json::Value::as_str) - .unwrap_or_default() - .to_string(); - let fields = ToolCallUpdateFields::new() + let content = tool_result_content(&payload, cwd); + let locations = tool_result_locations(&payload, cwd); + let mut fields = ToolCallUpdateFields::new() .status(status) - .content((!text.is_empty()).then(|| vec![ToolCallContent::from(text)])) + .content((!content.is_empty()).then_some(content)) .raw_output(payload); + if !locations.is_empty() { + fields = fields.locations(locations); + } let update = SessionUpdate::ToolCallUpdate(ToolCallUpdate::new(result.tool_call_id, fields)); connection .send_notification(SessionNotification::new(session_id.to_string(), update)) @@ -1268,23 +2449,54 @@ fn send_usage( session_id: &str, session: &AcpSession, used: usize, + cost: Option, + usage: Option, ) -> Result<(), Error> { - let Some(size) = session.context_window else { - return Ok(()); - }; - let update = SessionUpdate::UsageUpdate(UsageUpdate::new(used as u64, size as u64)); + let update = SessionUpdate::UsageUpdate(usage_update(session, used, cost, usage)); connection .send_notification(SessionNotification::new(session_id.to_string(), update)) .map_err(|_| internal_error()) } +fn usage_update( + session: &AcpSession, + used: usize, + cost: Option, + usage: Option, +) -> UsageUpdate { + let size = session + .context_window + .map(u64::from) + .unwrap_or_else(|| (used as u64).max(1)); + let mut meta = serde_json::Map::new(); + meta.insert( + "crabcode".to_string(), + serde_json::json!({ + "contextWindowKnown": session.context_window.is_some(), + "usage": usage.map(|usage| serde_json::json!({ + "inputTokens": usage.input_tokens, + "outputTokens": usage.output_tokens, + "cacheReadTokens": usage.cache_read_tokens, + "cacheWriteTokens": usage.cache_write_tokens, + })), + }), + ); + UsageUpdate::new(used as u64, size) + .cost(cost.map(|amount| AcpCost::new(amount, "USD"))) + .meta(meta) +} + fn replay_messages( connection: &ConnectionTo, session_id: &str, messages: &[crate::session::types::Message], + cwd: &Path, ) -> Result<(), Error> { - for (message_index, message) in messages.iter().enumerate() { - let message_id = format!("{session_id}:message:{message_index}"); + for message in messages { + if crate::session::compaction::is_compaction_display_item(message) { + continue; + } + let message_id = message.id.clone(); match message.role { crate::session::types::MessageRole::User => { if !message.content.is_empty() { @@ -1325,8 +2537,8 @@ fn replay_messages( )?; } } - "tool_call" => replay_tool_call(connection, session_id, part)?, - "tool_result" => replay_tool_result(connection, session_id, part)?, + "tool_call" => replay_tool_call(connection, session_id, part, cwd)?, + "tool_result" => replay_tool_result(connection, session_id, part, cwd)?, _ => {} } } @@ -1363,6 +2575,7 @@ fn replay_tool_call( connection: &ConnectionTo, session_id: &str, part: &crate::session::types::MessagePart, + cwd: &Path, ) -> Result<(), Error> { let Some(tool_call_id) = part.tool_id() else { return Ok(()); @@ -1383,6 +2596,7 @@ fn replay_tool_call( ToolCall::new(tool_call_id.to_string(), tool_title(name, &input)) .kind(tool_kind(name)) .status(status) + .locations(tool_locations(name, &input, cwd)) .raw_input(input), ); connection @@ -1394,6 +2608,7 @@ fn replay_tool_result( connection: &ConnectionTo, session_id: &str, part: &crate::session::types::MessagePart, + cwd: &Path, ) -> Result<(), Error> { let Some(tool_call_id) = part.tool_id() else { return Ok(()); @@ -1413,6 +2628,7 @@ fn replay_tool_result( name: part.tool_name().unwrap_or("tool").to_string(), content, }, + cwd, ) } @@ -1494,6 +2710,30 @@ mod tests { model } + fn test_session() -> AcpSession { + AcpSession { + cwd: PathBuf::from("/tmp"), + config: crate::config::configuration::LoadedConfig { + merged_config: crate::config::configuration::MergedConfig::default(), + raw_merged: serde_json::Value::Null, + diagnostics: Default::default(), + inventory: Default::default(), + project_root: PathBuf::from("/tmp"), + cwd: PathBuf::from("/tmp"), + xdg_config_home: PathBuf::from("/tmp"), + }, + skills: crate::skill::SkillStore::load(Path::new("/tmp"), Path::new("/tmp")), + models: vec![model("example", "Example", "chat", "Chat")], + provider: "example".to_string(), + model: "chat".to_string(), + agent: "Build".to_string(), + reasoning_selection: crate::model::reasoning::ReasoningEffort::High, + reasoning: None, + context_window: None, + cancellation: None, + } + } + #[test] fn maps_crabcode_tools_to_acp_kinds() { assert_eq!(tool_kind("bash"), ToolKind::Execute); @@ -1502,6 +2742,78 @@ mod tests { assert_eq!(tool_kind("unknown"), ToolKind::Other); } + #[test] + fn acp_usage_update_includes_cumulative_usd_cost() { + let usage = crate::aisdk::chunk::LanguageModelUsage { + input_tokens: 800, + output_tokens: 200, + cache_read_tokens: 500, + cache_write_tokens: 100, + }; + let update = usage_update(&test_session(), 1_000, Some(0.125), Some(usage)); + assert_eq!(update.cost.as_ref().map(|cost| cost.amount), Some(0.125)); + assert_eq!( + update.cost.as_ref().map(|cost| cost.currency.as_str()), + Some("USD") + ); + assert_eq!(update.size, 1_000); + let meta = update.meta.expect("usage metadata"); + assert_eq!(meta["crabcode"]["contextWindowKnown"], false); + assert_eq!(meta["crabcode"]["usage"]["inputTokens"], 800); + assert_eq!(meta["crabcode"]["usage"]["cacheWriteTokens"], 100); + } + + #[test] + fn session_info_includes_hierarchy_metadata() { + let now = std::time::SystemTime::now(); + let info = acp_session_info( + crate::session::manager::SessionInfo { + id: "child".to_string(), + parent_id: Some("parent".to_string()), + title: "Child".to_string(), + created_at: now, + updated_at: now, + message_count: 0, + workspace_id: 1, + workspace_path: "/tmp".to_string(), + workspace_name: "tmp".to_string(), + workspace_sort_order: 0, + status: crate::session::types::SessionStatus::Idle, + pinned_at: None, + archived_at: None, + }, + Some("root".to_string()), + ); + let meta = info.meta.expect("session metadata"); + assert_eq!(meta["crabcode"]["parentSessionId"], "parent"); + assert_eq!(meta["crabcode"]["rootSessionId"], "root"); + } + + #[test] + fn session_list_cursor_pages_all_sessions() { + assert_eq!( + session_page(None, 250).unwrap(), + (0, 100, Some("100".to_string())) + ); + assert_eq!( + session_page(Some("100"), 250).unwrap(), + (100, 200, Some("200".to_string())) + ); + assert_eq!(session_page(Some("200"), 250).unwrap(), (200, 250, None)); + assert!(session_page(Some("invalid"), 250).is_err()); + } + + #[test] + fn terminal_support_tracks_client_capability() { + let service = AcpService::new(Path::new("/tmp")).unwrap(); + assert!(!service.supports_terminals()); + + service.set_client_capabilities( + agent_client_protocol::schema::v1::ClientCapabilities::new().terminal(true), + ); + assert!(service.supports_terminals()); + } + #[test] fn builds_titles_from_tool_input() { assert_eq!( @@ -1514,6 +2826,365 @@ mod tests { ); } + #[test] + fn acp_tool_result_prefers_full_output() { + let payload = serde_json::json!({ + "output": "complete tool output", + "output_preview": "short preview", + }); + + assert_eq!(tool_result_text(&payload), "complete tool output"); + } + + #[test] + fn acp_tool_result_supports_legacy_preview_payloads() { + let payload = serde_json::json!({"output_preview": "legacy output"}); + + assert_eq!(tool_result_text(&payload), "legacy output"); + } + + #[test] + fn acp_tool_locations_normalize_multi_file_and_patch_paths() { + let cwd = Path::new("/tmp/workspace"); + let write_locations = tool_locations( + "write_files", + &serde_json::json!({ + "files": [ + {"file_path": "src/a.rs", "content": "a"}, + {"file_path": "/tmp/b.rs", "content": "b"} + ] + }), + cwd, + ); + assert_eq!( + write_locations[0].path, + PathBuf::from("/tmp/workspace/src/a.rs") + ); + assert_eq!(write_locations[1].path, PathBuf::from("/tmp/b.rs")); + + let patch_locations = tool_locations( + "apply_patch", + &serde_json::json!({ + "patch": "*** Begin Patch\n*** Update File: src/a.rs\n*** Add File: src/b.rs\n*** End Patch" + }), + cwd, + ); + assert_eq!(patch_locations.len(), 2); + assert_eq!( + patch_locations[1].path, + PathBuf::from("/tmp/workspace/src/b.rs") + ); + } + + #[test] + fn acp_tool_result_emits_diff_location_and_image_content() { + let payload = serde_json::json!({ + "output": "updated", + "metadata": { + "path": "src/main.rs", + "line_number": 4, + "old_text": "fn old() {}", + "new_text": "fn new() {}" + }, + "images": [{ + "data_url": "data:image/png;base64,aGk=", + "media_type": "image/png" + }] + }); + let cwd = Path::new("/tmp/workspace"); + + let locations = tool_result_locations(&payload, cwd); + assert_eq!( + locations[0].path, + PathBuf::from("/tmp/workspace/src/main.rs") + ); + assert_eq!(locations[0].line, Some(4)); + + let content = tool_result_content(&payload, cwd); + let diff = content.iter().find_map(|item| match item { + ToolCallContent::Diff(diff) => Some(diff), + _ => None, + }); + let diff = diff.expect("diff content"); + assert_eq!(diff.path, PathBuf::from("/tmp/workspace/src/main.rs")); + assert_eq!(diff.old_text.as_deref(), Some("fn old() {}")); + assert_eq!(diff.new_text, "fn new() {}"); + assert!(content.iter().any(|item| matches!( + item, + ToolCallContent::Content(content) + if matches!(&content.content, ContentBlock::Image(image) if image.data == "aGk=" && image.mime_type == "image/png") + ))); + } + + #[test] + fn acp_permission_prefers_originating_tool_call_id() { + assert_eq!(permission_tool_call_id(Some("call_123")), "call_123"); + } + + #[test] + fn acp_permission_generates_fallback_id_without_origin() { + assert!(permission_tool_call_id(None).starts_with("permission:")); + } + + #[test] + fn acp_permission_edit_includes_preflight_diff() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("main.rs"); + std::fs::write(&path, "fn old() {}\n").unwrap(); + let (response_tx, _response_rx) = tokio::sync::oneshot::channel(); + let prompt = crate::tools::PermissionPrompt { + tool_call_id: Some("call_edit".to_string()), + tool_id: "edit".to_string(), + action: crate::tools::PermissionAction::Write, + permission: "edit".to_string(), + patterns: vec![path.to_string_lossy().into_owned()], + target: Some(path.to_string_lossy().into_owned()), + command: None, + workdir: None, + workspace: dir.path().to_string_lossy().into_owned(), + reason: "approval".to_string(), + raw_input: serde_json::json!({ + "file_path": path, + "old_string": "old", + "new_string": "new", + "replace_all": false + }), + response_tx, + }; + + let content = permission_tool_content(&prompt, dir.path()); + let ToolCallContent::Diff(diff) = &content[0] else { + panic!("expected preflight diff"); + }; + assert_eq!(diff.old_text.as_deref(), Some("fn old() {}\n")); + assert_eq!(diff.new_text, "fn new() {}\n"); + } + + #[test] + fn acp_permission_apply_patch_includes_preflight_diff() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("note.txt"); + std::fs::write(&path, "before\n").unwrap(); + let (response_tx, _response_rx) = tokio::sync::oneshot::channel(); + let prompt = crate::tools::PermissionPrompt { + tool_call_id: Some("call_patch".to_string()), + tool_id: "apply_patch".to_string(), + action: crate::tools::PermissionAction::Write, + permission: "edit".to_string(), + patterns: vec![path.to_string_lossy().into_owned()], + target: Some(path.to_string_lossy().into_owned()), + command: None, + workdir: None, + workspace: dir.path().to_string_lossy().into_owned(), + reason: "approval".to_string(), + raw_input: serde_json::json!({ + "patch": format!( + "*** Begin Patch\n*** Update File: {}\n@@\n-before\n+after\n*** End Patch\n", + path.display() + ) + }), + response_tx, + }; + + let content = permission_tool_content(&prompt, dir.path()); + let ToolCallContent::Diff(diff) = &content[0] else { + panic!("expected apply_patch preflight diff"); + }; + assert_eq!(diff.path, path); + assert_eq!(diff.old_text.as_deref(), Some("before\n")); + assert_eq!(diff.new_text, "after\n"); + } + + #[test] + fn acp_question_form_preserves_single_multi_custom_and_scope() { + let form = acp_question_form( + "session_1", + Some("question_call_1"), + &serde_json::json!([ + { + "question": "Pick one", + "header": "Single", + "options": [ + {"label": "A", "description": "First"}, + {"label": "B", "description": "Second"} + ] + }, + { + "question": "Pick several", + "header": "Multiple", + "multiple": true, + "options": [ + {"label": "X", "description": "First"}, + {"label": "Y", "description": "Second"} + ] + } + ]), + ); + let wire = serde_json::to_value(&form.request).expect("elicitation request"); + + assert_eq!(wire["mode"], "form"); + assert_eq!(wire["sessionId"], "session_1"); + assert_eq!(wire["toolCallId"], "question_call_1"); + assert_eq!( + wire["requestedSchema"]["properties"]["question_0"]["type"], + "string" + ); + assert_eq!( + wire["requestedSchema"]["properties"]["question_0"]["oneOf"][0]["title"], + "A" + ); + assert_eq!( + wire["requestedSchema"]["properties"]["question_1"]["type"], + "array" + ); + assert_eq!( + wire["requestedSchema"]["properties"]["question_1_custom"]["type"], + "string" + ); + } + + #[test] + fn acp_question_answers_restore_labels_and_custom_text() { + let form = acp_question_form( + "session_1", + None, + &serde_json::json!([ + { + "question": "Pick one", + "options": [{"label": "A"}, {"label": "B"}] + }, + { + "question": "Pick several", + "multiple": true, + "options": [{"label": "X"}, {"label": "Y"}] + } + ]), + ); + let mut content = std::collections::BTreeMap::new(); + content.insert( + "question_0".to_string(), + ElicitationContentValue::String("q0_option_1".to_string()), + ); + content.insert( + "question_1".to_string(), + ElicitationContentValue::StringArray(vec![ + "q1_option_0".to_string(), + "q1_option_1".to_string(), + ]), + ); + content.insert( + "question_1_custom".to_string(), + ElicitationContentValue::String("Other choice".to_string()), + ); + let action = ElicitationAction::Accept( + agent_client_protocol::schema::v1::ElicitationAcceptAction::new().content(content), + ); + + assert_eq!( + acp_question_answers(&form.fields, action), + serde_json::json!([["B"], ["X", "Y", "Other choice"]]) + ); + assert_eq!( + acp_question_answers(&form.fields, ElicitationAction::Cancel), + serde_json::json!([[], []]) + ); + + let mut custom_content = std::collections::BTreeMap::new(); + custom_content.insert( + "question_0".to_string(), + ElicitationContentValue::String("q0_option_0".to_string()), + ); + custom_content.insert( + "question_0_custom".to_string(), + ElicitationContentValue::String("Custom only".to_string()), + ); + let custom_action = ElicitationAction::Accept( + agent_client_protocol::schema::v1::ElicitationAcceptAction::new() + .content(custom_content), + ); + assert_eq!( + acp_question_answers(&form.fields, custom_action), + serde_json::json!([["Custom only"], []]) + ); + } + + #[test] + fn maps_typed_turn_stop_reasons_to_acp() { + assert_eq!( + acp_stop_reason(Some(crate::llm::TurnStopReason::MaxTokens)), + StopReason::MaxTokens + ); + assert_eq!( + acp_stop_reason(Some(crate::llm::TurnStopReason::Refusal)), + StopReason::Refusal + ); + assert_eq!( + acp_stop_reason(Some(crate::llm::TurnStopReason::MaxTurnRequests)), + StopReason::MaxTurnRequests + ); + assert_eq!( + acp_stop_reason(Some(crate::llm::TurnStopReason::MaxTurnRequests)), + StopReason::MaxTurnRequests + ); + assert_eq!(acp_stop_reason(None), StopReason::EndTurn); + } + + #[test] + fn recognizes_only_exact_compact_control_command() { + assert_eq!(compact_command("/compact").unwrap(), true); + assert_eq!(compact_command(" /compact ").unwrap(), true); + assert_eq!(compact_command("/compactness").unwrap(), false); + assert_eq!(compact_command("hello").unwrap(), false); + assert!(compact_command("/compact extra").is_err()); + } + + #[test] + fn advertises_compact_as_no_input_command() { + let session = test_session(); + let command = available_commands(&session) + .into_iter() + .find(|command| command.name == "compact") + .expect("compact command"); + assert_eq!( + command.description, + "Summarize this session to reduce context" + ); + assert!(command.input.is_none()); + } + + #[test] + fn builds_smaller_soft_compaction_for_acp() { + let session = test_session(); + let messages = vec![ + crate::session::types::Message::user("u".repeat(8_000)), + crate::session::types::Message::assistant("a".repeat(8_000)), + crate::session::types::Message::user("recent"), + ]; + let selection = crate::session::compaction::select_messages_for_compaction_with_min( + &messages, + crate::session::compaction::DEFAULT_TAIL_TURNS, + 0, + ) + .expect("compaction selection"); + let before_tokens = crate::session::compaction::total_context_tokens(&messages); + let before_messages = + crate::session::compaction::filter_messages_for_context(&messages).len(); + + let (compacted, stats) = compacted_messages( + &messages, + &selection, + "short handoff", + &session, + before_tokens, + before_messages, + ) + .expect("smaller compaction"); + + assert!(stats.after_tokens < stats.before_tokens); + assert!(crate::session::compaction::latest_compaction_stats(&compacted).is_some()); + assert_eq!(compacted[0].id, messages[0].id); + } + #[test] fn flattens_text_and_embedded_context() { let text = prompt_text(vec![ @@ -1533,25 +3204,87 @@ mod tests { } #[test] - fn writes_supported_acp_image_to_temp_file() { + fn writes_supported_acp_image_to_managed_session_storage() { + let session_id = format!("acp-image-{}", cuid2::create_id()); let image = agent_client_protocol::schema::v1::ImageContent::new("aGk=", "image/png"); - let path = write_prompt_image(&image).expect("image file"); + let path = write_prompt_image(&session_id, &image).expect("image file"); + assert!(Path::new(&path).starts_with(crate::persistence::attachments::root_dir())); assert_eq!(std::fs::read(&path).expect("image bytes"), b"hi"); - let _ = std::fs::remove_file(path); + crate::persistence::attachments::cleanup_session(&session_id).unwrap(); } #[test] - fn writes_acp_clipboard_image_data_uri_to_temp_file() { + fn writes_supported_acp_audio_to_managed_session_storage() { + let session_id = format!("acp-audio-{}", cuid2::create_id()); + let audio = agent_client_protocol::schema::v1::AudioContent::new("YXVkaW8=", "audio/wav"); + let path = write_prompt_audio(&session_id, &audio).expect("audio file"); + + assert!(path.ends_with(".wav")); + assert_eq!(std::fs::read(&path).unwrap(), b"audio"); + crate::persistence::attachments::cleanup_session(&session_id).unwrap(); + } + + #[test] + fn rejects_audio_for_models_without_audio_modality() { + let session_id = format!("acp-audio-{}", cuid2::create_id()); + let result = prompt_content( + vec![ContentBlock::Audio( + agent_client_protocol::schema::v1::AudioContent::new("YXVkaW8=", "audio/wav"), + )], + false, + false, + &session_id, + &test_session(), + ); + + assert!(result.is_err()); + crate::persistence::attachments::cleanup_session(&session_id).unwrap(); + } + + #[test] + fn writes_acp_clipboard_image_data_uri_to_managed_storage() { + let session_id = format!("acp-image-{}", cuid2::create_id()); let image = agent_client_protocol::schema::v1::ImageContent::new( "data:image/png;base64,aGk=", "application/octet-stream", ); - let path = write_prompt_image(&image).expect("image file"); + let path = write_prompt_image(&session_id, &image).expect("image file"); assert!(path.ends_with(".png")); assert_eq!(std::fs::read(&path).expect("image bytes"), b"hi"); - let _ = std::fs::remove_file(path); + crate::persistence::attachments::cleanup_session(&session_id).unwrap(); + } + + #[test] + fn prompt_image_failure_rolls_back_prior_managed_files() { + let session_id = format!("acp-image-{}", cuid2::create_id()); + let result = prompt_content( + vec![ + ContentBlock::Image(agent_client_protocol::schema::v1::ImageContent::new( + "aGk=", + "image/png", + )), + ContentBlock::Image(agent_client_protocol::schema::v1::ImageContent::new( + "not-base64", + "image/png", + )), + ], + true, + false, + &session_id, + &test_session(), + ); + + assert!(result.is_err()); + let directory = crate::persistence::attachments::session_dir(&session_id).unwrap(); + assert!( + !directory.exists() + || std::fs::read_dir(&directory) + .unwrap() + .all(|entry| entry.is_err()) + ); + crate::persistence::attachments::cleanup_session(&session_id).unwrap(); } #[test] @@ -1561,14 +3294,14 @@ mod tests { "image/png", ); - assert!(write_prompt_image(&image).is_err()); + assert!(write_prompt_image("test", &image).is_err()); } #[test] fn rejects_unsupported_acp_image_mime_type() { let image = agent_client_protocol::schema::v1::ImageContent::new("aGk=", "image/tiff"); - assert!(write_prompt_image(&image).is_err()); + assert!(write_prompt_image("test", &image).is_err()); } fn config_with_command(command: crate::command::custom::CustomCommand) -> LoadedConfig { @@ -1633,14 +3366,77 @@ mod tests { )); } + #[test] + fn acp_tool_result_restores_native_mcp_resources() { + let metadata = serde_json::json!({ + "mcp_result": { + "content": [ + { + "type": "image", + "data": "aGk=", + "mimeType": "image/png", + "annotations": { "priority": 0.5 } + }, + { + "type": "resource_link", + "uri": "file:///tmp/readme.md", + "name": "readme", + "mimeType": "text/markdown", + "annotations": { "priority": 0.8 } + } + ] + } + }); + + let content = mcp_native_tool_content(&metadata); + assert_eq!(content.len(), 2); + assert!(matches!( + &content[0], + ToolCallContent::Content(content) + if matches!(&content.content, ContentBlock::Image(image) + if image.data == "aGk=" && image.annotations.is_some()) + )); + assert!(matches!( + &content[1], + ToolCallContent::Content(content) + if matches!(&content.content, ContentBlock::ResourceLink(link) + if link.uri == "file:///tmp/readme.md") + )); + } + + #[test] + fn acp_tool_result_restores_mcp_audio_and_preserves_unknown_blocks() { + let metadata = serde_json::json!({ + "mcp_result": { + "content": [ + { "type": "audio", "data": "YXVkaW8=", "mimeType": "audio/wav" }, + { "type": "future_media", "payload": { "value": 1 } } + ] + } + }); + let content = mcp_native_tool_content(&metadata); + assert!(matches!( + &content[0], + ToolCallContent::Content(content) + if matches!(&content.content, ContentBlock::Audio(audio) + if audio.data == "YXVkaW8=" && audio.mime_type == "audio/wav") + )); + assert!(matches!( + &content[1], + ToolCallContent::Content(content) + if matches!(&content.content, ContentBlock::Text(text) + if text.text.contains("future_media") && text.text.contains("payload")) + )); + } + #[tokio::test] async fn expands_custom_slash_command_before_prompting() { let config = config_with_command(crate::command::custom::CustomCommand { name: "review".to_string(), description: None, template: "Review this carefully: $ARGUMENTS".to_string(), - agent: None, - model: None, + agent: Some("build".to_string()), + model: Some("openai/gpt-5".to_string()), subtask: Some(false), source: crate::command::custom::CustomCommandSource::Config(PathBuf::from( "/tmp/opencode.jsonc", @@ -1653,7 +3449,14 @@ mod tests { .await .expect("expanded command"); - assert_eq!(prompt, "Review this carefully: src/acp/service.rs"); + assert_eq!( + prompt, + SlashExpansion::Prompt { + prompt: "Review this carefully: src/acp/service.rs".to_string(), + agent: Some("build".to_string()), + model: Some("openai/gpt-5".to_string()), + } + ); } #[tokio::test] @@ -1688,6 +3491,9 @@ mod tests { let prompt = expand_slash_command(&session, "/reviewer src/lib.rs") .await .expect("expanded skill"); + let SlashExpansion::Prompt { prompt, .. } = prompt else { + panic!("expected skill prompt"); + }; assert!(prompt.contains("Inspect correctness and risks.")); assert!(prompt.contains("src/lib.rs")); } @@ -1723,7 +3529,10 @@ mod tests { let prompt = expand_slash_command(&session, "/mcp") .await .expect("mcp status"); - assert!(prompt.contains("filesystem (local, enabled)")); + let SlashExpansion::LocalResult(prompt) = prompt else { + panic!("expected local MCP result"); + }; + assert!(prompt.contains("filesystem (local, connecting)")); } #[test] @@ -1810,28 +3619,7 @@ mod tests { #[test] fn preserves_selected_reasoning_effort_when_model_cannot_apply_it() { - let model = model("example", "Example", "chat", "Chat"); - let session = AcpSession { - cwd: PathBuf::from("/tmp"), - config: crate::config::configuration::LoadedConfig { - merged_config: crate::config::configuration::MergedConfig::default(), - raw_merged: serde_json::Value::Null, - diagnostics: Default::default(), - inventory: Default::default(), - project_root: PathBuf::from("/tmp"), - cwd: PathBuf::from("/tmp"), - xdg_config_home: PathBuf::from("/tmp"), - }, - skills: crate::skill::SkillStore::load(Path::new("/tmp"), Path::new("/tmp")), - models: vec![model], - provider: "example".to_string(), - model: "chat".to_string(), - agent: "Build".to_string(), - reasoning_selection: crate::model::reasoning::ReasoningEffort::High, - reasoning: None, - context_window: None, - cancellation: None, - }; + let session = test_session(); let option = reasoning_config_option(&session); assert_eq!(option.id.to_string(), "effort"); assert_eq!(option.name, "Effort"); diff --git a/src/aisdk/README.md b/src/aisdk/README.md index 8ba7f3cf..64ac4eda 100644 --- a/src/aisdk/README.md +++ b/src/aisdk/README.md @@ -37,5 +37,8 @@ Done for packaging/host hooks: - Absolute `crate::{chunk,error,...}` paths are crate-root-shaped (host re-exports them today) - Product-leaky debug path renamed/feature-gated - Product-flavored comments/tests scrubbed +- Typed terminal stop reasons include normal completion, max tokens, refusal, hooks, and errors +- Normalized provider usage events retain input, output, cache-read, and cache-write token accounting across multi-step turns +- User messages support typed image and WAV/MP3 audio inputs; audio is serialized through verified Chat Completions `input_audio` content parts Keep app glue outside this tree (`src/tools/aisdk_bridge.rs`, `src/llm/*`). diff --git a/src/aisdk/chunk.rs b/src/aisdk/chunk.rs index c062ccd7..faaf430e 100644 --- a/src/aisdk/chunk.rs +++ b/src/aisdk/chunk.rs @@ -20,6 +20,7 @@ pub enum ChunkType { end_turn: Option, reasoning_items: Vec, doom_loop_triggers: Vec, + usage: Option, }, Retry(crate::retry::RetryStatus), StreamRollback { @@ -28,6 +29,8 @@ pub enum ChunkType { }, Warning(String), Metadata(String), + /// Provider-reported token usage for one model request. + Usage(LanguageModelUsage), End { reason: Option, }, @@ -37,6 +40,36 @@ pub enum ChunkType { NotSupported(String), } +/// Normalized provider usage. `input_tokens` includes cached input; cache +/// fields describe subsets used for pricing and observability. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct LanguageModelUsage { + pub input_tokens: u64, + pub output_tokens: u64, + pub cache_read_tokens: u64, + pub cache_write_tokens: u64, +} + +impl LanguageModelUsage { + pub fn is_empty(self) -> bool { + self.input_tokens == 0 + && self.output_tokens == 0 + && self.cache_read_tokens == 0 + && self.cache_write_tokens == 0 + } +} + +impl std::ops::AddAssign for LanguageModelUsage { + fn add_assign(&mut self, rhs: Self) { + self.input_tokens = self.input_tokens.saturating_add(rhs.input_tokens); + self.output_tokens = self.output_tokens.saturating_add(rhs.output_tokens); + self.cache_read_tokens = self.cache_read_tokens.saturating_add(rhs.cache_read_tokens); + self.cache_write_tokens = self + .cache_write_tokens + .saturating_add(rhs.cache_write_tokens); + } +} + #[derive(Debug, Clone, PartialEq, Eq, Default)] pub struct ReasoningReplayItem { pub id: Option, @@ -58,6 +91,7 @@ impl ChunkType { end_turn, reasoning_items: Vec::new(), doom_loop_triggers: Vec::new(), + usage: None, } } } @@ -88,6 +122,7 @@ impl FinishReason { "tool_calls" | "function_call" => Self::ToolCalls, "length" => Self::Length, "content_filter" => Self::ContentFilter, + "refusal" => Self::Refusal, other => Self::Unknown(other.to_string()), } } @@ -126,3 +161,16 @@ impl FinishReason { matches!(self, Self::Stop | Self::StopSequence) } } + +#[cfg(test)] +mod tests { + use super::FinishReason; + + #[test] + fn compatible_refusal_is_typed() { + assert_eq!( + FinishReason::from_openai_compatible("refusal"), + FinishReason::Refusal + ); + } +} diff --git a/src/aisdk/message.rs b/src/aisdk/message.rs index ffdfc222..50459f92 100644 --- a/src/aisdk/message.rs +++ b/src/aisdk/message.rs @@ -5,6 +5,13 @@ pub(crate) fn is_prefixed_response_item_id(id: &str) -> bool { .is_some_and(|(prefix, suffix)| !prefix.is_empty() && !suffix.is_empty()) } +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AudioContent { + pub data: String, + pub format: String, + pub media_type: String, +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(tag = "role")] pub enum Message { @@ -35,6 +42,7 @@ impl Message { Self::User(UserMessage { content: content.into(), images: Vec::new(), + audios: Vec::new(), }) } @@ -42,6 +50,19 @@ impl Message { Self::User(UserMessage { content: content.into(), images, + audios: Vec::new(), + }) + } + + pub fn user_with_attachments( + content: impl Into, + images: Vec, + audios: Vec, + ) -> Self { + Self::User(UserMessage { + content: content.into(), + images, + audios, }) } @@ -159,6 +180,8 @@ pub struct UserMessage { pub content: String, #[serde(default)] pub images: Vec, + #[serde(default)] + pub audios: Vec, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -231,6 +254,7 @@ impl From for UserMessage { Self { content, images: Vec::new(), + audios: Vec::new(), } } } @@ -240,6 +264,7 @@ impl From<&str> for UserMessage { Self { content: content.to_string(), images: Vec::new(), + audios: Vec::new(), } } } diff --git a/src/aisdk/providers/anthropic.rs b/src/aisdk/providers/anthropic.rs index c34c5858..be1689cc 100644 --- a/src/aisdk/providers/anthropic.rs +++ b/src/aisdk/providers/anthropic.rs @@ -20,12 +20,47 @@ pub struct Anthropic { reasoning_effort: Option, } +fn anthropic_input_usage(usage: &serde_json::Value) -> Option { + let mut normalized = anthropic_usage(usage)?; + normalized.output_tokens = 0; + Some(normalized) +} + +fn anthropic_output_usage(usage: &serde_json::Value) -> Option { + let output_tokens = usage + .get("output_tokens") + .and_then(|value| value.as_u64())?; + Some(crate::chunk::LanguageModelUsage { + output_tokens, + ..Default::default() + }) +} + impl Anthropic { pub fn builder() -> AnthropicBuilder { AnthropicBuilder::default() } } +fn anthropic_stream_chunks(event_type: &str, value: &serde_json::Value) -> Vec> { + let usage = match event_type { + "message_start" => value + .get("message") + .and_then(|message| message.get("usage")) + .and_then(anthropic_input_usage), + "message_delta" => value.get("usage").and_then(anthropic_output_usage), + _ => None, + }; + let mut chunks = Vec::new(); + if let Some(usage) = usage { + chunks.push(Ok(ChunkType::Usage(usage))); + } + if let Some(chunk) = anthropic_stream_chunk(event_type, value) { + chunks.push(chunk); + } + chunks +} + #[derive(Default)] pub struct AnthropicBuilder { base_url: Option, @@ -199,29 +234,25 @@ impl Provider for Anthropic { let stream = response .bytes_stream() .eventsource() - .filter_map(|ev| match ev { + .map(|ev| match ev { Ok(event) => { let event_type = event.event.as_str(); let data = &event.data; if data.is_empty() { - return futures::future::ready(None); + return Vec::new(); } match serde_json::from_str::(data) { - Ok(value) => { - futures::future::ready(anthropic_stream_chunk(event_type, &value)) - } - Err(e) => futures::future::ready(Some(Ok(ChunkType::Failed(format!( - "Invalid SSE data: {}", - e - ))))), + Ok(value) => anthropic_stream_chunks(event_type, &value), + Err(e) => vec![Ok(ChunkType::Failed(format!("Invalid SSE data: {}", e)))], } } - Err(e) => futures::future::ready(Some(Ok(ChunkType::RetryableFailure( - RetryError::from_message(format!("SSE error: {}", e)), - )))), + Err(e) => vec![Ok(ChunkType::RetryableFailure(RetryError::from_message( + format!("SSE error: {}", e), + )))], }) + .flat_map(futures::stream::iter) .boxed(); Ok(stream) @@ -233,13 +264,7 @@ fn anthropic_stream_chunk( value: &serde_json::Value, ) -> Option> { match event_type { - "message_start" => { - // Partial usage early in the stream (cache fields may already appear). - if let Some(usage) = value.get("message").and_then(|m| m.get("usage")) { - log_anthropic_usage(usage); - } - None - } + "message_start" => None, "content_block_start" => { if let Some(payload) = anthropic_hosted_search_start(value) { Some(Ok(ChunkType::ProviderToolCall(payload))) @@ -252,13 +277,7 @@ fn anthropic_stream_chunk( } } "content_block_delta" => anthropic_content_block_delta(value).map(Ok), - "message_delta" => { - // Final usage wins for cache_read / cache_creation. - if let Some(usage) = value.get("usage") { - log_anthropic_usage(usage); - } - anthropic_message_delta(value).map(Ok) - } + "message_delta" => anthropic_message_delta(value).map(Ok), "message_stop" => Some(Ok(ChunkType::End { reason: None })), "error" => { let error_msg = value["error"]["message"] @@ -272,7 +291,7 @@ fn anthropic_stream_chunk( /// Log Anthropic usage via the host logger so cache hits are verifiable. /// Note: `input_tokens` is non-cached only; total input ≈ input + cache_read + cache_creation. -fn log_anthropic_usage(usage: &serde_json::Value) { +fn anthropic_usage(usage: &serde_json::Value) -> Option { let input = usage.get("input_tokens").and_then(|v| v.as_u64()); let output = usage.get("output_tokens").and_then(|v| v.as_u64()); let cache_read = usage @@ -286,7 +305,7 @@ fn log_anthropic_usage(usage: &serde_json::Value) { // Skip empty/partial early frames with no signal. if input.is_none() && output.is_none() && cache_read == 0 && cache_creation == 0 { - return; + return None; } let input_v = input.unwrap_or(0); @@ -308,6 +327,12 @@ fn log_anthropic_usage(usage: &serde_json::Value) { total_input, hit_pct )); + Some(crate::chunk::LanguageModelUsage { + input_tokens: total_input, + output_tokens: output.unwrap_or(0), + cache_read_tokens: cache_read, + cache_write_tokens: cache_creation, + }) } fn anthropic_content_block_delta(value: &serde_json::Value) -> Option { @@ -337,13 +362,9 @@ fn anthropic_message_delta(value: &serde_json::Value) -> Option { .and_then(|delta| delta.get("stop_reason")) .and_then(|stop_reason| stop_reason.as_str())?; - match stop_reason { - "max_tokens" => Some(ChunkType::Incomplete("stop_reason=max_tokens".to_string())), - "refusal" => Some(ChunkType::Failed("stop_reason=refusal".to_string())), - reason => Some(ChunkType::End { - reason: Some(FinishReason::from_anthropic(reason)), - }), - } + Some(ChunkType::End { + reason: Some(FinishReason::from_anthropic(stop_reason)), + }) } fn anthropic_hosted_search_start(value: &serde_json::Value) -> Option { @@ -810,7 +831,7 @@ mod tests { } #[test] - fn max_tokens_stop_reason_emits_incomplete_chunk() { + fn max_tokens_stop_reason_emits_terminal_reason() { let value = serde_json::json!({ "type": "message_delta", "delta": { @@ -821,7 +842,57 @@ mod tests { .expect("event should produce a chunk") .expect("chunk should parse"); - assert!(matches!(chunk, ChunkType::Incomplete(_))); + assert!(matches!( + chunk, + ChunkType::End { + reason: Some(FinishReason::Length) + } + )); + } + + #[test] + fn message_delta_emits_usage_and_terminal_reason() { + let chunks = anthropic_stream_chunks( + "message_delta", + &serde_json::json!({ + "usage": { "output_tokens": 40 }, + "delta": { "stop_reason": "end_turn" } + }), + ); + + assert!(matches!( + chunks.first(), + Some(Ok(ChunkType::Usage(crate::chunk::LanguageModelUsage { + output_tokens: 40, + .. + }))) + )); + assert!(matches!( + chunks.get(1), + Some(Ok(ChunkType::End { + reason: Some(FinishReason::EndTurn) + })) + )); + } + + #[test] + fn refusal_stop_reason_emits_terminal_reason() { + let value = serde_json::json!({ + "type": "message_delta", + "delta": { + "stop_reason": "refusal", + }, + }); + let chunk = anthropic_stream_chunk("message_delta", &value) + .expect("event should produce a chunk") + .expect("chunk should parse"); + + assert!(matches!( + chunk, + ChunkType::End { + reason: Some(FinishReason::Refusal) + } + )); } #[test] diff --git a/src/aisdk/providers/compatible.rs b/src/aisdk/providers/compatible.rs index bbd0aa32..a0415c04 100644 --- a/src/aisdk/providers/compatible.rs +++ b/src/aisdk/providers/compatible.rs @@ -238,7 +238,7 @@ impl Provider for OpenAICompatible { } fn openai_compatible_user_content(user: &crate::message::UserMessage) -> serde_json::Value { - if user.images.is_empty() { + if user.images.is_empty() && user.audios.is_empty() { return serde_json::json!(user.content); } @@ -257,6 +257,15 @@ fn openai_compatible_user_content(user: &crate::message::UserMessage) -> serde_j }, }) })); + parts.extend(user.audios.iter().map(|audio| { + serde_json::json!({ + "type": "input_audio", + "input_audio": { + "data": audio.data, + "format": audio.format, + }, + }) + })); serde_json::Value::Array(parts) } @@ -436,9 +445,15 @@ fn debug_log(msg: &str) { /// Log OpenAI-compatible / AI Gateway usage via the host logger. /// Looks for `prompt_tokens_details.cached_tokens` and Anthropic-style fields /// that some gateways forward. -fn log_openai_compatible_usage(usage: &serde_json::Value) { - let prompt = usage.get("prompt_tokens").and_then(|v| v.as_u64()); - let completion = usage.get("completion_tokens").and_then(|v| v.as_u64()); +fn openai_compatible_usage(usage: &serde_json::Value) -> Option { + let prompt = usage + .get("prompt_tokens") + .or_else(|| usage.get("input_tokens")) + .and_then(|v| v.as_u64()); + let completion = usage + .get("completion_tokens") + .or_else(|| usage.get("output_tokens")) + .and_then(|v| v.as_u64()); let cached = usage .pointer("/prompt_tokens_details/cached_tokens") .and_then(|v| v.as_u64()) @@ -450,6 +465,7 @@ fn log_openai_compatible_usage(usage: &serde_json::Value) { .unwrap_or(0); let cache_creation = usage .get("cache_creation_input_tokens") + .or_else(|| usage.get("cache_write_input_tokens")) .and_then(|v| v.as_u64()) .unwrap_or(0); @@ -459,7 +475,7 @@ fn log_openai_compatible_usage(usage: &serde_json::Value) { && cache_read == 0 && cache_creation == 0 { - return; + return None; } // Prefer OpenAI-style cached_tokens; fall back to Anthropic-style cache_read. @@ -489,10 +505,27 @@ fn log_openai_compatible_usage(usage: &serde_json::Value) { cache_creation, hit_pct )); + + let anthropic_shape = cached == 0 && (cache_read > 0 || cache_creation > 0); + let input_tokens = if anthropic_shape { + prompt + .unwrap_or(0) + .saturating_add(cache_read) + .saturating_add(cache_creation) + } else { + prompt.unwrap_or(0) + }; + Some(crate::chunk::LanguageModelUsage { + input_tokens, + output_tokens: completion.unwrap_or(0), + cache_read_tokens: effective_cached, + cache_write_tokens: cache_creation, + }) } fn process_sse_data(data: &str) -> Vec> { let data = data.trim(); + let mut chunks = Vec::new(); if data == "[DONE]" { debug_log("[SSE] Terminal: [DONE]"); @@ -501,7 +534,7 @@ fn process_sse_data(data: &str) -> Vec> { if data.is_empty() || is_sse_metadata_line(data) { debug_log("[SSE] Ignored: empty or metadata/comment"); - return vec![]; + return chunks; } debug_log(&format!("[SSE] Raw data: {}", data)); @@ -525,8 +558,9 @@ fn process_sse_data(data: &str) -> Vec> { // Final usage often arrives on a choices-empty (or choices-missing) chunk. // Log cache-related fields so gateway Anthropic hits are verifiable. - if let Some(usage) = value.get("usage") { - log_openai_compatible_usage(usage); + let usage = value.get("usage").and_then(openai_compatible_usage); + if let Some(usage) = usage { + chunks.push(Ok(ChunkType::Usage(usage))); } let Some(choices) = value["choices"].as_array() else { @@ -534,18 +568,16 @@ fn process_sse_data(data: &str) -> Vec> { "[SSE] No choices array. JSON keys: {:?}", value.as_object().map(|o| o.keys().collect::>()) )); - return vec![]; + return chunks; }; if choices.is_empty() { debug_log("[SSE] choices array is empty"); - return vec![]; + return chunks; } let choice = &choices[0]; let finish_reason = choice["finish_reason"].as_str().unwrap_or(""); - let mut chunks = Vec::new(); - // Log the full choice structure for debugging debug_log(&format!( "[SSE] Choice JSON: {}", @@ -604,12 +636,6 @@ fn process_sse_data(data: &str) -> Vec> { match finish_reason { "" => {} - "length" => chunks.push(Ok(ChunkType::Incomplete( - "finish_reason=length".to_string(), - ))), - "content_filter" => chunks.push(Ok(ChunkType::Failed( - "finish_reason=content_filter".to_string(), - ))), _ => chunks.push(Ok(ChunkType::End { reason: Some(FinishReason::from_openai_compatible(finish_reason)), })), @@ -651,6 +677,60 @@ mod tests { assert!(provider.api_key.is_empty()); } + #[test] + fn serializes_audio_input_content_part() { + let user = crate::message::UserMessage { + content: "Describe this".to_string(), + images: Vec::new(), + audios: vec![crate::message::AudioContent { + data: "YXVkaW8=".to_string(), + format: "wav".to_string(), + media_type: "audio/wav".to_string(), + }], + }; + + assert_eq!( + openai_compatible_user_content(&user), + serde_json::json!([ + {"type": "text", "text": "Describe this"}, + {"type": "input_audio", "input_audio": {"data": "YXVkaW8=", "format": "wav"}} + ]) + ); + } + + #[test] + fn responses_style_usage_aliases_are_normalized() { + let chunks = process_sse_data( + r#"{"choices":[],"usage":{"input_tokens":120,"output_tokens":30,"cached_tokens":80,"cache_write_input_tokens":10}}"#, + ); + assert!(matches!( + chunks.as_slice(), + [Ok(ChunkType::Usage(crate::chunk::LanguageModelUsage { + input_tokens: 120, + output_tokens: 30, + cache_read_tokens: 80, + cache_write_tokens: 10, + }))] + )); + } + + #[test] + fn usage_only_chunk_emits_normalized_usage() { + let chunks = process_sse_data( + r#"{"choices":[],"usage":{"prompt_tokens":120,"completion_tokens":30,"prompt_tokens_details":{"cached_tokens":80}}}"#, + ); + + assert!(matches!( + chunks.as_slice(), + [Ok(ChunkType::Usage(crate::chunk::LanguageModelUsage { + input_tokens: 120, + output_tokens: 30, + cache_read_tokens: 80, + cache_write_tokens: 0, + }))] + )); + } + #[test] fn emits_tool_call_delta_without_finish_reason() { let data = r#"{"choices":[{"index":0,"delta":{"tool_calls":[{"id":"tool-1","index":0,"type":"function","function":{"name":"question","arguments":"{\"questions\":[{\"header\":\"Hobbies\",\"options\":[]}]}"}}]}}]}"#; @@ -780,14 +860,17 @@ mod tests { } #[test] - fn length_finish_reason_emits_incomplete_chunk() { + fn length_finish_reason_emits_terminal_reason() { let data = r#"{"choices":[{"index":0,"finish_reason":"length","delta":{"role":"assistant","content":""}}]}"#; let chunks = process_sse_data(data); - assert!(chunks - .iter() - .any(|chunk| matches!(chunk, Ok(ChunkType::Incomplete(_))))); + assert!(chunks.iter().any(|chunk| matches!( + chunk, + Ok(ChunkType::End { + reason: Some(FinishReason::Length) + }) + ))); } #[test] diff --git a/src/aisdk/providers/openai.rs b/src/aisdk/providers/openai.rs index 45e60fe1..8d5d8393 100644 --- a/src/aisdk/providers/openai.rs +++ b/src/aisdk/providers/openai.rs @@ -43,6 +43,25 @@ pub trait HttpResponseRetryPolicy: Send + Sync + std::fmt::Debug { ) -> Option; } +fn responses_incomplete_chunk(value: &serde_json::Value) -> ChunkType { + let reason = value + .get("response") + .and_then(|response| response.get("incomplete_details")) + .and_then(|details| details.get("reason")) + .and_then(serde_json::Value::as_str); + match reason { + Some("max_output_tokens" | "max_tokens") => ChunkType::End { + reason: Some(crate::chunk::FinishReason::Length), + }, + Some("content_filter" | "refusal" | "safety") => ChunkType::End { + reason: Some(crate::chunk::FinishReason::Refusal), + }, + _ => ChunkType::RetryableFailure(RetryError::from_message(responses_incomplete_message( + value, + ))), + } +} + #[derive(Debug, Clone)] pub struct OpenAI { base_url: String, @@ -1545,14 +1564,13 @@ fn response_sse_data_to_chunk(data: &str) -> Option> { return Some(Ok(responses_error_chunk(&value, event_type))); } let resp = &value["response"]; - if let Some(usage) = resp.get("usage") { - log_openai_responses_usage(usage); - } + let usage = resp.get("usage").and_then(openai_responses_usage); log_openai_responses_completed(resp); Some(Ok(ChunkType::ResponseCompleted { end_turn: resp.get("end_turn").and_then(|value| value.as_bool()), reasoning_items: reasoning_items_from_response_output(resp), doom_loop_triggers: doom_loop_triggers_from(resp), + usage, })) } // Grok Build / cli-chat-proxy: `response.doom_loop_check` with @@ -1565,9 +1583,7 @@ fn response_sse_data_to_chunk(data: &str) -> Option> { "doom_loop_check triggers={triggers}" )))) } - "response.incomplete" => Some(Ok(ChunkType::RetryableFailure(RetryError::from_message( - responses_incomplete_message(&value), - )))), + "response.incomplete" => Some(Ok(responses_incomplete_chunk(&value))), "response.failed" | "error" => Some(Ok(responses_error_chunk(&value, event_type))), _ => { if let Some(reasoning_item) = responses_reasoning_item_chunk(&value) { @@ -1589,7 +1605,7 @@ fn response_sse_data_to_chunk(data: &str) -> Option> { /// Log Responses API usage for prompt-cache visibility. /// Looks for `input_tokens_details.cached_tokens` (OpenAI/xAI shape). -fn log_openai_responses_usage(usage: &serde_json::Value) { +fn openai_responses_usage(usage: &serde_json::Value) -> Option { let input = usage .get("input_tokens") .or_else(|| usage.get("prompt_tokens")) @@ -1606,7 +1622,7 @@ fn log_openai_responses_usage(usage: &serde_json::Value) { .unwrap_or(0); if input.is_none() && output.is_none() && cached == 0 { - return; + return None; } let input_v = input.unwrap_or(0); @@ -1623,6 +1639,12 @@ fn log_openai_responses_usage(usage: &serde_json::Value) { cached, hit_pct )); + Some(crate::chunk::LanguageModelUsage { + input_tokens: input.unwrap_or(0), + output_tokens: output.unwrap_or(0), + cache_read_tokens: cached, + cache_write_tokens: 0, + }) } /// Attribute a `response.completed` payload: status, incomplete reason, and @@ -2493,6 +2515,59 @@ mod tests { )); } + #[test] + fn response_incomplete_safety_reasons_emit_refusal() { + for reason in ["refusal", "content_filter", "safety"] { + let chunk = response_sse_data_to_chunk(&format!( + r#"{{"type":"response.incomplete","response":{{"incomplete_details":{{"reason":"{reason}"}}}}}}"# + )) + .expect("expected incomplete chunk"); + assert!(matches!( + chunk, + Ok(ChunkType::End { + reason: Some(crate::chunk::FinishReason::Refusal) + }) + )); + } + } + + #[test] + fn response_completed_retains_provider_usage() { + let chunk = response_sse_data_to_chunk( + r#"{"type":"response.completed","response":{"usage":{"input_tokens":200,"output_tokens":50,"input_tokens_details":{"cached_tokens":150}}}}"#, + ) + .expect("expected completion chunk") + .expect("completion should parse"); + + let ChunkType::ResponseCompleted { usage, .. } = chunk else { + panic!("expected response completed"); + }; + assert_eq!( + usage, + Some(crate::chunk::LanguageModelUsage { + input_tokens: 200, + output_tokens: 50, + cache_read_tokens: 150, + cache_write_tokens: 0, + }) + ); + } + + #[test] + fn response_incomplete_max_output_tokens_emits_terminal_reason() { + let chunk = response_sse_data_to_chunk( + r#"{"type":"response.incomplete","response":{"incomplete_details":{"reason":"max_output_tokens"}}}"#, + ) + .expect("expected incomplete chunk"); + + assert!(matches!( + chunk, + Ok(ChunkType::End { + reason: Some(crate::chunk::FinishReason::Length) + }) + )); + } + #[test] fn retryable_failure_is_terminal_for_sse_eof_tracking() { let chunk = Ok(ChunkType::RetryableFailure( diff --git a/src/aisdk/response.rs b/src/aisdk/response.rs index cef392f7..0f46376c 100644 --- a/src/aisdk/response.rs +++ b/src/aisdk/response.rs @@ -13,6 +13,21 @@ use tokio::sync::mpsc; const PHASELESS_AMBIGUOUS_FOLLOW_UP_LIMIT: usize = 1; const PROVIDER_STEP_MAX_RETRIES: usize = 10; + +fn terminal_stop_reason(reason: Option<&FinishReason>) -> StopReason { + match reason { + Some(FinishReason::Length) => StopReason::MaxTokens, + Some(FinishReason::Refusal | FinishReason::ContentFilter) => StopReason::Refusal, + _ => StopReason::Finish, + } +} + +fn provider_reason_ends_turn(reason: Option<&FinishReason>) -> bool { + matches!( + reason, + Some(FinishReason::Length | FinishReason::Refusal | FinishReason::ContentFilter) + ) +} /// Grok Build only acts on `tail_repetition:{n}@thinking` from /// `response.doom_loop_check` — never on tool names across steps. /// `.devrefs/references/xai-org/grok-build/crates/codegen/xai-grok-sampler/src/doom_loop.rs` @@ -133,6 +148,7 @@ pub async fn stream_with_tools( let mut cached_repeatable_tool_results: HashMap = HashMap::new(); let mut phase_less_ambiguous_follow_ups = 0usize; let mut doom_loop = DoomLoopTracker::default(); + let mut total_usage = crate::chunk::LanguageModelUsage::default(); loop { step_idx += 1; @@ -260,7 +276,12 @@ pub async fn stream_with_tools( end_turn, reasoning_items, doom_loop_triggers, + usage, }) => { + if let Some(usage) = usage { + total_usage += usage; + let _ = tx_loop.send(ChunkType::Usage(total_usage)); + } saw_terminal_event = true; response_end_turn = end_turn; for item in reasoning_items { @@ -342,6 +363,10 @@ pub async fn stream_with_tools( } let _ = tx_loop.send(ChunkType::Metadata(msg)); } + Ok(ChunkType::Usage(usage)) => { + total_usage += usage; + let _ = tx_loop.send(ChunkType::Usage(total_usage)); + } Ok(ChunkType::Warning(msg)) => { let _ = tx_loop.send(ChunkType::Warning(msg)); } @@ -768,6 +793,7 @@ pub async fn stream_with_tools( && response_end_turn.is_none() && last_assistant_message_phase.is_none() && phase_less_ambiguous_follow_ups < PHASELESS_AMBIGUOUS_FOLLOW_UP_LIMIT + && !provider_reason_ends_turn(provider_finish_reason.as_ref()) && provider_finish_reason .as_ref() .is_some_and(|reason| !reason.is_final_assistant_stop()); @@ -808,7 +834,8 @@ pub async fn stream_with_tools( ))); continue; } - *stop_reason_arc.lock().await = Some(StopReason::Finish); + *stop_reason_arc.lock().await = + Some(terminal_stop_reason(provider_finish_reason.as_ref())); break; } @@ -2939,6 +2966,7 @@ mod tests { end_turn: None, reasoning_items: Vec::new(), doom_loop_triggers: vec!["tail_repetition:8@thinking".to_string()], + usage: None, }), ], 1 => vec![ @@ -4083,7 +4111,7 @@ mod tests { assert!(!empty_logged); assert_eq!(retries, 0); assert_eq!(provider.requests.load(Ordering::SeqCst), 1); - assert_eq!(response.stop_reason().await, Some(StopReason::Finish)); + assert_eq!(response.stop_reason().await, Some(StopReason::Refusal)); } #[tokio::test] @@ -4478,3 +4506,22 @@ mod tests { assert_eq!(calls[0].arguments["file_path"], "Cargo.toml"); } } +#[test] +fn terminal_provider_reasons_map_to_typed_stop_reasons() { + assert_eq!( + terminal_stop_reason(Some(&FinishReason::Length)), + StopReason::MaxTokens + ); + assert_eq!( + terminal_stop_reason(Some(&FinishReason::Refusal)), + StopReason::Refusal + ); + assert_eq!( + terminal_stop_reason(Some(&FinishReason::ContentFilter)), + StopReason::Refusal + ); + assert_eq!( + terminal_stop_reason(Some(&FinishReason::Stop)), + StopReason::Finish + ); +} diff --git a/src/aisdk/stop.rs b/src/aisdk/stop.rs index bbcf071d..c84d3d9f 100644 --- a/src/aisdk/stop.rs +++ b/src/aisdk/stop.rs @@ -3,6 +3,8 @@ use std::sync::Arc; #[derive(Debug, Clone, PartialEq)] pub enum StopReason { Finish, + MaxTokens, + Refusal, Hook, Error(String), Other(String), diff --git a/src/app.rs b/src/app.rs index b2132163..238a1d6f 100644 --- a/src/app.rs +++ b/src/app.rs @@ -7500,9 +7500,23 @@ impl App { .map(|session| fork_title_from_session_title(&session.title)) .unwrap_or_else(|| fork_title_from_session_title("fork")); - let _ = self.create_new_session(Some(fork_title)); - for msg in &messages_to_fork { - let _ = self.session_manager.add_message_to_current_session(msg); + let fork_id = self.create_new_session(Some(fork_title)); + let messages_to_fork = + match crate::persistence::attachments::clone_messages(&messages_to_fork, &fork_id) { + Ok(messages) => messages, + Err(error) => { + self.session_manager.delete_session(&fork_id); + self.push_command_error(format!("Failed to copy fork attachments: {error}")); + return false; + } + }; + if let Err(error) = self + .session_manager + .replace_session_messages(&fork_id, messages_to_fork.clone()) + { + self.session_manager.delete_session(&fork_id); + self.push_command_error(format!("Failed to persist fork: {error:?}")); + return false; } self.chat_state.chat.clear(); @@ -9740,7 +9754,21 @@ impl App { self.cancelled_streaming_session(session_id); false } - crate::llm::ChunkMessage::Metrics { .. } => true, + crate::llm::ChunkMessage::Metrics { + duration_ms, + usage, + cost, + .. + } => { + if let Some(usage) = usage { + if let Some(chat) = self.chat_for_session_mut(session_id) { + chat.apply_streaming_usage(usage, cost, duration_ms); + } + self.mark_streaming_snapshot_pending(session_id); + } + true + } + crate::llm::ChunkMessage::TurnStopReason(_) => true, crate::llm::ChunkMessage::ToolCalls(tool_calls) => { self.set_session_retry_status(session_id, None); // Close the generation sample as a tool-calls finish (excluded from @@ -9819,6 +9847,7 @@ impl App { crate::llm::ChunkMessage::QuestionRequest { questions, response_tx, + .. } => { self.maybe_persist_streaming_snapshot_for_session(session_id, true); let _ = self.session_manager.set_session_status( @@ -12467,6 +12496,7 @@ mod tests { let mut app = test_app(); let (permission_tx, _permission_rx) = tokio::sync::oneshot::channel(); app.permission_dialog_state.enqueue(PermissionPrompt { + tool_call_id: None, tool_id: "list".to_string(), action: PermissionAction::List, permission: "external_directory".to_string(), @@ -12474,7 +12504,9 @@ mod tests { target: Some("/tmp".to_string()), command: None, workdir: None, + workspace: "/tmp".to_string(), reason: "approval required".to_string(), + raw_input: serde_json::Value::Null, response_tx: permission_tx, }); let (question_tx, _question_rx) = tokio::sync::oneshot::channel(); @@ -12498,6 +12530,7 @@ mod tests { let mut app = test_app(); let (permission_tx, _permission_rx) = tokio::sync::oneshot::channel(); app.permission_dialog_state.enqueue(PermissionPrompt { + tool_call_id: None, tool_id: "list".to_string(), action: PermissionAction::List, permission: "external_directory".to_string(), @@ -12505,7 +12538,9 @@ mod tests { target: Some("/tmp".to_string()), command: None, workdir: None, + workspace: "/tmp".to_string(), reason: "approval required".to_string(), + raw_input: serde_json::Value::Null, response_tx: permission_tx, }); let (question_tx, _question_rx) = tokio::sync::oneshot::channel(); @@ -13123,6 +13158,7 @@ mod tests { app.chat_state.chat.scroll_offset = 0; let (permission_tx, _permission_rx) = tokio::sync::oneshot::channel(); app.permission_dialog_state.enqueue(PermissionPrompt { + tool_call_id: None, tool_id: "list".to_string(), action: PermissionAction::List, permission: "external_directory".to_string(), @@ -13130,7 +13166,9 @@ mod tests { target: Some("/tmp".to_string()), command: None, workdir: None, + workspace: "/tmp".to_string(), reason: "approval required".to_string(), + raw_input: serde_json::Value::Null, response_tx: permission_tx, }); app.overlay_focus = OverlayFocus::PermissionDialog; diff --git a/src/llm/client.rs b/src/llm/client.rs index ff30a074..3bd8722c 100644 --- a/src/llm/client.rs +++ b/src/llm/client.rs @@ -5,10 +5,10 @@ use crate::aisdk::core::{ stop::StopReason, Message as AisdkMessage, Tool, }; -use crate::aisdk::message::ImageContent; +use crate::aisdk::message::{AudioContent, ImageContent}; use crate::aisdk::{Anthropic, OpenAI, OpenAICompatible}; use futures::StreamExt; -use std::{collections::HashMap, time::Instant}; +use std::{collections::HashMap, path::Path, time::Instant}; use tokio_util::sync::CancellationToken; use crate::tools::aisdk_bridge::convert_to_aisdk_tools; @@ -43,11 +43,68 @@ struct ProviderRequestConfig { api_key: Option, reasoning_effort: Option, supports_image_input: bool, + supports_audio_input: bool, + pricing: Option, openai_options: OpenAIRequestOptions, /// Vercel AI Gateway: enable `providerOptions.gateway.caching = "auto"`. gateway_caching_auto: bool, } +fn messages_have_user_audio(messages: &[crate::session::types::Message]) -> bool { + messages.iter().any(|message| { + message.role == crate::session::types::MessageRole::User + && !message.local_audio_paths.is_empty() + }) +} + +fn provider_kind_for_model( + provider_name: &str, + npm_package: &str, + has_audio_input: bool, +) -> ProviderKind { + let kind = ProviderKind::from_provider(provider_name, npm_package); + if has_audio_input && kind == ProviderKind::OpenAI { + ProviderKind::OpenAICompatible + } else { + kind + } +} + +fn model_supports_audio_input(model: Option<&crate::model::discovery::Model>) -> bool { + model + .and_then(|model| model.modalities.as_ref()) + .is_some_and(|modalities| modalities.input.iter().any(|item| item == "audio")) +} + +fn usage_cost( + usage: crate::aisdk::chunk::LanguageModelUsage, + pricing: Option<&crate::model::discovery::Cost>, +) -> Option { + let pricing = pricing?; + let cached = usage.cache_read_tokens.min(usage.input_tokens); + let written = usage + .cache_write_tokens + .min(usage.input_tokens.saturating_sub(cached)); + let uncached = usage + .input_tokens + .saturating_sub(cached) + .saturating_sub(written); + let input_cost = uncached as f64 * pricing.input; + let cache_read_cost = cached as f64 * pricing.cache_read.unwrap_or(pricing.input); + let cache_write_cost = written as f64 * pricing.cache_write.unwrap_or(pricing.input); + let output_cost = usage.output_tokens as f64 * pricing.output; + Some((input_cost + cache_read_cost + cache_write_cost + output_cost) / 1_000_000.0) +} + +fn turn_stop_reason(stop_reason: Option<&StopReason>) -> Option { + match stop_reason { + Some(StopReason::MaxTokens) => Some(crate::llm::TurnStopReason::MaxTokens), + Some(StopReason::Hook) => Some(crate::llm::TurnStopReason::MaxTurnRequests), + Some(StopReason::Refusal) => Some(crate::llm::TurnStopReason::Refusal), + _ => None, + } +} + impl ProviderRequestConfig { fn new( kind: ProviderKind, @@ -66,6 +123,8 @@ impl ProviderRequestConfig { api_key, reasoning_effort, supports_image_input, + supports_audio_input: false, + pricing: None, openai_options: OpenAIRequestOptions::default(), gateway_caching_auto: false, } @@ -569,6 +628,7 @@ fn truncate_log_value(value: &str, max_chars: usize) -> String { struct StreamRelayResult { outcome: StreamRelayOutcome, stats: RelayStats, + usage: Option, } pub async fn stream_llm_with_cancellation( @@ -605,8 +665,14 @@ pub async fn stream_llm_with_cancellation( messages.len() ); let ui_model = model.clone(); - let request_config = - prepare_request_config(&provider_name, model, reasoning_effort, &sender).await?; + let request_config = prepare_request_config( + &provider_name, + model, + reasoning_effort, + messages_have_user_audio(&messages), + &sender, + ) + .await?; let mut request_config = request_config; let model_mismatch_warning = ui_vs_request_model_mismatch_warning(&ui_model, &request_config.model_name); @@ -672,9 +738,10 @@ pub async fn stream_llm_with_cancellation( ); } - let aisdk_messages = convert_messages_for_model( + let aisdk_messages = convert_messages_for_model_with_audio( &messages, request_config.supports_image_input, + request_config.supports_audio_input, show_vlm_agent_hint, ); // Stamp Build affinity *after* message conversion so turn_idx matches wire content. @@ -737,6 +804,7 @@ pub async fn stream_llm_with_cancellation( let start_time = Instant::now(); let mut token_count: usize = 0; + let pricing = request_config.pricing.clone(); let relay_result = match relay_stream_to_sender( &mut response.stream, @@ -746,6 +814,8 @@ pub async fn stream_llm_with_cancellation( &start_time, primary_log_context, model_mismatch_warning, + pricing.as_ref(), + None, ) .await .map_err(|err| err.to_string()) @@ -767,6 +837,9 @@ pub async fn stream_llm_with_cancellation( }; let stop_reason = response.stop_reason().await; + if let Some(reason) = turn_stop_reason(stop_reason.as_ref()) { + let _ = sender.send(crate::llm::ChunkMessage::TurnStopReason(reason)); + } let stream_outcome = relay_result.outcome; let primary_outcome_label = stream_outcome_label(stream_outcome, stop_reason.as_ref()); crate::emit_log!( @@ -835,6 +908,8 @@ pub async fn stream_llm_with_cancellation( &start_time, summary_log_context, None, + pricing.as_ref(), + relay_result.usage, ) .await .map_err(|err| err.to_string()) @@ -888,7 +963,7 @@ pub async fn build_subagent_llm_session( sender: &crate::llm::ChunkSender, ) -> Result { let request_config = - prepare_request_config(provider_name, model, reasoning_effort, sender).await?; + prepare_request_config(provider_name, model, reasoning_effort, false, sender).await?; Ok(crate::agent::config::LlmSessionConfig { provider_name: request_config.provider_name, model: request_config.model_name, @@ -926,8 +1001,14 @@ pub async fn summarize_for_compaction( } let (warning_sender, _warning_receiver) = tokio::sync::mpsc::unbounded_channel(); - let request_config = - prepare_request_config(&provider_name, model, reasoning_effort, &warning_sender).await?; + let request_config = prepare_request_config( + &provider_name, + model, + reasoning_effort, + false, + &warning_sender, + ) + .await?; let messages = vec![AisdkMessage::user(prompt)]; let mut response = stream_provider_request( &request_config, @@ -970,6 +1051,7 @@ pub async fn summarize_for_compaction( | ChunkType::RetryableFailure(_) | ChunkType::Warning(_) | ChunkType::Metadata(_) + | ChunkType::Usage(_) | ChunkType::Start | ChunkType::Incomplete(_) => {} ChunkType::StreamRollback { text, .. } => { @@ -999,7 +1081,7 @@ pub async fn generate_session_title( ) -> Result { let (warning_sender, _warning_receiver) = tokio::sync::mpsc::unbounded_channel(); let request_config = - prepare_request_config(&provider_name, model, None, &warning_sender).await?; + prepare_request_config(&provider_name, model, None, false, &warning_sender).await?; let prompt = format!( "Generate a concise chat title for this user request.\n\nRules:\n- Return only the title, no quotes or punctuation wrapper.\n- 3 to 7 words.\n- Use title case only when natural.\n- Do not end with a period.\n\nUser request:\n{}", user_message.trim() @@ -1029,6 +1111,7 @@ pub async fn generate_session_title( | ChunkType::RetryableFailure(_) | ChunkType::Warning(_) | ChunkType::Metadata(_) + | ChunkType::Usage(_) | ChunkType::Start | ChunkType::Incomplete(_) => {} ChunkType::StreamRollback { text, .. } => { @@ -1072,6 +1155,7 @@ async fn prepare_request_config( provider_name: &str, model: String, reasoning_effort: Option, + has_audio_input: bool, sender: &crate::llm::ChunkSender, ) -> Result { let auth_dao = crate::persistence::AuthDAO::new()?; @@ -1093,8 +1177,22 @@ async fn prepare_request_config( }; let supports_image_input = model_supports_image_input(&model, provider.models.get(&model)); + let supports_audio_input = model_supports_audio_input(provider.models.get(&model)); + if has_audio_input + && matches!( + auth_config, + Some(crate::persistence::AuthConfig::OAuth { .. }) + ) + && matches!(provider_name, "openai" | "xai") + { + return Err(anyhow::anyhow!( + "Audio input for {provider_name} requires API-key Chat Completions transport; the configured OAuth Responses transport does not support audio input" + ) + .into()); + } let model_route = resolve_model_route(&provider, model); - let provider_kind = ProviderKind::from_provider(provider_name, &model_route.npm_package); + let provider_kind = + provider_kind_for_model(provider_name, &model_route.npm_package, has_audio_input); let base_url = if provider_name == "xai" && model_route.api.trim().is_empty() { // models.dev currently ships empty api for xAI; default to the public endpoint. "https://api.x.ai".to_string() @@ -1116,6 +1214,11 @@ async fn prepare_request_config( reasoning_effort, supports_image_input, ); + request_config.pricing = provider + .models + .get(&model_route.model_name) + .and_then(|model| model.cost.clone()); + request_config.supports_audio_input = supports_audio_input; // Anthropic via AI Gateway needs explicit cache markers; gateway "auto" // inserts them. Without this, Anthropic traffic never cache-reads. if is_vercel_ai_gateway(provider_name, &model_route.npm_package) { @@ -1836,8 +1939,11 @@ async fn relay_stream_to_sender( start_time: &Instant, context: StreamLogContext<'_>, mut mismatch_warning: Option, + pricing: Option<&crate::model::discovery::Cost>, + base_usage: Option, ) -> Result { let mut stats = RelayStats::default(); + let mut stream_usage = None; crate::emit_log!( "[RELAY] relay_stream_to_sender started {}", context.describe() @@ -1940,11 +2046,15 @@ async fn relay_stream_to_sender( let _ = sender.send(crate::llm::ChunkMessage::Metrics { token_count: *token_count, duration_ms, + usage: combined_usage(base_usage, stream_usage), + cost: combined_usage(base_usage, stream_usage) + .and_then(|usage| usage_cost(usage, pricing)), }); let _ = sender.send(crate::llm::ChunkMessage::End); return Ok(StreamRelayResult { outcome: StreamRelayOutcome::Ended, stats, + usage: combined_usage(base_usage, stream_usage), }); } ChunkType::ResponseCompleted { end_turn, .. } => { @@ -1959,11 +2069,15 @@ async fn relay_stream_to_sender( let _ = sender.send(crate::llm::ChunkMessage::Metrics { token_count: *token_count, duration_ms, + usage: combined_usage(base_usage, stream_usage), + cost: combined_usage(base_usage, stream_usage) + .and_then(|usage| usage_cost(usage, pricing)), }); let _ = sender.send(crate::llm::ChunkMessage::End); return Ok(StreamRelayResult { outcome: StreamRelayOutcome::Ended, stats, + usage: combined_usage(base_usage, stream_usage), }); } ChunkType::AssistantMessagePhase { phase } => { @@ -1978,6 +2092,16 @@ async fn relay_stream_to_sender( stats.record_metadata(&message); crate::emit_log!("[RELAY] Metadata {}", message); } + ChunkType::Usage(usage) => { + stream_usage = Some(usage); + crate::emit_log!( + "[RELAY] Usage input={} output={} cache_read={} cache_write={}", + usage.input_tokens, + usage.output_tokens, + usage.cache_read_tokens, + usage.cache_write_tokens, + ); + } ChunkType::Retry(status) => { let elapsed_ms = start_time.elapsed().as_millis(); stats.record_chunk("Retry", elapsed_ms); @@ -2071,9 +2195,24 @@ async fn relay_stream_to_sender( Ok(StreamRelayResult { outcome: StreamRelayOutcome::Exhausted, stats, + usage: combined_usage(base_usage, stream_usage), }) } +fn combined_usage( + base: Option, + current: Option, +) -> Option { + match (base, current) { + (None, None) => None, + (Some(usage), None) | (None, Some(usage)) => Some(usage), + (Some(mut base), Some(current)) => { + base += current; + Some(base) + } + } +} + async fn reached_step_limit(agent_max_steps: Option, response: &StreamTextResponse) -> bool { agent_max_steps.is_some() && matches!(response.stop_reason().await, Some(StopReason::Hook)) } @@ -2082,6 +2221,33 @@ fn estimate_tokens(content: &str) -> usize { content.chars().count().max(1) / 4 } +fn audio_content_for_path(path: &Path) -> Option { + use base64::Engine as _; + let extension = path + .extension() + .and_then(|extension| extension.to_str())? + .to_ascii_lowercase(); + let media_type = match extension.as_str() { + "wav" => "audio/wav", + "mp3" => "audio/mpeg", + _ => { + crate::emit_log!("unsupported audio attachment format: {}", path.display()); + return None; + } + }; + match std::fs::read(path) { + Ok(data) => Some(AudioContent { + data: base64::engine::general_purpose::STANDARD.encode(data), + format: extension, + media_type: media_type.to_string(), + }), + Err(error) => { + crate::emit_log!("failed to attach audio {}: {}", path.display(), error); + None + } + } +} + fn convert_messages(messages: &[crate::session::types::Message]) -> Vec { convert_messages_for_model(messages, true, false) } @@ -2090,6 +2256,20 @@ fn convert_messages_for_model( messages: &[crate::session::types::Message], supports_image_input: bool, show_vlm_agent_hint: bool, +) -> Vec { + convert_messages_for_model_with_audio( + messages, + supports_image_input, + false, + show_vlm_agent_hint, + ) +} + +fn convert_messages_for_model_with_audio( + messages: &[crate::session::types::Message], + supports_image_input: bool, + supports_audio_input: bool, + show_vlm_agent_hint: bool, ) -> Vec { let mut aisdk_messages = Vec::new(); // Soft compaction keeps full UI history; only the active post-boundary @@ -2112,22 +2292,17 @@ fn convert_messages_for_model( aisdk_messages.push(AisdkMessage::system(content)); } crate::session::types::MessageRole::User => { - let content = crate::utils::sanitize::strip_legacy_image_descriptions(&msg.content); + let mut content = + crate::utils::sanitize::strip_legacy_image_descriptions(&msg.content); if !supports_image_input && !msg.local_image_paths.is_empty() { if show_vlm_agent_hint { - aisdk_messages.push(AisdkMessage::user(content_with_vlm_agent_hint( - &content, - &msg.local_image_paths, - ))); + content = content_with_vlm_agent_hint(&content, &msg.local_image_paths); } else { - aisdk_messages.push(AisdkMessage::user( - content_with_unsupported_image_note( - &content, - msg.local_image_paths.len(), - ), - )); + content = content_with_unsupported_image_note( + &content, + msg.local_image_paths.len(), + ); } - continue; } let images = msg @@ -2152,18 +2327,36 @@ fn convert_messages_for_model( }) .collect::>(); - // Empty user rows without images also pad the sticky prefix. - if content.trim().is_empty() && images.is_empty() { + let audios = if supports_audio_input { + msg.local_audio_paths + .iter() + .filter_map(|path| audio_content_for_path(Path::new(path))) + .collect::>() + } else { + if !msg.local_audio_paths.is_empty() { + content.push_str(&format!( + "\n\n[{} audio attachment(s) omitted because the selected model does not support audio input.]", + msg.local_audio_paths.len() + )); + } + Vec::new() + }; + + // Empty user rows without attachments also pad the sticky prefix. + if content.trim().is_empty() && images.is_empty() && audios.is_empty() { continue; } - if images.is_empty() { + if images.is_empty() && audios.is_empty() { aisdk_messages.push(AisdkMessage::user(content)); } else { - aisdk_messages.push(AisdkMessage::user_with_images( - content_with_vision_attached_image_hint(&content), - images, - )); + let content = if images.is_empty() { + content + } else { + content_with_vision_attached_image_hint(&content) + }; + aisdk_messages + .push(AisdkMessage::user_with_attachments(content, images, audios)); } } crate::session::types::MessageRole::Assistant => { @@ -2596,14 +2789,50 @@ fn normalize_anthropic_base_url(base_url: &str) -> String { mod tests { use super::{ apply_provider_request_defaults, convert_messages, convert_messages_for_model, - is_openai_oauth_model_allowed, maybe_apply_unauthenticated_free_provider_key, - model_supports_image_input, openai_oauth_default_originator, - openai_oauth_model_uses_responses_lite, openai_request_instructions, resolve_api_key, - resolve_model_route, ui_vs_request_model_mismatch_warning, vlm_agent_has_model, - AisdkMessage, OpenAIRequestOptions, ProviderKind, ProviderRequestConfig, + convert_messages_for_model_with_audio, is_openai_oauth_model_allowed, + maybe_apply_unauthenticated_free_provider_key, model_supports_image_input, + openai_oauth_default_originator, openai_oauth_model_uses_responses_lite, + openai_request_instructions, provider_kind_for_model, resolve_api_key, resolve_model_route, + ui_vs_request_model_mismatch_warning, vlm_agent_has_model, AisdkMessage, + OpenAIRequestOptions, ProviderKind, ProviderRequestConfig, }; use crate::persistence::AuthConfig; + use base64::Engine as _; + + #[test] + fn audio_model_receives_base64_audio_attachment() { + let dir = tempfile::tempdir().expect("temp dir"); + let path = dir.path().join("sample.wav"); + std::fs::write(&path, b"audio-bytes").unwrap(); + let mut user_message = crate::session::types::Message::user("transcribe"); + user_message.local_audio_paths = vec![path.to_string_lossy().into_owned()]; + + let messages = convert_messages_for_model_with_audio(&[user_message], false, true, false); + let AisdkMessage::User(message) = &messages[0] else { + panic!("expected user message"); + }; + assert_eq!(message.audios.len(), 1); + assert_eq!(message.audios[0].format, "wav"); + assert_eq!( + base64::engine::general_purpose::STANDARD + .decode(&message.audios[0].data) + .unwrap(), + b"audio-bytes" + ); + } + + #[test] + fn openai_audio_models_use_chat_completions_transport() { + assert_eq!( + provider_kind_for_model("openai", "@ai-sdk/openai", true), + ProviderKind::OpenAICompatible + ); + assert_eq!( + provider_kind_for_model("openai", "@ai-sdk/openai", false), + ProviderKind::OpenAI + ); + } #[test] fn stored_auth_takes_precedence_over_custom_provider_api_key() { @@ -3763,3 +3992,38 @@ fn content_with_vlm_agent_hint(content: &str, image_paths: &[String]) -> String format!("{content}\n\n{hint}") } } +#[test] +fn maps_runtime_stop_reasons_to_turn_events() { + assert_eq!( + turn_stop_reason(Some(&StopReason::MaxTokens)), + Some(crate::llm::TurnStopReason::MaxTokens) + ); + assert_eq!( + turn_stop_reason(Some(&StopReason::Refusal)), + Some(crate::llm::TurnStopReason::Refusal) + ); + assert_eq!( + turn_stop_reason(Some(&StopReason::Hook)), + Some(crate::llm::TurnStopReason::MaxTurnRequests) + ); + assert_eq!(turn_stop_reason(Some(&StopReason::Finish)), None); +} + +#[test] +fn computes_cache_aware_usage_cost() { + let usage = crate::aisdk::chunk::LanguageModelUsage { + input_tokens: 1_000_000, + output_tokens: 100_000, + cache_read_tokens: 600_000, + cache_write_tokens: 100_000, + }; + let pricing = crate::model::discovery::Cost { + input: 2.0, + output: 10.0, + cache_read: Some(0.2), + cache_write: Some(2.5), + }; + + let cost = usage_cost(usage, Some(&pricing)).unwrap(); + assert!((cost - 1.97).abs() < f64::EPSILON); +} diff --git a/src/llm/mod.rs b/src/llm/mod.rs index 780d79d3..26f43d1f 100644 --- a/src/llm/mod.rs +++ b/src/llm/mod.rs @@ -35,6 +35,7 @@ pub enum ChunkMessage { }, PermissionRequest(crate::tools::PermissionPrompt), QuestionRequest { + tool_call_id: Option, questions: serde_json::Value, response_tx: tokio::sync::oneshot::Sender, }, @@ -47,15 +48,25 @@ pub enum ChunkMessage { job_id: String, event: BackgroundJobEventKind, }, + TurnStopReason(TurnStopReason), End, Failed(String), Cancelled, Metrics { token_count: usize, duration_ms: u64, + usage: Option, + cost: Option, }, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TurnStopReason { + MaxTokens, + MaxTurnRequests, + Refusal, +} + #[derive(Debug, Clone)] pub enum BackgroundJobEventKind { Started { diff --git a/src/main.rs b/src/main.rs index c48eefb9..00c1f7f1 100644 --- a/src/main.rs +++ b/src/main.rs @@ -558,7 +558,8 @@ async fn run_print_mode( | crate::llm::ChunkMessage::SubagentStarted { .. } | crate::llm::ChunkMessage::SubagentChunk { .. } | crate::llm::ChunkMessage::TerminalSessionEvent { .. } - | crate::llm::ChunkMessage::BackgroundJobEvent { .. } => {} + | crate::llm::ChunkMessage::BackgroundJobEvent { .. } + | crate::llm::ChunkMessage::TurnStopReason(_) => {} crate::llm::ChunkMessage::End => { println!(); play_resolved_sound(&sounds, crate::sound::SoundEvent::Complete); diff --git a/src/mcp/mod.rs b/src/mcp/mod.rs index 6ae6d379..ef226117 100644 --- a/src/mcp/mod.rs +++ b/src/mcp/mod.rs @@ -292,16 +292,41 @@ impl McpManager { if result.is_error == Some(true) { return Err(ToolError::Execution(call_tool_result_text(&result))); } - let output = if let Some(structured) = result.structured_content { - serde_json::to_string_pretty(&structured).unwrap_or_else(|_| structured.to_string()) + Ok(mcp_tool_result(server_name, tool_name, &result)) + } +} + +fn mcp_tool_result( + server_name: &str, + tool_name: &str, + result: &rmcp::model::CallToolResult, +) -> ToolResult { + let text = call_tool_result_text(result); + let output = if let Some(structured) = result.structured_content.as_ref() { + let structured = + serde_json::to_string_pretty(structured).unwrap_or_else(|_| structured.to_string()); + if text.trim().is_empty() || text == structured { + structured } else { - call_tool_result_text(&result) - }; - Ok(ToolResult::new( - format!("MCP: {server_name}.{tool_name}"), - output, - )) + format!("{structured}\n\n{text}") + } + } else { + text + }; + let mut tool_result = ToolResult::new(format!("MCP: {server_name}.{tool_name}"), output) + .with_metadata( + "mcp_result", + serde_json::to_value(&result).unwrap_or(serde_json::Value::Null), + ); + for content in &result.content { + if let ContentBlock::Image(image) = content { + tool_result = tool_result.with_image( + format!("data:{};base64,{}", image.mime_type, image.data), + image.mime_type.clone(), + ); + } } + tool_result } type ConnectOutcome = Result<(RunningService, Vec), McpStatus>; @@ -754,6 +779,36 @@ mod tests { use super::*; use serde_json::json; + #[test] + fn mcp_result_preserves_images_resources_annotations_and_structured_content() { + let resource = rmcp::model::Resource::new("file:///tmp/readme.md", "readme") + .with_mime_type("text/markdown"); + let mut result = rmcp::model::CallToolResult::success(vec![ + ContentBlock::Image( + rmcp::model::ImageContent::new("aGk=", "image/png") + .with_annotations(rmcp::model::Annotations::default().with_priority(0.8)), + ), + ContentBlock::ResourceLink(resource), + ]); + result.structured_content = Some(json!({"answer": 42})); + + let converted = mcp_tool_result("docs", "lookup", &result); + assert_eq!(converted.images.len(), 1); + assert_eq!(converted.images[0].media_type, "image/png"); + assert_eq!( + converted.metadata["mcp_result"]["structuredContent"]["answer"], + 42 + ); + let priority = converted.metadata["mcp_result"]["content"][0]["annotations"]["priority"] + .as_f64() + .unwrap(); + assert!((priority - 0.8).abs() < 0.000_001); + assert_eq!( + converted.metadata["mcp_result"]["content"][1]["uri"], + "file:///tmp/readme.md" + ); + } + #[test] fn normalize_strips_root_anyof_with_non_object_branches() { // Mirrors cua-driver `browser_prepare`: object root + anyOf of required-only diff --git a/src/model/discovery.rs b/src/model/discovery.rs index 193d7969..be2f4fa1 100644 --- a/src/model/discovery.rs +++ b/src/model/discovery.rs @@ -784,6 +784,31 @@ impl Discovery { model.limit.as_ref().map(|l| l.context) } + pub fn model_supports_input_modality( + &self, + provider_id: &str, + model_id: &str, + modality: &str, + ) -> bool { + if self + .custom_providers + .as_ref() + .and_then(|providers| providers.get(&provider_id.trim().to_ascii_lowercase())) + .and_then(|provider| provider.models.get(model_id)) + .and_then(|model| model.modalities.as_ref()) + .is_some_and(|modalities| modalities.input.iter().any(|input| input == modality)) + { + return true; + } + self.load_cache_entry() + .ok() + .flatten() + .and_then(|entry| entry.data.get(provider_id).cloned()) + .and_then(|provider| provider.models.get(model_id).cloned()) + .and_then(|model| model.modalities) + .is_some_and(|modalities| modalities.input.iter().any(|input| input == modality)) + } + pub fn get_model_name(&self, provider_id: &str, model_id: &str) -> Option { if let Some(name) = self .custom_providers @@ -1263,6 +1288,40 @@ mod tests { assert_eq!(model.limit.as_ref().map(|limit| limit.output), Some(8192)); } + #[test] + fn custom_model_modalities_enable_audio_input_lookup() { + let custom_providers = HashMap::from([( + "openai".to_string(), + CustomProviderConfig { + name: None, + npm: Some("@ai-sdk/openai-compatible".to_string()), + base_url: Some("https://api.openai.com/v1".to_string()), + api_key: None, + models: HashMap::from([( + "audio-model".to_string(), + CustomModelConfig { + name: None, + context_window: None, + max_tokens: None, + attachment: None, + reasoning: None, + reasoning_options: None, + temperature: None, + tool_call: None, + modalities: Some(CustomModelModalities { + input: vec!["text".to_string(), "audio".to_string()], + output: vec!["text".to_string()], + }), + launch: false, + }, + )]), + }, + )]); + let discovery = Discovery::new_with_custom(Some(custom_providers)).unwrap(); + + assert!(discovery.model_supports_input_modality("openai", "audio-model", "audio")); + } + #[test] fn custom_model_attachment_flag_updates_modalities() { let mut providers = HashMap::new(); diff --git a/src/persistence/attachments.rs b/src/persistence/attachments.rs new file mode 100644 index 00000000..d433c150 --- /dev/null +++ b/src/persistence/attachments.rs @@ -0,0 +1,322 @@ +use anyhow::{anyhow, Context, Result}; +use std::path::{Component, Path, PathBuf}; + +pub fn root_dir() -> PathBuf { + if cfg!(test) || std::env::var_os("CRABCODE_TEST_MODE").is_some() { + PathBuf::from("/tmp/crabcode_test_data/attachments") + } else { + super::get_data_dir().join("attachments") + } +} + +pub struct AttachmentMigration { + pub messages: Vec, + pub created: Vec, +} + +pub fn migrate_messages( + messages: &[crate::session::types::Message], + session_id: &str, +) -> Result> { + let mut migrated = messages.to_vec(); + let mut created = Vec::new(); + let mut changed = false; + + for message in &mut migrated { + for paths in [ + &mut message.local_image_paths, + &mut message.local_audio_paths, + ] { + let mut unique = Vec::new(); + for attachment_path in std::mem::take(paths) { + if unique.contains(&attachment_path) { + changed = true; + continue; + } + let source = PathBuf::from(&attachment_path); + if is_managed_for_session(&source, session_id) { + unique.push(attachment_path); + continue; + } + let metadata = match std::fs::symlink_metadata(&source) { + Ok(metadata) if metadata.file_type().is_file() => metadata, + _ => { + unique.push(attachment_path); + continue; + } + }; + if metadata.file_type().is_symlink() { + unique.push(attachment_path); + continue; + } + let Some(extension) = source.extension().and_then(|extension| extension.to_str()) + else { + unique.push(attachment_path); + continue; + }; + let data = match std::fs::read(&source) { + Ok(data) => data, + Err(_) => { + unique.push(attachment_path); + continue; + } + }; + match write(session_id, extension, &data) { + Ok(path) => { + unique.push(path.to_string_lossy().into_owned()); + created.push(path); + changed = true; + } + Err(error) => { + for path in created { + remove_file(&path); + } + return Err(error); + } + } + } + *paths = unique; + } + } + + if changed { + Ok(Some(AttachmentMigration { + messages: migrated, + created, + })) + } else { + Ok(None) + } +} + +fn is_managed_for_session(path: &Path, session_id: &str) -> bool { + session_dir(session_id).is_ok_and(|dir| { + path.strip_prefix(dir).is_ok_and(|relative| { + relative.components().count() == 1 + && relative + .components() + .all(|component| matches!(component, Component::Normal(_))) + }) + }) +} +fn validate_session_id(session_id: &str) -> Result<()> { + let path = Path::new(session_id); + if session_id.is_empty() + || path.is_absolute() + || path.components().count() != 1 + || !matches!(path.components().next(), Some(Component::Normal(_))) + { + return Err(anyhow!("invalid attachment session id")); + } + Ok(()) +} + +pub fn session_dir(session_id: &str) -> Result { + validate_session_id(session_id)?; + Ok(root_dir().join(session_id)) +} + +pub fn ensure_session_dir(session_id: &str) -> Result { + let dir = session_dir(session_id)?; + super::create_private_dir_all(&dir)?; + Ok(dir) +} + +pub fn write(session_id: &str, extension: &str, data: &[u8]) -> Result { + if extension.is_empty() || !extension.chars().all(|ch| ch.is_ascii_alphanumeric()) { + return Err(anyhow!("invalid attachment extension")); + } + let dir = ensure_session_dir(session_id)?; + let id = cuid2::create_id(); + let final_path = dir.join(format!("{id}.{extension}")); + let temporary_path = dir.join(format!(".{id}.tmp")); + std::fs::write(&temporary_path, data) + .with_context(|| format!("failed to write attachment {}", temporary_path.display()))?; + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + std::fs::set_permissions(&temporary_path, std::fs::Permissions::from_mode(0o600))?; + } + if let Err(error) = std::fs::rename(&temporary_path, &final_path) { + let _ = std::fs::remove_file(&temporary_path); + return Err(error) + .with_context(|| format!("failed to finalize attachment {}", final_path.display())); + } + Ok(final_path) +} + +pub fn is_managed(path: &Path) -> bool { + path.strip_prefix(root_dir()).is_ok_and(|relative| { + relative + .components() + .all(|component| matches!(component, Component::Normal(_))) + }) +} + +pub fn remove_file(path: &Path) { + if is_managed(path) { + let _ = std::fs::remove_file(path); + } +} + +pub fn cleanup_session(session_id: &str) -> Result<()> { + let dir = session_dir(session_id)?; + match std::fs::remove_dir_all(&dir) { + Ok(()) => Ok(()), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()), + Err(error) => Err(error).with_context(|| format!("failed to remove {}", dir.display())), + } +} + +pub fn clone_messages( + messages: &[crate::session::types::Message], + destination_session_id: &str, +) -> Result> { + let mut cloned = messages.to_vec(); + let mut created = Vec::new(); + for message in &mut cloned { + message.id = cuid2::create_id(); + for paths in [ + &mut message.local_image_paths, + &mut message.local_audio_paths, + ] { + for attachment_path in paths { + let source = PathBuf::from(&*attachment_path); + if !is_managed(&source) { + continue; + } + if std::fs::symlink_metadata(&source)?.file_type().is_symlink() { + return Err(anyhow!("managed attachment cannot be a symlink")); + } + let extension = source + .extension() + .and_then(|extension| extension.to_str()) + .ok_or_else(|| anyhow!("managed attachment has no extension"))?; + let data = std::fs::read(&source) + .with_context(|| format!("failed to read attachment {}", source.display()))?; + match write(destination_session_id, extension, &data) { + Ok(path) => { + *attachment_path = path.to_string_lossy().into_owned(); + created.push(path); + } + Err(error) => { + for path in created { + remove_file(&path); + } + return Err(error); + } + } + } + } + } + Ok(cloned) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn managed_attachment_round_trip_and_cleanup() { + let session = format!("attachment-test-{}", cuid2::create_id()); + let path = write(&session, "png", b"png-data").unwrap(); + assert!(is_managed(&path)); + assert_eq!(std::fs::read(&path).unwrap(), b"png-data"); + cleanup_session(&session).unwrap(); + assert!(!path.exists()); + } + + #[test] + fn fork_clones_managed_files_and_regenerates_message_ids() { + let source_session = format!("attachment-source-{}", cuid2::create_id()); + let destination_session = format!("attachment-dest-{}", cuid2::create_id()); + let source = write(&source_session, "png", b"image").unwrap(); + let mut message = crate::session::types::Message::user("image"); + let original_id = message.id.clone(); + message.local_image_paths = vec![source.to_string_lossy().into_owned()]; + + let cloned = clone_messages(&[message], &destination_session).unwrap(); + assert_ne!(cloned[0].id, original_id); + assert_ne!(cloned[0].local_image_paths[0], source.to_string_lossy()); + assert_eq!( + std::fs::read(&cloned[0].local_image_paths[0]).unwrap(), + b"image" + ); + + cleanup_session(&source_session).unwrap(); + assert!(Path::new(&cloned[0].local_image_paths[0]).exists()); + cleanup_session(&destination_session).unwrap(); + } + + #[test] + fn traversal_path_is_not_managed() { + let traversal = root_dir().join("session").join("..").join("outside.png"); + assert!(!is_managed(&traversal)); + } + + #[test] + fn legacy_image_and_audio_paths_migrate_once() { + let legacy = tempfile::tempdir().unwrap(); + let image = legacy.path().join("image.png"); + let audio = legacy.path().join("audio.wav"); + std::fs::write(&image, b"image").unwrap(); + std::fs::write(&audio, b"audio").unwrap(); + let session = format!("attachment-migrate-{}", cuid2::create_id()); + let mut message = crate::session::types::Message::user("attachments"); + message.local_image_paths = vec![image.to_string_lossy().into_owned()]; + message.local_audio_paths = vec![audio.to_string_lossy().into_owned()]; + + let migration = migrate_messages(&[message], &session) + .unwrap() + .expect("legacy migration"); + assert_eq!(migration.created.len(), 2); + assert!(migration.messages[0] + .local_image_paths + .iter() + .all(|path| is_managed_for_session(Path::new(path), &session))); + assert!(migration.messages[0] + .local_audio_paths + .iter() + .all(|path| is_managed_for_session(Path::new(path), &session))); + assert!(migrate_messages(&migration.messages, &session) + .unwrap() + .is_none()); + + cleanup_session(&session).unwrap(); + } + + #[test] + fn missing_legacy_paths_remain_loadable() { + let session = format!("attachment-missing-{}", cuid2::create_id()); + let mut message = crate::session::types::Message::user("missing"); + message.local_image_paths = vec!["/definitely/missing/image.png".to_string()]; + + assert!(migrate_messages(&[message], &session).unwrap().is_none()); + } + + #[test] + fn fork_clones_mixed_managed_attachments() { + let source_session = format!("attachment-mixed-source-{}", cuid2::create_id()); + let destination_session = format!("attachment-mixed-dest-{}", cuid2::create_id()); + let image = write(&source_session, "png", b"image").unwrap(); + let audio = write(&source_session, "wav", b"audio").unwrap(); + let mut message = crate::session::types::Message::user("mixed"); + message.local_image_paths = vec![image.to_string_lossy().into_owned()]; + message.local_audio_paths = vec![audio.to_string_lossy().into_owned()]; + + let cloned = clone_messages(&[message], &destination_session).unwrap(); + assert_eq!( + std::fs::read(&cloned[0].local_image_paths[0]).unwrap(), + b"image" + ); + assert_eq!( + std::fs::read(&cloned[0].local_audio_paths[0]).unwrap(), + b"audio" + ); + + cleanup_session(&source_session).unwrap(); + assert!(Path::new(&cloned[0].local_image_paths[0]).exists()); + assert!(Path::new(&cloned[0].local_audio_paths[0]).exists()); + cleanup_session(&destination_session).unwrap(); + } +} diff --git a/src/persistence/conversions.rs b/src/persistence/conversions.rs index 54dbe02b..7c39edc1 100644 --- a/src/persistence/conversions.rs +++ b/src/persistence/conversions.rs @@ -22,6 +22,7 @@ impl From for Message { } else { msg.parts .into_iter() + .filter(|part| part.part_type != "local_image" && part.part_type != "local_audio") .map(|part| PersistenceMessagePart { part_type: part.part_type, data: part.data, @@ -44,6 +45,12 @@ impl From for Message { data: serde_json::json!({ "path": path }), }); } + for path in &msg.local_audio_paths { + parts.push(PersistenceMessagePart { + part_type: "local_audio".to_string(), + data: serde_json::json!({ "path": path }), + }); + } if let Some(stats) = msg.compaction_stats { if let Ok(data) = serde_json::to_value(stats) { @@ -71,7 +78,7 @@ impl From for Message { } Message { - id: cuid2::create_id(), + id: msg.id, session_id: 0, role: match msg.role { MessageRole::User => "user".to_string(), @@ -94,6 +101,11 @@ impl From for Message { t1_ms: msg.t1_ms.map(|v| v as i64), tn_ms: msg.tn_ms.map(|v| v as i64), output_tokens: msg.output_tokens.map(|v| v as i64), + input_tokens: msg.input_tokens.map(|v| v as i64), + cache_read_tokens: msg.cache_read_tokens.map(|v| v as i64), + cache_write_tokens: msg.cache_write_tokens.map(|v| v as i64), + cost: msg.cost, + usage_authoritative: msg.usage_authoritative, } } } @@ -110,6 +122,20 @@ impl TryFrom for SessionMessage { data: part.data.clone(), }) .collect(); + let local_audio_paths = session_parts + .iter() + .filter_map(|part| { + (part.part_type == "local_audio") + .then(|| part.data.get("path").and_then(|value| value.as_str())) + .flatten() + }) + .map(str::to_string) + .fold(Vec::new(), |mut paths, path| { + if !paths.contains(&path) { + paths.push(path); + } + paths + }); let content = session_parts .iter() @@ -146,7 +172,12 @@ impl TryFrom for SessionMessage { } }) .map(|path| path.to_string()) - .collect(); + .fold(Vec::new(), |mut paths, path| { + if !paths.contains(&path) { + paths.push(path); + } + paths + }); let compaction_stats = session_parts .iter() @@ -170,6 +201,7 @@ impl TryFrom for SessionMessage { }; Ok(SessionMessage { + id: msg.id, role, content, reasoning, @@ -200,10 +232,30 @@ impl TryFrom for SessionMessage { output_tokens: msg .output_tokens .and_then(|v| if v > 0 { Some(v as usize) } else { None }), + input_tokens: msg + .input_tokens + .and_then(|v| if v > 0 { Some(v as usize) } else { None }), + cache_read_tokens: msg.cache_read_tokens.and_then(|v| { + if v > 0 { + Some(v as usize) + } else { + None + } + }), + cache_write_tokens: msg.cache_write_tokens.and_then(|v| { + if v > 0 { + Some(v as usize) + } else { + None + } + }), + cost: msg.cost, + usage_authoritative: msg.usage_authoritative, tokens_per_sec: None, model: msg.model.clone(), provider: msg.provider.clone(), local_image_paths, + local_audio_paths, compaction_stats, was_interrupted, }) @@ -231,6 +283,95 @@ pub fn persistence_to_session( mod tests { use super::*; + #[test] + fn message_id_round_trips_through_persistence() { + let session_message = SessionMessage::assistant("hello"); + let id = session_message.id.clone(); + + let persistence_message: Message = session_message.into(); + assert_eq!(persistence_message.id, id); + + let restored = SessionMessage::try_from(persistence_message).unwrap(); + assert_eq!(restored.id, id); + } + + #[test] + fn authoritative_usage_round_trips_through_persistence() { + let mut session_message = SessionMessage::assistant("hello"); + session_message.apply_usage( + crate::aisdk::chunk::LanguageModelUsage { + input_tokens: 100, + output_tokens: 25, + cache_read_tokens: 60, + cache_write_tokens: 10, + }, + Some(0.0125), + ); + + let restored = SessionMessage::try_from(Message::from(session_message)).unwrap(); + assert_eq!(restored.input_tokens, Some(100)); + assert_eq!(restored.output_tokens, Some(25)); + assert_eq!(restored.cache_read_tokens, Some(60)); + assert_eq!(restored.cache_write_tokens, Some(10)); + assert_eq!(restored.cost, Some(0.0125)); + assert!(restored.usage_authoritative); + } + + #[test] + fn audio_paths_round_trip_through_persistence() { + let mut session_message = SessionMessage::user("listen"); + session_message.local_audio_paths = vec!["/tmp/audio.wav".to_string()]; + + let persisted = Message::from(session_message); + assert_eq!( + persisted + .parts + .iter() + .filter(|part| part.part_type == "local_audio") + .count(), + 1 + ); + let restored = SessionMessage::try_from(persisted).unwrap(); + assert_eq!(restored.local_audio_paths, vec!["/tmp/audio.wav"]); + } + + #[test] + fn duplicate_legacy_attachment_parts_are_deduplicated() { + let message = Message { + id: "message".to_string(), + session_id: 1, + role: "user".to_string(), + parts: vec![ + PersistenceMessagePart { + part_type: "local_audio".to_string(), + data: serde_json::json!({ "path": "/tmp/audio.wav" }), + }, + PersistenceMessagePart { + part_type: "local_audio".to_string(), + data: serde_json::json!({ "path": "/tmp/audio.wav" }), + }, + ], + timestamp: 0, + tokens_used: 0, + model: None, + provider: None, + agent_mode: None, + duration_ms: 0, + t0_ms: None, + t1_ms: None, + tn_ms: None, + output_tokens: None, + input_tokens: None, + cache_read_tokens: None, + cache_write_tokens: None, + cost: None, + usage_authoritative: false, + }; + + let restored = SessionMessage::try_from(message).unwrap(); + assert_eq!(restored.local_audio_paths, vec!["/tmp/audio.wav"]); + } + #[test] fn compaction_stats_round_trip_through_message_parts() { let stats = CompactionStats { diff --git a/src/persistence/history.rs b/src/persistence/history.rs index 4ac270ae..a631b599 100644 --- a/src/persistence/history.rs +++ b/src/persistence/history.rs @@ -15,6 +15,72 @@ pub struct Workspace { pub last_opened_at: i64, } +#[cfg(test)] +mod tests { + use super::*; + + fn test_dao() -> HistoryDAO { + let mut conn = Connection::open_in_memory().unwrap(); + run_migrations(&mut conn).unwrap(); + let workspace_id = ensure_workspace(&conn, "/tmp/workspace", "workspace").unwrap(); + HistoryDAO { + conn, + current_workspace_id: workspace_id, + current_workspace_path: "/tmp/workspace".to_string(), + current_workspace_name: "workspace".to_string(), + } + } + + #[test] + fn authoritative_usage_updates_message_and_session_totals() { + let dao = test_dao(); + let session_id = dao + .create_session("session", "Session".to_string()) + .unwrap(); + let message = Message { + id: "message".to_string(), + session_id, + role: "assistant".to_string(), + parts: Vec::new(), + timestamp: chrono::Utc::now().timestamp(), + tokens_used: 5, + model: Some("model".to_string()), + provider: Some("provider".to_string()), + agent_mode: None, + duration_ms: 10, + t0_ms: None, + t1_ms: None, + tn_ms: None, + output_tokens: Some(25), + input_tokens: Some(100), + cache_read_tokens: Some(60), + cache_write_tokens: Some(10), + cost: Some(0.0125), + usage_authoritative: true, + }; + + dao.add_message(&message).unwrap(); + let restored = dao.get_messages(session_id).unwrap(); + assert_eq!(restored[0].input_tokens, Some(100)); + assert_eq!(restored[0].cost, Some(0.0125)); + let session = dao.get_session(session_id).unwrap().unwrap(); + assert_eq!(session.total_tokens, 125); + assert!((session.total_cost - 0.0125).abs() < f64::EPSILON); + } +} + +fn message_total_tokens(message: &Message) -> i32 { + if message.usage_authoritative { + let total = message + .input_tokens + .unwrap_or(0) + .saturating_add(message.output_tokens.unwrap_or(0)); + i32::try_from(total).unwrap_or(i32::MAX) + } else { + message.tokens_used + } +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Session { pub id: i64, @@ -62,6 +128,11 @@ pub struct Message { pub t1_ms: Option, pub tn_ms: Option, pub output_tokens: Option, + pub input_tokens: Option, + pub cache_read_tokens: Option, + pub cache_write_tokens: Option, + pub cost: Option, + pub usage_authoritative: bool, } pub struct HistoryDAO { @@ -438,9 +509,10 @@ impl HistoryDAO { self.conn.execute( "INSERT INTO messages ( id, session_id, role, parts, timestamp, tokens_used, model, provider, agent_mode, duration_ms, - t0_ms, t1_ms, tn_ms, output_tokens + t0_ms, t1_ms, tn_ms, output_tokens, input_tokens, cache_read_tokens, + cache_write_tokens, cost, usage_authoritative ) - VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14)", + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18, ?19)", params![ &msg.id, msg.session_id, @@ -456,10 +528,21 @@ impl HistoryDAO { msg.t1_ms, msg.tn_ms, msg.output_tokens, + msg.input_tokens, + msg.cache_read_tokens, + msg.cache_write_tokens, + msg.cost, + msg.usage_authoritative, ], )?; - self.update_session_stats(msg.session_id, msg.tokens_used, 0.0, msg.timestamp)?; + let tokens = message_total_tokens(msg); + self.update_session_stats( + msg.session_id, + tokens, + msg.cost.unwrap_or(0.0), + msg.timestamp, + )?; Ok(()) } @@ -476,20 +559,23 @@ impl HistoryDAO { )?; let mut total_tokens: i64 = 0; + let mut total_cost = 0.0; let mut updated_at = chrono::Utc::now().timestamp(); { let mut insert = tx.prepare_cached( "INSERT INTO messages ( id, session_id, role, parts, timestamp, tokens_used, model, provider, agent_mode, duration_ms, - t0_ms, t1_ms, tn_ms, output_tokens + t0_ms, t1_ms, tn_ms, output_tokens, input_tokens, cache_read_tokens, + cache_write_tokens, cost, usage_authoritative ) - VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14)", + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18, ?19)", )?; for msg in messages { let parts_json = serde_json::to_string(&msg.parts)?; - total_tokens += msg.tokens_used as i64; + total_tokens += i64::from(message_total_tokens(msg)); + total_cost += msg.cost.unwrap_or(0.0); updated_at = msg.timestamp; insert.execute(params![ @@ -507,6 +593,11 @@ impl HistoryDAO { msg.t1_ms, msg.tn_ms, msg.output_tokens, + msg.input_tokens, + msg.cache_read_tokens, + msg.cache_write_tokens, + msg.cost, + msg.usage_authoritative, ])?; } } @@ -525,13 +616,14 @@ impl HistoryDAO { tx.execute( "UPDATE sessions SET total_tokens = ?1, - total_cost = 0, - total_time_sec = ?2, - avg_tokens_per_sec = ?3, - updated_at = ?4 - WHERE id = ?5", + total_cost = ?2, + total_time_sec = ?3, + avg_tokens_per_sec = ?4, + updated_at = ?5 + WHERE id = ?6", params![ total_tokens, + total_cost, total_time_sec, avg_tokens_per_sec, updated_at, @@ -547,7 +639,8 @@ impl HistoryDAO { pub fn get_messages(&self, session_id: i64) -> Result> { let mut stmt = self.conn.prepare( "SELECT id, session_id, role, parts, timestamp, tokens_used, model, provider, agent_mode, duration_ms, - t0_ms, t1_ms, tn_ms, output_tokens + t0_ms, t1_ms, tn_ms, output_tokens, input_tokens, cache_read_tokens, + cache_write_tokens, cost, usage_authoritative FROM messages WHERE session_id = ?1 ORDER BY timestamp ASC, rowid ASC", )?; @@ -570,6 +663,11 @@ impl HistoryDAO { t1_ms: row.get(11)?, tn_ms: row.get(12)?, output_tokens: row.get(13)?, + input_tokens: row.get(14)?, + cache_read_tokens: row.get(15)?, + cache_write_tokens: row.get(16)?, + cost: row.get(17)?, + usage_authoritative: row.get(18)?, }) })?; diff --git a/src/persistence/migrations.rs b/src/persistence/migrations.rs index 1fedbbf3..0aa7550c 100644 --- a/src/persistence/migrations.rs +++ b/src/persistence/migrations.rs @@ -16,9 +16,41 @@ pub fn run_migrations(db: &mut Connection) -> Result<()> { migrate_to_v3(db)?; } + if current_version < 4 { + migrate_to_v4(db)?; + } + Ok(()) } +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn migration_v4_adds_authoritative_usage_columns() { + let mut db = Connection::open_in_memory().unwrap(); + run_migrations(&mut db).unwrap(); + + let columns = db + .prepare("PRAGMA table_info(messages)") + .unwrap() + .query_map([], |row| row.get::<_, String>(1)) + .unwrap() + .collect::, _>>() + .unwrap(); + for column in [ + "input_tokens", + "cache_read_tokens", + "cache_write_tokens", + "cost", + "usage_authoritative", + ] { + assert!(columns.iter().any(|candidate| candidate == column)); + } + } +} + fn get_current_version(db: &Connection) -> Result { match db.prepare("SELECT MAX(version) FROM migrations") { Ok(mut stmt) => { @@ -108,6 +140,30 @@ fn migrate_to_v1(db: &mut Connection) -> Result<()> { Ok(()) } +fn migrate_to_v4(db: &mut Connection) -> Result<()> { + let tx = db.transaction()?; + let _ = tx.execute("ALTER TABLE messages ADD COLUMN input_tokens INTEGER", []); + let _ = tx.execute( + "ALTER TABLE messages ADD COLUMN cache_read_tokens INTEGER", + [], + ); + let _ = tx.execute( + "ALTER TABLE messages ADD COLUMN cache_write_tokens INTEGER", + [], + ); + let _ = tx.execute("ALTER TABLE messages ADD COLUMN cost REAL", []); + let _ = tx.execute( + "ALTER TABLE messages ADD COLUMN usage_authoritative INTEGER NOT NULL DEFAULT 0", + [], + ); + tx.execute( + "INSERT OR IGNORE INTO migrations (version, applied_at) VALUES (4, strftime('%s', 'now'))", + params![], + )?; + tx.commit()?; + Ok(()) +} + fn migrate_to_v2(db: &mut Connection) -> Result<()> { let tx = db.transaction()?; diff --git a/src/persistence/mod.rs b/src/persistence/mod.rs index 47129682..fc18bc58 100644 --- a/src/persistence/mod.rs +++ b/src/persistence/mod.rs @@ -2,6 +2,7 @@ use anyhow::Result; use std::ffi::OsString; use std::path::{Path, PathBuf}; +pub mod attachments; pub mod auth; pub mod conversions; pub mod db; @@ -55,7 +56,7 @@ fn resolve_state_home(xdg_state_home: Option, home_dir: Option Result<()> { +pub(crate) fn create_private_dir_all(dir: &Path) -> Result<()> { std::fs::create_dir_all(dir)?; restrict_dir_permissions(dir)?; Ok(()) diff --git a/src/remote/mod.rs b/src/remote/mod.rs index 02586dad..5d4a3b82 100644 --- a/src/remote/mod.rs +++ b/src/remote/mod.rs @@ -304,6 +304,7 @@ struct RemoteMessage { model: Option, provider: Option, local_image_paths: Vec, + local_audio_paths: Vec, was_interrupted: bool, parts: Vec, } @@ -3305,6 +3306,7 @@ fn remote_message(message: &Message) -> RemoteMessage { model: message.model.clone(), provider: message.provider.clone(), local_image_paths: message.local_image_paths.clone(), + local_audio_paths: message.local_audio_paths.clone(), was_interrupted: message.was_interrupted, parts: message.parts.clone(), } @@ -4121,10 +4123,12 @@ mod tests { fn remote_message_includes_local_image_paths() { let mut message = Message::user("see [Image #1]"); message.local_image_paths = vec!["/tmp/example.png".to_string()]; + message.local_audio_paths = vec!["/tmp/example.wav".to_string()]; let remote = remote_message(&message); assert_eq!(remote.local_image_paths, vec!["/tmp/example.png"]); + assert_eq!(remote.local_audio_paths, vec!["/tmp/example.wav"]); } #[test] diff --git a/src/session/compaction.rs b/src/session/compaction.rs index 2acd8474..b6968037 100644 --- a/src/session/compaction.rs +++ b/src/session/compaction.rs @@ -172,7 +172,6 @@ fn select_messages_with_budget( kept_tokens = kept_tokens.saturating_add(size); } } - // OpenCode: if no keep, or keep would start at work index 0, summarize all. let tail_start_work = match tail_start_work { Some(0) | None => work_indices.len(), @@ -556,6 +555,18 @@ fn message_content_for_prompt(message: &Message) -> String { } } + if !message.local_audio_paths.is_empty() { + if !content.trim().is_empty() { + content.push('\n'); + } + content.push_str("Attached local audio:\n"); + for path in &message.local_audio_paths { + content.push_str("- "); + content.push_str(path); + content.push('\n'); + } + } + content } diff --git a/src/session/manager.rs b/src/session/manager.rs index 47eb694b..cc106918 100644 --- a/src/session/manager.rs +++ b/src/session/manager.rs @@ -247,6 +247,28 @@ impl SessionManager { let _ = dao.replace_messages(db_id, &persistence_messages); } + if let Some(migration) = + crate::persistence::attachments::migrate_messages(&hydrated.messages, id)? + { + let persistence_messages: Vec = migration + .messages + .iter() + .cloned() + .map(|message| { + let mut db_message: crate::persistence::Message = message.into(); + db_message.session_id = db_id; + db_message + }) + .collect(); + if let Err(error) = dao.replace_messages(db_id, &persistence_messages) { + for path in migration.created { + crate::persistence::attachments::remove_file(&path); + } + return Err(SessionError::PersistenceError(error.to_string())); + } + hydrated.messages = migration.messages; + } + *existing = hydrated; self.message_counts .insert(id.to_string(), existing.messages.len()); @@ -730,11 +752,7 @@ impl SessionManager { session_id: &str, message: &crate::session::types::Message, ) -> Result<(), SessionError> { - if let Some(session) = self.sessions.get_mut(session_id) { - session.add_message(message.clone()); - self.message_counts - .insert(session_id.to_string(), session.messages.len()); - } else { + if !self.sessions.contains_key(session_id) { return Err(SessionError::NotFound(session_id.to_string())); } @@ -742,11 +760,18 @@ impl SessionManager { if let Some(db_id) = self.id_mapping.get(session_id) { let mut db_message: crate::persistence::Message = message.clone().into(); db_message.session_id = *db_id; - let _ = dao - .add_message(&db_message) - .map_err(|e| SessionError::PersistenceError(e.to_string())); + dao.add_message(&db_message) + .map_err(|e| SessionError::PersistenceError(e.to_string()))?; } } + + let session = self + .sessions + .get_mut(session_id) + .ok_or_else(|| SessionError::NotFound(session_id.to_string()))?; + session.add_message(message.clone()); + self.message_counts + .insert(session_id.to_string(), session.messages.len()); Ok(()) } @@ -755,19 +780,15 @@ impl SessionManager { session_id: &str, messages: Vec, ) -> Result<(), SessionError> { - if let Some(session) = self.sessions.get_mut(session_id) { - session.messages = messages.clone(); - session.updated_at = SystemTime::now(); - self.message_counts - .insert(session_id.to_string(), session.messages.len()); - } else { + if !self.sessions.contains_key(session_id) { return Err(SessionError::NotFound(session_id.to_string())); } if let Some(ref dao) = self.history_dao { if let Some(db_id) = self.id_mapping.get(session_id) { let persistence_messages: Vec = messages - .into_iter() + .iter() + .cloned() .map(|message| { let mut db_message: crate::persistence::Message = message.into(); db_message.session_id = *db_id; @@ -780,6 +801,15 @@ impl SessionManager { } } + let session = self + .sessions + .get_mut(session_id) + .ok_or_else(|| SessionError::NotFound(session_id.to_string()))?; + session.messages = messages; + session.updated_at = SystemTime::now(); + self.message_counts + .insert(session_id.to_string(), session.messages.len()); + Ok(()) } @@ -945,9 +975,17 @@ impl SessionManager { } pub fn delete_session(&mut self, id: &str) -> bool { + self.try_delete_session(id).unwrap_or_else(|error| { + crate::emit_log!("Failed to delete session {}: {:?}", id, error); + false + }) + } + + pub fn try_delete_session(&mut self, id: &str) -> Result { if let Some(db_id) = self.id_mapping.get(id) { if let Some(ref dao) = self.history_dao { - let _ = dao.delete_session(*db_id); + dao.delete_session(*db_id) + .map_err(|error| SessionError::PersistenceError(error.to_string()))?; } } @@ -969,9 +1007,12 @@ impl SessionManager { if self.current_session_id.as_ref() == Some(&id.to_string()) { self.current_session_id = None; } - true + if let Err(error) = crate::persistence::attachments::cleanup_session(id) { + crate::emit_log!("Failed to clean session attachments for {}: {}", id, error); + } + Ok(true) } else { - false + Ok(false) } } } diff --git a/src/session/types.rs b/src/session/types.rs index 2ace65a6..c2db52ba 100644 --- a/src/session/types.rs +++ b/src/session/types.rs @@ -158,6 +158,7 @@ impl CompactionStats { #[derive(Debug, Clone, PartialEq)] pub struct Message { + pub id: String, pub role: MessageRole, pub content: String, pub reasoning: Option, @@ -174,17 +175,38 @@ pub struct Message { pub t1_ms: Option, pub tn_ms: Option, pub output_tokens: Option, + pub input_tokens: Option, + pub cache_read_tokens: Option, + pub cache_write_tokens: Option, + pub cost: Option, + pub usage_authoritative: bool, /// Precomputed tokens/s (OpenCode inter-token aggregate). Prefer over /// recomputing `output_tokens / duration_ms`. pub tokens_per_sec: Option, pub model: Option, pub provider: Option, pub local_image_paths: Vec, + pub local_audio_paths: Vec, pub compaction_stats: Option, pub was_interrupted: bool, } impl Message { + pub fn apply_usage( + &mut self, + usage: crate::aisdk::chunk::LanguageModelUsage, + cost: Option, + ) { + self.input_tokens = Some(usize::try_from(usage.input_tokens).unwrap_or(usize::MAX)); + self.output_tokens = Some(usize::try_from(usage.output_tokens).unwrap_or(usize::MAX)); + self.cache_read_tokens = + Some(usize::try_from(usage.cache_read_tokens).unwrap_or(usize::MAX)); + self.cache_write_tokens = + Some(usize::try_from(usage.cache_write_tokens).unwrap_or(usize::MAX)); + self.cost = cost; + self.usage_authoritative = true; + } + pub fn new(role: MessageRole, content: impl Into) -> Self { let content = content.into(); let parts = if content.is_empty() { @@ -194,6 +216,7 @@ impl Message { }; Self { + id: cuid2::create_id(), role, content, reasoning: None, @@ -208,10 +231,16 @@ impl Message { t1_ms: None, tn_ms: None, output_tokens: None, + input_tokens: None, + cache_read_tokens: None, + cache_write_tokens: None, + cost: None, + usage_authoritative: false, tokens_per_sec: None, model: None, provider: None, local_image_paths: Vec::new(), + local_audio_paths: Vec::new(), compaction_stats: None, was_interrupted: false, } @@ -242,6 +271,7 @@ impl Message { }; Self { + id: cuid2::create_id(), role: MessageRole::Assistant, content, reasoning: None, @@ -256,10 +286,16 @@ impl Message { t1_ms: None, tn_ms: None, output_tokens: None, + input_tokens: None, + cache_read_tokens: None, + cache_write_tokens: None, + cost: None, + usage_authoritative: false, tokens_per_sec: None, model: None, provider: None, local_image_paths: Vec::new(), + local_audio_paths: Vec::new(), compaction_stats: None, was_interrupted: false, } diff --git a/src/tools/aisdk_bridge.rs b/src/tools/aisdk_bridge.rs index 39ce65b5..22db5238 100644 --- a/src/tools/aisdk_bridge.rs +++ b/src/tools/aisdk_bridge.rs @@ -146,7 +146,13 @@ pub async fn convert_to_aisdk_tools( } if let Err(e) = permissions - .preflight(&agent_mode, &tool_id_for_exec, &input, sender.as_ref()) + .preflight_for_call( + &agent_mode, + &tool_id_for_exec, + &input, + Some(&call_id), + sender.as_ref(), + ) .await { let err = format!("{}", e); @@ -228,24 +234,7 @@ pub async fn convert_to_aisdk_tools( }; if let Some(ref sender) = sender { - let preview = truncate_tool_output(&tool_result.output, TOOL_UI_PREVIEW_LIMIT); - - let line_count = tool_result.output.lines().count(); - let meta = serde_json::Value::Object( - tool_result - .metadata - .into_iter() - .collect::>(), - ); - - let payload = serde_json::json!({ - "status": "ok", - "title": tool_result.title, - "output_preview": preview, - "line_count": line_count, - "metadata": meta, - }) - .to_string(); + let payload = tool_success_payload(&tool_result); if sender .send(crate::llm::ChunkMessage::ToolResult( @@ -334,6 +323,28 @@ fn truncate_tool_output(output: &str, limit: usize) -> String { truncated } +fn tool_success_payload(tool_result: &crate::tools::ToolResult) -> String { + let preview = truncate_tool_output(&tool_result.output, TOOL_UI_PREVIEW_LIMIT); + let meta = serde_json::Value::Object( + tool_result + .metadata + .iter() + .map(|(key, value)| (key.clone(), value.clone())) + .collect(), + ); + + serde_json::json!({ + "status": "ok", + "title": tool_result.title, + "output": tool_result.output, + "output_preview": preview, + "line_count": tool_result.output.lines().count(), + "metadata": meta, + "images": tool_result.images, + }) + .to_string() +} + fn unsupported_image_input_note(image_count: usize) -> String { let image_label = if image_count == 1 { "image" } else { "images" }; format!( @@ -355,6 +366,7 @@ fn send_tool_error_result( let payload = serde_json::json!({ "status": "error", "title": "Tool failed", + "output": error, "output_preview": preview, "line_count": error.lines().count().max(1), "metadata": { @@ -401,7 +413,8 @@ fn param_to_json_schema(param_type: &crate::tools::ParameterType) -> serde_json: #[cfg(test)] mod tests { - use super::{send_tool_error_result, truncate_tool_output}; + use super::{send_tool_error_result, tool_success_payload, truncate_tool_output}; + use std::collections::HashMap; #[test] fn truncate_tool_output_bounds_large_results() { @@ -420,6 +433,43 @@ mod tests { assert_eq!(truncate_tool_output(output, 40_000), output); } + #[test] + fn tool_success_payload_retains_full_output_and_bounded_preview() { + let output = "a".repeat(5_000); + let result = crate::tools::ToolResult { + title: "Large result".to_string(), + output: output.clone(), + metadata: HashMap::new(), + images: Vec::new(), + }; + + let payload: serde_json::Value = + serde_json::from_str(&tool_success_payload(&result)).expect("payload should be json"); + + assert_eq!(payload["output"], output); + assert!(payload["output_preview"] + .as_str() + .is_some_and(|preview| preview.len() < 5_000)); + assert!(payload["output_preview"] + .as_str() + .is_some_and(|preview| preview.contains("tool output truncated to 4000 bytes"))); + } + + #[test] + fn tool_success_payload_retains_result_images() { + let result = crate::tools::ToolResult::new("Image", "viewed") + .with_image("data:image/png;base64,aGk=", "image/png"); + + let payload: serde_json::Value = + serde_json::from_str(&tool_success_payload(&result)).expect("payload should be json"); + + assert_eq!( + payload["images"][0]["data_url"], + "data:image/png;base64,aGk=" + ); + assert_eq!(payload["images"][0]["media_type"], "image/png"); + } + #[test] fn send_tool_error_result_emits_error_payload() { let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); @@ -443,6 +493,10 @@ mod tests { serde_json::from_str(&result.content).expect("payload should be json"); assert_eq!(payload["status"], "error"); assert_eq!(payload["title"], "Tool failed"); + assert_eq!( + payload["output"], + "Execution error: Could not find text to replace" + ); assert_eq!( payload["output_preview"], "Execution error: Could not find text to replace" diff --git a/src/tools/edit.rs b/src/tools/edit.rs index e37a732c..70eb658d 100644 --- a/src/tools/edit.rs +++ b/src/tools/edit.rs @@ -155,7 +155,10 @@ impl ToolHandler for EditTool { format!("Replaced at line {}", line_num.unwrap_or(1)) }, ) - .with_metadata("replace_count", serde_json::json!(count)); + .with_metadata("replace_count", serde_json::json!(count)) + .with_metadata("path", serde_json::json!(file_path)) + .with_metadata("old_text", serde_json::json!(content)) + .with_metadata("new_text", serde_json::json!(new_content)); if let Some(line_num) = line_num { result = result.with_metadata("line_number", serde_json::json!(line_num)); diff --git a/src/tools/fs/write.rs b/src/tools/fs/write.rs index 2ab70d39..2dc7caf1 100644 --- a/src/tools/fs/write.rs +++ b/src/tools/fs/write.rs @@ -63,6 +63,7 @@ impl ToolHandler for WriteTool { let content = get_string_param(¶ms, "content") .ok_or_else(|| ToolError::Validation("content is required".to_string()))?; + let old_text = std::fs::read_to_string(&file_path).ok(); let (is_new, bytes) = write_one_file(&file_path, &content)?; Ok(ToolResult::new( @@ -72,7 +73,10 @@ impl ToolHandler for WriteTool { } else { format!("Updated file with {} bytes", bytes) }, - )) + ) + .with_metadata("path", serde_json::json!(file_path)) + .with_metadata("old_text", serde_json::json!(old_text)) + .with_metadata("new_text", serde_json::json!(content))) } } @@ -137,6 +141,7 @@ impl ToolHandler for WriteFilesTool { .ok_or_else(|| ToolError::Validation("files must be an array".to_string()))?; let mut summaries = Vec::with_capacity(files.len()); + let mut changes = Vec::with_capacity(files.len()); for file in files { let file_path = file .get("file_path") @@ -146,16 +151,23 @@ impl ToolHandler for WriteFilesTool { .get("content") .and_then(Value::as_str) .ok_or_else(|| ToolError::Validation("content is required".to_string()))?; + let old_text = std::fs::read_to_string(file_path).ok(); let (is_new, bytes) = write_one_file(file_path, content)?; let action = if is_new { "created" } else { "updated" }; summaries.push(format!("{file_path}: {action} {bytes} bytes")); + changes.push(serde_json::json!({ + "path": file_path, + "old_text": old_text, + "new_text": content, + })); } Ok(ToolResult::new( format!("Write files: {}", summaries.len()), summaries.join("\n"), ) - .with_metadata("file_count", serde_json::json!(summaries.len()))) + .with_metadata("file_count", serde_json::json!(summaries.len())) + .with_metadata("changes", serde_json::json!(changes))) } } diff --git a/src/tools/patch.rs b/src/tools/patch.rs index eb431702..337694f8 100644 --- a/src/tools/patch.rs +++ b/src/tools/patch.rs @@ -5,7 +5,7 @@ use crate::tools::{ }; use async_trait::async_trait; use serde_json::Value; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::path::{Path, PathBuf}; pub struct ApplyPatchTool; @@ -75,21 +75,40 @@ impl ToolHandler for ApplyPatchTool { Ok(()) } - async fn execute(&self, params: Value, _ctx: &ToolContext) -> Result { + async fn execute(&self, params: Value, ctx: &ToolContext) -> Result { let patch = get_string_param(¶ms, "patch") .ok_or_else(|| ToolError::Validation("patch is required".to_string()))?; let patch = clean_patch_input(&patch); + let paths = patch_paths_as_pathbufs(¶ms, ctx.workdir()); + let before = paths + .iter() + .map(|path| (path.clone(), std::fs::read_to_string(path).ok())) + .collect::>(); let summary = if patch.trim_start().starts_with("*** Begin Patch") { apply_codex_patch(&patch)? } else { apply_unified_patch(&patch)? }; + let changes = before + .into_iter() + .filter_map(|(path, old_text)| { + let new_text = std::fs::read_to_string(&path).ok(); + (old_text != new_text).then(|| { + serde_json::json!({ + "path": path, + "old_text": old_text, + "new_text": new_text.unwrap_or_default(), + }) + }) + }) + .collect::>(); Ok(ToolResult::new( "Apply patch", format!("Applied patch: {}", summary.describe()), ) - .with_metadata("file_count", serde_json::json!(summary.touched()))) + .with_metadata("file_count", serde_json::json!(summary.touched())) + .with_metadata("changes", serde_json::json!(changes))) } } @@ -153,6 +172,259 @@ pub(crate) fn patch_paths_as_pathbufs(params: &Value, workdir: &Path) -> Vec, + pub new_text: String, +} + +#[derive(Default)] +struct PatchPreviewState { + original: HashMap>, + current: HashMap>, + order: Vec, +} + +impl PatchPreviewState { + fn resolve(workdir: &Path, path: &str) -> PathBuf { + let path = PathBuf::from(path); + if path.is_absolute() { + path + } else { + workdir.join(path) + } + } + + fn load(&mut self, path: &Path) -> Result, ToolError> { + if let Some(content) = self.current.get(path) { + return Ok(content.clone()); + } + let content = match std::fs::read(path) { + Ok(bytes) => Some(decode_utf8(&bytes)?), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => None, + Err(error) => { + return Err(ToolError::Execution(format!( + "Failed to read {}: {error}", + path.display() + ))); + } + }; + self.original.insert(path.to_path_buf(), content.clone()); + self.current.insert(path.to_path_buf(), content.clone()); + self.order.push(path.to_path_buf()); + Ok(content) + } + + fn required(&mut self, path: &Path) -> Result { + self.load(path)?.ok_or_else(|| { + ToolError::NotFound(format!("Patch source file not found: {}", path.display())) + }) + } + + fn create(&mut self, path: PathBuf, content: String) -> Result<(), ToolError> { + if self.load(&path)?.is_some() { + return Err(ToolError::Execution(format!( + "Refusing to overwrite existing file: {}", + path.display() + ))); + } + self.current.insert(path, Some(content)); + Ok(()) + } + + fn write(&mut self, path: PathBuf, content: String) -> Result<(), ToolError> { + self.required(&path)?; + self.current.insert(path, Some(content)); + Ok(()) + } + + fn delete(&mut self, path: PathBuf) -> Result<(), ToolError> { + self.required(&path)?; + self.current.insert(path, None); + Ok(()) + } + + fn finish(self) -> Vec { + self.order + .into_iter() + .filter_map(|path| { + let old_text = self.original.get(&path).cloned().flatten(); + let new_text = self.current.get(&path).cloned().flatten(); + (old_text != new_text).then_some(PatchPreview { + path, + old_text, + new_text: new_text.unwrap_or_default(), + }) + }) + .collect() + } +} + +pub(crate) fn preview_patch( + params: &Value, + workdir: &Path, +) -> Result, ToolError> { + let patch = get_string_param(params, "patch") + .ok_or_else(|| ToolError::Validation("patch is required".to_string()))?; + let patch = clean_patch_input(&patch); + let mut state = PatchPreviewState::default(); + if patch.trim_start().starts_with("*** Begin Patch") { + preview_codex_patch(&patch, workdir, &mut state)?; + } else { + preview_unified_patch(&patch, workdir, &mut state)?; + } + let changes = state.finish(); + if changes.is_empty() { + return Err(ToolError::Validation( + "Patch did not contain any file changes".to_string(), + )); + } + Ok(changes) +} + +fn preview_unified_patch( + patch: &str, + workdir: &Path, + state: &mut PatchPreviewState, +) -> Result<(), ToolError> { + let lines: Vec<&str> = patch.lines().collect(); + let mut index = 0; + while index < lines.len() { + if !lines[index].starts_with("--- ") { + index += 1; + continue; + } + let old_path = normalize_diff_path(lines[index].trim_start_matches("--- ")); + index += 1; + if index >= lines.len() || !lines[index].starts_with("+++ ") { + return Err(ToolError::Validation( + "Unified diff file header must include a +++ path".to_string(), + )); + } + let new_path = normalize_diff_path(lines[index].trim_start_matches("+++ ")); + index += 1; + let source = + (old_path != "/dev/null").then(|| PatchPreviewState::resolve(workdir, &old_path)); + let target = + (new_path != "/dev/null").then(|| PatchPreviewState::resolve(workdir, &new_path)); + let mut content = match &source { + Some(path) => state.required(path)?, + None => String::new(), + }; + while index < lines.len() + && !lines[index].starts_with("--- ") + && !lines[index].starts_with("diff --git ") + { + if !lines[index].starts_with("@@") { + index += 1; + continue; + } + index += 1; + let (old_text, new_text, next_index) = collect_hunk(&lines, index); + content = replace_hunk(&content, &old_text, &new_text)?; + index = next_index; + } + match (source, target) { + (None, Some(target)) => state.create(target, content)?, + (Some(source), None) => state.delete(source)?, + (Some(source), Some(target)) if source == target => state.write(source, content)?, + (Some(source), Some(target)) => { + state.create(target, content)?; + state.delete(source)?; + } + (None, None) => { + return Err(ToolError::Validation( + "Patch cannot use /dev/null for both paths".to_string(), + )); + } + } + } + Ok(()) +} + +fn preview_codex_patch( + patch: &str, + workdir: &Path, + state: &mut PatchPreviewState, +) -> Result<(), ToolError> { + let lines: Vec<&str> = patch.lines().collect(); + let mut index = 0; + if lines.get(index).map(|line| line.trim()) != Some("*** Begin Patch") { + return Err(ToolError::Validation( + "Codex patch must start with *** Begin Patch".to_string(), + )); + } + index += 1; + while index < lines.len() { + let line = lines[index].trim(); + if line == "*** End Patch" { + break; + } + if let Some(path) = line.strip_prefix("*** Add File: ") { + index += 1; + let mut file_lines = Vec::new(); + while index < lines.len() && !lines[index].starts_with("*** ") { + let Some(content) = lines[index].strip_prefix('+') else { + return Err(ToolError::Validation( + "Add File lines must start with +".to_string(), + )); + }; + file_lines.push(content.to_string()); + index += 1; + } + state.create( + PatchPreviewState::resolve(workdir, path), + join_hunk_lines(&file_lines), + )?; + continue; + } + if let Some(path) = line.strip_prefix("*** Delete File: ") { + state.delete(PatchPreviewState::resolve(workdir, path))?; + index += 1; + continue; + } + if let Some(path) = line.strip_prefix("*** Update File: ") { + let source = PatchPreviewState::resolve(workdir, path); + let mut content = state.required(&source)?; + index += 1; + let move_to = lines + .get(index) + .and_then(|line| line.trim().strip_prefix("*** Move to: ")) + .map(str::to_string); + if move_to.is_some() { + index += 1; + } + while index < lines.len() && !lines[index].starts_with("*** ") { + if !lines[index].starts_with("@@") { + index += 1; + continue; + } + index += 1; + let (old_text, new_text, next_index) = collect_hunk(&lines, index); + content = replace_hunk(&content, &old_text, &new_text)?; + index = next_index; + } + if let Some(target) = move_to { + let target = PatchPreviewState::resolve(workdir, &target); + if target == source { + state.write(source, content)?; + } else { + state.create(target, content)?; + state.delete(source)?; + } + } else { + state.write(source, content)?; + } + continue; + } + return Err(ToolError::Validation(format!( + "Unsupported patch directive: {line}" + ))); + } + Ok(()) +} + fn clean_patch_input(raw: &str) -> String { let trimmed = raw.trim(); let mut lines: Vec<&str> = trimmed.lines().collect(); @@ -718,6 +990,54 @@ mod tests { assert_eq!(std::fs::read_to_string(second).unwrap(), "alpha\ngamma\n"); assert!(result.output.contains("updated 2")); assert_eq!(result.metadata["file_count"], serde_json::json!(2)); + let changes = result.metadata["changes"] + .as_array() + .expect("patch changes"); + assert_eq!(changes.len(), 2); + assert_eq!(changes[0]["old_text"], "one\ntwo\n"); + assert_eq!(changes[0]["new_text"], "one\nthree\n"); + } + + #[test] + fn preview_patch_builds_full_file_changes_without_mutating_files() { + let dir = tempfile::tempdir().unwrap(); + let source = dir.path().join("source.txt"); + let deleted = dir.path().join("deleted.txt"); + std::fs::write(&source, "one\ntwo\n").unwrap(); + std::fs::write(&deleted, "remove me\n").unwrap(); + let patch = format!( + "*** Begin Patch\n*** Update File: {}\n*** Move to: moved.txt\n@@\n one\n-two\n+three\n*** Delete File: {}\n*** Add File: added.txt\n+new file\n*** End Patch\n", + source.display(), + deleted.display() + ); + + let changes = preview_patch(&serde_json::json!({ "patch": patch }), dir.path()).unwrap(); + + assert_eq!(changes.len(), 4); + assert!(changes.iter().any(|change| { + change.path == source + && change.old_text.as_deref() == Some("one\ntwo\n") + && change.new_text.is_empty() + })); + assert!(changes.iter().any(|change| { + change.path == dir.path().join("moved.txt") + && change.old_text.is_none() + && change.new_text == "one\nthree\n" + })); + assert!(changes.iter().any(|change| { + change.path == deleted + && change.old_text.as_deref() == Some("remove me\n") + && change.new_text.is_empty() + })); + assert!(changes.iter().any(|change| { + change.path == dir.path().join("added.txt") + && change.old_text.is_none() + && change.new_text == "new file\n" + })); + assert_eq!(std::fs::read_to_string(&source).unwrap(), "one\ntwo\n"); + assert_eq!(std::fs::read_to_string(&deleted).unwrap(), "remove me\n"); + assert!(!dir.path().join("moved.txt").exists()); + assert!(!dir.path().join("added.txt").exists()); } #[tokio::test] diff --git a/src/tools/permission.rs b/src/tools/permission.rs index f5706135..2414d407 100644 --- a/src/tools/permission.rs +++ b/src/tools/permission.rs @@ -98,6 +98,7 @@ pub type PermissionRules = Vec; #[derive(Debug)] pub struct PermissionPrompt { + pub tool_call_id: Option, pub tool_id: String, pub action: PermissionAction, pub permission: String, @@ -105,7 +106,9 @@ pub struct PermissionPrompt { pub target: Option, pub command: Option, pub workdir: Option, + pub workspace: String, pub reason: String, + pub raw_input: Value, pub response_tx: tokio::sync::oneshot::Sender, } @@ -323,6 +326,18 @@ impl ToolPermissions { tool_id: &str, params: &Value, sender: Option<&ChunkSender>, + ) -> Result<(), ToolError> { + self.preflight_for_call(agent_mode, tool_id, params, None, sender) + .await + } + + pub async fn preflight_for_call( + &self, + agent_mode: &str, + tool_id: &str, + params: &Value, + tool_call_id: Option<&str>, + sender: Option<&ChunkSender>, ) -> Result<(), ToolError> { if !self.is_tool_allowed_for_agent(agent_mode, tool_id) { return Err(ToolError::Permission(format!( @@ -363,6 +378,8 @@ impl ToolPermissions { PermissionReasonKind::ConfiguredAsk, path.as_deref(), command.clone(), + params, + tool_call_id, sender, ) .await; @@ -407,6 +424,8 @@ impl ToolPermissions { reason_kind, reason_path.as_deref().or(path.as_deref()), command.clone(), + params, + tool_call_id, sender, ) .await; @@ -442,6 +461,8 @@ impl ToolPermissions { reason_kind, path.as_deref(), command, + params, + tool_call_id, sender, ) .await; @@ -459,6 +480,8 @@ impl ToolPermissions { reason_kind: PermissionReasonKind, path: Option<&Path>, command: Option, + params: &Value, + tool_call_id: Option<&str>, sender: Option<&ChunkSender>, ) -> Result<(), ToolError> { let target = path @@ -515,6 +538,7 @@ impl ToolPermissions { let (response_tx, response_rx) = tokio::sync::oneshot::channel(); let prompt = PermissionPrompt { + tool_call_id: tool_call_id.map(str::to_string), tool_id: tool_id.to_string(), action, permission: grant.permission.clone(), @@ -522,7 +546,9 @@ impl ToolPermissions { target: prompt_target, command, workdir, + workspace: self.workdir.to_string_lossy().into_owned(), reason: reason_text, + raw_input: params.clone(), response_tx, }; @@ -1291,6 +1317,39 @@ mod tests { ); } + #[tokio::test] + async fn permission_prompt_retains_originating_tool_call_id() { + let perms = ToolPermissions::new("/tmp/workspace"); + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); + let params = serde_json::json!({ "file_path": "/tmp/elsewhere/file.txt" }); + + let pending = tokio::spawn({ + let perms = perms.clone(); + let tx = tx.clone(); + async move { + perms + .preflight_for_call("build", "read", ¶ms, Some("call_123"), Some(&tx)) + .await + } + }); + + let prompt = match rx.recv().await { + Some(ChunkMessage::PermissionRequest(prompt)) => prompt, + _ => panic!("Expected permission prompt"), + }; + assert_eq!(prompt.tool_call_id.as_deref(), Some("call_123")); + assert_eq!( + prompt.raw_input, + serde_json::json!({ "file_path": "/tmp/elsewhere/file.txt" }) + ); + assert_eq!(prompt.workspace, "/tmp/workspace"); + let _ = prompt.response_tx.send(PermissionResponse::Deny); + assert!(pending + .await + .expect("preflight task should complete") + .is_err()); + } + #[tokio::test] async fn allow_always_persists_for_same_request_fingerprint() { let perms = ToolPermissions::new("/tmp/workspace"); diff --git a/src/tools/question.rs b/src/tools/question.rs index e10e814f..89923256 100644 --- a/src/tools/question.rs +++ b/src/tools/question.rs @@ -74,16 +74,75 @@ fn parse_questions_param(params: &Value) -> Result { } }; - match parsed { - Value::Array(_) => Ok(normalize_questions(parsed)), - Value::Object(_) => Ok(normalize_questions(Value::Array(vec![parsed]))), + let normalized = match parsed { + Value::Array(items) if items.is_empty() => { + return Err(ToolError::Validation( + "questions array cannot be empty".to_string(), + )); + } + Value::Array(_) => normalize_questions(parsed), + Value::Object(_) => normalize_questions(Value::Array(vec![parsed])), Value::String(s) if !s.trim().is_empty() => { - Ok(normalize_questions(question_from_plain_text(params, &s))) + normalize_questions(question_from_plain_text(params, &s)) + } + _ => { + return Err(ToolError::Validation( + "questions JSON must decode to an array or object".to_string(), + )); + } + }; + validate_normalized_questions(&normalized)?; + Ok(normalized) +} + +fn validate_normalized_questions(questions: &Value) -> Result<(), ToolError> { + let items = questions + .as_array() + .ok_or_else(|| ToolError::Validation("questions must normalize to an array".to_string()))?; + for (index, item) in items.iter().enumerate() { + let object = item.as_object().ok_or_else(|| { + ToolError::Validation(format!("question {} must be an object", index + 1)) + })?; + let has_prompt = ["question", "header"] + .iter() + .filter_map(|key| object.get(*key).and_then(Value::as_str)) + .any(|value| !value.trim().is_empty()); + if !has_prompt { + return Err(ToolError::Validation(format!( + "question {} must include non-empty question or header text", + index + 1 + ))); + } + let options = object + .get("options") + .and_then(Value::as_array) + .ok_or_else(|| { + ToolError::Validation(format!("question {} options must be an array", index + 1)) + })?; + let mut labels = std::collections::HashSet::new(); + for (option_index, option) in options.iter().enumerate() { + let label = option + .get("label") + .and_then(Value::as_str) + .or_else(|| option.as_str()) + .map(str::trim) + .filter(|label| !label.is_empty()) + .ok_or_else(|| { + ToolError::Validation(format!( + "question {} option {} must include a non-empty label", + index + 1, + option_index + 1 + )) + })?; + if !labels.insert(label.to_string()) { + return Err(ToolError::Validation(format!( + "question {} contains duplicate option label: {label}", + index + 1 + ))); + } } - _ => Err(ToolError::Validation( - "questions JSON must decode to an array or object".to_string(), - )), } + Ok(()) } fn normalize_questions(value: Value) -> Value { @@ -385,6 +444,7 @@ impl ToolHandler for QuestionTool { sender .send(crate::llm::ChunkMessage::QuestionRequest { + tool_call_id: ctx.call_id.clone(), questions: questions.clone(), response_tx, }) @@ -545,6 +605,23 @@ mod tests { assert!(err.contains("questions parameter cannot be empty")); } + #[test] + fn parse_questions_rejects_empty_or_malformed_items() { + for params in [ + json!({ "questions": [] }), + json!({ "questions": [null] }), + json!({ "questions": [{ "question": "", "header": "" }] }), + json!({ + "questions": [{ + "question": "Pick", + "options": [{"label":"A"}, {"label":"A"}] + }] + }), + ] { + assert!(parse_questions_param(¶ms).is_err(), "{params}"); + } + } + #[test] fn model_output_includes_questions_and_answers() { let questions = json!([ @@ -589,4 +666,42 @@ mod tests { .unwrap() .contains("Do not call the question tool again")); } + + #[tokio::test] + async fn question_request_preserves_tool_call_id_and_answers() { + let (sender, mut receiver) = tokio::sync::mpsc::unbounded_channel(); + let tool = QuestionTool::new().with_sender(sender); + let (_abort_tx, abort_rx) = tokio::sync::watch::channel(false); + let ctx = ToolContext::new("session", "message", "Build", abort_rx) + .with_call_id("question_call_1"); + let task = tokio::spawn(async move { + tool.execute( + json!({ + "questions": [{ + "question": "Pick one", + "header": "Choice", + "options": [{"label": "A", "description": "First"}] + }] + }), + &ctx, + ) + .await + }); + + let Some(crate::llm::ChunkMessage::QuestionRequest { + tool_call_id, + questions, + response_tx, + }) = receiver.recv().await + else { + panic!("question request"); + }; + assert_eq!(tool_call_id.as_deref(), Some("question_call_1")); + assert_eq!(questions[0]["question"], "Pick one"); + response_tx.send(json!([["A"]])).expect("question response"); + + let result = task.await.expect("question task").expect("tool result"); + assert!(result.output.contains("\"status\":\"answered\"")); + assert_eq!(result.metadata["answers"], json!([["A"]])); + } } diff --git a/src/tools/terminal_session.rs b/src/tools/terminal_session.rs index 8860e214..447b48cc 100644 --- a/src/tools/terminal_session.rs +++ b/src/tools/terminal_session.rs @@ -36,9 +36,19 @@ pub struct TerminalSessionStart { #[derive(Debug, Clone, Serialize, Deserialize)] pub enum TerminalSessionControl { - Start { rows: u16, cols: u16 }, + Start { + rows: u16, + cols: u16, + }, Input(Vec), - Resize { rows: u16, cols: u16 }, + Resize { + rows: u16, + cols: u16, + }, + /// Complete the session using a terminal hosted by an external client. + ExternalResult(TerminalSessionResult), + /// Fail the session because an external terminal backend could not complete it. + ExternalError(String), Stop, } @@ -126,6 +136,28 @@ impl TranscriptState { } } +pub(crate) fn external_terminal_result( + start: &TerminalSessionStart, + output: &str, + truncated: bool, + exit_code: Option, + stopped_by_user: bool, +) -> TerminalSessionResult { + let mut transcript = TranscriptState::new(start.rows.max(1), start.cols.max(1)); + transcript.append(output.as_bytes()); + transcript.truncated |= truncated; + TerminalSessionResult { + session_id: start.session_id.clone(), + exit_code, + transcript_bytes: transcript.raw.len(), + transcript_truncated: transcript.truncated, + transcript_plain: transcript.plain_text(), + cols: transcript.cols, + rows: transcript.rows, + stopped_by_user, + } +} + /// Convert a PTY byte stream into display-safe text for chat history and model context. /// The live terminal still receives the original bytes through the VT parser. pub(crate) fn sanitize_terminal_output(raw: &[u8]) -> String { @@ -308,6 +340,10 @@ impl TerminalSessionTool { start.rows = rows.max(1); start.cols = cols.max(1); } + Some(TerminalSessionControl::ExternalResult(result)) => return Ok(result), + Some(TerminalSessionControl::ExternalError(error)) => { + return Err(ToolError::Execution(error)); + } Some(TerminalSessionControl::Stop) | None => { emit_event(&sender, &tool_call_id, TerminalSessionEvent::Stopped); return Ok(TerminalSessionResult { @@ -464,6 +500,8 @@ impl TerminalSessionTool { TerminalSessionEvent::Resized { rows, cols }, ); } + Some(TerminalSessionControl::ExternalResult(_)) => {} + Some(TerminalSessionControl::ExternalError(_)) => {} Some(TerminalSessionControl::Stop) | None => { stopped_by_user = true; if let Ok(mut guard) = child.lock() { @@ -730,6 +768,28 @@ mod tests { assert_eq!(state.cols, 120); } + #[test] + fn external_terminal_result_preserves_client_output_and_exit() { + let start = TerminalSessionStart { + session_id: "session".to_string(), + tool_call_id: "call".to_string(), + command: "echo hi".to_string(), + description: "test".to_string(), + workdir: None, + cols: 80, + rows: 24, + job_id: None, + }; + + let result = external_terminal_result(&start, "\x1b[31mhi\x1b[0m\n", true, Some(7), false); + + assert_eq!(result.session_id, "session"); + assert_eq!(result.transcript_plain, "hi\n"); + assert!(result.transcript_truncated); + assert_eq!(result.exit_code, Some(7)); + assert!(!result.stopped_by_user); + } + #[test] fn shell_command_builder_accepts_workdir() { let dir = PathBuf::from("/tmp/work"); diff --git a/src/ui/components/chat.rs b/src/ui/components/chat.rs index 49784036..9944ae27 100644 --- a/src/ui/components/chat.rs +++ b/src/ui/components/chat.rs @@ -1860,6 +1860,22 @@ impl Chat { self.invalidate_cache(); } + pub fn apply_streaming_usage( + &mut self, + usage: crate::aisdk::chunk::LanguageModelUsage, + cost: Option, + duration_ms: u64, + ) { + if let Some(message) = self + .messages + .iter_mut() + .rfind(|message| message.role == MessageRole::Assistant) + { + message.apply_usage(usage, cost); + message.duration_ms = Some(duration_ms); + } + } + pub fn truncate_messages(&mut self, len: usize) { self.messages.truncate(len); self.invalidate_cache(); @@ -2252,6 +2268,8 @@ impl Chat { std::mem::discriminant(&msg.role).hash(&mut h); msg.content.hash(&mut h); msg.reasoning.hash(&mut h); + msg.local_image_paths.hash(&mut h); + msg.local_audio_paths.hash(&mut h); for part in &msg.parts { part.part_type.hash(&mut h); part.data.to_string().hash(&mut h); @@ -2593,7 +2611,9 @@ impl Chat { .rposition(|m| m.role == MessageRole::Assistant) { if let Some(msg) = self.messages.get_mut(idx) { - msg.output_tokens = Some(token_count); + if !msg.usage_authoritative { + msg.output_tokens = Some(token_count); + } msg.token_count = Some(token_count); msg.duration_ms = Some(decode_duration_ms); msg.tokens_per_sec = final_tps; diff --git a/src/views/permission_dialog.rs b/src/views/permission_dialog.rs index 684c21cd..220043d6 100644 --- a/src/views/permission_dialog.rs +++ b/src/views/permission_dialog.rs @@ -556,6 +556,7 @@ mod tests { fn bash_detail_lines_show_command_and_workdir() { let (response_tx, _response_rx) = tokio::sync::oneshot::channel(); let prompt = PermissionPrompt { + tool_call_id: None, tool_id: "bash".to_string(), action: PermissionAction::Bash, permission: "bash".to_string(), @@ -563,7 +564,9 @@ mod tests { target: Some("cargo test".to_string()), command: Some("cargo test".to_string()), workdir: Some("/tmp/workspace".to_string()), + workspace: "/tmp/workspace".to_string(), reason: "Bash command execution requires permission".to_string(), + raw_input: serde_json::Value::Null, response_tx, }; let colors = Theme::load_builtin_default().get_colors(true); @@ -589,6 +592,7 @@ mod tests { fn external_directory_detail_target_shows_wildcard_scope() { let (response_tx, _response_rx) = tokio::sync::oneshot::channel(); let prompt = PermissionPrompt { + tool_call_id: None, tool_id: "read".to_string(), action: PermissionAction::Read, permission: "external_directory".to_string(), @@ -596,7 +600,9 @@ mod tests { target: Some("/Users/carlo/Desktop/Projects/sheetpilot".to_string()), command: None, workdir: None, + workspace: "/tmp".to_string(), reason: "Tool 'read' wants to access path outside working directory".to_string(), + raw_input: serde_json::Value::Null, response_tx, }; let colors = Theme::load_builtin_default().get_colors(true); @@ -614,6 +620,7 @@ mod tests { let (response_tx, _response_rx) = tokio::sync::oneshot::channel(); let mut state = PermissionDialogState::new(); state.enqueue(PermissionPrompt { + tool_call_id: None, tool_id: "bash".to_string(), action: PermissionAction::Bash, permission: "bash".to_string(), @@ -621,7 +628,9 @@ mod tests { target: Some("cargo test".to_string()), command: Some("cargo test".to_string()), workdir: Some("/tmp/workspace".to_string()), + workspace: "/tmp/workspace".to_string(), reason: "Bash command execution requires permission".to_string(), + raw_input: serde_json::Value::Null, response_tx, }); @@ -658,6 +667,7 @@ mod tests { let (response_tx, _response_rx) = tokio::sync::oneshot::channel(); let mut state = PermissionDialogState::new(); state.enqueue(PermissionPrompt { + tool_call_id: None, tool_id: "read".to_string(), action: PermissionAction::Read, permission: "read".to_string(), @@ -665,7 +675,9 @@ mod tests { target: Some("/tmp/file".to_string()), command: None, workdir: None, + workspace: "/tmp".to_string(), reason: "explicit approval required".to_string(), + raw_input: serde_json::Value::Null, response_tx, }); @@ -689,6 +701,7 @@ mod tests { let (response_tx, _response_rx) = tokio::sync::oneshot::channel(); let mut state = PermissionDialogState::new(); state.enqueue(PermissionPrompt { + tool_call_id: None, tool_id: "read".to_string(), action: PermissionAction::Read, permission: "external_directory".to_string(), @@ -696,7 +709,9 @@ mod tests { target: Some("/Users/carlo/Desktop/Projects/sheetpilot/README.md".to_string()), command: None, workdir: None, + workspace: "/tmp".to_string(), reason: "Tool 'read' wants to access path outside working directory: /Users/carlo/Desktop/Projects/sheetpilot/README.md".to_string(), + raw_input: serde_json::Value::Null, response_tx, }); let colors = Theme::load_builtin_default().get_colors(true); @@ -729,6 +744,7 @@ mod tests { let (response_tx, _response_rx) = tokio::sync::oneshot::channel(); let mut state = PermissionDialogState::new(); state.enqueue(PermissionPrompt { + tool_call_id: None, tool_id: "read".to_string(), action: PermissionAction::Read, permission: "read".to_string(), @@ -736,7 +752,9 @@ mod tests { target: Some("/tmp/file".to_string()), command: None, workdir: None, + workspace: "/tmp".to_string(), reason: "explicit approval required".to_string(), + raw_input: serde_json::Value::Null, response_tx, }); let colors = Theme::load_builtin_default().get_colors(true); @@ -770,6 +788,7 @@ mod tests { let (response_tx, _response_rx) = tokio::sync::oneshot::channel(); let mut state = PermissionDialogState::new(); state.enqueue(PermissionPrompt { + tool_call_id: None, tool_id: "read".to_string(), action: PermissionAction::Read, permission: "read".to_string(), @@ -777,7 +796,9 @@ mod tests { target: Some("/tmp/file".to_string()), command: None, workdir: None, + workspace: "/tmp".to_string(), reason: "explicit approval required".to_string(), + raw_input: serde_json::Value::Null, response_tx, }); let colors = Theme::load_builtin_default().get_colors(true); diff --git a/tests/acp_stdio.rs b/tests/acp_stdio.rs new file mode 100644 index 00000000..3e357dad --- /dev/null +++ b/tests/acp_stdio.rs @@ -0,0 +1,87 @@ +use std::io::{BufRead, BufReader, Write}; +use std::process::{Command, Stdio}; +use std::sync::mpsc; +use std::time::{Duration, Instant}; + +#[test] +fn initialize_over_stdio_and_shutdown_on_eof() { + let workspace = tempfile::tempdir().expect("workspace"); + let home = tempfile::tempdir().expect("home"); + let config = tempfile::tempdir().expect("config"); + let state = tempfile::tempdir().expect("state"); + let mut child = Command::new(env!("CARGO_BIN_EXE_crabcode")) + .args(["acp", "--cwd"]) + .arg(workspace.path()) + .env("HOME", home.path()) + .env("XDG_CONFIG_HOME", config.path()) + .env("XDG_STATE_HOME", state.path()) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .expect("spawn crabcode acp"); + + let stdout = child.stdout.take().expect("stdout"); + let (line_tx, line_rx) = mpsc::channel(); + std::thread::spawn(move || { + let mut line = String::new(); + let result = BufReader::new(stdout).read_line(&mut line).map(|_| line); + let _ = line_tx.send(result); + }); + + let request = serde_json::json!({ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": 1, + "clientCapabilities": {} + } + }); + let mut stdin = child.stdin.take().expect("stdin"); + writeln!(stdin, "{request}").expect("write initialize"); + stdin.flush().expect("flush initialize"); + + let line = match line_rx.recv_timeout(Duration::from_secs(10)) { + Ok(Ok(line)) => line, + Ok(Err(error)) => { + let _ = child.kill(); + panic!("failed reading ACP response: {error}"); + } + Err(_) => { + let _ = child.kill(); + panic!("timed out waiting for ACP initialize response"); + } + }; + let response: serde_json::Value = serde_json::from_str(line.trim()).unwrap_or_else(|error| { + let _ = child.kill(); + panic!("invalid protocol response {line:?}: {error}"); + }); + assert_eq!(response["jsonrpc"], "2.0"); + assert_eq!(response["id"], 1); + assert_eq!(response["result"]["protocolVersion"], 1); + assert_eq!(response["result"]["agentInfo"]["name"], "crabcode"); + assert_eq!( + response["result"]["agentInfo"]["version"], + env!("CARGO_PKG_VERSION") + ); + assert_eq!(response["result"]["agentCapabilities"]["loadSession"], true); + assert_eq!( + response["result"]["agentCapabilities"]["promptCapabilities"]["audio"], + true + ); + + drop(stdin); + let deadline = Instant::now() + Duration::from_secs(10); + let status = loop { + if let Some(status) = child.try_wait().expect("poll ACP process") { + break status; + } + if Instant::now() >= deadline { + let _ = child.kill(); + panic!("ACP process did not shut down after stdin EOF"); + } + std::thread::sleep(Duration::from_millis(20)); + }; + assert!(status.success(), "ACP exited with {status}"); +}