Preserve exact tool-call text during prompt replay

[howard0su]

The OpenAI-compatible dflash server was parsing assistant tool-call output into structured JSON and then rebuilding those turns through the chat template on later requests. That preserved the semantics of the tool call, but not the exact text the model originally emitted. Small formatting differences in the rebuilt assistant turn can change tokenization, which makes prefix and KV reuse less stable for tool-using conversations.

Fix this by keeping a small Python-side tool-memory store in the server path and using it during prompt reconstruction. The server now remembers the original assistant text for generated tool-call turns, keyed by tool-call IDs. When a later request sends those same tool calls back as structured history, the prompt tokenizer looks up the IDs and re-injects the original assistant text verbatim instead of re-rendering canonicalized tool-call objects. If a lookup is missing or inconsistent, the existing structured rendering path still applies.

This change stays intentionally simple and Python-first. It adds a focused ToolMemory helper plus regression coverage for chat-completions and responses round trips, without introducing a native radix-tree backend before profiling shows the server-side store is a real bottleneck.

[cubic-dev-ai]

2 issues found across 4 files

Prompt for AI agents (unresolved issues)

Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.


<file name="dflash/scripts/server.py">

<violation number="1" location="dflash/scripts/server.py:675">
P1: Startup crash: `prefill_cache_bytes` is undefined here, and `PrefixCache.init_full_cache()` does not accept `budget_bytes`, so enabling prefill will raise before the server starts.</violation>

<violation number="2" location="dflash/scripts/server.py:1255">
P1: Streaming tool-call replay can lose preceding reasoning text, so the remembered raw assistant turn is not exact.</violation>
</file>

_{Reply with feedback, questions, or to request a fix. Tag @cubic-dev-ai to re-run a review.}

[cubic-dev-ai]

1 issue found across 4 files

Prompt for AI agents (unresolved issues)

Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.


<file name="dflash/scripts/server.py">

<violation number="1" location="dflash/scripts/server.py:676">
P2: Process-wide ToolMemory can replay raw assistant text from unrelated requests when tool-call IDs are reused, causing cross-request data mixing/information disclosure.</violation>
</file>

_{Reply with feedback, questions, or to request a fix. Tag @cubic-dev-ai to re-run a review.}

[cubic-dev-ai]

P2: Process-wide ToolMemory can replay raw assistant text from unrelated requests when tool-call IDs are reused, causing cross-request data mixing/information disclosure.

Prompt for AI agents

Check if this issue is valid — if so, understand the root cause and fix it. At dflash/scripts/server.py, line 676:

<comment>Process-wide ToolMemory can replay raw assistant text from unrelated requests when tool-call IDs are reused, causing cross-request data mixing/information disclosure.</comment>

<file context>
@@ -672,6 +673,21 @@ def _resolve_kv_k_type():
     )
     if prefill_cfg is not None and prefill_cache_slots > 0:
         prefix_cache.init_full_cache(prefill_cache_slots)
+    tool_memory = ToolMemory(
+        max_entries=int(os.environ.get("DFLASH_TOOL_MEMORY_MAX_ENTRIES", "50000")),
+        max_bytes=int(os.environ.get("DFLASH_TOOL_MEMORY_MAX_BYTES", str(64 * 1024 * 1024))),
</file context>