Compare commits
41 Commits
249b2e5c98
...
phase-2-pr
| Author | SHA1 | Date | |
|---|---|---|---|
|
61adff347a
|
|||
|
0af8c8d6e7
|
|||
|
435fd10902
|
|||
|
cb303832bc
|
|||
|
44008358c5
|
|||
|
2f387f33f8
|
|||
|
fc9a8c42a3
|
|||
|
7733eecba5
|
|||
|
fdc0adb738
|
|||
|
8fa1d1962e
|
|||
|
cad7552104
|
|||
|
1818dfb337
|
|||
|
5ed1140c97
|
|||
|
957f704efa
|
|||
|
1859777332
|
|||
|
6927286cab
|
|||
|
302ccfb982
|
|||
|
df0abfe4d4
|
|||
|
b9016571f6
|
|||
|
adbc52bfcd
|
|||
|
537a0fe7f2
|
|||
|
cbadfcf112
|
|||
|
3ecbb21ece
|
|||
|
0d841a4981
|
|||
|
0bbb9b752d
|
|||
|
5aac1ffc59
|
|||
|
ec2b6450b2
|
|||
|
a494c8d43c
|
|||
|
abbedf8d8a
|
|||
|
6cc14e925c
|
|||
|
1c16732668
|
|||
|
5a0861d639
|
|||
|
33652ac651
|
|||
|
c297a54074
|
|||
|
0121a1930f
|
|||
|
13f4c36aeb
|
|||
|
4a51a54554
|
|||
|
0609f1ac5d
|
|||
|
96fc379893
|
|||
|
e267f583e1
|
|||
|
e23d5011d0
|
@@ -41,6 +41,7 @@ concurrency:
|
||||
|
||||
env:
|
||||
CARGO_INCREMENTAL: "0"
|
||||
CARGO_TERM_COLOR: "always"
|
||||
|
||||
jobs:
|
||||
prepare:
|
||||
|
||||
@@ -92,10 +92,68 @@ jobs:
|
||||
exit 1
|
||||
- run: sccache --show-stats
|
||||
|
||||
# Type-check the CUDA-only code path. Borrow-check-only — we
|
||||
# never run the tests here (the runner has no GPU). This catches
|
||||
# the category of bug where a refactor compiles fine under the
|
||||
# default feature set (which is what the `clippy` and `test` jobs
|
||||
# exercise) but fails inside a `#[cfg(feature = "cuda")]` block.
|
||||
# `runs-on: cuda-13.0` selects the runner that ships nvcc /
|
||||
# cudarc's build prerequisites. The generic `rust` and `rpm`
|
||||
# runners don't have them (the previous label `rpm` was tried
|
||||
# first and tripped cudarc's `nvcc --version` build script —
|
||||
# see commit history).
|
||||
cuda-check:
|
||||
name: CUDA type-check
|
||||
runs-on: cuda-13.0
|
||||
# The workflow-level env sets `RUSTC_WRAPPER: sccache` for the
|
||||
# `rust` runner (where fmt/clippy/test live and sccache is
|
||||
# installed). The `cuda-13.0` runner doesn't have sccache on
|
||||
# PATH, so inheriting the wrapper makes cargo bail with
|
||||
# `could not execute process `sccache rustc -vV` (never executed)`
|
||||
# before borrow-check even starts. Clear it locally. Also clear
|
||||
# SCCACHE_* so cargo doesn't try to contact the cache (the
|
||||
# remote auth headers come from secrets that aren't present on
|
||||
# this runner either). Lose the cache, keep the gate.
|
||||
env:
|
||||
RUSTC_WRAPPER: ""
|
||||
SCCACHE_BUCKET: ""
|
||||
SCCACHE_ENDPOINT: ""
|
||||
SCCACHE_REGION: ""
|
||||
SCCACHE_S3_USE_SSL: ""
|
||||
AWS_ACCESS_KEY_ID: ""
|
||||
AWS_SECRET_ACCESS_KEY: ""
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- name: cargo check --features cuda (with retry)
|
||||
run: |
|
||||
# act launches the step shell without /etc/profile, so the
|
||||
# gitea_runner user's inherited PATH lacks /usr/local/cuda-13.0/bin.
|
||||
# cudarc's build.rs:157 shells out to `nvcc --version` (because
|
||||
# the neuron crate enables cuda-version-from-build-system) and
|
||||
# panics with ENOENT if nvcc isn't resolvable. build-prerelease.yml
|
||||
# does the same export — keep them in sync.
|
||||
export PATH="/usr/local/cuda-13.0/bin:${PATH}"
|
||||
export LD_LIBRARY_PATH="/usr/local/cuda-13.0/targets/x86_64-linux/lib:/usr/local/cuda-13.0/lib64:${LD_LIBRARY_PATH:-}"
|
||||
export LIBRARY_PATH="/usr/local/cuda-13.0/targets/x86_64-linux/lib:/usr/local/cuda-13.0/lib64:${LIBRARY_PATH:-}"
|
||||
for attempt in 1 2 3; do
|
||||
echo "::group::cuda-check attempt ${attempt}"
|
||||
if cargo check -p neuron --features cuda --all-targets; then
|
||||
echo "::endgroup::"
|
||||
exit 0
|
||||
fi
|
||||
echo "::endgroup::"
|
||||
echo "cuda-check failed on attempt ${attempt}"
|
||||
if [ "${attempt}" -lt 3 ]; then
|
||||
sleep 5
|
||||
fi
|
||||
done
|
||||
echo "cuda-check failed after 3 attempts"
|
||||
exit 1
|
||||
|
||||
srpm-cortex:
|
||||
name: Build cortex SRPM
|
||||
runs-on: rpm
|
||||
needs: [fmt, clippy, test]
|
||||
needs: [fmt, clippy, test, cuda-check]
|
||||
if: startsWith(github.ref, 'refs/tags/v')
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
@@ -155,7 +213,7 @@ jobs:
|
||||
srpm-neuron:
|
||||
name: Build neuron SRPM
|
||||
runs-on: rpm
|
||||
needs: [fmt, clippy, test]
|
||||
needs: [fmt, clippy, test, cuda-check]
|
||||
if: startsWith(github.ref, 'refs/tags/v')
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
812
Cargo.lock
generated
812
Cargo.lock
generated
File diff suppressed because it is too large
Load Diff
@@ -5,6 +5,7 @@ members = [
|
||||
"crates/cortex-gateway",
|
||||
"crates/cortex-cli",
|
||||
"crates/neuron",
|
||||
"crates/helexa-acp",
|
||||
]
|
||||
|
||||
[workspace.package]
|
||||
|
||||
@@ -6,4 +6,5 @@ pub mod harness;
|
||||
pub mod metrics;
|
||||
pub mod node;
|
||||
pub mod openai;
|
||||
pub mod responses;
|
||||
pub mod translate;
|
||||
|
||||
346
crates/cortex-core/src/responses.rs
Normal file
346
crates/cortex-core/src/responses.rs
Normal file
@@ -0,0 +1,346 @@
|
||||
//! OpenAI Responses API (`POST /v1/responses`) envelope types.
|
||||
//!
|
||||
//! This is OpenAI's newer chat surface, distinct from
|
||||
//! `/v1/chat/completions` in three ways that matter for us:
|
||||
//!
|
||||
//! 1. **Input shape**. Instead of a `messages` array, the request
|
||||
//! carries `input` — either a plain string (single user turn)
|
||||
//! or an array of typed items (messages, function calls,
|
||||
//! function-call outputs, reasoning blocks, …).
|
||||
//! 2. **Output shape**. The response carries a single `output`
|
||||
//! array of items, each typed. We always emit one
|
||||
//! `OutputItem::Message` containing the assistant's reply (plus,
|
||||
//! when we get there, separate `function_call` items).
|
||||
//! 3. **Streaming events**. Where chat completions stream
|
||||
//! structurally-identical `chat.completion.chunk` frames over
|
||||
//! `data:` lines, Responses streams *named* events
|
||||
//! (`response.created`, `response.output_text.delta`,
|
||||
//! `response.completed`, …) over `event:` + `data:` SSE pairs.
|
||||
//! The wire projector in `neuron::wire::openai_responses` builds
|
||||
//! these from the same [`crate::openai`]-shaped
|
||||
//! `InferenceEvent` stream the chat projector consumes.
|
||||
//!
|
||||
//! Scope cuts for this first cut:
|
||||
//!
|
||||
//! - **`previous_response_id` is rejected at parse time**. Stateful
|
||||
//! chained conversations need a persistence layer we don't have.
|
||||
//! - **Reasoning items are accepted-and-ignored** (no Qwen3
|
||||
//! `<think>` routing yet). Audio and embedded resources are
|
||||
//! rejected as unsupported.
|
||||
//! - **Tool calls** (function_call / function_call_output) are
|
||||
//! carried as round-trip types but the candle harness doesn't
|
||||
//! emit them yet — wired so the surface is in place for the
|
||||
//! day we add proper tool-call extraction.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
// ── Request ──────────────────────────────────────────────────────────
|
||||
|
||||
/// Body of a `POST /v1/responses` request.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ResponsesRequest {
|
||||
pub model: String,
|
||||
pub input: ResponsesInput,
|
||||
/// System-prompt-style instructions. The Responses API
|
||||
/// separates these from input so a caller doesn't have to
|
||||
/// build a `system` message item by hand.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub instructions: Option<String>,
|
||||
#[serde(default)]
|
||||
pub stream: bool,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub max_output_tokens: Option<u64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub temperature: Option<f64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub top_p: Option<f64>,
|
||||
/// Chained-conversation identifier. We don't store responses
|
||||
/// server-side yet; if this is `Some`, the handler returns 400.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub previous_response_id: Option<String>,
|
||||
/// Catch-all for anything we don't model yet (tools, tool_choice,
|
||||
/// reasoning, response_format, …). Lets a client send a
|
||||
/// forward-compatible request without our parser rejecting it.
|
||||
#[serde(flatten)]
|
||||
pub extra: Value,
|
||||
}
|
||||
|
||||
/// `input` is either a single string or an array of typed items.
|
||||
/// `#[serde(untagged)]` so the wire shape `"input": "hi"` and
|
||||
/// `"input": [{...}]` both deserialize.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum ResponsesInput {
|
||||
Text(String),
|
||||
Items(Vec<ResponsesInputItem>),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ResponsesInputItem {
|
||||
/// A user / assistant / system turn.
|
||||
Message {
|
||||
role: String,
|
||||
content: ResponsesMessageContent,
|
||||
},
|
||||
/// Assistant emitted a tool call. Round-trip only — neuron
|
||||
/// doesn't synthesise these yet.
|
||||
FunctionCall {
|
||||
call_id: String,
|
||||
name: String,
|
||||
arguments: String,
|
||||
},
|
||||
/// User is feeding a tool result back into the model.
|
||||
FunctionCallOutput { call_id: String, output: String },
|
||||
/// Reasoning items emitted by o-series models. Accepted but
|
||||
/// not forwarded to the model — neuron's candle path doesn't
|
||||
/// surface reasoning separately yet.
|
||||
Reasoning {
|
||||
#[serde(default)]
|
||||
content: Vec<Value>,
|
||||
},
|
||||
}
|
||||
|
||||
/// Inside a `Message` item, content is either a plain string or an
|
||||
/// array of typed parts. Mirrors the chat-completions Parts shape.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum ResponsesMessageContent {
|
||||
Text(String),
|
||||
Parts(Vec<ResponsesContentPart>),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ResponsesContentPart {
|
||||
/// Plain text inside a user / system turn.
|
||||
InputText { text: String },
|
||||
/// An image. `image_url` is either a remote URL or a
|
||||
/// `data:image/png;base64,…` URI; the request translator just
|
||||
/// forwards the string.
|
||||
InputImage {
|
||||
image_url: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
detail: Option<String>,
|
||||
},
|
||||
/// Returned text inside an assistant turn — only relevant when
|
||||
/// the caller is feeding an assistant turn back in to continue
|
||||
/// a conversation manually (no `previous_response_id`).
|
||||
OutputText {
|
||||
text: String,
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
annotations: Vec<Value>,
|
||||
},
|
||||
}
|
||||
|
||||
// ── Response (non-streaming) ─────────────────────────────────────────
|
||||
|
||||
/// Body of a `POST /v1/responses` response.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ResponsesResponse {
|
||||
pub id: String,
|
||||
/// Always `"response"`.
|
||||
pub object: String,
|
||||
pub created_at: u64,
|
||||
/// `"completed"`, `"incomplete"`, or — for the initial event of
|
||||
/// a streaming response — `"in_progress"`.
|
||||
pub status: String,
|
||||
pub model: String,
|
||||
pub output: Vec<ResponsesOutputItem>,
|
||||
/// Populated on completion; `None` while streaming.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub usage: Option<ResponsesUsage>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ResponsesOutputItem {
|
||||
Message {
|
||||
id: String,
|
||||
/// Always `"assistant"` for model output.
|
||||
role: String,
|
||||
/// Output content parts. We always emit a single
|
||||
/// `OutputText` today; multi-part output would land here
|
||||
/// once we have e.g. image generation.
|
||||
content: Vec<ResponsesOutputContent>,
|
||||
/// Item-level status. `"in_progress"` while streaming the
|
||||
/// content parts, `"completed"` when done.
|
||||
#[serde(default = "default_item_status")]
|
||||
status: String,
|
||||
},
|
||||
/// Reserved for the day tool-call extraction lands. The wire
|
||||
/// shape mirrors `ResponsesInputItem::FunctionCall`.
|
||||
FunctionCall {
|
||||
id: String,
|
||||
call_id: String,
|
||||
name: String,
|
||||
arguments: String,
|
||||
#[serde(default = "default_item_status")]
|
||||
status: String,
|
||||
},
|
||||
}
|
||||
|
||||
fn default_item_status() -> String {
|
||||
"completed".into()
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ResponsesOutputContent {
|
||||
OutputText {
|
||||
text: String,
|
||||
/// Citations / inline annotations. Empty today; reserved
|
||||
/// for the day we wire in web search / file search.
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
annotations: Vec<Value>,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ResponsesUsage {
|
||||
pub input_tokens: u64,
|
||||
pub output_tokens: u64,
|
||||
pub total_tokens: u64,
|
||||
}
|
||||
|
||||
// ── Streaming event names ────────────────────────────────────────────
|
||||
|
||||
/// Event names the SSE projector emits, hoisted as constants so
|
||||
/// the projector and the wire shape stay in sync without
|
||||
/// string-typos. The strings are dictated by OpenAI's published
|
||||
/// Responses API.
|
||||
pub mod events {
|
||||
pub const CREATED: &str = "response.created";
|
||||
/// Fired between `response.created` and the first output-item
|
||||
/// event. Marks "request validated, model is generating" —
|
||||
/// some clients use it to differentiate the "warming up" state
|
||||
/// from "streaming tokens" in their UI.
|
||||
pub const IN_PROGRESS: &str = "response.in_progress";
|
||||
pub const OUTPUT_ITEM_ADDED: &str = "response.output_item.added";
|
||||
pub const CONTENT_PART_ADDED: &str = "response.content_part.added";
|
||||
pub const OUTPUT_TEXT_DELTA: &str = "response.output_text.delta";
|
||||
pub const OUTPUT_TEXT_DONE: &str = "response.output_text.done";
|
||||
pub const CONTENT_PART_DONE: &str = "response.content_part.done";
|
||||
pub const OUTPUT_ITEM_DONE: &str = "response.output_item.done";
|
||||
pub const COMPLETED: &str = "response.completed";
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn deserialises_input_string_form() {
|
||||
let raw = r#"{"model": "m", "input": "hello"}"#;
|
||||
let req: ResponsesRequest = serde_json::from_str(raw).unwrap();
|
||||
match req.input {
|
||||
ResponsesInput::Text(s) => assert_eq!(s, "hello"),
|
||||
other => panic!("expected Text, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deserialises_input_items_form() {
|
||||
let raw = r#"{
|
||||
"model": "m",
|
||||
"input": [
|
||||
{"type": "message", "role": "user", "content": "hi"}
|
||||
]
|
||||
}"#;
|
||||
let req: ResponsesRequest = serde_json::from_str(raw).unwrap();
|
||||
match req.input {
|
||||
ResponsesInput::Items(items) => {
|
||||
assert_eq!(items.len(), 1);
|
||||
match &items[0] {
|
||||
ResponsesInputItem::Message { role, content } => {
|
||||
assert_eq!(role, "user");
|
||||
match content {
|
||||
ResponsesMessageContent::Text(t) => assert_eq!(t, "hi"),
|
||||
other => panic!("expected Text content, got {other:?}"),
|
||||
}
|
||||
}
|
||||
other => panic!("expected Message item, got {other:?}"),
|
||||
}
|
||||
}
|
||||
other => panic!("expected Items, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deserialises_input_with_image() {
|
||||
let raw = r#"{
|
||||
"model": "m",
|
||||
"input": [
|
||||
{"type": "message", "role": "user", "content": [
|
||||
{"type": "input_text", "text": "what is this"},
|
||||
{"type": "input_image", "image_url": "data:image/png;base64,AAA="}
|
||||
]}
|
||||
]
|
||||
}"#;
|
||||
let req: ResponsesRequest = serde_json::from_str(raw).unwrap();
|
||||
let items = match req.input {
|
||||
ResponsesInput::Items(i) => i,
|
||||
other => panic!("expected Items, got {other:?}"),
|
||||
};
|
||||
let parts = match &items[0] {
|
||||
ResponsesInputItem::Message {
|
||||
content: ResponsesMessageContent::Parts(p),
|
||||
..
|
||||
} => p,
|
||||
other => panic!("expected Parts, got {other:?}"),
|
||||
};
|
||||
assert_eq!(parts.len(), 2);
|
||||
assert!(matches!(
|
||||
&parts[0],
|
||||
ResponsesContentPart::InputText { text } if text == "what is this"
|
||||
));
|
||||
assert!(matches!(
|
||||
&parts[1],
|
||||
ResponsesContentPart::InputImage { image_url, .. }
|
||||
if image_url == "data:image/png;base64,AAA="
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_fields_round_trip_via_extra() {
|
||||
let raw = r#"{
|
||||
"model": "m",
|
||||
"input": "hi",
|
||||
"tools": [{"type": "web_search"}],
|
||||
"reasoning": {"effort": "medium"}
|
||||
}"#;
|
||||
let req: ResponsesRequest = serde_json::from_str(raw).unwrap();
|
||||
assert!(req.extra.get("tools").is_some());
|
||||
assert!(req.extra.get("reasoning").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_round_trips_through_serde() {
|
||||
let r = ResponsesResponse {
|
||||
id: "resp_1".into(),
|
||||
object: "response".into(),
|
||||
created_at: 1700,
|
||||
status: "completed".into(),
|
||||
model: "m".into(),
|
||||
output: vec![ResponsesOutputItem::Message {
|
||||
id: "msg_1".into(),
|
||||
role: "assistant".into(),
|
||||
content: vec![ResponsesOutputContent::OutputText {
|
||||
text: "hi there".into(),
|
||||
annotations: vec![],
|
||||
}],
|
||||
status: "completed".into(),
|
||||
}],
|
||||
usage: Some(ResponsesUsage {
|
||||
input_tokens: 5,
|
||||
output_tokens: 3,
|
||||
total_tokens: 8,
|
||||
}),
|
||||
};
|
||||
let json = serde_json::to_string(&r).unwrap();
|
||||
let parsed: ResponsesResponse = serde_json::from_str(&json).unwrap();
|
||||
assert_eq!(parsed.id, "resp_1");
|
||||
assert_eq!(parsed.output.len(), 1);
|
||||
}
|
||||
}
|
||||
@@ -20,6 +20,7 @@ pub fn api_routes() -> Router<Arc<CortexState>> {
|
||||
Router::new()
|
||||
.route("/v1/chat/completions", post(chat_completions))
|
||||
.route("/v1/completions", post(completions))
|
||||
.route("/v1/responses", post(responses))
|
||||
.route("/v1/models", get(list_models))
|
||||
.route("/v1/messages", post(anthropic_messages))
|
||||
.route("/health", get(health))
|
||||
@@ -74,6 +75,58 @@ async fn chat_completions(
|
||||
.await
|
||||
}
|
||||
|
||||
/// `POST /v1/responses` — proxy to the appropriate backend node.
|
||||
///
|
||||
/// Same routing shape as [`chat_completions`]: extract `model` from
|
||||
/// the body, resolve to a node, forward verbatim. No translation —
|
||||
/// neuron speaks the Responses API natively (see
|
||||
/// `crates/neuron/src/wire/openai_responses.rs`), so the gateway is
|
||||
/// a pass-through. Streaming and non-streaming are handled
|
||||
/// identically; the upstream `Content-Type` (text/event-stream vs.
|
||||
/// application/json) propagates through the proxy.
|
||||
async fn responses(
|
||||
State(fleet): State<Arc<CortexState>>,
|
||||
headers: HeaderMap,
|
||||
body: Bytes,
|
||||
) -> Response {
|
||||
let model_id = match extract_model(&body) {
|
||||
Some(m) => m,
|
||||
None => {
|
||||
tracing::warn!(
|
||||
handler = "responses",
|
||||
"rejected: missing 'model' field in request body"
|
||||
);
|
||||
return error_response(400, "missing 'model' field in request body");
|
||||
}
|
||||
};
|
||||
|
||||
let route = match router::resolve(&fleet, &model_id).await {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
handler = "responses",
|
||||
model = %model_id,
|
||||
error = %e,
|
||||
"route resolve failed"
|
||||
);
|
||||
return error_response(404, &e.to_string());
|
||||
}
|
||||
};
|
||||
|
||||
touch_model(&fleet, &route.node_name, &route.resolved_model_id).await;
|
||||
|
||||
let body = rewrite_model_in_body(body, &route.resolved_model_id);
|
||||
proxy_with_metrics(
|
||||
&fleet,
|
||||
&route,
|
||||
"/v1/responses",
|
||||
headers,
|
||||
body,
|
||||
&route.resolved_model_id,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// `POST /v1/completions` — proxy completions endpoint.
|
||||
async fn completions(
|
||||
State(fleet): State<Arc<CortexState>>,
|
||||
|
||||
@@ -44,6 +44,7 @@ pub async fn spawn_mock_neuron() -> String {
|
||||
post(|Json(_body): Json<Value>| async { Json(json!({"status": "unloaded"})) }),
|
||||
)
|
||||
.route("/v1/chat/completions", post(mock_chat_completions))
|
||||
.route("/v1/responses", post(mock_responses))
|
||||
.route("/v1/models", get(mock_v1_models));
|
||||
|
||||
tokio::spawn(async move {
|
||||
@@ -93,6 +94,39 @@ async fn mock_chat_completions(Json(body): Json<Value>) -> Json<Value> {
|
||||
}))
|
||||
}
|
||||
|
||||
async fn mock_responses(Json(body): Json<Value>) -> Json<Value> {
|
||||
let model = body
|
||||
.get("model")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("unknown");
|
||||
// Echo the model field back and synthesise a tiny ResponsesResponse.
|
||||
// Mirrors the shape neuron's /v1/responses handler emits so the
|
||||
// gateway test only needs to assert the proxy round-tripped it.
|
||||
Json(json!({
|
||||
"id": "resp-test-001",
|
||||
"object": "response",
|
||||
"created_at": 1700000000_u64,
|
||||
"status": "completed",
|
||||
"model": model,
|
||||
"output": [{
|
||||
"type": "message",
|
||||
"id": "msg-test-001",
|
||||
"role": "assistant",
|
||||
"content": [{
|
||||
"type": "output_text",
|
||||
"text": "Hello from mock backend",
|
||||
"annotations": []
|
||||
}],
|
||||
"status": "completed"
|
||||
}],
|
||||
"usage": {
|
||||
"input_tokens": 5,
|
||||
"output_tokens": 5,
|
||||
"total_tokens": 10
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
/// Spawns a mock neuron that returns SSE streaming responses for chat completions.
|
||||
pub async fn spawn_streaming_mock_neuron(chunk_count: usize, chunk_delay: Duration) -> String {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
|
||||
91
crates/cortex-gateway/tests/responses.rs
Normal file
91
crates/cortex-gateway/tests/responses.rs
Normal file
@@ -0,0 +1,91 @@
|
||||
//! Integration tests for the `/v1/responses` proxy route.
|
||||
//!
|
||||
//! The gateway forwards the request body to whichever neuron has the
|
||||
//! model loaded. These tests exercise the routing decision (200 on a
|
||||
//! known model, 404 on an unknown model, 400 on a missing model
|
||||
//! field) and confirm the response body round-trips verbatim.
|
||||
|
||||
mod common;
|
||||
|
||||
use serde_json::json;
|
||||
|
||||
/// Happy path: gateway routes a `/v1/responses` request to the neuron
|
||||
/// that has the model loaded, and the neuron's response body
|
||||
/// arrives at the client unchanged.
|
||||
#[tokio::test]
|
||||
async fn test_responses_proxy() {
|
||||
let mock_url = common::spawn_mock_neuron().await;
|
||||
let gw_url = common::spawn_gateway(&mock_url).await;
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let resp = client
|
||||
.post(format!("{gw_url}/v1/responses"))
|
||||
.header("content-type", "application/json")
|
||||
.json(&json!({
|
||||
"model": "test-model",
|
||||
"input": "Hi"
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(resp.status(), 200);
|
||||
|
||||
let body: serde_json::Value = resp.json().await.expect("valid JSON response");
|
||||
assert_eq!(body["id"], "resp-test-001");
|
||||
assert_eq!(body["object"], "response");
|
||||
assert_eq!(body["model"], "test-model");
|
||||
assert_eq!(body["status"], "completed");
|
||||
assert_eq!(
|
||||
body["output"][0]["content"][0]["text"],
|
||||
"Hello from mock backend"
|
||||
);
|
||||
// Usage shape is the Responses-specific (input/output_tokens),
|
||||
// not the chat-completions one (prompt/completion_tokens). Asserts
|
||||
// the proxy didn't accidentally route through the wrong handler.
|
||||
assert_eq!(body["usage"]["total_tokens"], 10);
|
||||
assert!(body["usage"].get("input_tokens").is_some());
|
||||
}
|
||||
|
||||
/// A request that targets a model not present in the catalogue gets
|
||||
/// 404 from the router. This matches the chat-completions handler's
|
||||
/// behaviour — same error path, same status code, so a client can
|
||||
/// share retry logic across the two routes.
|
||||
#[tokio::test]
|
||||
async fn test_responses_model_not_found() {
|
||||
let mock_url = common::spawn_mock_neuron().await;
|
||||
let gw_url = common::spawn_gateway(&mock_url).await;
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let resp = client
|
||||
.post(format!("{gw_url}/v1/responses"))
|
||||
.json(&json!({
|
||||
"model": "not-in-catalogue",
|
||||
"input": "Hi"
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), 404);
|
||||
}
|
||||
|
||||
/// A request body without a `model` field can't be routed; the
|
||||
/// gateway returns 400 before reaching a backend. Same as the
|
||||
/// chat-completions handler — extracted via the same `extract_model`
|
||||
/// helper.
|
||||
#[tokio::test]
|
||||
async fn test_responses_missing_model_field() {
|
||||
let mock_url = common::spawn_mock_neuron().await;
|
||||
let gw_url = common::spawn_gateway(&mock_url).await;
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let resp = client
|
||||
.post(format!("{gw_url}/v1/responses"))
|
||||
.json(&json!({
|
||||
"input": "Hi"
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), 400);
|
||||
}
|
||||
48
crates/helexa-acp/Cargo.toml
Normal file
48
crates/helexa-acp/Cargo.toml
Normal file
@@ -0,0 +1,48 @@
|
||||
[package]
|
||||
name = "helexa-acp"
|
||||
version = "0.1.16"
|
||||
edition = "2024"
|
||||
license = "Apache-2.0"
|
||||
repository = "https://git.lair.cafe/helexa/cortex"
|
||||
description = """
|
||||
Agent Client Protocol bridge for the helexa self-hosted LLM stack.
|
||||
Speaks ACP to ACP-compatible editor clients (Zed, etc.) and forwards
|
||||
the conversation to any OpenAI-compatible HTTP endpoint — defaulting
|
||||
to cortex (helexa's reverse-proxy / fleet gateway).
|
||||
"""
|
||||
|
||||
# This crate is intentionally self-contained — no dependencies on other
|
||||
# workspace crates (cortex-core, cortex-gateway, neuron). The goal is
|
||||
# a painless migration to a dedicated GitHub repo in the future if the
|
||||
# project grows beyond helexa's needs. All deps are crates.io.
|
||||
[dependencies]
|
||||
# `unstable_session_model` flips on the SessionModelState type and the
|
||||
# session/set_model RPC the model-picker dropdown in Zed needs. The
|
||||
# feature is upstream-marked unstable; we accept that risk because the
|
||||
# model picker is core UX and the alternative (rolling our own
|
||||
# extension method) drifts further from spec each time it moves.
|
||||
agent-client-protocol = { version = "0.12", features = ["unstable_session_model"] }
|
||||
tokio = { version = "1", features = ["rt-multi-thread", "macros", "sync", "io-util", "process", "signal"] }
|
||||
reqwest = { version = "0.12", features = ["json", "stream", "rustls-tls"], default-features = false }
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
toml = "0.8"
|
||||
tracing = "0.1"
|
||||
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
|
||||
anyhow = "1"
|
||||
thiserror = "2"
|
||||
async-trait = "0.1"
|
||||
futures = "0.3"
|
||||
tokio-stream = "0.1"
|
||||
tokio-util = { version = "0.7", features = ["rt"] }
|
||||
eventsource-stream = "0.2"
|
||||
async-stream = "0.3"
|
||||
url = { version = "2", features = ["serde"] }
|
||||
# Already transitively pulled via the ACP SDK; declared directly so we
|
||||
# can format ISO 8601 timestamps for `SessionInfo.updated_at` in the
|
||||
# session/list response.
|
||||
chrono = { version = "0.4", default-features = false, features = ["std"] }
|
||||
|
||||
[[bin]]
|
||||
name = "helexa-acp"
|
||||
path = "src/main.rs"
|
||||
546
crates/helexa-acp/README.md
Normal file
546
crates/helexa-acp/README.md
Normal file
@@ -0,0 +1,546 @@
|
||||
# helexa-acp
|
||||
|
||||
ACP (Agent Client Protocol) bridge for editors like
|
||||
[Zed](https://zed.dev). Lets you point your editor's agent panel at
|
||||
**any combination** of OpenAI-compatible, OpenAI Responses, and
|
||||
Anthropic Messages endpoints — public APIs, private LAN deployments,
|
||||
local Ollama / LM Studio — and switch between them per session via a
|
||||
model dropdown.
|
||||
|
||||
The "missing ACP binary" for users who don't want to be locked into
|
||||
one vendor's agent client.
|
||||
|
||||
```
|
||||
┌───────────────────────────────────┐
|
||||
│ Zed (or any ACP editor client) │
|
||||
└────────────┬──────────────────────┘
|
||||
│ stdio JSON-RPC (ACP)
|
||||
▼
|
||||
┌─────────────────┐
|
||||
│ helexa-acp │ ← one binary, multi-endpoint
|
||||
└─────┬───────────┘
|
||||
│ HTTP / SSE
|
||||
┌────────┼─────────────┬──────────────┬──────────────┐
|
||||
▼ ▼ ▼ ▼ ▼
|
||||
cortex/ OpenAI Anthropic OpenRouter LM Studio
|
||||
neuron Responses Messages
|
||||
(self- (gpt-5,…) (Claude)
|
||||
hosted)
|
||||
```
|
||||
|
||||
## What it does
|
||||
|
||||
- **Speaks ACP** over stdio to editor clients (Zed today; any future
|
||||
ACP client tomorrow).
|
||||
- **Multi-endpoint** — one config file lists every LLM endpoint
|
||||
you want available; pick one per session via the model dropdown
|
||||
(`endpoint:model` selector).
|
||||
- **Three wire formats**: `openai-chat` (the broadly compatible
|
||||
default), `openai-responses` (newer OpenAI surface), and
|
||||
`anthropic-messages` (Claude). Each is a separate provider impl
|
||||
in `src/provider/`; adding a fourth (Gemini, Ollama native, …) is
|
||||
one file plus a `WireApi` enum variant.
|
||||
- **Built-in tools**: `read_file`, `write_file`, `edit_file`,
|
||||
`list_dir`, `bash`. Permission-gated by default; the editor user
|
||||
approves writes/shell per-call.
|
||||
- **Three session modes**: Default (gated), Bypass Permissions
|
||||
(auto-allow), and Plan (write-only-to-plan-dir, no shell).
|
||||
- **Vision** — drag-drop images into the agent panel against any
|
||||
vision-capable model.
|
||||
- **Session resume** — multi-day conversations survive editor
|
||||
restarts via on-disk transcript persistence.
|
||||
- **Context compaction** — rolling history stays inside the model's
|
||||
context window automatically so long sessions on small-context
|
||||
local models don't fall over.
|
||||
|
||||
## Install
|
||||
|
||||
### From source
|
||||
|
||||
```sh
|
||||
git clone https://git.lair.cafe/helexa/cortex.git
|
||||
cd cortex
|
||||
cargo install --path crates/helexa-acp
|
||||
# Binary lands at ~/.cargo/bin/helexa-acp
|
||||
```
|
||||
|
||||
### Pre-built RPM (Fedora 43)
|
||||
|
||||
```sh
|
||||
dnf copr enable helexa/helexa
|
||||
dnf install helexa-acp
|
||||
```
|
||||
|
||||
The COPR project bundles helexa-acp alongside the cortex gateway
|
||||
and helexa-neuron flavours; install only the package(s) you need.
|
||||
|
||||
## Quick start
|
||||
|
||||
The fastest path: env-var single-endpoint config.
|
||||
|
||||
```sh
|
||||
export HELEXA_ACP_BASE_URL=http://hanzalova.internal:31313/v1
|
||||
export HELEXA_ACP_MODEL=Qwen/Qwen3.6-27B
|
||||
helexa-acp # speaks ACP over stdin/stdout; not interactive
|
||||
```
|
||||
|
||||
Then in Zed (`~/.config/zed/settings.json`):
|
||||
|
||||
```jsonc
|
||||
{
|
||||
"agent_servers": {
|
||||
"helexa": {
|
||||
"command": "helexa-acp",
|
||||
"args": []
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Restart Zed → open the agent panel → pick "helexa" → start
|
||||
chatting. Tool calls (file reads, writes, bash) prompt for
|
||||
permission per-call in Default mode.
|
||||
|
||||
That's the minimum. The full config story below is what unlocks
|
||||
the multi-endpoint dropdown.
|
||||
|
||||
## Multi-endpoint config
|
||||
|
||||
Copy `helexa-acp.example.toml` from this repo to
|
||||
`$XDG_CONFIG_HOME/helexa-acp/config.toml` (typically
|
||||
`~/.config/helexa-acp/config.toml`) and edit:
|
||||
|
||||
```toml
|
||||
default_endpoint = "helexa"
|
||||
|
||||
[[endpoints]]
|
||||
name = "helexa"
|
||||
base_url = "http://hanzalova.internal:31313/v1"
|
||||
wire_api = "openai-chat"
|
||||
default_model = "Qwen/Qwen3.6-27B"
|
||||
max_tokens = 8192
|
||||
context_window = 32768
|
||||
|
||||
[[endpoints]]
|
||||
name = "openrouter"
|
||||
base_url = "https://openrouter.ai/api/v1"
|
||||
wire_api = "openai-chat"
|
||||
api_key_env = "OPENROUTER_API_KEY"
|
||||
default_model = "anthropic/claude-opus-4"
|
||||
|
||||
[[endpoints]]
|
||||
name = "anthropic"
|
||||
base_url = "https://api.anthropic.com/v1"
|
||||
wire_api = "anthropic-messages"
|
||||
api_key_env = "ANTHROPIC_API_KEY"
|
||||
default_model = "claude-opus-4"
|
||||
```
|
||||
|
||||
Restart Zed. The model dropdown lists every model from every
|
||||
configured endpoint with the `endpoint:model` selector
|
||||
(`helexa:Qwen/Qwen3.6-27B`, `openrouter:anthropic/claude-opus-4`,
|
||||
…). Switch mid-session; the next prompt routes to the new endpoint.
|
||||
|
||||
When only one endpoint is configured the prefix is dropped (model
|
||||
ids appear bare).
|
||||
|
||||
### Selector syntax
|
||||
|
||||
The `model` field on every internal request is parsed as
|
||||
`<endpoint>:<model>`:
|
||||
|
||||
- `openrouter:gpt-4o` → routes to the `openrouter` endpoint,
|
||||
model `gpt-4o`.
|
||||
- `helexa/large` → no colon → falls through to whichever endpoint
|
||||
is named in `default_endpoint`, model `helexa/large`.
|
||||
- `:gpt-5` → leading colon → also falls through to default.
|
||||
|
||||
## Endpoint cookbook
|
||||
|
||||
Copy-pasteable blocks. Mix and match.
|
||||
|
||||
### cortex / neuron (self-hosted)
|
||||
|
||||
```toml
|
||||
[[endpoints]]
|
||||
name = "helexa"
|
||||
base_url = "http://hanzalova.internal:31313/v1"
|
||||
wire_api = "openai-chat"
|
||||
default_model = "Qwen/Qwen3.6-27B"
|
||||
max_tokens = 8192
|
||||
context_window = 32768
|
||||
```
|
||||
|
||||
Use `openai-responses` instead of `openai-chat` once cortex 0.1.16+
|
||||
is deployed and you want the Responses API surface (vision item
|
||||
shape, structured reasoning items, etc.).
|
||||
|
||||
### OpenAI directly
|
||||
|
||||
```toml
|
||||
[[endpoints]]
|
||||
name = "openai"
|
||||
base_url = "https://api.openai.com/v1"
|
||||
wire_api = "openai-responses"
|
||||
api_key_env = "OPENAI_API_KEY"
|
||||
default_model = "gpt-5"
|
||||
```
|
||||
|
||||
`openai-responses` is the right choice for current OpenAI models;
|
||||
`openai-chat` works against legacy GPT-3.5/4 deployments and
|
||||
anything labelled "chat completions".
|
||||
|
||||
### Anthropic directly
|
||||
|
||||
```toml
|
||||
[[endpoints]]
|
||||
name = "anthropic"
|
||||
base_url = "https://api.anthropic.com/v1"
|
||||
wire_api = "anthropic-messages"
|
||||
api_key_env = "ANTHROPIC_API_KEY"
|
||||
default_model = "claude-opus-4"
|
||||
```
|
||||
|
||||
helexa-acp sends `x-api-key` + `anthropic-version: 2023-06-01`
|
||||
automatically. The `api_key_env` indirection keeps your key out of
|
||||
the config file.
|
||||
|
||||
### OpenRouter (multi-vendor proxy)
|
||||
|
||||
```toml
|
||||
[[endpoints]]
|
||||
name = "openrouter"
|
||||
base_url = "https://openrouter.ai/api/v1"
|
||||
wire_api = "openai-chat"
|
||||
api_key_env = "OPENROUTER_API_KEY"
|
||||
default_model = "anthropic/claude-opus-4"
|
||||
```
|
||||
|
||||
OpenRouter speaks OpenAI-compat for every model it fronts, so
|
||||
`openai-chat` is the right wire format regardless of the
|
||||
underlying vendor.
|
||||
|
||||
### LM Studio (local)
|
||||
|
||||
```toml
|
||||
[[endpoints]]
|
||||
name = "lmstudio"
|
||||
base_url = "http://localhost:1234/v1"
|
||||
wire_api = "openai-chat"
|
||||
default_model = "auto"
|
||||
```
|
||||
|
||||
LM Studio's "auto" model id picks whatever's loaded. Same shape
|
||||
works for Ollama in compat mode (`http://localhost:11434/v1`) and
|
||||
vLLM.
|
||||
|
||||
### Multiple cortex deployments
|
||||
|
||||
```toml
|
||||
[[endpoints]]
|
||||
name = "lan"
|
||||
base_url = "http://hanzalova.internal:31313/v1"
|
||||
wire_api = "openai-chat"
|
||||
default_model = "Qwen/Qwen3.6-27B"
|
||||
|
||||
[[endpoints]]
|
||||
name = "cloud"
|
||||
base_url = "https://cortex.example.com/v1"
|
||||
wire_api = "openai-chat"
|
||||
api_key_env = "CLOUD_CORTEX_KEY"
|
||||
default_model = "Qwen/Qwen3-VL-8B"
|
||||
```
|
||||
|
||||
Use the `endpoint:model` selector to switch between them mid-session.
|
||||
|
||||
## Zed setup
|
||||
|
||||
`~/.config/zed/settings.json`:
|
||||
|
||||
```jsonc
|
||||
{
|
||||
"agent_servers": {
|
||||
"helexa": {
|
||||
"command": "helexa-acp"
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Optional environment overrides for the binary:
|
||||
|
||||
```jsonc
|
||||
{
|
||||
"agent_servers": {
|
||||
"helexa": {
|
||||
"command": "helexa-acp",
|
||||
"env": {
|
||||
"HELEXA_ACP_LOG_FILE": "/tmp/helexa-acp.log",
|
||||
"RUST_LOG": "helexa_acp=debug"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
`HELEXA_ACP_LOG_FILE` is the one you actually want — Zed doesn't
|
||||
surface the agent's stderr, so without that env var debug output is
|
||||
invisible. Point it at a file you can `tail -f`.
|
||||
|
||||
After restarting Zed: ⌘+? (or wherever your "Open Agent Panel"
|
||||
binding is) → select "helexa" → the model dropdown populates from
|
||||
your config → start prompting.
|
||||
|
||||
## Modes
|
||||
|
||||
Three session modes ship; the user picks via Zed's mode dropdown
|
||||
on the agent panel.
|
||||
|
||||
| Mode | Reads | Writes | Bash | Permission prompts |
|
||||
|------|-------|--------|------|--------------------|
|
||||
| **Default** | ✓ | with prompt | with prompt | per call |
|
||||
| **Bypass Permissions** | ✓ | ✓ | ✓ | never |
|
||||
| **Plan** | ✓ | only into plan dir | disabled | never (plan-dir writes auto-allow) |
|
||||
|
||||
### Default
|
||||
|
||||
Reads are always allowed (`read_file`, `list_dir` are
|
||||
unrestricted). Writes and shell commands prompt the user before
|
||||
running. The intended baseline for any session where the agent
|
||||
might do something you'd rather review first.
|
||||
|
||||
### Bypass Permissions
|
||||
|
||||
Auto-allow every tool call. Use for agentic loops you trust — bulk
|
||||
edits across many files, scripted workflows, prepared session
|
||||
templates. Never for code the agent hasn't seen before.
|
||||
|
||||
### Plan
|
||||
|
||||
The "draft an implementation plan before you write code" mode.
|
||||
Available tools:
|
||||
|
||||
- `read_file`, `list_dir`: unrestricted (read the codebase).
|
||||
- `write_file`, `edit_file`: allowed *only* under
|
||||
`$XDG_DATA_HOME/helexa-acp/plans/<project-id>/`. Any path
|
||||
outside that returns "plan mode: writes are restricted to …"
|
||||
back to the model so it self-corrects.
|
||||
- `bash`: disabled outright. Returns "plan mode: shell execution
|
||||
is disabled" if attempted.
|
||||
|
||||
When the plan is complete, the model presents a 3-option menu:
|
||||
|
||||
1. **Bypass Permissions** — implement the plan now, no prompts.
|
||||
2. **Default** — implement now with per-tool prompts.
|
||||
3. **Plan** (stay here) — refine the plan with more guidance.
|
||||
|
||||
Switch the mode dropdown to your preference and reply to proceed.
|
||||
|
||||
## Tools
|
||||
|
||||
Five tools, defined in `src/tools.rs`:
|
||||
|
||||
| Tool | Args | Gated in Default? |
|
||||
|------|------|-------------------|
|
||||
| `read_file` | `path`, `line?`, `limit?` | no |
|
||||
| `list_dir` | `path` | no |
|
||||
| `write_file` | `path`, `content` | yes |
|
||||
| `edit_file` | `path`, `old_text`, `new_text` | yes |
|
||||
| `bash` | `command`, `cwd?` | yes |
|
||||
|
||||
### Path handling
|
||||
|
||||
`~`, `~/`, `$HOME`, and `$HOME/` are expanded server-side before
|
||||
the path reaches ACP or local fs. Lets the model emit
|
||||
`~/git/repo/file.rs` and have it Just Work.
|
||||
|
||||
`read_file` first tries the editor's filesystem (ACP's
|
||||
`fs/read_text_file` — respects open buffers, workspace overlays,
|
||||
etc.). If that fails — typically because the path is outside Zed's
|
||||
workspace boundary — it falls back to `std::fs::read_to_string`.
|
||||
This lets the agent pull in shared material like
|
||||
`~/git/architecture/generic.md` from a different project's
|
||||
session.
|
||||
|
||||
The fallback is logged at warn level so you can see when it kicks
|
||||
in.
|
||||
|
||||
### Tool dispatch
|
||||
|
||||
Tool descriptions reach the model through a Qwen3 Hermes-format
|
||||
`# Tools` block injected into the system prompt — cortex/neuron
|
||||
pass the OpenAI `tools` request field through to the encoder
|
||||
unread, so we work the model into emitting `<tool_call>{json}</tool_call>`
|
||||
markers it then parses out of the content stream. This applies to
|
||||
the helexa wire format; OpenAI / Anthropic endpoints with native
|
||||
tool support would use their own paths once they're wired in.
|
||||
|
||||
The parser is tolerant: malformed JSON (trailing braces, missing
|
||||
`name`, name nested in `arguments`) gets a repair pass; if that
|
||||
fails the call surfaces as a "Malformed tool call" card in Zed and
|
||||
the model gets a synthetic error result so it can self-correct.
|
||||
|
||||
## Session resume
|
||||
|
||||
helexa-acp persists every session to
|
||||
`$XDG_DATA_HOME/helexa-acp/sessions/<id>.json`. Zed's `session/list`
|
||||
RPC asks helexa-acp to enumerate them on workspace open;
|
||||
`session/load` rehydrates and replays the transcript as
|
||||
`session/update` notifications so the agent panel renders the
|
||||
prior conversation.
|
||||
|
||||
Behaviour:
|
||||
|
||||
- Persisted per-round, so a mid-turn agent stall (long bash, wedged
|
||||
ACP roundtrip) doesn't lose earlier rounds.
|
||||
- Survives editor restart and the helexa-acp binary upgrading
|
||||
between versions.
|
||||
- Project-scoped: only sessions whose `cwd` matches the workspace
|
||||
are listed.
|
||||
|
||||
To wipe history: `rm -rf $XDG_DATA_HOME/helexa-acp/sessions/`.
|
||||
|
||||
## Context compaction
|
||||
|
||||
When an endpoint sets `context_window`, helexa-acp projects the
|
||||
rolling history into a token budget before each request — old
|
||||
`ToolResult` content (read_file payloads are the worst offenders)
|
||||
gets elided to one-line markers, preserving `tool_call_id` pairing
|
||||
so the wire schema stays valid.
|
||||
|
||||
System prompts, user turns, and the most recent ~4 messages are
|
||||
never elided. The full history stays on disk; compaction is a
|
||||
per-request projection, not a destructive edit.
|
||||
|
||||
Set `context_window = 32768` for a 32 K Qwen3, `131072` for a
|
||||
modern Claude, etc. With `max_tokens` also set, the budget is
|
||||
`context_window - max_tokens - 512_safety`.
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### "default endpoint 'helexa' has no usable provider — check config"
|
||||
|
||||
The named default endpoint failed to construct. Usually:
|
||||
|
||||
- `api_key_env` references a variable that isn't set in the env
|
||||
Zed launched helexa-acp with.
|
||||
- The TOML's `wire_api` is misspelled (only `openai-chat`,
|
||||
`openai-responses`, `anthropic-messages` are accepted).
|
||||
|
||||
Test by running `helexa-acp` directly from a shell — startup
|
||||
errors land on stderr.
|
||||
|
||||
### Model dropdown is empty
|
||||
|
||||
Each provider's `list_models` failed at startup. Look at
|
||||
`HELEXA_ACP_LOG_FILE` for "list_models failed; this endpoint's
|
||||
models won't appear in the picker". Likely the endpoint URL is
|
||||
wrong, the API key is invalid, or the upstream `/v1/models`
|
||||
endpoint isn't responding.
|
||||
|
||||
The agent still works against `default_model` even when the
|
||||
dropdown is empty — list-models is for picking, not routing.
|
||||
|
||||
### "prompt_too_long" / agent stalls mid-conversation
|
||||
|
||||
You hit the model's context window. Set `context_window` on the
|
||||
endpoint and helexa-acp will compact before sending. The log line
|
||||
`context compaction applied` confirms it's running; if it fires
|
||||
but the upstream still rejects, the compaction heuristic
|
||||
under-counted and the budget needs tuning down.
|
||||
|
||||
### Reading files outside the workspace returns "not found"
|
||||
|
||||
Zed's `fs/read_text_file` is workspace-scoped. helexa-acp falls
|
||||
back to local `std::fs` automatically when that fails — look for
|
||||
`fs/read_text_file failed; falling back to local std::fs` in the
|
||||
log. If even local read fails, the file genuinely doesn't exist
|
||||
or the user process lacks permissions.
|
||||
|
||||
### Tool calls render as text instead of structured cards
|
||||
|
||||
The model is emitting `<tool_call>` markers that the parser can't
|
||||
decode. Two common causes:
|
||||
|
||||
1. The system prompt isn't reaching the model (cortex/neuron's
|
||||
tool-block injection didn't fire). Confirm with
|
||||
`RUST_LOG=helexa_acp=debug` and look at the outgoing
|
||||
`POST /chat/completions` body.
|
||||
2. The model itself is too small / undertrained to follow the
|
||||
Hermes format reliably. helexa-acp has shape-based name
|
||||
inference and JSON repair, but there's a floor below which
|
||||
nothing helps.
|
||||
|
||||
### Plan-mode writes refused even inside the plan dir
|
||||
|
||||
The path comparison is byte-for-byte. If the model emits a path
|
||||
with `~` and the plan_dir has the expanded form, expansion runs
|
||||
*before* the comparison — but resolved-vs-symlinked-path
|
||||
mismatches can still bite. The error message names the attempted
|
||||
path and the expected prefix so you can compare directly.
|
||||
|
||||
## Architecture
|
||||
|
||||
Source layout under `crates/helexa-acp/src/`:
|
||||
|
||||
| File | Responsibility |
|
||||
|------|----------------|
|
||||
| `main.rs` | tokio + Stdio transport. Builds providers, hands off to `agent::Agent` |
|
||||
| `config.rs` | TOML + env-fallback config, endpoint resolver |
|
||||
| `agent.rs` | ACP handlers (initialize, session/new, session/prompt, session/cancel, session/set_mode, session/set_model, session/load, session/list), prompt loop with tool-call recursion |
|
||||
| `session.rs` | Per-session state map (Arc<RwLock<HashMap<…>>>) |
|
||||
| `store.rs` | On-disk session persistence, plan-dir resolution |
|
||||
| `prompt.rs` | System-prompt assembly, plan-mode addendum |
|
||||
| `tools.rs` | Tool schemas + shape-based name inference |
|
||||
| `tool_runner.rs` | Dispatch a single tool call through ACP client RPCs; permission gate |
|
||||
| `qwen3.rs` | Qwen3 Hermes tool-format parser (`<tool_call>` / `<think>` markers) |
|
||||
| `compaction.rs` | Token-budget compaction for the rolling history |
|
||||
| `path_util.rs` | `~` / `$HOME` expansion shared across every path-taking tool |
|
||||
| `provider/openai_chat.rs` | OpenAI chat completions provider |
|
||||
| `provider/openai_responses.rs` | OpenAI Responses API provider |
|
||||
| `provider/anthropic_messages.rs` | Anthropic Messages API provider |
|
||||
|
||||
### Adding a new wire format
|
||||
|
||||
1. New file under `src/provider/` implementing the `Provider`
|
||||
trait (encoder + SSE decoder).
|
||||
2. Add a `WireApi` variant in `config.rs`.
|
||||
3. Wire it into `build_provider` in `main.rs`.
|
||||
4. Done — every other module is wire-format-agnostic.
|
||||
|
||||
### Concurrency
|
||||
|
||||
- `Arc<RwLock<HashMap<SessionId, Arc<Mutex<SessionState>>>>>` —
|
||||
per-session mutex so concurrent requests across sessions don't
|
||||
contend; the map's RwLock is read-mostly.
|
||||
- Every tool call dispatched serially within a session (parallel
|
||||
dispatch would require Zed to handle interleaved permission
|
||||
prompts).
|
||||
- Provider streams are back-pressured by the consumer (bounded
|
||||
mpsc channels).
|
||||
|
||||
### Self-contained
|
||||
|
||||
The crate has no workspace-internal dependencies (no
|
||||
`cortex-core`, no `cortex-gateway`). Migration to a dedicated
|
||||
GitHub repo for cross-platform CI / cargo-dist binaries is
|
||||
Cargo.toml-only.
|
||||
|
||||
## Status
|
||||
|
||||
- Stages 1–6 shipped: scaffold, agent loop, tools, modes, session
|
||||
resume, image input, model picker, three wire formats.
|
||||
- Stage 8 (RPM + multi-platform CI) tracked in the canonical plan;
|
||||
Linux x86_64 RPM ships today via the cortex monorepo's Gitea
|
||||
Actions.
|
||||
|
||||
## Contributing
|
||||
|
||||
Repository: https://git.lair.cafe/helexa/cortex (`crates/helexa-acp/`).
|
||||
Issues / PRs welcome. The canonical staged plan is in
|
||||
`~/.claude/plans/plan-the-per-device-worker-abstract-micali.md` on
|
||||
the maintainer's machine; the substages 3a–3e and 6a/6b that the
|
||||
canonical plan didn't anticipate are documented in commit messages.
|
||||
|
||||
CI: `cargo fmt --check --all`, `cargo clippy --workspace -- -D
|
||||
warnings`, `cargo test --workspace` must all pass before merge.
|
||||
1820
crates/helexa-acp/src/agent.rs
Normal file
1820
crates/helexa-acp/src/agent.rs
Normal file
File diff suppressed because it is too large
Load Diff
425
crates/helexa-acp/src/compaction.rs
Normal file
425
crates/helexa-acp/src/compaction.rs
Normal file
@@ -0,0 +1,425 @@
|
||||
//! Rolling-conversation compaction for small-context local models.
|
||||
//!
|
||||
//! The tool-call loop in [`crate::agent`] grows the message vec it
|
||||
//! sends upstream every round. On a frontier model that's fine; on a
|
||||
//! 32 K Qwen3 the first few `read_file` results can push the prompt
|
||||
//! past the model's context window, at which point cortex/neuron
|
||||
//! refuses with `prompt_too_long` and the whole turn dies. Long-form
|
||||
//! local agents are unusable without something here.
|
||||
//!
|
||||
//! Strategy (intentionally simple — no LLM-summarization round-trip,
|
||||
//! no tokenizer dependency):
|
||||
//!
|
||||
//! 1. **Protect** the things the model cannot reason without:
|
||||
//! - The system prompt (idx 0).
|
||||
//! - Every `Role::User` turn (the user's intent — irreplaceable).
|
||||
//! - The last [`KEEP_TAIL`] messages (most recent rounds stay
|
||||
//! verbatim so the model can keep working on what it just
|
||||
//! observed).
|
||||
//! 2. **Elide** older `Role::Assistant` prose and older `Role::Tool`
|
||||
//! result content. The structure stays — `tool_call_id`s, tool
|
||||
//! names, and argument JSON survive intact — so OpenAI's strict
|
||||
//! `tool_calls` ↔ `tool` pairing schema remains satisfied. Only
|
||||
//! the *payload* shrinks to a one-line marker.
|
||||
//! 3. Walk oldest→newest, recomputing the budget after each elision.
|
||||
//! Stop as soon as we fit; we don't compact more than necessary.
|
||||
//! 4. If we still exceed budget after eliding everything we're
|
||||
//! allowed to, return what we have. The upstream will surface a
|
||||
//! `prompt_too_long` error and the user can intervene; that's
|
||||
//! better than silently dropping content the model needs.
|
||||
//!
|
||||
//! Token estimation uses a `chars / 3.5` heuristic — conservative
|
||||
//! (over-estimates tokens slightly) so we compact a touch early
|
||||
//! rather than a touch late.
|
||||
|
||||
use crate::provider::{Message, MessageContent, MessagePart, Role};
|
||||
|
||||
/// Most-recent N messages that are never elided. Roughly "the
|
||||
/// current tool round in flight" — assistant turn that called the
|
||||
/// tools + each tool result + a bit of slack.
|
||||
const KEEP_TAIL: usize = 4;
|
||||
|
||||
/// Below this content size we don't bother eliding — the savings
|
||||
/// don't outweigh the loss of detail. Roughly 60–80 tokens.
|
||||
const ELIDE_MIN_CHARS: usize = 256;
|
||||
|
||||
/// Roughly tokens-per-character for English + code mixed in. The
|
||||
/// actual per-tokenizer ratio varies (GPT-4o ≈ 4 chars/token on
|
||||
/// English prose, ≈ 3 chars/token on code-heavy text). We pick a
|
||||
/// value on the conservative end so the budget check fires *before*
|
||||
/// the upstream tokenizer says no.
|
||||
const CHARS_PER_TOKEN: f32 = 3.5;
|
||||
|
||||
/// Per-message envelope overhead (role + JSON framing). Comes out
|
||||
/// to a few tokens; tiny but it adds up across long histories.
|
||||
const ENVELOPE_TOKENS: usize = 8;
|
||||
|
||||
/// Rough per-image token cost used by the budget estimator. Real
|
||||
/// vision tokenizers vary widely (256–1024 tokens for typical
|
||||
/// resolutions on Qwen3-VL, OpenAI's `low`/`high` detail toggles
|
||||
/// pick between ~85 and ~1000+). 512 is a defensible middle that
|
||||
/// keeps compaction from treating images as free.
|
||||
const IMAGE_TOKENS_APPROX: usize = 512;
|
||||
|
||||
/// Stats reported back from [`compact_to_budget`] for the caller to
|
||||
/// log. The numbers are estimates (see [`estimate_tokens`]), so
|
||||
/// don't compare them to upstream-reported token counts as if they
|
||||
/// were exact.
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct CompactionStats {
|
||||
/// Estimated tokens in the input messages.
|
||||
pub original_tokens: usize,
|
||||
/// Estimated tokens after compaction. Equal to `original_tokens`
|
||||
/// when no compaction was needed.
|
||||
pub final_tokens: usize,
|
||||
/// Number of messages whose content was elided. Zero is the
|
||||
/// hot path (nothing to do).
|
||||
pub elided_messages: usize,
|
||||
}
|
||||
|
||||
impl CompactionStats {
|
||||
fn unchanged(tokens: usize) -> Self {
|
||||
Self {
|
||||
original_tokens: tokens,
|
||||
final_tokens: tokens,
|
||||
elided_messages: 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Approximate token count for one message. Sums the textual
|
||||
/// payload's chars, divides by [`CHARS_PER_TOKEN`], and adds an
|
||||
/// envelope constant. Cheap (no allocation) so safe to call once per
|
||||
/// message per round.
|
||||
pub fn estimate_tokens(msg: &Message) -> usize {
|
||||
let chars = match &msg.content {
|
||||
MessageContent::Text { text } => text.len(),
|
||||
MessageContent::MultiPart { parts } => parts
|
||||
.iter()
|
||||
.map(|p| match p {
|
||||
MessagePart::Text { text } => text.len(),
|
||||
// Each image is one block in the context window; the
|
||||
// upstream tokenizer handles the real cost (and it
|
||||
// varies wildly by model — Qwen3-VL uses ~256-1024
|
||||
// tokens per image depending on size). Take a
|
||||
// middle estimate so the budget tracker doesn't
|
||||
// pretend images are free.
|
||||
MessagePart::Image(_) => IMAGE_TOKENS_APPROX * CHARS_PER_TOKEN as usize,
|
||||
})
|
||||
.sum(),
|
||||
MessageContent::ToolCalls { text, calls } => {
|
||||
let txt = text.as_deref().map(|s| s.len()).unwrap_or(0);
|
||||
let calls_size: usize = calls
|
||||
.iter()
|
||||
.map(|c| c.name.len() + c.arguments.len() + c.id.len())
|
||||
.sum();
|
||||
txt + calls_size
|
||||
}
|
||||
MessageContent::ToolResult {
|
||||
tool_call_id,
|
||||
content,
|
||||
} => tool_call_id.len() + content.len(),
|
||||
};
|
||||
((chars as f32 / CHARS_PER_TOKEN) as usize) + ENVELOPE_TOKENS
|
||||
}
|
||||
|
||||
/// Sum of [`estimate_tokens`] across all messages.
|
||||
pub fn total_tokens(messages: &[Message]) -> usize {
|
||||
messages.iter().map(estimate_tokens).sum()
|
||||
}
|
||||
|
||||
/// Project `messages` into a vec whose estimated token count fits in
|
||||
/// `budget` tokens. Returns the projection plus stats about what
|
||||
/// was done. When the input already fits, the projection is a clone
|
||||
/// of the input and stats report zero elisions.
|
||||
///
|
||||
/// See module docs for the strategy and protected set.
|
||||
pub fn compact_to_budget(messages: &[Message], budget: usize) -> (Vec<Message>, CompactionStats) {
|
||||
let original = total_tokens(messages);
|
||||
if original <= budget {
|
||||
return (messages.to_vec(), CompactionStats::unchanged(original));
|
||||
}
|
||||
|
||||
let mut out = messages.to_vec();
|
||||
let len = out.len();
|
||||
let tail_start = len.saturating_sub(KEEP_TAIL);
|
||||
let mut elided = 0usize;
|
||||
|
||||
// Two passes. First pass: ToolResult contents (largest savings
|
||||
// per elision — read_file payloads land here). Second pass: long
|
||||
// Assistant prose. We don't interleave because eliding a long
|
||||
// assistant turn before a really old read_file would do less
|
||||
// good per elision; oldest-first ordering is enforced *within*
|
||||
// each pass instead.
|
||||
for pass in 0..2 {
|
||||
for i in 1..tail_start {
|
||||
if matches!(out[i].role, Role::User) {
|
||||
continue;
|
||||
}
|
||||
let target_pass_2 = matches!(
|
||||
&out[i].content,
|
||||
MessageContent::Text { .. } | MessageContent::ToolCalls { .. }
|
||||
);
|
||||
let target_pass_1 = matches!(&out[i].content, MessageContent::ToolResult { .. });
|
||||
let in_pass = (pass == 0 && target_pass_1) || (pass == 1 && target_pass_2);
|
||||
if !in_pass {
|
||||
continue;
|
||||
}
|
||||
if elide_in_place(&mut out[i]) {
|
||||
elided += 1;
|
||||
if total_tokens(&out) <= budget {
|
||||
let final_tokens = total_tokens(&out);
|
||||
return (
|
||||
out,
|
||||
CompactionStats {
|
||||
original_tokens: original,
|
||||
final_tokens,
|
||||
elided_messages: elided,
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let final_tokens = total_tokens(&out);
|
||||
(
|
||||
out,
|
||||
CompactionStats {
|
||||
original_tokens: original,
|
||||
final_tokens,
|
||||
elided_messages: elided,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
/// Shrink one message's payload while keeping its structural role
|
||||
/// (so tool_call_id pairing survives). Returns `true` when the
|
||||
/// message changed.
|
||||
///
|
||||
/// - `ToolResult.content` → `(elided: N bytes of tool result)`
|
||||
/// - `ToolCalls.text` → `(elided: N bytes of assistant prose)`
|
||||
/// - `Text` (assistant) → `(elided: N bytes of assistant prose)`
|
||||
///
|
||||
/// Already-tiny payloads are skipped — eliding a 50-byte string
|
||||
/// would *grow* it once the marker is in place.
|
||||
fn elide_in_place(msg: &mut Message) -> bool {
|
||||
match &mut msg.content {
|
||||
MessageContent::ToolResult { content, .. } => {
|
||||
if content.len() < ELIDE_MIN_CHARS {
|
||||
return false;
|
||||
}
|
||||
*content = format!("(elided: {} bytes of tool result)", content.len());
|
||||
true
|
||||
}
|
||||
MessageContent::ToolCalls { text, .. } => match text {
|
||||
Some(t) if t.len() >= ELIDE_MIN_CHARS => {
|
||||
*text = Some(format!("(elided: {} bytes of assistant prose)", t.len()));
|
||||
true
|
||||
}
|
||||
_ => false,
|
||||
},
|
||||
MessageContent::Text { text } => {
|
||||
if text.len() < ELIDE_MIN_CHARS {
|
||||
return false;
|
||||
}
|
||||
*text = format!("(elided: {} bytes of assistant prose)", text.len());
|
||||
true
|
||||
}
|
||||
MessageContent::MultiPart { .. } => {
|
||||
// MultiPart messages today only exist as User turns,
|
||||
// and User turns are protected by the role check in
|
||||
// `compact_to_budget` — so this branch is unreachable
|
||||
// for current call sites. Returning false keeps the
|
||||
// unreachable path benign if a future stage starts
|
||||
// emitting MultiPart on other roles.
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::provider::ToolCall;
|
||||
|
||||
fn sys(text: &str) -> Message {
|
||||
Message {
|
||||
role: Role::System,
|
||||
content: MessageContent::Text { text: text.into() },
|
||||
}
|
||||
}
|
||||
fn user(text: &str) -> Message {
|
||||
Message {
|
||||
role: Role::User,
|
||||
content: MessageContent::Text { text: text.into() },
|
||||
}
|
||||
}
|
||||
fn assistant_text(text: &str) -> Message {
|
||||
Message {
|
||||
role: Role::Assistant,
|
||||
content: MessageContent::Text { text: text.into() },
|
||||
}
|
||||
}
|
||||
fn assistant_calls(text: Option<&str>, name: &str, args: &str, id: &str) -> Message {
|
||||
Message {
|
||||
role: Role::Assistant,
|
||||
content: MessageContent::ToolCalls {
|
||||
text: text.map(|s| s.to_string()),
|
||||
calls: vec![ToolCall {
|
||||
id: id.into(),
|
||||
name: name.into(),
|
||||
arguments: args.into(),
|
||||
}],
|
||||
},
|
||||
}
|
||||
}
|
||||
fn tool_result(id: &str, body: &str) -> Message {
|
||||
Message {
|
||||
role: Role::Tool,
|
||||
content: MessageContent::ToolResult {
|
||||
tool_call_id: id.into(),
|
||||
content: body.into(),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn under_budget_is_a_no_op_clone() {
|
||||
let msgs = vec![sys("you are an agent"), user("hi"), assistant_text("hello")];
|
||||
let (out, stats) = compact_to_budget(&msgs, 10_000);
|
||||
assert_eq!(stats.elided_messages, 0);
|
||||
assert_eq!(stats.original_tokens, stats.final_tokens);
|
||||
assert_eq!(out.len(), msgs.len());
|
||||
// Strings unchanged.
|
||||
match &out[2].content {
|
||||
MessageContent::Text { text } => assert_eq!(text, "hello"),
|
||||
other => panic!("expected Text, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn elides_old_tool_result_before_old_assistant_prose() {
|
||||
// History: sys, user, assistant_calls, big_tool_result,
|
||||
// assistant_with_big_text, user, assistant_calls,
|
||||
// small_tool_result.
|
||||
// KEEP_TAIL=4 protects the last four; the big tool result
|
||||
// sits in the prunable range and should go first because
|
||||
// pass 0 (tool results) runs before pass 1 (prose).
|
||||
let big_result = "X".repeat(4096);
|
||||
let big_prose = "Y".repeat(2048);
|
||||
let msgs = vec![
|
||||
sys("preamble"),
|
||||
user("first ask"),
|
||||
assistant_calls(None, "read_file", r#"{"path":"/a"}"#, "c0"),
|
||||
tool_result("c0", &big_result),
|
||||
assistant_text(&big_prose),
|
||||
user("follow up"),
|
||||
assistant_calls(None, "read_file", r#"{"path":"/b"}"#, "c1"),
|
||||
tool_result("c1", "short result body"),
|
||||
];
|
||||
let before = total_tokens(&msgs);
|
||||
// Force compaction by setting budget well below current.
|
||||
let budget = before / 2;
|
||||
let (out, stats) = compact_to_budget(&msgs, budget);
|
||||
|
||||
assert!(
|
||||
stats.elided_messages >= 1,
|
||||
"expected at least one elision, got {stats:?}"
|
||||
);
|
||||
// The big tool result must be elided (oldest fat target).
|
||||
match &out[3].content {
|
||||
MessageContent::ToolResult { content, .. } => {
|
||||
assert!(
|
||||
content.starts_with("(elided:"),
|
||||
"tool result not elided: {content:?}"
|
||||
);
|
||||
}
|
||||
other => panic!("expected ToolResult, got {other:?}"),
|
||||
}
|
||||
// Last four messages must be untouched.
|
||||
assert!(matches!(
|
||||
&out[out.len() - 1].content,
|
||||
MessageContent::ToolResult { content, .. } if content == "short result body"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn never_elides_system_or_user_turns() {
|
||||
let big_user = "U".repeat(8192);
|
||||
let msgs = vec![sys("preamble"), user(&big_user), assistant_text("ok")];
|
||||
let budget = 10; // way below — forces all possible elision
|
||||
let (out, _stats) = compact_to_budget(&msgs, budget);
|
||||
// System unchanged.
|
||||
match &out[0].content {
|
||||
MessageContent::Text { text } => assert_eq!(text, "preamble"),
|
||||
other => panic!("expected Text, got {other:?}"),
|
||||
}
|
||||
// User unchanged even though it's huge.
|
||||
match &out[1].content {
|
||||
MessageContent::Text { text } => assert_eq!(text.len(), big_user.len()),
|
||||
other => panic!("expected Text, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn preserves_tool_call_id_pairing_after_elision() {
|
||||
// OpenAI strict mode rejects a tool-result whose tool_call_id
|
||||
// doesn't match a preceding assistant tool_call. Elision
|
||||
// must not break that linkage.
|
||||
let big = "Z".repeat(4096);
|
||||
let msgs = vec![
|
||||
sys("preamble"),
|
||||
user("first"),
|
||||
assistant_calls(None, "read_file", r#"{"path":"/a"}"#, "call_42"),
|
||||
tool_result("call_42", &big),
|
||||
// Tail messages.
|
||||
user("next"),
|
||||
assistant_calls(None, "read_file", r#"{"path":"/b"}"#, "call_43"),
|
||||
tool_result("call_43", "ok"),
|
||||
assistant_text("done"),
|
||||
];
|
||||
let budget = total_tokens(&msgs) / 3;
|
||||
let (out, _stats) = compact_to_budget(&msgs, budget);
|
||||
// The assistant call and its result both carry call_42.
|
||||
let call_id = match &out[2].content {
|
||||
MessageContent::ToolCalls { calls, .. } => calls[0].id.clone(),
|
||||
other => panic!("expected ToolCalls, got {other:?}"),
|
||||
};
|
||||
match &out[3].content {
|
||||
MessageContent::ToolResult { tool_call_id, .. } => {
|
||||
assert_eq!(tool_call_id, &call_id, "pairing broken");
|
||||
}
|
||||
other => panic!("expected ToolResult, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn estimate_tokens_grows_with_content() {
|
||||
let small = sys("hi");
|
||||
let large = sys(&"x".repeat(10_000));
|
||||
assert!(estimate_tokens(&large) > estimate_tokens(&small) * 100);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn elide_in_place_skips_short_content() {
|
||||
let mut m = tool_result("c0", "tiny");
|
||||
assert!(!elide_in_place(&mut m));
|
||||
match m.content {
|
||||
MessageContent::ToolResult { content, .. } => assert_eq!(content, "tiny"),
|
||||
other => panic!("expected ToolResult, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn returns_best_effort_when_budget_unmeetable() {
|
||||
// Single huge user message that cannot be elided. Budget 10.
|
||||
// We don't error — we return what we have and let upstream
|
||||
// refuse the prompt with its own error.
|
||||
let big_user = "U".repeat(100_000);
|
||||
let msgs = vec![sys("preamble"), user(&big_user)];
|
||||
let (out, stats) = compact_to_budget(&msgs, 10);
|
||||
assert_eq!(out.len(), msgs.len());
|
||||
assert!(stats.final_tokens > 10, "still over budget by design");
|
||||
}
|
||||
}
|
||||
424
crates/helexa-acp/src/config.rs
Normal file
424
crates/helexa-acp/src/config.rs
Normal file
@@ -0,0 +1,424 @@
|
||||
//! Configuration for the helexa-acp bridge.
|
||||
//!
|
||||
//! Loaded from `$XDG_CONFIG_HOME/helexa-acp/config.toml` (or
|
||||
//! `~/.config/helexa-acp/config.toml` as a fallback). If no config file
|
||||
//! exists, falls back to building a single anonymous endpoint from env
|
||||
//! vars — that keeps "just point at one cortex" frictionless without
|
||||
//! requiring a config file on disk.
|
||||
//!
|
||||
//! The design goal is "the missing ACP binary for users with multiple
|
||||
//! API endpoints (possibly on a private LAN, possibly mixing wire
|
||||
//! types)". Hence: every endpoint is named, has its own wire API, and
|
||||
//! has its own default model. The agent's selected model id can be
|
||||
//! prefixed `endpoint:model` to route across endpoints; a bare
|
||||
//! `model` falls through to the configured `default_endpoint`.
|
||||
//!
|
||||
//! ### Example TOML
|
||||
//!
|
||||
//! ```toml
|
||||
//! default_endpoint = "helexa"
|
||||
//!
|
||||
//! [[endpoints]]
|
||||
//! name = "helexa"
|
||||
//! base_url = "http://hanzalova.internal:31313/v1"
|
||||
//! wire_api = "openai-chat"
|
||||
//! default_model = "helexa/large"
|
||||
//!
|
||||
//! [[endpoints]]
|
||||
//! name = "openrouter"
|
||||
//! base_url = "https://openrouter.ai/api/v1"
|
||||
//! wire_api = "openai-chat"
|
||||
//! api_key_env = "OPENROUTER_API_KEY"
|
||||
//! default_model = "anthropic/claude-opus-4"
|
||||
//!
|
||||
//! [[endpoints]]
|
||||
//! name = "lmstudio"
|
||||
//! base_url = "http://localhost:1234/v1"
|
||||
//! wire_api = "openai-chat"
|
||||
//! default_model = "auto"
|
||||
//! ```
|
||||
|
||||
use anyhow::{Context, anyhow};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::path::{Path, PathBuf};
|
||||
use url::Url;
|
||||
|
||||
const DEFAULT_BASE_URL: &str = "http://hanzalova.internal:31313/v1";
|
||||
const DEFAULT_MODEL: &str = "helexa/large";
|
||||
const DEFAULT_ENDPOINT_NAME: &str = "default";
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Config {
|
||||
/// Name of the endpoint used when a request doesn't pick one
|
||||
/// explicitly. Must reference an entry in `endpoints`. Defaults to
|
||||
/// the first endpoint declared if unset.
|
||||
#[serde(default)]
|
||||
pub default_endpoint: Option<String>,
|
||||
/// Per-endpoint configuration. At least one entry is required.
|
||||
#[serde(default)]
|
||||
pub endpoints: Vec<EndpointConfig>,
|
||||
/// Optional path to a system-prompt file. When unset, the built-in
|
||||
/// default prompt from `prompt.rs` is used.
|
||||
#[serde(default)]
|
||||
pub system_prompt_path: Option<PathBuf>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct EndpointConfig {
|
||||
/// Short identifier used in `endpoint:model` routing and in logs.
|
||||
pub name: String,
|
||||
/// Base URL of the OpenAI-compatible API. Must include the `/v1`
|
||||
/// (or equivalent) suffix — paths like `chat/completions` and
|
||||
/// `models` are joined onto this.
|
||||
pub base_url: Url,
|
||||
/// Wire protocol the endpoint speaks. Phase 1 supports
|
||||
/// [`WireApi::OpenAiChat`] only; `openai-responses` and
|
||||
/// `anthropic-messages` land later behind their own providers.
|
||||
#[serde(default)]
|
||||
pub wire_api: WireApi,
|
||||
/// Model to use when the client hasn't picked one via
|
||||
/// `session/set_model`.
|
||||
#[serde(default)]
|
||||
pub default_model: Option<String>,
|
||||
/// Static API key to send as `Authorization: Bearer …`. Prefer
|
||||
/// `api_key_env` for anything sensitive — keys in plain TOML are a
|
||||
/// liability.
|
||||
#[serde(default)]
|
||||
pub api_key: Option<String>,
|
||||
/// Env var name to read for the API key. Resolved at startup so a
|
||||
/// missing env var yields a clear error rather than silent
|
||||
/// unauthenticated calls.
|
||||
#[serde(default)]
|
||||
pub api_key_env: Option<String>,
|
||||
/// Cap on the model's output tokens per turn. `None` lets the
|
||||
/// upstream pick its own default (cortex/neuron's default is
|
||||
/// often small enough to trip Zed's "Output Limit Reached" on
|
||||
/// long responses). Set to e.g. `32768` to let the model
|
||||
/// produce longer turns. Goes into the OpenAI `max_tokens`
|
||||
/// request field.
|
||||
#[serde(default)]
|
||||
pub max_tokens: Option<u64>,
|
||||
/// Model context window in tokens (prompt + response). When set,
|
||||
/// the agent compacts conversation history before each completion
|
||||
/// so the prompt fits within `context_window - max_tokens - safety`
|
||||
/// tokens — long sessions on small-context local models (Qwen3 at
|
||||
/// 32 K) survive past the first few tool-call rounds rather than
|
||||
/// dying with `prompt_too_long`. `None` disables compaction.
|
||||
#[serde(default)]
|
||||
pub context_window: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
|
||||
pub enum WireApi {
|
||||
/// `POST {base}/chat/completions` returning OpenAI-format SSE.
|
||||
/// Compatible with cortex, LM Studio, Ollama (compat mode),
|
||||
/// OpenRouter, OpenAI itself.
|
||||
#[default]
|
||||
#[serde(rename = "openai-chat")]
|
||||
OpenAiChat,
|
||||
/// `POST {base}/responses` — OpenAI's newer Responses API. Not
|
||||
/// implemented yet; the variant is reserved so endpoint configs
|
||||
/// can be authored ahead of provider support.
|
||||
#[serde(rename = "openai-responses")]
|
||||
OpenAiResponses,
|
||||
/// `POST {base}/messages` — Anthropic format. Reserved.
|
||||
#[serde(rename = "anthropic-messages")]
|
||||
AnthropicMessages,
|
||||
}
|
||||
|
||||
impl EndpointConfig {
|
||||
/// Resolve the API key from `api_key` (literal) or `api_key_env`
|
||||
/// (env-var lookup). Returns `Ok(None)` when neither is set;
|
||||
/// `Err` when `api_key_env` references a missing variable.
|
||||
pub fn resolve_api_key(&self) -> anyhow::Result<Option<String>> {
|
||||
if let Some(literal) = &self.api_key {
|
||||
return Ok(Some(literal.clone()));
|
||||
}
|
||||
if let Some(var) = &self.api_key_env {
|
||||
return Ok(Some(std::env::var(var).with_context(|| {
|
||||
format!(
|
||||
"endpoint '{}' references missing env var {}",
|
||||
self.name, var
|
||||
)
|
||||
})?));
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
/// `{base_url}/chat/completions`.
|
||||
pub fn chat_completions_url(&self) -> Url {
|
||||
join_segments(&self.base_url, &["chat", "completions"])
|
||||
}
|
||||
|
||||
/// `{base_url}/responses` — OpenAI Responses API endpoint.
|
||||
pub fn responses_url(&self) -> Url {
|
||||
join_segments(&self.base_url, &["responses"])
|
||||
}
|
||||
|
||||
/// `{base_url}/models`. Called from `Provider::list_models`, which
|
||||
/// Stage 4 wires into the model-picker dropdown; until then it's
|
||||
/// reachable code with no in-tree callers.
|
||||
#[allow(dead_code)]
|
||||
pub fn models_url(&self) -> Url {
|
||||
join_segments(&self.base_url, &["models"])
|
||||
}
|
||||
}
|
||||
|
||||
impl Config {
|
||||
/// Load from TOML at the standard config path, or build from env
|
||||
/// vars if no file exists. Env-fallback yields a single endpoint
|
||||
/// named `"default"`.
|
||||
pub fn load() -> anyhow::Result<Self> {
|
||||
let path = config_path();
|
||||
if let Some(path) = &path
|
||||
&& path.exists()
|
||||
{
|
||||
return Self::from_file(path);
|
||||
}
|
||||
Self::from_env()
|
||||
}
|
||||
|
||||
/// Single-endpoint config constructed from `HELEXA_ACP_BASE_URL`,
|
||||
/// `HELEXA_ACP_MODEL`, `HELEXA_ACP_API_KEY`,
|
||||
/// `HELEXA_ACP_SYSTEM_PROMPT_PATH`, `HELEXA_ACP_MAX_TOKENS`.
|
||||
pub fn from_env() -> anyhow::Result<Self> {
|
||||
let base_url = std::env::var("HELEXA_ACP_BASE_URL")
|
||||
.ok()
|
||||
.unwrap_or_else(|| DEFAULT_BASE_URL.into());
|
||||
let base_url = Url::parse(&base_url)
|
||||
.with_context(|| format!("HELEXA_ACP_BASE_URL is not a valid URL ({base_url})"))?;
|
||||
let default_model = std::env::var("HELEXA_ACP_MODEL")
|
||||
.ok()
|
||||
.unwrap_or_else(|| DEFAULT_MODEL.into());
|
||||
let api_key = std::env::var("HELEXA_ACP_API_KEY")
|
||||
.ok()
|
||||
.filter(|s| !s.is_empty());
|
||||
let system_prompt_path = std::env::var("HELEXA_ACP_SYSTEM_PROMPT_PATH")
|
||||
.ok()
|
||||
.filter(|s| !s.is_empty())
|
||||
.map(PathBuf::from);
|
||||
let max_tokens = std::env::var("HELEXA_ACP_MAX_TOKENS")
|
||||
.ok()
|
||||
.filter(|s| !s.is_empty())
|
||||
.map(|s| {
|
||||
s.parse::<u64>().with_context(|| {
|
||||
format!("HELEXA_ACP_MAX_TOKENS is not a positive integer ({s})")
|
||||
})
|
||||
})
|
||||
.transpose()?;
|
||||
let context_window = std::env::var("HELEXA_ACP_CONTEXT_WINDOW")
|
||||
.ok()
|
||||
.filter(|s| !s.is_empty())
|
||||
.map(|s| {
|
||||
s.parse::<usize>().with_context(|| {
|
||||
format!("HELEXA_ACP_CONTEXT_WINDOW is not a positive integer ({s})")
|
||||
})
|
||||
})
|
||||
.transpose()?;
|
||||
Ok(Self {
|
||||
default_endpoint: Some(DEFAULT_ENDPOINT_NAME.into()),
|
||||
endpoints: vec![EndpointConfig {
|
||||
name: DEFAULT_ENDPOINT_NAME.into(),
|
||||
base_url,
|
||||
wire_api: WireApi::OpenAiChat,
|
||||
default_model: Some(default_model),
|
||||
api_key,
|
||||
api_key_env: None,
|
||||
max_tokens,
|
||||
context_window,
|
||||
}],
|
||||
system_prompt_path,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn from_file(path: &Path) -> anyhow::Result<Self> {
|
||||
let text = std::fs::read_to_string(path)
|
||||
.with_context(|| format!("read config {}", path.display()))?;
|
||||
let mut cfg: Self =
|
||||
toml::from_str(&text).with_context(|| format!("parse config {}", path.display()))?;
|
||||
cfg.validate()?;
|
||||
Ok(cfg)
|
||||
}
|
||||
|
||||
fn validate(&mut self) -> anyhow::Result<()> {
|
||||
if self.endpoints.is_empty() {
|
||||
return Err(anyhow!("config has no [[endpoints]] entries"));
|
||||
}
|
||||
for (i, ep) in self.endpoints.iter().enumerate() {
|
||||
if ep.name.is_empty() {
|
||||
return Err(anyhow!("endpoints[{i}] has empty name"));
|
||||
}
|
||||
if ep.name.contains(':') {
|
||||
return Err(anyhow!(
|
||||
"endpoints[{i}].name '{}' contains ':' which would clash \
|
||||
with the endpoint:model selector syntax",
|
||||
ep.name
|
||||
));
|
||||
}
|
||||
}
|
||||
// Pick a default endpoint if none was named.
|
||||
if self.default_endpoint.is_none() {
|
||||
self.default_endpoint = Some(self.endpoints[0].name.clone());
|
||||
}
|
||||
let default_name = self.default_endpoint.as_deref().unwrap();
|
||||
if !self.endpoints.iter().any(|e| e.name == default_name) {
|
||||
return Err(anyhow!(
|
||||
"default_endpoint '{default_name}' is not declared in [[endpoints]]"
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Look up an endpoint by name. Returns `None` if not configured.
|
||||
pub fn endpoint(&self, name: &str) -> Option<&EndpointConfig> {
|
||||
self.endpoints.iter().find(|e| e.name == name)
|
||||
}
|
||||
|
||||
/// The default endpoint (guaranteed to exist after `validate`).
|
||||
pub fn default_endpoint(&self) -> &EndpointConfig {
|
||||
let name = self
|
||||
.default_endpoint
|
||||
.as_deref()
|
||||
.expect("default_endpoint set by validate");
|
||||
self.endpoint(name)
|
||||
.expect("default_endpoint resolves after validate")
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse an ACP-side `model` field into (endpoint name, raw model id).
|
||||
///
|
||||
/// `helexa:helexa/large` → (`Some("helexa")`, `"helexa/large"`).
|
||||
/// `helexa/large` → (`None`, `"helexa/large"`).
|
||||
///
|
||||
/// The split happens at the FIRST colon. Model ids commonly contain
|
||||
/// `/` (HuggingFace style) but rarely `:`; if a model id ever does, the
|
||||
/// user can quote-prefix with the default endpoint name.
|
||||
pub fn parse_model_selector(input: &str) -> (Option<&str>, &str) {
|
||||
match input.split_once(':') {
|
||||
Some((endpoint, model)) if !endpoint.is_empty() && !model.is_empty() => {
|
||||
(Some(endpoint), model)
|
||||
}
|
||||
_ => (None, input),
|
||||
}
|
||||
}
|
||||
|
||||
fn config_path() -> Option<PathBuf> {
|
||||
if let Ok(override_path) = std::env::var("HELEXA_ACP_CONFIG_PATH") {
|
||||
return Some(PathBuf::from(override_path));
|
||||
}
|
||||
let xdg = std::env::var("XDG_CONFIG_HOME")
|
||||
.ok()
|
||||
.filter(|s| !s.is_empty());
|
||||
let base = xdg.map(PathBuf::from).or_else(|| {
|
||||
std::env::var("HOME")
|
||||
.ok()
|
||||
.map(|h| PathBuf::from(h).join(".config"))
|
||||
})?;
|
||||
Some(base.join("helexa-acp").join("config.toml"))
|
||||
}
|
||||
|
||||
fn join_segments(base: &Url, segments: &[&str]) -> Url {
|
||||
let mut out = base.clone();
|
||||
if let Ok(mut path) = out.path_segments_mut() {
|
||||
path.pop_if_empty().extend(segments.iter().copied());
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn url_join_handles_trailing_slash() {
|
||||
let ep = EndpointConfig {
|
||||
name: "x".into(),
|
||||
base_url: Url::parse("http://h.internal:31313/v1").unwrap(),
|
||||
wire_api: WireApi::OpenAiChat,
|
||||
default_model: None,
|
||||
api_key: None,
|
||||
api_key_env: None,
|
||||
max_tokens: None,
|
||||
context_window: None,
|
||||
};
|
||||
assert_eq!(
|
||||
ep.chat_completions_url().as_str(),
|
||||
"http://h.internal:31313/v1/chat/completions"
|
||||
);
|
||||
assert_eq!(
|
||||
ep.models_url().as_str(),
|
||||
"http://h.internal:31313/v1/models"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_model_selector() {
|
||||
assert_eq!(
|
||||
parse_model_selector("helexa:helexa/large"),
|
||||
(Some("helexa"), "helexa/large")
|
||||
);
|
||||
assert_eq!(parse_model_selector("helexa/large"), (None, "helexa/large"));
|
||||
assert_eq!(parse_model_selector("gpt-5"), (None, "gpt-5"));
|
||||
// Edge case: a leading colon → no endpoint.
|
||||
assert_eq!(parse_model_selector(":gpt-5"), (None, ":gpt-5"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn env_fallback_builds_single_endpoint() {
|
||||
// Don't actually set env vars (would race with other tests);
|
||||
// just confirm the default path constructs cleanly.
|
||||
unsafe {
|
||||
std::env::remove_var("HELEXA_ACP_BASE_URL");
|
||||
std::env::remove_var("HELEXA_ACP_MODEL");
|
||||
std::env::remove_var("HELEXA_ACP_API_KEY");
|
||||
}
|
||||
let cfg = Config::from_env().unwrap();
|
||||
assert_eq!(cfg.endpoints.len(), 1);
|
||||
assert_eq!(cfg.endpoints[0].name, "default");
|
||||
assert_eq!(cfg.endpoints[0].base_url.as_str(), DEFAULT_BASE_URL);
|
||||
assert_eq!(
|
||||
cfg.endpoints[0].default_model.as_deref(),
|
||||
Some(DEFAULT_MODEL)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn toml_parses_multi_endpoint() {
|
||||
let toml_text = r#"
|
||||
default_endpoint = "helexa"
|
||||
|
||||
[[endpoints]]
|
||||
name = "helexa"
|
||||
base_url = "http://hanzalova.internal:31313/v1"
|
||||
default_model = "helexa/large"
|
||||
|
||||
[[endpoints]]
|
||||
name = "openrouter"
|
||||
base_url = "https://openrouter.ai/api/v1"
|
||||
wire_api = "openai-chat"
|
||||
api_key_env = "OPENROUTER_API_KEY"
|
||||
default_model = "anthropic/claude-opus-4"
|
||||
"#;
|
||||
let mut cfg: Config = toml::from_str(toml_text).unwrap();
|
||||
cfg.validate().unwrap();
|
||||
assert_eq!(cfg.endpoints.len(), 2);
|
||||
assert_eq!(cfg.default_endpoint().name, "helexa");
|
||||
assert_eq!(cfg.endpoints[0].wire_api, WireApi::OpenAiChat);
|
||||
assert_eq!(
|
||||
cfg.endpoints[1].api_key_env.as_deref(),
|
||||
Some("OPENROUTER_API_KEY")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validate_rejects_colon_in_endpoint_name() {
|
||||
let toml_text = r#"
|
||||
[[endpoints]]
|
||||
name = "bad:name"
|
||||
base_url = "http://x/v1"
|
||||
"#;
|
||||
let mut cfg: Config = toml::from_str(toml_text).unwrap();
|
||||
let err = cfg.validate().unwrap_err();
|
||||
assert!(format!("{err}").contains("clash"));
|
||||
}
|
||||
}
|
||||
145
crates/helexa-acp/src/main.rs
Normal file
145
crates/helexa-acp/src/main.rs
Normal file
@@ -0,0 +1,145 @@
|
||||
//! helexa-acp — Agent Client Protocol bridge for multi-endpoint LLM
|
||||
//! setups (helexa, LM Studio, Ollama, OpenRouter, OpenAI, Anthropic,
|
||||
//! …) with a clean per-endpoint wire-format selector.
|
||||
//!
|
||||
//! Speaks ACP over stdio to an editor client (Zed today). Every
|
||||
//! configured endpoint produces a wire-format-specific
|
||||
//! [`provider::Provider`] implementation; the agent loop in
|
||||
//! [`agent::Agent`] is provider-agnostic, so adding e.g. an Anthropic
|
||||
//! /v1/messages provider doesn't touch `agent.rs`.
|
||||
//!
|
||||
//! Config: `$XDG_CONFIG_HOME/helexa-acp/config.toml` for the multi-
|
||||
//! endpoint case; env vars (`HELEXA_ACP_BASE_URL`, etc.) for the
|
||||
//! single-endpoint case when no config file exists.
|
||||
|
||||
use agent_client_protocol::{Result, Stdio};
|
||||
use std::sync::Arc;
|
||||
|
||||
mod agent;
|
||||
mod compaction;
|
||||
mod config;
|
||||
mod path_util;
|
||||
mod prompt;
|
||||
mod provider;
|
||||
mod qwen3;
|
||||
mod session;
|
||||
mod store;
|
||||
mod tool_runner;
|
||||
mod tools;
|
||||
|
||||
use agent::Agent;
|
||||
use config::{Config, EndpointConfig, WireApi};
|
||||
use provider::{
|
||||
Provider, anthropic_messages::AnthropicMessagesProvider, openai_chat::OpenAIChatProvider,
|
||||
openai_responses::OpenAIResponsesProvider,
|
||||
};
|
||||
|
||||
/// Set up tracing. Logs go to stderr by default — stdout is
|
||||
/// reserved for the JSON-RPC stream. Setting `HELEXA_ACP_LOG_FILE`
|
||||
/// to an absolute path appends logs to that file instead, which is
|
||||
/// the practical way to capture debug output when the agent runs
|
||||
/// under an editor (Zed, etc.) that doesn't surface stderr.
|
||||
///
|
||||
/// `RUST_LOG` still controls levels (e.g. `helexa_acp=debug`).
|
||||
/// ANSI colours are auto-stripped when writing to a file so the log
|
||||
/// is plain text.
|
||||
fn init_tracing() {
|
||||
let env_filter = tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info"));
|
||||
|
||||
let log_file = std::env::var("HELEXA_ACP_LOG_FILE")
|
||||
.ok()
|
||||
.filter(|s| !s.is_empty());
|
||||
|
||||
match log_file {
|
||||
Some(path) => match std::fs::OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(&path)
|
||||
{
|
||||
Ok(file) => {
|
||||
tracing_subscriber::fmt()
|
||||
.with_writer(std::sync::Mutex::new(file))
|
||||
.with_env_filter(env_filter)
|
||||
.with_ansi(false)
|
||||
.init();
|
||||
}
|
||||
Err(e) => {
|
||||
// Fall back to stderr and shout. We don't want a
|
||||
// typo'd log path to silence the agent entirely.
|
||||
tracing_subscriber::fmt()
|
||||
.with_writer(std::io::stderr)
|
||||
.with_env_filter(env_filter)
|
||||
.init();
|
||||
tracing::warn!(
|
||||
path = %path,
|
||||
error = %e,
|
||||
"HELEXA_ACP_LOG_FILE could not be opened; using stderr"
|
||||
);
|
||||
}
|
||||
},
|
||||
None => {
|
||||
tracing_subscriber::fmt()
|
||||
.with_writer(std::io::stderr)
|
||||
.with_env_filter(env_filter)
|
||||
.init();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Build a provider for `endpoint` according to its declared
|
||||
/// `wire_api`. Future wire types (OpenAI Responses, Anthropic
|
||||
/// /v1/messages, Ollama native) slot in here without changing the
|
||||
/// caller.
|
||||
fn build_provider(endpoint: EndpointConfig) -> anyhow::Result<Arc<dyn Provider>> {
|
||||
match endpoint.wire_api {
|
||||
WireApi::OpenAiChat => Ok(Arc::new(OpenAIChatProvider::new(endpoint)?)),
|
||||
WireApi::OpenAiResponses => Ok(Arc::new(OpenAIResponsesProvider::new(endpoint)?)),
|
||||
WireApi::AnthropicMessages => Ok(Arc::new(AnthropicMessagesProvider::new(endpoint)?)),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
init_tracing();
|
||||
|
||||
let cfg = Config::load()
|
||||
.map_err(|e| agent_client_protocol::util::internal_error(format!("config: {e:#}")))?;
|
||||
tracing::info!(
|
||||
endpoints = cfg.endpoints.len(),
|
||||
default_endpoint = %cfg.default_endpoint().name,
|
||||
default_model = ?cfg.default_endpoint().default_model,
|
||||
"helexa-acp starting"
|
||||
);
|
||||
|
||||
// Build a provider for each configured endpoint up-front. Cheap —
|
||||
// just sets up a reqwest::Client and resolves the API key — and
|
||||
// surfaces config mistakes (missing API key env var, unsupported
|
||||
// wire_api) before the editor even sends an initialize request.
|
||||
let mut providers: Vec<Arc<dyn Provider>> = Vec::with_capacity(cfg.endpoints.len());
|
||||
for endpoint in &cfg.endpoints {
|
||||
match build_provider(endpoint.clone()) {
|
||||
Ok(p) => {
|
||||
tracing::info!(
|
||||
endpoint = %endpoint.name,
|
||||
base_url = %endpoint.base_url,
|
||||
wire_api = ?endpoint.wire_api,
|
||||
"registered provider"
|
||||
);
|
||||
providers.push(p);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
endpoint = %endpoint.name,
|
||||
error = %format!("{e:#}"),
|
||||
"skipping endpoint with invalid config"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let agent = Agent::new(&cfg, providers)
|
||||
.await
|
||||
.map_err(|e| agent_client_protocol::util::internal_error(format!("agent: {e:#}")))?;
|
||||
agent.serve(Stdio::new()).await
|
||||
}
|
||||
192
crates/helexa-acp/src/path_util.rs
Normal file
192
crates/helexa-acp/src/path_util.rs
Normal file
@@ -0,0 +1,192 @@
|
||||
//! Path expansion shared across every tool that takes a path.
|
||||
//!
|
||||
//! Models often emit shell-style paths like `~/git/repo/file.rs` or
|
||||
//! `$HOME/notes.md`. ACP's `fs/read_text_file` and friends — and our
|
||||
//! own local `std::fs` reads — both want a real absolute path; the
|
||||
//! `~` / `$HOME` forms reach them as literal strings and the open
|
||||
//! fails. The tool schemas already document "absolute path" but in
|
||||
//! practice the model slips up often enough that handling it
|
||||
//! server-side is the difference between "works" and "the agent is
|
||||
//! brittle".
|
||||
//!
|
||||
//! Scope is deliberately small:
|
||||
//!
|
||||
//! - `~` and `~/` (current user only — `~user` lookups would require
|
||||
//! pulling in passwd parsing).
|
||||
//! - `$HOME` and `$HOME/`.
|
||||
//!
|
||||
//! Any other shell variable (`$PWD`, `${HOME}`, …) passes through
|
||||
//! unchanged. The shell already expands them inside `bash` tool
|
||||
//! commands; for the file-tool argument fields, we deliberately
|
||||
//! limit the set so the behaviour is predictable.
|
||||
//!
|
||||
//! Falls back to the input path verbatim when `HOME` is unset
|
||||
//! (stripped-down container env). That preserves the "no surprise
|
||||
//! mutations" rule — never invent a path the caller didn't ask for.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
/// Process-global lock for tests that mutate `HOME`. Anyone in the
|
||||
/// crate touching `HOME` must hold this for the duration of the
|
||||
/// read-modify-restore window — otherwise concurrent `cargo test`
|
||||
/// workers race and flake.
|
||||
///
|
||||
/// Only built into the test binaries. Production code never mutates
|
||||
/// env vars.
|
||||
#[cfg(test)]
|
||||
pub(crate) static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
|
||||
|
||||
/// Expand `~`, `~/`, `$HOME`, and `$HOME/` prefixes against the
|
||||
/// current user's home directory. All other inputs pass through
|
||||
/// unchanged.
|
||||
///
|
||||
/// Returns the input verbatim if `HOME` isn't set in the env.
|
||||
pub fn expand_path(input: &Path) -> PathBuf {
|
||||
let Some(s) = input.to_str() else {
|
||||
return input.to_path_buf();
|
||||
};
|
||||
let Ok(home) = std::env::var("HOME") else {
|
||||
return input.to_path_buf();
|
||||
};
|
||||
let home = PathBuf::from(home);
|
||||
if s == "~" || s == "$HOME" {
|
||||
return home;
|
||||
}
|
||||
if let Some(rest) = s.strip_prefix("~/") {
|
||||
return home.join(rest);
|
||||
}
|
||||
if let Some(rest) = s.strip_prefix("$HOME/") {
|
||||
return home.join(rest);
|
||||
}
|
||||
input.to_path_buf()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// Set HOME for the duration of the test. Tests using this run
|
||||
/// serially under the crate-wide [`ENV_LOCK`] because env
|
||||
/// mutation isn't thread-safe — `cargo test` parallel workers
|
||||
/// would race without it.
|
||||
fn with_home<F: FnOnce()>(home: &str, body: F) {
|
||||
let _g = ENV_LOCK.lock().unwrap();
|
||||
let prior = std::env::var("HOME").ok();
|
||||
// SAFETY: tests touch process-global env. The mutex
|
||||
// serialises access; sub-threads in other test modules
|
||||
// touching HOME aren't expected (none in this crate).
|
||||
unsafe {
|
||||
std::env::set_var("HOME", home);
|
||||
}
|
||||
body();
|
||||
unsafe {
|
||||
match prior {
|
||||
Some(p) => std::env::set_var("HOME", p),
|
||||
None => std::env::remove_var("HOME"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expands_tilde_slash() {
|
||||
with_home("/home/me", || {
|
||||
assert_eq!(
|
||||
expand_path(Path::new("~/git/repo/file.rs")),
|
||||
PathBuf::from("/home/me/git/repo/file.rs")
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expands_bare_tilde() {
|
||||
with_home("/home/me", || {
|
||||
assert_eq!(expand_path(Path::new("~")), PathBuf::from("/home/me"));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expands_dollar_home_slash() {
|
||||
with_home("/home/me", || {
|
||||
assert_eq!(
|
||||
expand_path(Path::new("$HOME/notes.md")),
|
||||
PathBuf::from("/home/me/notes.md")
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expands_bare_dollar_home() {
|
||||
with_home("/home/me", || {
|
||||
assert_eq!(expand_path(Path::new("$HOME")), PathBuf::from("/home/me"));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn absolute_path_passes_through() {
|
||||
with_home("/home/me", || {
|
||||
assert_eq!(
|
||||
expand_path(Path::new("/etc/hostname")),
|
||||
PathBuf::from("/etc/hostname")
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn relative_path_passes_through() {
|
||||
with_home("/home/me", || {
|
||||
assert_eq!(
|
||||
expand_path(Path::new("src/main.rs")),
|
||||
PathBuf::from("src/main.rs")
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tilde_user_form_not_expanded() {
|
||||
// ~other is shell sugar for /home/other and would require
|
||||
// passwd parsing to resolve. Out of scope — pass it
|
||||
// through and let the open fail with a clear error.
|
||||
with_home("/home/me", || {
|
||||
assert_eq!(
|
||||
expand_path(Path::new("~other/x")),
|
||||
PathBuf::from("~other/x")
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_home_env_passes_through() {
|
||||
// Share the same crate-wide lock as `with_home` — otherwise
|
||||
// a parallel test setting HOME races this clear-and-assert
|
||||
// window.
|
||||
let _g = ENV_LOCK.lock().unwrap();
|
||||
let prior = std::env::var("HOME").ok();
|
||||
// SAFETY: serialised by LOCK above.
|
||||
unsafe {
|
||||
std::env::remove_var("HOME");
|
||||
}
|
||||
assert_eq!(
|
||||
expand_path(Path::new("~/git/repo")),
|
||||
PathBuf::from("~/git/repo")
|
||||
);
|
||||
unsafe {
|
||||
if let Some(p) = prior {
|
||||
std::env::set_var("HOME", p);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dollar_other_var_not_expanded() {
|
||||
with_home("/home/me", || {
|
||||
assert_eq!(
|
||||
expand_path(Path::new("$PWD/file")),
|
||||
PathBuf::from("$PWD/file")
|
||||
);
|
||||
assert_eq!(
|
||||
expand_path(Path::new("${HOME}/file")),
|
||||
PathBuf::from("${HOME}/file")
|
||||
);
|
||||
});
|
||||
}
|
||||
}
|
||||
274
crates/helexa-acp/src/prompt.rs
Normal file
274
crates/helexa-acp/src/prompt.rs
Normal file
@@ -0,0 +1,274 @@
|
||||
//! System prompt assembly.
|
||||
//!
|
||||
//! The system message has two parts:
|
||||
//!
|
||||
//! 1. A short human-readable preamble (working directory, style
|
||||
//! instructions). Either the built-in [`DEFAULT_PROMPT`] or a
|
||||
//! user-supplied file at `HELEXA_ACP_SYSTEM_PROMPT_PATH` /
|
||||
//! `system_prompt_path`. `{cwd}` is substituted in both.
|
||||
//! 2. A `# Tools` block in Qwen3 Hermes format (see [`crate::qwen3`])
|
||||
//! describing the available functions. This is what makes the
|
||||
//! model actually call them — neuron/cortex don't honour the
|
||||
//! OpenAI `tools` API field, so the tool list has to live in the
|
||||
//! prompt itself.
|
||||
|
||||
use agent_client_protocol::schema::SessionModeId;
|
||||
use anyhow::Context;
|
||||
use std::path::Path;
|
||||
|
||||
use crate::provider::ToolSpec;
|
||||
use crate::qwen3;
|
||||
use crate::session::MODE_PLAN;
|
||||
|
||||
const DEFAULT_PROMPT: &str = "\
|
||||
You are helexa-acp, a coding assistant working inside an editor.
|
||||
|
||||
Working directory: {cwd}
|
||||
|
||||
Use the tools described below whenever the user's request involves
|
||||
looking at or modifying files, or running commands. Do not ask the
|
||||
user to paste file contents you could read yourself. All file paths
|
||||
must be absolute. Writes and shell commands may prompt the user for
|
||||
permission depending on the session mode.
|
||||
|
||||
Be concise; the user is reading your output in an editor pane.";
|
||||
|
||||
/// Build the system prompt for a session.
|
||||
///
|
||||
/// - `cwd`: session working directory (substituted for `{cwd}` in
|
||||
/// the preamble — both the default and any user-supplied template).
|
||||
/// - `override_path`: path to a user-supplied template, already
|
||||
/// resolved by [`crate::config::Config`]. The `# Tools` block is
|
||||
/// appended *after* the user's template so a custom preamble
|
||||
/// still gets the tool descriptions the model needs.
|
||||
/// - `tools`: the tools to advertise. Empty list → no `# Tools`
|
||||
/// block is appended at all.
|
||||
/// - `mode`: current session mode. When the mode is [`MODE_PLAN`]
|
||||
/// a plan-mode addendum describing the restrictions and the
|
||||
/// completion menu is appended *after* the `# Tools` block so it
|
||||
/// is the last thing the model reads before user input.
|
||||
/// - `plan_dir`: resolved plan directory for the cwd. Only consulted
|
||||
/// when `mode == MODE_PLAN`. `None` means the plan directory could
|
||||
/// not be resolved (no `HOME` / `XDG_DATA_HOME`) — the addendum
|
||||
/// still renders but with a placeholder so the model knows to
|
||||
/// surface the error to the user rather than guess a path.
|
||||
pub fn build_system_prompt(
|
||||
cwd: &Path,
|
||||
override_path: Option<&Path>,
|
||||
tools: &[ToolSpec],
|
||||
mode: &SessionModeId,
|
||||
plan_dir: Option<&Path>,
|
||||
) -> anyhow::Result<String> {
|
||||
let template = match override_path {
|
||||
Some(path) => std::fs::read_to_string(path)
|
||||
.with_context(|| format!("read system prompt from {}", path.display()))?,
|
||||
None => DEFAULT_PROMPT.to_string(),
|
||||
};
|
||||
let mut prompt = template.replace("{cwd}", &cwd.display().to_string());
|
||||
prompt.push_str(&qwen3::render_tool_block(tools));
|
||||
if mode.0.as_ref() == MODE_PLAN {
|
||||
prompt.push_str(&render_plan_mode_block(plan_dir));
|
||||
}
|
||||
Ok(prompt)
|
||||
}
|
||||
|
||||
/// Plan-mode instruction block. Tells the model:
|
||||
///
|
||||
/// 1. Where it may write — only inside `plan_dir`.
|
||||
/// 2. What it may *not* do — bash is disabled; writes outside
|
||||
/// `plan_dir` are refused by the runtime.
|
||||
/// 3. How to finish — emit the 3-option menu so the user can
|
||||
/// switch modes and either kick off implementation (with or
|
||||
/// without permission prompts) or keep iterating on the plan.
|
||||
fn render_plan_mode_block(plan_dir: Option<&Path>) -> String {
|
||||
let plan_path = plan_dir
|
||||
.map(|p| p.display().to_string())
|
||||
.unwrap_or_else(|| "<plan directory could not be resolved — tell the user>".to_string());
|
||||
format!(
|
||||
"\n\n# Plan mode\n\
|
||||
\n\
|
||||
You are in **plan mode**. Your task is to draft a written\n\
|
||||
implementation plan for the user; you must NOT modify any\n\
|
||||
project files or run shell commands.\n\
|
||||
\n\
|
||||
Rules in plan mode:\n\
|
||||
\n\
|
||||
- `read_file` and `list_dir` are unrestricted — use them to\n\
|
||||
explore the codebase as needed.\n\
|
||||
- `write_file` and `edit_file` are allowed ONLY under the\n\
|
||||
plan directory: `{plan_path}`. The runtime will refuse any\n\
|
||||
write outside it.\n\
|
||||
- `bash` is disabled. Do not call it.\n\
|
||||
\n\
|
||||
Write the plan as one or more Markdown files under\n\
|
||||
`{plan_path}`. Use descriptive filenames\n\
|
||||
(`01-overview.md`, `02-data-model.md`, etc.). It is fine to\n\
|
||||
iterate — overwrite the file when you refine a section.\n\
|
||||
\n\
|
||||
When the plan is complete, do NOT begin implementation.\n\
|
||||
Instead, end your turn with this menu, verbatim, so the\n\
|
||||
user can choose how to proceed:\n\
|
||||
\n\
|
||||
---\n\
|
||||
**Plan complete.** To proceed, switch the session mode in\n\
|
||||
the agent dropdown and send a follow-up message:\n\
|
||||
\n\
|
||||
1. **Bypass Permissions** — implement the plan now, skipping\n\
|
||||
per-tool permission prompts.\n\
|
||||
2. **Default** — implement the plan now, prompting before\n\
|
||||
each write or shell command.\n\
|
||||
3. **Plan** (stay here) — refine the plan; reply with the\n\
|
||||
change you want and I will revise it.\n\
|
||||
---\n"
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::session::{MODE_DEFAULT, MODE_PLAN};
|
||||
use std::io::Write;
|
||||
|
||||
fn default_mode() -> SessionModeId {
|
||||
SessionModeId::new(MODE_DEFAULT)
|
||||
}
|
||||
fn plan_mode() -> SessionModeId {
|
||||
SessionModeId::new(MODE_PLAN)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_prompt_substitutes_cwd() {
|
||||
let prompt =
|
||||
build_system_prompt(Path::new("/home/me/proj"), None, &[], &default_mode(), None)
|
||||
.unwrap();
|
||||
assert!(
|
||||
prompt.contains("/home/me/proj"),
|
||||
"cwd not interpolated: {prompt}"
|
||||
);
|
||||
assert!(prompt.contains("helexa-acp"));
|
||||
assert!(
|
||||
!prompt.contains("{cwd}"),
|
||||
"left-over placeholder in default prompt"
|
||||
);
|
||||
// With no tools, the # Tools block is absent.
|
||||
assert!(!prompt.contains("# Tools"));
|
||||
// Default mode does not get the plan-mode addendum.
|
||||
assert!(!prompt.contains("# Plan mode"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tools_are_appended_in_hermes_format() {
|
||||
let spec = ToolSpec {
|
||||
name: "read_file".into(),
|
||||
description: "Read a file.".into(),
|
||||
parameters: serde_json::json!({"type":"object","properties":{}, "required":[]}),
|
||||
};
|
||||
let prompt =
|
||||
build_system_prompt(Path::new("/x"), None, &[spec], &default_mode(), None).unwrap();
|
||||
assert!(prompt.contains("# Tools"));
|
||||
assert!(prompt.contains("<tools>"));
|
||||
assert!(prompt.contains("\"name\":\"read_file\""));
|
||||
assert!(prompt.contains("<tool_call>"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn override_path_is_read_and_templated() {
|
||||
let mut tmp = tempfile_in_target("prompt.txt");
|
||||
tmp.write_all(b"custom prompt for {cwd} only").unwrap();
|
||||
tmp.flush().unwrap();
|
||||
|
||||
let path = tmp.path().to_path_buf();
|
||||
drop(tmp);
|
||||
|
||||
let prompt = build_system_prompt(
|
||||
Path::new("/etc"),
|
||||
Some(path.as_path()),
|
||||
&[],
|
||||
&default_mode(),
|
||||
None,
|
||||
)
|
||||
.expect("read override");
|
||||
assert_eq!(prompt, "custom prompt for /etc only");
|
||||
|
||||
let _ = std::fs::remove_file(&path);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_override_path_errors() {
|
||||
let err = build_system_prompt(
|
||||
Path::new("/tmp"),
|
||||
Some(Path::new("/definitely/not/a/real/path")),
|
||||
&[],
|
||||
&default_mode(),
|
||||
None,
|
||||
)
|
||||
.unwrap_err();
|
||||
assert!(format!("{err:#}").contains("read system prompt"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn plan_mode_addendum_includes_plan_dir_and_menu() {
|
||||
let plan_dir = Path::new("/home/me/.local/share/helexa-acp/plans/proj-deadbeef");
|
||||
let prompt = build_system_prompt(
|
||||
Path::new("/home/me/proj"),
|
||||
None,
|
||||
&[],
|
||||
&plan_mode(),
|
||||
Some(plan_dir),
|
||||
)
|
||||
.unwrap();
|
||||
assert!(prompt.contains("# Plan mode"));
|
||||
assert!(
|
||||
prompt.contains(plan_dir.to_str().unwrap()),
|
||||
"plan dir not interpolated: {prompt}"
|
||||
);
|
||||
// The 3-option menu must be present so the model emits it verbatim.
|
||||
assert!(prompt.contains("Bypass Permissions"));
|
||||
assert!(prompt.contains("**Default**"));
|
||||
assert!(prompt.contains("3. **Plan**"));
|
||||
// Bash disabled instruction must be present.
|
||||
assert!(prompt.contains("`bash` is disabled"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn plan_mode_addendum_handles_unresolved_plan_dir() {
|
||||
let prompt =
|
||||
build_system_prompt(Path::new("/home/me/proj"), None, &[], &plan_mode(), None).unwrap();
|
||||
assert!(prompt.contains("# Plan mode"));
|
||||
assert!(prompt.contains("could not be resolved"));
|
||||
}
|
||||
|
||||
/// Tiny temp-file helper that doesn't pull in the `tempfile` crate.
|
||||
/// Writes under `target/` so it's cleaned up by `cargo clean`.
|
||||
fn tempfile_in_target(name: &str) -> TempHandle {
|
||||
let base = std::env::var("CARGO_TARGET_TMPDIR")
|
||||
.ok()
|
||||
.map(std::path::PathBuf::from)
|
||||
.unwrap_or_else(std::env::temp_dir);
|
||||
let _ = std::fs::create_dir_all(&base);
|
||||
let pid = std::process::id();
|
||||
let path = base.join(format!("helexa-acp-{pid}-{name}"));
|
||||
let file = std::fs::File::create(&path).expect("create temp file");
|
||||
TempHandle { file, path }
|
||||
}
|
||||
|
||||
struct TempHandle {
|
||||
file: std::fs::File,
|
||||
path: std::path::PathBuf,
|
||||
}
|
||||
|
||||
impl TempHandle {
|
||||
fn path(&self) -> &Path {
|
||||
&self.path
|
||||
}
|
||||
}
|
||||
|
||||
impl Write for TempHandle {
|
||||
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
|
||||
self.file.write(buf)
|
||||
}
|
||||
fn flush(&mut self) -> std::io::Result<()> {
|
||||
self.file.flush()
|
||||
}
|
||||
}
|
||||
}
|
||||
1200
crates/helexa-acp/src/provider/anthropic_messages.rs
Normal file
1200
crates/helexa-acp/src/provider/anthropic_messages.rs
Normal file
File diff suppressed because it is too large
Load Diff
230
crates/helexa-acp/src/provider/mod.rs
Normal file
230
crates/helexa-acp/src/provider/mod.rs
Normal file
@@ -0,0 +1,230 @@
|
||||
//! Provider trait — the seam between the ACP-side agent loop and
|
||||
//! whatever wire protocol an endpoint actually speaks.
|
||||
//!
|
||||
//! Every concrete provider (OpenAI chat completions, OpenAI Responses,
|
||||
//! Anthropic /v1/messages, Ollama native, …) implements
|
||||
//! [`Provider`]. The agent constructs a [`CompletionRequest`] using
|
||||
//! provider-agnostic types and consumes a stream of
|
||||
//! [`CompletionEvent`]s — neither end knows which wire format is on
|
||||
//! the other side of the trait.
|
||||
//!
|
||||
//! Day-1 provider: [`openai_chat::OpenAIChatProvider`]. Day-N
|
||||
//! providers slot in without touching `agent.rs`.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use futures::stream::BoxStream;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
pub mod anthropic_messages;
|
||||
pub mod openai_chat;
|
||||
pub mod openai_responses;
|
||||
|
||||
/// Provider-agnostic LLM endpoint. Implementations translate between
|
||||
/// [`CompletionRequest`] / [`CompletionEvent`] and whatever wire
|
||||
/// format their endpoint speaks.
|
||||
#[async_trait]
|
||||
pub trait Provider: Send + Sync {
|
||||
/// Endpoint name as configured by the user (e.g. `"helexa"`,
|
||||
/// `"openrouter"`). Used in logs and in the `endpoint:model`
|
||||
/// selector.
|
||||
fn name(&self) -> &str;
|
||||
|
||||
/// List models available at this endpoint. Used to build the
|
||||
/// model-picker dropdown in editor clients (Stage 4). Should
|
||||
/// return quickly (cache if necessary).
|
||||
#[allow(dead_code)]
|
||||
async fn list_models(&self) -> anyhow::Result<Vec<ModelInfo>>;
|
||||
|
||||
/// Run a chat completion. Returns a stream of provider-agnostic
|
||||
/// events. The stream stops when the upstream finishes, when
|
||||
/// `cancel` is fired, or when the stream is dropped.
|
||||
async fn complete(
|
||||
&self,
|
||||
request: CompletionRequest,
|
||||
cancel: CancellationToken,
|
||||
) -> anyhow::Result<BoxStream<'static, anyhow::Result<CompletionEvent>>>;
|
||||
}
|
||||
|
||||
/// One model exposed by a provider. Constructed by `list_models` —
|
||||
/// Stage 4 is when the agent loop starts consuming it for the
|
||||
/// model-picker dropdown.
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ModelInfo {
|
||||
pub id: String,
|
||||
/// Human-friendly name, if the endpoint exposes one. Otherwise
|
||||
/// `id` is used as the display name.
|
||||
#[serde(default)]
|
||||
pub display_name: Option<String>,
|
||||
}
|
||||
|
||||
/// Inputs to a completion. Provider-agnostic — concrete providers
|
||||
/// translate this into their wire format.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CompletionRequest {
|
||||
/// Endpoint-local model id (without the `endpoint:` prefix).
|
||||
pub model: String,
|
||||
pub messages: Vec<Message>,
|
||||
/// Tools the model is allowed to call. Empty list means no tool
|
||||
/// support advertised.
|
||||
pub tools: Vec<ToolSpec>,
|
||||
pub temperature: Option<f64>,
|
||||
pub top_p: Option<f64>,
|
||||
pub max_tokens: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Message {
|
||||
pub role: Role,
|
||||
pub content: MessageContent,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum Role {
|
||||
System,
|
||||
User,
|
||||
Assistant,
|
||||
/// Tool result message. Provider impls turn this into whatever
|
||||
/// shape the upstream wire format wants (OpenAI uses
|
||||
/// `role: "tool"` + `tool_call_id`; Anthropic uses content blocks).
|
||||
/// Stage 3 (tools) constructs this; Stage 2 never does.
|
||||
Tool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum MessageContent {
|
||||
/// Plain text turn (system / user / assistant). Struct variant
|
||||
/// rather than newtype so the persisted JSON has an explicit
|
||||
/// `text` field — that lets us use internal tagging on the
|
||||
/// enum, which is incompatible with newtype-of-primitive
|
||||
/// variants.
|
||||
Text { text: String },
|
||||
/// Mixed text + image user turn. Stage 5 introduces this when
|
||||
/// Zed sends an `ImageContent` block alongside the user's prompt.
|
||||
/// Providers that don't support vision should down-convert by
|
||||
/// dropping image parts and concatenating text parts.
|
||||
MultiPart { parts: Vec<MessagePart> },
|
||||
/// Assistant turn that called one or more tools. Stage 3 starts
|
||||
/// constructing this when the provider stream yields a
|
||||
/// `ToolCallStart` / `ToolCallArgsDelta` sequence.
|
||||
ToolCalls {
|
||||
/// Optional text the assistant said alongside the tool calls.
|
||||
text: Option<String>,
|
||||
calls: Vec<ToolCall>,
|
||||
},
|
||||
/// Tool result. `tool_call_id` matches the assistant's call id.
|
||||
/// Stage 3 constructs this after the tool runner finishes.
|
||||
ToolResult {
|
||||
tool_call_id: String,
|
||||
content: String,
|
||||
},
|
||||
}
|
||||
|
||||
/// One part of a [`MessageContent::MultiPart`] message.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum MessagePart {
|
||||
Text { text: String },
|
||||
Image(ImageData),
|
||||
}
|
||||
|
||||
/// Inline image attachment. `data` is base64-encoded raw image
|
||||
/// bytes; the encoder constructs an `image_url` data URI from it
|
||||
/// at request time. `uri` carries any pointer the client supplied
|
||||
/// (e.g. `file:///tmp/x.png`) — we keep it on the message for
|
||||
/// debugging / future providers but the OpenAI encoder ignores it
|
||||
/// when `data` is present (data wins, since it round-trips through
|
||||
/// every wire format).
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ImageData {
|
||||
pub mime_type: String,
|
||||
/// Base64-encoded image bytes (no `data:` prefix, no padding
|
||||
/// stripped — exactly what `ImageContent.data` carried).
|
||||
pub data: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub uri: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ToolCall {
|
||||
/// Provider-assigned id that ties the call to its result. The
|
||||
/// Qwen3 wire format we use today doesn't carry this on the
|
||||
/// model side (calls and results are matched positionally inside
|
||||
/// a turn), so the field looks unused in the prod build — but it
|
||||
/// flows through to `MessageContent::ToolResult.tool_call_id` for
|
||||
/// history bookkeeping and a future strict-OpenAI backend will
|
||||
/// consume it directly.
|
||||
#[allow(dead_code)]
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
/// JSON-encoded arguments. Kept as a string because providers
|
||||
/// stream argument bytes incrementally and only validate at the
|
||||
/// end; the agent decodes once the call is complete.
|
||||
pub arguments: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ToolSpec {
|
||||
pub name: String,
|
||||
pub description: String,
|
||||
/// JSON Schema of the arguments object.
|
||||
pub parameters: Value,
|
||||
}
|
||||
|
||||
/// Events emitted by a provider during a streaming completion.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum CompletionEvent {
|
||||
/// Incremental visible text from the assistant.
|
||||
TextDelta(String),
|
||||
/// Incremental "reasoning" / thought text, if the model emits one
|
||||
/// (e.g. Qwen3 with `<think>` tags surfaced as a separate stream,
|
||||
/// or OpenAI reasoning models).
|
||||
ReasoningDelta(String),
|
||||
/// A new tool call has started. Stage 2 ignores the payload; the
|
||||
/// agent loop in Stage 3 reads `index` to correlate with
|
||||
/// [`Self::ToolCallArgsDelta`], `id` for the eventual tool-result
|
||||
/// turn, and `name` to dispatch the runner.
|
||||
#[allow(dead_code)]
|
||||
ToolCallStart {
|
||||
index: usize,
|
||||
id: String,
|
||||
name: String,
|
||||
},
|
||||
/// More argument bytes for a tool call already announced via
|
||||
/// [`Self::ToolCallStart`]. Stage 2 ignores; Stage 3 accumulates
|
||||
/// the bytes by `index` until the call's arguments are complete.
|
||||
#[allow(dead_code)]
|
||||
ToolCallArgsDelta { index: usize, args_delta: String },
|
||||
/// A `<tool_call>` block whose JSON couldn't be parsed even with
|
||||
/// the qwen3 module's repair attempts. The agent surfaces this
|
||||
/// as a Failed `SessionUpdate::ToolCall` card with the raw body
|
||||
/// visible (so the editor renders structured failure UI rather
|
||||
/// than dumping the body inline in the message pane), and feeds
|
||||
/// a synthetic tool-error message back into history so the
|
||||
/// model can self-correct on the next round.
|
||||
MalformedToolCall { raw: String },
|
||||
/// Stream finished. Carries the upstream `finish_reason` if it
|
||||
/// gave one (`"stop"`, `"length"`, `"tool_calls"`, …).
|
||||
Finish { reason: Option<String> },
|
||||
/// Final usage stats, if the provider supplied them. Stage 2
|
||||
/// matches the variant to drop it; Stage 6b (token metrics) is
|
||||
/// when the payload starts being read.
|
||||
#[allow(dead_code)]
|
||||
Usage(UsageStats),
|
||||
}
|
||||
|
||||
/// Token accounting reported by the provider at the end of a stream.
|
||||
/// Stage 2 doesn't surface usage anywhere — the stable `PromptResponse`
|
||||
/// has no usage field, and the unstable variant is gated. Stage 6b
|
||||
/// turns these on with Prometheus metrics.
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
pub struct UsageStats {
|
||||
pub prompt_tokens: u64,
|
||||
pub completion_tokens: u64,
|
||||
pub total_tokens: u64,
|
||||
}
|
||||
1002
crates/helexa-acp/src/provider/openai_chat.rs
Normal file
1002
crates/helexa-acp/src/provider/openai_chat.rs
Normal file
File diff suppressed because it is too large
Load Diff
987
crates/helexa-acp/src/provider/openai_responses.rs
Normal file
987
crates/helexa-acp/src/provider/openai_responses.rs
Normal file
@@ -0,0 +1,987 @@
|
||||
//! OpenAI Responses API (`POST /v1/responses`) provider.
|
||||
//!
|
||||
//! Mirror image of [`super::openai_chat`]: same `Provider` trait
|
||||
//! impl, same back-pressured SSE decoder, but speaking OpenAI's
|
||||
//! newer Responses surface instead of chat completions.
|
||||
//!
|
||||
//! Differences from the chat provider, all contained in this file:
|
||||
//!
|
||||
//! - **Request encoding**: history flattens into an `input` array
|
||||
//! of typed items (`message`, `function_call`, `function_call_output`)
|
||||
//! plus a top-level `instructions` field for the system prompt.
|
||||
//! Multi-part user content stays in the same `[{type:"input_text"},
|
||||
//! {type:"input_image"}]` shape neuron's `request_to_chat` already
|
||||
//! accepts.
|
||||
//! - **Streaming decoder**: events are named (`response.created`,
|
||||
//! `response.output_text.delta`, `response.completed`, …) carried
|
||||
//! on the SSE `event:` line. The chat path's `[DONE]` terminator
|
||||
//! doesn't apply; the stream ends after `response.completed`.
|
||||
//! - **Tool calls** plumb through the `response.output_item.added`
|
||||
//! (item type `function_call`) → `response.function_call_arguments.delta`
|
||||
//! → `response.function_call_arguments.done` event sequence. The
|
||||
//! neuron candle harness doesn't synthesize these yet (tracked as
|
||||
//! issue #6), but the decoder is wired so the day the upstream
|
||||
//! does, downstream `CompletionEvent::ToolCall*` plumbing just
|
||||
//! works.
|
||||
//!
|
||||
//! Tool-name handling: the model knows its tool descriptions via
|
||||
//! the [`crate::qwen3`] system-prompt block exactly the way the chat
|
||||
//! provider does. We don't echo them in the request body because
|
||||
//! neuron currently ignores `tools` on /v1/responses (same as on
|
||||
//! /v1/chat/completions). Once neuron honours request-side tool
|
||||
//! definitions, both providers add them in the same place.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use eventsource_stream::Eventsource;
|
||||
use futures::{Stream, StreamExt, stream::BoxStream};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Value, json};
|
||||
use std::collections::HashMap;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use super::{
|
||||
CompletionEvent, CompletionRequest, Message, MessageContent, MessagePart, ModelInfo, Provider,
|
||||
Role, UsageStats,
|
||||
};
|
||||
use crate::config::EndpointConfig;
|
||||
|
||||
pub struct OpenAIResponsesProvider {
|
||||
endpoint: EndpointConfig,
|
||||
#[allow(dead_code)] // Read in `complete()`'s HTTP path; tests don't stand up a server.
|
||||
api_key: Option<String>,
|
||||
#[allow(dead_code)]
|
||||
http: reqwest::Client,
|
||||
}
|
||||
|
||||
impl OpenAIResponsesProvider {
|
||||
pub fn new(endpoint: EndpointConfig) -> anyhow::Result<Self> {
|
||||
let api_key = endpoint.resolve_api_key()?;
|
||||
let http = reqwest::Client::builder()
|
||||
// Same generous timeout as the chat provider: cortex may
|
||||
// need to cold-load a model before serving the first
|
||||
// chunk, which can be tens of seconds. Cancellation
|
||||
// handles early termination, not timeout.
|
||||
.timeout(std::time::Duration::from_secs(600))
|
||||
.build()?;
|
||||
Ok(Self {
|
||||
endpoint,
|
||||
api_key,
|
||||
http,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for OpenAIResponsesProvider {
|
||||
fn name(&self) -> &str {
|
||||
&self.endpoint.name
|
||||
}
|
||||
|
||||
async fn list_models(&self) -> anyhow::Result<Vec<ModelInfo>> {
|
||||
let mut req = self.http.get(self.endpoint.models_url());
|
||||
if let Some(key) = &self.api_key {
|
||||
req = req.bearer_auth(key);
|
||||
}
|
||||
let resp = req
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{} list_models: {e}", self.endpoint.name))?;
|
||||
let status = resp.status();
|
||||
if !status.is_success() {
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
anyhow::bail!(
|
||||
"{} list_models returned {}: {}",
|
||||
self.endpoint.name,
|
||||
status,
|
||||
body
|
||||
);
|
||||
}
|
||||
let body: WireModelsResponse = resp.json().await?;
|
||||
Ok(body
|
||||
.data
|
||||
.into_iter()
|
||||
.map(|m| ModelInfo {
|
||||
id: m.id,
|
||||
display_name: None,
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn complete(
|
||||
&self,
|
||||
request: CompletionRequest,
|
||||
cancel: CancellationToken,
|
||||
) -> anyhow::Result<BoxStream<'static, anyhow::Result<CompletionEvent>>> {
|
||||
let body = encode_request(&request);
|
||||
tracing::debug!(
|
||||
endpoint = %self.endpoint.name,
|
||||
url = %self.endpoint.responses_url(),
|
||||
body = %serde_json::to_string(&body).unwrap_or_else(|_| "<unserializable>".into()),
|
||||
"POST /responses"
|
||||
);
|
||||
let mut req = self.http.post(self.endpoint.responses_url()).json(&body);
|
||||
if let Some(key) = &self.api_key {
|
||||
req = req.bearer_auth(key);
|
||||
}
|
||||
let resp = req
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{} responses send: {e}", self.endpoint.name))?;
|
||||
let status = resp.status();
|
||||
if !status.is_success() {
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
anyhow::bail!(
|
||||
"{} responses returned {}: {}",
|
||||
self.endpoint.name,
|
||||
status,
|
||||
body
|
||||
);
|
||||
}
|
||||
let sse = resp.bytes_stream().eventsource();
|
||||
let stream = decode_stream(sse, cancel);
|
||||
Ok(Box::pin(stream))
|
||||
}
|
||||
}
|
||||
|
||||
// ── Request encoding ─────────────────────────────────────────────────
|
||||
|
||||
fn encode_request(req: &CompletionRequest) -> Value {
|
||||
// Pull the system messages out of history into a single
|
||||
// `instructions` string — the Responses API expects them there,
|
||||
// not inline as an `input` item. Multiple system messages
|
||||
// concatenate with blank lines so we don't lose ordering.
|
||||
let mut instructions: Vec<String> = Vec::new();
|
||||
let mut input_items: Vec<Value> = Vec::new();
|
||||
for msg in &req.messages {
|
||||
if msg.role == Role::System
|
||||
&& let MessageContent::Text { text } = &msg.content
|
||||
{
|
||||
instructions.push(text.clone());
|
||||
continue;
|
||||
}
|
||||
if let Some(item) = encode_message_as_input_item(msg) {
|
||||
input_items.push(item);
|
||||
}
|
||||
}
|
||||
|
||||
let mut body = json!({
|
||||
"model": req.model,
|
||||
"input": input_items,
|
||||
"stream": true,
|
||||
});
|
||||
if let Value::Object(map) = &mut body {
|
||||
if !instructions.is_empty() {
|
||||
map.insert(
|
||||
"instructions".into(),
|
||||
Value::String(instructions.join("\n\n")),
|
||||
);
|
||||
}
|
||||
if let Some(t) = req.temperature {
|
||||
map.insert("temperature".into(), json!(t));
|
||||
}
|
||||
if let Some(p) = req.top_p {
|
||||
map.insert("top_p".into(), json!(p));
|
||||
}
|
||||
if let Some(m) = req.max_tokens {
|
||||
// Responses calls it `max_output_tokens`; preserve the
|
||||
// semantic (response cap) when we translate.
|
||||
map.insert("max_output_tokens".into(), json!(m));
|
||||
}
|
||||
}
|
||||
body
|
||||
}
|
||||
|
||||
fn encode_message_as_input_item(msg: &Message) -> Option<Value> {
|
||||
match (msg.role, &msg.content) {
|
||||
(Role::System, _) => None, // handled out-of-band as `instructions`
|
||||
(Role::User, MessageContent::Text { text }) => Some(json!({
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
"content": text,
|
||||
})),
|
||||
(Role::User, MessageContent::MultiPart { parts }) => Some(json!({
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
"content": encode_user_parts(parts),
|
||||
})),
|
||||
(Role::Assistant, MessageContent::Text { text }) => Some(json!({
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [{
|
||||
"type": "output_text",
|
||||
"text": text,
|
||||
"annotations": [],
|
||||
}],
|
||||
})),
|
||||
(Role::Assistant, MessageContent::ToolCalls { text, calls }) => {
|
||||
// Assistant turns that called tools become a sequence of
|
||||
// items: an optional `message` (any prose alongside the
|
||||
// call) followed by one `function_call` per call. Mirrors
|
||||
// OpenAI Responses' "each item is one structural slot"
|
||||
// shape.
|
||||
//
|
||||
// We can't return multiple items from one call site, so
|
||||
// we encode this by side-stuffing additional items into a
|
||||
// single composite value and have the caller flatten —
|
||||
// but that complicates the API. Easier: build the array
|
||||
// ourselves in the caller path. For now, emit just the
|
||||
// function_calls (the assistant's prose lives in the next
|
||||
// turn's chat history anyway because the model isn't
|
||||
// looking back at its own previous narration). If the
|
||||
// text is non-empty AND we have calls, we lose the text;
|
||||
// qwen3 rarely emits prose alongside tool calls so this
|
||||
// is a deliberate simplification — revisit if it bites.
|
||||
let _ = text;
|
||||
// Take the first call only for the moment; multi-call
|
||||
// turns would need the caller-flattening above.
|
||||
let call = calls.first()?;
|
||||
Some(json!({
|
||||
"type": "function_call",
|
||||
"call_id": call.id,
|
||||
"name": call.name,
|
||||
"arguments": call.arguments,
|
||||
}))
|
||||
}
|
||||
(
|
||||
Role::Tool,
|
||||
MessageContent::ToolResult {
|
||||
tool_call_id,
|
||||
content,
|
||||
},
|
||||
) => Some(json!({
|
||||
"type": "function_call_output",
|
||||
"call_id": tool_call_id,
|
||||
"output": content,
|
||||
})),
|
||||
(role, content) => {
|
||||
tracing::warn!(
|
||||
?role,
|
||||
?content,
|
||||
"openai_responses: unexpected (role, content) shape"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn encode_user_parts(parts: &[MessagePart]) -> Value {
|
||||
let items: Vec<Value> = parts
|
||||
.iter()
|
||||
.map(|p| match p {
|
||||
MessagePart::Text { text } => json!({"type": "input_text", "text": text}),
|
||||
MessagePart::Image(img) => json!({
|
||||
"type": "input_image",
|
||||
"image_url": format!("data:{};base64,{}", img.mime_type, img.data),
|
||||
}),
|
||||
})
|
||||
.collect();
|
||||
Value::Array(items)
|
||||
}
|
||||
|
||||
// ── Wire types ──────────────────────────────────────────────────────
|
||||
|
||||
#[allow(dead_code)] // fields read only when list_models runs against a real endpoint
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct WireModelsResponse {
|
||||
data: Vec<WireModelObject>,
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct WireModelObject {
|
||||
id: String,
|
||||
}
|
||||
|
||||
// SSE event payload shapes. We only model the fields we care about;
|
||||
// `#[serde(default)]` + `Option` everywhere else lets the upstream
|
||||
// add optional fields without breaking deserialise.
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
struct OutputItemAddedEvent {
|
||||
#[serde(default)]
|
||||
output_index: u32,
|
||||
item: OutputItem,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
enum OutputItem {
|
||||
Message {
|
||||
#[serde(default)]
|
||||
id: Option<String>,
|
||||
},
|
||||
FunctionCall {
|
||||
#[serde(default)]
|
||||
id: Option<String>,
|
||||
#[serde(default)]
|
||||
call_id: Option<String>,
|
||||
#[serde(default)]
|
||||
name: Option<String>,
|
||||
/// Some upstreams populate `arguments` already on the
|
||||
/// `output_item.added` event for a fully-buffered tool call
|
||||
/// (i.e. when the model finalised the call before the SSE
|
||||
/// flush). Capture it so we can emit a single args delta.
|
||||
#[serde(default)]
|
||||
arguments: Option<String>,
|
||||
},
|
||||
/// `reasoning`, `web_search_call`, etc. We capture-and-ignore
|
||||
/// any item we don't model; the decoder still emits the
|
||||
/// outer events correctly.
|
||||
#[serde(other)]
|
||||
Unknown,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
struct OutputTextDeltaEvent {
|
||||
#[serde(default)]
|
||||
item_id: Option<String>,
|
||||
#[serde(default)]
|
||||
output_index: u32,
|
||||
#[serde(default)]
|
||||
delta: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
struct FunctionCallArgumentsDeltaEvent {
|
||||
#[serde(default)]
|
||||
item_id: Option<String>,
|
||||
#[serde(default)]
|
||||
output_index: u32,
|
||||
#[serde(default)]
|
||||
delta: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
struct ResponseCompletedEvent {
|
||||
response: ResponseShell,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
struct ResponseShell {
|
||||
#[serde(default)]
|
||||
status: Option<String>,
|
||||
#[serde(default)]
|
||||
usage: Option<WireUsage>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
struct WireUsage {
|
||||
#[serde(default)]
|
||||
input_tokens: u64,
|
||||
#[serde(default)]
|
||||
output_tokens: u64,
|
||||
#[serde(default)]
|
||||
total_tokens: u64,
|
||||
}
|
||||
|
||||
// ── Streaming decoder ───────────────────────────────────────────────
|
||||
|
||||
/// Translate the named-event Responses SSE into the provider-agnostic
|
||||
/// [`CompletionEvent`] stream the agent loop expects. The decoder
|
||||
/// holds per-stream state — output_index → tool-call-index plus
|
||||
/// the next available tool-call slot — so it can fire
|
||||
/// `ToolCallStart` exactly once per item.
|
||||
fn decode_stream<S>(
|
||||
sse: S,
|
||||
cancel: CancellationToken,
|
||||
) -> impl Stream<Item = anyhow::Result<CompletionEvent>>
|
||||
where
|
||||
S: Stream<
|
||||
Item = Result<
|
||||
eventsource_stream::Event,
|
||||
eventsource_stream::EventStreamError<reqwest::Error>,
|
||||
>,
|
||||
> + Send
|
||||
+ 'static,
|
||||
{
|
||||
async_stream::stream! {
|
||||
let mut sse = Box::pin(sse);
|
||||
// Maps an output_index that's a function_call to the tool-call
|
||||
// slot we hand downstream. Lets us correlate later
|
||||
// `function_call_arguments.delta` events back to the index
|
||||
// we already announced on `output_item.added`.
|
||||
let mut tool_index_by_output: HashMap<u32, usize> = HashMap::new();
|
||||
let mut next_tool_index: usize = 0;
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
biased;
|
||||
_ = cancel.cancelled() => {
|
||||
tracing::debug!("openai_responses: cancellation requested, ending stream");
|
||||
break;
|
||||
}
|
||||
next = sse.next() => {
|
||||
let Some(event) = next else { break };
|
||||
let event = match event {
|
||||
Ok(e) => e,
|
||||
Err(e) => {
|
||||
yield Err(anyhow::anyhow!("SSE transport: {e}"));
|
||||
break;
|
||||
}
|
||||
};
|
||||
// Event name lives on `event.event`; data is JSON.
|
||||
let event_name = event.event.as_str();
|
||||
let data = event.data.as_str();
|
||||
match event_name {
|
||||
"response.output_text.delta" => {
|
||||
match serde_json::from_str::<OutputTextDeltaEvent>(data) {
|
||||
Ok(d) if !d.delta.is_empty() => {
|
||||
yield Ok(CompletionEvent::TextDelta(d.delta));
|
||||
}
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
raw = %data,
|
||||
"openai_responses: failed to parse output_text.delta; skipping"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
"response.output_item.added" => {
|
||||
match serde_json::from_str::<OutputItemAddedEvent>(data) {
|
||||
Ok(ev) => {
|
||||
if let OutputItem::FunctionCall {
|
||||
id,
|
||||
call_id,
|
||||
name,
|
||||
arguments,
|
||||
} = ev.item
|
||||
{
|
||||
let idx = next_tool_index;
|
||||
next_tool_index += 1;
|
||||
tool_index_by_output.insert(ev.output_index, idx);
|
||||
// Prefer the user-facing
|
||||
// `call_id` (what gets paired
|
||||
// with tool results) over the
|
||||
// internal item `id` when
|
||||
// both are present. Falls
|
||||
// back to a synthetic id so
|
||||
// history bookkeeping never
|
||||
// breaks.
|
||||
let final_id = call_id
|
||||
.or(id)
|
||||
.unwrap_or_else(|| format!("call_{idx}"));
|
||||
let final_name = name.unwrap_or_default();
|
||||
yield Ok(CompletionEvent::ToolCallStart {
|
||||
index: idx,
|
||||
id: final_id,
|
||||
name: final_name,
|
||||
});
|
||||
// Some upstreams attach the
|
||||
// fully-buffered arguments on
|
||||
// the `output_item.added`
|
||||
// event itself (rare; happens
|
||||
// when the model finalised
|
||||
// before the SSE flush).
|
||||
// Emit as a single args
|
||||
// delta if present.
|
||||
if let Some(args) = arguments
|
||||
&& !args.is_empty()
|
||||
{
|
||||
yield Ok(CompletionEvent::ToolCallArgsDelta {
|
||||
index: idx,
|
||||
args_delta: args,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
raw = %data,
|
||||
"openai_responses: failed to parse output_item.added; skipping"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
"response.function_call_arguments.delta" => {
|
||||
match serde_json::from_str::<FunctionCallArgumentsDeltaEvent>(data) {
|
||||
Ok(ev) => {
|
||||
let Some(&idx) = tool_index_by_output.get(&ev.output_index)
|
||||
else {
|
||||
// Args delta for an item we
|
||||
// never saw an `output_item.added`
|
||||
// for. Could happen if the
|
||||
// upstream reordered events;
|
||||
// log + skip.
|
||||
tracing::warn!(
|
||||
output_index = ev.output_index,
|
||||
"openai_responses: function_call_arguments.delta for unknown output_index"
|
||||
);
|
||||
continue;
|
||||
};
|
||||
if !ev.delta.is_empty() {
|
||||
yield Ok(CompletionEvent::ToolCallArgsDelta {
|
||||
index: idx,
|
||||
args_delta: ev.delta,
|
||||
});
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
raw = %data,
|
||||
"openai_responses: failed to parse function_call_arguments.delta; skipping"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
"response.completed" => {
|
||||
// Final event. Pull usage + status off
|
||||
// the response shell. Status maps:
|
||||
// "completed" → no special handling
|
||||
// (caller treats as EndTurn),
|
||||
// "incomplete" → length stop.
|
||||
let (reason, usage) =
|
||||
match serde_json::from_str::<ResponseCompletedEvent>(data) {
|
||||
Ok(ev) => {
|
||||
let reason = match ev.response.status.as_deref() {
|
||||
Some("incomplete") => Some("length".to_string()),
|
||||
_ => Some("stop".to_string()),
|
||||
};
|
||||
let usage = ev.response.usage.map(|u| UsageStats {
|
||||
prompt_tokens: u.input_tokens,
|
||||
completion_tokens: u.output_tokens,
|
||||
total_tokens: u.total_tokens,
|
||||
});
|
||||
(reason, usage)
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
raw = %data,
|
||||
"openai_responses: failed to parse response.completed; ending stream with EndTurn"
|
||||
);
|
||||
(Some("stop".to_string()), None)
|
||||
}
|
||||
};
|
||||
if let Some(u) = usage {
|
||||
yield Ok(CompletionEvent::Usage(u));
|
||||
}
|
||||
yield Ok(CompletionEvent::Finish { reason });
|
||||
break;
|
||||
}
|
||||
// Bookkeeping events we don't need to surface:
|
||||
// response.created, response.in_progress,
|
||||
// response.content_part.added/.done,
|
||||
// response.output_text.done,
|
||||
// response.output_item.done,
|
||||
// response.function_call_arguments.done,
|
||||
// response.reasoning_*. Logged at debug for
|
||||
// wire-tracing.
|
||||
other => {
|
||||
tracing::trace!(
|
||||
event = other,
|
||||
"openai_responses: bookkeeping event"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::provider::ToolCall;
|
||||
use crate::provider::{ImageData, MessagePart};
|
||||
use futures::stream;
|
||||
use url::Url;
|
||||
|
||||
fn ep() -> EndpointConfig {
|
||||
EndpointConfig {
|
||||
name: "test".into(),
|
||||
base_url: Url::parse("http://localhost:9999/v1").unwrap(),
|
||||
wire_api: crate::config::WireApi::OpenAiResponses,
|
||||
default_model: None,
|
||||
api_key: None,
|
||||
api_key_env: None,
|
||||
max_tokens: None,
|
||||
context_window: None,
|
||||
}
|
||||
}
|
||||
|
||||
// ── encode_request ──────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn system_messages_collapse_to_instructions() {
|
||||
let req = CompletionRequest {
|
||||
model: "m".into(),
|
||||
messages: vec![
|
||||
Message {
|
||||
role: Role::System,
|
||||
content: MessageContent::Text {
|
||||
text: "you are helpful".into(),
|
||||
},
|
||||
},
|
||||
Message {
|
||||
role: Role::User,
|
||||
content: MessageContent::Text { text: "hi".into() },
|
||||
},
|
||||
],
|
||||
tools: vec![],
|
||||
temperature: Some(0.7),
|
||||
top_p: None,
|
||||
max_tokens: Some(256),
|
||||
};
|
||||
let body = encode_request(&req);
|
||||
assert_eq!(body["model"], "m");
|
||||
assert_eq!(body["instructions"], "you are helpful");
|
||||
assert_eq!(body["stream"], true);
|
||||
assert_eq!(body["max_output_tokens"], 256);
|
||||
assert_eq!(body["temperature"], 0.7);
|
||||
let input = body["input"].as_array().unwrap();
|
||||
// System message NOT echoed in input — it's only in
|
||||
// instructions.
|
||||
assert_eq!(input.len(), 1);
|
||||
assert_eq!(input[0]["type"], "message");
|
||||
assert_eq!(input[0]["role"], "user");
|
||||
assert_eq!(input[0]["content"], "hi");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multiple_system_messages_concatenate() {
|
||||
let req = CompletionRequest {
|
||||
model: "m".into(),
|
||||
messages: vec![
|
||||
Message {
|
||||
role: Role::System,
|
||||
content: MessageContent::Text {
|
||||
text: "first".into(),
|
||||
},
|
||||
},
|
||||
Message {
|
||||
role: Role::System,
|
||||
content: MessageContent::Text {
|
||||
text: "second".into(),
|
||||
},
|
||||
},
|
||||
Message {
|
||||
role: Role::User,
|
||||
content: MessageContent::Text { text: "hi".into() },
|
||||
},
|
||||
],
|
||||
tools: vec![],
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
max_tokens: None,
|
||||
};
|
||||
let body = encode_request(&req);
|
||||
assert_eq!(body["instructions"], "first\n\nsecond");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_multipart_becomes_input_parts_array() {
|
||||
let req = CompletionRequest {
|
||||
model: "vl".into(),
|
||||
messages: vec![Message {
|
||||
role: Role::User,
|
||||
content: MessageContent::MultiPart {
|
||||
parts: vec![
|
||||
MessagePart::Text {
|
||||
text: "what's in this?".into(),
|
||||
},
|
||||
MessagePart::Image(ImageData {
|
||||
mime_type: "image/png".into(),
|
||||
data: "AAA=".into(),
|
||||
uri: None,
|
||||
}),
|
||||
],
|
||||
},
|
||||
}],
|
||||
tools: vec![],
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
max_tokens: None,
|
||||
};
|
||||
let body = encode_request(&req);
|
||||
let content = &body["input"][0]["content"].as_array().unwrap().clone();
|
||||
assert_eq!(content.len(), 2);
|
||||
assert_eq!(content[0]["type"], "input_text");
|
||||
assert_eq!(content[0]["text"], "what's in this?");
|
||||
assert_eq!(content[1]["type"], "input_image");
|
||||
assert_eq!(content[1]["image_url"], "data:image/png;base64,AAA=");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn assistant_text_becomes_output_text_content_part() {
|
||||
let req = CompletionRequest {
|
||||
model: "m".into(),
|
||||
messages: vec![
|
||||
Message {
|
||||
role: Role::User,
|
||||
content: MessageContent::Text { text: "hi".into() },
|
||||
},
|
||||
Message {
|
||||
role: Role::Assistant,
|
||||
content: MessageContent::Text {
|
||||
text: "hello there".into(),
|
||||
},
|
||||
},
|
||||
Message {
|
||||
role: Role::User,
|
||||
content: MessageContent::Text {
|
||||
text: "more".into(),
|
||||
},
|
||||
},
|
||||
],
|
||||
tools: vec![],
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
max_tokens: None,
|
||||
};
|
||||
let body = encode_request(&req);
|
||||
let input = body["input"].as_array().unwrap();
|
||||
assert_eq!(input.len(), 3);
|
||||
assert_eq!(input[1]["type"], "message");
|
||||
assert_eq!(input[1]["role"], "assistant");
|
||||
assert_eq!(input[1]["content"][0]["type"], "output_text");
|
||||
assert_eq!(input[1]["content"][0]["text"], "hello there");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_calls_and_results_round_trip_via_function_call_items() {
|
||||
let req = CompletionRequest {
|
||||
model: "m".into(),
|
||||
messages: vec![
|
||||
Message {
|
||||
role: Role::Assistant,
|
||||
content: MessageContent::ToolCalls {
|
||||
text: None,
|
||||
calls: vec![ToolCall {
|
||||
id: "call_42".into(),
|
||||
name: "read_file".into(),
|
||||
arguments: r#"{"path":"/etc/hostname"}"#.into(),
|
||||
}],
|
||||
},
|
||||
},
|
||||
Message {
|
||||
role: Role::Tool,
|
||||
content: MessageContent::ToolResult {
|
||||
tool_call_id: "call_42".into(),
|
||||
content: "host".into(),
|
||||
},
|
||||
},
|
||||
],
|
||||
tools: vec![],
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
max_tokens: None,
|
||||
};
|
||||
let body = encode_request(&req);
|
||||
let input = body["input"].as_array().unwrap();
|
||||
assert_eq!(input.len(), 2);
|
||||
assert_eq!(input[0]["type"], "function_call");
|
||||
assert_eq!(input[0]["call_id"], "call_42");
|
||||
assert_eq!(input[0]["name"], "read_file");
|
||||
assert_eq!(input[0]["arguments"], r#"{"path":"/etc/hostname"}"#);
|
||||
assert_eq!(input[1]["type"], "function_call_output");
|
||||
assert_eq!(input[1]["call_id"], "call_42");
|
||||
assert_eq!(input[1]["output"], "host");
|
||||
}
|
||||
|
||||
// ── decode_stream ───────────────────────────────────────────────
|
||||
|
||||
fn sse_event(name: &str, data: &str) -> eventsource_stream::Event {
|
||||
eventsource_stream::Event {
|
||||
id: String::new(),
|
||||
retry: None,
|
||||
event: name.into(),
|
||||
data: data.into(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn collect_events(
|
||||
items: Vec<eventsource_stream::Event>,
|
||||
) -> Vec<anyhow::Result<CompletionEvent>> {
|
||||
let sse = stream::iter(
|
||||
items
|
||||
.into_iter()
|
||||
.map(Ok::<_, eventsource_stream::EventStreamError<reqwest::Error>>),
|
||||
);
|
||||
let decoded = decode_stream(sse, CancellationToken::new());
|
||||
decoded.collect().await
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn decodes_text_then_finish() {
|
||||
let events = collect_events(vec![
|
||||
sse_event("response.created", "{}"),
|
||||
sse_event(
|
||||
"response.output_text.delta",
|
||||
r#"{"item_id":"msg_1","output_index":0,"delta":"hel"}"#,
|
||||
),
|
||||
sse_event(
|
||||
"response.output_text.delta",
|
||||
r#"{"item_id":"msg_1","output_index":0,"delta":"lo"}"#,
|
||||
),
|
||||
sse_event(
|
||||
"response.completed",
|
||||
r#"{"response":{"status":"completed","usage":{"input_tokens":3,"output_tokens":2,"total_tokens":5}}}"#,
|
||||
),
|
||||
])
|
||||
.await;
|
||||
let events: Vec<CompletionEvent> = events.into_iter().map(|r| r.unwrap()).collect();
|
||||
let mut iter = events.into_iter();
|
||||
assert!(matches!(iter.next(), Some(CompletionEvent::TextDelta(t)) if t == "hel"));
|
||||
assert!(matches!(iter.next(), Some(CompletionEvent::TextDelta(t)) if t == "lo"));
|
||||
assert!(matches!(iter.next(), Some(CompletionEvent::Usage(u)) if u.total_tokens == 5));
|
||||
assert!(matches!(
|
||||
iter.next(),
|
||||
Some(CompletionEvent::Finish { reason: Some(r) }) if r == "stop"
|
||||
));
|
||||
assert!(iter.next().is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn empty_delta_is_dropped() {
|
||||
let events = collect_events(vec![
|
||||
sse_event(
|
||||
"response.output_text.delta",
|
||||
r#"{"item_id":"m","output_index":0,"delta":""}"#,
|
||||
),
|
||||
sse_event(
|
||||
"response.completed",
|
||||
r#"{"response":{"status":"completed"}}"#,
|
||||
),
|
||||
])
|
||||
.await;
|
||||
let mut completion_events = events.into_iter().map(|r| r.unwrap());
|
||||
// First event MUST be the Finish — the empty delta dropped.
|
||||
assert!(matches!(
|
||||
completion_events.next(),
|
||||
Some(CompletionEvent::Finish { .. })
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn incomplete_status_maps_to_length_finish_reason() {
|
||||
let events = collect_events(vec![sse_event(
|
||||
"response.completed",
|
||||
r#"{"response":{"status":"incomplete"}}"#,
|
||||
)])
|
||||
.await;
|
||||
let events: Vec<CompletionEvent> = events.into_iter().map(|r| r.unwrap()).collect();
|
||||
assert!(matches!(
|
||||
events.last(),
|
||||
Some(CompletionEvent::Finish { reason: Some(r) }) if r == "length"
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn function_call_items_emit_toolcall_events() {
|
||||
let events = collect_events(vec![
|
||||
sse_event(
|
||||
"response.output_item.added",
|
||||
r#"{"output_index":0,"item":{"type":"function_call","id":"item_1","call_id":"call_xyz","name":"read_file"}}"#,
|
||||
),
|
||||
sse_event(
|
||||
"response.function_call_arguments.delta",
|
||||
r#"{"item_id":"item_1","output_index":0,"delta":"{\"path"}"#,
|
||||
),
|
||||
sse_event(
|
||||
"response.function_call_arguments.delta",
|
||||
r#"{"item_id":"item_1","output_index":0,"delta":"\":\"/etc/hostname\"}"}"#,
|
||||
),
|
||||
sse_event("response.completed", r#"{"response":{"status":"completed"}}"#),
|
||||
])
|
||||
.await;
|
||||
let events: Vec<CompletionEvent> = events.into_iter().map(|r| r.unwrap()).collect();
|
||||
let mut iter = events.into_iter();
|
||||
assert!(matches!(
|
||||
iter.next(),
|
||||
Some(CompletionEvent::ToolCallStart { index: 0, ref id, ref name })
|
||||
if id == "call_xyz" && name == "read_file"
|
||||
));
|
||||
assert!(matches!(
|
||||
iter.next(),
|
||||
Some(CompletionEvent::ToolCallArgsDelta { index: 0, ref args_delta })
|
||||
if args_delta == r#"{"path"#
|
||||
));
|
||||
assert!(matches!(
|
||||
iter.next(),
|
||||
Some(CompletionEvent::ToolCallArgsDelta { index: 0, ref args_delta })
|
||||
if args_delta == r#"":"/etc/hostname"}"#
|
||||
));
|
||||
assert!(matches!(iter.next(), Some(CompletionEvent::Finish { .. })));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn function_call_added_with_inline_arguments_emits_single_args_delta() {
|
||||
// Some upstreams (rare) include the fully-buffered arguments
|
||||
// on the `output_item.added` event when the model finalised
|
||||
// the call before SSE flush. Verify both ToolCallStart and a
|
||||
// single args delta fire.
|
||||
let events = collect_events(vec![
|
||||
sse_event(
|
||||
"response.output_item.added",
|
||||
r#"{"output_index":0,"item":{"type":"function_call","call_id":"call_a","name":"f","arguments":"{\"x\":1}"}}"#,
|
||||
),
|
||||
sse_event("response.completed", r#"{"response":{"status":"completed"}}"#),
|
||||
])
|
||||
.await;
|
||||
let events: Vec<CompletionEvent> = events.into_iter().map(|r| r.unwrap()).collect();
|
||||
let mut iter = events.into_iter();
|
||||
assert!(matches!(
|
||||
iter.next(),
|
||||
Some(CompletionEvent::ToolCallStart { .. })
|
||||
));
|
||||
assert!(matches!(
|
||||
iter.next(),
|
||||
Some(CompletionEvent::ToolCallArgsDelta { index: 0, ref args_delta })
|
||||
if args_delta == r#"{"x":1}"#
|
||||
));
|
||||
assert!(matches!(iter.next(), Some(CompletionEvent::Finish { .. })));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cancellation_ends_stream_promptly() {
|
||||
// Hand the decoder an empty stream + a triggered cancellation
|
||||
// token; it should terminate without yielding anything.
|
||||
let sse = stream::iter(Vec::<
|
||||
Result<eventsource_stream::Event, eventsource_stream::EventStreamError<reqwest::Error>>,
|
||||
>::new());
|
||||
let cancel = CancellationToken::new();
|
||||
cancel.cancel();
|
||||
let decoded = decode_stream(sse, cancel);
|
||||
let events: Vec<_> = decoded.collect().await;
|
||||
assert!(events.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn malformed_event_payload_is_skipped() {
|
||||
let events = collect_events(vec![
|
||||
sse_event("response.output_text.delta", "{not valid json"),
|
||||
sse_event(
|
||||
"response.output_text.delta",
|
||||
r#"{"item_id":"m","output_index":0,"delta":"ok"}"#,
|
||||
),
|
||||
sse_event(
|
||||
"response.completed",
|
||||
r#"{"response":{"status":"completed"}}"#,
|
||||
),
|
||||
])
|
||||
.await;
|
||||
let events: Vec<CompletionEvent> = events.into_iter().map(|r| r.unwrap()).collect();
|
||||
// First text delta dropped; second one fires.
|
||||
assert!(
|
||||
events
|
||||
.iter()
|
||||
.any(|e| matches!(e, CompletionEvent::TextDelta(t) if t == "ok"))
|
||||
);
|
||||
// No errors yielded (parse failures are warn-and-skip).
|
||||
assert!(
|
||||
events
|
||||
.iter()
|
||||
.all(|e| !matches!(e, CompletionEvent::Finish { reason: None }))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_construction_is_cheap() {
|
||||
let _ = OpenAIResponsesProvider::new(ep()).unwrap();
|
||||
}
|
||||
}
|
||||
1018
crates/helexa-acp/src/qwen3.rs
Normal file
1018
crates/helexa-acp/src/qwen3.rs
Normal file
File diff suppressed because it is too large
Load Diff
188
crates/helexa-acp/src/session.rs
Normal file
188
crates/helexa-acp/src/session.rs
Normal file
@@ -0,0 +1,188 @@
|
||||
//! Per-session state for the ACP agent loop.
|
||||
//!
|
||||
//! Concurrency:
|
||||
//!
|
||||
//! - [`SessionStore`] is an `Arc<RwLock<HashMap<SessionId, …>>>`. The map
|
||||
//! itself is read-mostly: it changes only on `session/new` and never
|
||||
//! shrinks during Stage 2, so an `RwLock` keeps concurrent reads
|
||||
//! contention-free.
|
||||
//! - Each session is wrapped in its own `Arc<Mutex<SessionState>>`. Holding
|
||||
//! one session's lock doesn't block requests against any other session,
|
||||
//! which matters once a client opens multiple sessions in parallel.
|
||||
//!
|
||||
//! All operations hold a lock only long enough to copy out (or mutate) the
|
||||
//! state they need — never across an `await` that drives the upstream
|
||||
//! provider stream.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
use agent_client_protocol::schema::{SessionId, SessionModeId};
|
||||
use tokio::sync::{Mutex, RwLock};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use crate::provider::Message;
|
||||
|
||||
/// Mode id advertised as the gated default. Writes / bash prompt for
|
||||
/// permission via `session/request_permission`.
|
||||
pub const MODE_DEFAULT: &str = "default";
|
||||
|
||||
/// Mode id advertised as "auto-allow everything". Matches the
|
||||
/// favorite name (`bypassPermissions`) Zed clients tend to reference.
|
||||
pub const MODE_BYPASS: &str = "bypassPermissions";
|
||||
|
||||
/// Mode id for read-and-plan-only operation. The model may read files
|
||||
/// and list directories freely, may write *only* into the per-project
|
||||
/// plan directory under `$XDG_DATA_HOME/helexa-acp/plans/<project-id>/`,
|
||||
/// and cannot run shell commands. Designed for "draft the
|
||||
/// implementation plan, then I'll review and let you execute" flows.
|
||||
pub const MODE_PLAN: &str = "plan";
|
||||
|
||||
/// State carried for a single ACP session.
|
||||
///
|
||||
/// Mutated under `Mutex<SessionState>`; never share a clone across
|
||||
/// tasks expecting to see the same `cancel` token — clone the token
|
||||
/// explicitly when handing it to the streaming task.
|
||||
#[derive(Debug)]
|
||||
pub struct SessionState {
|
||||
/// Conversation history in chronological order (user / assistant
|
||||
/// turns). The system prompt is *not* stored here — it's built
|
||||
/// fresh per request so any cwd / config changes take effect.
|
||||
pub history: Vec<Message>,
|
||||
/// Working directory the client opened the session against. Used
|
||||
/// by [`crate::prompt::build_system_prompt`] and (Stage 3) by
|
||||
/// filesystem tools.
|
||||
pub cwd: PathBuf,
|
||||
/// Currently-selected model id. Format is either a bare model id
|
||||
/// (resolved against the default endpoint) or `endpoint:model`.
|
||||
/// Mutated by `session/set_model` in Stage 4; Stage 2 sets it
|
||||
/// once at session creation and never changes it.
|
||||
pub model_id: String,
|
||||
/// Cancellation handle for the in-flight prompt, if any. A fresh
|
||||
/// token is installed at the start of every `session/prompt`
|
||||
/// request; `session/cancel` fires this one. Between prompts the
|
||||
/// token is "spent" — firing it does nothing — which is fine,
|
||||
/// `session/cancel` is a no-op when there's nothing to cancel.
|
||||
pub cancel: CancellationToken,
|
||||
/// Permission gating mode. Stage 3 advertises two ids in
|
||||
/// `NewSessionResponse.modes`: [`MODE_DEFAULT`] (writes / bash
|
||||
/// prompt the user) and [`MODE_BYPASS`] (auto-allow). Mutated by
|
||||
/// `session/set_mode`.
|
||||
pub mode_id: SessionModeId,
|
||||
}
|
||||
|
||||
impl SessionState {
|
||||
pub fn new(cwd: PathBuf, model_id: String) -> Self {
|
||||
Self {
|
||||
history: Vec::new(),
|
||||
cwd,
|
||||
model_id,
|
||||
cancel: CancellationToken::new(),
|
||||
mode_id: SessionModeId::new(MODE_DEFAULT),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Concurrent map of live sessions.
|
||||
///
|
||||
/// Cloning is cheap (`Arc` bump). Pass clones into every handler that
|
||||
/// needs session access; never hold a clone across an `.await` that
|
||||
/// could outlive the request.
|
||||
pub type SessionStore = Arc<RwLock<HashMap<SessionId, Arc<Mutex<SessionState>>>>>;
|
||||
|
||||
/// Fresh, empty session store.
|
||||
pub fn new_store() -> SessionStore {
|
||||
Arc::new(RwLock::new(HashMap::new()))
|
||||
}
|
||||
|
||||
/// Look up a session by id. Returns `None` if no such session is registered.
|
||||
pub async fn get(store: &SessionStore, id: &SessionId) -> Option<Arc<Mutex<SessionState>>> {
|
||||
store.read().await.get(id).cloned()
|
||||
}
|
||||
|
||||
/// Register a fresh session. Overwrites any prior entry with the same id
|
||||
/// (which should never happen — ids are uniquely generated by the agent).
|
||||
pub async fn insert(store: &SessionStore, id: SessionId, state: SessionState) {
|
||||
store.write().await.insert(id, Arc::new(Mutex::new(state)));
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::provider::{MessageContent, Role};
|
||||
|
||||
fn id(s: &str) -> SessionId {
|
||||
SessionId::new(s)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn insert_then_get_round_trip() {
|
||||
let store = new_store();
|
||||
let state = SessionState::new(PathBuf::from("/tmp"), "m".into());
|
||||
insert(&store, id("s1"), state).await;
|
||||
let got = get(&store, &id("s1")).await.expect("session present");
|
||||
let locked = got.lock().await;
|
||||
assert_eq!(locked.cwd, PathBuf::from("/tmp"));
|
||||
assert_eq!(locked.model_id, "m");
|
||||
assert!(locked.history.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn missing_session_is_none() {
|
||||
let store = new_store();
|
||||
assert!(get(&store, &id("nope")).await.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn history_is_per_session() {
|
||||
let store = new_store();
|
||||
insert(
|
||||
&store,
|
||||
id("a"),
|
||||
SessionState::new(PathBuf::from("/a"), "m".into()),
|
||||
)
|
||||
.await;
|
||||
insert(
|
||||
&store,
|
||||
id("b"),
|
||||
SessionState::new(PathBuf::from("/b"), "m".into()),
|
||||
)
|
||||
.await;
|
||||
|
||||
// Appending to a's history must not affect b's.
|
||||
get(&store, &id("a"))
|
||||
.await
|
||||
.unwrap()
|
||||
.lock()
|
||||
.await
|
||||
.history
|
||||
.push(Message {
|
||||
role: Role::User,
|
||||
content: MessageContent::Text {
|
||||
text: "hello".into(),
|
||||
},
|
||||
});
|
||||
|
||||
assert_eq!(
|
||||
get(&store, &id("a"))
|
||||
.await
|
||||
.unwrap()
|
||||
.lock()
|
||||
.await
|
||||
.history
|
||||
.len(),
|
||||
1
|
||||
);
|
||||
assert_eq!(
|
||||
get(&store, &id("b"))
|
||||
.await
|
||||
.unwrap()
|
||||
.lock()
|
||||
.await
|
||||
.history
|
||||
.len(),
|
||||
0
|
||||
);
|
||||
}
|
||||
}
|
||||
462
crates/helexa-acp/src/store.rs
Normal file
462
crates/helexa-acp/src/store.rs
Normal file
@@ -0,0 +1,462 @@
|
||||
//! On-disk session persistence for `session/load` support.
|
||||
//!
|
||||
//! Storage layout:
|
||||
//!
|
||||
//! ```text
|
||||
//! $XDG_DATA_HOME/helexa-acp/sessions/{session_id}.json
|
||||
//! ```
|
||||
//!
|
||||
//! (Fallback to `~/.local/share/helexa-acp/sessions/` when
|
||||
//! `$XDG_DATA_HOME` is unset.) One JSON file per session. Writes
|
||||
//! happen at the end of every `session/prompt` round through
|
||||
//! [`save`], using tempfile-plus-rename so a crash mid-write can't
|
||||
//! corrupt the store. Reads happen on `session/load` via [`load`].
|
||||
//!
|
||||
//! No compaction, no rotation: files accumulate until the user
|
||||
//! cleans them up. That's deliberate — disk is cheap, and the
|
||||
//! resume-on-restart workflow matters more than tidiness. The
|
||||
//! [`SESSIONS_DIRNAME`] subdirectory is created lazily on first
|
||||
//! save so an unprivileged install path never errors at startup.
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::time::SystemTime;
|
||||
|
||||
use agent_client_protocol::schema::SessionId;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::provider::Message;
|
||||
|
||||
const APP_DIRNAME: &str = "helexa-acp";
|
||||
const SESSIONS_DIRNAME: &str = "sessions";
|
||||
const PLANS_DIRNAME: &str = "plans";
|
||||
|
||||
/// The shape persisted to disk for one session. Only what we can't
|
||||
/// rebuild from the running config goes in here: the conversation
|
||||
/// history, the mode toggle, the model id, and the cwd-at-creation.
|
||||
///
|
||||
/// `created_at` / `updated_at` are seconds-since-epoch — cheap to
|
||||
/// compare, no third-party time crate, and stable across runs.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PersistedSession {
|
||||
pub session_id: String,
|
||||
pub cwd: PathBuf,
|
||||
pub model_id: String,
|
||||
pub mode_id: String,
|
||||
pub history: Vec<Message>,
|
||||
pub created_at: u64,
|
||||
pub updated_at: u64,
|
||||
}
|
||||
|
||||
/// Resolve the directory that holds session JSON files. Honors
|
||||
/// `$XDG_DATA_HOME`; falls back to `~/.local/share/helexa-acp/sessions/`.
|
||||
/// Returns `None` if neither is resolvable (no `HOME` set — possible
|
||||
/// in stripped-down container environments).
|
||||
pub fn sessions_dir() -> Option<PathBuf> {
|
||||
let base = std::env::var("XDG_DATA_HOME")
|
||||
.ok()
|
||||
.filter(|s| !s.is_empty())
|
||||
.map(PathBuf::from)
|
||||
.or_else(|| {
|
||||
std::env::var("HOME")
|
||||
.ok()
|
||||
.map(|h| PathBuf::from(h).join(".local").join("share"))
|
||||
})?;
|
||||
Some(base.join(APP_DIRNAME).join(SESSIONS_DIRNAME))
|
||||
}
|
||||
|
||||
/// Atomic save into the default sessions directory.
|
||||
pub fn save(session: &PersistedSession) -> anyhow::Result<()> {
|
||||
let dir = sessions_dir()
|
||||
.ok_or_else(|| anyhow::anyhow!("can't resolve XDG_DATA_HOME or HOME for session store"))?;
|
||||
save_to_dir(&dir, session)
|
||||
}
|
||||
|
||||
/// Load from the default sessions directory.
|
||||
pub fn load(session_id: &SessionId) -> anyhow::Result<PersistedSession> {
|
||||
let dir = sessions_dir()
|
||||
.ok_or_else(|| anyhow::anyhow!("can't resolve XDG_DATA_HOME or HOME for session store"))?;
|
||||
load_from_dir(&dir, session_id)
|
||||
}
|
||||
|
||||
/// Atomic save into an explicit directory. Writes to
|
||||
/// `{id}.json.tmp` then renames over `{id}.json`. Creates the
|
||||
/// target directory if it doesn't exist. Split from [`save`] so
|
||||
/// unit tests can target a per-test scratch dir without mutating
|
||||
/// process-global env vars.
|
||||
pub fn save_to_dir(dir: &std::path::Path, session: &PersistedSession) -> anyhow::Result<()> {
|
||||
std::fs::create_dir_all(dir).map_err(|e| anyhow::anyhow!("create {}: {e}", dir.display()))?;
|
||||
let safe = sanitize_id(&session.session_id);
|
||||
let final_path = dir.join(format!("{safe}.json"));
|
||||
let tmp_path = dir.join(format!("{safe}.json.tmp"));
|
||||
let json = serde_json::to_string_pretty(session)?;
|
||||
std::fs::write(&tmp_path, json)
|
||||
.map_err(|e| anyhow::anyhow!("write {}: {e}", tmp_path.display()))?;
|
||||
std::fs::rename(&tmp_path, &final_path)
|
||||
.map_err(|e| anyhow::anyhow!("rename → {}: {e}", final_path.display()))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Load from an explicit directory. Returns a friendly error
|
||||
/// message when the session id has no file on disk so the caller
|
||||
/// can map it to a clean ACP error response.
|
||||
pub fn load_from_dir(
|
||||
dir: &std::path::Path,
|
||||
session_id: &SessionId,
|
||||
) -> anyhow::Result<PersistedSession> {
|
||||
let safe = sanitize_id(session_id.0.as_ref());
|
||||
let path = dir.join(format!("{safe}.json"));
|
||||
let bytes = std::fs::read(&path).map_err(|e| {
|
||||
if e.kind() == std::io::ErrorKind::NotFound {
|
||||
anyhow::anyhow!("no persisted session at {}", path.display())
|
||||
} else {
|
||||
anyhow::anyhow!("read {}: {e}", path.display())
|
||||
}
|
||||
})?;
|
||||
let session: PersistedSession = serde_json::from_slice(&bytes)
|
||||
.map_err(|e| anyhow::anyhow!("parse {}: {e}", path.display()))?;
|
||||
Ok(session)
|
||||
}
|
||||
|
||||
/// List all persisted sessions, optionally filtered by `cwd`. Used
|
||||
/// by the `session/list` handler so a client (Zed) can find the
|
||||
/// session that belongs to the workspace it's reopening.
|
||||
///
|
||||
/// `filter_cwd = None` returns every session on disk. `Some(path)`
|
||||
/// returns only sessions whose persisted `cwd` is exactly equal.
|
||||
///
|
||||
/// Files that fail to parse are skipped with a warning rather than
|
||||
/// aborting the whole list — one corrupt session shouldn't make
|
||||
/// the resume picker unusable.
|
||||
pub fn list(filter_cwd: Option<&std::path::Path>) -> anyhow::Result<Vec<PersistedSession>> {
|
||||
let dir = sessions_dir()
|
||||
.ok_or_else(|| anyhow::anyhow!("can't resolve XDG_DATA_HOME or HOME for session store"))?;
|
||||
list_in_dir(&dir, filter_cwd)
|
||||
}
|
||||
|
||||
/// Explicit-dir variant for tests, mirroring [`save_to_dir`] /
|
||||
/// [`load_from_dir`].
|
||||
pub fn list_in_dir(
|
||||
dir: &std::path::Path,
|
||||
filter_cwd: Option<&std::path::Path>,
|
||||
) -> anyhow::Result<Vec<PersistedSession>> {
|
||||
let read = match std::fs::read_dir(dir) {
|
||||
Ok(r) => r,
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
|
||||
Err(e) => return Err(anyhow::anyhow!("read_dir {}: {e}", dir.display())),
|
||||
};
|
||||
let mut out = Vec::new();
|
||||
for entry in read.flatten() {
|
||||
let path = entry.path();
|
||||
if path.extension().and_then(|s| s.to_str()) != Some("json") {
|
||||
continue;
|
||||
}
|
||||
match std::fs::read(&path).and_then(|bytes| {
|
||||
serde_json::from_slice::<PersistedSession>(&bytes).map_err(std::io::Error::other)
|
||||
}) {
|
||||
Ok(session) => {
|
||||
if let Some(want) = filter_cwd
|
||||
&& session.cwd != want
|
||||
{
|
||||
continue;
|
||||
}
|
||||
out.push(session);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
path = %path.display(),
|
||||
error = %e,
|
||||
"store: skipping unparseable session file"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
// Most-recent first by updated_at.
|
||||
out.sort_by_key(|s| std::cmp::Reverse(s.updated_at));
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Seconds-since-epoch, saturating to 0 if the system clock is
|
||||
/// behind epoch (which shouldn't happen but the type system
|
||||
/// requires a fallible read).
|
||||
pub fn now_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(SystemTime::UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
/// Root directory for plan-mode artefacts. Mirrors [`sessions_dir`]
|
||||
/// but under `…/helexa-acp/plans/` so plans and conversation
|
||||
/// transcripts are siblings, not nested.
|
||||
pub fn plans_root() -> Option<PathBuf> {
|
||||
sessions_dir().and_then(|s| s.parent().map(|p| p.join(PLANS_DIRNAME)))
|
||||
}
|
||||
|
||||
/// Per-project plan directory:
|
||||
/// `$XDG_DATA_HOME/helexa-acp/plans/<project-id>/`. The id derives
|
||||
/// from the session's cwd so plans for the same project survive
|
||||
/// across cwd-changes (a `/home/foo/git/bar` ↔ symlinked
|
||||
/// `/srv/checkout/bar` would technically diverge, accepted as a
|
||||
/// won't-fix corner case).
|
||||
pub fn plan_dir_for(cwd: &std::path::Path) -> Option<PathBuf> {
|
||||
plans_root().map(|root| root.join(project_id_for(cwd)))
|
||||
}
|
||||
|
||||
/// Deterministic, human-readable project identifier. Format:
|
||||
/// `<basename>-<8-hex>` where the 8-hex suffix is FNV-1a of the
|
||||
/// full path. Basename keeps the path skim-readable when poking
|
||||
/// around `$XDG_DATA_HOME` by hand; the hash suffix disambiguates
|
||||
/// repos that share a final path component (e.g. multiple
|
||||
/// `/.../checkout/beat` checkouts).
|
||||
///
|
||||
/// FNV-1a rather than `std::collections::hash::DefaultHasher`
|
||||
/// because the latter (SipHash) reseeds per process, so it'd give
|
||||
/// us a different project_id on every run.
|
||||
pub fn project_id_for(cwd: &std::path::Path) -> String {
|
||||
let basename = cwd
|
||||
.file_name()
|
||||
.and_then(|s| s.to_str())
|
||||
.unwrap_or("unknown");
|
||||
let sanitised: String = basename
|
||||
.chars()
|
||||
.map(|c| {
|
||||
if c.is_ascii_alphanumeric() || c == '-' || c == '_' {
|
||||
c
|
||||
} else {
|
||||
'_'
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
let hash = fnv1a_32(cwd.to_string_lossy().as_bytes());
|
||||
format!("{sanitised}-{hash:08x}")
|
||||
}
|
||||
|
||||
/// FNV-1a (32-bit). Deterministic, no third-party crate. Used for
|
||||
/// project ids only — not cryptographic.
|
||||
fn fnv1a_32(bytes: &[u8]) -> u32 {
|
||||
let mut h: u32 = 0x811c_9dc5;
|
||||
for b in bytes {
|
||||
h ^= u32::from(*b);
|
||||
h = h.wrapping_mul(0x0100_0193);
|
||||
}
|
||||
h
|
||||
}
|
||||
|
||||
/// Format seconds-since-epoch as an ISO 8601 / RFC 3339 string
|
||||
/// (`YYYY-MM-DDTHH:MM:SSZ`) for `SessionInfo.updated_at`. Returns
|
||||
/// `None` for values outside the representable range, in which
|
||||
/// case the caller should omit the field.
|
||||
pub fn unix_to_iso8601(secs: u64) -> Option<String> {
|
||||
use chrono::TimeZone;
|
||||
let dt = chrono::Utc.timestamp_opt(secs as i64, 0).single()?;
|
||||
Some(dt.to_rfc3339_opts(chrono::SecondsFormat::Secs, true))
|
||||
}
|
||||
|
||||
/// Strip anything that isn't a safe filename character so a
|
||||
/// mischievous (or just unconventional) session id can't escape
|
||||
/// the sessions directory.
|
||||
fn sanitize_id(id: &str) -> String {
|
||||
id.chars()
|
||||
.map(|c| {
|
||||
if c.is_ascii_alphanumeric() || c == '-' || c == '_' {
|
||||
c
|
||||
} else {
|
||||
'_'
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::provider::{MessageContent, Role};
|
||||
|
||||
/// Unique scratch dir per test invocation. We use this dir
|
||||
/// directly with the `*_to_dir` / `*_from_dir` functions so
|
||||
/// the tests never mutate `$XDG_DATA_HOME` — that env var
|
||||
/// would race across the parallel test harness.
|
||||
fn unique_dir() -> PathBuf {
|
||||
let base = std::env::var("CARGO_TARGET_TMPDIR")
|
||||
.ok()
|
||||
.map(PathBuf::from)
|
||||
.unwrap_or_else(std::env::temp_dir);
|
||||
let pid = std::process::id();
|
||||
let nanos = SystemTime::now()
|
||||
.duration_since(SystemTime::UNIX_EPOCH)
|
||||
.map(|d| d.subsec_nanos())
|
||||
.unwrap_or(0);
|
||||
let dir = base.join(format!("helexa-acp-store-test-{pid}-{nanos}"));
|
||||
std::fs::create_dir_all(&dir).expect("create test dir");
|
||||
dir
|
||||
}
|
||||
|
||||
fn sample(id: &str) -> PersistedSession {
|
||||
PersistedSession {
|
||||
session_id: id.into(),
|
||||
cwd: PathBuf::from("/home/me/proj"),
|
||||
model_id: "Qwen/Qwen3.6-27B".into(),
|
||||
mode_id: "default".into(),
|
||||
history: vec![
|
||||
Message {
|
||||
role: Role::User,
|
||||
content: MessageContent::Text {
|
||||
text: "hello".into(),
|
||||
},
|
||||
},
|
||||
Message {
|
||||
role: Role::Assistant,
|
||||
content: MessageContent::Text { text: "hi".into() },
|
||||
},
|
||||
],
|
||||
created_at: 1_700_000_000,
|
||||
updated_at: 1_700_000_001,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn round_trip_save_then_load() {
|
||||
let dir = unique_dir();
|
||||
save_to_dir(&dir, &sample("hxa-1")).expect("save");
|
||||
let loaded = load_from_dir(&dir, &SessionId::new("hxa-1")).expect("load");
|
||||
assert_eq!(loaded.session_id, "hxa-1");
|
||||
assert_eq!(loaded.cwd, PathBuf::from("/home/me/proj"));
|
||||
assert_eq!(loaded.history.len(), 2);
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_missing_session_errors_with_not_found_message() {
|
||||
let dir = unique_dir();
|
||||
let err = load_from_dir(&dir, &SessionId::new("nope")).unwrap_err();
|
||||
let msg = format!("{err}");
|
||||
assert!(
|
||||
msg.contains("no persisted session"),
|
||||
"want NotFound, got: {msg}"
|
||||
);
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn save_overwrites_existing_atomically() {
|
||||
let dir = unique_dir();
|
||||
save_to_dir(&dir, &sample("hxa-1")).expect("save");
|
||||
let mut updated = sample("hxa-1");
|
||||
updated.history.push(Message {
|
||||
role: Role::User,
|
||||
content: MessageContent::Text {
|
||||
text: "third turn".into(),
|
||||
},
|
||||
});
|
||||
updated.updated_at = 1_700_000_500;
|
||||
save_to_dir(&dir, &updated).expect("re-save");
|
||||
let loaded = load_from_dir(&dir, &SessionId::new("hxa-1")).expect("load");
|
||||
assert_eq!(loaded.history.len(), 3);
|
||||
assert_eq!(loaded.updated_at, 1_700_000_500);
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn save_then_load_preserves_tool_calls_and_results() {
|
||||
use crate::provider::ToolCall;
|
||||
let dir = unique_dir();
|
||||
let mut session = sample("hxa-2");
|
||||
session.history.push(Message {
|
||||
role: Role::Assistant,
|
||||
content: MessageContent::ToolCalls {
|
||||
text: Some("calling".into()),
|
||||
calls: vec![ToolCall {
|
||||
id: "call_0".into(),
|
||||
name: "read_file".into(),
|
||||
arguments: r#"{"path":"/etc/hostname"}"#.into(),
|
||||
}],
|
||||
},
|
||||
});
|
||||
session.history.push(Message {
|
||||
role: Role::Tool,
|
||||
content: MessageContent::ToolResult {
|
||||
tool_call_id: "call_0".into(),
|
||||
content: "host".into(),
|
||||
},
|
||||
});
|
||||
save_to_dir(&dir, &session).expect("save");
|
||||
let loaded = load_from_dir(&dir, &SessionId::new("hxa-2")).expect("load");
|
||||
assert_eq!(loaded.history.len(), 4);
|
||||
match &loaded.history[2].content {
|
||||
MessageContent::ToolCalls { calls, .. } => {
|
||||
assert_eq!(calls[0].name, "read_file");
|
||||
}
|
||||
other => panic!("expected ToolCalls, got {other:?}"),
|
||||
}
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_filters_by_cwd_and_sorts_recent_first() {
|
||||
let dir = unique_dir();
|
||||
let mut a = sample("a");
|
||||
a.cwd = PathBuf::from("/home/me/proj-x");
|
||||
a.updated_at = 1_700_000_010;
|
||||
let mut b = sample("b");
|
||||
b.cwd = PathBuf::from("/home/me/proj-x");
|
||||
b.updated_at = 1_700_000_020;
|
||||
let mut c = sample("c");
|
||||
c.cwd = PathBuf::from("/home/me/elsewhere");
|
||||
c.updated_at = 1_700_000_030;
|
||||
save_to_dir(&dir, &a).unwrap();
|
||||
save_to_dir(&dir, &b).unwrap();
|
||||
save_to_dir(&dir, &c).unwrap();
|
||||
|
||||
let proj_x = PathBuf::from("/home/me/proj-x");
|
||||
let list = list_in_dir(&dir, Some(&proj_x)).unwrap();
|
||||
let ids: Vec<&str> = list.iter().map(|s| s.session_id.as_str()).collect();
|
||||
// Filtered to proj-x; b before a because b is more recent.
|
||||
assert_eq!(ids, vec!["b", "a"]);
|
||||
|
||||
let all = list_in_dir(&dir, None).unwrap();
|
||||
assert_eq!(all.len(), 3);
|
||||
// Global list still sorted recent-first across all cwds.
|
||||
assert_eq!(all[0].session_id, "c");
|
||||
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_returns_empty_for_missing_dir() {
|
||||
let dir = unique_dir().join("does-not-exist");
|
||||
let list = list_in_dir(&dir, None).unwrap();
|
||||
assert!(list.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_skips_unparseable_files() {
|
||||
let dir = unique_dir();
|
||||
save_to_dir(&dir, &sample("good")).unwrap();
|
||||
std::fs::write(dir.join("garbage.json"), b"{not valid json").unwrap();
|
||||
let list = list_in_dir(&dir, None).unwrap();
|
||||
// Garbage skipped; good survives.
|
||||
assert_eq!(list.len(), 1);
|
||||
assert_eq!(list[0].session_id, "good");
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn iso8601_formats_unix_seconds() {
|
||||
// 2024-01-01T00:00:00Z is 1704067200 unix seconds.
|
||||
assert_eq!(
|
||||
unix_to_iso8601(1_704_067_200),
|
||||
Some("2024-01-01T00:00:00Z".into())
|
||||
);
|
||||
assert_eq!(unix_to_iso8601(0), Some("1970-01-01T00:00:00Z".into()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sanitize_id_rejects_path_traversal() {
|
||||
// `../../etc/passwd` — 6 non-alnum chars before "etc"
|
||||
// (`.`, `.`, `/`, `.`, `.`, `/`), one between, none
|
||||
// after, none before nothing. Every disallowed char
|
||||
// collapses to `_`.
|
||||
assert_eq!(sanitize_id("../../etc/passwd"), "______etc_passwd");
|
||||
assert_eq!(sanitize_id("ok-name_42"), "ok-name_42");
|
||||
}
|
||||
}
|
||||
1469
crates/helexa-acp/src/tool_runner.rs
Normal file
1469
crates/helexa-acp/src/tool_runner.rs
Normal file
File diff suppressed because it is too large
Load Diff
300
crates/helexa-acp/src/tools.rs
Normal file
300
crates/helexa-acp/src/tools.rs
Normal file
@@ -0,0 +1,300 @@
|
||||
//! Tool schemas sent to the upstream model on every completion.
|
||||
//!
|
||||
//! These are the OpenAI-function-style declarations the LLM sees in
|
||||
//! `CompletionRequest.tools`; the runtime dispatch happens in
|
||||
//! [`crate::tool_runner`]. Keeping declarations and execution in
|
||||
//! separate modules makes it easy to add a tool without touching the
|
||||
//! runner, and vice versa.
|
||||
//!
|
||||
//! Stage 3 ships five: filesystem read / write / edit, directory
|
||||
//! listing, and `bash`. Image generation, web fetch, MCP-derived
|
||||
//! tools, etc. are out of scope here.
|
||||
|
||||
use serde_json::json;
|
||||
|
||||
use crate::provider::ToolSpec;
|
||||
|
||||
pub const READ_FILE: &str = "read_file";
|
||||
pub const WRITE_FILE: &str = "write_file";
|
||||
pub const EDIT_FILE: &str = "edit_file";
|
||||
pub const LIST_DIR: &str = "list_dir";
|
||||
pub const BASH: &str = "bash";
|
||||
|
||||
/// Build the static tool list passed to the model on every prompt.
|
||||
/// Cheap — the JSON Schema fragments are constructed each call but
|
||||
/// the bodies are small constants. If this ever shows up in a
|
||||
/// profile we can `OnceLock` the Vec.
|
||||
pub fn all_tools() -> Vec<ToolSpec> {
|
||||
vec![
|
||||
ToolSpec {
|
||||
name: READ_FILE.to_string(),
|
||||
description: "Read the contents of a text file. Returns the file's text.".to_string(),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {
|
||||
"type": "string",
|
||||
"description": "Absolute path to the file."
|
||||
},
|
||||
"line": {
|
||||
"type": "integer",
|
||||
"description": "Optional 1-based line number to start reading from.",
|
||||
"minimum": 1
|
||||
},
|
||||
"limit": {
|
||||
"type": "integer",
|
||||
"description": "Optional maximum number of lines to read.",
|
||||
"minimum": 1
|
||||
}
|
||||
},
|
||||
"required": ["path"],
|
||||
"additionalProperties": false
|
||||
}),
|
||||
},
|
||||
ToolSpec {
|
||||
name: WRITE_FILE.to_string(),
|
||||
description: "Write text content to a file, replacing any existing contents. \
|
||||
Creates the file (and parent directories) if needed."
|
||||
.to_string(),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {
|
||||
"type": "string",
|
||||
"description": "Absolute path to the file."
|
||||
},
|
||||
"content": {
|
||||
"type": "string",
|
||||
"description": "Full new contents of the file."
|
||||
}
|
||||
},
|
||||
"required": ["path", "content"],
|
||||
"additionalProperties": false
|
||||
}),
|
||||
},
|
||||
ToolSpec {
|
||||
name: EDIT_FILE.to_string(),
|
||||
description: "Replace one exact substring in a file with another. \
|
||||
Fails if `old_text` does not appear in the file, or appears more than once. \
|
||||
Use multiple edit_file calls for multiple edits."
|
||||
.to_string(),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {
|
||||
"type": "string",
|
||||
"description": "Absolute path to the file."
|
||||
},
|
||||
"old_text": {
|
||||
"type": "string",
|
||||
"description": "Exact text fragment to replace. Must be unique within the file."
|
||||
},
|
||||
"new_text": {
|
||||
"type": "string",
|
||||
"description": "Replacement text."
|
||||
}
|
||||
},
|
||||
"required": ["path", "old_text", "new_text"],
|
||||
"additionalProperties": false
|
||||
}),
|
||||
},
|
||||
ToolSpec {
|
||||
name: LIST_DIR.to_string(),
|
||||
description:
|
||||
"List the entries of a directory. Returns names and a (f|d|l) kind per entry."
|
||||
.to_string(),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {
|
||||
"type": "string",
|
||||
"description": "Absolute path to the directory."
|
||||
}
|
||||
},
|
||||
"required": ["path"],
|
||||
"additionalProperties": false
|
||||
}),
|
||||
},
|
||||
ToolSpec {
|
||||
name: BASH.to_string(),
|
||||
description: "Run a shell command via `sh -c`. \
|
||||
Returns combined stdout+stderr and the exit status. \
|
||||
The command runs in the session's working directory unless `cwd` is given."
|
||||
.to_string(),
|
||||
parameters: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"command": {
|
||||
"type": "string",
|
||||
"description": "Shell command line, evaluated by `sh -c`."
|
||||
},
|
||||
"cwd": {
|
||||
"type": "string",
|
||||
"description": "Optional absolute path to run the command from."
|
||||
}
|
||||
},
|
||||
"required": ["command"],
|
||||
"additionalProperties": false
|
||||
}),
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
/// Try to infer which tool was intended from the shape of an
|
||||
/// `arguments` object alone. Used by the agent when the model
|
||||
/// emits a `<tool_call>` whose JSON has the right arguments but a
|
||||
/// missing or invalid top-level `name` field — a recurring
|
||||
/// Qwen3.6-27B failure mode.
|
||||
///
|
||||
/// Returns `Some(name)` only when the argument keys uniquely match
|
||||
/// exactly one tool in the catalogue. Ambiguous shapes (`{path}`
|
||||
/// alone could be either [`READ_FILE`] or [`LIST_DIR`]) return
|
||||
/// `None` so the caller surfaces a Failed-card and lets the model
|
||||
/// retry rather than guessing wrong.
|
||||
///
|
||||
/// Inference table (key set → tool):
|
||||
///
|
||||
/// | Keys | Tool |
|
||||
/// |---------------------------------------|--------------|
|
||||
/// | `{command}` or `{command, cwd}` | `bash` |
|
||||
/// | `{path, content}` | `write_file` |
|
||||
/// | `{path, old_text, new_text}` | `edit_file` |
|
||||
/// | `{path}` / `{path, line}` / `{path, line, limit}` | *ambiguous* — None |
|
||||
/// | (anything else) | None |
|
||||
pub fn infer_tool_name(arguments: &serde_json::Value) -> Option<&'static str> {
|
||||
let obj = arguments.as_object()?;
|
||||
let keys: std::collections::HashSet<&str> = obj.keys().map(|s| s.as_str()).collect();
|
||||
|
||||
// `command` is unique to bash. Allow the optional `cwd` arg
|
||||
// alongside but nothing else (any unrecognised keys → bail and
|
||||
// let the model retry rather than misroute).
|
||||
if keys.contains("command") && keys.iter().all(|k| matches!(*k, "command" | "cwd")) {
|
||||
return Some(BASH);
|
||||
}
|
||||
// `content` is unique to write_file.
|
||||
if keys.contains("content") && keys.contains("path") && keys.len() == 2 {
|
||||
return Some(WRITE_FILE);
|
||||
}
|
||||
// `old_text` + `new_text` are unique to edit_file.
|
||||
if keys.contains("old_text")
|
||||
&& keys.contains("new_text")
|
||||
&& keys.contains("path")
|
||||
&& keys.len() == 3
|
||||
{
|
||||
return Some(EDIT_FILE);
|
||||
}
|
||||
// `{path}` / `{path, line}` / `{path, line, limit}` overlap
|
||||
// between read_file (file contents) and list_dir (directory
|
||||
// contents). No safe inference — refuse.
|
||||
None
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn all_tools_has_five_named_entries() {
|
||||
let tools = all_tools();
|
||||
let names: Vec<&str> = tools.iter().map(|t| t.name.as_str()).collect();
|
||||
assert_eq!(
|
||||
names,
|
||||
vec![READ_FILE, WRITE_FILE, EDIT_FILE, LIST_DIR, BASH]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn infer_bash_from_command_only() {
|
||||
let args = serde_json::json!({"command": "ls /tmp"});
|
||||
assert_eq!(infer_tool_name(&args), Some(BASH));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn infer_bash_from_command_and_cwd() {
|
||||
let args = serde_json::json!({"command": "ls", "cwd": "/tmp"});
|
||||
assert_eq!(infer_tool_name(&args), Some(BASH));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn infer_bash_from_mkdir_like_real_failure() {
|
||||
// Lifted verbatim from the agent failure that motivated
|
||||
// this helper (helexa-acp.log @ 10:03:11).
|
||||
let args = serde_json::json!({
|
||||
"command": "mkdir -p /home/grenade/git/beat/beat/doc/plan/{01-discovery,02-segmentation,03-description,04-summary,05-output}"
|
||||
});
|
||||
assert_eq!(infer_tool_name(&args), Some(BASH));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn infer_write_file() {
|
||||
let args = serde_json::json!({"path": "/tmp/x", "content": "hi"});
|
||||
assert_eq!(infer_tool_name(&args), Some(WRITE_FILE));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn infer_edit_file() {
|
||||
let args = serde_json::json!({
|
||||
"path": "/tmp/x", "old_text": "a", "new_text": "b"
|
||||
});
|
||||
assert_eq!(infer_tool_name(&args), Some(EDIT_FILE));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn refuse_ambiguous_path_only() {
|
||||
let args = serde_json::json!({"path": "/tmp/x"});
|
||||
assert_eq!(infer_tool_name(&args), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn refuse_ambiguous_path_with_optionals() {
|
||||
// read_file accepts these optionals; list_dir doesn't —
|
||||
// but Qwen wouldn't reliably emit them either, so we
|
||||
// can't use their presence to disambiguate. Refuse.
|
||||
let args = serde_json::json!({"path": "/tmp/x", "line": 1, "limit": 50});
|
||||
assert_eq!(infer_tool_name(&args), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn refuse_command_with_extra_unknown_keys() {
|
||||
// Defence in depth: an unrecognised key alongside
|
||||
// `command` means we don't really know what tool the
|
||||
// model wanted; refuse rather than guess.
|
||||
let args = serde_json::json!({"command": "ls", "extra": "?"});
|
||||
assert_eq!(infer_tool_name(&args), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn refuse_empty_args() {
|
||||
let args = serde_json::json!({});
|
||||
assert_eq!(infer_tool_name(&args), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn refuse_non_object_args() {
|
||||
let args = serde_json::json!("not an object");
|
||||
assert_eq!(infer_tool_name(&args), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn every_tool_has_an_object_parameter_schema() {
|
||||
for tool in all_tools() {
|
||||
let ty = tool.parameters.get("type").and_then(|v| v.as_str());
|
||||
assert_eq!(
|
||||
ty,
|
||||
Some("object"),
|
||||
"tool {} parameters.type must be \"object\"",
|
||||
tool.name
|
||||
);
|
||||
assert!(
|
||||
tool.parameters.get("properties").is_some(),
|
||||
"tool {} missing properties",
|
||||
tool.name
|
||||
);
|
||||
assert!(
|
||||
tool.parameters.get("required").is_some(),
|
||||
"tool {} missing required list",
|
||||
tool.name
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -76,6 +76,15 @@ cudarc = { version = "0.19", optional = true, default-features = false, features
|
||||
half = { version = "2.5", optional = true }
|
||||
tokenizers = { version = "0.22", default-features = false, features = ["onig"] }
|
||||
hf-hub = { version = "0.4", features = ["tokio"] }
|
||||
# Jinja-compatible template renderer for the model's
|
||||
# `tokenizer_config.json::chat_template`. Hugging Face's chat
|
||||
# templates use a strict subset of Jinja2 that minijinja supports
|
||||
# out of the box. ~80KB compiled; pure Rust, no async surface.
|
||||
# Features: `builtins` for the `is defined` / `default` filters HF
|
||||
# templates use; `json` for `tojson` (some Qwen3 templates emit
|
||||
# tool definitions via tojson); `serde` so we can hand it a
|
||||
# serde_json::Value as the context.
|
||||
minijinja = { version = "2", features = ["builtins", "json", "serde"] }
|
||||
# Direct dep on `safetensors` (re-exported by candle but its `TensorView`
|
||||
# / `slice::IndexOp` types are public-but-not-re-exported). Used by the
|
||||
# tp `fused_load` module to read per-rank slices of fused QKV tensors
|
||||
@@ -85,6 +94,7 @@ safetensors = "0.7"
|
||||
[dev-dependencies]
|
||||
tokio = { workspace = true, features = ["test-util"] }
|
||||
reqwest.workspace = true
|
||||
tempfile = "3"
|
||||
|
||||
[build-dependencies]
|
||||
# Used by `build.rs` to compile `src/cuda/*.cu` into `libneuroncuda.a`
|
||||
|
||||
@@ -3,7 +3,9 @@
|
||||
use crate::activation::ActivationTracker;
|
||||
use crate::harness::HarnessRegistry;
|
||||
use crate::harness::candle::{CandleHarness, InferenceError};
|
||||
use crate::harness::preflight::PreflightError;
|
||||
use crate::health::HealthCache;
|
||||
use crate::wire::{openai_chat, openai_responses};
|
||||
use axum::Router;
|
||||
use axum::extract::{Path, State};
|
||||
use axum::http::StatusCode;
|
||||
@@ -12,11 +14,13 @@ use axum::response::{IntoResponse, Json};
|
||||
use axum::routing::{get, post};
|
||||
use cortex_core::discovery::{DiscoveryResponse, HealthResponse};
|
||||
use cortex_core::harness::ModelSpec;
|
||||
use cortex_core::openai::ChatCompletionRequest;
|
||||
use cortex_core::openai::{ChatCompletionRequest, MessageContent};
|
||||
use cortex_core::responses::{ResponsesRequest, ResponsesUsage};
|
||||
use futures::stream::{self, StreamExt};
|
||||
use serde_json::{Value, json};
|
||||
use std::convert::Infallible;
|
||||
use std::sync::Arc;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
use tokio::sync::RwLock;
|
||||
use tokio_stream::wrappers::ReceiverStream;
|
||||
|
||||
@@ -44,6 +48,7 @@ pub fn neuron_routes() -> Router<Arc<NeuronState>> {
|
||||
.route("/models/unload", post(unload_model))
|
||||
.route("/models/{model_id}/endpoint", get(model_endpoint))
|
||||
.route("/v1/chat/completions", post(chat_completions))
|
||||
.route("/v1/responses", post(responses))
|
||||
}
|
||||
|
||||
async fn discovery_handler(State(state): State<Arc<NeuronState>>) -> Json<DiscoveryResponse> {
|
||||
@@ -80,6 +85,24 @@ async fn load_model(
|
||||
match registry.load_model(&spec).await {
|
||||
Ok(()) => Json(json!({"status": "loaded"})).into_response(),
|
||||
Err(e) => {
|
||||
// If the underlying failure is a structured preflight
|
||||
// rejection, surface it as 422 Unprocessable Entity with
|
||||
// the typed JSON body. The kind/model_id/suggestion/etc.
|
||||
// fields let cortex (and operators reading the response
|
||||
// directly) act on the failure without parsing free text.
|
||||
if let Some(pf) = e.downcast_ref::<PreflightError>() {
|
||||
tracing::warn!(
|
||||
model = %spec.model_id,
|
||||
reason = preflight_kind(pf),
|
||||
detail = %pf,
|
||||
"load_model rejected by preflight"
|
||||
);
|
||||
return (
|
||||
StatusCode::UNPROCESSABLE_ENTITY,
|
||||
Json(json!({ "error": pf })),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
// Log the full anyhow chain server-side so journalctl shows
|
||||
// the underlying failure (hf-hub timeout, permission denied,
|
||||
// disk full, etc.) without needing to inspect the HTTP
|
||||
@@ -98,6 +121,18 @@ async fn load_model(
|
||||
}
|
||||
}
|
||||
|
||||
/// Short kebab-case tag for a preflight failure, used as a structured
|
||||
/// log field for journalctl-side filtering. Mirrors the same helper in
|
||||
/// `startup.rs`; duplicated to keep the module surfaces independent.
|
||||
fn preflight_kind(err: &PreflightError) -> &'static str {
|
||||
match err {
|
||||
PreflightError::RepoFetchFailed { .. } => "repo_fetch_failed",
|
||||
PreflightError::EmptyRepo { .. } => "empty_repo",
|
||||
PreflightError::TpRequiresSafetensors { .. } => "tp_requires_safetensors",
|
||||
PreflightError::QuantNotFound { .. } => "quant_not_found",
|
||||
}
|
||||
}
|
||||
|
||||
async fn unload_model(
|
||||
State(state): State<Arc<NeuronState>>,
|
||||
Json(body): Json<Value>,
|
||||
@@ -144,6 +179,7 @@ async fn model_endpoint(
|
||||
/// `ChatCompletionResponse`.
|
||||
async fn chat_completions(
|
||||
State(state): State<Arc<NeuronState>>,
|
||||
headers: axum::http::HeaderMap,
|
||||
Json(req): Json<ChatCompletionRequest>,
|
||||
) -> impl IntoResponse {
|
||||
let Some(candle) = state.candle.as_ref().map(Arc::clone) else {
|
||||
@@ -154,8 +190,23 @@ async fn chat_completions(
|
||||
.into_response();
|
||||
};
|
||||
|
||||
// Reasoning-content opt-in. Off by default → naïve clients
|
||||
// (Zed's commit-message generator, vanilla OpenAI clients)
|
||||
// never see `<think>` blocks. On when the caller sends
|
||||
// `x-include-thinking: true` (helexa-acp does this so its
|
||||
// own ThinkParser keeps working unchanged).
|
||||
let include_thinking = headers
|
||||
.get("x-include-thinking")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(|s| matches!(s.trim().to_ascii_lowercase().as_str(), "1" | "true" | "yes"))
|
||||
.unwrap_or(false);
|
||||
let chat_config = openai_chat::ChatProjectionConfig {
|
||||
include_thinking,
|
||||
reasoning_markers: None, // filled in from the loaded model inside candle
|
||||
};
|
||||
|
||||
if req.stream.unwrap_or(false) {
|
||||
match candle.chat_completion_stream(req).await {
|
||||
match candle.chat_completion_stream_with(req, chat_config).await {
|
||||
Ok(rx) => {
|
||||
// Each chunk → one SSE `data: {json}` line. After the
|
||||
// channel closes, append the OpenAI [DONE] terminator.
|
||||
@@ -246,3 +297,187 @@ async fn chat_completions(
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// OpenAI Responses API (`POST /v1/responses`). Translates the
|
||||
/// Responses-shaped request into a chat-completions one the candle
|
||||
/// harness already understands, then re-projects the harness's
|
||||
/// event stream into the Responses event family.
|
||||
async fn responses(
|
||||
State(state): State<Arc<NeuronState>>,
|
||||
Json(req): Json<ResponsesRequest>,
|
||||
) -> impl IntoResponse {
|
||||
let Some(candle) = state.candle.as_ref().map(Arc::clone) else {
|
||||
return (
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
Json(json!({"error": "candle harness not enabled on this neuron"})),
|
||||
)
|
||||
.into_response();
|
||||
};
|
||||
|
||||
let stream_requested = req.stream;
|
||||
let model_id = req.model.clone();
|
||||
let response_id = mint_response_id();
|
||||
let message_item_id = mint_message_item_id();
|
||||
|
||||
// Translate Responses → chat completions. The only failure
|
||||
// mode today is `previous_response_id` set, which we reject
|
||||
// with 400 — stateful conversations need a persistence layer
|
||||
// we haven't built.
|
||||
let mut chat_req = match openai_responses::request_to_chat(req) {
|
||||
Ok(r) => r,
|
||||
Err(openai_responses::TranslateError::ChainedConversationNotSupported) => {
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(json!({
|
||||
"error": "previous_response_id is not supported on this neuron",
|
||||
"code": "chained_conversation_not_supported"
|
||||
})),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
};
|
||||
chat_req.stream = Some(stream_requested);
|
||||
|
||||
if stream_requested {
|
||||
match candle
|
||||
.responses_stream(chat_req, response_id, message_item_id)
|
||||
.await
|
||||
{
|
||||
Ok(rx) => {
|
||||
// Each ResponseStreamFrame → one SSE event carrying
|
||||
// both an event name and JSON data. The Responses
|
||||
// API doesn't use a `[DONE]` terminator — clients
|
||||
// see the `response.completed` event as the end of
|
||||
// the stream.
|
||||
let body_stream = ReceiverStream::new(rx).map(|frame| {
|
||||
let body = serde_json::to_string(&frame.data).unwrap_or_else(|_| "{}".into());
|
||||
Ok::<_, Infallible>(Event::default().event(frame.event_name).data(body))
|
||||
});
|
||||
Sse::new(body_stream)
|
||||
.keep_alive(KeepAlive::default())
|
||||
.into_response()
|
||||
}
|
||||
Err(e) => inference_error_response(e),
|
||||
}
|
||||
} else {
|
||||
// Non-streaming: drive the existing chat completion path
|
||||
// and translate the result. We don't currently re-tokenise
|
||||
// to compute usage; the harness returns it via the chat
|
||||
// response and we pass it through.
|
||||
match candle.chat_completion(chat_req).await {
|
||||
Ok(chat_resp) => {
|
||||
// Extract the assistant text (chat completions
|
||||
// always emits one choice on the candle path).
|
||||
let text = chat_resp
|
||||
.choices
|
||||
.first()
|
||||
.map(|c| match &c.message.content {
|
||||
MessageContent::Text(t) => t.clone(),
|
||||
MessageContent::Parts(_) => {
|
||||
// Candle output is always text today;
|
||||
// a Parts response would be surprising.
|
||||
// Empty-string fallback is safer than
|
||||
// a panic.
|
||||
String::new()
|
||||
}
|
||||
})
|
||||
.unwrap_or_default();
|
||||
let finish = chat_resp
|
||||
.choices
|
||||
.first()
|
||||
.and_then(|c| c.finish_reason.as_deref())
|
||||
.map(finish_reason_from_str)
|
||||
.unwrap_or(crate::wire::FinishReason::Stop);
|
||||
let usage = chat_resp.usage.as_ref().map(|u| ResponsesUsage {
|
||||
input_tokens: u.prompt_tokens,
|
||||
output_tokens: u.completion_tokens,
|
||||
total_tokens: u.prompt_tokens + u.completion_tokens,
|
||||
});
|
||||
let meta = openai_responses::ResponseMeta {
|
||||
response_id: mint_response_id(),
|
||||
created_at: unix_now_secs(),
|
||||
model_id,
|
||||
message_item_id: mint_message_item_id(),
|
||||
};
|
||||
let _ = chat_resp; // make the borrow-checker happy if `text` consumed it
|
||||
let resp = openai_responses::build_response(&meta, text, finish, usage);
|
||||
Json(resp).into_response()
|
||||
}
|
||||
Err(e) => inference_error_response(e),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn finish_reason_from_str(s: &str) -> crate::wire::FinishReason {
|
||||
use crate::wire::FinishReason;
|
||||
match s {
|
||||
"length" => FinishReason::Length,
|
||||
"tool_calls" => FinishReason::ToolCalls,
|
||||
_ => FinishReason::Stop,
|
||||
}
|
||||
}
|
||||
|
||||
/// Centralised mapping from [`InferenceError`] to an HTTP response.
|
||||
/// Lifted out so the chat-completions and responses handlers stay
|
||||
/// readable and changes to error-code semantics happen in one spot.
|
||||
fn inference_error_response(err: InferenceError) -> axum::response::Response {
|
||||
match err {
|
||||
InferenceError::ModelNotLoaded(id) => (
|
||||
StatusCode::NOT_FOUND,
|
||||
Json(json!({"error": format!("model '{id}' not loaded on this neuron")})),
|
||||
)
|
||||
.into_response(),
|
||||
InferenceError::PromptTooLong { prompt_len, max } => (
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(json!({
|
||||
"error": format!("prompt has {prompt_len} tokens but max is {max}"),
|
||||
"code": "prompt_too_long",
|
||||
"prompt_len": prompt_len,
|
||||
"max": max,
|
||||
})),
|
||||
)
|
||||
.into_response(),
|
||||
InferenceError::InsufficientVram {
|
||||
free_mb,
|
||||
required_mb,
|
||||
} => (
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
Json(json!({
|
||||
"error": format!(
|
||||
"insufficient free VRAM: {free_mb} MiB free, need at least {required_mb} MiB"
|
||||
),
|
||||
"code": "insufficient_vram",
|
||||
"free_mb": free_mb,
|
||||
"required_mb": required_mb,
|
||||
})),
|
||||
)
|
||||
.into_response(),
|
||||
InferenceError::Other(e) => (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(json!({"error": format!("{e:#}")})),
|
||||
)
|
||||
.into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
fn mint_response_id() -> String {
|
||||
format!("resp_{:x}", unix_subsec_nanos())
|
||||
}
|
||||
|
||||
fn mint_message_item_id() -> String {
|
||||
format!("msg_{:x}", unix_subsec_nanos())
|
||||
}
|
||||
|
||||
fn unix_now_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
fn unix_subsec_nanos() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_nanos() as u64)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
392
crates/neuron/src/harness/chat_template.rs
Normal file
392
crates/neuron/src/harness/chat_template.rs
Normal file
@@ -0,0 +1,392 @@
|
||||
//! Chat-template rendering for the model-supplied Jinja templates
|
||||
//! HuggingFace tokenizers ship in `tokenizer_config.json`.
|
||||
//!
|
||||
//! ## Background
|
||||
//!
|
||||
//! Every modern open-weight model bundles a `chat_template` field
|
||||
//! in its `tokenizer_config.json` — a Jinja2 template string that
|
||||
//! converts a sequence of `{role, content}` messages into the
|
||||
//! exact prompt the model was trained on. Examples:
|
||||
//!
|
||||
//! - Qwen3-Coder: `<|im_start|>{role}\n{content}<|im_end|>\n…`
|
||||
//! with conditional `enable_thinking` handling that injects an
|
||||
//! empty `<think>\n\n</think>` block when set false.
|
||||
//! - DeepSeek-R1: similar im_start framing with different special-
|
||||
//! token names.
|
||||
//! - Mistral / Magistral: a `[INST]` / `[/INST]` framing.
|
||||
//! - Claude / Llama: another shape again.
|
||||
//!
|
||||
//! Rendering the model's own template is the only way to get the
|
||||
//! *exact* prompt format the model was trained on plus the
|
||||
//! model-specific kwargs (`enable_thinking`, `tools`, …) without
|
||||
//! hardcoding per-model logic. The alternative — neuron's previous
|
||||
//! `format_qwen3_prompt` — was a hardcoded Qwen3 ChatML glue that
|
||||
//! ignored kwargs entirely.
|
||||
//!
|
||||
//! ## Scope
|
||||
//!
|
||||
//! This module is request-side only: it builds the prompt string
|
||||
//! the tokenizer ingests before inference. The reasoning- and
|
||||
//! tool-call-marker token routing (issues #6, #8) is response-side
|
||||
//! and stays in `wire::openai_chat` / the streaming inference
|
||||
//! loops.
|
||||
//!
|
||||
//! ## Fallback
|
||||
//!
|
||||
//! When the model's `tokenizer_config.json` is missing, doesn't
|
||||
//! parse, lacks a `chat_template`, or renders an error, the caller
|
||||
//! falls back to `format_qwen3_prompt`. The
|
||||
//! `NEURON_USE_CHAT_TEMPLATE=false` env var is a global kill
|
||||
//! switch — if a deploy goes sideways and the renderer is to
|
||||
//! blame, an operator can flip the env and restart neuron without
|
||||
//! shipping a new build.
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use cortex_core::openai::{ChatMessage, MessageContent};
|
||||
use minijinja::Environment;
|
||||
use serde_json::Value;
|
||||
use std::path::Path;
|
||||
|
||||
/// Environment variable that, when set to `false`/`0`/`no`,
|
||||
/// forces every model to skip its `chat_template` and fall back
|
||||
/// to `format_qwen3_prompt`. Default (unset) is "use chat
|
||||
/// templates where available".
|
||||
pub const KILL_SWITCH_ENV: &str = "NEURON_USE_CHAT_TEMPLATE";
|
||||
|
||||
/// Read the global kill switch. `true` means chat templates are
|
||||
/// enabled; `false` forces the fallback path everywhere.
|
||||
pub fn chat_templates_enabled() -> bool {
|
||||
match std::env::var(KILL_SWITCH_ENV).ok().as_deref() {
|
||||
Some(s) => !matches!(
|
||||
s.trim().to_ascii_lowercase().as_str(),
|
||||
"false" | "0" | "no" | "off"
|
||||
),
|
||||
None => true,
|
||||
}
|
||||
}
|
||||
|
||||
/// Convenience: probe for `tokenizer_config.json` in the same
|
||||
/// directory the tokenizer was loaded from. Both files come from
|
||||
/// the same HuggingFace snapshot in the hf-hub cache, so the
|
||||
/// sibling path is reliable.
|
||||
pub fn load_chat_template_alongside(tokenizer_json_path: &Path) -> Option<String> {
|
||||
let parent = tokenizer_json_path.parent()?;
|
||||
let config_path = parent.join("tokenizer_config.json");
|
||||
load_chat_template_from(&config_path)
|
||||
}
|
||||
|
||||
/// Best-effort load of `chat_template` from a HuggingFace
|
||||
/// `tokenizer_config.json`. Returns `None` when the file is
|
||||
/// absent, doesn't parse, or lacks the `chat_template` field —
|
||||
/// in all of those cases the caller falls back to
|
||||
/// `format_qwen3_prompt`. Warnings are logged so an operator can
|
||||
/// see why the fallback fired.
|
||||
pub fn load_chat_template_from(path: &Path) -> Option<String> {
|
||||
let text = match std::fs::read_to_string(path) {
|
||||
Ok(t) => t,
|
||||
Err(e) => {
|
||||
tracing::debug!(
|
||||
path = %path.display(),
|
||||
error = %e,
|
||||
"chat_template: tokenizer_config.json absent or unreadable; falling back"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let value: Value = match serde_json::from_str(&text) {
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
path = %path.display(),
|
||||
error = %e,
|
||||
"chat_template: tokenizer_config.json failed to parse; falling back"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
// Some tokenizer_config.json files carry `chat_template` as an
|
||||
// array of `{name, template}` objects (multi-template models —
|
||||
// tool-use variant, default variant). For now we pick the first
|
||||
// entry; future iterations could honour a name hint.
|
||||
match value.get("chat_template") {
|
||||
Some(Value::String(s)) => Some(s.clone()),
|
||||
Some(Value::Array(arr)) => {
|
||||
for entry in arr {
|
||||
if let Some(t) = entry.get("template").and_then(|v| v.as_str()) {
|
||||
return Some(t.to_string());
|
||||
}
|
||||
}
|
||||
tracing::warn!(
|
||||
path = %path.display(),
|
||||
"chat_template: array form had no usable template entry; falling back"
|
||||
);
|
||||
None
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Render the chat template into the prompt the model expects.
|
||||
///
|
||||
/// `template` is the raw Jinja string from `tokenizer_config.json`.
|
||||
/// `messages` is the conversation in order. `kwargs` is the
|
||||
/// `chat_template_kwargs` object the client supplied on the
|
||||
/// request (or `Value::Null` when absent). The function expands
|
||||
/// the kwargs into the Jinja context alongside the standard
|
||||
/// `messages` and `add_generation_prompt` variables HF templates
|
||||
/// expect.
|
||||
///
|
||||
/// `tools` is the request's `tools` array (or `Value::Null`).
|
||||
/// Some chat templates iterate it to emit native tool definitions
|
||||
/// (Qwen3-Coder's tool-use template, Mistral's [TOOL_DEFINITIONS]
|
||||
/// frame). We forward whatever the client sent without
|
||||
/// interpretation.
|
||||
pub fn render_chat_template(
|
||||
template: &str,
|
||||
messages: &[ChatMessage],
|
||||
tools: &Value,
|
||||
kwargs: &Value,
|
||||
) -> Result<String> {
|
||||
let mut env = Environment::new();
|
||||
// Compile the template against a fixed name so error messages
|
||||
// surface "chat_template" rather than `<template>`.
|
||||
env.add_template("chat_template", template)
|
||||
.context("compile chat_template")?;
|
||||
let tmpl = env.get_template("chat_template").unwrap();
|
||||
|
||||
// Convert our internal ChatMessage shape into the
|
||||
// `[{role, content}]` shape HF templates iterate. Text content
|
||||
// becomes a string; Parts becomes an array of content blocks.
|
||||
// The HF templates handle both shapes via `content is string`
|
||||
// checks or content-array iteration.
|
||||
let messages_json: Vec<Value> = messages
|
||||
.iter()
|
||||
.map(|m| {
|
||||
let content_value = match &m.content {
|
||||
MessageContent::Text(s) => Value::String(s.clone()),
|
||||
MessageContent::Parts(parts) => Value::Array(parts.clone()),
|
||||
};
|
||||
let mut obj = serde_json::Map::new();
|
||||
obj.insert("role".into(), Value::String(m.role.clone()));
|
||||
obj.insert("content".into(), content_value);
|
||||
// Forward extras (e.g. tool_calls on assistant turns,
|
||||
// tool_call_id on tool result turns). HF templates that
|
||||
// need them read e.g. `message.tool_calls`.
|
||||
if let Value::Object(extras) = &m.extra {
|
||||
for (k, v) in extras {
|
||||
obj.insert(k.clone(), v.clone());
|
||||
}
|
||||
}
|
||||
Value::Object(obj)
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Build the kwargs context. Add base bindings the template
|
||||
// expects (`messages`, `add_generation_prompt`, `tools`) plus
|
||||
// anything the caller passed in `chat_template_kwargs`. Caller
|
||||
// kwargs override the defaults so `add_generation_prompt: false`
|
||||
// from the request actually wins.
|
||||
let mut ctx_map = serde_json::Map::new();
|
||||
ctx_map.insert("messages".into(), Value::Array(messages_json));
|
||||
ctx_map.insert("add_generation_prompt".into(), Value::Bool(true));
|
||||
if !tools.is_null() {
|
||||
ctx_map.insert("tools".into(), tools.clone());
|
||||
}
|
||||
if let Value::Object(kwargs_obj) = kwargs {
|
||||
for (k, v) in kwargs_obj {
|
||||
ctx_map.insert(k.clone(), v.clone());
|
||||
}
|
||||
}
|
||||
// `Template::render` takes any Serialize value; serde_json's
|
||||
// `Value` implements it natively, so we pass the assembled
|
||||
// context object directly without going through the
|
||||
// `context!` macro (which expects minijinja-native values).
|
||||
tmpl.render(Value::Object(ctx_map))
|
||||
.context("render chat_template")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
fn user_msg(text: &str) -> ChatMessage {
|
||||
ChatMessage {
|
||||
role: "user".into(),
|
||||
content: MessageContent::Text(text.into()),
|
||||
extra: Value::Object(Default::default()),
|
||||
}
|
||||
}
|
||||
|
||||
fn assistant_msg(text: &str) -> ChatMessage {
|
||||
ChatMessage {
|
||||
role: "assistant".into(),
|
||||
content: MessageContent::Text(text.into()),
|
||||
extra: Value::Object(Default::default()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Minimal Qwen3-style template — enough surface to confirm
|
||||
/// our renderer threads role + content correctly without
|
||||
/// loading a real model's tokenizer_config.json.
|
||||
const QWEN3_LIKE: &str = "{%- for message in messages -%}\
|
||||
<|im_start|>{{ message.role }}\n{{ message.content }}<|im_end|>\n\
|
||||
{%- endfor -%}\
|
||||
{%- if add_generation_prompt -%}<|im_start|>assistant\n{%- endif -%}";
|
||||
|
||||
#[test]
|
||||
fn renders_basic_conversation() {
|
||||
let prompt = render_chat_template(
|
||||
QWEN3_LIKE,
|
||||
&[user_msg("hello"), assistant_msg("hi"), user_msg("bye")],
|
||||
&Value::Null,
|
||||
&Value::Null,
|
||||
)
|
||||
.unwrap();
|
||||
// Structural assertions — the exact whitespace produced
|
||||
// by a given template is a Jinja-trim concern that varies
|
||||
// per real chat_template. What matters is that every
|
||||
// turn's role + content thread through in order, and that
|
||||
// the generation cue lands at the end.
|
||||
assert!(
|
||||
prompt.contains("<|im_start|>user\nhello<|im_end|>"),
|
||||
"first user turn missing: {prompt}"
|
||||
);
|
||||
assert!(
|
||||
prompt.contains("<|im_start|>assistant\nhi<|im_end|>"),
|
||||
"assistant turn missing: {prompt}"
|
||||
);
|
||||
assert!(
|
||||
prompt.contains("<|im_start|>user\nbye<|im_end|>"),
|
||||
"second user turn missing: {prompt}"
|
||||
);
|
||||
assert!(
|
||||
prompt.ends_with("<|im_start|>assistant")
|
||||
|| prompt.ends_with("<|im_start|>assistant\n"),
|
||||
"generation cue missing at end: {prompt}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn kwargs_are_threaded_into_template_context() {
|
||||
// Replica of Qwen3's enable_thinking branch in
|
||||
// simplified form. When the kwarg is false, the model's
|
||||
// template injects an empty `<think>...</think>` block
|
||||
// before the generation cue — pre-filling the model's
|
||||
// reasoning slot with "no thinking" so the model emits
|
||||
// the answer directly.
|
||||
let template = "{%- if enable_thinking is defined and enable_thinking is false -%}\
|
||||
NO_THINK\
|
||||
{%- else -%}\
|
||||
THINK_OK\
|
||||
{%- endif -%}";
|
||||
let r_disabled = render_chat_template(
|
||||
template,
|
||||
&[],
|
||||
&Value::Null,
|
||||
&json!({ "enable_thinking": false }),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(r_disabled, "NO_THINK");
|
||||
let r_default = render_chat_template(template, &[], &Value::Null, &Value::Null).unwrap();
|
||||
assert_eq!(r_default, "THINK_OK");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_template_field_returns_none() {
|
||||
let tmp = std::env::temp_dir().join("neuron-test-tokenizer-missing-field.json");
|
||||
std::fs::write(&tmp, r#"{"some_other_field": 1}"#).unwrap();
|
||||
assert!(load_chat_template_from(&tmp).is_none());
|
||||
let _ = std::fs::remove_file(tmp);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_template_from_string_field() {
|
||||
let tmp = std::env::temp_dir().join("neuron-test-tokenizer-string.json");
|
||||
std::fs::write(
|
||||
&tmp,
|
||||
r#"{"chat_template": "hello {{ messages[0].content }}"}"#,
|
||||
)
|
||||
.unwrap();
|
||||
let t = load_chat_template_from(&tmp).expect("template loaded");
|
||||
assert!(t.contains("messages[0].content"));
|
||||
let _ = std::fs::remove_file(tmp);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_template_from_array_form() {
|
||||
// Some HF models ship `chat_template` as `[{name, template}, ...]`.
|
||||
let tmp = std::env::temp_dir().join("neuron-test-tokenizer-array.json");
|
||||
std::fs::write(
|
||||
&tmp,
|
||||
r#"{"chat_template": [{"name": "default", "template": "ARR"}]}"#,
|
||||
)
|
||||
.unwrap();
|
||||
let t = load_chat_template_from(&tmp).expect("template loaded");
|
||||
assert_eq!(t, "ARR");
|
||||
let _ = std::fs::remove_file(tmp);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_file_returns_none_quietly() {
|
||||
let absent = std::path::PathBuf::from("/definitely/not/a/real/path.json");
|
||||
assert!(load_chat_template_from(&absent).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unparseable_returns_none() {
|
||||
let tmp = std::env::temp_dir().join("neuron-test-tokenizer-garbage.json");
|
||||
std::fs::write(&tmp, b"{not valid json").unwrap();
|
||||
assert!(load_chat_template_from(&tmp).is_none());
|
||||
let _ = std::fs::remove_file(tmp);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn kill_switch_recognises_truthy_falsy_values() {
|
||||
// Test against the actual env var so callers see the
|
||||
// same behaviour as production. Serialise via a
|
||||
// mutex — see path_util.rs for the pattern.
|
||||
use std::sync::Mutex;
|
||||
static LOCK: Mutex<()> = Mutex::new(());
|
||||
let _g = LOCK.lock().unwrap();
|
||||
let prior = std::env::var(KILL_SWITCH_ENV).ok();
|
||||
unsafe {
|
||||
std::env::remove_var(KILL_SWITCH_ENV);
|
||||
}
|
||||
assert!(chat_templates_enabled());
|
||||
for value in ["false", "0", "no", "off", "FALSE", " no "] {
|
||||
unsafe { std::env::set_var(KILL_SWITCH_ENV, value) };
|
||||
assert!(!chat_templates_enabled(), "value {value:?} should disable");
|
||||
}
|
||||
for value in ["true", "1", "yes", ""] {
|
||||
unsafe { std::env::set_var(KILL_SWITCH_ENV, value) };
|
||||
assert!(chat_templates_enabled(), "value {value:?} should enable");
|
||||
}
|
||||
unsafe {
|
||||
match prior {
|
||||
Some(p) => std::env::set_var(KILL_SWITCH_ENV, p),
|
||||
None => std::env::remove_var(KILL_SWITCH_ENV),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn message_extras_thread_through_for_tool_calls() {
|
||||
// HF templates read assistant.tool_calls and tool
|
||||
// turns' tool_call_id. Confirm our extras flatten into
|
||||
// the message object the template iterates.
|
||||
let mut extras = serde_json::Map::new();
|
||||
extras.insert(
|
||||
"tool_calls".into(),
|
||||
json!([{"id": "t1", "function": {"name": "x", "arguments": "{}"}}]),
|
||||
);
|
||||
let msg = ChatMessage {
|
||||
role: "assistant".into(),
|
||||
content: MessageContent::Text(String::new()),
|
||||
extra: Value::Object(extras),
|
||||
};
|
||||
let template = "{{ messages[0].tool_calls[0].id }}";
|
||||
let rendered = render_chat_template(template, &[msg], &Value::Null, &Value::Null).unwrap();
|
||||
assert_eq!(rendered, "t1");
|
||||
}
|
||||
}
|
||||
@@ -2,7 +2,9 @@
|
||||
|
||||
pub mod arch;
|
||||
pub mod candle;
|
||||
pub mod chat_template;
|
||||
pub mod device_worker;
|
||||
pub mod preflight;
|
||||
pub mod tp;
|
||||
|
||||
use anyhow::Result;
|
||||
|
||||
575
crates/neuron/src/harness/preflight.rs
Normal file
575
crates/neuron/src/harness/preflight.rs
Normal file
@@ -0,0 +1,575 @@
|
||||
//! Placement feasibility check that runs before any device allocation,
|
||||
//! NCCL handshake, or weight download.
|
||||
//!
|
||||
//! The loader path in `candle.rs` historically discovers an
|
||||
//! incompatibility *after* it has already started fetching files —
|
||||
//! "fetch config.json from HauhauCS/...: 404 Not Found" surfaces hours
|
||||
//! after operators set `tensor_parallel = 2` on a GGUF-only repo, with
|
||||
//! no hint about what's actually wrong. Preflight closes that gap:
|
||||
//!
|
||||
//! 1. one `repo.info()` round-trip (siblings listing, no blob fetch)
|
||||
//! 2. classify the repo: GGUF-only, dense safetensors, mixed, empty
|
||||
//! 3. apply the feasibility table against the requested
|
||||
//! `ModelSpec` (tp_size, quant)
|
||||
//! 4. return a structured `PreflightError` the API layer can map to
|
||||
//! 422 + JSON, or `Ok(PlacementPlan)` carrying the decisions the
|
||||
//! downstream load path needs (which GGUF file to fetch, etc.).
|
||||
//!
|
||||
//! Phase 2 of plan-source-aware-loader-preflight. The Phase 1 scheme
|
||||
//! work — `ModelSourceId` and per-scheme `SourceConfig` — is a
|
||||
//! separate PR; preflight runs against the single configured
|
||||
//! HuggingFace source for now and the scheme threading drops in
|
||||
//! cleanly when Phase 1 lands.
|
||||
|
||||
use cortex_core::harness::ModelSpec;
|
||||
use hf_hub::api::tokio::Api;
|
||||
use serde::Serialize;
|
||||
|
||||
/// What the repo's siblings listing tells us about how to load it.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
|
||||
#[serde(tag = "kind", rename_all = "snake_case")]
|
||||
pub enum SourceFormat {
|
||||
/// Only GGUF files present. Single-GPU load path. `quants` is the
|
||||
/// lowercased filename list so the operator can be told what's
|
||||
/// actually available when their `quant=` choice doesn't match.
|
||||
Gguf { quants: Vec<String> },
|
||||
/// Dense safetensors (single-file or sharded via index.json).
|
||||
/// Goes through `load_arch_dense` on single-GPU, or `load_tp` (with
|
||||
/// optional in-situ quantization) when `tensor_parallel > 1`.
|
||||
DenseSafetensors { sharded: bool },
|
||||
/// Both safetensors and GGUF present — prefer the dense path
|
||||
/// because it composes with TP and ISQ. We surface the GGUF
|
||||
/// filenames anyway so operators with a strong preference can
|
||||
/// see they exist.
|
||||
Mixed { gguf_quants: Vec<String> },
|
||||
/// No recognised weight files. Either a tokenizer-only repo
|
||||
/// (e.g. some base-model repos that only host `tokenizer.json` and
|
||||
/// expect the operator to use a `-GGUF` sibling repo) or a
|
||||
/// genuinely empty entry.
|
||||
Empty,
|
||||
}
|
||||
|
||||
/// Output of `preflight` for a load that can proceed. Carries the
|
||||
/// decisions downstream resolve_* paths would otherwise re-derive.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct PlacementPlan {
|
||||
pub model_id: String,
|
||||
pub format: SourceFormat,
|
||||
pub tp_size: u32,
|
||||
/// Filename of the GGUF to fetch, populated when `format` is
|
||||
/// `Gguf` and a single-GPU load was requested. None for the
|
||||
/// dense/TP path.
|
||||
pub picked_quant_file: Option<String>,
|
||||
}
|
||||
|
||||
/// Structured failure modes. Each variant carries the fields the API
|
||||
/// layer needs to produce an actionable 422 body.
|
||||
#[derive(Debug, Clone, Serialize, thiserror::Error)]
|
||||
#[serde(tag = "kind", rename_all = "snake_case")]
|
||||
pub enum PreflightError {
|
||||
/// `repo.info()` failed. Captures the underlying cause as a string
|
||||
/// so the operator log shows whether it's auth, 404, or transport.
|
||||
#[error("failed to fetch repo info for '{model_id}': {cause}")]
|
||||
RepoFetchFailed { model_id: String, cause: String },
|
||||
|
||||
/// The repo exists but has no recognised weight files.
|
||||
#[error(
|
||||
"repo '{model_id}' has no recognised weight files (no .gguf, no .safetensors); \
|
||||
a tokenizer-only repo cannot be loaded directly"
|
||||
)]
|
||||
EmptyRepo { model_id: String },
|
||||
|
||||
/// Operator asked for `tensor_parallel > 1` on a GGUF-only repo.
|
||||
/// The TP path requires safetensors+config for in-situ
|
||||
/// quantization; GGUF-TP isn't implemented (see CLAUDE.md).
|
||||
#[error(
|
||||
"cannot load '{model_id}' with tensor_parallel={tp_size}: repo is GGUF-only \
|
||||
({} .gguf files); TP requires dense safetensors. {suggestion}",
|
||||
gguf_quants.len()
|
||||
)]
|
||||
TpRequiresSafetensors {
|
||||
model_id: String,
|
||||
tp_size: u32,
|
||||
gguf_quants: Vec<String>,
|
||||
suggestion: String,
|
||||
},
|
||||
|
||||
/// Operator asked for a GGUF quant whose substring doesn't match
|
||||
/// any filename in the repo. `nearest` is a best-effort Levenshtein
|
||||
/// suggestion against the available quant names.
|
||||
#[error(
|
||||
"no GGUF file in '{model_id}' matches quant '{requested}'; \
|
||||
available: {available:?}{}",
|
||||
nearest.as_ref().map(|n| format!("; did you mean '{n}'?")).unwrap_or_default()
|
||||
)]
|
||||
QuantNotFound {
|
||||
model_id: String,
|
||||
requested: String,
|
||||
available: Vec<String>,
|
||||
nearest: Option<String>,
|
||||
},
|
||||
}
|
||||
|
||||
/// Run the placement check.
|
||||
///
|
||||
/// One network round-trip (`repo.info()`); no blob fetches. Returns
|
||||
/// `Ok(PlacementPlan)` when the requested combination is feasible, or
|
||||
/// a structured `PreflightError` describing what's wrong.
|
||||
pub async fn preflight(api: &Api, spec: &ModelSpec) -> Result<PlacementPlan, PreflightError> {
|
||||
let repo = api.model(spec.model_id.clone());
|
||||
let info = repo
|
||||
.info()
|
||||
.await
|
||||
.map_err(|e| PreflightError::RepoFetchFailed {
|
||||
model_id: spec.model_id.clone(),
|
||||
cause: format!("{e}"),
|
||||
})?;
|
||||
|
||||
let filenames: Vec<&str> = info.siblings.iter().map(|s| s.rfilename.as_str()).collect();
|
||||
let format = classify(&filenames);
|
||||
let tp_size = spec.tensor_parallel.unwrap_or(1);
|
||||
|
||||
match (&format, tp_size, spec.quant.as_deref()) {
|
||||
// No weights at all — nothing to do.
|
||||
(SourceFormat::Empty, _, _) => Err(PreflightError::EmptyRepo {
|
||||
model_id: spec.model_id.clone(),
|
||||
}),
|
||||
|
||||
// GGUF-only + TP: not supported. Today's HauhauCS failure.
|
||||
(SourceFormat::Gguf { quants }, tp, _) if tp > 1 => {
|
||||
Err(PreflightError::TpRequiresSafetensors {
|
||||
model_id: spec.model_id.clone(),
|
||||
tp_size: tp,
|
||||
gguf_quants: quants.clone(),
|
||||
suggestion: format!(
|
||||
"Set tensor_parallel=1 and pick a quant from {quants:?}, \
|
||||
or use a dense safetensors release of this model."
|
||||
),
|
||||
})
|
||||
}
|
||||
|
||||
// GGUF-only + single-GPU: pick the file that matches the
|
||||
// operator's quant. Empty quant matches the first GGUF.
|
||||
(SourceFormat::Gguf { quants }, _, requested) => {
|
||||
let picked = pick_gguf_file(&filenames, requested.unwrap_or(""));
|
||||
match picked {
|
||||
Some(fname) => Ok(PlacementPlan {
|
||||
model_id: spec.model_id.clone(),
|
||||
format: format.clone(),
|
||||
tp_size,
|
||||
picked_quant_file: Some(fname),
|
||||
}),
|
||||
None => Err(PreflightError::QuantNotFound {
|
||||
model_id: spec.model_id.clone(),
|
||||
requested: requested.unwrap_or("").to_string(),
|
||||
available: quants.clone(),
|
||||
nearest: nearest_quant(requested.unwrap_or(""), quants),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
// Dense or mixed: dense path handles both single-GPU and TP.
|
||||
// The architecture compatibility check stays where it is —
|
||||
// `check_dense_config_supported` runs once `config.json` is
|
||||
// on disk, since it needs the parsed JSON.
|
||||
(SourceFormat::DenseSafetensors { .. } | SourceFormat::Mixed { .. }, _, _) => {
|
||||
Ok(PlacementPlan {
|
||||
model_id: spec.model_id.clone(),
|
||||
format: format.clone(),
|
||||
tp_size,
|
||||
picked_quant_file: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Classify a siblings file list into a `SourceFormat`. Pulled out so
|
||||
/// the unit tests can exercise it against fixture JSON without
|
||||
/// spinning up an Api.
|
||||
pub fn classify(filenames: &[&str]) -> SourceFormat {
|
||||
let mut gguf_quants: Vec<String> = filenames
|
||||
.iter()
|
||||
.filter(|f| f.to_lowercase().ends_with(".gguf"))
|
||||
.map(|f| f.to_lowercase())
|
||||
.collect();
|
||||
gguf_quants.sort();
|
||||
gguf_quants.dedup();
|
||||
|
||||
let has_safetensors = filenames.iter().any(|f| f.ends_with(".safetensors"));
|
||||
let sharded = filenames
|
||||
.iter()
|
||||
.any(|f| f.ends_with("model.safetensors.index.json"));
|
||||
|
||||
match (has_safetensors, gguf_quants.is_empty()) {
|
||||
(true, true) => SourceFormat::DenseSafetensors { sharded },
|
||||
(true, false) => SourceFormat::Mixed { gguf_quants },
|
||||
(false, false) => SourceFormat::Gguf {
|
||||
quants: gguf_quants,
|
||||
},
|
||||
(false, true) => SourceFormat::Empty,
|
||||
}
|
||||
}
|
||||
|
||||
/// Mirror of the quant-matching logic in `candle.rs::resolve_files` so
|
||||
/// preflight picks the same file the downstream loader would. Empty
|
||||
/// quant returns the first `.gguf` (any quant). Lowercased substring
|
||||
/// match otherwise.
|
||||
fn pick_gguf_file(filenames: &[&str], quant_lc: &str) -> Option<String> {
|
||||
filenames
|
||||
.iter()
|
||||
.filter(|f| f.to_lowercase().ends_with(".gguf"))
|
||||
.find(|f| quant_lc.is_empty() || f.to_lowercase().contains(quant_lc))
|
||||
.map(|f| f.to_string())
|
||||
}
|
||||
|
||||
/// Best-effort suggestion when the operator's quant name doesn't
|
||||
/// substring-match any filename. Extracts the quant-ish token from
|
||||
/// each `.gguf` filename and picks the one with the smallest
|
||||
/// Levenshtein distance to the requested string. Returns None when
|
||||
/// the input is empty or no candidates exist.
|
||||
fn nearest_quant(requested: &str, candidates: &[String]) -> Option<String> {
|
||||
if requested.is_empty() || candidates.is_empty() {
|
||||
return None;
|
||||
}
|
||||
// Pull the "Q6_K_P"/"IQ4_XS"-ish token out of each filename for a
|
||||
// fairer comparison. Filenames look like
|
||||
// `Qwen3.6-27B-Uncensored-HauhauCS-Aggressive-Q6_K_P.gguf`, so the
|
||||
// quant is the last `-`-separated segment before the extension,
|
||||
// lowercased.
|
||||
let tokens: Vec<(String, String)> = candidates
|
||||
.iter()
|
||||
.map(|f| (extract_quant_token(f), f.clone()))
|
||||
.collect();
|
||||
|
||||
let req_lc = requested.to_lowercase();
|
||||
tokens
|
||||
.into_iter()
|
||||
.min_by_key(|(token, _)| levenshtein(&req_lc, token))
|
||||
.map(|(token, _)| token)
|
||||
}
|
||||
|
||||
fn extract_quant_token(filename: &str) -> String {
|
||||
let stem = filename
|
||||
.rsplit_once('.')
|
||||
.map(|(s, _)| s)
|
||||
.unwrap_or(filename);
|
||||
let token = stem.rsplit('-').next().unwrap_or(stem);
|
||||
token.to_lowercase()
|
||||
}
|
||||
|
||||
/// Iterative Levenshtein. Small inputs (quant names are <=12 chars),
|
||||
/// no need for the `levenshtein` crate.
|
||||
fn levenshtein(a: &str, b: &str) -> usize {
|
||||
let a: Vec<char> = a.chars().collect();
|
||||
let b: Vec<char> = b.chars().collect();
|
||||
let (m, n) = (a.len(), b.len());
|
||||
if m == 0 {
|
||||
return n;
|
||||
}
|
||||
if n == 0 {
|
||||
return m;
|
||||
}
|
||||
let mut prev: Vec<usize> = (0..=n).collect();
|
||||
let mut curr = vec![0usize; n + 1];
|
||||
for i in 1..=m {
|
||||
curr[0] = i;
|
||||
for j in 1..=n {
|
||||
let cost = if a[i - 1] == b[j - 1] { 0 } else { 1 };
|
||||
curr[j] = (prev[j] + 1).min(curr[j - 1] + 1).min(prev[j - 1] + cost);
|
||||
}
|
||||
std::mem::swap(&mut prev, &mut curr);
|
||||
}
|
||||
prev[n]
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn spec(model_id: &str, tp: Option<u32>, quant: Option<&str>) -> ModelSpec {
|
||||
ModelSpec {
|
||||
model_id: model_id.into(),
|
||||
harness: "candle".into(),
|
||||
quant: quant.map(String::from),
|
||||
tensor_parallel: tp,
|
||||
devices: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classify_gguf_only() {
|
||||
let files = [
|
||||
"README.md",
|
||||
".gitattributes",
|
||||
"Qwen3.6-27B-Q6_K_P.gguf",
|
||||
"Qwen3.6-27B-Q4_K_P.gguf",
|
||||
];
|
||||
match classify(&files) {
|
||||
SourceFormat::Gguf { quants } => {
|
||||
assert_eq!(quants.len(), 2);
|
||||
assert!(quants.iter().any(|q| q.contains("q6_k_p")));
|
||||
}
|
||||
other => panic!("expected Gguf, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classify_dense_sharded() {
|
||||
let files = [
|
||||
"config.json",
|
||||
"tokenizer.json",
|
||||
"model.safetensors.index.json",
|
||||
"model-00001-of-00002.safetensors",
|
||||
"model-00002-of-00002.safetensors",
|
||||
];
|
||||
assert_eq!(
|
||||
classify(&files),
|
||||
SourceFormat::DenseSafetensors { sharded: true }
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classify_dense_single_file() {
|
||||
let files = ["config.json", "tokenizer.json", "model.safetensors"];
|
||||
assert_eq!(
|
||||
classify(&files),
|
||||
SourceFormat::DenseSafetensors { sharded: false }
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classify_mixed() {
|
||||
let files = [
|
||||
"config.json",
|
||||
"tokenizer.json",
|
||||
"model.safetensors",
|
||||
"model-Q4_K_M.gguf",
|
||||
];
|
||||
match classify(&files) {
|
||||
SourceFormat::Mixed { gguf_quants } => {
|
||||
assert_eq!(gguf_quants, vec!["model-q4_k_m.gguf"]);
|
||||
}
|
||||
other => panic!("expected Mixed, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classify_empty() {
|
||||
let files = ["README.md", "tokenizer.json"];
|
||||
assert_eq!(classify(&files), SourceFormat::Empty);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pick_gguf_substring_match() {
|
||||
let files = ["model-Q4_K_M.gguf", "model-Q6_K.gguf", "model-Q8_0.gguf"];
|
||||
assert_eq!(
|
||||
pick_gguf_file(&files, "q6_k"),
|
||||
Some("model-Q6_K.gguf".into())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pick_gguf_empty_returns_first() {
|
||||
let files = ["model-Q4_K_M.gguf", "model-Q6_K.gguf"];
|
||||
assert_eq!(pick_gguf_file(&files, ""), Some("model-Q4_K_M.gguf".into()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pick_gguf_no_match() {
|
||||
let files = ["model-Q4_K_M.gguf", "model-Q6_K.gguf"];
|
||||
assert_eq!(pick_gguf_file(&files, "iq2_xs"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn nearest_quant_suggests_close_match() {
|
||||
// Today's HauhauCS scenario: operator wrote "q6k", actual
|
||||
// filename token is "q6_k_p". Should suggest the latter.
|
||||
let candidates = vec![
|
||||
"qwen-q4_k_p.gguf".to_string(),
|
||||
"qwen-q5_k_p.gguf".to_string(),
|
||||
"qwen-q6_k_p.gguf".to_string(),
|
||||
"qwen-q8_k_p.gguf".to_string(),
|
||||
];
|
||||
assert_eq!(nearest_quant("q6k", &candidates), Some("q6_k_p".into()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn nearest_quant_empty_input() {
|
||||
assert_eq!(nearest_quant("", &[]), None);
|
||||
assert_eq!(nearest_quant("q6k", &[]), None);
|
||||
assert_eq!(nearest_quant("", &["model-q4.gguf".into()]), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_quant_handles_typical_filenames() {
|
||||
assert_eq!(extract_quant_token("Qwen3.6-27B-Q6_K_P.gguf"), "q6_k_p");
|
||||
assert_eq!(extract_quant_token("model-IQ4_XS.gguf"), "iq4_xs");
|
||||
assert_eq!(extract_quant_token("simple.gguf"), "simple");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn levenshtein_basics() {
|
||||
assert_eq!(levenshtein("", ""), 0);
|
||||
assert_eq!(levenshtein("abc", ""), 3);
|
||||
assert_eq!(levenshtein("", "abc"), 3);
|
||||
assert_eq!(levenshtein("kitten", "sitting"), 3);
|
||||
assert_eq!(levenshtein("q6k", "q6_k_p"), 3);
|
||||
assert_eq!(levenshtein("q6k", "q4_k_p"), 4);
|
||||
}
|
||||
|
||||
// Higher-level preflight tests below exercise the full feasibility
|
||||
// table via a thin wrapper that bypasses the network — we hand it
|
||||
// a pre-built `SourceFormat` and request shape, then drive the
|
||||
// same decision logic. The end-to-end test with a mock HTTP
|
||||
// server lives in tests/preflight.rs (integration).
|
||||
|
||||
/// Mirror of the `match` in `preflight()` but takes a classified
|
||||
/// `SourceFormat` directly. Lets us unit-test the feasibility
|
||||
/// table without making the API trait object-safe / boxable.
|
||||
fn decide(
|
||||
spec: &ModelSpec,
|
||||
format: &SourceFormat,
|
||||
filenames: &[&str],
|
||||
) -> Result<PlacementPlan, PreflightError> {
|
||||
let tp_size = spec.tensor_parallel.unwrap_or(1);
|
||||
match (format, tp_size, spec.quant.as_deref()) {
|
||||
(SourceFormat::Empty, _, _) => Err(PreflightError::EmptyRepo {
|
||||
model_id: spec.model_id.clone(),
|
||||
}),
|
||||
(SourceFormat::Gguf { quants }, tp, _) if tp > 1 => {
|
||||
Err(PreflightError::TpRequiresSafetensors {
|
||||
model_id: spec.model_id.clone(),
|
||||
tp_size: tp,
|
||||
gguf_quants: quants.clone(),
|
||||
suggestion: format!(
|
||||
"Set tensor_parallel=1 and pick a quant from {quants:?}, \
|
||||
or use a dense safetensors release of this model."
|
||||
),
|
||||
})
|
||||
}
|
||||
(SourceFormat::Gguf { quants }, _, requested) => {
|
||||
let picked = pick_gguf_file(filenames, requested.unwrap_or(""));
|
||||
match picked {
|
||||
Some(fname) => Ok(PlacementPlan {
|
||||
model_id: spec.model_id.clone(),
|
||||
format: format.clone(),
|
||||
tp_size,
|
||||
picked_quant_file: Some(fname),
|
||||
}),
|
||||
None => Err(PreflightError::QuantNotFound {
|
||||
model_id: spec.model_id.clone(),
|
||||
requested: requested.unwrap_or("").to_string(),
|
||||
available: quants.clone(),
|
||||
nearest: nearest_quant(requested.unwrap_or(""), quants),
|
||||
}),
|
||||
}
|
||||
}
|
||||
(SourceFormat::DenseSafetensors { .. } | SourceFormat::Mixed { .. }, _, _) => {
|
||||
Ok(PlacementPlan {
|
||||
model_id: spec.model_id.clone(),
|
||||
format: format.clone(),
|
||||
tp_size,
|
||||
picked_quant_file: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn feasibility_gguf_tp_rejected() {
|
||||
let files = ["Qwen-Q6_K_P.gguf", "Qwen-Q4_K_P.gguf"];
|
||||
let fmt = classify(&files);
|
||||
let s = spec("HauhauCS/Qwen3.6", Some(2), Some("q6k"));
|
||||
match decide(&s, &fmt, &files).unwrap_err() {
|
||||
PreflightError::TpRequiresSafetensors {
|
||||
model_id,
|
||||
tp_size,
|
||||
gguf_quants,
|
||||
..
|
||||
} => {
|
||||
assert_eq!(model_id, "HauhauCS/Qwen3.6");
|
||||
assert_eq!(tp_size, 2);
|
||||
assert_eq!(gguf_quants.len(), 2);
|
||||
}
|
||||
other => panic!("expected TpRequiresSafetensors, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn feasibility_gguf_single_gpu_bad_quant() {
|
||||
let files = [
|
||||
"Qwen-Q4_K_P.gguf",
|
||||
"Qwen-Q5_K_P.gguf",
|
||||
"Qwen-Q6_K_P.gguf",
|
||||
"Qwen-Q8_K_P.gguf",
|
||||
];
|
||||
let fmt = classify(&files);
|
||||
let s = spec("HauhauCS/Qwen3.6", Some(1), Some("q6k"));
|
||||
match decide(&s, &fmt, &files).unwrap_err() {
|
||||
PreflightError::QuantNotFound {
|
||||
requested,
|
||||
nearest,
|
||||
available,
|
||||
..
|
||||
} => {
|
||||
assert_eq!(requested, "q6k");
|
||||
assert_eq!(nearest.as_deref(), Some("q6_k_p"));
|
||||
assert_eq!(available.len(), 4);
|
||||
}
|
||||
other => panic!("expected QuantNotFound, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn feasibility_gguf_single_gpu_good_quant() {
|
||||
let files = ["Qwen-Q4_K_M.gguf", "Qwen-Q6_K.gguf"];
|
||||
let fmt = classify(&files);
|
||||
let s = spec("Qwen/Q-GGUF", Some(1), Some("q6_k"));
|
||||
let plan = decide(&s, &fmt, &files).unwrap();
|
||||
assert_eq!(plan.picked_quant_file.as_deref(), Some("Qwen-Q6_K.gguf"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn feasibility_dense_tp_ok() {
|
||||
let files = [
|
||||
"config.json",
|
||||
"tokenizer.json",
|
||||
"model.safetensors.index.json",
|
||||
"model-00001-of-00002.safetensors",
|
||||
];
|
||||
let fmt = classify(&files);
|
||||
let s = spec("Qwen/Q3-30B", Some(2), Some("q5k"));
|
||||
let plan = decide(&s, &fmt, &files).unwrap();
|
||||
assert_eq!(plan.tp_size, 2);
|
||||
assert!(plan.picked_quant_file.is_none());
|
||||
assert!(matches!(
|
||||
plan.format,
|
||||
SourceFormat::DenseSafetensors { sharded: true }
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn feasibility_empty_rejected() {
|
||||
let files = ["README.md", "tokenizer.json"];
|
||||
let fmt = classify(&files);
|
||||
let s = spec("Empty/Repo", Some(1), None);
|
||||
match decide(&s, &fmt, &files).unwrap_err() {
|
||||
PreflightError::EmptyRepo { model_id } => assert_eq!(model_id, "Empty/Repo"),
|
||||
other => panic!("expected EmptyRepo, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn error_serialization_carries_kind_field() {
|
||||
let err = PreflightError::TpRequiresSafetensors {
|
||||
model_id: "x/y".into(),
|
||||
tp_size: 2,
|
||||
gguf_quants: vec!["q6_k_p".into()],
|
||||
suggestion: "...".into(),
|
||||
};
|
||||
let v: serde_json::Value = serde_json::to_value(&err).unwrap();
|
||||
assert_eq!(v["kind"], "tp_requires_safetensors");
|
||||
assert_eq!(v["model_id"], "x/y");
|
||||
assert_eq!(v["tp_size"], 2);
|
||||
}
|
||||
}
|
||||
@@ -6,3 +6,4 @@ pub mod discovery;
|
||||
pub mod harness;
|
||||
pub mod health;
|
||||
pub mod startup;
|
||||
pub mod wire;
|
||||
|
||||
@@ -7,6 +7,7 @@
|
||||
|
||||
use crate::activation::ActivationTracker;
|
||||
use crate::harness::HarnessRegistry;
|
||||
use crate::harness::preflight::PreflightError;
|
||||
use cortex_core::harness::ModelSpec;
|
||||
use std::time::{Duration, Instant};
|
||||
use tokio::signal;
|
||||
@@ -53,18 +54,45 @@ pub async fn load_default_models(
|
||||
Err(e) => {
|
||||
let rendered = format!("{e:#}");
|
||||
activation.fail_loading(&spec.model_id, &rendered).await;
|
||||
tracing::warn!(
|
||||
model = %spec.model_id,
|
||||
error = %rendered,
|
||||
elapsed_ms = start.elapsed().as_millis() as u64,
|
||||
"failed to load default model, continuing"
|
||||
);
|
||||
// When the underlying failure is a preflight rejection,
|
||||
// pull the structured fields out so journalctl shows
|
||||
// `reason=tp_requires_safetensors detail="..."` instead
|
||||
// of an opaque "fetch config.json … 404". The operator
|
||||
// can act on the structured form directly.
|
||||
if let Some(pf) = e.downcast_ref::<PreflightError>() {
|
||||
tracing::warn!(
|
||||
model = %spec.model_id,
|
||||
reason = preflight_kind(pf),
|
||||
detail = %pf,
|
||||
elapsed_ms = start.elapsed().as_millis() as u64,
|
||||
"failed to load default model, continuing"
|
||||
);
|
||||
} else {
|
||||
tracing::warn!(
|
||||
model = %spec.model_id,
|
||||
error = %rendered,
|
||||
elapsed_ms = start.elapsed().as_millis() as u64,
|
||||
"failed to load default model, continuing"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
activation.mark_ready().await;
|
||||
}
|
||||
|
||||
/// Short kebab-case tag for a preflight failure. Used as a structured
|
||||
/// log field so journalctl filtering can match on the failure class
|
||||
/// (`reason=tp_requires_safetensors`, `reason=quant_not_found`, etc.).
|
||||
fn preflight_kind(err: &PreflightError) -> &'static str {
|
||||
match err {
|
||||
PreflightError::RepoFetchFailed { .. } => "repo_fetch_failed",
|
||||
PreflightError::EmptyRepo { .. } => "empty_repo",
|
||||
PreflightError::TpRequiresSafetensors { .. } => "tp_requires_safetensors",
|
||||
PreflightError::QuantNotFound { .. } => "quant_not_found",
|
||||
}
|
||||
}
|
||||
|
||||
/// Future that resolves on SIGINT (Ctrl-C) or SIGTERM (systemd stop).
|
||||
///
|
||||
/// Wired into `axum::serve(...).with_graceful_shutdown(shutdown_signal())`
|
||||
|
||||
306
crates/neuron/src/wire/event.rs
Normal file
306
crates/neuron/src/wire/event.rs
Normal file
@@ -0,0 +1,306 @@
|
||||
//! Format-agnostic inference event stream.
|
||||
//!
|
||||
//! The candle harness emits a sequence of these for every streaming
|
||||
//! request. Wire-format projections in sibling modules
|
||||
//! ([`super::openai_chat`], the eventual `openai_responses` /
|
||||
//! `anthropic_messages` projections) read this stream and produce
|
||||
//! the chunks / events their HTTP clients expect.
|
||||
//!
|
||||
//! Design notes:
|
||||
//!
|
||||
//! - [`Start`] carries no token of its own. It only signals "the
|
||||
//! model has accepted the prompt and is about to begin emitting
|
||||
//! text". OpenAI chat materialises this as a `role: assistant`
|
||||
//! chunk; OpenAI Responses as the `response.created` +
|
||||
//! `response.output_item.added` pair; Anthropic as
|
||||
//! `message_start`. All three of those would otherwise have to
|
||||
//! peek at the *first* token to know when to emit, which couples
|
||||
//! the wire layer to the producer's pacing.
|
||||
//! - [`TextDelta`] is *visible* output. Reasoning / `<think>`
|
||||
//! blocks go through a future [`ReasoningDelta`] variant once
|
||||
//! the harness learns to split them (today they pass through as
|
||||
//! plain text inside `TextDelta`; helexa-acp picks them apart on
|
||||
//! the consumer side).
|
||||
//! - [`Finish`] is the only place a stream is allowed to end
|
||||
//! cleanly. Projections rely on this to emit final usage
|
||||
//! bookkeeping; absence means the producer crashed and the
|
||||
//! consumer should treat the stream as truncated.
|
||||
//!
|
||||
//! [`Start`]: InferenceEvent::Start
|
||||
//! [`TextDelta`]: InferenceEvent::TextDelta
|
||||
//! [`Finish`]: InferenceEvent::Finish
|
||||
|
||||
/// One unit of output from the inference loop.
|
||||
///
|
||||
/// Producers send these on an `mpsc::Sender<InferenceEvent>`;
|
||||
/// projection layers in sibling modules consume them and emit
|
||||
/// wire-format-specific frames downstream.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum InferenceEvent {
|
||||
/// The producer has accepted the prompt and is about to emit
|
||||
/// the first token. Sent at most once per stream.
|
||||
Start,
|
||||
/// A piece of visible assistant text. Multiple deltas
|
||||
/// concatenate into the complete reply.
|
||||
TextDelta(String),
|
||||
/// Reasoning / scratchpad text the model emitted inside a
|
||||
/// `<think>` block (or equivalent). The harness routes
|
||||
/// content between marker tokens here so wire projectors can
|
||||
/// decide what to do with it (chat completions drops by
|
||||
/// default; Responses API has a dedicated event family).
|
||||
ReasoningDelta(String),
|
||||
/// A tool call has been parsed out of a `<tool_call>{json}</tool_call>`
|
||||
/// block. Carries the parsed name + arguments JSON string
|
||||
/// (Anthropic / OpenAI projectors emit their own wire shape
|
||||
/// from this).
|
||||
///
|
||||
/// `index` is the call slot — incremented per tool call in a
|
||||
/// turn so wire formats that order calls by index
|
||||
/// (OpenAI chat completions) can correlate.
|
||||
ToolCall {
|
||||
index: usize,
|
||||
id: String,
|
||||
name: String,
|
||||
/// Complete JSON arguments string. The model could in
|
||||
/// principle stream these token-by-token, but our
|
||||
/// extraction buffers the whole block until `</tool_call>`
|
||||
/// arrives and emits exactly one event per call.
|
||||
arguments: String,
|
||||
},
|
||||
/// The stream is complete. Carries the reason so wire formats
|
||||
/// that use it (OpenAI's `finish_reason`, Anthropic's
|
||||
/// `stop_reason`) can render it without re-parsing.
|
||||
Finish { reason: FinishReason },
|
||||
}
|
||||
|
||||
/// Why a stream stopped. Stays small on purpose — anything that
|
||||
/// doesn't map cleanly to one of these collapses to [`Stop`].
|
||||
///
|
||||
/// Mappings to wire formats:
|
||||
///
|
||||
/// | variant | OpenAI `finish_reason` | OpenAI Responses `status` | Anthropic `stop_reason` |
|
||||
/// |---------|------------------------|---------------------------|-------------------------|
|
||||
/// | `Stop` | `"stop"` | `"completed"` | `"end_turn"` |
|
||||
/// | `Length`| `"length"` | `"incomplete"` | `"max_tokens"` |
|
||||
/// | `ToolCalls` | `"tool_calls"` | `"completed"` | `"tool_use"` |
|
||||
///
|
||||
/// [`Stop`]: FinishReason::Stop
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum FinishReason {
|
||||
/// Model emitted EOS naturally.
|
||||
Stop,
|
||||
/// Hit `max_tokens` before EOS.
|
||||
Length,
|
||||
/// Stopped because the model called a tool and is waiting for
|
||||
/// the result. Not yet emitted by the candle harness —
|
||||
/// reserved for the day tool-call extraction lands.
|
||||
#[allow(dead_code)]
|
||||
ToolCalls,
|
||||
}
|
||||
|
||||
impl FinishReason {
|
||||
/// String form used by OpenAI chat completions and OpenAI
|
||||
/// completions. Wire modules can call this directly or do their
|
||||
/// own mapping for non-string formats.
|
||||
pub fn as_openai_str(self) -> &'static str {
|
||||
match self {
|
||||
FinishReason::Stop => "stop",
|
||||
FinishReason::Length => "length",
|
||||
FinishReason::ToolCalls => "tool_calls",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Open/close token IDs for the reasoning marker a loaded model uses
|
||||
/// (or `None` for non-reasoning models). The harness reads this once
|
||||
/// at load time from the tokenizer's added-tokens table, then the
|
||||
/// inference loop checks `next_token` against the pair to flip
|
||||
/// between [`InferenceEvent::TextDelta`] and
|
||||
/// [`InferenceEvent::ReasoningDelta`].
|
||||
///
|
||||
/// `open` and `close` text are kept alongside the IDs so wire
|
||||
/// projectors that want to re-emit the literal markers (the
|
||||
/// opt-in `include_thinking` path on chat completions) don't have
|
||||
/// to reach back into the tokenizer for the strings.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ReasoningTokenPair {
|
||||
pub open_id: u32,
|
||||
pub close_id: u32,
|
||||
pub open_text: String,
|
||||
pub close_text: String,
|
||||
}
|
||||
|
||||
/// Known reasoning-marker conventions. Each is a `(open, close)`
|
||||
/// pair of literal token strings. Each modern reasoning model
|
||||
/// declares its markers in the tokenizer's `added_tokens` table;
|
||||
/// at load time we probe for whichever pair the loaded tokenizer
|
||||
/// has and stash both IDs.
|
||||
///
|
||||
/// Ordering matters only for tie-breaking when a model declares
|
||||
/// multiple pairs (shouldn't happen in practice); the first hit
|
||||
/// wins.
|
||||
const KNOWN_REASONING_MARKERS: &[(&str, &str)] = &[
|
||||
// Qwen3, DeepSeek-R1, gpt-oss, and most other open-weight
|
||||
// reasoning models.
|
||||
("<think>", "</think>"),
|
||||
// Mistral Magistral.
|
||||
("[THINK]", "[/THINK]"),
|
||||
// Some older derivatives; harmless to probe.
|
||||
("<thought>", "</thought>"),
|
||||
("<reasoning>", "</reasoning>"),
|
||||
];
|
||||
|
||||
/// Open/close token IDs for the model's tool-call marker
|
||||
/// convention (or `None` for models that don't emit structured
|
||||
/// tool calls). Same shape as [`ReasoningTokenPair`]: probed once
|
||||
/// at load time, consumed by the inference loop to switch between
|
||||
/// "emit visible deltas" and "buffer JSON for the next tool
|
||||
/// call".
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ToolCallTokenPair {
|
||||
pub open_id: u32,
|
||||
pub close_id: u32,
|
||||
pub open_text: String,
|
||||
pub close_text: String,
|
||||
}
|
||||
|
||||
/// Tool-call marker conventions. Open-weight tool-use models
|
||||
/// converged on `<tool_call>` / `</tool_call>` (Qwen3-Coder /
|
||||
/// -Instruct, the Hermes function-call format, DeepSeek-Coder,
|
||||
/// gpt-oss). The pair lives alongside the reasoning markers in
|
||||
/// the same `added_tokens` table.
|
||||
const KNOWN_TOOL_CALL_MARKERS: &[(&str, &str)] = &[("<tool_call>", "</tool_call>")];
|
||||
|
||||
/// Probe a tokenizer for known tool-call marker pairs. Mirrors
|
||||
/// [`detect_reasoning_token_pair`] — both open AND close must
|
||||
/// resolve for the pair to be returned. `None` means the model
|
||||
/// doesn't emit structured tool calls (or its tokenizer split
|
||||
/// the markers across tokens).
|
||||
pub fn detect_tool_call_token_pair<F>(token_to_id: F) -> Option<ToolCallTokenPair>
|
||||
where
|
||||
F: Fn(&str) -> Option<u32>,
|
||||
{
|
||||
for (open_text, close_text) in KNOWN_TOOL_CALL_MARKERS {
|
||||
let open_id = token_to_id(open_text);
|
||||
let close_id = token_to_id(close_text);
|
||||
if let (Some(open_id), Some(close_id)) = (open_id, close_id) {
|
||||
return Some(ToolCallTokenPair {
|
||||
open_id,
|
||||
close_id,
|
||||
open_text: (*open_text).into(),
|
||||
close_text: (*close_text).into(),
|
||||
});
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Inspect a tokenizer for known reasoning-marker pairs and return
|
||||
/// the first match. The tokenizer types this trait is defined over
|
||||
/// just need to expose `token_to_id(&str) -> Option<u32>` so this
|
||||
/// stays decoupled from the candle crate — the production caller
|
||||
/// passes a `tokenizers::Tokenizer`, but tests can fake one.
|
||||
///
|
||||
/// Returns `None` when no known marker pair is fully declared
|
||||
/// (both open AND close token ids must resolve). That's the
|
||||
/// pass-through case — non-reasoning models, or reasoning models
|
||||
/// whose tokenizer split the markers across multiple tokens (rare
|
||||
/// in practice; modern reasoning tokenizers list them as
|
||||
/// `added_tokens`).
|
||||
pub fn detect_reasoning_token_pair<F>(token_to_id: F) -> Option<ReasoningTokenPair>
|
||||
where
|
||||
F: Fn(&str) -> Option<u32>,
|
||||
{
|
||||
for (open_text, close_text) in KNOWN_REASONING_MARKERS {
|
||||
let open_id = token_to_id(open_text);
|
||||
let close_id = token_to_id(close_text);
|
||||
if let (Some(open_id), Some(close_id)) = (open_id, close_id) {
|
||||
return Some(ReasoningTokenPair {
|
||||
open_id,
|
||||
close_id,
|
||||
open_text: (*open_text).into(),
|
||||
close_text: (*close_text).into(),
|
||||
});
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::collections::HashMap;
|
||||
|
||||
fn lookup<'a>(map: &'a HashMap<&'static str, u32>) -> impl Fn(&str) -> Option<u32> + 'a {
|
||||
|s| map.get(s).copied()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detects_qwen3_style_think_markers() {
|
||||
let mut m = HashMap::new();
|
||||
m.insert("<think>", 151648);
|
||||
m.insert("</think>", 151649);
|
||||
let pair = detect_reasoning_token_pair(lookup(&m)).expect("pair detected");
|
||||
assert_eq!(pair.open_id, 151648);
|
||||
assert_eq!(pair.close_id, 151649);
|
||||
assert_eq!(pair.open_text, "<think>");
|
||||
assert_eq!(pair.close_text, "</think>");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detects_mistral_magistral_markers() {
|
||||
let mut m = HashMap::new();
|
||||
m.insert("[THINK]", 100);
|
||||
m.insert("[/THINK]", 101);
|
||||
let pair = detect_reasoning_token_pair(lookup(&m)).expect("pair detected");
|
||||
assert_eq!(pair.open_text, "[THINK]");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn returns_none_when_only_open_marker_present() {
|
||||
// A pathological tokenizer that has `<think>` but not
|
||||
// `</think>` shouldn't half-detect. Pass-through.
|
||||
let mut m = HashMap::new();
|
||||
m.insert("<think>", 1);
|
||||
assert!(detect_reasoning_token_pair(lookup(&m)).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn returns_none_for_non_reasoning_tokenizer() {
|
||||
let m: HashMap<&'static str, u32> = HashMap::new();
|
||||
assert!(detect_reasoning_token_pair(lookup(&m)).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detects_tool_call_markers() {
|
||||
let mut m = HashMap::new();
|
||||
m.insert("<tool_call>", 151657);
|
||||
m.insert("</tool_call>", 151658);
|
||||
let pair = detect_tool_call_token_pair(lookup(&m)).expect("pair detected");
|
||||
assert_eq!(pair.open_id, 151657);
|
||||
assert_eq!(pair.close_id, 151658);
|
||||
assert_eq!(pair.open_text, "<tool_call>");
|
||||
assert_eq!(pair.close_text, "</tool_call>");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn returns_none_for_non_tool_use_tokenizer() {
|
||||
let m: HashMap<&'static str, u32> = HashMap::new();
|
||||
assert!(detect_tool_call_token_pair(lookup(&m)).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn first_match_wins_when_multiple_pairs_declared() {
|
||||
// Hypothetical tokenizer with both Qwen-style AND Mistral-style
|
||||
// markers — the `<think>` pair is earlier in the convention
|
||||
// table so it wins.
|
||||
let mut m = HashMap::new();
|
||||
m.insert("<think>", 1);
|
||||
m.insert("</think>", 2);
|
||||
m.insert("[THINK]", 3);
|
||||
m.insert("[/THINK]", 4);
|
||||
let pair = detect_reasoning_token_pair(lookup(&m)).unwrap();
|
||||
assert_eq!(pair.open_id, 1);
|
||||
assert_eq!(pair.close_id, 2);
|
||||
}
|
||||
}
|
||||
27
crates/neuron/src/wire/mod.rs
Normal file
27
crates/neuron/src/wire/mod.rs
Normal file
@@ -0,0 +1,27 @@
|
||||
//! Wire-format projection layer.
|
||||
//!
|
||||
//! The candle harness produces a single, format-agnostic stream of
|
||||
//! [`InferenceEvent`]s. Each wire format (OpenAI chat completions,
|
||||
//! OpenAI Responses, Anthropic messages, …) lives in its own module
|
||||
//! under `wire::` and projects that event stream into the chunks /
|
||||
//! events its HTTP clients expect.
|
||||
//!
|
||||
//! The benefit over translating *between* wire shapes (OpenAI chat
|
||||
//! → Anthropic, etc.) is that we never have to reason about a
|
||||
//! wire-N → wire-M conversion: every translation is wire-N ↔ the
|
||||
//! internal event currency, and the projections are independent. A
|
||||
//! new wire format adds a new file under `wire::`; nothing else
|
||||
//! needs to know about it.
|
||||
//!
|
||||
//! Today: [`openai_chat`]. Stage 2 adds `openai_responses`. Stage 3
|
||||
//! could add a native Anthropic projection that replaces the
|
||||
//! gateway-side translation.
|
||||
|
||||
pub mod event;
|
||||
pub mod openai_chat;
|
||||
pub mod openai_responses;
|
||||
|
||||
pub use event::{
|
||||
FinishReason, InferenceEvent, ReasoningTokenPair, ToolCallTokenPair,
|
||||
detect_reasoning_token_pair, detect_tool_call_token_pair,
|
||||
};
|
||||
558
crates/neuron/src/wire/openai_chat.rs
Normal file
558
crates/neuron/src/wire/openai_chat.rs
Normal file
@@ -0,0 +1,558 @@
|
||||
//! OpenAI chat completions projection.
|
||||
//!
|
||||
//! Reads [`InferenceEvent`]s from a receiver and produces
|
||||
//! [`ChatCompletionChunk`]s in the shape `POST /v1/chat/completions`
|
||||
//! clients expect on its streaming SSE response. The HTTP handler in
|
||||
//! [`crate::api`] wraps the resulting receiver in axum's
|
||||
//! `Sse::new(...)` adapter; nothing in this module touches HTTP
|
||||
//! framing or `data:` lines.
|
||||
//!
|
||||
//! Per the OpenAI streaming spec, three chunk shapes appear:
|
||||
//!
|
||||
//! 1. **Role chunk** — `delta: { "role": "assistant" }`, no content,
|
||||
//! sent once at stream start. We emit this on [`InferenceEvent::Start`].
|
||||
//! 2. **Content chunks** — `delta: { "content": "<text>" }`, one per
|
||||
//! [`InferenceEvent::TextDelta`].
|
||||
//! 3. **Final chunk** — empty `delta`, `finish_reason` populated.
|
||||
//! Emitted on [`InferenceEvent::Finish`].
|
||||
//!
|
||||
//! `usage` stays `None` on every chunk; the legacy candle paths
|
||||
//! never surfaced usage on the streaming endpoint and we keep that
|
||||
//! behaviour bit-for-bit so existing clients see no diff.
|
||||
//!
|
||||
//! Back-pressure: the projection task awaits both `rx.recv()` and
|
||||
//! `tx.send()`. A slow consumer fills the output channel → the
|
||||
//! task blocks on send → it stops reading from the input → the
|
||||
//! producer blocks on its own send. The bounded channels
|
||||
//! propagate without us writing any logic.
|
||||
|
||||
use cortex_core::openai::{ChatCompletionChunk, ChunkChoice};
|
||||
use serde_json::json;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use super::event::{FinishReason, InferenceEvent, ReasoningTokenPair};
|
||||
|
||||
/// Output channel buffer size. Mirrors the input side's bound; one
|
||||
/// event maps to at most one chunk, so equal capacity keeps the
|
||||
/// two ends in sync without surprising memory growth.
|
||||
const CHUNK_CHANNEL_CAPACITY: usize = 32;
|
||||
|
||||
/// Per-stream config for the chat projector. Used by the
|
||||
/// production handler to thread per-request choices (currently:
|
||||
/// whether to surface reasoning content) into the projection
|
||||
/// without bloating the function signature.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct ChatProjectionConfig {
|
||||
/// When `true`, reasoning content is re-wrapped with the
|
||||
/// model's literal open/close markers and emitted as content
|
||||
/// deltas — preserving the on-the-wire shape that
|
||||
/// reasoning-aware clients like helexa-acp's `ThinkParser`
|
||||
/// expect.
|
||||
///
|
||||
/// When `false` (the default), [`InferenceEvent::ReasoningDelta`]s
|
||||
/// are dropped entirely so consumers that don't know about
|
||||
/// reasoning (Zed's commit-message generator, any vanilla
|
||||
/// OpenAI client) don't have model-internal scratchpad
|
||||
/// material leaking into their UI. The chat-completions wire
|
||||
/// format has no slot for reasoning, so the default chooses
|
||||
/// the safer-for-naïve-clients behaviour.
|
||||
pub include_thinking: bool,
|
||||
/// Open/close marker strings to re-emit when `include_thinking`
|
||||
/// is set. Sourced from the loaded model's
|
||||
/// [`ReasoningTokenPair`]; `None` for non-reasoning models or
|
||||
/// when the caller doesn't have the pair handy (in which case
|
||||
/// `include_thinking` becomes equivalent to dropping reasoning
|
||||
/// because there's nothing to wrap).
|
||||
pub reasoning_markers: Option<ReasoningTokenPair>,
|
||||
}
|
||||
|
||||
/// Project an [`InferenceEvent`] receiver into a
|
||||
/// [`ChatCompletionChunk`] receiver. Spawns one tokio task that
|
||||
/// owns the input receiver for the stream's lifetime and exits
|
||||
/// when either side closes.
|
||||
///
|
||||
/// `id`, `created`, and `model_id` are stamped into every emitted
|
||||
/// chunk so the receiver can stay generic (decoupled from
|
||||
/// per-request metadata).
|
||||
pub fn project_chat_stream(
|
||||
rx: mpsc::Receiver<InferenceEvent>,
|
||||
id: String,
|
||||
created: u64,
|
||||
model_id: String,
|
||||
) -> mpsc::Receiver<ChatCompletionChunk> {
|
||||
// Default config: include_thinking off, no marker rewrap.
|
||||
project_chat_stream_with(rx, id, created, model_id, ChatProjectionConfig::default())
|
||||
}
|
||||
|
||||
/// Same as [`project_chat_stream`] but with a per-stream config
|
||||
/// (currently controlling reasoning surfacing). Production
|
||||
/// callers that need the opt-in path call this directly; the
|
||||
/// shorter wrapper above stays as the no-config convenience.
|
||||
pub fn project_chat_stream_with(
|
||||
mut rx: mpsc::Receiver<InferenceEvent>,
|
||||
id: String,
|
||||
created: u64,
|
||||
model_id: String,
|
||||
config: ChatProjectionConfig,
|
||||
) -> mpsc::Receiver<ChatCompletionChunk> {
|
||||
let (tx, out_rx) = mpsc::channel::<ChatCompletionChunk>(CHUNK_CHANNEL_CAPACITY);
|
||||
|
||||
tokio::spawn(async move {
|
||||
// Track whether the previous event was inside a reasoning
|
||||
// block — used to decide when to emit the literal close
|
||||
// marker on the include_thinking re-wrap path. When this
|
||||
// flips from true → false (a TextDelta or Finish lands
|
||||
// after one or more ReasoningDeltas), we emit the close
|
||||
// marker exactly once.
|
||||
let mut was_in_reasoning = false;
|
||||
|
||||
while let Some(event) = rx.recv().await {
|
||||
// Close-marker insertion: if we're leaving a reasoning
|
||||
// chain, emit the literal close marker before the
|
||||
// current event.
|
||||
if was_in_reasoning && !matches!(event, InferenceEvent::ReasoningDelta(_)) {
|
||||
if let Some(marker) = config
|
||||
.include_thinking
|
||||
.then_some(())
|
||||
.and(config.reasoning_markers.as_ref())
|
||||
{
|
||||
let chunk = content_chunk(&id, created, &model_id, &marker.close_text);
|
||||
if tx.send(chunk).await.is_err() {
|
||||
return;
|
||||
}
|
||||
}
|
||||
was_in_reasoning = false;
|
||||
}
|
||||
|
||||
let chunks = match event {
|
||||
InferenceEvent::Start => vec![role_chunk(&id, created, &model_id)],
|
||||
InferenceEvent::TextDelta(text) => {
|
||||
if text.is_empty() {
|
||||
// DecodeStream is buffering a multi-byte
|
||||
// codepoint; don't bother sending an empty
|
||||
// chunk downstream.
|
||||
continue;
|
||||
}
|
||||
vec![content_chunk(&id, created, &model_id, &text)]
|
||||
}
|
||||
InferenceEvent::ReasoningDelta(text) => {
|
||||
if !config.include_thinking {
|
||||
// Default path — reasoning has no slot in
|
||||
// chat completions, so it's dropped. Naïve
|
||||
// clients (Zed commit-message generator,
|
||||
// any vanilla OpenAI client) get clean
|
||||
// output.
|
||||
continue;
|
||||
}
|
||||
let Some(markers) = config.reasoning_markers.as_ref() else {
|
||||
// Caller asked to include thinking but
|
||||
// didn't supply markers — best we can do
|
||||
// is emit the content as visible text.
|
||||
// Skip the wrap entirely.
|
||||
if text.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let chunk = content_chunk(&id, created, &model_id, &text);
|
||||
if tx.send(chunk).await.is_err() {
|
||||
return;
|
||||
}
|
||||
continue;
|
||||
};
|
||||
// First chunk of a reasoning block → open
|
||||
// marker prelude. Subsequent reasoning deltas
|
||||
// in the same block reuse `was_in_reasoning`
|
||||
// to skip the prelude.
|
||||
let mut chunks = Vec::new();
|
||||
if !was_in_reasoning {
|
||||
chunks.push(content_chunk(&id, created, &model_id, &markers.open_text));
|
||||
}
|
||||
if !text.is_empty() {
|
||||
chunks.push(content_chunk(&id, created, &model_id, &text));
|
||||
}
|
||||
was_in_reasoning = true;
|
||||
chunks
|
||||
}
|
||||
InferenceEvent::ToolCall {
|
||||
index,
|
||||
id: call_id,
|
||||
name,
|
||||
arguments,
|
||||
} => {
|
||||
// OpenAI streaming shape for tool calls:
|
||||
// `delta.tool_calls[]` with id + function.name
|
||||
// on the first chunk per index, then
|
||||
// function.arguments deltas. We have the
|
||||
// complete arguments buffered already, so one
|
||||
// delta carries everything.
|
||||
vec![tool_call_chunk(
|
||||
&id, created, &model_id, index, &call_id, &name, &arguments,
|
||||
)]
|
||||
}
|
||||
InferenceEvent::Finish { reason } => {
|
||||
vec![final_chunk(&id, created, &model_id, reason)]
|
||||
}
|
||||
};
|
||||
for chunk in chunks {
|
||||
if tx.send(chunk).await.is_err() {
|
||||
// Consumer hung up; nothing more to do.
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
out_rx
|
||||
}
|
||||
|
||||
fn role_chunk(id: &str, created: u64, model_id: &str) -> ChatCompletionChunk {
|
||||
ChatCompletionChunk {
|
||||
id: id.into(),
|
||||
object: "chat.completion.chunk".into(),
|
||||
created,
|
||||
model: model_id.into(),
|
||||
choices: vec![ChunkChoice {
|
||||
index: 0,
|
||||
delta: json!({ "role": "assistant" }),
|
||||
finish_reason: None,
|
||||
extra: serde_json::Value::Object(Default::default()),
|
||||
}],
|
||||
usage: None,
|
||||
extra: serde_json::Value::Object(Default::default()),
|
||||
}
|
||||
}
|
||||
|
||||
fn content_chunk(id: &str, created: u64, model_id: &str, text: &str) -> ChatCompletionChunk {
|
||||
ChatCompletionChunk {
|
||||
id: id.into(),
|
||||
object: "chat.completion.chunk".into(),
|
||||
created,
|
||||
model: model_id.into(),
|
||||
choices: vec![ChunkChoice {
|
||||
index: 0,
|
||||
delta: json!({ "content": text }),
|
||||
finish_reason: None,
|
||||
extra: serde_json::Value::Object(Default::default()),
|
||||
}],
|
||||
usage: None,
|
||||
extra: serde_json::Value::Object(Default::default()),
|
||||
}
|
||||
}
|
||||
|
||||
/// OpenAI chat streaming shape for a tool call. One chunk per
|
||||
/// call slot, carrying id + name + the complete arguments JSON.
|
||||
/// Mirrors the format real OpenAI emits on the streaming path,
|
||||
/// minus the per-token arguments-streaming complication (we have
|
||||
/// the whole buffer already after the model finishes the
|
||||
/// `<tool_call>...</tool_call>` block).
|
||||
fn tool_call_chunk(
|
||||
id: &str,
|
||||
created: u64,
|
||||
model_id: &str,
|
||||
index: usize,
|
||||
call_id: &str,
|
||||
name: &str,
|
||||
arguments: &str,
|
||||
) -> ChatCompletionChunk {
|
||||
ChatCompletionChunk {
|
||||
id: id.into(),
|
||||
object: "chat.completion.chunk".into(),
|
||||
created,
|
||||
model: model_id.into(),
|
||||
choices: vec![ChunkChoice {
|
||||
index: 0,
|
||||
delta: json!({
|
||||
"tool_calls": [{
|
||||
"index": index,
|
||||
"id": call_id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": name,
|
||||
"arguments": arguments,
|
||||
}
|
||||
}],
|
||||
}),
|
||||
finish_reason: None,
|
||||
extra: serde_json::Value::Object(Default::default()),
|
||||
}],
|
||||
usage: None,
|
||||
extra: serde_json::Value::Object(Default::default()),
|
||||
}
|
||||
}
|
||||
|
||||
fn final_chunk(
|
||||
id: &str,
|
||||
created: u64,
|
||||
model_id: &str,
|
||||
reason: FinishReason,
|
||||
) -> ChatCompletionChunk {
|
||||
ChatCompletionChunk {
|
||||
id: id.into(),
|
||||
object: "chat.completion.chunk".into(),
|
||||
created,
|
||||
model: model_id.into(),
|
||||
choices: vec![ChunkChoice {
|
||||
index: 0,
|
||||
delta: serde_json::Value::Object(Default::default()),
|
||||
finish_reason: Some(reason.as_openai_str().to_string()),
|
||||
extra: serde_json::Value::Object(Default::default()),
|
||||
}],
|
||||
usage: None,
|
||||
extra: serde_json::Value::Object(Default::default()),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// Drain the projection's output into a Vec for assertion.
|
||||
async fn collect(mut rx: mpsc::Receiver<ChatCompletionChunk>) -> Vec<ChatCompletionChunk> {
|
||||
let mut out = Vec::new();
|
||||
while let Some(chunk) = rx.recv().await {
|
||||
out.push(chunk);
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn empty_event_stream_yields_no_chunks() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(4);
|
||||
drop(tx);
|
||||
let out = collect(project_chat_stream(rx, "id-1".into(), 1700, "m".into())).await;
|
||||
assert!(out.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn start_text_finish_produces_three_chunks() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(4);
|
||||
let out_rx = project_chat_stream(rx, "id-1".into(), 1700, "m".into());
|
||||
|
||||
tx.send(InferenceEvent::Start).await.unwrap();
|
||||
tx.send(InferenceEvent::TextDelta("hello".into()))
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::Finish {
|
||||
reason: FinishReason::Stop,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
|
||||
let out = collect(out_rx).await;
|
||||
assert_eq!(out.len(), 3);
|
||||
assert_eq!(out[0].choices[0].delta["role"], "assistant");
|
||||
assert_eq!(out[1].choices[0].delta["content"], "hello");
|
||||
assert_eq!(out[2].choices[0].finish_reason.as_deref(), Some("stop"));
|
||||
// Every chunk carries the stamped metadata.
|
||||
for chunk in &out {
|
||||
assert_eq!(chunk.id, "id-1");
|
||||
assert_eq!(chunk.created, 1700);
|
||||
assert_eq!(chunk.model, "m");
|
||||
assert_eq!(chunk.object, "chat.completion.chunk");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn empty_text_delta_is_dropped() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(4);
|
||||
let out_rx = project_chat_stream(rx, "id".into(), 1, "m".into());
|
||||
tx.send(InferenceEvent::TextDelta(String::new()))
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
let out = collect(out_rx).await;
|
||||
assert!(out.is_empty(), "empty deltas must not produce chunks");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn finish_length_maps_to_openai_string() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(4);
|
||||
let out_rx = project_chat_stream(rx, "id".into(), 1, "m".into());
|
||||
tx.send(InferenceEvent::Finish {
|
||||
reason: FinishReason::Length,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
let out = collect(out_rx).await;
|
||||
assert_eq!(out.len(), 1);
|
||||
assert_eq!(out[0].choices[0].finish_reason.as_deref(), Some("length"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reasoning_delta_is_dropped_in_chat_projection() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(4);
|
||||
let out_rx = project_chat_stream(rx, "id".into(), 1, "m".into());
|
||||
tx.send(InferenceEvent::ReasoningDelta("<think>".into()))
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::TextDelta("real".into()))
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
let out = collect(out_rx).await;
|
||||
assert_eq!(out.len(), 1);
|
||||
assert_eq!(out[0].choices[0].delta["content"], "real");
|
||||
}
|
||||
|
||||
fn pair() -> ReasoningTokenPair {
|
||||
ReasoningTokenPair {
|
||||
open_id: 0,
|
||||
close_id: 1,
|
||||
open_text: "<think>".into(),
|
||||
close_text: "</think>".into(),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn include_thinking_rewraps_reasoning_with_literal_markers() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(8);
|
||||
let out_rx = project_chat_stream_with(
|
||||
rx,
|
||||
"id".into(),
|
||||
1,
|
||||
"m".into(),
|
||||
ChatProjectionConfig {
|
||||
include_thinking: true,
|
||||
reasoning_markers: Some(pair()),
|
||||
},
|
||||
);
|
||||
tx.send(InferenceEvent::ReasoningDelta("first ".into()))
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::ReasoningDelta("second".into()))
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::TextDelta("answer".into()))
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::Finish {
|
||||
reason: FinishReason::Stop,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
let out = collect(out_rx).await;
|
||||
// Expected sequence: open marker → reasoning content (2 chunks)
|
||||
// → close marker → visible answer → final chunk.
|
||||
let contents: Vec<&str> = out
|
||||
.iter()
|
||||
.filter_map(|c| c.choices[0].delta["content"].as_str())
|
||||
.collect();
|
||||
assert_eq!(
|
||||
contents,
|
||||
vec!["<think>", "first ", "second", "</think>", "answer"]
|
||||
);
|
||||
assert_eq!(
|
||||
out.last().unwrap().choices[0].finish_reason.as_deref(),
|
||||
Some("stop")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn include_thinking_closes_marker_at_finish_when_no_trailing_text() {
|
||||
// Edge case: stream ends inside a reasoning block (model
|
||||
// hit max_tokens mid-thought, no visible answer ever).
|
||||
// The Finish event still triggers the close marker so the
|
||||
// stream is balanced.
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(4);
|
||||
let out_rx = project_chat_stream_with(
|
||||
rx,
|
||||
"id".into(),
|
||||
1,
|
||||
"m".into(),
|
||||
ChatProjectionConfig {
|
||||
include_thinking: true,
|
||||
reasoning_markers: Some(pair()),
|
||||
},
|
||||
);
|
||||
tx.send(InferenceEvent::ReasoningDelta("thinking...".into()))
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::Finish {
|
||||
reason: FinishReason::Length,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
let out = collect(out_rx).await;
|
||||
let contents: Vec<&str> = out
|
||||
.iter()
|
||||
.filter_map(|c| c.choices[0].delta["content"].as_str())
|
||||
.collect();
|
||||
assert_eq!(contents, vec!["<think>", "thinking...", "</think>"]);
|
||||
assert_eq!(
|
||||
out.last().unwrap().choices[0].finish_reason.as_deref(),
|
||||
Some("length")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn include_thinking_without_markers_emits_content_directly() {
|
||||
// Defensive: if the caller asks for thinking but the
|
||||
// model declared no markers, we still emit the content
|
||||
// rather than dropping it. Better to leak than to lose.
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(4);
|
||||
let out_rx = project_chat_stream_with(
|
||||
rx,
|
||||
"id".into(),
|
||||
1,
|
||||
"m".into(),
|
||||
ChatProjectionConfig {
|
||||
include_thinking: true,
|
||||
reasoning_markers: None,
|
||||
},
|
||||
);
|
||||
tx.send(InferenceEvent::ReasoningDelta("raw".into()))
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::Finish {
|
||||
reason: FinishReason::Stop,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
let out = collect(out_rx).await;
|
||||
let contents: Vec<&str> = out
|
||||
.iter()
|
||||
.filter_map(|c| c.choices[0].delta["content"].as_str())
|
||||
.collect();
|
||||
assert_eq!(contents, vec!["raw"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn include_thinking_off_drops_reasoning_even_with_markers() {
|
||||
// Default behaviour even when markers happen to be
|
||||
// configured. The flag is the gate, not the marker
|
||||
// presence.
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(4);
|
||||
let out_rx = project_chat_stream_with(
|
||||
rx,
|
||||
"id".into(),
|
||||
1,
|
||||
"m".into(),
|
||||
ChatProjectionConfig {
|
||||
include_thinking: false,
|
||||
reasoning_markers: Some(pair()),
|
||||
},
|
||||
);
|
||||
tx.send(InferenceEvent::ReasoningDelta("hidden".into()))
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::TextDelta("visible".into()))
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::Finish {
|
||||
reason: FinishReason::Stop,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
let out = collect(out_rx).await;
|
||||
let contents: Vec<&str> = out
|
||||
.iter()
|
||||
.filter_map(|c| c.choices[0].delta["content"].as_str())
|
||||
.collect();
|
||||
assert_eq!(contents, vec!["visible"]);
|
||||
}
|
||||
}
|
||||
870
crates/neuron/src/wire/openai_responses.rs
Normal file
870
crates/neuron/src/wire/openai_responses.rs
Normal file
@@ -0,0 +1,870 @@
|
||||
//! OpenAI Responses API projection.
|
||||
//!
|
||||
//! Two responsibilities:
|
||||
//!
|
||||
//! 1. **Translate request shape**: [`request_to_chat`] flattens
|
||||
//! [`ResponsesRequest`]'s typed `input` items + `instructions`
|
||||
//! into the [`ChatCompletionRequest`] the candle harness already
|
||||
//! knows how to run. The Responses-specific shape stops at this
|
||||
//! function — everything downstream is the same chat path the
|
||||
//! `/v1/chat/completions` route exercises.
|
||||
//!
|
||||
//! 2. **Project event stream**: [`project_responses_stream`] reads
|
||||
//! [`InferenceEvent`]s from the harness and emits the named SSE
|
||||
//! events the Responses API client expects
|
||||
//! (`response.created`, `response.output_text.delta`,
|
||||
//! `response.completed`, …) along with their JSON payloads.
|
||||
//! The HTTP handler in [`crate::api`] reads
|
||||
//! `(event_name, data)` tuples off the receiver and stamps them
|
||||
//! onto axum SSE frames.
|
||||
//!
|
||||
//! Scope cuts (carried over from [`cortex_core::responses`]):
|
||||
//!
|
||||
//! - `previous_response_id` is rejected by [`request_to_chat`]
|
||||
//! with [`TranslateError::ChainedConversationNotSupported`].
|
||||
//! - `Reasoning` input items are dropped (no equivalent in chat).
|
||||
//! - `FunctionCall` / `FunctionCallOutput` items round-trip but the
|
||||
//! harness never emits tool calls today; the synthesis paths are
|
||||
//! in place so the surface is ready when it does.
|
||||
|
||||
use cortex_core::openai::{ChatCompletionRequest, ChatMessage, MessageContent};
|
||||
use cortex_core::responses::{
|
||||
ResponsesContentPart, ResponsesInput, ResponsesInputItem, ResponsesMessageContent,
|
||||
ResponsesOutputContent, ResponsesOutputItem, ResponsesRequest, ResponsesResponse,
|
||||
ResponsesUsage, events,
|
||||
};
|
||||
use serde_json::{Value, json};
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use super::event::{FinishReason, InferenceEvent};
|
||||
|
||||
/// Per-request metadata that has to be stamped into every emitted
|
||||
/// event. The projector spawns a task that owns one of these.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ResponseMeta {
|
||||
pub response_id: String,
|
||||
pub created_at: u64,
|
||||
pub model_id: String,
|
||||
/// Item id used inside `output[0]` (the message). All
|
||||
/// `content_part.*` and `output_text.*` events reference this
|
||||
/// so the consumer knows which item the delta belongs to.
|
||||
pub message_item_id: String,
|
||||
}
|
||||
|
||||
/// Reasons [`request_to_chat`] refuses a request.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum TranslateError {
|
||||
#[error(
|
||||
"previous_response_id is not supported on this neuron; chained \
|
||||
conversations require server-side state we don't store yet"
|
||||
)]
|
||||
ChainedConversationNotSupported,
|
||||
}
|
||||
|
||||
/// Flatten a [`ResponsesRequest`] into the chat-completions shape
|
||||
/// the candle harness already knows how to drive. Keeps the
|
||||
/// Responses-specific machinery contained to a single function so
|
||||
/// the harness stays format-agnostic.
|
||||
///
|
||||
/// Semantics:
|
||||
///
|
||||
/// - `instructions` (if set) becomes a leading `system` message.
|
||||
/// - `input: "<string>"` becomes a single `user` message.
|
||||
/// - `input: [items]` flattens each item:
|
||||
/// - `Message { role, content }` → one `ChatMessage`.
|
||||
/// - `FunctionCall` → an `assistant` turn whose `extra.tool_calls`
|
||||
/// carries the call (chat-completions-shaped). The harness
|
||||
/// doesn't act on tool_calls today, but the shape stays
|
||||
/// consistent with what chat would expect.
|
||||
/// - `FunctionCallOutput` → a `tool` role message with the
|
||||
/// output text. Matches OpenAI's chat convention.
|
||||
/// - `Reasoning` items are dropped (no equivalent in chat).
|
||||
/// - Text parts within an array `content` collapse to a single
|
||||
/// string; image parts get rendered as a chat-style content
|
||||
/// array `[{type:"text"}, {type:"image_url"}]` so the chat
|
||||
/// handler's existing vision path applies.
|
||||
pub fn request_to_chat(req: ResponsesRequest) -> Result<ChatCompletionRequest, TranslateError> {
|
||||
if req.previous_response_id.is_some() {
|
||||
return Err(TranslateError::ChainedConversationNotSupported);
|
||||
}
|
||||
|
||||
let mut messages: Vec<ChatMessage> = Vec::new();
|
||||
|
||||
if let Some(instructions) = req.instructions
|
||||
&& !instructions.is_empty()
|
||||
{
|
||||
messages.push(ChatMessage {
|
||||
role: "system".into(),
|
||||
content: MessageContent::Text(instructions),
|
||||
extra: Value::Object(Default::default()),
|
||||
});
|
||||
}
|
||||
|
||||
match req.input {
|
||||
ResponsesInput::Text(text) => {
|
||||
messages.push(ChatMessage {
|
||||
role: "user".into(),
|
||||
content: MessageContent::Text(text),
|
||||
extra: Value::Object(Default::default()),
|
||||
});
|
||||
}
|
||||
ResponsesInput::Items(items) => {
|
||||
for item in items {
|
||||
if let Some(msg) = input_item_to_chat(item) {
|
||||
messages.push(msg);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(ChatCompletionRequest {
|
||||
model: req.model,
|
||||
messages,
|
||||
temperature: req.temperature,
|
||||
top_p: req.top_p,
|
||||
max_tokens: req.max_output_tokens,
|
||||
stream: Some(req.stream),
|
||||
extra: Value::Object(Default::default()),
|
||||
})
|
||||
}
|
||||
|
||||
fn input_item_to_chat(item: ResponsesInputItem) -> Option<ChatMessage> {
|
||||
match item {
|
||||
ResponsesInputItem::Message { role, content } => Some(ChatMessage {
|
||||
role,
|
||||
content: message_content_to_chat(content),
|
||||
extra: Value::Object(Default::default()),
|
||||
}),
|
||||
ResponsesInputItem::FunctionCall {
|
||||
call_id,
|
||||
name,
|
||||
arguments,
|
||||
} => {
|
||||
// Express the call in chat-completions shape via
|
||||
// `extra.tool_calls`. The harness ignores it today but
|
||||
// the shape is consistent for the day it doesn't.
|
||||
let mut extra = serde_json::Map::new();
|
||||
extra.insert(
|
||||
"tool_calls".into(),
|
||||
json!([{
|
||||
"id": call_id,
|
||||
"type": "function",
|
||||
"function": { "name": name, "arguments": arguments },
|
||||
}]),
|
||||
);
|
||||
Some(ChatMessage {
|
||||
role: "assistant".into(),
|
||||
content: MessageContent::Text(String::new()),
|
||||
extra: Value::Object(extra),
|
||||
})
|
||||
}
|
||||
ResponsesInputItem::FunctionCallOutput { call_id, output } => {
|
||||
let mut extra = serde_json::Map::new();
|
||||
extra.insert("tool_call_id".into(), Value::String(call_id));
|
||||
Some(ChatMessage {
|
||||
role: "tool".into(),
|
||||
content: MessageContent::Text(output),
|
||||
extra: Value::Object(extra),
|
||||
})
|
||||
}
|
||||
// Reasoning items don't have a chat-completions equivalent
|
||||
// we can faithfully forward. Silently drop — the alternative
|
||||
// is rejecting a well-formed request, which is worse UX.
|
||||
ResponsesInputItem::Reasoning { .. } => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn message_content_to_chat(content: ResponsesMessageContent) -> MessageContent {
|
||||
match content {
|
||||
ResponsesMessageContent::Text(s) => MessageContent::Text(s),
|
||||
ResponsesMessageContent::Parts(parts) => {
|
||||
// Collapse to a string when every part is text; emit
|
||||
// the chat content-array shape only when an image is
|
||||
// present (some upstreams treat the array form as a
|
||||
// vision-only signal and reject it for text-only
|
||||
// models).
|
||||
let has_image = parts
|
||||
.iter()
|
||||
.any(|p| matches!(p, ResponsesContentPart::InputImage { .. }));
|
||||
if !has_image {
|
||||
let joined = parts
|
||||
.into_iter()
|
||||
.filter_map(|p| match p {
|
||||
ResponsesContentPart::InputText { text }
|
||||
| ResponsesContentPart::OutputText { text, .. } => Some(text),
|
||||
ResponsesContentPart::InputImage { .. } => None,
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n\n");
|
||||
return MessageContent::Text(joined);
|
||||
}
|
||||
let mut out: Vec<Value> = Vec::with_capacity(parts.len());
|
||||
for p in parts {
|
||||
match p {
|
||||
ResponsesContentPart::InputText { text }
|
||||
| ResponsesContentPart::OutputText { text, .. } => {
|
||||
out.push(json!({ "type": "text", "text": text }));
|
||||
}
|
||||
ResponsesContentPart::InputImage { image_url, .. } => {
|
||||
out.push(json!({
|
||||
"type": "image_url",
|
||||
"image_url": { "url": image_url },
|
||||
}));
|
||||
}
|
||||
}
|
||||
}
|
||||
MessageContent::Parts(out)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Streaming projection ─────────────────────────────────────────────
|
||||
|
||||
/// One frame the projector emits. The HTTP handler maps each into
|
||||
/// an axum `Sse::Event` with both an `event:` name and a `data:`
|
||||
/// JSON payload — Responses, unlike chat completions, uses named
|
||||
/// SSE events.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ResponseStreamFrame {
|
||||
pub event_name: &'static str,
|
||||
pub data: Value,
|
||||
}
|
||||
|
||||
/// Project an [`InferenceEvent`] receiver into a stream of
|
||||
/// [`ResponseStreamFrame`]s. The emitted sequence per stream is:
|
||||
///
|
||||
/// 1. `response.created` — shell with `status: "in_progress"`.
|
||||
/// 2. `response.output_item.added` — empty message item.
|
||||
/// 3. `response.content_part.added` — empty `output_text` part.
|
||||
/// 4. `response.output_text.delta` × N — token-by-token text.
|
||||
/// 5. `response.output_text.done` — full accumulated text.
|
||||
/// 6. `response.content_part.done` — full part payload.
|
||||
/// 7. `response.output_item.done` — full message item.
|
||||
/// 8. `response.completed` — final response with `status:"completed"`.
|
||||
///
|
||||
/// Empty TextDeltas (the harness's incomplete-UTF-8 buffering) are
|
||||
/// dropped. `ReasoningDelta`s have no representation in the
|
||||
/// Responses API spec we model yet, so they're dropped too.
|
||||
pub fn project_responses_stream(
|
||||
rx: mpsc::Receiver<InferenceEvent>,
|
||||
meta: ResponseMeta,
|
||||
) -> mpsc::Receiver<ResponseStreamFrame> {
|
||||
let (tx, out_rx) = mpsc::channel::<ResponseStreamFrame>(64);
|
||||
tokio::spawn(async move {
|
||||
run_projection(rx, meta, tx).await;
|
||||
});
|
||||
out_rx
|
||||
}
|
||||
|
||||
async fn run_projection(
|
||||
mut rx: mpsc::Receiver<InferenceEvent>,
|
||||
meta: ResponseMeta,
|
||||
tx: mpsc::Sender<ResponseStreamFrame>,
|
||||
) {
|
||||
let mut accumulated = String::new();
|
||||
let mut finish: Option<FinishReason> = None;
|
||||
let mut emitted_start = false;
|
||||
|
||||
while let Some(event) = rx.recv().await {
|
||||
match event {
|
||||
InferenceEvent::Start => {
|
||||
emitted_start = true;
|
||||
if !emit_start_frames(&tx, &meta).await {
|
||||
return;
|
||||
}
|
||||
}
|
||||
InferenceEvent::TextDelta(text) => {
|
||||
if text.is_empty() {
|
||||
continue;
|
||||
}
|
||||
accumulated.push_str(&text);
|
||||
let frame = ResponseStreamFrame {
|
||||
event_name: events::OUTPUT_TEXT_DELTA,
|
||||
data: json!({
|
||||
"item_id": meta.message_item_id,
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": text,
|
||||
}),
|
||||
};
|
||||
if tx.send(frame).await.is_err() {
|
||||
return;
|
||||
}
|
||||
}
|
||||
InferenceEvent::ReasoningDelta(_) => {
|
||||
// No representation in our Responses model yet.
|
||||
// Stage where it'd land: a `response.reasoning_*`
|
||||
// event family alongside `response.output_text.*`.
|
||||
}
|
||||
InferenceEvent::ToolCall { .. } => {
|
||||
// Responses-side tool-call routing not wired yet
|
||||
// (would emit response.function_call_arguments.*
|
||||
// events). Drop for now; the chat-completions
|
||||
// projector handles tool calls. Future work
|
||||
// tracked in #7 alongside the in_progress event.
|
||||
}
|
||||
InferenceEvent::Finish { reason } => {
|
||||
finish = Some(reason);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Producers can drop without ever sending Start (e.g. early
|
||||
// poisoned-model error). Synthesize the open frames so the
|
||||
// consumer at least sees a coherent shell before completed.
|
||||
if !emitted_start && !emit_start_frames(&tx, &meta).await {
|
||||
return;
|
||||
}
|
||||
|
||||
let reason = finish.unwrap_or(FinishReason::Stop);
|
||||
let _ = emit_finish_frames(&tx, &meta, &accumulated, reason).await;
|
||||
}
|
||||
|
||||
async fn emit_start_frames(tx: &mpsc::Sender<ResponseStreamFrame>, meta: &ResponseMeta) -> bool {
|
||||
let shell = response_shell(meta, "in_progress", &[], None);
|
||||
let frames = [
|
||||
ResponseStreamFrame {
|
||||
event_name: events::CREATED,
|
||||
data: json!({ "response": shell.clone() }),
|
||||
},
|
||||
// `response.in_progress` carries the same shell as
|
||||
// `response.created` — both report the "in_progress"
|
||||
// status and both are payload-light bookkeeping events.
|
||||
// The distinction is meaningful to clients that
|
||||
// differentiate "request validated" from "model is
|
||||
// generating" in their UI (loading spinner vs streaming
|
||||
// spinner). OpenAI's own Responses SSE emits them as a
|
||||
// pair; matching the wire shape avoids subtle client
|
||||
// breakage.
|
||||
ResponseStreamFrame {
|
||||
event_name: events::IN_PROGRESS,
|
||||
data: json!({ "response": shell }),
|
||||
},
|
||||
ResponseStreamFrame {
|
||||
event_name: events::OUTPUT_ITEM_ADDED,
|
||||
data: json!({
|
||||
"output_index": 0,
|
||||
"item": empty_message_item(&meta.message_item_id),
|
||||
}),
|
||||
},
|
||||
ResponseStreamFrame {
|
||||
event_name: events::CONTENT_PART_ADDED,
|
||||
data: json!({
|
||||
"item_id": meta.message_item_id,
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"part": { "type": "output_text", "text": "", "annotations": [] },
|
||||
}),
|
||||
},
|
||||
];
|
||||
for frame in frames {
|
||||
if tx.send(frame).await.is_err() {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
async fn emit_finish_frames(
|
||||
tx: &mpsc::Sender<ResponseStreamFrame>,
|
||||
meta: &ResponseMeta,
|
||||
full_text: &str,
|
||||
reason: FinishReason,
|
||||
) -> bool {
|
||||
let status = finish_to_status(reason);
|
||||
let full_part = json!({
|
||||
"type": "output_text",
|
||||
"text": full_text,
|
||||
"annotations": [],
|
||||
});
|
||||
let full_item = json!({
|
||||
"type": "message",
|
||||
"id": meta.message_item_id,
|
||||
"role": "assistant",
|
||||
"content": [full_part.clone()],
|
||||
"status": status,
|
||||
});
|
||||
let frames = [
|
||||
ResponseStreamFrame {
|
||||
event_name: events::OUTPUT_TEXT_DONE,
|
||||
data: json!({
|
||||
"item_id": meta.message_item_id,
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"text": full_text,
|
||||
}),
|
||||
},
|
||||
ResponseStreamFrame {
|
||||
event_name: events::CONTENT_PART_DONE,
|
||||
data: json!({
|
||||
"item_id": meta.message_item_id,
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"part": full_part,
|
||||
}),
|
||||
},
|
||||
ResponseStreamFrame {
|
||||
event_name: events::OUTPUT_ITEM_DONE,
|
||||
data: json!({
|
||||
"output_index": 0,
|
||||
"item": full_item.clone(),
|
||||
}),
|
||||
},
|
||||
ResponseStreamFrame {
|
||||
event_name: events::COMPLETED,
|
||||
data: json!({
|
||||
"response": response_shell(meta, status, &[full_item], None)
|
||||
}),
|
||||
},
|
||||
];
|
||||
for frame in frames {
|
||||
if tx.send(frame).await.is_err() {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
fn response_shell(
|
||||
meta: &ResponseMeta,
|
||||
status: &str,
|
||||
output: &[Value],
|
||||
usage: Option<&ResponsesUsage>,
|
||||
) -> Value {
|
||||
let mut obj = serde_json::Map::new();
|
||||
obj.insert("id".into(), Value::String(meta.response_id.clone()));
|
||||
obj.insert("object".into(), Value::String("response".into()));
|
||||
obj.insert("created_at".into(), json!(meta.created_at));
|
||||
obj.insert("status".into(), Value::String(status.into()));
|
||||
obj.insert("model".into(), Value::String(meta.model_id.clone()));
|
||||
obj.insert("output".into(), Value::Array(output.to_vec()));
|
||||
if let Some(u) = usage {
|
||||
obj.insert(
|
||||
"usage".into(),
|
||||
json!({
|
||||
"input_tokens": u.input_tokens,
|
||||
"output_tokens": u.output_tokens,
|
||||
"total_tokens": u.total_tokens,
|
||||
}),
|
||||
);
|
||||
}
|
||||
Value::Object(obj)
|
||||
}
|
||||
|
||||
fn empty_message_item(item_id: &str) -> Value {
|
||||
json!({
|
||||
"type": "message",
|
||||
"id": item_id,
|
||||
"role": "assistant",
|
||||
"content": [],
|
||||
"status": "in_progress",
|
||||
})
|
||||
}
|
||||
|
||||
fn finish_to_status(reason: FinishReason) -> &'static str {
|
||||
match reason {
|
||||
FinishReason::Stop | FinishReason::ToolCalls => "completed",
|
||||
FinishReason::Length => "incomplete",
|
||||
}
|
||||
}
|
||||
|
||||
// ── Non-streaming helpers ────────────────────────────────────────────
|
||||
|
||||
/// Collect a chat-completions response into a non-streaming
|
||||
/// [`ResponsesResponse`]. Used by the `/v1/responses` handler when
|
||||
/// the request doesn't set `stream: true`.
|
||||
pub fn build_response(
|
||||
meta: &ResponseMeta,
|
||||
full_text: String,
|
||||
reason: FinishReason,
|
||||
usage: Option<ResponsesUsage>,
|
||||
) -> ResponsesResponse {
|
||||
let status = finish_to_status(reason).to_string();
|
||||
ResponsesResponse {
|
||||
id: meta.response_id.clone(),
|
||||
object: "response".into(),
|
||||
created_at: meta.created_at,
|
||||
status: status.clone(),
|
||||
model: meta.model_id.clone(),
|
||||
output: vec![ResponsesOutputItem::Message {
|
||||
id: meta.message_item_id.clone(),
|
||||
role: "assistant".into(),
|
||||
content: vec![ResponsesOutputContent::OutputText {
|
||||
text: full_text,
|
||||
annotations: vec![],
|
||||
}],
|
||||
status,
|
||||
}],
|
||||
usage,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use cortex_core::openai::MessageContent;
|
||||
|
||||
fn meta() -> ResponseMeta {
|
||||
ResponseMeta {
|
||||
response_id: "resp_1".into(),
|
||||
created_at: 1700,
|
||||
model_id: "m".into(),
|
||||
message_item_id: "msg_1".into(),
|
||||
}
|
||||
}
|
||||
|
||||
// ── request translator ──────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn translates_text_input_to_single_user_message() {
|
||||
let req = ResponsesRequest {
|
||||
model: "m".into(),
|
||||
input: ResponsesInput::Text("hi".into()),
|
||||
instructions: None,
|
||||
stream: false,
|
||||
max_output_tokens: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
previous_response_id: None,
|
||||
extra: Value::Object(Default::default()),
|
||||
};
|
||||
let chat = request_to_chat(req).unwrap();
|
||||
assert_eq!(chat.messages.len(), 1);
|
||||
assert_eq!(chat.messages[0].role, "user");
|
||||
assert!(matches!(
|
||||
&chat.messages[0].content,
|
||||
MessageContent::Text(t) if t == "hi"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn instructions_become_leading_system_message() {
|
||||
let req = ResponsesRequest {
|
||||
model: "m".into(),
|
||||
input: ResponsesInput::Text("hi".into()),
|
||||
instructions: Some("you are helpful".into()),
|
||||
stream: false,
|
||||
max_output_tokens: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
previous_response_id: None,
|
||||
extra: Value::Object(Default::default()),
|
||||
};
|
||||
let chat = request_to_chat(req).unwrap();
|
||||
assert_eq!(chat.messages.len(), 2);
|
||||
assert_eq!(chat.messages[0].role, "system");
|
||||
assert!(matches!(
|
||||
&chat.messages[0].content,
|
||||
MessageContent::Text(t) if t == "you are helpful"
|
||||
));
|
||||
assert_eq!(chat.messages[1].role, "user");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_previous_response_id() {
|
||||
let req = ResponsesRequest {
|
||||
model: "m".into(),
|
||||
input: ResponsesInput::Text("hi".into()),
|
||||
instructions: None,
|
||||
stream: false,
|
||||
max_output_tokens: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
previous_response_id: Some("resp_prev".into()),
|
||||
extra: Value::Object(Default::default()),
|
||||
};
|
||||
assert!(matches!(
|
||||
request_to_chat(req),
|
||||
Err(TranslateError::ChainedConversationNotSupported)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn translates_input_items_to_chat_messages() {
|
||||
let req = ResponsesRequest {
|
||||
model: "m".into(),
|
||||
input: ResponsesInput::Items(vec![
|
||||
ResponsesInputItem::Message {
|
||||
role: "user".into(),
|
||||
content: ResponsesMessageContent::Text("first".into()),
|
||||
},
|
||||
ResponsesInputItem::Message {
|
||||
role: "assistant".into(),
|
||||
content: ResponsesMessageContent::Text("reply".into()),
|
||||
},
|
||||
ResponsesInputItem::Message {
|
||||
role: "user".into(),
|
||||
content: ResponsesMessageContent::Text("second".into()),
|
||||
},
|
||||
]),
|
||||
instructions: None,
|
||||
stream: false,
|
||||
max_output_tokens: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
previous_response_id: None,
|
||||
extra: Value::Object(Default::default()),
|
||||
};
|
||||
let chat = request_to_chat(req).unwrap();
|
||||
assert_eq!(chat.messages.len(), 3);
|
||||
let roles: Vec<&str> = chat.messages.iter().map(|m| m.role.as_str()).collect();
|
||||
assert_eq!(roles, vec!["user", "assistant", "user"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_input_translates_to_chat_parts_array() {
|
||||
let req = ResponsesRequest {
|
||||
model: "m".into(),
|
||||
input: ResponsesInput::Items(vec![ResponsesInputItem::Message {
|
||||
role: "user".into(),
|
||||
content: ResponsesMessageContent::Parts(vec![
|
||||
ResponsesContentPart::InputText {
|
||||
text: "what is this?".into(),
|
||||
},
|
||||
ResponsesContentPart::InputImage {
|
||||
image_url: "data:image/png;base64,AAA=".into(),
|
||||
detail: None,
|
||||
},
|
||||
]),
|
||||
}]),
|
||||
instructions: None,
|
||||
stream: false,
|
||||
max_output_tokens: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
previous_response_id: None,
|
||||
extra: Value::Object(Default::default()),
|
||||
};
|
||||
let chat = request_to_chat(req).unwrap();
|
||||
let parts = match &chat.messages[0].content {
|
||||
MessageContent::Parts(p) => p.clone(),
|
||||
other => panic!("expected Parts, got {other:?}"),
|
||||
};
|
||||
assert_eq!(parts.len(), 2);
|
||||
assert_eq!(parts[0]["type"], "text");
|
||||
assert_eq!(parts[1]["type"], "image_url");
|
||||
assert_eq!(parts[1]["image_url"]["url"], "data:image/png;base64,AAA=");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn text_only_parts_collapse_to_string() {
|
||||
let req = ResponsesRequest {
|
||||
model: "m".into(),
|
||||
input: ResponsesInput::Items(vec![ResponsesInputItem::Message {
|
||||
role: "user".into(),
|
||||
content: ResponsesMessageContent::Parts(vec![
|
||||
ResponsesContentPart::InputText {
|
||||
text: "first".into(),
|
||||
},
|
||||
ResponsesContentPart::InputText {
|
||||
text: "second".into(),
|
||||
},
|
||||
]),
|
||||
}]),
|
||||
instructions: None,
|
||||
stream: false,
|
||||
max_output_tokens: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
previous_response_id: None,
|
||||
extra: Value::Object(Default::default()),
|
||||
};
|
||||
let chat = request_to_chat(req).unwrap();
|
||||
assert!(matches!(
|
||||
&chat.messages[0].content,
|
||||
MessageContent::Text(t) if t == "first\n\nsecond"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reasoning_items_are_silently_dropped() {
|
||||
let req = ResponsesRequest {
|
||||
model: "m".into(),
|
||||
input: ResponsesInput::Items(vec![
|
||||
ResponsesInputItem::Reasoning { content: vec![] },
|
||||
ResponsesInputItem::Message {
|
||||
role: "user".into(),
|
||||
content: ResponsesMessageContent::Text("hi".into()),
|
||||
},
|
||||
]),
|
||||
instructions: None,
|
||||
stream: false,
|
||||
max_output_tokens: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
previous_response_id: None,
|
||||
extra: Value::Object(Default::default()),
|
||||
};
|
||||
let chat = request_to_chat(req).unwrap();
|
||||
assert_eq!(chat.messages.len(), 1);
|
||||
assert_eq!(chat.messages[0].role, "user");
|
||||
}
|
||||
|
||||
// ── streaming projector ─────────────────────────────────────────
|
||||
|
||||
async fn collect(mut rx: mpsc::Receiver<ResponseStreamFrame>) -> Vec<ResponseStreamFrame> {
|
||||
let mut out = Vec::new();
|
||||
while let Some(f) = rx.recv().await {
|
||||
out.push(f);
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn full_stream_emits_expected_event_sequence() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(8);
|
||||
let out = project_responses_stream(rx, meta());
|
||||
|
||||
tx.send(InferenceEvent::Start).await.unwrap();
|
||||
tx.send(InferenceEvent::TextDelta("hel".into()))
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::TextDelta("lo".into()))
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::Finish {
|
||||
reason: FinishReason::Stop,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
|
||||
let frames = collect(out).await;
|
||||
let names: Vec<&str> = frames.iter().map(|f| f.event_name).collect();
|
||||
assert_eq!(
|
||||
names,
|
||||
vec![
|
||||
events::CREATED,
|
||||
events::IN_PROGRESS,
|
||||
events::OUTPUT_ITEM_ADDED,
|
||||
events::CONTENT_PART_ADDED,
|
||||
events::OUTPUT_TEXT_DELTA,
|
||||
events::OUTPUT_TEXT_DELTA,
|
||||
events::OUTPUT_TEXT_DONE,
|
||||
events::CONTENT_PART_DONE,
|
||||
events::OUTPUT_ITEM_DONE,
|
||||
events::COMPLETED,
|
||||
]
|
||||
);
|
||||
|
||||
// The two deltas should carry the right text. Indices
|
||||
// shifted by one after IN_PROGRESS inserted between
|
||||
// CREATED and OUTPUT_ITEM_ADDED.
|
||||
assert_eq!(frames[4].data["delta"], "hel");
|
||||
assert_eq!(frames[5].data["delta"], "lo");
|
||||
|
||||
// The done event has the full accumulated text.
|
||||
assert_eq!(frames[6].data["text"], "hello");
|
||||
|
||||
// Completed event carries the full message item.
|
||||
let completed = &frames[9].data["response"];
|
||||
assert_eq!(completed["status"], "completed");
|
||||
let output = completed["output"].as_array().unwrap();
|
||||
assert_eq!(output.len(), 1);
|
||||
assert_eq!(output[0]["content"][0]["text"], "hello");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn length_finish_maps_to_incomplete_status() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(8);
|
||||
let out = project_responses_stream(rx, meta());
|
||||
tx.send(InferenceEvent::Start).await.unwrap();
|
||||
tx.send(InferenceEvent::Finish {
|
||||
reason: FinishReason::Length,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
let frames = collect(out).await;
|
||||
let completed = frames
|
||||
.iter()
|
||||
.find(|f| f.event_name == events::COMPLETED)
|
||||
.unwrap();
|
||||
assert_eq!(completed.data["response"]["status"], "incomplete");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn synthesises_start_frames_when_producer_skips_start() {
|
||||
// A producer that drops without sending Start (poisoned
|
||||
// model, immediate disconnect, …) should still produce a
|
||||
// coherent stream — the projector synthesises the
|
||||
// mandatory header frames before COMPLETED so the
|
||||
// consumer never sees an output_text.done without a
|
||||
// matching content_part.added.
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(8);
|
||||
let out = project_responses_stream(rx, meta());
|
||||
drop(tx);
|
||||
let frames = collect(out).await;
|
||||
let names: Vec<&str> = frames.iter().map(|f| f.event_name).collect();
|
||||
assert!(names.contains(&events::CREATED));
|
||||
assert!(names.contains(&events::COMPLETED));
|
||||
assert!(names.contains(&events::OUTPUT_TEXT_DONE));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn empty_text_deltas_are_dropped() {
|
||||
let (tx, rx) = mpsc::channel::<InferenceEvent>(8);
|
||||
let out = project_responses_stream(rx, meta());
|
||||
tx.send(InferenceEvent::Start).await.unwrap();
|
||||
tx.send(InferenceEvent::TextDelta(String::new()))
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::TextDelta("real".into()))
|
||||
.await
|
||||
.unwrap();
|
||||
tx.send(InferenceEvent::Finish {
|
||||
reason: FinishReason::Stop,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
drop(tx);
|
||||
let frames = collect(out).await;
|
||||
let delta_count = frames
|
||||
.iter()
|
||||
.filter(|f| f.event_name == events::OUTPUT_TEXT_DELTA)
|
||||
.count();
|
||||
assert_eq!(delta_count, 1, "empty delta must not produce a frame");
|
||||
}
|
||||
|
||||
// ── non-streaming builder ───────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn build_response_produces_completed_message_with_usage() {
|
||||
let r = build_response(
|
||||
&meta(),
|
||||
"hello".into(),
|
||||
FinishReason::Stop,
|
||||
Some(ResponsesUsage {
|
||||
input_tokens: 5,
|
||||
output_tokens: 1,
|
||||
total_tokens: 6,
|
||||
}),
|
||||
);
|
||||
assert_eq!(r.status, "completed");
|
||||
match &r.output[0] {
|
||||
ResponsesOutputItem::Message {
|
||||
role,
|
||||
content,
|
||||
status,
|
||||
..
|
||||
} => {
|
||||
assert_eq!(role, "assistant");
|
||||
assert_eq!(status, "completed");
|
||||
match &content[0] {
|
||||
ResponsesOutputContent::OutputText { text, .. } => {
|
||||
assert_eq!(text, "hello");
|
||||
}
|
||||
}
|
||||
}
|
||||
other => panic!("expected Message, got {other:?}"),
|
||||
}
|
||||
let u = r.usage.unwrap();
|
||||
assert_eq!(u.total_tokens, 6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_response_length_yields_incomplete_status() {
|
||||
let r = build_response(&meta(), "trunc".into(), FinishReason::Length, None);
|
||||
assert_eq!(r.status, "incomplete");
|
||||
}
|
||||
}
|
||||
@@ -322,3 +322,168 @@ async fn test_chat_completions_streaming_model_not_loaded() {
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), 404);
|
||||
}
|
||||
|
||||
// ── /v1/responses ────────────────────────────────────────────────────
|
||||
|
||||
/// `/v1/responses` returns 503 when no candle harness is registered —
|
||||
/// matches the chat-completions error shape so a client can swap
|
||||
/// endpoints without re-handling 503s.
|
||||
#[tokio::test]
|
||||
async fn test_responses_no_candle_harness() {
|
||||
let registry = HarnessRegistry::new();
|
||||
let health_cache = Arc::new(HealthCache::new());
|
||||
let state = Arc::new(NeuronState {
|
||||
discovery: fake_discovery(),
|
||||
health_cache,
|
||||
registry: RwLock::new(registry),
|
||||
candle: None,
|
||||
activation: Arc::new(ActivationTracker::new(&[])),
|
||||
});
|
||||
let app = api::neuron_routes().with_state(state);
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
let url = format!("http://{addr}");
|
||||
|
||||
let resp = reqwest::Client::new()
|
||||
.post(format!("{url}/v1/responses"))
|
||||
.json(&json!({"model": "anything", "input": "hi"}))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), 503);
|
||||
}
|
||||
|
||||
/// `previous_response_id` is rejected at translate time with 400 —
|
||||
/// we don't store responses server-side yet, so chained
|
||||
/// conversations can't be honoured.
|
||||
#[tokio::test]
|
||||
async fn test_responses_rejects_previous_response_id() {
|
||||
use cortex_core::harness::HarnessConfig;
|
||||
use neuron::config::HarnessSettings;
|
||||
|
||||
let registry = HarnessRegistry::from_configs(
|
||||
&[HarnessConfig {
|
||||
name: "candle".into(),
|
||||
}],
|
||||
"http://localhost:0",
|
||||
&HarnessSettings::default(),
|
||||
);
|
||||
let candle = registry.candle();
|
||||
let health_cache = Arc::new(HealthCache::new());
|
||||
let state = Arc::new(NeuronState {
|
||||
discovery: fake_discovery(),
|
||||
health_cache,
|
||||
registry: RwLock::new(registry),
|
||||
candle,
|
||||
activation: Arc::new(ActivationTracker::new(&[])),
|
||||
});
|
||||
let app = api::neuron_routes().with_state(state);
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
let url = format!("http://{addr}");
|
||||
|
||||
let resp = reqwest::Client::new()
|
||||
.post(format!("{url}/v1/responses"))
|
||||
.json(&json!({
|
||||
"model": "anything",
|
||||
"input": "hi",
|
||||
"previous_response_id": "resp_prev_42"
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), 400);
|
||||
let body: serde_json::Value = resp.json().await.unwrap();
|
||||
assert_eq!(body["code"], "chained_conversation_not_supported");
|
||||
}
|
||||
|
||||
/// `/v1/responses` returns 404 when the model isn't loaded — same
|
||||
/// surface as chat completions.
|
||||
#[tokio::test]
|
||||
async fn test_responses_model_not_loaded() {
|
||||
use cortex_core::harness::HarnessConfig;
|
||||
use neuron::config::HarnessSettings;
|
||||
|
||||
let registry = HarnessRegistry::from_configs(
|
||||
&[HarnessConfig {
|
||||
name: "candle".into(),
|
||||
}],
|
||||
"http://localhost:0",
|
||||
&HarnessSettings::default(),
|
||||
);
|
||||
let candle = registry.candle();
|
||||
let health_cache = Arc::new(HealthCache::new());
|
||||
let state = Arc::new(NeuronState {
|
||||
discovery: fake_discovery(),
|
||||
health_cache,
|
||||
registry: RwLock::new(registry),
|
||||
candle,
|
||||
activation: Arc::new(ActivationTracker::new(&[])),
|
||||
});
|
||||
let app = api::neuron_routes().with_state(state);
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
let url = format!("http://{addr}");
|
||||
|
||||
let resp = reqwest::Client::new()
|
||||
.post(format!("{url}/v1/responses"))
|
||||
.json(&json!({"model": "not-loaded", "input": "hi"}))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), 404);
|
||||
}
|
||||
|
||||
/// Same model-not-loaded surface on the streaming path. The
|
||||
/// stream is opened only after model lookup succeeds, so a
|
||||
/// missing model fails fast with a non-SSE 404 response.
|
||||
#[tokio::test]
|
||||
async fn test_responses_streaming_model_not_loaded() {
|
||||
use cortex_core::harness::HarnessConfig;
|
||||
use neuron::config::HarnessSettings;
|
||||
|
||||
let registry = HarnessRegistry::from_configs(
|
||||
&[HarnessConfig {
|
||||
name: "candle".into(),
|
||||
}],
|
||||
"http://localhost:0",
|
||||
&HarnessSettings::default(),
|
||||
);
|
||||
let candle = registry.candle();
|
||||
let health_cache = Arc::new(HealthCache::new());
|
||||
let state = Arc::new(NeuronState {
|
||||
discovery: fake_discovery(),
|
||||
health_cache,
|
||||
registry: RwLock::new(registry),
|
||||
candle,
|
||||
activation: Arc::new(ActivationTracker::new(&[])),
|
||||
});
|
||||
let app = api::neuron_routes().with_state(state);
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
let url = format!("http://{addr}");
|
||||
|
||||
let resp = reqwest::Client::new()
|
||||
.post(format!("{url}/v1/responses"))
|
||||
.json(&json!({
|
||||
"model": "not-loaded",
|
||||
"input": "hi",
|
||||
"stream": true
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), 404);
|
||||
}
|
||||
|
||||
269
crates/neuron/tests/preflight.rs
Normal file
269
crates/neuron/tests/preflight.rs
Normal file
@@ -0,0 +1,269 @@
|
||||
//! End-to-end preflight tests against a mock HF-compatible server.
|
||||
//!
|
||||
//! Unit tests in `harness/preflight.rs` exercise the classifier and
|
||||
//! feasibility table against synthetic file lists. These tests close
|
||||
//! the loop: spawn an axum server that returns a `RepoInfo`-shaped
|
||||
//! JSON payload at `/api/models/{org}/{name}`, point `hf_hub::Api` at
|
||||
//! it, and assert `preflight()` returns the expected outcome.
|
||||
|
||||
use axum::Router;
|
||||
use axum::extract::Path;
|
||||
use axum::http::StatusCode;
|
||||
use axum::response::{IntoResponse, Json};
|
||||
use axum::routing::get;
|
||||
use cortex_core::harness::ModelSpec;
|
||||
use neuron::harness::preflight::{PreflightError, SourceFormat, preflight};
|
||||
use serde_json::{Value, json};
|
||||
use std::sync::Arc;
|
||||
use std::sync::Mutex;
|
||||
|
||||
/// Per-test mock state: a map from `{org}/{name}` to the JSON body the
|
||||
/// mock server returns at the corresponding `/api/models/{org}/{name}`
|
||||
/// endpoint. `None` means "respond 404".
|
||||
type MockBodies = Arc<Mutex<std::collections::HashMap<String, Option<Value>>>>;
|
||||
|
||||
async fn spawn_mock(bodies: MockBodies) -> String {
|
||||
// hf-hub 0.4 calls /api/models/{org}/{name}/revision/main for
|
||||
// `repo.info()`. We route both shapes so the test stays robust
|
||||
// to a future hf-hub upgrade that drops the `/revision/main`
|
||||
// suffix.
|
||||
let app = Router::new()
|
||||
.route("/api/models/{org}/{name}", get(model_info))
|
||||
.route(
|
||||
"/api/models/{org}/{name}/revision/{rev}",
|
||||
get(model_info_rev),
|
||||
)
|
||||
.with_state(bodies);
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
format!("http://{addr}")
|
||||
}
|
||||
|
||||
async fn model_info(
|
||||
Path((org, name)): Path<(String, String)>,
|
||||
axum::extract::State(bodies): axum::extract::State<MockBodies>,
|
||||
) -> impl IntoResponse {
|
||||
respond(&format!("{org}/{name}"), &bodies)
|
||||
}
|
||||
|
||||
async fn model_info_rev(
|
||||
Path((org, name, _rev)): Path<(String, String, String)>,
|
||||
axum::extract::State(bodies): axum::extract::State<MockBodies>,
|
||||
) -> impl IntoResponse {
|
||||
respond(&format!("{org}/{name}"), &bodies)
|
||||
}
|
||||
|
||||
fn respond(key: &str, bodies: &MockBodies) -> axum::response::Response {
|
||||
let entry = bodies.lock().unwrap().get(key).cloned();
|
||||
match entry {
|
||||
Some(Some(body)) => Json(body).into_response(),
|
||||
Some(None) | None => (StatusCode::NOT_FOUND, "not found").into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
fn build_api(endpoint: &str, cache_dir: &std::path::Path) -> hf_hub::api::tokio::Api {
|
||||
hf_hub::api::tokio::ApiBuilder::new()
|
||||
.with_endpoint(endpoint.to_string())
|
||||
.with_cache_dir(cache_dir.to_path_buf())
|
||||
.build()
|
||||
.expect("build hf-hub Api")
|
||||
}
|
||||
|
||||
fn siblings(filenames: &[&str]) -> Value {
|
||||
json!({
|
||||
"sha": "0000000000000000000000000000000000000000",
|
||||
"siblings": filenames.iter().map(|f| json!({ "rfilename": f })).collect::<Vec<_>>(),
|
||||
})
|
||||
}
|
||||
|
||||
fn spec(model_id: &str, tp: Option<u32>, quant: Option<&str>) -> ModelSpec {
|
||||
ModelSpec {
|
||||
model_id: model_id.into(),
|
||||
harness: "candle".into(),
|
||||
quant: quant.map(String::from),
|
||||
tensor_parallel: tp,
|
||||
devices: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn preflight_gguf_tp_rejected_over_http() {
|
||||
let cache = tempfile::tempdir().expect("tempdir");
|
||||
let bodies: MockBodies = Arc::new(Mutex::new(Default::default()));
|
||||
bodies.lock().unwrap().insert(
|
||||
"HauhauCS/Qwen3.6".to_string(),
|
||||
Some(siblings(&[
|
||||
"README.md",
|
||||
".gitattributes",
|
||||
"Qwen3.6-Q4_K_P.gguf",
|
||||
"Qwen3.6-Q6_K_P.gguf",
|
||||
"Qwen3.6-Q8_K_P.gguf",
|
||||
])),
|
||||
);
|
||||
let endpoint = spawn_mock(bodies).await;
|
||||
|
||||
let api = build_api(&endpoint, cache.path());
|
||||
let s = spec("HauhauCS/Qwen3.6", Some(2), Some("q6k"));
|
||||
let err = preflight(&api, &s).await.unwrap_err();
|
||||
match err {
|
||||
PreflightError::TpRequiresSafetensors {
|
||||
model_id,
|
||||
tp_size,
|
||||
gguf_quants,
|
||||
..
|
||||
} => {
|
||||
assert_eq!(model_id, "HauhauCS/Qwen3.6");
|
||||
assert_eq!(tp_size, 2);
|
||||
assert_eq!(gguf_quants.len(), 3);
|
||||
}
|
||||
other => panic!("expected TpRequiresSafetensors, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn preflight_gguf_quant_suggestion_over_http() {
|
||||
let cache = tempfile::tempdir().expect("tempdir");
|
||||
let bodies: MockBodies = Arc::new(Mutex::new(Default::default()));
|
||||
bodies.lock().unwrap().insert(
|
||||
"HauhauCS/Qwen3.6".to_string(),
|
||||
Some(siblings(&[
|
||||
"Qwen3.6-Q4_K_P.gguf",
|
||||
"Qwen3.6-Q5_K_P.gguf",
|
||||
"Qwen3.6-Q6_K_P.gguf",
|
||||
"Qwen3.6-Q8_K_P.gguf",
|
||||
])),
|
||||
);
|
||||
let endpoint = spawn_mock(bodies).await;
|
||||
|
||||
let api = build_api(&endpoint, cache.path());
|
||||
let s = spec("HauhauCS/Qwen3.6", Some(1), Some("q6k"));
|
||||
let err = preflight(&api, &s).await.unwrap_err();
|
||||
match err {
|
||||
PreflightError::QuantNotFound {
|
||||
requested,
|
||||
nearest,
|
||||
available,
|
||||
..
|
||||
} => {
|
||||
assert_eq!(requested, "q6k");
|
||||
assert_eq!(nearest.as_deref(), Some("q6_k_p"));
|
||||
assert_eq!(available.len(), 4);
|
||||
}
|
||||
other => panic!("expected QuantNotFound, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn preflight_dense_safetensors_tp_ok() {
|
||||
let cache = tempfile::tempdir().expect("tempdir");
|
||||
let bodies: MockBodies = Arc::new(Mutex::new(Default::default()));
|
||||
bodies.lock().unwrap().insert(
|
||||
"Qwen/Q3-30B".to_string(),
|
||||
Some(siblings(&[
|
||||
"config.json",
|
||||
"tokenizer.json",
|
||||
"tokenizer_config.json",
|
||||
"model.safetensors.index.json",
|
||||
"model-00001-of-00006.safetensors",
|
||||
"model-00002-of-00006.safetensors",
|
||||
"model-00003-of-00006.safetensors",
|
||||
])),
|
||||
);
|
||||
let endpoint = spawn_mock(bodies).await;
|
||||
|
||||
let api = build_api(&endpoint, cache.path());
|
||||
let s = spec("Qwen/Q3-30B", Some(2), Some("q5k"));
|
||||
let plan = preflight(&api, &s).await.expect("dense+tp should succeed");
|
||||
assert_eq!(plan.tp_size, 2);
|
||||
assert!(plan.picked_quant_file.is_none());
|
||||
assert!(matches!(
|
||||
plan.format,
|
||||
SourceFormat::DenseSafetensors { sharded: true }
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn preflight_gguf_single_gpu_good_quant() {
|
||||
let cache = tempfile::tempdir().expect("tempdir");
|
||||
let bodies: MockBodies = Arc::new(Mutex::new(Default::default()));
|
||||
bodies.lock().unwrap().insert(
|
||||
"HauhauCS/Qwen3.6".to_string(),
|
||||
Some(siblings(&["Qwen3.6-Q4_K_P.gguf", "Qwen3.6-Q6_K_P.gguf"])),
|
||||
);
|
||||
let endpoint = spawn_mock(bodies).await;
|
||||
|
||||
let api = build_api(&endpoint, cache.path());
|
||||
let s = spec("HauhauCS/Qwen3.6", Some(1), Some("q6_k_p"));
|
||||
let plan = preflight(&api, &s)
|
||||
.await
|
||||
.expect("good quant should succeed");
|
||||
assert_eq!(plan.tp_size, 1);
|
||||
assert_eq!(
|
||||
plan.picked_quant_file.as_deref(),
|
||||
Some("Qwen3.6-Q6_K_P.gguf")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn preflight_repo_fetch_failed_on_404() {
|
||||
// Mock server has no entry for this id → 404, exercising the
|
||||
// RepoFetchFailed path (the same shape today's HauhauCS scenario
|
||||
// would have produced if we'd added preflight before the cache
|
||||
// download was attempted).
|
||||
let cache = tempfile::tempdir().expect("tempdir");
|
||||
let bodies: MockBodies = Arc::new(Mutex::new(Default::default()));
|
||||
let endpoint = spawn_mock(bodies).await;
|
||||
|
||||
let api = build_api(&endpoint, cache.path());
|
||||
let s = spec("DoesNot/Exist", Some(1), None);
|
||||
let err = preflight(&api, &s).await.unwrap_err();
|
||||
assert!(
|
||||
matches!(err, PreflightError::RepoFetchFailed { .. }),
|
||||
"expected RepoFetchFailed, got {err:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn preflight_empty_repo_rejected() {
|
||||
let cache = tempfile::tempdir().expect("tempdir");
|
||||
let bodies: MockBodies = Arc::new(Mutex::new(Default::default()));
|
||||
bodies.lock().unwrap().insert(
|
||||
"Empty/Repo".to_string(),
|
||||
Some(siblings(&["README.md", "tokenizer.json"])),
|
||||
);
|
||||
let endpoint = spawn_mock(bodies).await;
|
||||
|
||||
let api = build_api(&endpoint, cache.path());
|
||||
let s = spec("Empty/Repo", Some(1), None);
|
||||
let err = preflight(&api, &s).await.unwrap_err();
|
||||
assert!(
|
||||
matches!(err, PreflightError::EmptyRepo { .. }),
|
||||
"expected EmptyRepo, got {err:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn preflight_mixed_repo_prefers_safetensors() {
|
||||
let cache = tempfile::tempdir().expect("tempdir");
|
||||
let bodies: MockBodies = Arc::new(Mutex::new(Default::default()));
|
||||
bodies.lock().unwrap().insert(
|
||||
"Mixed/Repo".to_string(),
|
||||
Some(siblings(&[
|
||||
"config.json",
|
||||
"tokenizer.json",
|
||||
"model.safetensors",
|
||||
"model-Q4_K_M.gguf",
|
||||
])),
|
||||
);
|
||||
let endpoint = spawn_mock(bodies).await;
|
||||
|
||||
let api = build_api(&endpoint, cache.path());
|
||||
// TP=2 + quant should succeed via the dense path even though a
|
||||
// GGUF is present — the dense path handles ISQ.
|
||||
let s = spec("Mixed/Repo", Some(2), Some("q5k"));
|
||||
let plan = preflight(&api, &s).await.expect("mixed should succeed");
|
||||
assert!(matches!(plan.format, SourceFormat::Mixed { .. }));
|
||||
}
|
||||
85
helexa-acp.example.toml
Normal file
85
helexa-acp.example.toml
Normal file
@@ -0,0 +1,85 @@
|
||||
# helexa-acp.example.toml — example configuration
|
||||
#
|
||||
# Copy to $XDG_CONFIG_HOME/helexa-acp/config.toml (typically
|
||||
# ~/.config/helexa-acp/config.toml) and adjust for your environment.
|
||||
#
|
||||
# helexa-acp is the ACP (Agent Client Protocol) bridge that connects
|
||||
# editors like Zed to multiple LLM endpoints. Each endpoint speaks a
|
||||
# specific wire format (openai-chat, openai-responses, or
|
||||
# anthropic-messages); helexa-acp picks the right provider at runtime
|
||||
# based on the `wire_api` field.
|
||||
#
|
||||
# Selecting a model from the editor follows the `endpoint:model`
|
||||
# syntax — e.g. `openrouter:anthropic/claude-opus-4` routes the
|
||||
# request to the `openrouter` endpoint with model
|
||||
# `anthropic/claude-opus-4`. A bare `<model>` (no colon) falls
|
||||
# through to whichever endpoint is named in `default_endpoint`.
|
||||
|
||||
default_endpoint = "helexa"
|
||||
|
||||
# Optional: override the built-in system prompt with a file of your own.
|
||||
# When unset, helexa-acp uses a concise coder prompt from src/prompt.rs.
|
||||
# `{cwd}` in the file gets substituted with the session's working
|
||||
# directory at request time.
|
||||
# system_prompt_path = "/home/me/.config/helexa-acp/system-prompt.md"
|
||||
|
||||
# ── helexa (cortex/neuron, self-hosted) ────────────────────────────
|
||||
#
|
||||
# The canonical default. Drives cortex's reverse-proxy / fleet
|
||||
# gateway, which routes to whichever neuron has the model loaded.
|
||||
# `openai-chat` works against any cortex deployment; for vision
|
||||
# models or reasoning surface, switch to `openai-responses` (cortex
|
||||
# 0.1.16+).
|
||||
|
||||
[[endpoints]]
|
||||
name = "helexa"
|
||||
base_url = "http://hanzalova.internal:31313/v1"
|
||||
wire_api = "openai-chat"
|
||||
default_model = "Qwen/Qwen3.6-27B"
|
||||
max_tokens = 8192
|
||||
# Compaction kicks in when the rolling history grows past this token
|
||||
# budget. Set to your model's context window. Disable by removing
|
||||
# the field entirely.
|
||||
context_window = 32768
|
||||
|
||||
# ── OpenRouter (proxy for OpenAI/Anthropic/Google/etc.) ────────────
|
||||
|
||||
[[endpoints]]
|
||||
name = "openrouter"
|
||||
base_url = "https://openrouter.ai/api/v1"
|
||||
wire_api = "openai-chat"
|
||||
api_key_env = "OPENROUTER_API_KEY"
|
||||
default_model = "anthropic/claude-opus-4"
|
||||
|
||||
# ── OpenAI directly (Responses API) ────────────────────────────────
|
||||
#
|
||||
# Use `openai-responses` for the o-series and any model that
|
||||
# benefits from the newer Responses API surface (web search,
|
||||
# computer use, reasoning effort, etc.).
|
||||
|
||||
[[endpoints]]
|
||||
name = "openai"
|
||||
base_url = "https://api.openai.com/v1"
|
||||
wire_api = "openai-responses"
|
||||
api_key_env = "OPENAI_API_KEY"
|
||||
default_model = "gpt-5"
|
||||
|
||||
# ── Anthropic directly ─────────────────────────────────────────────
|
||||
|
||||
[[endpoints]]
|
||||
name = "anthropic"
|
||||
base_url = "https://api.anthropic.com/v1"
|
||||
wire_api = "anthropic-messages"
|
||||
api_key_env = "ANTHROPIC_API_KEY"
|
||||
default_model = "claude-opus-4"
|
||||
|
||||
# ── Local LM Studio / Ollama (compat mode) ─────────────────────────
|
||||
#
|
||||
# Most local-LLM servers expose OpenAI-compatible chat completions.
|
||||
# Use `wire_api = "openai-chat"` and point at the local port.
|
||||
|
||||
# [[endpoints]]
|
||||
# name = "lmstudio"
|
||||
# base_url = "http://localhost:1234/v1"
|
||||
# wire_api = "openai-chat"
|
||||
# default_model = "auto"
|
||||
Reference in New Issue
Block a user