1 //! Inference backend abstraction.
2 //!
3 //! The agent loop only needs to send a turn (optionally with tools) and return
4 //! tool results. This module defines that seam as the `InferenceBackend` trait
5 //! plus a few neutral types, with two implementations:
6 //!
7 //! - `LocalBackend` runs on-device through the `onde` crate (`ChatEngine`).
8 //! - `OpenAiBackend` talks to any OpenAI-compatible HTTP endpoint, configured by
9 //! `base_url`, `api_key`, and `model`.
10 //!
11 //! The trait exposes neither `onde` nor OpenAI types, so the loop does not depend
12 //! on a specific backend.
13 //!
14 //! The seam is consumed by both surfaces: the interactive client (`#[cfg(unix)]`,
15 //! see `run_interactive` in `main.rs` and `mod tui` in `chat.rs`) and the ACP
16 //! server's prompt loop. Some items are still reached only through the
17 //! Unix-only interactive paths, so the dead-code lint stays suppressed on
18 //! non-Unix targets only — Unix builds keep full coverage.
19 #![cfg_attr(not(unix), allow(dead_code))]
20
21 use std::sync::Arc;
22
23 use async_trait::async_trait;
24 use onde::inference::{ChatEngine, ChatMessage, ChatRole, ToolDefinition};
25 use serde::Deserialize;
26 use tokio::sync::Mutex;
27
28 // ── Neutral types ───────────────────────────────────────────────────────────────
29
30 /// A tool the model may call, in a provider-neutral form. `parameters_schema` is
31 /// a JSON Schema encoded as a string (matching how siGit already declares tools).
32 #[derive(Debug, Clone)]
33 pub struct ToolSpec {
34 pub name: String,
35 pub description: String,
36 pub parameters_schema: String,
37 }
38
39 /// A tool call requested by the model.
40 #[derive(Debug, Clone)]
41 pub struct ToolCall {
42 pub id: String,
43 pub name: String,
44 /// Arguments as a JSON-encoded string.
45 pub arguments: String,
46 }
47
48 /// The output of executing one tool call, fed back to the model.
49 #[derive(Debug, Clone)]
50 pub struct ToolResult {
51 pub tool_call_id: String,
52 pub content: String,
53 }
54
55 /// The result of one assistant turn: free text and/or tool calls.
56 #[derive(Debug, Clone, Default)]
57 pub struct TurnResult {
58 pub text: String,
59 pub tool_calls: Vec<ToolCall>,
60 }
61
62 /// Backend errors are plain strings. Callers map them to ACP errors.
63 pub type BackendError = String;
64
65 /// Rough context budget for a conversation, in estimated tokens (see
66 /// [`estimate_tokens`]). When a snapshot exceeds this, the agent loops compact
67 /// history before the next tool round.
68 pub const DEFAULT_CONTEXT_TOKEN_BUDGET: usize = 24_000;
69
70 /// How many trailing messages survive a compaction verbatim (the rest are
71 /// folded into the summary).
72 pub const COMPACT_KEEP_LAST: usize = 6;
73
74 /// The summarization request sent to the model when compacting history.
75 const SUMMARIZE_PROMPT: &str = "Summarize this coding session so far: decisions made, \
76 files touched, current state, open items. Be concise and factual.";
77
78 /// Crude token estimate for a history snapshot: serialized characters / 4.
79 /// Deliberately model-agnostic — it only needs to be in the right ballpark to
80 /// decide when compaction is worth an extra inference round.
81 pub fn estimate_tokens(history: &[serde_json::Value]) -> usize {
82 let chars: usize = history
83 .iter()
84 .map(|message| message.to_string().chars().count())
85 .sum();
86 chars / 4
87 }
88
89 /// A sink for streaming assistant text deltas to the UI as they are produced.
90 ///
91 /// When a caller passes `Some(sink)`, a streaming-capable backend forwards each
92 /// text fragment through it as the model emits it; the returned [`TurnResult`]
93 /// still carries the fully assembled text (and any tool calls). When the sink is
94 /// `None`, the backend runs in non-streaming mode. Unbounded so the inference
95 /// task never blocks on a slow consumer.
96 pub type TokenSink = tokio::sync::mpsc::UnboundedSender<String>;
97
98 // ── The trait ───────────────────────────────────────────────────────────────────
99
100 /// A swappable inference backend driving siGit Code's agent loop.
101 #[async_trait]
102 pub trait InferenceBackend: Send + Sync {
103 /// Start an assistant turn from a new user message, offering `tools`.
104 ///
105 /// If `sink` is `Some`, text is streamed through it as it is generated. A
106 /// backend may decline to stream a given round (for example, on-device
107 /// inference cannot stream while it is still deciding whether to call a
108 /// tool); in that case the text is delivered only via the returned result.
109 async fn send_message_with_tools(
110 &self,
111 text: &str,
112 tools: &[ToolSpec],
113 sink: Option<&TokenSink>,
114 ) -> Result<TurnResult, BackendError>;
115
116 /// Continue the turn by returning tool results. `tools` may be `None` on the
117 /// final round to force a text answer. `sink` streams that text when set.
118 async fn send_tool_results(
119 &self,
120 results: Vec<ToolResult>,
121 tools: Option<&[ToolSpec]>,
122 sink: Option<&TokenSink>,
123 ) -> Result<TurnResult, BackendError>;
124
125 /// Record tool results in the conversation history *without* asking the
126 /// model to continue the turn. Used when a turn is abandoned mid-round
127 /// (the user cancelled at the permission gate): by then the assistant
128 /// message carrying the tool calls is already in history, and leaving them
129 /// unanswered makes strict OpenAI-compatible endpoints reject every later
130 /// request in the session.
131 async fn record_cancelled_tool_results(&self, results: Vec<ToolResult>);
132
133 /// Whether inference runs over the network (a configured provider) rather
134 /// than on-device. Drives UI labelling so the displayed model can't claim a
135 /// local model while requests actually go to the cloud.
136 fn is_remote(&self) -> bool;
137
138 /// A serializable snapshot of the conversation history, one JSON object per
139 /// message (`{"role": ..., "content": ...}` at minimum). The snapshot is
140 /// what the session store persists; it includes any seeded system message
141 /// so [`InferenceBackend::restore_history`] can replace state wholesale.
142 async fn history_snapshot(&self) -> Vec<serde_json::Value>;
143
144 /// Replace the conversation history with a previously saved snapshot.
145 /// Backends that cannot represent every entry (e.g. on-device history has
146 /// no tool-call structure) flatten what they can and drop the rest.
147 async fn restore_history(&self, history: Vec<serde_json::Value>);
148
149 /// Shrink the conversation history: summarize everything so far with one
150 /// extra (non-streaming) inference round, then rebuild history as
151 /// `[system message, summary, last keep_last non-system messages]`. On
152 /// error the original history is left in place.
153 async fn compact_history(&self, keep_last: usize) -> Result<(), BackendError>;
154 }
155
156 // ── Local backend (onde ChatEngine) ──────────────────────────────────────────────
157
158 /// On-device inference. A thin adapter over `onde::ChatEngine`.
159 pub struct LocalBackend {
160 engine: Arc<ChatEngine>,
161 }
162
163 impl LocalBackend {
164 pub fn new(engine: Arc<ChatEngine>) -> Self {
165 Self { engine }
166 }
167 }
168
169 fn to_onde_tools(tools: &[ToolSpec]) -> Vec<ToolDefinition> {
170 tools
171 .iter()
172 .map(|tool| ToolDefinition {
173 name: tool.name.clone(),
174 description: tool.description.clone(),
175 parameters_schema: tool.parameters_schema.clone(),
176 })
177 .collect()
178 }
179
180 #[async_trait]
181 impl InferenceBackend for LocalBackend {
182 async fn send_message_with_tools(
183 &self,
184 text: &str,
185 tools: &[ToolSpec],
186 sink: Option<&TokenSink>,
187 ) -> Result<TurnResult, BackendError> {
188 // onde's tool-aware path is non-streaming: it has to buffer the whole
189 // reply to detect tool calls. We can only stream when no tools are on
190 // offer (a plain answer), which is exactly the tools-disabled case.
191 if let Some(sink) = sink
192 && tools.is_empty()
193 {
194 let rx = self
195 .engine
196 .stream_message(text)
197 .await
198 .map_err(|error| error.to_string())?;
199 return drain_onde_stream(rx, sink).await;
200 }
201
202 let onde_tools = to_onde_tools(tools);
203 let result = self
204 .engine
205 .send_message_with_tools(text, &onde_tools)
206 .await
207 .map_err(|error| error.to_string())?;
208 Ok(onde_result_to_turn(result))
209 }
210
211 async fn send_tool_results(
212 &self,
213 results: Vec<ToolResult>,
214 tools: Option<&[ToolSpec]>,
215 sink: Option<&TokenSink>,
216 ) -> Result<TurnResult, BackendError> {
217 let onde_results: Vec<onde::inference::ToolResult> = results
218 .into_iter()
219 .map(|result| onde::inference::ToolResult {
220 tool_call_id: result.tool_call_id,
221 content: result.content,
222 })
223 .collect();
224
225 // The final round passes `tools = None` to force a text answer; that's
226 // the only round onde can stream, since no further tool calls are parsed.
227 if let Some(sink) = sink
228 && tools.is_none()
229 {
230 let rx = self
231 .engine
232 .stream_tool_results(onde_results, None)
233 .await
234 .map_err(|error| error.to_string())?;
235 return drain_onde_stream(rx, sink).await;
236 }
237
238 let onde_tools = tools.map(to_onde_tools);
239 let result = self
240 .engine
241 .send_tool_results(onde_results, onde_tools.as_deref())
242 .await
243 .map_err(|error| error.to_string())?;
244 Ok(onde_result_to_turn(result))
245 }
246
247 async fn record_cancelled_tool_results(&self, _results: Vec<ToolResult>) {
248 // onde's public API cannot append tool-result history entries without
249 // running another inference round, so the dangling tool call stays in
250 // its history. The chat template replays it as-is, which local models
251 // tolerate — worst case the model re-issues the call next turn.
252 }
253
254 fn is_remote(&self) -> bool {
255 false
256 }
257
258 async fn history_snapshot(&self) -> Vec<serde_json::Value> {
259 // onde's `history()` already flattens tool entries: assistant tool
260 // calls become plain assistant text and tool results are omitted, so
261 // the snapshot is lossy for tool-heavy turns (acceptable in this MVP).
262 self.engine
263 .history()
264 .await
265 .iter()
266 .map(|message| {
267 serde_json::json!({
268 "role": message.role.to_string(),
269 "content": message.content,
270 })
271 })
272 .collect()
273 }
274
275 async fn restore_history(&self, history: Vec<serde_json::Value>) {
276 self.engine.clear_history().await;
277 for entry in history {
278 let role = entry["role"].as_str().unwrap_or("");
279 let content = entry["content"].as_str().unwrap_or("").to_string();
280 // Tool-call-only assistant entries and empty tool results carry no
281 // text a plain chat history can replay; drop them.
282 if content.is_empty() && role != "user" && role != "system" {
283 continue;
284 }
285 let message = match role {
286 "system" => ChatMessage::system(content),
287 "user" => ChatMessage::user(content),
288 "assistant" => ChatMessage::assistant(content),
289 // Tool results flatten to plain text (MVP; acceptable loss).
290 "tool" => ChatMessage::user(format!("[tool result]\n{content}")),
291 _ => continue,
292 };
293 self.engine.push_history(message).await;
294 }
295 }
296
297 async fn compact_history(&self, keep_last: usize) -> Result<(), BackendError> {
298 let snapshot = self.engine.history().await;
299 // One plain (tool-free) inference round produces the summary. On error
300 // history is untouched — send_message only mutates it on success, and
301 // whatever it appended is wiped by the clear below anyway.
302 let result = self
303 .engine
304 .send_message(SUMMARIZE_PROMPT)
305 .await
306 .map_err(|error| error.to_string())?;
307 // Local models may reason in <think> blocks; keep only the visible part.
308 let (_think, summary) = crate::chat::strip_think_blocks(&result.text);
309
310 self.engine.clear_history().await;
311 // Leading system messages carry the session context; keep them all.
312 for message in snapshot
313 .iter()
314 .take_while(|message| message.role == ChatRole::System)
315 {
316 self.engine.push_history(message.clone()).await;
317 }
318 self.engine
319 .push_history(ChatMessage::user(format!(
320 "[Conversation summary]\n{summary}"
321 )))
322 .await;
323 let non_system: Vec<&ChatMessage> = snapshot
324 .iter()
325 .filter(|message| message.role != ChatRole::System)
326 .collect();
327 let tail_start = non_system.len().saturating_sub(keep_last);
328 for message in &non_system[tail_start..] {
329 self.engine.push_history((*message).clone()).await;
330 }
331 Ok(())
332 }
333 }
334
335 /// Drain an onde streaming receiver, forwarding each token to `sink` and
336 /// assembling the full text. onde reports stream failures as a final chunk whose
337 /// `finish_reason` is `"error: …"`; surface those as a backend error.
338 async fn drain_onde_stream(
339 mut rx: tokio::sync::mpsc::Receiver<onde::inference::StreamChunk>,
340 sink: &TokenSink,
341 ) -> Result<TurnResult, BackendError> {
342 let mut text = String::new();
343 while let Some(chunk) = rx.recv().await {
344 if !chunk.delta.is_empty() {
345 text.push_str(&chunk.delta);
346 // The receiver is the UI; if it's gone the turn is being cancelled,
347 // so stop assembling rather than spinning the model to completion.
348 if sink.send(chunk.delta).is_err() {
349 break;
350 }
351 }
352 if chunk.done {
353 if let Some(reason) = chunk.finish_reason
354 && let Some(message) = reason.strip_prefix("error: ")
355 {
356 return Err(message.to_string());
357 }
358 break;
359 }
360 }
361 Ok(TurnResult {
362 text,
363 tool_calls: Vec::new(),
364 })
365 }
366
367 /// Convert an `onde` tool-aware result into the neutral [`TurnResult`].
368 fn onde_result_to_turn(result: onde::inference::ToolAwareResult) -> TurnResult {
369 TurnResult {
370 text: result.text,
371 tool_calls: result
372 .tool_calls
373 .into_iter()
374 .map(|call| ToolCall {
375 id: call.id,
376 name: call.function_name,
377 arguments: call.arguments,
378 })
379 .collect(),
380 }
381 }
382
383 // ── OpenAI-compatible backend ─────────────────────────────────────────────────────
384
385 /// Inference against any OpenAI-compatible Chat Completions endpoint.
386 ///
387 /// Conversation state is held client-side and replayed on every request, so the
388 /// endpoint can be stateless. Standard OpenAI function-calling is used end to
389 /// end (`tools`, `choices[].message.tool_calls`, `role: "tool"` follow-ups).
390 pub struct OpenAiBackend {
391 base_url: String,
392 api_key: String,
393 model: String,
394 http: reqwest::Client,
395 /// The full message list sent on each request (system + turns + tool results).
396 history: Mutex<Vec<serde_json::Value>>,
397 }
398
399 impl OpenAiBackend {
400 /// Build a backend for `{base_url, api_key, model}`, seeding the optional
401 /// system prompt. `base_url` should include the API root (e.g. ending in
402 /// `/v1`); the chat path is appended.
403 pub fn new(
404 base_url: impl Into<String>,
405 api_key: impl Into<String>,
406 model: impl Into<String>,
407 system_prompt: Option<String>,
408 ) -> Self {
409 let mut history = Vec::new();
410 if let Some(prompt) = system_prompt {
411 history.push(serde_json::json!({ "role": "system", "content": prompt }));
412 }
413 Self {
414 base_url: base_url.into(),
415 api_key: api_key.into(),
416 model: model.into(),
417 http: reqwest::Client::new(),
418 history: Mutex::new(history),
419 }
420 }
421
422 fn tools_json(tools: &[ToolSpec]) -> Vec<serde_json::Value> {
423 tools
424 .iter()
425 .map(|tool| {
426 // parameters_schema is a JSON string; parse it, defaulting to an
427 // empty object schema if malformed.
428 let parameters: serde_json::Value = serde_json::from_str(&tool.parameters_schema)
429 .unwrap_or_else(|_| serde_json::json!({ "type": "object", "properties": {} }));
430 serde_json::json!({
431 "type": "function",
432 "function": {
433 "name": tool.name,
434 "description": tool.description,
435 "parameters": parameters,
436 }
437 })
438 })
439 .collect()
440 }
441
442 /// POST the current history (plus `tools`) and apply the assistant reply to
443 /// history, returning the neutral turn result. Streams via SSE when `sink`
444 /// is set; otherwise reads a single JSON response.
445 async fn complete(
446 &self,
447 tools: Option<&[ToolSpec]>,
448 sink: Option<&TokenSink>,
449 ) -> Result<TurnResult, BackendError> {
450 let url = format!("{}/chat/completions", self.base_url.trim_end_matches('/'));
451 let streaming = sink.is_some();
452
453 let mut body = serde_json::json!({
454 "model": self.model,
455 "messages": *self.history.lock().await,
456 "stream": streaming,
457 });
458 if let Some(tools) = tools
459 && !tools.is_empty()
460 {
461 body["tools"] = serde_json::Value::Array(Self::tools_json(tools));
462 }
463
464 let response = self
465 .http
466 .post(&url)
467 .bearer_auth(&self.api_key)
468 .json(&body)
469 .send()
470 .await
471 .map_err(|error| format!("request to {url} failed: {error}"))?;
472
473 if !response.status().is_success() {
474 let status = response.status();
475 let detail = response.text().await.unwrap_or_default();
476 return Err(format!("endpoint returned {status}: {detail}"));
477 }
478
479 if let Some(sink) = sink {
480 self.consume_stream(response, sink).await
481 } else {
482 self.consume_json(response).await
483 }
484 }
485
486 /// Parse a single non-streaming chat-completion response.
487 async fn consume_json(&self, response: reqwest::Response) -> Result<TurnResult, BackendError> {
488 let parsed: ChatCompletion = response
489 .json()
490 .await
491 .map_err(|error| format!("response parse error: {error}"))?;
492
493 let message = parsed
494 .choices
495 .into_iter()
496 .next()
497 .map(|choice| choice.message)
498 .ok_or_else(|| "endpoint returned no choices".to_string())?;
499
500 let text = message.content.clone().unwrap_or_default();
501 let tool_calls: Vec<ToolCall> = message
502 .tool_calls
503 .iter()
504 .flatten()
505 .map(|call| ToolCall {
506 id: call.id.clone(),
507 name: call.function.name.clone(),
508 arguments: call.function.arguments.clone(),
509 })
510 .collect();
511
512 // Record the assistant turn so later tool results have context.
513 self.history.lock().await.push(message.into_history_value());
514
515 Ok(TurnResult { text, tool_calls })
516 }
517
518 /// Consume an OpenAI Server-Sent Events stream, forwarding content deltas to
519 /// `sink` and reassembling any tool calls (which arrive fragmented across
520 /// chunks, keyed by `index`).
521 async fn consume_stream(
522 &self,
523 response: reqwest::Response,
524 sink: &TokenSink,
525 ) -> Result<TurnResult, BackendError> {
526 use futures::StreamExt;
527
528 let mut stream = response.bytes_stream();
529 // Newlines are ASCII, so splitting raw bytes on `\n` never bisects a
530 // multibyte UTF-8 sequence; we only lossily decode whole lines.
531 let mut buffer: Vec<u8> = Vec::new();
532 let mut text = String::new();
533 let mut tool_accum: Vec<StreamingToolCall> = Vec::new();
534 let mut done = false;
535
536 while let Some(item) = stream.next().await {
537 let bytes = item.map_err(|error| format!("stream read error: {error}"))?;
538 buffer.extend_from_slice(&bytes);
539
540 while let Some(pos) = buffer.iter().position(|&b| b == b'\n') {
541 let line: Vec<u8> = buffer.drain(..=pos).collect();
542 let line = String::from_utf8_lossy(&line);
543 let line = line.trim();
544
545 let Some(data) = line.strip_prefix("data:") else {
546 continue;
547 };
548 let data = data.trim();
549 if data == "[DONE]" {
550 done = true;
551 break;
552 }
553 if data.is_empty() {
554 continue;
555 }
556
557 let chunk: StreamCompletion = match serde_json::from_str(data) {
558 Ok(chunk) => chunk,
559 // Skip keep-alive comments and anything we can't parse rather
560 // than aborting a turn over one malformed frame.
561 Err(_) => continue,
562 };
563
564 let Some(choice) = chunk.choices.into_iter().next() else {
565 continue;
566 };
567 if let Some(content) = choice.delta.content
568 && !content.is_empty()
569 {
570 text.push_str(&content);
571 if sink.send(content).is_err() {
572 // Consumer dropped (turn cancelled) — stop reading.
573 done = true;
574 break;
575 }
576 }
577 for delta in choice.delta.tool_calls.into_iter().flatten() {
578 let index = delta.index.unwrap_or(0) as usize;
579 if tool_accum.len() <= index {
580 tool_accum.resize_with(index + 1, StreamingToolCall::default);
581 }
582 let slot = &mut tool_accum[index];
583 if let Some(id) = delta.id {
584 slot.id = id;
585 }
586 if let Some(function) = delta.function {
587 if let Some(name) = function.name {
588 slot.name = name;
589 }
590 if let Some(arguments) = function.arguments {
591 slot.arguments.push_str(&arguments);
592 }
593 }
594 }
595 }
596
597 if done {
598 break;
599 }
600 }
601
602 let tool_calls: Vec<ToolCall> = tool_accum
603 .iter()
604 .filter(|call| !call.name.is_empty())
605 .enumerate()
606 .map(|(index, call)| ToolCall {
607 id: if call.id.is_empty() {
608 format!("call_{index}")
609 } else {
610 call.id.clone()
611 },
612 name: call.name.clone(),
613 arguments: call.arguments.clone(),
614 })
615 .collect();
616
617 // Record the assistant turn so later tool results have context.
618 self.history
619 .lock()
620 .await
621 .push(streamed_assistant_history(&text, &tool_calls));
622
623 Ok(TurnResult { text, tool_calls })
624 }
625 }
626
627 /// One tool call being reassembled from streamed deltas.
628 #[derive(Default)]
629 struct StreamingToolCall {
630 id: String,
631 name: String,
632 arguments: String,
633 }
634
635 /// Rebuild the assistant message for replay in history after a streamed turn,
636 /// preserving any tool calls so the follow-up request is well-formed. Mirrors
637 /// [`ResponseMessage::into_history_value`] for the non-streaming path.
638 fn streamed_assistant_history(text: &str, tool_calls: &[ToolCall]) -> serde_json::Value {
639 let mut message = serde_json::json!({ "role": "assistant" });
640 message["content"] = if text.is_empty() {
641 serde_json::Value::Null
642 } else {
643 serde_json::Value::String(text.to_string())
644 };
645 if !tool_calls.is_empty() {
646 message["tool_calls"] = serde_json::json!(
647 tool_calls
648 .iter()
649 .map(|call| serde_json::json!({
650 "id": call.id,
651 "type": "function",
652 "function": {
653 "name": call.name,
654 "arguments": call.arguments,
655 }
656 }))
657 .collect::<Vec<_>>()
658 );
659 }
660 message
661 }
662
663 #[async_trait]
664 impl InferenceBackend for OpenAiBackend {
665 async fn send_message_with_tools(
666 &self,
667 text: &str,
668 tools: &[ToolSpec],
669 sink: Option<&TokenSink>,
670 ) -> Result<TurnResult, BackendError> {
671 self.history
672 .lock()
673 .await
674 .push(serde_json::json!({ "role": "user", "content": text }));
675 self.complete(Some(tools), sink).await
676 }
677
678 async fn send_tool_results(
679 &self,
680 results: Vec<ToolResult>,
681 tools: Option<&[ToolSpec]>,
682 sink: Option<&TokenSink>,
683 ) -> Result<TurnResult, BackendError> {
684 {
685 let mut history = self.history.lock().await;
686 for result in results {
687 history.push(serde_json::json!({
688 "role": "tool",
689 "tool_call_id": result.tool_call_id,
690 "content": result.content,
691 }));
692 }
693 }
694 self.complete(tools, sink).await
695 }
696
697 async fn record_cancelled_tool_results(&self, results: Vec<ToolResult>) {
698 let mut history = self.history.lock().await;
699 for result in results {
700 history.push(serde_json::json!({
701 "role": "tool",
702 "tool_call_id": result.tool_call_id,
703 "content": result.content,
704 }));
705 }
706 }
707
708 fn is_remote(&self) -> bool {
709 true
710 }
711
712 async fn history_snapshot(&self) -> Vec<serde_json::Value> {
713 self.history.lock().await.clone()
714 }
715
716 async fn restore_history(&self, history: Vec<serde_json::Value>) {
717 // The snapshot includes the seeded system message, so a wholesale
718 // replacement restores exactly what was saved.
719 *self.history.lock().await = history;
720 }
721
722 async fn compact_history(&self, keep_last: usize) -> Result<(), BackendError> {
723 let snapshot: Vec<serde_json::Value> = self.history.lock().await.clone();
724
725 // Ask the endpoint for a summary of the conversation so far, through
726 // the ordinary completion machinery (non-streaming).
727 self.history
728 .lock()
729 .await
730 .push(serde_json::json!({ "role": "user", "content": SUMMARIZE_PROMPT }));
731 let summary = match self.complete(None, None).await {
732 Ok(result) => result.text,
733 Err(error) => {
734 // Roll back the summarization request; the turn never happened.
735 *self.history.lock().await = snapshot;
736 return Err(error);
737 }
738 };
739
740 let system = snapshot
741 .first()
742 .filter(|message| message["role"] == "system")
743 .cloned();
744 let non_system: Vec<serde_json::Value> = snapshot
745 .iter()
746 .filter(|message| message["role"] != "system")
747 .cloned()
748 .collect();
749 let tail_start = non_system.len().saturating_sub(keep_last);
750 let mut tail = non_system[tail_start..].to_vec();
751 // Drop leading tool results whose assistant tool-call message was
752 // summarized away — strict endpoints reject orphaned `role: "tool"`
753 // entries on the very next request.
754 while tail
755 .first()
756 .is_some_and(|message| message["role"] == "tool")
757 {
758 tail.remove(0);
759 }
760
761 let mut rebuilt = Vec::new();
762 if let Some(system) = system {
763 rebuilt.push(system);
764 }
765 rebuilt.push(serde_json::json!({
766 "role": "user",
767 "content": format!("[Conversation summary]\n{summary}"),
768 }));
769 rebuilt.extend(tail);
770 *self.history.lock().await = rebuilt;
771 Ok(())
772 }
773 }
774
775 // ── OpenAI response shapes ────────────────────────────────────────────────────────
776
777 #[derive(Debug, Deserialize)]
778 struct ChatCompletion {
779 #[serde(default)]
780 choices: Vec<CompletionChoice>,
781 }
782
783 #[derive(Debug, Deserialize)]
784 struct CompletionChoice {
785 message: ResponseMessage,
786 }
787
788 #[derive(Debug, Deserialize)]
789 struct ResponseMessage {
790 #[serde(default)]
791 content: Option<String>,
792 #[serde(default)]
793 tool_calls: Option<Vec<ResponseToolCall>>,
794 }
795
796 impl ResponseMessage {
797 /// Reconstruct the assistant message for replay in history, preserving any
798 /// tool calls so the follow-up request is well-formed.
799 fn into_history_value(self) -> serde_json::Value {
800 let mut message = serde_json::json!({ "role": "assistant" });
801 message["content"] = match self.content {
802 Some(text) => serde_json::Value::String(text),
803 None => serde_json::Value::Null,
804 };
805 if let Some(tool_calls) = self.tool_calls {
806 message["tool_calls"] = serde_json::json!(
807 tool_calls
808 .into_iter()
809 .map(|call| serde_json::json!({
810 "id": call.id,
811 "type": "function",
812 "function": {
813 "name": call.function.name,
814 "arguments": call.function.arguments,
815 }
816 }))
817 .collect::<Vec<_>>()
818 );
819 }
820 message
821 }
822 }
823
824 #[derive(Debug, Deserialize)]
825 struct ResponseToolCall {
826 id: String,
827 function: ResponseFunction,
828 }
829
830 #[derive(Debug, Deserialize)]
831 struct ResponseFunction {
832 name: String,
833 #[serde(default)]
834 arguments: String,
835 }
836
837 // ── OpenAI streaming (SSE) chunk shapes ─────────────────────────────────────────
838
839 #[derive(Debug, Deserialize)]
840 struct StreamCompletion {
841 #[serde(default)]
842 choices: Vec<StreamChoice>,
843 }
844
845 #[derive(Debug, Deserialize)]
846 struct StreamChoice {
847 #[serde(default)]
848 delta: StreamDelta,
849 }
850
851 #[derive(Debug, Default, Deserialize)]
852 struct StreamDelta {
853 #[serde(default)]
854 content: Option<String>,
855 #[serde(default)]
856 tool_calls: Option<Vec<StreamToolCallDelta>>,
857 }
858
859 #[derive(Debug, Deserialize)]
860 struct StreamToolCallDelta {
861 #[serde(default)]
862 index: Option<u32>,
863 #[serde(default)]
864 id: Option<String>,
865 #[serde(default)]
866 function: Option<StreamFunctionDelta>,
867 }
868
869 #[derive(Debug, Deserialize)]
870 struct StreamFunctionDelta {
871 #[serde(default)]
872 name: Option<String>,
873 #[serde(default)]
874 arguments: Option<String>,
875 }
876
877 #[cfg(test)]
878 mod tests {
879 use super::*;
880
881 #[test]
882 fn tools_json_wraps_function_schema() {
883 let tools = vec![ToolSpec {
884 name: "read_file".to_string(),
885 description: "Read a file".to_string(),
886 parameters_schema: r#"{"type":"object","properties":{"path":{"type":"string"}}}"#
887 .to_string(),
888 }];
889 let json = OpenAiBackend::tools_json(&tools);
890 assert_eq!(json[0]["type"], "function");
891 assert_eq!(json[0]["function"]["name"], "read_file");
892 assert_eq!(
893 json[0]["function"]["parameters"]["properties"]["path"]["type"],
894 "string"
895 );
896 }
897
898 #[test]
899 fn malformed_schema_falls_back_to_empty_object() {
900 let tools = vec![ToolSpec {
901 name: "x".to_string(),
902 description: String::new(),
903 parameters_schema: "not json".to_string(),
904 }];
905 let json = OpenAiBackend::tools_json(&tools);
906 assert_eq!(json[0]["function"]["parameters"]["type"], "object");
907 }
908
909 #[test]
910 fn streamed_assistant_history_omits_empty_tool_calls() {
911 let value = streamed_assistant_history("hello", &[]);
912 assert_eq!(value["role"], "assistant");
913 assert_eq!(value["content"], "hello");
914 assert!(value.get("tool_calls").is_none());
915 }
916
917 #[test]
918 fn streamed_assistant_history_preserves_tool_calls() {
919 let calls = vec![ToolCall {
920 id: "call_0".to_string(),
921 name: "read_file".to_string(),
922 arguments: r#"{"path":"a.rs"}"#.to_string(),
923 }];
924 let value = streamed_assistant_history("", &calls);
925 assert!(value["content"].is_null());
926 assert_eq!(value["tool_calls"][0]["id"], "call_0");
927 assert_eq!(value["tool_calls"][0]["type"], "function");
928 assert_eq!(value["tool_calls"][0]["function"]["name"], "read_file");
929 assert_eq!(
930 value["tool_calls"][0]["function"]["arguments"],
931 r#"{"path":"a.rs"}"#
932 );
933 }
934
935 #[tokio::test]
936 async fn cancelled_tool_results_close_out_history() {
937 let backend = OpenAiBackend::new("http://localhost", "", "test-model", None);
938 backend
939 .history
940 .lock()
941 .await
942 .push(streamed_assistant_history(
943 "",
944 &[ToolCall {
945 id: "call_9".to_string(),
946 name: "run_command".to_string(),
947 arguments: r#"{"command":"ls"}"#.to_string(),
948 }],
949 ));
950
951 backend
952 .record_cancelled_tool_results(vec![ToolResult {
953 tool_call_id: "call_9".to_string(),
954 content: "cancelled by the user".to_string(),
955 }])
956 .await;
957
958 let history = backend.history.lock().await;
959 let last = history.last().unwrap();
960 assert_eq!(last["role"], "tool");
961 assert_eq!(last["tool_call_id"], "call_9");
962 assert_eq!(last["content"], "cancelled by the user");
963 }
964
965 #[test]
966 fn estimate_tokens_scales_with_serialized_size() {
967 assert_eq!(estimate_tokens(&[]), 0);
968
969 let short = vec![serde_json::json!({ "role": "user", "content": "hi" })];
970 let long = vec![serde_json::json!({ "role": "user", "content": "x".repeat(4_000) })];
971 let short_estimate = estimate_tokens(&short);
972 let long_estimate = estimate_tokens(&long);
973
974 assert!(short_estimate > 0, "non-empty history estimates > 0 tokens");
975 assert!(long_estimate > short_estimate, "longer history costs more");
976 // 4,000 content chars / 4 ≈ 1,000 tokens, plus a little JSON framing.
977 assert!((1_000..1_100).contains(&long_estimate), "{long_estimate}");
978 }
979
980 #[tokio::test]
981 async fn openai_snapshot_restore_round_trips_exactly() {
982 let backend = OpenAiBackend::new("http://localhost", "", "m", Some("be helpful".into()));
983 {
984 let mut history = backend.history.lock().await;
985 history.push(serde_json::json!({ "role": "user", "content": "hello" }));
986 history.push(streamed_assistant_history(
987 "",
988 &[ToolCall {
989 id: "call_1".to_string(),
990 name: "read_file".to_string(),
991 arguments: r#"{"path":"a.rs"}"#.to_string(),
992 }],
993 ));
994 history.push(serde_json::json!({
995 "role": "tool", "tool_call_id": "call_1", "content": "fn main() {}",
996 }));
997 history.push(serde_json::json!({ "role": "assistant", "content": "done" }));
998 }
999 let snapshot = backend.history_snapshot().await;
1000 assert_eq!(
1001 snapshot[0]["role"], "system",
1002 "snapshot keeps the system message"
1003 );
1004
1005 // Restoring into a backend seeded with a *different* system prompt must
1006 // replace everything, including that seed.
1007 let restored = OpenAiBackend::new("http://localhost", "", "m", Some("other seed".into()));
1008 restored.restore_history(snapshot.clone()).await;
1009 assert_eq!(restored.history_snapshot().await, snapshot);
1010 }
1011
1012 /// Minimal scripted OpenAI-compatible endpoint: accepts one HTTP request on
1013 /// a std listener and answers with a fixed non-streaming completion.
1014 fn spawn_completion_stub(summary: &str) -> std::net::SocketAddr {
1015 use std::io::{Read, Write};
1016
1017 let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
1018 let addr = listener.local_addr().unwrap();
1019 let body = serde_json::json!({
1020 "choices": [{ "message": { "role": "assistant", "content": summary } }]
1021 })
1022 .to_string();
1023 std::thread::spawn(move || {
1024 let (mut stream, _) = listener.accept().unwrap();
1025 // Read until the full request (headers + content-length body) is in.
1026 let mut request = Vec::new();
1027 let mut chunk = [0u8; 4096];
1028 loop {
1029 let n = stream.read(&mut chunk).unwrap_or(0);
1030 if n == 0 {
1031 break;
1032 }
1033 request.extend_from_slice(&chunk[..n]);
1034 if let Some(headers_end) =
1035 request.windows(4).position(|window| window == b"\r\n\r\n")
1036 {
1037 let headers = String::from_utf8_lossy(&request[..headers_end]);
1038 let content_length = headers
1039 .lines()
1040 .find_map(|line| {
1041 line.to_ascii_lowercase()
1042 .strip_prefix("content-length:")
1043 .map(|value| value.trim().parse::<usize>().unwrap_or(0))
1044 })
1045 .unwrap_or(0);
1046 if request.len() >= headers_end + 4 + content_length {
1047 break;
1048 }
1049 }
1050 }
1051 let response = format!(
1052 "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\n\
1053 content-length: {}\r\nconnection: close\r\n\r\n{}",
1054 body.len(),
1055 body
1056 );
1057 let _ = stream.write_all(response.as_bytes());
1058 });
1059 addr
1060 }
1061
1062 #[tokio::test]
1063 async fn compact_history_rebuilds_system_summary_and_tail() {
1064 let addr = spawn_completion_stub("We refactored backend.rs; tests pass.");
1065 let backend = OpenAiBackend::new(
1066 format!("http://{addr}/v1"),
1067 "test-key",
1068 "test-model",
1069 Some("be helpful".into()),
1070 );
1071 {
1072 let mut history = backend.history.lock().await;
1073 for i in 0..5 {
1074 let role = if i % 2 == 0 { "user" } else { "assistant" };
1075 history.push(serde_json::json!({
1076 "role": role, "content": format!("message {i}"),
1077 }));
1078 }
1079 }
1080
1081 backend.compact_history(2).await.unwrap();
1082
1083 let history = backend.history_snapshot().await;
1084 assert_eq!(history.len(), 4, "system + summary + last 2: {history:?}");
1085 assert_eq!(history[0]["role"], "system");
1086 assert_eq!(history[0]["content"], "be helpful");
1087 assert_eq!(history[1]["role"], "user");
1088 let summary_text = history[1]["content"].as_str().unwrap();
1089 assert!(summary_text.starts_with("[Conversation summary]\n"));
1090 assert!(summary_text.contains("We refactored backend.rs; tests pass."));
1091 assert_eq!(
1092 history[2],
1093 serde_json::json!({ "role": "assistant", "content": "message 3" })
1094 );
1095 assert_eq!(
1096 history[3],
1097 serde_json::json!({ "role": "user", "content": "message 4" })
1098 );
1099 }
1100
1101 #[tokio::test]
1102 async fn compact_history_failure_leaves_history_intact() {
1103 // No listener at this address: the summarization request fails, and
1104 // history must roll back to exactly what it was.
1105 let backend =
1106 OpenAiBackend::new("http://127.0.0.1:9", "", "test-model", Some("sys".into()));
1107 backend
1108 .history
1109 .lock()
1110 .await
1111 .push(serde_json::json!({ "role": "user", "content": "hello" }));
1112 let before = backend.history_snapshot().await;
1113
1114 assert!(backend.compact_history(2).await.is_err());
1115 assert_eq!(backend.history_snapshot().await, before);
1116 }
1117
1118 #[test]
1119 fn assistant_message_with_tool_calls_round_trips() {
1120 let message = ResponseMessage {
1121 content: None,
1122 tool_calls: Some(vec![ResponseToolCall {
1123 id: "call_1".to_string(),
1124 function: ResponseFunction {
1125 name: "read_file".to_string(),
1126 arguments: r#"{"path":"a.rs"}"#.to_string(),
1127 },
1128 }]),
1129 };
1130 let value = message.into_history_value();
1131 assert_eq!(value["role"], "assistant");
1132 assert!(value["content"].is_null());
1133 assert_eq!(value["tool_calls"][0]["id"], "call_1");
1134 assert_eq!(value["tool_calls"][0]["type"], "function");
1135 assert_eq!(value["tool_calls"][0]["function"]["name"], "read_file");
1136 }
1137 }