Complete local HTTP endpoint parity

src/server/tools.rs

@@ -1,5 +1,9 @@
 use super::*;
 
+const TOOL_MEMORY_FILE: &str = "tool-replay.json";
+const TOOL_MEMORY_MAX_IDS: usize = 100_000;
+const TOOL_MEMORY_MAX_BYTES: u64 = 512 * 1024 * 1024;
+
 enum ToolProjectionState {
     Seeking,
     Invokes,
@@ -226,6 +230,15 @@ impl ToolProjector {
         events
     }
 
+    pub(super) fn finish(&mut self, prefix: &str) -> Vec<ToolProjectionEvent> {
+        if let Some(repaired) = repair_generated_tools(&self.raw) {
+            let suffix = repaired[self.raw.len()..].to_owned();
+            self.push(&suffix, true, prefix)
+        } else {
+            self.push("", true, prefix)
+        }
+    }
+
     fn emit_value(&mut self, end: usize, events: &mut Vec<ToolProjectionEvent>) {
         if end <= self.position {
             return;
@@ -506,6 +519,17 @@ pub(super) fn parse_generated_tools_with_ids(
     text: &str,
     protocol: Protocol,
     streamed_ids: &[String],
+) -> (String, Vec<ApiToolCall>) {
+    let repaired = repair_generated_tools(text);
+    let text = repaired.as_deref().unwrap_or(text);
+    parse_generated_tools_once(state, text, protocol, streamed_ids)
+}
+
+fn parse_generated_tools_once(
+    state: &State,
+    text: &str,
+    protocol: Protocol,
+    streamed_ids: &[String],
 ) -> (String, Vec<ApiToolCall>) {
     let Some((start, syntax)) = TOOL_SYNTAXES
         .iter()
@@ -575,16 +599,95 @@ pub(super) fn parse_generated_tools_with_ids(
         .tool_memory
         .lock()
         .unwrap_or_else(|poisoned| poisoned.into_inner());
-    // ponytail: one process-local replay table; add LRU eviction if 100k live tool ids is measured insufficient.
-    if memory.len() >= 100_000 {
+    // ponytail: clear-at-cap is cheaper than an LRU; add ordered eviction if real clients retain 100k live ids.
+    if memory.len() >= TOOL_MEMORY_MAX_IDS {
         memory.clear();
     }
     for call in &calls {
         memory.insert(call.id.clone(), raw.to_owned());
     }
+    write_tool_memory(&state.cache_path, &memory);
     (content.to_owned(), calls)
 }
 
+pub(super) fn read_tool_memory(directory: &std::path::Path) -> HashMap<String, String> {
+    if directory.as_os_str().is_empty() {
+        return HashMap::new();
+    }
+    let path = directory.join(TOOL_MEMORY_FILE);
+    if std::fs::metadata(&path).is_ok_and(|metadata| metadata.len() > TOOL_MEMORY_MAX_BYTES) {
+        return HashMap::new();
+    }
+    let Ok(file) = File::open(path) else {
+        return HashMap::new();
+    };
+    let Ok(memory) = serde_json::from_reader::<_, HashMap<String, String>>(file) else {
+        return HashMap::new();
+    };
+    if memory.len() <= TOOL_MEMORY_MAX_IDS {
+        memory
+    } else {
+        HashMap::new()
+    }
+}
+
+fn write_tool_memory(directory: &std::path::Path, memory: &HashMap<String, String>) {
+    if directory.as_os_str().is_empty() || std::fs::create_dir_all(directory).is_err() {
+        return;
+    }
+    let path = directory.join(TOOL_MEMORY_FILE);
+    let temporary = path.with_extension("json.tmp");
+    let Ok(mut file) = File::create(&temporary) else {
+        return;
+    };
+    if serde_json::to_writer(&mut file, memory).is_ok()
+        && file.sync_all().is_ok()
+        && file
+            .metadata()
+            .is_ok_and(|metadata| metadata.len() <= TOOL_MEMORY_MAX_BYTES)
+    {
+        let _ = std::fs::rename(&temporary, path);
+    } else {
+        let _ = std::fs::remove_file(temporary);
+    }
+}
+
+fn repair_generated_tools(text: &str) -> Option<String> {
+    let scan_start = text
+        .rfind("</think>")
+        .map_or(0, |position| position + "</think>".len());
+    let scan = &text[scan_start..];
+    let syntax = TOOL_SYNTAXES
+        .iter()
+        .filter_map(|syntax| scan.find(syntax.tool_start).map(|start| (start, *syntax)))
+        .min_by_key(|(start, _)| *start)?
+        .1;
+    let tool_open = scan.matches(syntax.tool_start).count();
+    let tool_close = scan.matches(syntax.tool_end).count();
+    let invoke_open = scan.matches(syntax.invoke_start).count();
+    let invoke_close = scan.matches(syntax.invoke_end).count();
+    let parameter_open = scan.matches(syntax.parameter_start).count();
+    let parameter_close = scan.matches(syntax.parameter_end).count();
+    if (tool_open, invoke_open, parameter_open) == (tool_close, invoke_close, parameter_close)
+        || tool_close > tool_open
+        || invoke_close > invoke_open
+        || parameter_close > parameter_open
+    {
+        return None;
+    }
+    let mut repaired = text.to_owned();
+    for _ in parameter_close..parameter_open {
+        repaired.push_str(syntax.parameter_end);
+    }
+    for _ in invoke_close..invoke_open {
+        repaired.push_str(syntax.invoke_end);
+    }
+    for _ in tool_close..tool_open {
+        repaired.push_str(syntax.tool_end);
+    }
+    Some(repaired)
+}
+
 #[derive(Clone, Copy)]
 pub(super) struct ToolSyntax {
     pub(super) tool_start: &'static str,