//! ACP agent loop with tools and session modes (Stage 3). //! //! Handlers: //! //! | ACP method | Behaviour | //! |-----------------------|-------------------------------------------------------------| //! | `initialize` | echo protocol version, advertise capabilities | //! | `session/new` | mint id, register state, advertise [Default, Bypass] modes | //! | `session/prompt` | tool-call loop: stream → dispatch tools → re-enter, repeat | //! | `session/cancel` | fire the session's cancellation token | //! | `session/set_mode` | mutate the session's mode (gated vs. bypass-permissions) | //! | `session/set_model` | switch the session's active model (endpoint:model selector) | //! | (anything else) | "not implemented yet" error | //! //! Stage 5 flipped on image content. Stage 6 starts adding new wire //! protocols (Anthropic /v1/messages first). use std::path::PathBuf; use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; use agent_client_protocol::schema::{ AgentCapabilities, CancelNotification, ContentBlock, InitializeRequest, InitializeResponse, ListSessionsRequest, ListSessionsResponse, LoadSessionRequest, LoadSessionResponse, ModelId, ModelInfo as AcpModelInfo, NewSessionRequest, NewSessionResponse, PromptCapabilities, PromptRequest, PromptResponse, SessionCapabilities, SessionId, SessionInfo, SessionListCapabilities, SessionMode, SessionModeId, SessionModeState, SessionModelState, SessionNotification, SessionUpdate, SetSessionModeRequest, SetSessionModeResponse, SetSessionModelRequest, SetSessionModelResponse, StopReason, TextContent, }; use agent_client_protocol::{Agent as AgentRole, Client, ConnectionTo, Dispatch, Stdio}; use futures::StreamExt; use std::collections::BTreeMap; use tokio_util::sync::CancellationToken; use crate::compaction; use crate::config::{Config, parse_model_selector}; use crate::prompt::build_system_prompt; use crate::provider::{ CompletionEvent, CompletionRequest, Message, MessageContent, Provider, Role, ToolCall, }; use crate::session::{self, MODE_BYPASS, MODE_DEFAULT, MODE_PLAN, SessionState, SessionStore}; use crate::store::{self, PersistedSession}; use crate::tool_runner::{AcpClientOps, ToolCallEvent, dispatch_tool_call}; use crate::tools; /// Maximum number of provider→tool→provider round-trips per /// `session/prompt` request. Bound exists to keep a runaway model /// from looping forever; the spec maps this to /// [`StopReason::MaxTurnRequests`]. const MAX_TOOL_ROUNDS: usize = 25; /// Public entry point. Wraps an `Arc` so handlers can clone /// it cheaply into every closure. pub struct Agent { inner: Arc, } struct AgentInner { /// Every successfully-built provider, indexed positionally. We look /// providers up by name (`endpoint:` prefix) rather than by index. providers: Vec>, /// Name of the endpoint used when a request omits the /// `endpoint:model` prefix. default_endpoint_name: String, /// Default model for the default endpoint, if configured. Required /// for Stage 2 because session/set_model lands in Stage 4 — a /// session with no model can't prompt anything. default_model: Option, /// Per-endpoint `max_tokens` override. Looked up by endpoint /// name after resolution. `None` (or an absent entry) means the /// upstream picks its own default. max_tokens: std::collections::HashMap, /// Per-endpoint model context window in tokens. When set, the /// agent compacts history before each completion so the prompt /// fits inside `context_window - max_tokens - safety` tokens. /// Absent entry → no compaction (legacy behaviour). context_window: std::collections::HashMap, /// Aggregated list of selectable models across every configured /// endpoint, computed once at startup. With a single endpoint /// the model ids appear bare; with multiple endpoints every id /// carries the `endpoint:` prefix so the picker is unambiguous. /// Empty when every provider's `list_models` failed at startup — /// the dropdown then shows nothing and the session keeps using /// the configured `default_model`. available_models: Vec, sessions: SessionStore, system_prompt_path: Option, /// Monotonic counter for minting session ids. The wire format is /// `hxa-{n}` — short, debuggable, and the protocol doesn't require /// UUIDs for session ids (it only requires them for message ids /// behind an unstable flag). next_session_id: AtomicU64, } impl Agent { /// Construct an agent from a validated [`Config`] and the providers /// that were successfully built for each endpoint. /// /// `async` because we call `Provider::list_models` on every /// provider up-front so the model-picker dropdown is populated /// from the very first `session/new`. Per-endpoint failure /// warns and skips rather than aborting startup — a single /// unreachable endpoint shouldn't take down the agent. pub async fn new(cfg: &Config, providers: Vec>) -> anyhow::Result { if providers.is_empty() { anyhow::bail!("no usable providers"); } let default = cfg.default_endpoint(); // The default endpoint's provider must have built successfully — // otherwise we can't honour `model = "bare-model-id"` requests. // (If only a non-default endpoint is usable, the operator should // promote it to `default_endpoint` in the TOML.) if !providers.iter().any(|p| p.name() == default.name) { anyhow::bail!( "default endpoint '{}' has no usable provider — check config", default.name ); } let max_tokens = cfg .endpoints .iter() .filter_map(|ep| ep.max_tokens.map(|m| (ep.name.clone(), m))) .collect(); let context_window = cfg .endpoints .iter() .filter_map(|ep| ep.context_window.map(|w| (ep.name.clone(), w))) .collect(); let available_models = aggregate_models(&providers).await; tracing::info!( models = available_models.len(), endpoints = providers.len(), "model catalogue assembled" ); Ok(Self { inner: Arc::new(AgentInner { providers, default_endpoint_name: default.name.clone(), default_model: default.default_model.clone(), max_tokens, context_window, available_models, sessions: session::new_store(), system_prompt_path: cfg.system_prompt_path.clone(), next_session_id: AtomicU64::new(1), }), }) } /// Run the agent against an ACP transport (typically [`Stdio`]). /// Returns when the transport closes or a handler errors. pub async fn serve(self, transport: Stdio) -> agent_client_protocol::Result<()> { let inner = self.inner; AgentRole .builder() .name("helexa-acp") .on_receive_request( async move |req: InitializeRequest, responder, _cx| { responder.respond(initialize_response(&req)) }, agent_client_protocol::on_receive_request!(), ) .on_receive_request( { let inner = inner.clone(); async move |req: NewSessionRequest, responder, _cx| { let result = handle_new_session(&inner, req).await; match result { Ok(resp) => responder.respond(resp), Err(e) => responder.respond_with_internal_error(format!("{e:#}")), } } }, agent_client_protocol::on_receive_request!(), ) .on_receive_request( { let inner = inner.clone(); async move |req: LoadSessionRequest, responder, cx: ConnectionTo| { let session_id = req.session_id.clone(); match handle_load_session(&inner, req).await { Ok((resp, history)) => { let send_result = responder.respond(resp); // History replay happens off the // dispatch loop so the load reply // returns immediately. Zed receives // the response, then sees a stream // of session/update events that // repopulate the chat panel. let cx_clone = cx.clone(); let _ = cx.spawn(async move { replay_history(&cx_clone, &session_id, &history); Ok(()) }); send_result } Err(e) => responder.respond_with_internal_error(format!("{e:#}")), } } }, agent_client_protocol::on_receive_request!(), ) .on_receive_request( async move |req: ListSessionsRequest, responder, _cx| match handle_list_sessions( req, ) { Ok(resp) => responder.respond(resp), Err(e) => responder.respond_with_internal_error(format!("{e:#}")), }, agent_client_protocol::on_receive_request!(), ) .on_receive_request( { let inner = inner.clone(); async move |req: PromptRequest, responder, cx: ConnectionTo| { spawn_prompt(inner.clone(), cx, req, responder) } }, agent_client_protocol::on_receive_request!(), ) .on_receive_request( { let inner = inner.clone(); async move |req: SetSessionModeRequest, responder, _cx| { match handle_set_session_mode(&inner, req).await { Ok(()) => responder.respond(SetSessionModeResponse::new()), Err(e) => responder.respond_with_internal_error(format!("{e:#}")), } } }, agent_client_protocol::on_receive_request!(), ) .on_receive_request( { let inner = inner.clone(); async move |req: SetSessionModelRequest, responder, _cx| { match handle_set_session_model(&inner, req).await { Ok(()) => responder.respond(SetSessionModelResponse::new()), Err(e) => responder.respond_with_internal_error(format!("{e:#}")), } } }, agent_client_protocol::on_receive_request!(), ) .on_receive_notification( { let inner = inner.clone(); async move |notif: CancelNotification, _cx: ConnectionTo| { handle_cancel(&inner, notif).await; Ok(()) } }, agent_client_protocol::on_receive_notification!(), ) .on_receive_dispatch( async move |message: Dispatch, cx: ConnectionTo| { // `Dispatch` has three variants. For Request and // Notification we want the "not implemented yet" // error response. For *Response* we MUST forward // the result to its awaiting `ResponseRouter` — // otherwise our own outbound ACP calls // (`fs/read_text_file`, `session/request_permission`, // `terminal/*`, …) get their replies silently // overwritten with whatever error we'd send a // peer for an unknown method. That's how Stage 3 // tool dispatches were appearing as // "Internal error: not implemented yet" results // to the model. match message { Dispatch::Response(result, router) => router.respond_with_result(result), other => { tracing::warn!( method = ?other.method(), "unhandled ACP message" ); other.respond_with_error( agent_client_protocol::util::internal_error("not implemented yet"), cx, ) } } }, agent_client_protocol::on_receive_dispatch!(), ) .connect_to(transport) .await } } fn initialize_response(req: &InitializeRequest) -> InitializeResponse { // Stage 5: image is on (Zed clipboard / drag-drop). Audio and // embedded resources flip on in later stages. let prompt_caps = PromptCapabilities::default().image(true); // Stage 3b: advertise both the top-level `load_session` flag and // the `session/list` sub-capability. Zed (and other ACP clients) // uses `session/list` to discover the session id that belongs to // a workspace before sending `session/load` — without it, the // client only knows how to mint new sessions and resume never // fires regardless of what's on disk. let session_caps = SessionCapabilities::default().list(Some(SessionListCapabilities::default())); InitializeResponse::new(req.protocol_version).agent_capabilities( AgentCapabilities::new() .prompt_capabilities(prompt_caps) .session_capabilities(session_caps) .load_session(true), ) } async fn handle_new_session( inner: &AgentInner, req: NewSessionRequest, ) -> anyhow::Result { if !req.cwd.is_absolute() { anyhow::bail!("session cwd must be absolute, got {}", req.cwd.display()); } let model_id = inner .default_model .clone() .ok_or_else(|| anyhow::anyhow!( "default endpoint '{}' has no default_model — set one in config or wait for Stage 4 set_model", inner.default_endpoint_name ))?; let n = inner.next_session_id.fetch_add(1, Ordering::Relaxed); let session_id = SessionId::new(format!("hxa-{n}")); let cwd_display = req.cwd.display().to_string(); let log_model = model_id.clone(); let state = SessionState::new(req.cwd, model_id); session::insert(&inner.sessions, session_id.clone(), state).await; tracing::info!( session_id = %session_id.0, model_id = %log_model, cwd = %cwd_display, "session created" ); let resp = NewSessionResponse::new(session_id).modes(default_mode_state()); let resp = match session_model_state(inner, &log_model) { Some(models) => resp.models(Some(models)), None => resp, }; Ok(resp) } /// Rehydrate a session from disk. /// /// Behaviour: /// /// - Reads the persisted JSON from /// `$XDG_DATA_HOME/helexa-acp/sessions/{id}.json`. Missing file → /// error (Zed falls back to `session/new`). /// - Overwrites the persisted `cwd` with the one the client just /// sent. The user may have moved or symlinked the repo since /// the session was first created; the *current* cwd is the /// right place to root subsequent tool dispatches. /// - Materialises an in-memory `SessionState` with the persisted /// model + mode + history. /// - Returns `LoadSessionResponse` carrying the same mode list as /// `session/new`, plus the persisted `current_mode_id` so the /// client renders the mode dropdown in the correct state. async fn handle_load_session( inner: &AgentInner, req: LoadSessionRequest, ) -> anyhow::Result<(LoadSessionResponse, Vec)> { if !req.cwd.is_absolute() { anyhow::bail!("session cwd must be absolute, got {}", req.cwd.display()); } let persisted = store::load(&req.session_id)?; // Snapshot the values we need for logging, the response, and // the post-load history replay before we move pieces of // `persisted` into `state`. let model_id = persisted.model_id.clone(); let mode_id = persisted.mode_id.clone(); let history_for_replay = persisted.history.clone(); let history_turns = persisted.history.len(); let mut state = SessionState::new(req.cwd.clone(), persisted.model_id); state.history = persisted.history; state.mode_id = SessionModeId::new(persisted.mode_id); session::insert(&inner.sessions, req.session_id.clone(), state).await; tracing::info!( session_id = %req.session_id.0, model_id = %model_id, mode = %mode_id, cwd = %req.cwd.display(), history_turns, "session loaded from disk" ); let modes = SessionModeState::new( SessionModeId::new(mode_id), default_mode_state().available_modes, ); let resp = LoadSessionResponse::new().modes(modes); let resp = match session_model_state(inner, &model_id) { Some(models) => resp.models(Some(models)), None => resp, }; Ok((resp, history_for_replay)) } /// Re-emit a session's persisted history as `session/update` /// notifications so an ACP client (Zed) can render the prior chat /// after a `session/load`. Without this, even a successful load /// leaves the agent panel blank because Zed doesn't cache the /// transcript client-side for custom agent_servers entries — that /// caching only happens for first-party agents where Zed itself /// owns the conversation state. /// /// Mapping: /// /// - `Role::User` text → `SessionUpdate::UserMessageChunk` /// - `Role::Assistant` text → `SessionUpdate::AgentMessageChunk` /// - `Role::Assistant` with tool calls → text chunk (if any) plus /// one `ToolCall` event per call. We emit each with status = /// `Completed` because the call already ran; the matching /// `Role::Tool` result message is folded into the card's /// content via a subsequent `ToolCallUpdate`. /// - `Role::Tool` (tool result) → `ToolCallUpdate` carrying the /// result text, keyed by `tool_call_id` so it lands on the /// right card. /// - `Role::System` → skipped; system prompts aren't rendered. fn replay_history(cx: &ConnectionTo, session_id: &SessionId, history: &[Message]) { use agent_client_protocol::schema::{ Content, ToolCall as AcpToolCall, ToolCallContent, ToolCallId, ToolCallStatus, ToolCallUpdate, ToolCallUpdateFields, }; fn tool_kind_for(name: &str) -> agent_client_protocol::schema::ToolKind { use agent_client_protocol::schema::ToolKind; match name { "read_file" | "list_dir" => ToolKind::Read, "write_file" | "edit_file" => ToolKind::Edit, "bash" => ToolKind::Execute, _ => ToolKind::Other, } } fn title_for(name: &str, args_json: &str) -> String { match ( name, serde_json::from_str::(args_json).ok(), ) { ("read_file", Some(v)) => format!( "Read {}", v.get("path").and_then(|p| p.as_str()).unwrap_or("?") ), ("write_file", Some(v)) => format!( "Write {}", v.get("path").and_then(|p| p.as_str()).unwrap_or("?") ), ("edit_file", Some(v)) => format!( "Edit {}", v.get("path").and_then(|p| p.as_str()).unwrap_or("?") ), ("list_dir", Some(v)) => format!( "List {}", v.get("path").and_then(|p| p.as_str()).unwrap_or("?") ), ("bash", Some(v)) => { let cmd = v.get("command").and_then(|p| p.as_str()).unwrap_or("?"); let snippet = if cmd.len() > 60 { format!("{}…", &cmd[..60]) } else { cmd.to_string() }; format!("Run: {snippet}") } (other, _) => format!("Tool: {other}"), } } let send = |update: SessionUpdate| { let notif = SessionNotification::new(session_id.clone(), update); if let Err(e) = cx.send_notification(notif) { tracing::warn!( error = %format!("{e:#}"), "replay: failed to forward history event" ); } }; let mut total_events: usize = 0; for msg in history { match (msg.role, &msg.content) { (Role::User, MessageContent::Text { text }) => { send(SessionUpdate::UserMessageChunk(text_chunk(text.clone()))); total_events += 1; } (Role::User, MessageContent::MultiPart { parts }) => { // We can re-emit text parts as UserMessageChunks. // Images get a placeholder line so the user sees // *that* an image was attached; re-replaying the // image bytes themselves through ACP would require // a path round-trip we don't currently keep. for part in parts { match part { crate::provider::MessagePart::Text { text } => { send(SessionUpdate::UserMessageChunk(text_chunk(text.clone()))); total_events += 1; } crate::provider::MessagePart::Image(img) => { let label = match &img.uri { Some(u) => format!("[image: {u}]"), None => { format!("[image: {} ({} bytes)]", img.mime_type, img.data.len()) } }; send(SessionUpdate::UserMessageChunk(text_chunk(label))); total_events += 1; } } } } (Role::Assistant, MessageContent::Text { text }) => { send(SessionUpdate::AgentMessageChunk(text_chunk(text.clone()))); total_events += 1; } (Role::Assistant, MessageContent::ToolCalls { text, calls }) => { if let Some(t) = text && !t.is_empty() { send(SessionUpdate::AgentMessageChunk(text_chunk(t.clone()))); total_events += 1; } for call in calls { let raw_input = serde_json::from_str::(&call.arguments) .unwrap_or_else(|_| serde_json::Value::String(call.arguments.clone())); let card = AcpToolCall::new( ToolCallId::new(call.id.clone()), title_for(&call.name, &call.arguments), ) .kind(tool_kind_for(&call.name)) .status(ToolCallStatus::Completed) .raw_input(raw_input); send(SessionUpdate::ToolCall(card)); total_events += 1; } } ( Role::Tool, MessageContent::ToolResult { tool_call_id, content, }, ) => { let update = ToolCallUpdate::new( ToolCallId::new(tool_call_id.clone()), ToolCallUpdateFields::new() .status(ToolCallStatus::Completed) .content(vec![ToolCallContent::Content(Content::new( ContentBlock::Text(TextContent::new(content.clone())), ))]), ); send(SessionUpdate::ToolCallUpdate(update)); total_events += 1; } (Role::System, _) => { // System prompts aren't shown in the chat panel. } (role, content) => { tracing::debug!( ?role, ?content, "replay: unrecognised (role, content) shape; skipping" ); } } } tracing::info!( session_id = %session_id.0, events = total_events, history_turns = history.len(), "session history replayed to client" ); } /// Enumerate persisted sessions for the `session/list` ACP method. /// /// Zed calls this on workspace open to find the session belonging /// to the cwd it's reopening — without it, even though `session/load` /// works, the client has no way to discover the session_id and /// always falls back to `session/new`. That's exactly the /// "history didn't survive the restart" symptom. /// /// Cursor pagination from the request is accepted but ignored: /// helexa-acp's session counts are too small to need it. We always /// return the whole filtered list with `next_cursor = None`. fn handle_list_sessions(req: ListSessionsRequest) -> anyhow::Result { let sessions = store::list(req.cwd.as_deref())?; let infos: Vec = sessions .into_iter() .map(|s| { let mut info = SessionInfo::new(SessionId::new(s.session_id), s.cwd); info = info.title(derive_session_title(&s.history)); info = info.updated_at(store::unix_to_iso8601(s.updated_at)); info }) .collect(); tracing::info!( cwd = ?req.cwd, count = infos.len(), "session/list responded" ); Ok(ListSessionsResponse::new(infos)) } /// Best-effort human-readable title for a session, derived from the /// first user turn's text (truncated to ~60 chars). Empty string /// becomes `None` so Zed can fall back to its own placeholder. fn derive_session_title(history: &[Message]) -> Option { use crate::provider::MessagePart; history .iter() .find_map(|msg| match (msg.role, &msg.content) { (Role::User, MessageContent::Text { text }) => Some(text.clone()), (Role::User, MessageContent::MultiPart { parts }) => parts.iter().find_map(|p| { if let MessagePart::Text { text } = p { Some(text.clone()) } else { None } }), _ => None, }) .map(|s| { let trimmed = s.trim(); if trimmed.chars().count() > 60 { let prefix: String = trimmed.chars().take(60).collect(); format!("{prefix}…") } else { trimmed.to_string() } }) .filter(|s| !s.is_empty()) } /// Build the model catalogue advertised in `NewSessionResponse.models` /// (and the resume equivalent). Walks every provider, calls /// `list_models`, and prefixes ids with `endpoint:` when the user has /// more than one endpoint configured. A failing endpoint logs and /// contributes nothing — losing one endpoint must not blank the whole /// dropdown. async fn aggregate_models(providers: &[Arc]) -> Vec { let multi_endpoint = providers.len() > 1; let mut out: Vec = Vec::new(); for provider in providers { let endpoint = provider.name().to_string(); match provider.list_models().await { Ok(models) => { tracing::info!( endpoint = %endpoint, count = models.len(), "fetched models from endpoint" ); for m in models { let id = if multi_endpoint { format!("{endpoint}:{}", m.id) } else { m.id.clone() }; let display = m.display_name.unwrap_or_else(|| m.id.clone()); let info = AcpModelInfo::new(ModelId::new(id), display) .description(Some(format!("endpoint: {endpoint}"))); out.push(info); } } Err(e) => { tracing::warn!( endpoint = %endpoint, error = %format!("{e:#}"), "list_models failed; this endpoint's models won't appear in the picker" ); } } } out } /// Build the `SessionModelState` that Zed renders as the /// model-picker dropdown. The current model id is exactly what /// the session is using right now (already in `endpoint:model` /// form if it was set that way). Returns `None` when the /// catalogue is empty — no point showing an empty dropdown. fn session_model_state(inner: &AgentInner, current: &str) -> Option { if inner.available_models.is_empty() { return None; } Some(SessionModelState::new( ModelId::new(current.to_string()), inner.available_models.clone(), )) } async fn handle_set_session_model( inner: &AgentInner, req: SetSessionModelRequest, ) -> anyhow::Result<()> { let Some(state) = session::get(&inner.sessions, &req.session_id).await else { anyhow::bail!("unknown session id {}", req.session_id.0); }; let target = req.model_id.0.as_ref().to_string(); // Validate the requested model id resolves to a configured // provider. We don't require it to appear in `available_models` // because the catalogue may be stale (endpoint added a model // after startup) and rejecting unknown ids would be too rigid. // Provider lookup is the actual source of truth. let (_, _) = resolve_provider(&inner.providers, &inner.default_endpoint_name, &target) .map_err(|e| anyhow::anyhow!("set_session_model: {e:#}"))?; // Persist the new model id on the session under the mutex, // then snapshot for disk persistence outside the lock. let snapshot = { let mut s = state.lock().await; s.model_id = target.clone(); PersistedSession { session_id: req.session_id.0.as_ref().to_string(), cwd: s.cwd.clone(), model_id: s.model_id.clone(), mode_id: s.mode_id.0.as_ref().to_string(), history: s.history.clone(), created_at: store::now_secs(), updated_at: store::now_secs(), } }; if let Err(e) = store::save(&snapshot) { tracing::warn!( session_id = %req.session_id.0, error = %format!("{e:#}"), "session persist after set_model failed; on-disk model id stays stale" ); } tracing::info!( session_id = %req.session_id.0, model_id = %target, "session model changed" ); Ok(()) } /// The three modes every Stage 3 session advertises: /// /// - **Default** — writes / bash prompt the user. /// - **Bypass Permissions** — auto-allow. /// - **Plan** — read-and-plan-only. Writes are restricted to a /// per-project plan directory under `$XDG_DATA_HOME/helexa-acp/plans/` /// and bash is disabled. Designed for "draft the implementation /// plan, then I'll review and let you execute" flows. fn default_mode_state() -> SessionModeState { SessionModeState::new( SessionModeId::new(MODE_DEFAULT), vec![ SessionMode::new(SessionModeId::new(MODE_DEFAULT), "Default") .description("Prompt for permission before writes or shell commands."), SessionMode::new(SessionModeId::new(MODE_BYPASS), "Bypass Permissions") .description("Auto-allow all tool calls. Use with care."), SessionMode::new(SessionModeId::new(MODE_PLAN), "Plan") .description("Write plans to the plan directory; no shell, no writes outside it."), ], ) } async fn handle_set_session_mode( inner: &AgentInner, req: SetSessionModeRequest, ) -> anyhow::Result<()> { let Some(state) = session::get(&inner.sessions, &req.session_id).await else { anyhow::bail!("unknown session id {}", req.session_id.0); }; let accepted = matches!( req.mode_id.0.as_ref(), MODE_DEFAULT | MODE_BYPASS | MODE_PLAN ); if !accepted { anyhow::bail!( "unknown mode '{}' — must be one of: {}, {}, {}", req.mode_id.0, MODE_DEFAULT, MODE_BYPASS, MODE_PLAN ); } state.lock().await.mode_id = req.mode_id.clone(); tracing::info!( session_id = %req.session_id.0, mode = %req.mode_id.0, "session mode changed" ); Ok(()) } async fn handle_cancel(inner: &AgentInner, notif: CancelNotification) { let Some(state) = session::get(&inner.sessions, ¬if.session_id).await else { tracing::debug!(session_id = %notif.session_id.0, "cancel for unknown session, ignoring"); return; }; let cancel = state.lock().await.cancel.clone(); tracing::info!(session_id = %notif.session_id.0, "cancellation requested"); cancel.cancel(); } /// Kick the prompt off on a spawned task so the event loop is free to /// dispatch the matching `session/cancel`. The handler itself returns /// `Ok(())` immediately (= `Handled::Yes`); the spawned task is what /// eventually consumes `responder`. fn spawn_prompt( inner: Arc, cx: ConnectionTo, req: PromptRequest, responder: agent_client_protocol::Responder, ) -> agent_client_protocol::Result<()> { let task_cx = cx.clone(); cx.spawn(async move { if let Err(e) = drive_prompt(inner, task_cx, req, responder).await { // `drive_prompt` already consumed the responder on the // error paths it produces; this branch only fires if the // task itself errored before reaching responder.respond. // Log and swallow — propagating the error would tear down // the whole connection, which is too violent for one // failed prompt. tracing::error!(error = %format!("{e:#}"), "prompt task failed"); } Ok(()) })?; Ok(()) } async fn drive_prompt( inner: Arc, cx: ConnectionTo, req: PromptRequest, responder: agent_client_protocol::Responder, ) -> anyhow::Result<()> { let session_id = req.session_id.clone(); let Some(session_arc) = session::get(&inner.sessions, &session_id).await else { let _ = responder.respond_with_internal_error(format!("unknown session id {}", session_id.0)); return Ok(()); }; // Snapshot the inputs under the session lock, then drop the lock // before any `await` that touches the network. `mode_id` is // refreshed at the top of every round (the user can toggle modes // mid-turn and we want the next round's streaming + tool gating // to reflect that). let (existing_history, model_id, cwd, cancel, mut mode_id) = { let mut state = session_arc.lock().await; // Fire the session's current cancel before installing a new // one. If a previous prompt task is still in-flight (model // stalled mid-stream, a long-running bash, a wedged ACP // roundtrip), this lets it observe is_cancelled() at the // next .await and unwind cleanly — instead of two tasks // racing each other to mutate session.history and to // persist the same file. state.cancel.cancel(); let cancel = CancellationToken::new(); state.cancel = cancel.clone(); let user_content = flatten_prompt(&req.prompt); state.history.push(Message { role: Role::User, content: user_content, }); ( state.history.clone(), state.model_id.clone(), state.cwd.clone(), cancel, state.mode_id.clone(), ) }; let tool_specs = tools::all_tools(); // Plan-mode write target. Resolved once because the cwd doesn't // change for a session's lifetime; the directory is created // lazily by the runtime when a write lands inside it. let plan_dir = store::plan_dir_for(&cwd); let (provider, local_model) = match resolve_provider(&inner.providers, &inner.default_endpoint_name, &model_id) { Ok(pair) => pair, Err(e) => { let _ = responder.respond_with_internal_error(format!("{e:#}")); return Ok(()); } }; tracing::info!( session_id = %session_id.0, endpoint = %provider.name(), model = %local_model, mode = %mode_id.0, history_turns = existing_history.len(), "sending prompt upstream" ); let ops = AcpClientOps::new(cx.clone()); // `messages` is the rolling conversation we send to the provider // each round. Slot 0 is the system prompt — rebuilt at the top // of every round so a mid-turn mode toggle takes effect. We seed // a placeholder here and overwrite it on the first iteration. let mut messages: Vec = Vec::with_capacity(existing_history.len() + 1); messages.push(Message { role: Role::System, content: MessageContent::Text { text: String::new(), }, }); messages.extend(existing_history); // Buffer for turns produced this round. Flushed into // session.history *and* persisted at the end of every iteration // (and once more after the loop). Per-round persistence means // a stall later in the conversation doesn't lose earlier rounds. let mut new_turns: Vec = Vec::new(); // Monotonic counter for synthetic ids assigned to unparseable // blocks across all rounds of this prompt. let mut next_malformed_index: usize = 0; let mut stop_reason = StopReason::EndTurn; for round in 0..MAX_TOOL_ROUNDS { if cancel.is_cancelled() { stop_reason = StopReason::Cancelled; break; } // Refresh mode + rebuild system prompt at the top of every // round. Cheap (one mutex acquisition + one string build); // the win is that if the user flips the mode dropdown // mid-turn — particularly the Plan ↔ Bypass transitions // the plan-mode menu invites them to make — the new mode // gates both this round's streaming *and* its tool // dispatch. mode_id = session_arc.lock().await.mode_id.clone(); let system_prompt = build_system_prompt( &cwd, inner.system_prompt_path.as_deref(), &tool_specs, &mode_id, plan_dir.as_deref(), ) .map_err(|e| anyhow::anyhow!("build system prompt: {e:#}"))?; messages[0] = Message { role: Role::System, content: MessageContent::Text { text: system_prompt, }, }; tracing::info!( session_id = %session_id.0, round = round + 1, of = MAX_TOOL_ROUNDS, mode = %mode_id.0, history_turns = messages.len(), "prompt round: streaming" ); // Project history into the model's context window when the // endpoint advertises one. Compaction is a per-request // *projection* — `messages` (and the persisted session // history downstream) stay intact; only what we send // upstream shrinks. Without this, a 32 K Qwen3 dies after // the first few `read_file` results pile up in history. let provider_max_tokens = inner.max_tokens.get(provider.name()).copied(); let messages_for_provider = match inner.context_window.get(provider.name()).copied() { Some(ctx) => { let budget = prompt_budget(ctx, provider_max_tokens); let (compacted, stats) = compaction::compact_to_budget(&messages, budget); if stats.elided_messages > 0 { tracing::info!( session_id = %session_id.0, round = round + 1, context_window = ctx, budget, original_tokens = stats.original_tokens, final_tokens = stats.final_tokens, elided = stats.elided_messages, "context compaction applied" ); } compacted } None => messages.clone(), }; // Tool descriptions reach the model via the Qwen3 `# Tools` // block in the system prompt, not via the OpenAI `tools` // request field — cortex/neuron pass that field through to // the encoder unread, and including it would double-describe // tools once a strict-OpenAI backend lands. Leave empty. let completion_req = CompletionRequest { model: local_model.clone(), messages: messages_for_provider, tools: vec![], temperature: None, top_p: None, max_tokens: provider_max_tokens, }; let mut stream = match provider.complete(completion_req, cancel.clone()).await { Ok(s) => s, Err(e) => { let _ = responder .respond_with_internal_error(format!("{} complete: {e:#}", provider.name())); return Ok(()); } }; let mut assistant_text = String::new(); let mut finish_reason: Option = None; // `BTreeMap` keyed by the provider's tool-call index keeps // insertion order while allowing arg deltas to mutate any // bucket — `ToolCallStart` may arrive interleaved with // `ToolCallArgsDelta` for different indices. let mut tool_buckets: BTreeMap = BTreeMap::new(); // blocks whose JSON couldn't be parsed even with // qwen3's repair pass. We surface each as a Failed // ToolCall card and feed a synthetic error result back to // the model so it can retry on the next round. let mut malformed_calls: Vec = Vec::new(); while let Some(event) = stream.next().await { let event = match event { Ok(e) => e, Err(e) => { tracing::warn!(error = %format!("{e:#}"), "stream error; ending round"); break; } }; match event { CompletionEvent::TextDelta(t) => { assistant_text.push_str(&t); send_chunk( &cx, &session_id, SessionUpdate::AgentMessageChunk(text_chunk(t)), ); } CompletionEvent::ReasoningDelta(t) => { send_chunk( &cx, &session_id, SessionUpdate::AgentThoughtChunk(text_chunk(t)), ); } CompletionEvent::ToolCallStart { index, id, name } => { tool_buckets.insert( index, ToolCallBucket { id, name, arguments: String::new(), }, ); } CompletionEvent::ToolCallArgsDelta { index, args_delta } => { tool_buckets .entry(index) .or_default() .arguments .push_str(&args_delta); } CompletionEvent::MalformedToolCall { raw } => { malformed_calls.push(raw); } CompletionEvent::Finish { reason } => finish_reason = reason, CompletionEvent::Usage(_) => {} } } if cancel.is_cancelled() { stop_reason = StopReason::Cancelled; // Persist any partial text so the next turn has context. if !assistant_text.is_empty() { new_turns.push(Message { role: Role::Assistant, content: MessageContent::Text { text: assistant_text, }, }); } break; } // Recovery pass before deciding "is there work to do?". // For each malformed body, try shape-based inference // against the tool catalogue (handles the "model emitted // `arguments` but forgot `name`" case). Successes get // promoted to real tool buckets; failures stay in // `malformed_calls` for the Failed-card path below. malformed_calls.retain(|raw| match try_repair_missing_name(raw) { Some((name, args_json)) => { let idx = tool_buckets .keys() .max() .copied() .map(|m| m + 1) .unwrap_or(0); tracing::debug!( inferred_name = %name, index = idx, "qwen3: recovered missing-name tool call via shape inference" ); tool_buckets.insert( idx, ToolCallBucket { id: format!("call_recovered_{idx}"), name, arguments: args_json, }, ); false } None => true, }); let has_tool_calls = !tool_buckets.is_empty(); let has_malformed = !malformed_calls.is_empty(); if !has_tool_calls && !has_malformed { // Terminal turn: just text. Save and finish. if !assistant_text.is_empty() { new_turns.push(Message { role: Role::Assistant, content: MessageContent::Text { text: assistant_text, }, }); } stop_reason = map_finish_reason(finish_reason.as_deref()); break; } // Assistant turn carrying any successfully-parsed tool calls // (malformed ones are handled separately so each gets its // own Failed card with its raw body intact). let calls: Vec = tool_buckets .values() .map(|b| ToolCall { id: b.id.clone(), name: b.name.clone(), arguments: b.arguments.clone(), }) .collect(); if has_tool_calls || !assistant_text.is_empty() { let assistant_turn = Message { role: Role::Assistant, content: if has_tool_calls { MessageContent::ToolCalls { text: (!assistant_text.is_empty()).then_some(assistant_text), calls, } } else { MessageContent::Text { text: assistant_text, } }, }; new_turns.push(assistant_turn.clone()); messages.push(assistant_turn); } // Dispatch every tool call sequentially. Parallelism is // tempting but would require Zed to handle interleaved // permission prompts; serial is friendlier. for bucket in tool_buckets.into_values() { if cancel.is_cancelled() { stop_reason = StopReason::Cancelled; break; } let event = ToolCallEvent { id: bucket.id, name: bucket.name, arguments: bucket.arguments, }; tracing::info!( session_id = %session_id.0, tool = %event.name, tool_call_id = %event.id, "dispatch tool" ); let result = dispatch_tool_call(&ops, &session_id, &mode_id, &cwd, event, &cancel).await; tracing::info!( session_id = %session_id.0, tool_call_id = %result.tool_call_id, is_error = result.is_error, "dispatch tool complete" ); let result_turn = Message { role: Role::Tool, content: MessageContent::ToolResult { tool_call_id: result.tool_call_id, content: result.content, }, }; new_turns.push(result_turn.clone()); messages.push(result_turn); } // Handle malformed calls last — each becomes a Failed // SessionUpdate::ToolCall card (so Zed renders structured // failure UI instead of dumping raw JSON inline) plus a // synthetic tool-result message so the model gets concrete // feedback for self-correction on the next round. for raw in malformed_calls.drain(..) { if cancel.is_cancelled() { stop_reason = StopReason::Cancelled; break; } let synthetic_id = next_synthetic_id(&mut next_malformed_index); emit_malformed_tool_card(&cx, &session_id, &synthetic_id, &raw); let (call_turn, result_turn) = synthesize_malformed_history(&synthetic_id, &raw); new_turns.push(call_turn.clone()); messages.push(call_turn); new_turns.push(result_turn.clone()); messages.push(result_turn); } if cancel.is_cancelled() { stop_reason = StopReason::Cancelled; break; } if round + 1 == MAX_TOOL_ROUNDS { tracing::warn!( session_id = %session_id.0, rounds = MAX_TOOL_ROUNDS, "hit MAX_TOOL_ROUNDS, returning MaxTurnRequests" ); stop_reason = StopReason::MaxTurnRequests; } // Per-round flush: push this round's turns into the in-memory // history and persist to disk. If the model stalls in a later // round (long bash, upstream SSE that never finishes, etc.) // earlier rounds still survive a binary restart. if !new_turns.is_empty() { let drained = std::mem::take(&mut new_turns); tracing::info!( session_id = %session_id.0, round = round + 1, turns = drained.len(), "prompt round complete; persisting" ); extend_and_persist(&session_arc, &session_id, drained).await; } } // Final flush for whatever the break paths above left behind. // No-op when the per-round flush already drained new_turns. if !new_turns.is_empty() { extend_and_persist(&session_arc, &session_id, new_turns).await; } tracing::info!( session_id = %session_id.0, ?stop_reason, "prompt complete" ); let _ = responder.respond(PromptResponse::new(stop_reason)); Ok(()) } /// Push `new_turns` into the session's in-memory history under the /// session lock, then snapshot the full state and write it to disk /// *outside* the lock. Used by `drive_prompt` at the end of every /// tool-call round (so partial progress survives a stall) and once /// more after the loop (catching any turns the break paths left /// behind). /// /// Persistence failures log a warning and don't propagate — losing /// a save shouldn't tear down a live conversation. async fn extend_and_persist( session_arc: &Arc>, session_id: &SessionId, new_turns: Vec, ) { let snapshot = { let mut state = session_arc.lock().await; state.history.extend(new_turns); PersistedSession { session_id: session_id.0.as_ref().to_string(), cwd: state.cwd.clone(), model_id: state.model_id.clone(), mode_id: state.mode_id.0.as_ref().to_string(), history: state.history.clone(), // `created_at` ought to be preserved across saves — // currently SessionState doesn't carry it, so every // save refreshes both timestamps. Acceptable for // resume; future work: thread `created_at` through. created_at: store::now_secs(), updated_at: store::now_secs(), } }; if let Err(e) = store::save(&snapshot) { tracing::warn!( session_id = %session_id.0, error = %format!("{e:#}"), "session/persist failed; resume from disk will miss this round" ); } } /// Accumulator for one streamed tool call: the OpenAI wire format /// sends `id` + `name` once (in the first chunk for that index) and /// then argument bytes piecemeal. We gather them all before /// dispatching. #[derive(Debug, Default)] struct ToolCallBucket { id: String, name: String, arguments: String, } fn send_chunk(cx: &ConnectionTo, session_id: &SessionId, update: SessionUpdate) { let notif = SessionNotification::new(session_id.clone(), update); if let Err(e) = cx.send_notification(notif) { tracing::warn!(error = %format!("{e:#}"), "failed to forward session update"); } } fn text_chunk(text: String) -> agent_client_protocol::schema::ContentChunk { use agent_client_protocol::schema::ContentChunk; ContentChunk::new(ContentBlock::Text(TextContent::new(text))) } /// Mint a synthetic tool_call_id for a malformed `` block. /// The format mirrors successful calls (`call_`) but uses its own /// counter so the ids don't collide. fn next_synthetic_id(counter: &mut usize) -> String { let id = format!("call_malformed_{}", *counter); *counter += 1; id } /// Emit a `SessionUpdate::ToolCall` with `Failed` status so Zed /// renders the malformed call as a structured failure card (raw /// body visible inside the card) instead of leaving it as inline /// text in the message pane. fn emit_malformed_tool_card( cx: &ConnectionTo, session_id: &SessionId, tool_call_id: &str, raw: &str, ) { use agent_client_protocol::schema::{ Content, ToolCall as AcpToolCall, ToolCallContent, ToolCallId, ToolCallStatus, ToolKind, }; let body = format!( "Tool call JSON could not be parsed. Raw body:\n\n```\n{raw}\n```\n\n\ Expected schema:\n\n```json\n{{\"name\": \"\", \"arguments\": {{...}}}}\n```", ); let card = AcpToolCall::new(ToolCallId::new(tool_call_id), "Malformed tool call") .kind(ToolKind::Other) .status(ToolCallStatus::Failed) .raw_input(serde_json::Value::String(raw.to_string())) .content(vec![ToolCallContent::Content(Content::new( ContentBlock::Text(TextContent::new(body)), ))]); send_chunk(cx, session_id, SessionUpdate::ToolCall(card)); } /// Build the assistant-turn / tool-result pair for a malformed /// ``. The assistant turn carries the raw body verbatim /// (so the model sees its own previous output), and the tool /// result spells out *why* it failed with the expected schema — /// enough for a competent model to self-correct on the next round. /// Last-chance repair for a malformed `` body: if the /// model emitted a structurally-valid JSON object with `arguments` /// but a missing `name`, infer the intended tool from the /// arguments' shape (see [`tools::infer_tool_name`]). Returns /// `Some((name, arguments_json))` only when the inference is /// unambiguous; ambiguous or unrecognised shapes return `None` /// so the caller surfaces a Failed card. /// /// We don't try to repair anything qwen3.rs already gave up on for /// structural reasons (truncation, free-form prose) — those stay /// Failed and the model retries. fn try_repair_missing_name(raw: &str) -> Option<(String, String)> { let value: serde_json::Value = serde_json::from_str(raw.trim()).ok()?; // If a `name` exists at the top level, the parser's own // earlier repair passes already had a shot at this and decided // it was malformed for some other reason. Don't second-guess // them here. if value.get("name").is_some() { return None; } let arguments = value.get("arguments")?; let name = tools::infer_tool_name(arguments)?; let args_json = serde_json::to_string(arguments).ok()?; Some((name.to_string(), args_json)) } fn synthesize_malformed_history(tool_call_id: &str, raw: &str) -> (Message, Message) { let call = Message { role: Role::Assistant, content: MessageContent::ToolCalls { text: None, calls: vec![ToolCall { id: tool_call_id.to_string(), // Real tool names never start with `<` — using this // placeholder makes the malformed call's identity // unambiguous in the rendered transcript. name: "".to_string(), arguments: raw.to_string(), }], }, }; let result = Message { role: Role::Tool, content: MessageContent::ToolResult { tool_call_id: tool_call_id.to_string(), content: format!( "ERROR: previous body was not valid JSON. Body was:\n{raw}\n\n\ Retry with the schema: {{\"name\": \"\", \"arguments\": {{…}}}}" ), }, }; (call, result) } /// Compute the prompt token budget for an endpoint given its /// `context_window` and `max_tokens` settings. The model needs room /// for both the prompt and its response inside the context window, /// so the prompt budget is the remainder after subtracting the /// response cap (defaulting to a conservative 2048 when the endpoint /// didn't set one) and a small safety margin for tokenizer /// disagreement. /// /// The safety margin matters because our per-character estimate in /// [`compaction`] can drift a few percent from any given upstream /// tokenizer; we'd rather under-fill the context window than have a /// well-compacted history still trip `prompt_too_long`. fn prompt_budget(context_window: usize, max_tokens: Option) -> usize { const SAFETY_MARGIN: usize = 512; let max_tokens = max_tokens.unwrap_or(2048) as usize; context_window .saturating_sub(max_tokens) .saturating_sub(SAFETY_MARGIN) } fn map_finish_reason(reason: Option<&str>) -> StopReason { match reason { Some("length") => StopReason::MaxTokens, Some("refusal") => StopReason::Refusal, // "stop", "tool_calls" (no tools in Stage 2 — degrade to // EndTurn so we don't surface a bogus reason), missing, or // anything else → EndTurn. _ => StopReason::EndTurn, } } /// Pure helper — turn a prompt's ContentBlocks into the user-message /// content that goes into history. /// /// - All-text prompts collapse to [`MessageContent::Text`] (cheaper /// to encode upstream — many OpenAI-compatible servers prefer the /// string form when there's no reason to use the array form). /// - Anything with at least one image becomes /// [`MessageContent::MultiPart`], preserving block order so the /// user's "this image, then this text" pacing reaches the model. /// - `ResourceLink` is rendered as inline text so the model knows /// it was referenced. Audio and embedded resources aren't /// advertised as supported in [`PromptCapabilities`]; drop with a /// warning if a non-conformant client sends one. fn flatten_prompt(blocks: &[ContentBlock]) -> MessageContent { use crate::provider::{ImageData, MessagePart}; let mut parts: Vec = Vec::new(); let mut text_buf = String::new(); let flush_text = |buf: &mut String, parts: &mut Vec| { if !buf.is_empty() { parts.push(MessagePart::Text { text: std::mem::take(buf), }); } }; for block in blocks { match block { ContentBlock::Text(t) => { if !text_buf.is_empty() { text_buf.push_str("\n\n"); } text_buf.push_str(&t.text); } ContentBlock::ResourceLink(link) => { if !text_buf.is_empty() { text_buf.push_str("\n\n"); } text_buf.push_str(&format!("[resource link: {}]", link.uri)); } ContentBlock::Image(img) => { flush_text(&mut text_buf, &mut parts); parts.push(MessagePart::Image(ImageData { mime_type: img.mime_type.clone(), data: img.data.clone(), uri: img.uri.clone(), })); } other => { tracing::warn!(?other, "ignoring unsupported content block"); } } } flush_text(&mut text_buf, &mut parts); // Collapse to plain Text when there's no image part — the // OpenAI string-form is friendlier to non-vision endpoints // (some treat the array form as a vision-only path). let has_image = parts.iter().any(|p| matches!(p, MessagePart::Image(_))); if !has_image { let text = parts .into_iter() .filter_map(|p| match p { MessagePart::Text { text } => Some(text), MessagePart::Image(_) => None, }) .collect::>() .join("\n\n"); return MessageContent::Text { text }; } MessageContent::MultiPart { parts } } /// Pure helper — pick which provider handles a session's `model_id`. /// Returns the matching provider plus the endpoint-local model id /// (i.e. with any `endpoint:` prefix stripped). fn resolve_provider( providers: &[Arc], default_endpoint: &str, model_id: &str, ) -> anyhow::Result<(Arc, String)> { let (endpoint_hint, local_model) = parse_model_selector(model_id); let target_endpoint = endpoint_hint.unwrap_or(default_endpoint); let provider = providers .iter() .find(|p| p.name() == target_endpoint) .ok_or_else(|| anyhow::anyhow!("no provider for endpoint '{target_endpoint}'"))?; Ok((provider.clone(), local_model.to_string())) } #[cfg(test)] mod tests { use super::*; use agent_client_protocol::schema::ResourceLink; use async_trait::async_trait; use futures::stream::BoxStream; // ── flatten_prompt ────────────────────────────────────────────── fn expect_text(content: &MessageContent) -> &str { match content { MessageContent::Text { text } => text.as_str(), other => panic!("expected MessageContent::Text, got {other:?}"), } } #[test] fn flatten_empty_prompt_is_empty() { assert_eq!(expect_text(&flatten_prompt(&[])), ""); } #[test] fn flatten_joins_text_blocks_with_blank_line() { let blocks = vec![ ContentBlock::Text(TextContent::new("first")), ContentBlock::Text(TextContent::new("second")), ]; assert_eq!(expect_text(&flatten_prompt(&blocks)), "first\n\nsecond"); } #[test] fn flatten_resource_link_becomes_reference_line() { let blocks = vec![ContentBlock::ResourceLink(ResourceLink::new( "readme", "file:///tmp/x", ))]; assert_eq!( expect_text(&flatten_prompt(&blocks)), "[resource link: file:///tmp/x]" ); } #[test] fn flatten_text_and_image_produces_multipart_in_order() { use crate::provider::MessagePart; let blocks = vec![ ContentBlock::Text(TextContent::new("describe:")), ContentBlock::Image(agent_client_protocol::schema::ImageContent::new( "iVBORw0KGgo=", "image/png", )), ContentBlock::Text(TextContent::new("…in detail.")), ]; let content = flatten_prompt(&blocks); match content { MessageContent::MultiPart { parts } => { assert_eq!(parts.len(), 3); assert!(matches!(&parts[0], MessagePart::Text { text } if text == "describe:")); assert!(matches!(&parts[1], MessagePart::Image(img) if img.mime_type == "image/png" && img.data == "iVBORw0KGgo=")); assert!(matches!(&parts[2], MessagePart::Text { text } if text == "…in detail.")); } other => panic!("expected MultiPart, got {other:?}"), } } #[test] fn flatten_image_only_still_produces_multipart() { use crate::provider::MessagePart; let blocks = vec![ContentBlock::Image( agent_client_protocol::schema::ImageContent::new("AAAA", "image/jpeg"), )]; match flatten_prompt(&blocks) { MessageContent::MultiPart { parts } => { assert_eq!(parts.len(), 1); assert!(matches!(&parts[0], MessagePart::Image(img) if img.mime_type == "image/jpeg")); } other => panic!("expected MultiPart, got {other:?}"), } } // ── resolve_provider ──────────────────────────────────────────── /// Minimal Provider stub; just records its name. The trait methods /// aren't exercised by resolve_provider so we leave them /// unimplemented. struct StubProvider(&'static str); #[async_trait] impl Provider for StubProvider { fn name(&self) -> &str { self.0 } async fn list_models(&self) -> anyhow::Result> { unimplemented!() } async fn complete( &self, _request: CompletionRequest, _cancel: CancellationToken, ) -> anyhow::Result>> { unimplemented!() } } /// Provider stub whose `list_models` returns canned results. /// Used by the `aggregate_models` tests. struct ModelProvider { name: &'static str, models: anyhow::Result>, } impl ModelProvider { fn ok(name: &'static str, ids: &[&str]) -> Arc { let models = ids .iter() .map(|id| crate::provider::ModelInfo { id: (*id).to_string(), display_name: None, }) .collect(); Arc::new(Self { name, models: Ok(models), }) } fn err(name: &'static str, msg: &'static str) -> Arc { Arc::new(Self { name, models: Err(anyhow::anyhow!(msg)), }) } } #[async_trait] impl Provider for ModelProvider { fn name(&self) -> &str { self.name } async fn list_models(&self) -> anyhow::Result> { match &self.models { Ok(v) => Ok(v.clone()), Err(e) => Err(anyhow::anyhow!("{e:#}")), } } async fn complete( &self, _request: CompletionRequest, _cancel: CancellationToken, ) -> anyhow::Result>> { unimplemented!() } } fn providers() -> Vec> { vec![ Arc::new(StubProvider("helexa")), Arc::new(StubProvider("openrouter")), ] } #[test] fn bare_model_routes_to_default() { let (p, m) = resolve_provider(&providers(), "helexa", "helexa/large").unwrap(); assert_eq!(p.name(), "helexa"); assert_eq!(m, "helexa/large"); } #[test] fn prefixed_model_routes_by_endpoint() { let (p, m) = resolve_provider(&providers(), "helexa", "openrouter:anthropic/claude-opus-4").unwrap(); assert_eq!(p.name(), "openrouter"); assert_eq!(m, "anthropic/claude-opus-4"); } #[test] fn unknown_endpoint_errors() { // `Arc` doesn't impl Debug, which rules out // `.unwrap_err()` (it requires T: Debug). Pattern-match instead. match resolve_provider(&providers(), "helexa", "ghost:gpt-9") { Ok(_) => panic!("expected error for unknown endpoint"), Err(e) => assert!(format!("{e}").contains("ghost")), } } // ── map_finish_reason ─────────────────────────────────────────── // ── prompt_budget ─────────────────────────────────────────────── #[test] fn prompt_budget_reserves_response_and_safety() { // 32K window, 8K response cap → 32768 - 8192 - 512 = 24064. assert_eq!(prompt_budget(32_768, Some(8_192)), 24_064); } #[test] fn prompt_budget_uses_default_when_max_tokens_unset() { // Default response cap = 2048; safety = 512. assert_eq!(prompt_budget(32_768, None), 32_768 - 2_048 - 512); } #[test] fn prompt_budget_saturates_when_window_too_small() { // Pathological config: window smaller than response + safety. // Don't underflow — return zero so compaction tries hardest // and upstream surfaces the inevitable error. assert_eq!(prompt_budget(1_000, Some(8_192)), 0); } // ── aggregate_models ──────────────────────────────────────────── #[tokio::test] async fn aggregate_models_single_endpoint_has_bare_ids() { let providers = vec![ModelProvider::ok( "helexa", &["helexa/large", "helexa/small"], )]; let models = aggregate_models(&providers).await; let ids: Vec<&str> = models.iter().map(|m| m.model_id.0.as_ref()).collect(); assert_eq!(ids, vec!["helexa/large", "helexa/small"]); } #[tokio::test] async fn aggregate_models_multi_endpoint_prefixes_every_id() { let providers = vec![ ModelProvider::ok("helexa", &["helexa/large"]), ModelProvider::ok("openrouter", &["anthropic/claude-opus-4"]), ]; let models = aggregate_models(&providers).await; let ids: Vec<&str> = models.iter().map(|m| m.model_id.0.as_ref()).collect(); assert_eq!( ids, vec!["helexa:helexa/large", "openrouter:anthropic/claude-opus-4"] ); } #[tokio::test] async fn aggregate_models_skips_failing_endpoint() { let providers = vec![ ModelProvider::err("flaky", "boom"), ModelProvider::ok("openrouter", &["gpt-9"]), ]; let models = aggregate_models(&providers).await; let ids: Vec<&str> = models.iter().map(|m| m.model_id.0.as_ref()).collect(); // Multi-endpoint case → prefix survives even when one // endpoint dropped out. flaky's models are absent, not // null-filled. assert_eq!(ids, vec!["openrouter:gpt-9"]); } #[test] fn maps_known_finish_reasons() { assert!(matches!( map_finish_reason(Some("length")), StopReason::MaxTokens )); assert!(matches!( map_finish_reason(Some("refusal")), StopReason::Refusal )); assert!(matches!( map_finish_reason(Some("stop")), StopReason::EndTurn )); assert!(matches!( map_finish_reason(Some("tool_calls")), StopReason::EndTurn )); assert!(matches!(map_finish_reason(None), StopReason::EndTurn)); } }