max-tokens responses containing tool calls can permanently break pi-ai session replay

[youkongling]

Summary

A successful stream that terminates with max-tokens and contains a tool call can persist an assistant message whose durable content no longer matches its adapter replay state. On the next request, pi-ai history reconstruction rejects that message with:

invalid pi-ai replay state: block count does not match assistant content

The request fails during local history reconstruction, before it reaches the provider. Once such a message is stored, every subsequent “continue” for that session fails the same way.

Impact

• The unsafe/truncated tool call is correctly not executed.
• Safe assistant content is retained.
• However, the original replay state still describes the removed tool-call block.
• The persisted assistant content and replayState.blocks therefore have different lengths.
• Existing sessions containing this deterministic mismatch become permanently unreplayable.

I reproduced this with a real session using a pi-ai-backed provider. The truncated block was a large tool call; the durable assistant content retained reasoning/text blocks while replay state retained reasoning/text/tool-call blocks.

Root cause

BlockAssembler.blocks() intentionally removes tool calls on max-token truncation:

blocks(): ContentBlock[] {
  const blocks = this.order.map(index => this.assemble(this.mustGet(index), index))
  return this.finish.kind === 'max-tokens'
    ? blocks.filter(block => block.type !== 'tool-call')
    : blocks
}

But its replay-state accessor still returns the original terminal state:

get replayState(): unknown {
  return this._replayState
}

The agent loop then persists:

content: assembler.blocks(),
source: {
  provider: request.provider,
  model: request.model,
  ...assembler.replayState !== undefined
    ? { replayState: assembler.replayState }
    : {},
}

For example, the stored message can contain [reasoning, text] while replay state contains [reasoning, text, tool-call].

On resume, pi-ai replay validation correctly enforces:

if (state.blocks.length !== message.content.length) {
  return invalidReplay('block count does not match assistant content')
}

That invariant is right; the writer violated it.

Proposed fix

Patch branch: https://github.com/youkongling/deepseek-harness/tree/fix/max-tokens-replay-state
Patch commit: youkongling@f6bbb24
Raw patch: https://github.com/youkongling/deepseek-harness/commit/f6bbb24a439c10c1901566b8f8b96ce08c50da36.patch

1. Keep new writes aligned

When max-token safety filtering removes at least one tool call, BlockAssembler.replayState now returns undefined. The persisted assistant content and persisted replay metadata therefore remain consistent.

Max-token responses without tool calls retain replay state, as do normally completed responses.

2. Recover the known historical corruption without weakening validation

Sessions written by the old behavior may already contain the mismatch. The pi-ai conversion layer now recognizes only the exact legacy shape:

• replay stopReason is length;
• replay has more blocks than durable content;
• durable content contains no tool calls;
• after removing replay tool-call blocks, the remaining replay block sequence aligns with durable content.

Such a message is translated as provider-neutral history. Other malformed states, provider/model mismatches, unsupported versions, and unrelated block mismatches continue to fail explicitly with INVALID_REPLAY_STATE.

This preserves fail-loud replay validation while unbricking sessions corrupted by this specific writer bug.

Tests added

• BlockAssembler drops replay state when max-token filtering removes a tool call.
• BlockAssembler retains max-token replay state when no tool call was filtered.
• Agent-loop max-token regression verifies that tool calls are not executed and the logged assistant source has no replay state.
• pi-ai replay recovers the legacy length + dropped-tool-call corruption as provider-neutral history.
• pi-ai replay keeps valid max-token replay state when no tool call was dropped.
• pi-ai replay still rejects unrelated block-count mismatches.

Validation on top of current master (47f943859bef60e4160492346772ded9b24f765a):

vitest run packages/llm/llm packages/core/agent-loop packages/llm/llm-pi-ai
49 test files, 955 tests passed

tsc -b tsconfig.host.json
oxlint <changed source/test files>
verify-translation-pairing packages/llm/llm-pi-ai/README.md
git diff --check

The same vulnerable logic is present in current master and in the published @deepseek-ai/dsh-llm / @deepseek-ai/dsh-llm-pi-ai runtime packages I inspected.

中文说明

这是一个 DSH 核心写入路径的持久化一致性 bug。

当模型输出达到 max-tokens 且包含 tool-call 时，DSH 会安全过滤掉可能被截断的 tool-call，避免执行不完整工具调用；但 BlockAssembler 仍会把描述原始 tool-call 的 replayState 返回给 agent-loop，并随 assistant message 一起持久化。

结果是持久化后的 assistant content 与 replayState.blocks 数量不一致。例如：

assistant content: reasoning + text
replay blocks:    reasoning + text + tool-call

下次继续该会话时，pi-ai 历史重建会校验失败：

invalid pi-ai replay state: block count does not match assistant content

这个失败发生在本地历史重建阶段，请求不会真正到达模型，因此该会话后续“继续”会反复失败，表现为永久卡死。

修复思路：

1. 新写入路径：max-tokens 安全过滤移除 tool-call 时，同步丢弃 replayState，保证持久化内容与回放元数据一致。
2. 旧损坏会话：只窄化恢复确定性损坏特征——stopReason = length、replay blocks 多于 durable content、durable content 无 tool-call、去掉 replay 中的 tool-call 后两边完全对齐；此时转换为 provider-neutral 历史继续使用。
3. 其他 replay 版本错误、格式错误、provider/model 不匹配、无关 block mismatch 仍然保持 fail-loud，继续抛 INVALID_REPLAY_STATE。

修复分支：

https://github.com/youkongling/deepseek-harness/tree/fix/max-tokens-replay-state

修复提交：

youkongling@f6bbb24