DS4Server/src/server/tools.rs

use super::*;
pub(super) use crate::dsml::{SYNTAXES as TOOL_SYNTAXES, Syntax as ToolSyntax};

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,
    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 = if prefix.is_empty() {
                        String::new()
                    } else {
                        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
    }

    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;
        }
        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,
                tool: false,
                system: false,
                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,
                        tool: false,
                        system: false,
                        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,
                    tool: false,
                    system: 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()
        .unwrap_or_else(|poisoned| poisoned.into_inner());
    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.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 {
    let memory = state
        .tool_memory
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner());
    if 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 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()
        .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));
    }
    let mut memory = state
        .tool_memory
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner());
    // 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)
}

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;", "&")
}