Split Rust code into domain modules

2026-07-25 11:22:59 +02:00
parent 7843592400
commit 4292e0d3f9
22 changed files with 7650 additions and 7536 deletions
--- a/src/server/tools.rs
+++ b/src/server/tools.rs
@@ -0,0 +1,733 @@
+use super::*;
+
+enum ToolProjectionState {
+    Seeking,
+    Invokes,
+    Parameters,
+    Value,
+    Done,
+    Failed,
+}
+
+pub(super) enum ToolProjectionEvent {
+    Text(String),
+    Start {
+        index: usize,
+        id: String,
+        name: String,
+    },
+    Arguments {
+        index: usize,
+        fragment: String,
+    },
+    End {
+        index: usize,
+    },
+}
+
+pub(super) struct ToolProjector {
+    pub(super) raw: String,
+    position: usize,
+    text_emitted: usize,
+    state: ToolProjectionState,
+    index: usize,
+    pub(super) ids: Vec<String>,
+    first_parameter: bool,
+    string_parameter: bool,
+    syntax: Option<ToolSyntax>,
+}
+
+impl ToolProjector {
+    pub(super) fn new() -> Self {
+        Self {
+            raw: String::new(),
+            position: 0,
+            text_emitted: 0,
+            state: ToolProjectionState::Seeking,
+            index: 0,
+            ids: Vec::new(),
+            first_parameter: true,
+            string_parameter: false,
+            syntax: None,
+        }
+    }
+
+    pub(super) fn push(
+        &mut self,
+        chunk: &str,
+        final_chunk: bool,
+        prefix: &str,
+    ) -> Vec<ToolProjectionEvent> {
+        self.raw.push_str(chunk);
+        let mut events = Vec::new();
+        loop {
+            match self.state {
+                ToolProjectionState::Seeking => {
+                    if let Some((start, syntax)) = TOOL_SYNTAXES
+                        .iter()
+                        .filter_map(|syntax| {
+                            self.raw
+                                .find(syntax.tool_start)
+                                .map(|start| (start, *syntax))
+                        })
+                        .min_by_key(|(start, _)| *start)
+                    {
+                        if start > self.text_emitted {
+                            let text = self.raw[self.text_emitted..start].trim_end();
+                            if !text.is_empty() {
+                                events.push(ToolProjectionEvent::Text(text.to_owned()));
+                            }
+                        }
+                        self.position = start + syntax.tool_start.len();
+                        self.text_emitted = start;
+                        self.syntax = Some(syntax);
+                        self.state = ToolProjectionState::Invokes;
+                    } else {
+                        let limit = if final_chunk {
+                            self.raw.len()
+                        } else {
+                            TOOL_SYNTAXES
+                                .iter()
+                                .map(|syntax| {
+                                    safe_before_partial_marker(&self.raw, syntax.tool_start)
+                                })
+                                .min()
+                                .unwrap_or(self.raw.len())
+                        };
+                        if limit > self.text_emitted {
+                            let text = &self.raw[self.text_emitted..limit];
+                            if !text.trim().is_empty() {
+                                events.push(ToolProjectionEvent::Text(text.to_owned()));
+                                self.text_emitted = limit;
+                            }
+                        }
+                        break;
+                    }
+                }
+                ToolProjectionState::Invokes => {
+                    let syntax = self.syntax.unwrap();
+                    self.skip_whitespace();
+                    if self.full_at(syntax.tool_end) {
+                        self.position += syntax.tool_end.len();
+                        self.state = ToolProjectionState::Done;
+                        break;
+                    }
+                    if self.partial_at(syntax.tool_end) || self.partial_at(syntax.invoke_start) {
+                        break;
+                    }
+                    if !self.full_at(syntax.invoke_start) {
+                        self.state = ToolProjectionState::Failed;
+                        break;
+                    }
+                    let Some(tag_end) = self.raw[self.position..].find('>') else {
+                        break;
+                    };
+                    let tag_end = self.position + tag_end + 1;
+                    let Some(name) = dsml_attribute(&self.raw[self.position..tag_end], "name")
+                    else {
+                        self.state = ToolProjectionState::Failed;
+                        break;
+                    };
+                    let id = random_tool_id(prefix);
+                    self.ids.push(id.clone());
+                    events.push(ToolProjectionEvent::Start {
+                        index: self.index,
+                        id,
+                        name,
+                    });
+                    events.push(ToolProjectionEvent::Arguments {
+                        index: self.index,
+                        fragment: "{".into(),
+                    });
+                    self.position = tag_end;
+                    self.first_parameter = true;
+                    self.state = ToolProjectionState::Parameters;
+                }
+                ToolProjectionState::Parameters => {
+                    let syntax = self.syntax.unwrap();
+                    self.skip_whitespace();
+                    if self.full_at(syntax.invoke_end) {
+                        events.push(ToolProjectionEvent::Arguments {
+                            index: self.index,
+                            fragment: "}".into(),
+                        });
+                        events.push(ToolProjectionEvent::End { index: self.index });
+                        self.position += syntax.invoke_end.len();
+                        self.index += 1;
+                        self.state = ToolProjectionState::Invokes;
+                        continue;
+                    }
+                    if self.partial_at(syntax.invoke_end) || self.partial_at(syntax.parameter_start)
+                    {
+                        break;
+                    }
+                    if !self.full_at(syntax.parameter_start) {
+                        self.state = ToolProjectionState::Failed;
+                        break;
+                    }
+                    let Some(tag_end) = self.raw[self.position..].find('>') else {
+                        break;
+                    };
+                    let tag_end = self.position + tag_end + 1;
+                    let tag = &self.raw[self.position..tag_end];
+                    let Some(name) = dsml_attribute(tag, "name") else {
+                        self.state = ToolProjectionState::Failed;
+                        break;
+                    };
+                    self.string_parameter =
+                        dsml_attribute(tag, "string").as_deref() != Some("false");
+                    let mut fragment = if self.first_parameter {
+                        String::new()
+                    } else {
+                        ",".into()
+                    };
+                    self.first_parameter = false;
+                    fragment
+                        .push_str(&serde_json::to_string(&name).unwrap_or_else(|_| "\"\"".into()));
+                    fragment.push(':');
+                    if self.string_parameter {
+                        fragment.push('"');
+                    }
+                    events.push(ToolProjectionEvent::Arguments {
+                        index: self.index,
+                        fragment,
+                    });
+                    self.position = tag_end;
+                    self.state = ToolProjectionState::Value;
+                }
+                ToolProjectionState::Value => {
+                    let syntax = self.syntax.unwrap();
+                    if let Some(relative_end) = self.raw[self.position..].find(syntax.parameter_end)
+                    {
+                        let end = self.position + relative_end;
+                        self.emit_value(end, &mut events);
+                        if self.string_parameter {
+                            events.push(ToolProjectionEvent::Arguments {
+                                index: self.index,
+                                fragment: "\"".into(),
+                            });
+                        }
+                        self.position = end + syntax.parameter_end.len();
+                        self.state = ToolProjectionState::Parameters;
+                        continue;
+                    }
+                    let limit = safe_parameter_value_limit(
+                        &self.raw,
+                        self.position,
+                        syntax.parameter_end,
+                        self.string_parameter,
+                    );
+                    self.emit_value(limit, &mut events);
+                    break;
+                }
+                ToolProjectionState::Done | ToolProjectionState::Failed => break,
+            }
+        }
+        events
+    }
+
+    fn emit_value(&mut self, end: usize, events: &mut Vec<ToolProjectionEvent>) {
+        if end <= self.position {
+            return;
+        }
+        let raw = &self.raw[self.position..end];
+        let fragment = if self.string_parameter {
+            let value = unescape_dsml(raw);
+            let encoded = serde_json::to_string(&value).unwrap_or_else(|_| "\"\"".into());
+            encoded[1..encoded.len() - 1].to_owned()
+        } else {
+            raw.to_owned()
+        };
+        events.push(ToolProjectionEvent::Arguments {
+            index: self.index,
+            fragment,
+        });
+        self.position = end;
+    }
+
+    fn skip_whitespace(&mut self) {
+        while self.raw[self.position..]
+            .chars()
+            .next()
+            .is_some_and(char::is_whitespace)
+        {
+            self.position += self.raw[self.position..].chars().next().unwrap().len_utf8();
+        }
+    }
+
+    fn full_at(&self, marker: &str) -> bool {
+        self.raw.as_bytes()[self.position..].starts_with(marker.as_bytes())
+    }
+
+    fn partial_at(&self, marker: &str) -> bool {
+        let tail = &self.raw.as_bytes()[self.position..];
+        tail.len() < marker.len() && marker.as_bytes().starts_with(tail)
+    }
+}
+
+fn dsml_attribute(tag: &str, name: &str) -> Option<String> {
+    let start = tag.find(&format!("{name}=\""))? + name.len() + 2;
+    let end = start + tag[start..].find('"')?;
+    Some(unescape_dsml(&tag[start..end]))
+}
+
+fn safe_before_partial_marker(text: &str, marker: &str) -> usize {
+    let mut limit = text.len().saturating_sub(marker.len().saturating_sub(1));
+    while !text.is_char_boundary(limit) {
+        limit -= 1;
+    }
+    limit
+}
+
+fn safe_parameter_value_limit(text: &str, start: usize, end_marker: &str, string: bool) -> usize {
+    let bytes = text.as_bytes();
+    let marker = end_marker.as_bytes();
+    let mut limit = bytes.len();
+    for length in (1..marker.len().min(bytes.len().saturating_sub(start) + 1)).rev() {
+        if bytes[start..].ends_with(&marker[..length]) {
+            limit -= length;
+            break;
+        }
+    }
+    if string {
+        for entity in ["&amp;", "&lt;", "&gt;", "&quot;", "&apos;"] {
+            let entity = entity.as_bytes();
+            for length in 1..entity.len() {
+                if bytes[start..limit].ends_with(&entity[..length]) {
+                    limit -= length;
+                    return limit;
+                }
+            }
+        }
+    }
+    limit
+}
+
+pub(super) fn render_messages(
+    state: &State,
+    messages: &[ApiMessage],
+    tools: &[Value],
+    tool_schemas: &[String],
+    tools_enabled: bool,
+    protocol: Protocol,
+) -> Result<(String, Vec<ChatTurn>), (u16, String)> {
+    validate_tool_results(state, messages, protocol)?;
+    let preserve_reasoning = tools_enabled
+        || messages.iter().any(|message| {
+            matches!(message.role.as_str(), "tool" | "function") || !message.tool_calls.is_empty()
+        });
+    let mut system = String::new();
+    if tools_enabled {
+        system.push_str(TOOLS_PROMPT);
+        if tool_schemas.is_empty() {
+            for tool in tools {
+                let schema = tool.get("function").unwrap_or(tool);
+                if !system.ends_with("\n\n") {
+                    system.push('\n');
+                }
+                system.push_str(
+                    &serde_json::to_string(schema)
+                        .map_err(|error| (400, format!("invalid tool schema: {error}")))?,
+                );
+                system.push('\n');
+            }
+        } else {
+            for schema in tool_schemas {
+                if !system.ends_with("\n\n") {
+                    system.push('\n');
+                }
+                system.push_str(schema);
+                system.push('\n');
+            }
+        }
+        system.push_str(
+            "\nYou MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls. Use the exact parameter names from the schemas.",
+        );
+    }
+
+    let mut turns = Vec::<ChatTurn>::new();
+    for message in messages {
+        let content = content_text(&message.content);
+        match message.role.as_str() {
+            "system" | "developer" => {
+                if !system.is_empty() {
+                    system.push_str("\n\n");
+                }
+                system.push_str(&content);
+            }
+            "user" => turns.push(ChatTurn {
+                user: true,
+                skip_previous_eos: false,
+                reasoning: None,
+                reasoning_complete: true,
+                content,
+            }),
+            "tool" | "function" => {
+                let wrapped = format!(
+                    "<tool_result>{}</tool_result>",
+                    escape_tool_result(&content)
+                );
+                if let Some(previous) = turns.last_mut()
+                    && previous.user
+                    && previous.content.starts_with("<tool_result>")
+                {
+                    previous.content.push_str(&wrapped);
+                } else {
+                    turns.push(ChatTurn {
+                        user: true,
+                        skip_previous_eos: protocol == Protocol::Responses,
+                        reasoning: None,
+                        reasoning_complete: true,
+                        content: wrapped,
+                    });
+                }
+            }
+            "assistant" => {
+                let mut content = content;
+                if !message.tool_calls.is_empty() {
+                    content.push_str(&replayed_or_canonical_tools(state, &message.tool_calls));
+                }
+                let reasoning = content_text(&message.reasoning_content);
+                turns.push(ChatTurn {
+                    user: false,
+                    skip_previous_eos: false,
+                    reasoning: (preserve_reasoning && !reasoning.is_empty()).then_some(reasoning),
+                    reasoning_complete: true,
+                    content,
+                });
+            }
+            role => return Err((400, format!("unsupported message role: {role}"))),
+        }
+    }
+    Ok((system, turns))
+}
+
+pub(super) fn validate_tool_results(
+    state: &State,
+    messages: &[ApiMessage],
+    protocol: Protocol,
+) -> Result<(), (u16, String)> {
+    if !matches!(protocol, Protocol::Anthropic | Protocol::Responses) {
+        return Ok(());
+    }
+    let memory = state.tool_memory.lock().ok();
+    for (index, message) in messages.iter().enumerate() {
+        if !matches!(message.role.as_str(), "tool" | "function") || message.tool_call_id.is_empty()
+        {
+            continue;
+        }
+        let id = &message.tool_call_id;
+        let live = memory
+            .as_ref()
+            .is_some_and(|memory| memory.contains_key(id));
+        let replayed = messages[..index].iter().any(|message| {
+            message.role == "assistant" && message.tool_calls.iter().any(|call| call.id == *id)
+        });
+        if live || replayed {
+            continue;
+        }
+        let message = match protocol {
+            Protocol::Anthropic => format!(
+                "Anthropic continuation state is not available for tool_use_id {id}; retry by replaying the full messages history"
+            ),
+            Protocol::Responses => format!(
+                "Responses continuation state is not available for call_id {id}; retry by replaying the full input history"
+            ),
+            _ => unreachable!(),
+        };
+        return Err((400, message));
+    }
+    Ok(())
+}
+
+fn replayed_or_canonical_tools(state: &State, calls: &[ApiToolCall]) -> String {
+    if let Ok(memory) = state.tool_memory.lock()
+        && let Some(raw) = calls.iter().find_map(|call| {
+            (!call.id.is_empty())
+                .then(|| memory.get(&call.id))
+                .flatten()
+        })
+    {
+        return raw.clone();
+    }
+    canonical_tools(calls)
+}
+
+pub(super) fn canonical_tools(calls: &[ApiToolCall]) -> String {
+    let mut output = String::from("\n\n<｜DSML｜tool_calls>\n");
+    for call in calls {
+        output.push_str("<｜DSML｜invoke name=\"");
+        output.push_str(&escape_attribute(&call.function.name));
+        output.push_str("\">\n");
+        match serde_json::from_str::<Value>(&call.function.arguments) {
+            Ok(Value::Object(arguments)) => {
+                for (name, value) in arguments {
+                    output.push_str("<｜DSML｜parameter name=\"");
+                    output.push_str(&escape_attribute(&name));
+                    let string = value.as_str();
+                    output.push_str(if string.is_some() {
+                        "\" string=\"true\">"
+                    } else {
+                        "\" string=\"false\">"
+                    });
+                    if let Some(value) = string {
+                        output.push_str(&escape_parameter(value));
+                    } else {
+                        output.push_str(&escape_json_parameter(&value.to_string()));
+                    }
+                    output.push_str("</｜DSML｜parameter>\n");
+                }
+            }
+            _ => {
+                output.push_str("<｜DSML｜parameter name=\"arguments\" string=\"true\">");
+                output.push_str(&escape_parameter(&call.function.arguments));
+                output.push_str("</｜DSML｜parameter>\n");
+            }
+        }
+        output.push_str("</｜DSML｜invoke>\n");
+    }
+    output.push_str("</｜DSML｜tool_calls>");
+    output
+}
+
+pub(super) fn parse_generated_tools(
+    state: &State,
+    text: &str,
+    protocol: Protocol,
+) -> (String, Vec<ApiToolCall>) {
+    parse_generated_tools_with_ids(state, text, protocol, &[])
+}
+
+pub(super) fn parse_generated_tools_with_ids(
+    state: &State,
+    text: &str,
+    protocol: Protocol,
+    streamed_ids: &[String],
+) -> (String, Vec<ApiToolCall>) {
+    let Some((start, syntax)) = TOOL_SYNTAXES
+        .iter()
+        .filter_map(|syntax| text.find(syntax.tool_start).map(|start| (start, *syntax)))
+        .min_by_key(|(start, _)| *start)
+    else {
+        return (text.to_owned(), Vec::new());
+    };
+    let Some(relative_end) = text[start..].find(syntax.tool_end) else {
+        return (text.to_owned(), Vec::new());
+    };
+    let end = start + relative_end + syntax.tool_end.len();
+    let content = text[..start].trim_end();
+    let raw = &text[content.len()..end];
+    let mut calls = Vec::new();
+    let mut cursor = raw.find(syntax.tool_start).unwrap() + syntax.tool_start.len();
+    loop {
+        skip_text_whitespace(raw, &mut cursor);
+        if raw[cursor..].starts_with(syntax.tool_end) {
+            break;
+        }
+        if !raw[cursor..].starts_with(syntax.invoke_start) {
+            return (text.to_owned(), Vec::new());
+        }
+        let Some(tag_end) = raw[cursor..].find('>').map(|end| cursor + end + 1) else {
+            return (text.to_owned(), Vec::new());
+        };
+        let Some(name) = dsml_attribute(&raw[cursor..tag_end], "name") else {
+            return (text.to_owned(), Vec::new());
+        };
+        cursor = tag_end;
+        let mut arguments = Map::new();
+        loop {
+            skip_text_whitespace(raw, &mut cursor);
+            if raw[cursor..].starts_with(syntax.invoke_end) {
+                cursor += syntax.invoke_end.len();
+                break;
+            }
+            let Some((name, value)) = parse_tool_parameter(raw, &mut cursor, syntax) else {
+                return (text.to_owned(), Vec::new());
+            };
+            arguments.insert(name, value);
+        }
+        calls.push(ApiToolCall {
+            id: String::new(),
+            function: ApiFunction {
+                name,
+                arguments: Value::Object(arguments).to_string(),
+            },
+        });
+    }
+    if calls.is_empty() {
+        return (text.to_owned(), Vec::new());
+    }
+    let prefix = if protocol == Protocol::Anthropic {
+        "toolu_"
+    } else {
+        "call_"
+    };
+    for (index, call) in calls.iter_mut().enumerate() {
+        call.id = streamed_ids
+            .get(index)
+            .cloned()
+            .unwrap_or_else(|| random_tool_id(prefix));
+    }
+    if let Ok(mut memory) = state.tool_memory.lock() {
+        // ponytail: one process-local replay table; add LRU eviction if 100k live tool ids is measured insufficient.
+        if memory.len() >= 100_000 {
+            memory.clear();
+        }
+        for call in &calls {
+            memory.insert(call.id.clone(), raw.to_owned());
+        }
+    }
+    (content.to_owned(), calls)
+}
+
+#[derive(Clone, Copy)]
+pub(super) struct ToolSyntax {
+    pub(super) tool_start: &'static str,
+    pub(super) tool_end: &'static str,
+    pub(super) invoke_start: &'static str,
+    pub(super) invoke_end: &'static str,
+    pub(super) parameter_start: &'static str,
+    pub(super) parameter_end: &'static str,
+}
+
+pub(super) const TOOL_SYNTAXES: [ToolSyntax; 3] = [
+    ToolSyntax {
+        tool_start: "<｜DSML｜tool_calls>",
+        tool_end: "</｜DSML｜tool_calls>",
+        invoke_start: "<｜DSML｜invoke",
+        invoke_end: "</｜DSML｜invoke>",
+        parameter_start: "<｜DSML｜parameter",
+        parameter_end: "</｜DSML｜parameter>",
+    },
+    ToolSyntax {
+        tool_start: "<DSML｜tool_calls>",
+        tool_end: "</DSML｜tool_calls>",
+        invoke_start: "<DSML｜invoke",
+        invoke_end: "</DSML｜invoke>",
+        parameter_start: "<DSML｜parameter",
+        parameter_end: "</DSML｜parameter>",
+    },
+    ToolSyntax {
+        tool_start: "<tool_calls>",
+        tool_end: "</tool_calls>",
+        invoke_start: "<invoke",
+        invoke_end: "</invoke>",
+        parameter_start: "<parameter",
+        parameter_end: "</parameter>",
+    },
+];
+
+fn parse_tool_parameter(
+    text: &str,
+    cursor: &mut usize,
+    syntax: ToolSyntax,
+) -> Option<(String, Value)> {
+    if !text[*cursor..].starts_with(syntax.parameter_start) {
+        return None;
+    }
+    let tag_end = text[*cursor..].find('>').map(|end| *cursor + end + 1)?;
+    let tag = &text[*cursor..tag_end];
+    let name = dsml_attribute(tag, "name")?;
+    let is_string = dsml_attribute(tag, "string");
+    *cursor = tag_end;
+    let mut nested_start = *cursor;
+    skip_text_whitespace(text, &mut nested_start);
+    if is_string.is_none() && text[nested_start..].starts_with(syntax.parameter_start) {
+        *cursor = nested_start;
+        let mut nested = Map::new();
+        loop {
+            skip_text_whitespace(text, cursor);
+            if !text[*cursor..].starts_with(syntax.parameter_start) {
+                break;
+            }
+            let (name, value) = parse_tool_parameter(text, cursor, syntax)?;
+            nested.insert(name, value);
+        }
+        skip_text_whitespace(text, cursor);
+        if !text[*cursor..].starts_with(syntax.parameter_end) {
+            return None;
+        }
+        *cursor += syntax.parameter_end.len();
+        return Some((name, Value::Object(nested)));
+    }
+    let value_end = text[*cursor..]
+        .find(syntax.parameter_end)
+        .map(|end| *cursor + end)?;
+    let raw = &text[*cursor..value_end];
+    *cursor = value_end + syntax.parameter_end.len();
+    let value = if is_string.as_deref().unwrap_or("true") == "true" {
+        Value::String(unescape_dsml(raw))
+    } else {
+        serde_json::from_str(raw).unwrap_or(Value::Null)
+    };
+    Some((name, value))
+}
+
+fn skip_text_whitespace(text: &str, cursor: &mut usize) {
+    while text[*cursor..]
+        .chars()
+        .next()
+        .is_some_and(char::is_whitespace)
+    {
+        *cursor += text[*cursor..].chars().next().unwrap().len_utf8();
+    }
+}
+
+pub(super) fn tool_calls_json(calls: &[ApiToolCall]) -> Value {
+    Value::Array(
+        calls
+            .iter()
+            .map(|call| {
+                json!({
+                    "id": call.id,
+                    "type": "function",
+                    "function": {
+                        "name": call.function.name,
+                        "arguments": call.function.arguments,
+                    }
+                })
+            })
+            .collect(),
+    )
+}
+
+pub(super) fn content_text(value: &Value) -> String {
+    match value {
+        Value::String(text) => text.clone(),
+        Value::Array(parts) => parts
+            .iter()
+            .filter_map(|part| match part {
+                Value::String(text) => Some(text.as_str()),
+                Value::Object(object) => object.get("text").and_then(Value::as_str),
+                _ => None,
+            })
+            .collect(),
+        _ => String::new(),
+    }
+}
+
+pub(super) fn escape_attribute(text: &str) -> String {
+    text.replace('&', "&amp;")
+        .replace('<', "&lt;")
+        .replace('>', "&gt;")
+        .replace('"', "&quot;")
+}
+
+pub(super) fn escape_parameter(text: &str) -> String {
+    text.replace("</｜DSML｜parameter>", "&lt;/｜DSML｜parameter>")
+}
+
+pub(super) fn escape_json_parameter(text: &str) -> String {
+    text.replace("</｜DSML｜parameter>", "\\u003c/｜DSML｜parameter>")
+}
+
+pub(super) fn escape_tool_result(text: &str) -> String {
+    text.replace("</tool_result>", "&lt;/tool_result>")
+}
+
+pub(super) fn unescape_dsml(text: &str) -> String {
+    text.replace("&quot;", "\"")
+        .replace("&gt;", ">")
+        .replace("&lt;", "<")
+        .replace("&amp;", "&")
+}