How DS2API Prevents Tool Call Anti-Leakage in Streaming Responses

DS2API guarantees tool call anti-leakage in streaming responses by buffering potential XML tool blocks in a single-pass sieve, ensuring that partial tags, fenced code examples, or malformed blocks never appear in the text output sent to clients.

DS2API is a Go-based API gateway, implemented at CJackHwang/ds2api, that streams large language model outputs while maintaining strict separation between conversational text and executable tool instructions. The system solves the critical problem of accidental tool call leakage—where fragments of XML tool blocks might escape into the user-visible text stream—through a stateful buffering mechanism that validates complete blocks before emission.

The Tool-Sieve Streaming Pipeline

DS2API processes incoming model tokens through a tool-sieve located in the internal/toolstream package. This sieve operates in a single pass over arbitrarily chunked data, maintaining state across fragmented network packets.

Chunk Ingestion and Buffering

Every received chunk is appended to state.pending. The sieve then scans this buffer for potential tool call starts using findToolSegmentStart, which searches for opening tags defined in xmlToolTagsToDetect (<tool_calls>, <invoke, and others). Critically, this scan consults insideCodeFenceWithState to skip any tags that appear inside Markdown code fences, preventing documentation examples from being intercepted.

Transition to Capturing Mode

When findToolSegmentStart locates a valid opening tag:

1. Text preceding the tag is emitted immediately as a Content event.
2. The remaining buffer moves to state.capture.
3. The sieve enters capturing mode, isolating subsequent chunks from the text stream.

XML Validation and Event Emission

While capturing, consumeXMLToolCapture attempts to extract a complete, parsable XML block. It uses findXMLCloseOutsideCDATA to locate the matching closing tag while skipping content inside CDATA sections or HTML comments. Once a full block is captured:

• toolcall.ParseToolCalls parses the XML into structured tool calls.
• A ToolCalls event is emitted containing the executable payload.
• Any surrounding text discovered during parsing is emitted as separate Content events.

If the closing tag has not arrived (hasOpenXMLToolTag), the buffer remains in state.capture until more data arrives or the stream terminates.

Defensive Mechanisms Against Leakage

DS2API implements multiple layers of protection to prevent tool call fragments from leaking into the text stream under edge-case conditions.

Partial Tag Hold-Back

When a chunk ends with an incomplete tag (e.g., "Hello <too"), findPartialXMLToolTagStart isolates the fragment and returns it to the hold buffer. The incomplete fragment is not emitted as text until the tag is completed or the stream is finalized, preventing accidental exposure of the opening marker.

Fenced Code Block Protection

The splitSafeContentForToolDetection helper, combined with insideCodeFenceWithState, detects when potential XML tags reside within Markdown fences (triple backticks). If detected, the entire fragment passes through as ordinary text, preserving code examples without triggering the sieve.

CDATA-Aware Closing Tag Search

To handle payloads where CDATA sections contain literal tool-like XML (such as file content carrying <tool_calls>), findXMLCloseOutsideCDATA walks the string while explicitly skipping CDATA boundaries. This ensures the parser matches the true structural closing tag rather than an embedded literal.

Malformed XML Fallback

If consumeXMLToolCapture cannot parse a captured block via toolcall.ParseToolCalls, the system falls back to emitting the entire buffered content as plain text. This prevents silent dropping of content when models generate malformed or incomplete XML.

Server-Side Implementation

Below is a minimal illustration of how a server-side handler streams a response while guaranteeing anti-leakage:

