1use std::collections::{BTreeMap, HashSet};
12use std::io::{IsTerminal, Write};
13use std::path::{Path, PathBuf};
14use std::sync::{Arc, atomic::AtomicBool};
15use std::time::{Duration, Instant};
16use std::{error::Error, fmt};
17
18use anyhow::{Context, Result, anyhow};
19use serde_json::{Value, json};
20
21use crate::color::Style;
22use crate::config::ModelConfig;
23use crate::interrupt;
24use crate::model_timeout::ModelTurnDeadline;
25use crate::tools;
26
27struct Spinner {
30 handle: Option<tokio::task::JoinHandle<()>>,
31}
32
33impl Spinner {
34 fn start(label: &str, color: Style) -> Spinner {
35 if !std::io::stderr().is_terminal() {
36 return Spinner { handle: None };
37 }
38 let label = label.to_string();
39 let handle = tokio::spawn(async move {
40 let frames = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"];
41 let mut i = 0usize;
42 loop {
43 eprint!(
44 "\r{} {}",
45 color.dim(frames[i % frames.len()]),
46 color.dim(&label)
47 );
48 let _ = std::io::stderr().flush();
49 i += 1;
50 tokio::time::sleep(Duration::from_millis(80)).await;
51 }
52 });
53 Spinner {
54 handle: Some(handle),
55 }
56 }
57
58 fn stop(&mut self) {
59 if let Some(h) = self.handle.take() {
60 h.abort();
61 eprint!("\r\x1b[K");
62 let _ = std::io::stderr().flush();
63 }
64 }
65}
66
67impl Drop for Spinner {
68 fn drop(&mut self) {
69 self.stop();
70 }
71}
72
73fn system_message(system: &str, cache: bool) -> Value {
78 if cache {
79 json!({
80 "role": "system",
81 "content": [{"type": "text", "text": system, "cache_control": {"type": "ephemeral"}}],
82 })
83 } else {
84 json!({"role": "system", "content": system})
85 }
86}
87
88fn path_with_line_range(path: &str, v: &Value) -> String {
89 let start = v.get("start_line").and_then(|x| x.as_u64());
90 let end = v.get("end_line").and_then(|x| x.as_u64());
91 match (start, end) {
92 (None, None) => format!("{path}:1-EOF"),
93 (Some(start), None) => format!("{path}:{start}-EOF"),
94 (None, Some(end)) => format!("{path}:1-{end}"),
95 (Some(start), Some(end)) => format!("{path}:{start}-{end}"),
96 }
97}
98
99fn tool_call_label(name: &str, args: &str) -> String {
102 let hint = serde_json::from_str::<Value>(args)
103 .ok()
104 .and_then(|v| {
105 let scheduler = v.get("scheduler").and_then(|x| x.as_str());
106 if name == "read_file"
107 && let Some(path) = v.get("path").and_then(|x| x.as_str())
108 {
109 return Some(path_with_line_range(path, &v));
110 }
111 if name == "read_scheduler_file"
112 && let (Some(scheduler), Some(path)) =
113 (scheduler, v.get("path").and_then(|x| x.as_str()))
114 {
115 return Some(format!("{scheduler}/{}", path_with_line_range(path, &v)));
116 }
117 if name == "grep_schedulers"
118 && let (Some(scheduler), Some(pattern)) =
119 (scheduler, v.get("pattern").and_then(|x| x.as_str()))
120 {
121 return Some(format!("{scheduler}:{pattern}"));
122 }
123 v.get("path")
124 .or_else(|| v.get("pattern"))
125 .or_else(|| v.get("url"))
126 .and_then(|x| x.as_str())
127 .map(|s| s.to_string())
128 })
129 .unwrap_or_else(|| preview(args, 80).replace('\n', " "));
130 format!("{name}({hint})")
131}
132
133pub struct ToolLoopConfig {
135 pub sandbox: PathBuf,
136 pub scheds_root: Option<PathBuf>,
137 pub allow_edit: bool,
138 pub max_iterations: u32,
139}
140
141#[derive(Debug)]
143pub struct ApiStatusError {
144 status: reqwest::StatusCode,
145 body: String,
146}
147
148impl fmt::Display for ApiStatusError {
149 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
150 write!(f, "API error {}: {}", self.status, self.body)
151 }
152}
153
154impl Error for ApiStatusError {}
155
156pub fn is_api_error(err: &anyhow::Error) -> bool {
159 err.downcast_ref::<ApiStatusError>().is_some()
160 || err
161 .chain()
162 .any(|cause| cause.downcast_ref::<reqwest::Error>().is_some())
163}
164
165pub fn error_summary(err: &anyhow::Error) -> String {
167 let raw = err
168 .downcast_ref::<ApiStatusError>()
169 .map(|e| e.to_string())
170 .unwrap_or_else(|| err.to_string());
171 let one_line = raw.split_whitespace().collect::<Vec<_>>().join(" ");
172 const MAX: usize = 500;
173 if one_line.chars().count() > MAX {
174 format!("{}...", one_line.chars().take(MAX).collect::<String>())
175 } else {
176 one_line
177 }
178}
179
180fn apply_bearer(req: reqwest::RequestBuilder, api_key: &str) -> reqwest::RequestBuilder {
181 if api_key.is_empty() {
182 req
183 } else {
184 req.header("Authorization", format!("Bearer {api_key}"))
185 }
186}
187
188fn arguments_to_string(v: &Value) -> String {
190 match v {
191 Value::String(s) => s.clone(),
192 other => other.to_string(),
193 }
194}
195
196fn normalize_tool_call_arguments(tc: &mut Value) {
198 let Some(func) = tc.get_mut("function").and_then(|f| f.as_object_mut()) else {
199 return;
200 };
201 let args = func
202 .get("arguments")
203 .map(arguments_to_string)
204 .unwrap_or_else(|| "{}".to_string());
205 func.insert("arguments".to_string(), Value::String(args));
206}
207
208fn normalize_tool_calls(mut calls: Vec<Value>) -> Vec<Value> {
209 for tc in &mut calls {
210 normalize_tool_call_arguments(tc);
211 }
212 calls
213}
214
215fn sanitize_tool_calls_for_history(calls: &[Value]) -> Vec<Value> {
227 calls
228 .iter()
229 .map(|tc| {
230 let mut tc = tc.clone();
231 if let Some(func) = tc.get_mut("function").and_then(|f| f.as_object_mut()) {
232 let valid = func
233 .get("arguments")
234 .and_then(|a| a.as_str())
235 .is_some_and(|s| serde_json::from_str::<Value>(s).is_ok());
236 if !valid {
237 func.insert("arguments".to_string(), Value::String("{}".to_string()));
238 }
239 }
240 tc
241 })
242 .collect()
243}
244
245fn message_content_to_string(message: &Value) -> String {
246 match message.get("content") {
247 Some(Value::String(s)) => s.clone(),
248 Some(Value::Array(parts)) => parts
249 .iter()
250 .filter_map(|p| p.get("text").and_then(|t| t.as_str()))
251 .collect::<Vec<_>>()
252 .join(""),
253 _ => String::new(),
254 }
255}
256
257const DUMP_MAX_CHARS: usize = 8000;
259const SLOW_TOOL_LOG_AFTER: Duration = Duration::from_millis(500);
260
261fn preview(s: &str, max: usize) -> String {
263 let n = s.chars().count();
264 if n > max {
265 let head: String = s.chars().take(max).collect();
266 format!("{head}\n [... {} more chars elided ...]", n - max)
267 } else {
268 s.to_string()
269 }
270}
271
272const TOOL_CALL_OPEN: &str = "<tool_call>";
273const TOOL_CALL_CLOSE: &str = "</tool_call>";
274
275fn parse_one_text_call(inner: &str, idx: usize) -> Option<Value> {
278 let v: Value = serde_json::from_str(inner.trim()).ok()?;
279 let name = v.get("name").and_then(|n| n.as_str())?;
280 let args = v
281 .get("arguments")
282 .or_else(|| v.get("parameters"))
283 .cloned()
284 .unwrap_or_else(|| json!({}));
285 Some(json!({
286 "id": format!("text_call_{idx}"),
287 "type": "function",
288 "function": {"name": name, "arguments": arguments_to_string(&args)},
289 }))
290}
291
292fn parse_text_tool_calls(content: &str) -> (Vec<Value>, String) {
297 let mut calls = Vec::new();
298 let mut stripped = String::new();
299 let mut rest = content;
300 let mut idx = 0usize;
301 while let Some(start) = rest.find(TOOL_CALL_OPEN) {
302 stripped.push_str(&rest[..start]);
303 let after = &rest[start + TOOL_CALL_OPEN.len()..];
304 let (inner, consumed) = match after.find(TOOL_CALL_CLOSE) {
305 Some(end) => (
306 &after[..end],
307 start + TOOL_CALL_OPEN.len() + end + TOOL_CALL_CLOSE.len(),
308 ),
309 None => (after, rest.len()),
310 };
311 if let Some(call) = parse_one_text_call(inner, idx) {
312 calls.push(call);
313 idx += 1;
314 }
315 rest = &rest[consumed..];
316 }
317 stripped.push_str(rest);
318 (calls, stripped.trim().to_string())
319}
320
321#[derive(Default)]
322struct PendingToolCall {
323 id: String,
324 ty: String,
325 name: String,
326 arguments: String,
327}
328
329#[derive(Default)]
330struct StreamedChat {
331 content: String,
332 tool_calls: BTreeMap<usize, PendingToolCall>,
333 usage: Option<crate::usage::Usage>,
334 raw_events: String,
335}
336
337impl StreamedChat {
338 fn into_message(self) -> Value {
339 let content = if self.content.is_empty() {
340 Value::Null
341 } else {
342 Value::String(self.content)
343 };
344 let mut message = json!({
345 "role": "assistant",
346 "content": content,
347 });
348 if !self.tool_calls.is_empty() {
349 let calls: Vec<Value> = self
350 .tool_calls
351 .into_iter()
352 .map(|(idx, call)| {
353 json!({
354 "id": if call.id.is_empty() { format!("stream_call_{idx}") } else { call.id },
355 "type": if call.ty.is_empty() { "function".to_string() } else { call.ty },
356 "function": {
357 "name": call.name,
358 "arguments": if call.arguments.is_empty() { "{}".to_string() } else { call.arguments },
359 },
360 })
361 })
362 .collect();
363 message.as_object_mut().unwrap().insert(
364 "tool_calls".to_string(),
365 Value::Array(normalize_tool_calls(calls)),
366 );
367 }
368 message
369 }
370}
371
372enum StreamPrintState {
373 Visible,
374 InToolCall,
375}
376
377struct StreamTextPrinter {
378 enabled: bool,
379 started: bool,
380 saw_tool_call: bool,
381 last_was_newline: bool,
382 state: StreamPrintState,
383 buf: String,
384}
385
386fn floor_char_boundary(s: &str, idx: usize) -> usize {
387 let mut idx = idx.min(s.len());
388 while idx > 0 && !s.is_char_boundary(idx) {
389 idx -= 1;
390 }
391 idx
392}
393
394impl StreamTextPrinter {
395 fn new(enabled: bool) -> Self {
396 Self {
397 enabled,
398 started: false,
399 saw_tool_call: false,
400 last_was_newline: false,
401 state: StreamPrintState::Visible,
402 buf: String::new(),
403 }
404 }
405
406 fn emit(&mut self, text: &str) -> Result<()> {
407 if text.is_empty() {
408 return Ok(());
409 }
410 self.started = true;
411 print!("{text}");
412 std::io::stdout()
413 .flush()
414 .context("flush streamed assistant text")?;
415 self.last_was_newline = text.ends_with('\n');
416 Ok(())
417 }
418
419 fn note_tool_call(&mut self) {
420 self.saw_tool_call = true;
421 }
422
423 fn emit_before_tool_call(&mut self, text: &str) -> Result<()> {
424 let visible = text.trim_end();
425 if !visible.is_empty() {
426 self.emit(visible)?;
427 }
428 Ok(())
429 }
430
431 fn push(&mut self, text: &str) -> Result<()> {
432 if !self.enabled {
433 return Ok(());
434 }
435 self.buf.push_str(text);
436 loop {
437 match self.state {
438 StreamPrintState::Visible => {
439 if let Some(pos) = self.buf.find(TOOL_CALL_OPEN) {
440 let visible = self.buf[..pos].to_string();
441 self.note_tool_call();
442 self.emit_before_tool_call(&visible)?;
443 self.buf.drain(..pos + TOOL_CALL_OPEN.len());
444 self.state = StreamPrintState::InToolCall;
445 } else {
446 let keep = TOOL_CALL_OPEN.len().saturating_sub(1);
447 let emit_len = self.buf.len().saturating_sub(keep);
448 if emit_len == 0 {
449 break;
450 }
451 let emit_len = floor_char_boundary(&self.buf, emit_len);
452 if emit_len == 0 {
453 break;
454 }
455 let visible = self.buf[..emit_len].to_string();
456 if self.started || !visible.trim().is_empty() {
457 self.emit(&visible)?;
458 }
459 self.buf.drain(..emit_len);
460 break;
461 }
462 }
463 StreamPrintState::InToolCall => {
464 if let Some(pos) = self.buf.find(TOOL_CALL_CLOSE) {
465 self.buf.drain(..pos + TOOL_CALL_CLOSE.len());
466 self.state = StreamPrintState::Visible;
467 } else {
468 let keep = TOOL_CALL_CLOSE.len().saturating_sub(1);
469 if self.buf.len() > keep {
470 let drop_len = self.buf.len() - keep;
471 let drop_len = floor_char_boundary(&self.buf, drop_len);
472 if drop_len > 0 {
473 self.buf.drain(..drop_len);
474 }
475 }
476 break;
477 }
478 }
479 }
480 }
481 Ok(())
482 }
483
484 fn finish(&mut self) -> Result<()> {
485 if !self.enabled {
486 return Ok(());
487 }
488 if matches!(self.state, StreamPrintState::Visible) && !self.buf.is_empty() {
489 let visible = std::mem::take(&mut self.buf);
490 if !visible.trim().is_empty() {
491 if self.saw_tool_call {
492 self.emit(visible.trim_end())?;
493 } else {
494 self.emit(&visible)?;
495 }
496 }
497 }
498 if self.started && !self.last_was_newline {
499 println!();
500 }
501 Ok(())
502 }
503}
504
505fn append_json_text(dst: &mut String, v: &Value) {
506 match v {
507 Value::String(s) => dst.push_str(s),
508 Value::Null => {}
509 other => dst.push_str(&other.to_string()),
510 }
511}
512
513fn apply_stream_delta(delta: &Value, streamed: &mut StreamedChat) -> Result<Option<String>> {
514 let mut visible = String::new();
515 if let Some(content) = delta.get("content")
516 && let Some(s) = content.as_str()
517 {
518 streamed.content.push_str(s);
519 visible.push_str(s);
520 }
521
522 if let Some(calls) = delta.get("tool_calls").and_then(|v| v.as_array()) {
523 for (fallback_idx, tc) in calls.iter().enumerate() {
524 let idx = tc
525 .get("index")
526 .and_then(|v| v.as_u64())
527 .and_then(|n| usize::try_from(n).ok())
528 .unwrap_or(fallback_idx);
529 let entry = streamed.tool_calls.entry(idx).or_default();
530 if let Some(id) = tc.get("id").and_then(|v| v.as_str()) {
531 if entry.id.is_empty() {
532 entry.id = id.to_string();
533 } else if entry.id != id {
534 entry.id.push_str(id);
535 }
536 }
537 if let Some(ty) = tc.get("type").and_then(|v| v.as_str()) {
538 if entry.ty.is_empty() {
539 entry.ty = ty.to_string();
540 } else if entry.ty != ty {
541 entry.ty.push_str(ty);
542 }
543 }
544 if let Some(func) = tc.get("function") {
545 if let Some(name) = func.get("name").and_then(|v| v.as_str()) {
546 entry.name.push_str(name);
547 }
548 if let Some(args) = func.get("arguments") {
549 append_json_text(&mut entry.arguments, args);
550 }
551 }
552 }
553 }
554
555 if let Some(func) = delta.get("function_call") {
556 let entry = streamed.tool_calls.entry(0).or_default();
557 if entry.id.is_empty() {
558 entry.id = "stream_call_0".to_string();
559 }
560 if entry.ty.is_empty() {
561 entry.ty = "function".to_string();
562 }
563 if let Some(name) = func.get("name").and_then(|v| v.as_str()) {
564 entry.name.push_str(name);
565 }
566 if let Some(args) = func.get("arguments") {
567 append_json_text(&mut entry.arguments, args);
568 }
569 }
570
571 Ok((!visible.is_empty()).then_some(visible))
572}
573
574fn find_sse_delimiter(buf: &[u8]) -> Option<(usize, usize)> {
575 let lf = buf.windows(2).position(|w| w == b"\n\n").map(|p| (p, 2));
576 let crlf = buf
577 .windows(4)
578 .position(|w| w == b"\r\n\r\n")
579 .map(|p| (p, 4));
580 match (lf, crlf) {
581 (Some(a), Some(b)) => Some(if a.0 <= b.0 { a } else { b }),
582 (Some(a), None) => Some(a),
583 (None, Some(b)) => Some(b),
584 (None, None) => None,
585 }
586}
587
588fn sse_data(event: &[u8]) -> Result<Option<String>> {
589 let text = std::str::from_utf8(event).context("parse streaming SSE as UTF-8")?;
590 let data = text
591 .lines()
592 .filter_map(|line| line.strip_prefix("data:").map(str::trim_start))
593 .collect::<Vec<_>>();
594 if data.is_empty() {
595 Ok(None)
596 } else {
597 Ok(Some(data.join("\n")))
598 }
599}
600
601async fn read_streamed_chat(
602 mut resp: reqwest::Response,
603 stream_stdout: bool,
604 spinner: &mut Spinner,
605) -> Result<StreamedChat> {
606 let mut streamed = StreamedChat::default();
607 let mut printer = StreamTextPrinter::new(stream_stdout);
608 let mut buf = Vec::<u8>::new();
609
610 while let Some(chunk) = resp.chunk().await.context("read streaming chat chunk")? {
611 buf.extend_from_slice(&chunk);
612 while let Some((pos, delim_len)) = find_sse_delimiter(&buf) {
613 let event = buf.drain(..pos).collect::<Vec<_>>();
614 buf.drain(..delim_len);
615 let Some(data) = sse_data(&event)? else {
616 continue;
617 };
618 if data.trim() == "[DONE]" {
619 printer.finish()?;
620 return Ok(streamed);
621 }
622 streamed.raw_events.push_str(&data);
623 streamed.raw_events.push('\n');
624 let parsed: Value = serde_json::from_str(&data)
625 .with_context(|| format!("parse streaming chat JSON event: {data}"))?;
626 if let Some(u) = parsed.get("usage").filter(|u| !u.is_null()) {
627 streamed.usage = Some(crate::usage::Usage::from_openai(u));
628 }
629 let Some(delta) = parsed
630 .get("choices")
631 .and_then(|c| c.as_array())
632 .and_then(|a| a.first())
633 .and_then(|c| c.get("delta"))
634 else {
635 continue;
636 };
637 let has_tool_delta =
638 delta.get("tool_calls").is_some() || delta.get("function_call").is_some();
639 if delta.get("content").is_some() {
640 spinner.stop();
641 }
642 if has_tool_delta {
643 printer.note_tool_call();
644 }
645 if let Some(text) = apply_stream_delta(delta, &mut streamed)? {
646 printer.push(&text)?;
647 }
648 }
649 }
650
651 printer.finish()?;
652 Ok(streamed)
653}
654
655pub async fn chat(
659 client: &reqwest::Client,
660 model: &ModelConfig,
661 system: &str,
662 user: &str,
663 tool_loop: Option<&ToolLoopConfig>,
664 verbose: bool,
665 color: Style,
666 stream_stdout: bool,
667 usage: &mut crate::usage::Usage,
668 turn_timeout: Duration,
669 interrupted: Arc<AtomicBool>,
670) -> Result<String> {
671 let deadline = ModelTurnDeadline::new(turn_timeout);
672 let url = format!("{}/chat/completions", model.base_url.trim_end_matches('/'));
673 let mut use_cache = true;
678 let mut use_stream_options = true;
679 let mut messages: Vec<Value> = vec![
680 system_message(system, use_cache),
681 json!({"role": "user", "content": user}),
682 ];
683
684 if verbose {
685 eprintln!(
686 "\n{} {}",
687 color.dim("[api]"),
688 color.bold(format!(
689 "new turn (model={}, system prompt={} chars)",
690 model.model_id,
691 system.len()
692 ))
693 );
694 eprintln!(
695 "{} {}\n{}",
696 color.dim("[api]"),
697 color.blue("user:"),
698 preview(user, DUMP_MAX_CHARS)
699 );
700 }
701
702 let mut tool_iterations = 0u32;
703 let mut edits_applied = 0usize;
704 let mut disable_tools = false;
705 let mut force_edit = false;
706 let mut force_attempts = 0u32;
707 let mut request_n = 0u32;
708 let mut seen_tool_calls: HashSet<String> = HashSet::new();
713
714 loop {
715 deadline.check()?;
716 if interrupt::requested(&interrupted) {
717 return Err(interrupt::err());
718 }
719
720 const MAX_FORCE_ATTEMPTS: u32 = 3;
724 if force_edit {
725 force_attempts += 1;
726 if force_attempts > MAX_FORCE_ATTEMPTS {
727 force_edit = false;
728 disable_tools = true;
729 }
730 }
731
732 let mut body = serde_json::Map::new();
733 body.insert("model".into(), json!(model.model_id));
734 body.insert("messages".into(), json!(&messages));
735 body.insert("stream".into(), json!(true));
736 if use_stream_options {
737 body.insert(
738 "stream_options".into(),
739 json!({
740 "include_usage": true,
741 }),
742 );
743 }
744 if let Some(cfg) = tool_loop {
745 body.insert(
746 "tools".into(),
747 tools::openai_tools_json(cfg.allow_edit, cfg.scheds_root.is_some()),
748 );
749 let tool_choice = if force_edit && cfg.allow_edit {
752 json!({"type": "function", "function": {"name": "edit_file"}})
753 } else if disable_tools {
754 json!("none")
755 } else {
756 json!("auto")
757 };
758 body.insert("tool_choice".into(), tool_choice);
759 }
760 let body = Value::Object(body);
761
762 request_n += 1;
763 if verbose {
764 let tc = if tool_loop.is_none() {
765 ""
766 } else if force_edit {
767 " tool_choice=edit_file(forced)"
768 } else if disable_tools {
769 " tool_choice=none"
770 } else {
771 " tool_choice=auto"
772 };
773 eprintln!(
774 "{} {}",
775 color.dim("[api]"),
776 color.cyan(format!(
777 "POST {url} (request #{request_n}, {} messages{tc})",
778 messages.len()
779 ))
780 );
781 }
782
783 let mut spinner = Spinner::start("thinking...", color);
784 let resp = tokio::select! {
785 resp = apply_bearer(client.post(&url), &model.api_key).json(&body).send() => {
786 resp.context("POST chat/completions")?
787 }
788 _ = interrupt::wait(interrupted.clone()) => {
789 spinner.stop();
790 return Err(interrupt::err());
791 }
792 _ = tokio::time::sleep(deadline.remaining()) => {
793 spinner.stop();
794 return Err(deadline.timeout().into());
795 }
796 };
797 let status = resp.status();
798 if !status.is_success() {
799 spinner.stop();
800 let text = tokio::select! {
801 text = resp.text() => text.context("read chat/completions body")?,
802 _ = interrupt::wait(interrupted.clone()) => return Err(interrupt::err()),
803 _ = tokio::time::sleep(deadline.remaining()) => {
804 return Err(deadline.timeout().into());
805 }
806 };
807 if use_cache {
810 use_cache = false;
811 messages[0] = system_message(system, false);
812 if verbose {
813 eprintln!(
814 "{} {}",
815 color.dim("[api]"),
816 color.yellow(format!(
817 "request rejected ({status}); retrying without prompt-cache markers"
818 ))
819 );
820 }
821 continue;
822 }
823 if use_stream_options
824 && (text.contains("stream_options") || text.contains("include_usage"))
825 {
826 use_stream_options = false;
827 if verbose {
828 eprintln!(
829 "{} {}",
830 color.dim("[api]"),
831 color.yellow("stream usage option rejected; retrying without it")
832 );
833 }
834 continue;
835 }
836 return Err(ApiStatusError { status, body: text }.into());
837 }
838
839 let streamed = tokio::select! {
840 streamed = read_streamed_chat(resp, stream_stdout, &mut spinner) => streamed?,
841 _ = interrupt::wait(interrupted.clone()) => {
842 spinner.stop();
843 return Err(interrupt::err());
844 }
845 _ = tokio::time::sleep(deadline.remaining()) => {
846 spinner.stop();
847 return Err(deadline.timeout().into());
848 }
849 };
850 spinner.stop();
851 let raw_events = streamed.raw_events.clone();
852 let streamed_usage = streamed.usage;
853 let message = streamed.into_message();
854
855 if verbose {
856 eprintln!(
857 "{} {}\n{}",
858 color.dim("[api]"),
859 color.blue(format!(
860 "raw streamed response ({} bytes):",
861 raw_events.len()
862 )),
863 preview(&raw_events, 20000)
864 );
865 }
866
867 if verbose {
868 let content = message_content_to_string(&message);
869 if !content.trim().is_empty() {
870 eprintln!(
871 "{} {}\n{}",
872 color.dim("[api]"),
873 color.blue("assistant:"),
874 preview(&content, DUMP_MAX_CHARS)
875 );
876 }
877 }
878 if let Some(u) = streamed_usage {
879 usage.add(&u);
880 if verbose {
881 eprintln!(
882 "{} {}",
883 color.dim("[api]"),
884 color.dim(format!("usage: {}", u.footer_line()))
885 );
886 }
887 }
888
889 let tool_calls: Option<Vec<Value>> = message
890 .get("tool_calls")
891 .and_then(|tc| tc.as_array())
892 .filter(|a| !a.is_empty())
893 .cloned()
894 .map(normalize_tool_calls);
895
896 let (tool_calls, assistant_message) = if tool_calls.is_none() && tool_loop.is_some() {
901 let (parsed, stripped) = parse_text_tool_calls(&message_content_to_string(&message));
902 if parsed.is_empty() {
903 (tool_calls, message)
904 } else {
905 if verbose {
906 eprintln!(
907 "{} {}",
908 color.dim("[api]"),
909 color.yellow(format!(
910 "{} tool call(s) parsed from <tool_call> text",
911 parsed.len()
912 ))
913 );
914 }
915 let assistant = json!({
916 "role": "assistant",
917 "content": if stripped.is_empty() { Value::Null } else { Value::String(stripped) },
918 "tool_calls": sanitize_tool_calls_for_history(&parsed),
919 });
920 (Some(parsed), assistant)
921 }
922 } else {
923 let assistant = match &tool_calls {
924 Some(calls) => {
925 let mut assistant = message;
926 if let Some(obj) = assistant.as_object_mut() {
927 obj.insert(
928 "tool_calls".to_string(),
929 Value::Array(sanitize_tool_calls_for_history(calls)),
930 );
931 }
932 assistant
933 }
934 None => message,
935 };
936 (tool_calls, assistant)
937 };
938
939 if let (Some(arr), Some(cfg)) = (tool_calls, tool_loop) {
940 messages.push(assistant_message);
941
942 if disable_tools {
945 for tc in &arr {
946 if let Some(id) = tc.get("id").and_then(|x| x.as_str()) {
947 messages.push(json!({
948 "role": "tool",
949 "tool_call_id": id,
950 "content": "Tools are disabled now; reply with your one-line summary of the change you made."
951 }));
952 }
953 }
954 continue;
955 }
956
957 if tool_iterations >= cfg.max_iterations && !force_edit {
961 let nudge = if cfg.allow_edit && edits_applied == 0 {
962 force_edit = true;
963 "Exploration budget reached. You MUST make at least one edit_file change now."
964 } else {
965 disable_tools = true;
966 "Exploration budget reached; reply with your one-line summary of the change you made."
967 };
968 for tc in &arr {
969 if let Some(id) = tc.get("id").and_then(|x| x.as_str()) {
970 messages
971 .push(json!({"role": "tool", "tool_call_id": id, "content": nudge}));
972 }
973 }
974 continue;
975 }
976 tool_iterations += 1;
977
978 for tc in &arr {
979 let id = tc
980 .get("id")
981 .and_then(|x| x.as_str())
982 .ok_or_else(|| anyhow!("tool_calls[].id missing"))?;
983 let func = tc
984 .get("function")
985 .ok_or_else(|| anyhow!("tool_calls[].function missing"))?;
986 let name = func
987 .get("name")
988 .and_then(|x| x.as_str())
989 .ok_or_else(|| anyhow!("tool_calls[].function.name missing"))?;
990 let args = func
991 .get("arguments")
992 .map(arguments_to_string)
993 .unwrap_or_else(|| "{}".to_string());
994
995 eprintln!(
998 "{} {} ({})",
999 color.dim(" ->"),
1000 color.cyan(tool_call_label(name, &args)),
1001 color.dim(usage.footer_line())
1002 );
1003
1004 if !tools::is_write_tool(name)
1010 && !seen_tool_calls.insert(format!("{name}\u{0}{args}"))
1011 {
1012 eprintln!(
1013 "{} {} {}",
1014 color.dim(" <-"),
1015 color.cyan(tool_call_label(name, &args)),
1016 color.yellow("skipped: identical call already made this turn")
1017 );
1018 messages.push(json!({
1019 "role": "tool",
1020 "tool_call_id": id,
1021 "content": "Identical call already executed earlier in this turn; \
1022 its result has not changed and is shown above. Do not \
1023 repeat it - use that result, read a different location, \
1024 or make your edit_file change now.",
1025 }));
1026 continue;
1027 }
1028
1029 let tool_started = Instant::now();
1030 let out = match tokio::select! {
1031 out = tools::execute_tool_async(
1032 &cfg.sandbox,
1033 cfg.scheds_root.as_deref(),
1034 name,
1035 &args,
1036 cfg.allow_edit,
1037 ) => out,
1038 _ = interrupt::wait(interrupted.clone()) => return Err(interrupt::err()),
1039 _ = tokio::time::sleep(deadline.remaining()) => {
1040 return Err(deadline.timeout().into());
1041 }
1042 } {
1043 Ok(out) => {
1044 let elapsed = tool_started.elapsed();
1045 if verbose || elapsed >= SLOW_TOOL_LOG_AFTER {
1046 eprintln!(
1047 "{} {} done in {:.1}s ({} bytes)",
1048 color.dim(" <-"),
1049 color.cyan(tool_call_label(name, &args)),
1050 elapsed.as_secs_f64(),
1051 out.len()
1052 );
1053 }
1054 if tools::is_write_tool(name) {
1055 edits_applied += 1;
1056 }
1057 out
1058 }
1059 Err(e) => {
1064 let msg = format!("ERROR: {e:#}");
1067 let shown = format!("{e:#}");
1070 eprintln!(
1071 "{} {} failed after {:.1}s: {}",
1072 color.dim(" <-"),
1073 color.cyan(tool_call_label(name, &args)),
1074 tool_started.elapsed().as_secs_f64(),
1075 color.yellow(preview(shown.lines().next().unwrap_or(&shown), 300))
1076 );
1077 if verbose {
1078 eprintln!(
1079 " {}",
1080 color.yellow(preview(shown.lines().next().unwrap_or(&shown), 300))
1081 );
1082 }
1083 msg
1084 }
1085 };
1086 if verbose {
1087 eprintln!(
1088 "{} {}\n{}",
1089 color.dim("[api]"),
1090 color.blue(format!("tool result [{name}]:")),
1091 preview(&out, 2000)
1092 );
1093 }
1094 messages.push(json!({
1095 "role": "tool",
1096 "tool_call_id": id,
1097 "content": out,
1098 }));
1099 }
1100 if force_edit && edits_applied > 0 {
1105 force_edit = false;
1106 }
1107 continue;
1108 }
1109
1110 if let Some(cfg) = tool_loop
1114 && cfg.allow_edit
1115 && edits_applied == 0
1116 && !disable_tools
1117 {
1118 force_edit = true;
1119 messages.push(assistant_message);
1120 messages.push(json!({
1121 "role": "user",
1122 "content": "You have not made any edit yet. Make your edit_file change(s) now - actually call the tool, do not just describe the change."
1123 }));
1124 continue;
1125 }
1126
1127 let content = message_content_to_string(&assistant_message);
1128 return Ok(content);
1129 }
1130}
1131
1132pub fn tool_loop(
1134 sandbox: &Path,
1135 scheds_root: Option<&Path>,
1136 allow_edit: bool,
1137 max_iterations: u32,
1138) -> ToolLoopConfig {
1139 ToolLoopConfig {
1140 sandbox: sandbox.to_path_buf(),
1141 scheds_root: scheds_root.map(Path::to_path_buf),
1142 allow_edit,
1143 max_iterations,
1144 }
1145}
1146
1147#[cfg(test)]
1148mod tests {
1149 use super::*;
1150
1151 #[test]
1152 fn system_message_cache_markers() {
1153 let cached = system_message("SYS", true);
1154 assert_eq!(cached["content"][0]["text"], "SYS");
1155 assert_eq!(cached["content"][0]["cache_control"]["type"], "ephemeral");
1156
1157 let plain = system_message("SYS", false);
1158 assert_eq!(plain["content"], "SYS");
1159 assert!(plain["content"].is_string());
1160 }
1161
1162 #[test]
1163 fn api_status_errors_are_classified() {
1164 let err: anyhow::Error = ApiStatusError {
1165 status: reqwest::StatusCode::BAD_REQUEST,
1166 body: "Already borrowed".to_string(),
1167 }
1168 .into();
1169
1170 assert!(is_api_error(&err));
1171 assert!(error_summary(&err).contains("Already borrowed"));
1172 }
1173
1174 #[test]
1175 fn tool_call_label_includes_scheduler_for_scheduler_file_reads() {
1176 let args = json!({
1177 "scheduler": "scx_rusty",
1178 "path": "src/bpf/main.bpf.c",
1179 "start_line": 10,
1180 "end_line": 40,
1181 })
1182 .to_string();
1183 assert_eq!(
1184 tool_call_label("read_scheduler_file", &args),
1185 "read_scheduler_file(scx_rusty/src/bpf/main.bpf.c:10-40)"
1186 );
1187 }
1188
1189 #[test]
1190 fn tool_call_label_includes_range_for_file_reads() {
1191 let bounded = json!({
1192 "path": "src/main.rs",
1193 "start_line": 20,
1194 "end_line": 80,
1195 })
1196 .to_string();
1197 assert_eq!(
1198 tool_call_label("read_file", &bounded),
1199 "read_file(src/main.rs:20-80)"
1200 );
1201
1202 let unbounded = json!({
1203 "path": "src/main.rs",
1204 })
1205 .to_string();
1206 assert_eq!(
1207 tool_call_label("read_file", &unbounded),
1208 "read_file(src/main.rs:1-EOF)"
1209 );
1210 }
1211
1212 #[test]
1213 fn tool_call_label_includes_scheduler_for_scheduler_grep() {
1214 let args = json!({
1215 "scheduler": "scx_lavd",
1216 "pattern": "vtime",
1217 })
1218 .to_string();
1219 assert_eq!(
1220 tool_call_label("grep_schedulers", &args),
1221 "grep_schedulers(scx_lavd:vtime)"
1222 );
1223 }
1224
1225 #[test]
1226 fn parses_single_text_tool_call() {
1227 let content = "I'll lower the slice.\n<tool_call>\n{\"name\": \"edit_file\", \"arguments\": {\"path\": \"src/bpf/main.bpf.c\", \"old_string\": \"a\", \"new_string\": \"b\"}}\n</tool_call>";
1228 let (calls, stripped) = parse_text_tool_calls(content);
1229 assert_eq!(calls.len(), 1);
1230 assert_eq!(calls[0]["function"]["name"], "edit_file");
1231 let args: Value =
1232 serde_json::from_str(calls[0]["function"]["arguments"].as_str().unwrap()).unwrap();
1233 assert_eq!(args["path"], "src/bpf/main.bpf.c");
1234 assert_eq!(calls[0]["type"], "function");
1235 assert!(calls[0]["id"].as_str().unwrap().starts_with("text_call_"));
1236 assert_eq!(stripped, "I'll lower the slice.");
1237 }
1238
1239 #[test]
1240 fn parses_multiple_calls_and_unique_ids() {
1241 let content = "<tool_call>{\"name\":\"grep\",\"arguments\":{\"pattern\":\"x\"}}</tool_call><tool_call>{\"name\":\"read_file\",\"arguments\":{\"path\":\"y\"}}</tool_call>";
1242 let (calls, stripped) = parse_text_tool_calls(content);
1243 assert_eq!(calls.len(), 2);
1244 assert_eq!(calls[0]["function"]["name"], "grep");
1245 assert_eq!(calls[1]["function"]["name"], "read_file");
1246 assert!(calls[0]["function"]["arguments"].is_string());
1247 assert!(calls[1]["function"]["arguments"].is_string());
1248 assert_ne!(calls[0]["id"], calls[1]["id"]);
1249 assert_eq!(stripped, "");
1250 }
1251
1252 #[test]
1253 fn accepts_parameters_key_and_missing_close_tag() {
1254 let content = "<tool_call>{\"name\":\"list_dir\",\"parameters\":{\"path\":\".\"}}";
1256 let (calls, _) = parse_text_tool_calls(content);
1257 assert_eq!(calls.len(), 1);
1258 assert_eq!(calls[0]["function"]["name"], "list_dir");
1259 let args: Value =
1260 serde_json::from_str(calls[0]["function"]["arguments"].as_str().unwrap()).unwrap();
1261 assert_eq!(args["path"], ".");
1262 }
1263
1264 #[test]
1265 fn normalizes_structured_tool_call_arguments_for_replay() {
1266 let calls = normalize_tool_calls(vec![json!({
1267 "id": "call_0",
1268 "type": "function",
1269 "function": {
1270 "name": "list_dir",
1271 "arguments": {"path": "."},
1272 },
1273 })]);
1274
1275 assert_eq!(calls.len(), 1);
1276 assert!(calls[0]["function"]["arguments"].is_string());
1277 let args: Value =
1278 serde_json::from_str(calls[0]["function"]["arguments"].as_str().unwrap()).unwrap();
1279 assert_eq!(args["path"], ".");
1280 }
1281
1282 #[test]
1283 fn normalizes_missing_tool_call_arguments_to_empty_object_string() {
1284 let calls = normalize_tool_calls(vec![json!({
1285 "id": "call_0",
1286 "type": "function",
1287 "function": {"name": "list_schedulers"},
1288 })]);
1289
1290 assert_eq!(calls[0]["function"]["arguments"], "{}");
1291 }
1292
1293 #[test]
1294 fn sanitize_replaces_malformed_tool_call_arguments_but_keeps_valid_ones() {
1295 let calls = vec![
1299 json!({
1300 "id": "bad",
1301 "type": "function",
1302 "function": {
1303 "name": "edit_file",
1304 "arguments": "{\"path\": \"src/main.rs\", \"old_string\": \"void f(",
1305 },
1306 }),
1307 json!({
1308 "id": "good",
1309 "type": "function",
1310 "function": {
1311 "name": "read_file",
1312 "arguments": "{\"path\": \"src/main.rs\"}",
1313 },
1314 }),
1315 ];
1316
1317 let sanitized = sanitize_tool_calls_for_history(&calls);
1318 assert_eq!(sanitized[0]["function"]["arguments"], "{}");
1319 assert_ne!(calls[0]["function"]["arguments"], "{}");
1321 assert_eq!(
1323 sanitized[1]["function"]["arguments"],
1324 "{\"path\": \"src/main.rs\"}"
1325 );
1326 }
1327
1328 #[test]
1329 fn assembles_streamed_tool_call_deltas() {
1330 let mut streamed = StreamedChat::default();
1331
1332 let visible = apply_stream_delta(
1333 &json!({
1334 "content": "I will inspect it.",
1335 "tool_calls": [{
1336 "index": 0,
1337 "id": "call_1",
1338 "type": "function",
1339 "function": {"name": "read_", "arguments": "{\"path\":\"src"}
1340 }]
1341 }),
1342 &mut streamed,
1343 )
1344 .unwrap();
1345 assert_eq!(visible.as_deref(), Some("I will inspect it."));
1346
1347 apply_stream_delta(
1348 &json!({
1349 "tool_calls": [{
1350 "index": 0,
1351 "function": {"name": "file", "arguments": "/bpf/main.bpf.c\"}"}
1352 }]
1353 }),
1354 &mut streamed,
1355 )
1356 .unwrap();
1357
1358 let message = streamed.into_message();
1359 assert_eq!(message["content"], "I will inspect it.");
1360 assert_eq!(message["tool_calls"][0]["id"], "call_1");
1361 assert_eq!(message["tool_calls"][0]["function"]["name"], "read_file");
1362 let args: Value = serde_json::from_str(
1363 message["tool_calls"][0]["function"]["arguments"]
1364 .as_str()
1365 .unwrap(),
1366 )
1367 .unwrap();
1368 assert_eq!(args["path"], "src/bpf/main.bpf.c");
1369 }
1370
1371 #[test]
1372 fn stream_printer_handles_multibyte_visible_prefix() {
1373 let mut printer = StreamTextPrinter::new(true);
1374 printer.push("或许是 **").unwrap();
1375 printer.finish().unwrap();
1376 }
1377
1378 #[test]
1379 fn stream_printer_handles_multibyte_hidden_tool_call_buffer() {
1380 let mut printer = StreamTextPrinter::new(true);
1381 printer.push("<tool_call>或许是 **").unwrap();
1382 printer.finish().unwrap();
1383 }
1384
1385 #[test]
1386 fn stream_printer_suppresses_whitespace_before_text_tool_call() {
1387 let mut printer = StreamTextPrinter::new(true);
1388 printer
1389 .push("\n\n<tool_call>{\"name\":\"list_dir\"}")
1390 .unwrap();
1391 printer.finish().unwrap();
1392
1393 assert!(!printer.started);
1394 }
1395
1396 #[test]
1397 fn stream_printer_suppresses_whitespace_only_structured_tool_call_text() {
1398 let mut printer = StreamTextPrinter::new(true);
1399 printer.push("\n\n").unwrap();
1400 printer.note_tool_call();
1401 printer.finish().unwrap();
1402
1403 assert!(!printer.started);
1404 }
1405
1406 #[test]
1407 fn stream_printer_trims_blank_lines_before_tool_call_after_text() {
1408 let mut printer = StreamTextPrinter::new(true);
1409 printer
1410 .push("Let me inspect:\n\n<tool_call>{\"name\":\"list_dir\"}")
1411 .unwrap();
1412 printer.finish().unwrap();
1413
1414 assert!(printer.started);
1415 assert!(!printer.last_was_newline);
1416 }
1417
1418 #[test]
1419 fn floor_char_boundary_does_not_split_utf8() {
1420 let text = "或许是 **";
1421 assert_eq!(floor_char_boundary(text, 0), 0);
1422 assert_eq!(floor_char_boundary(text, 1), 0);
1423 assert_eq!(floor_char_boundary(text, 2), 0);
1424 assert_eq!(floor_char_boundary(text, 3), 3);
1425 }
1426
1427 #[test]
1428 fn parses_sse_data_events() {
1429 let event = b"event: ignored\ndata: {\"a\":1}\n\n";
1430 assert_eq!(sse_data(event).unwrap().as_deref(), Some("{\"a\":1}"));
1431 let (pos, len) = find_sse_delimiter(event).unwrap();
1432 assert_eq!(&event[pos..pos + len], b"\n\n");
1433 }
1434
1435 #[test]
1436 fn no_tool_call_text_yields_no_calls() {
1437 let (calls, stripped) = parse_text_tool_calls("just a plain summary line");
1438 assert!(calls.is_empty());
1439 assert_eq!(stripped, "just a plain summary line");
1440 }
1441
1442 #[test]
1443 fn skips_malformed_block() {
1444 let (calls, _) = parse_text_tool_calls("<tool_call>not json</tool_call>");
1445 assert!(calls.is_empty());
1446 }
1447}