rob thijssen
df0abfe4d4
All checks were successful
build-prerelease / Resolve version stamps (push) Successful in 34s
CI / Format (push) Successful in 37s
CI / Clippy (push) Successful in 2m33s
CI / Test (push) Successful in 5m4s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
build-prerelease / Build neuron-blackwell (push) Successful in 6m2s
build-prerelease / Build neuron-ampere (push) Successful in 7m49s
build-prerelease / Build neuron-ada (push) Successful in 5m27s
build-prerelease / Build cortex binary (push) Successful in 4m16s
build-prerelease / Package cortex RPM (push) Successful in 1m19s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 3m2s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 3m10s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 3m47s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 1m2s
Stage 5. Zed clipboard/DnD images get forwarded as OpenAI
content-array messages on user turns.
- New MessageContent::MultiPart variant + MessagePart (Text|Image)
+ ImageData struct (mime_type, base64 data, optional uri).
- flatten_prompt now produces structured content: collapses to
Text when every block is text (some upstreams treat array-form
as vision-only and refuse on text-only models), otherwise
produces MultiPart preserving block order.
- OpenAI encoder emits `[{type:"text",text:…}, {type:"image_url",
image_url:{url:"data:{mime};base64,{data}"}}]` for MultiPart user
messages. Data URIs are used over remote `uri` because they
round-trip through every upstream we care about.
- prompt_capabilities.image = true at initialize so Zed actually
sends image blocks.
- compaction estimates ~512 tokens per image (the middle of the
Qwen3-VL / OpenAI detail range) so the budget tracker doesn't
pretend images are free.
- session/load replays image-bearing user turns by surfacing the
text parts verbatim and rendering each image as a "[image: {mime}
({n} bytes)]" placeholder chunk — Zed can show the prior text
context even though re-uploading the bytes through ACP isn't
meaningful for resume.
- 4 new tests: flatten produces MultiPart in block order, image-only
prompts still flatten to MultiPart, encoder emits the correct
array shape, text-only encoding stays as the string form.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
1821 lines
72 KiB
Rust
1821 lines
72 KiB
Rust
//! 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<AgentInner>` so handlers can clone
|
|
/// it cheaply into every closure.
|
|
pub struct Agent {
|
|
inner: Arc<AgentInner>,
|
|
}
|
|
|
|
struct AgentInner {
|
|
/// Every successfully-built provider, indexed positionally. We look
|
|
/// providers up by name (`endpoint:` prefix) rather than by index.
|
|
providers: Vec<Arc<dyn Provider>>,
|
|
/// 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<String>,
|
|
/// 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<String, u64>,
|
|
/// 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<String, usize>,
|
|
/// 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<AcpModelInfo>,
|
|
sessions: SessionStore,
|
|
system_prompt_path: Option<PathBuf>,
|
|
/// 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<Arc<dyn Provider>>) -> anyhow::Result<Self> {
|
|
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<Client>| {
|
|
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<Client>| {
|
|
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<Client>| {
|
|
handle_cancel(&inner, notif).await;
|
|
Ok(())
|
|
}
|
|
},
|
|
agent_client_protocol::on_receive_notification!(),
|
|
)
|
|
.on_receive_dispatch(
|
|
async move |message: Dispatch, cx: ConnectionTo<Client>| {
|
|
// `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<NewSessionResponse> {
|
|
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<Message>)> {
|
|
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<Client>, 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::<serde_json::Value>(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::<serde_json::Value>(&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<ListSessionsResponse> {
|
|
let sessions = store::list(req.cwd.as_deref())?;
|
|
let infos: Vec<SessionInfo> = 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<String> {
|
|
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<dyn Provider>]) -> Vec<AcpModelInfo> {
|
|
let multi_endpoint = providers.len() > 1;
|
|
let mut out: Vec<AcpModelInfo> = 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<SessionModelState> {
|
|
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<AgentInner>,
|
|
cx: ConnectionTo<Client>,
|
|
req: PromptRequest,
|
|
responder: agent_client_protocol::Responder<PromptResponse>,
|
|
) -> 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<AgentInner>,
|
|
cx: ConnectionTo<Client>,
|
|
req: PromptRequest,
|
|
responder: agent_client_protocol::Responder<PromptResponse>,
|
|
) -> 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<Message> = 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<Message> = Vec::new();
|
|
// Monotonic counter for synthetic ids assigned to unparseable
|
|
// <tool_call> 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<String> = 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<usize, ToolCallBucket> = BTreeMap::new();
|
|
// <tool_call> 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<String> = 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<ToolCall> = 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<tokio::sync::Mutex<SessionState>>,
|
|
session_id: &SessionId,
|
|
new_turns: Vec<Message>,
|
|
) {
|
|
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<Client>, 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 `<tool_call>` block.
|
|
/// The format mirrors successful calls (`call_<n>`) 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<Client>,
|
|
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\": \"<function>\", \"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
|
|
/// `<tool_call>`. 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 `<tool_call>` 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: "<invalid>".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 <tool_call> body was not valid JSON. Body was:\n{raw}\n\n\
|
|
Retry with the schema: {{\"name\": \"<function>\", \"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<u64>) -> 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<MessagePart> = Vec::new();
|
|
let mut text_buf = String::new();
|
|
let flush_text = |buf: &mut String, parts: &mut Vec<MessagePart>| {
|
|
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::<Vec<_>>()
|
|
.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<dyn Provider>],
|
|
default_endpoint: &str,
|
|
model_id: &str,
|
|
) -> anyhow::Result<(Arc<dyn Provider>, 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<Vec<crate::provider::ModelInfo>> {
|
|
unimplemented!()
|
|
}
|
|
async fn complete(
|
|
&self,
|
|
_request: CompletionRequest,
|
|
_cancel: CancellationToken,
|
|
) -> anyhow::Result<BoxStream<'static, anyhow::Result<CompletionEvent>>> {
|
|
unimplemented!()
|
|
}
|
|
}
|
|
|
|
/// Provider stub whose `list_models` returns canned results.
|
|
/// Used by the `aggregate_models` tests.
|
|
struct ModelProvider {
|
|
name: &'static str,
|
|
models: anyhow::Result<Vec<crate::provider::ModelInfo>>,
|
|
}
|
|
|
|
impl ModelProvider {
|
|
fn ok(name: &'static str, ids: &[&str]) -> Arc<dyn Provider> {
|
|
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<dyn Provider> {
|
|
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<Vec<crate::provider::ModelInfo>> {
|
|
match &self.models {
|
|
Ok(v) => Ok(v.clone()),
|
|
Err(e) => Err(anyhow::anyhow!("{e:#}")),
|
|
}
|
|
}
|
|
async fn complete(
|
|
&self,
|
|
_request: CompletionRequest,
|
|
_cancel: CancellationToken,
|
|
) -> anyhow::Result<BoxStream<'static, anyhow::Result<CompletionEvent>>> {
|
|
unimplemented!()
|
|
}
|
|
}
|
|
|
|
fn providers() -> Vec<Arc<dyn Provider>> {
|
|
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<dyn Provider>` 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));
|
|
}
|
|
}
|