// Assume `state` is a *toolstream.State that lives for the duration of the stream.
func streamResponse(w http.ResponseWriter, modelStream <-chan string, toolNames []string) {
    // Send SSE headers (or any streaming protocol you prefer)
    fmt.Fprint(w, "data: ")

    for chunk := range modelStream {
        // Process each incoming chunk
        evts := toolstream.ProcessChunk(state, chunk, toolNames)

        // Emit events in order
        for _, ev := range evts {
            if ev.Content != "" {
                fmt.Fprintf(w, "%s", ev.Content)          // plain text part
            }
            if len(ev.ToolCalls) > 0 {
                // Serialize tool calls as the DS2API wire format
                payload, _ := json.Marshal(map[string]any{
                    "tool_calls": ev.ToolCalls,
                })
                fmt.Fprintf(w, "%s", payload)              // tool‑call part
            }
        }
        // Flush the writer so the client receives data immediately
        if f, ok := w.(http.Flusher); ok {
            f.Flush()
        }
    }

    // End of stream – make sure any leftover buffered data is emitted
    final := toolstream.Flush(state, toolNames)
    for _, ev := range final {
        if ev.Content != "" {
            fmt.Fprintf(w, "%s", ev.Content)
        }
        if len(ev.ToolCalls) > 0 {
            payload, _ := json.Marshal(map[string]any{
                "tool_calls": ev.ToolCalls,
            })
            fmt.Fprintf(w, "%s", payload)
        }
    }
}

The ProcessChunk and Flush calls guarantee that the Content emitted never contains fragments of a tool call, even if the model streams the XML in many tiny pieces.

Core Source Files

• internal/toolstream/tool_sieve_core.go – Orchestrates chunk processing, buffering, state.pending and state.capture management, and event emission.
• internal/toolstream/tool_sieve_xml.go – Implements XML-specific detection, CDATA-aware closing tag search via findXMLCloseOutsideCDATA, and partial-tag handling with findPartialXMLToolTagStart.
• internal/toolstream/tool_sieve_xml_test.go – Contains exhaustive tests asserting that tool calls never leak under edge cases including long payloads, CDATA, leading prose, fenced examples, malformed XML, and partial tags.
• internal/toolcall/toolcalls.go – Parses complete XML blocks into ParsedToolCall structures (invoked by consumeXMLToolCapture).
• api/chat-stream.js – Public entry point that forwards to the Go implementation via the JavaScript wrapper.

Summary

• DS2API uses a single-pass tool-sieve that processes chunks incrementally without backtracking.
• Partial tag detection via findPartialXMLToolTagStart prevents emission of incomplete XML fragments.
• Markdown code fence awareness ensures documentation examples are never misidentified as executable calls.
• CDATA-aware parsing prevents premature truncation when tool-like strings appear inside escaped content blocks.
• Malformed XML fallback converts invalid tool blocks to plain text rather than dropping them or leaking fragments.
• The toolstream package exposes ProcessChunk and Flush functions that guarantee Content events contain zero tool call leakage.

Frequently Asked Questions

What happens if a streaming chunk ends mid-tag?

When a chunk terminates with an incomplete opening tag (e.g., "<tool_cal"), findPartialXMLToolTagStart identifies the fragment and returns it to the hold buffer. The system withholds this partial content from the text stream until either the tag completes and validates as a tool call, or the stream ends without completion, at which point it is flushed as plain text.

How does DS2API distinguish between executable tool calls and examples in Markdown code blocks?

The sieve consults insideCodeFenceWithState during the findToolSegmentStart scan. If a potential XML tag resides between triple backticks, splitSafeContentForToolDetection routes the entire fenced block through as ordinary text content, bypassing tool detection entirely. This preserves code examples while preventing false positive execution.

Can the sieve handle CDATA sections containing tool-like XML?

Yes. The findXMLCloseOutsideCDATA function explicitly walks the buffer while skipping content inside CDATA and HTML comment boundaries. This ensures that when a tool call wraps file content containing literal <tool_calls> strings, the parser correctly identifies the outer structural closing tag rather than matching the inner literal.

What occurs when the XML is malformed or incomplete?

If consumeXMLToolCapture captures a block but toolcall.ParseToolCalls fails to parse it as valid executable XML, the system emits the entire buffered content—including the ostensible tags—as a standard Content event rather than a ToolCalls event. This graceful degradation prevents both data loss and leakage of unvalidated tool instructions.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}