Skip to content

fix(offload): invalidate token cache after stripping deleted tool_use blocks#234

Open
akhilesharora wants to merge 1 commit into
TencentCloud:mainfrom
akhilesharora:fix/offload-stale-token-cache
Open

fix(offload): invalidate token cache after stripping deleted tool_use blocks#234
akhilesharora wants to merge 1 commit into
TencentCloud:mainfrom
akhilesharora:fix/offload-stale-token-cache

Conversation

@akhilesharora

Copy link
Copy Markdown
Contributor

Description | 描述

context-token-tracker.ts caches per-message token counts in a WeakMap keyed on message identity plus _offloaded. Two offload hooks strip deleted tool_use blocks out of mixed assistant messages in place with content.splice(...) but never call invalidateTokenCache(msg) and never set _offloaded:

  • before-prompt-build.ts (splice at line 109)
  • llm-input-l3.ts (splice at line 1377)

So the next buildTiktokenContextSnapshot returns the stale pre-strip count, and since the strip runs just before the guard and emergency snapshots, the inflated total can cross the thresholds and over-delete real messages. Every other in-place mutation already invalidates the cache (l3-helpers.ts:219/257/302, llm-input-l3.ts:1019); these two were the only ones that skipped it. This calls invalidateTokenCache(msg) after each strip loop, only when a block was actually removed, and adds the import to before-prompt-build.ts.

Related Issue | 关联 Issue

Closes #233

Change Type | 修改类型

  • Bug fix | Bug 修复
  • New feature | 新功能
  • Documentation update | 文档更新
  • Code optimization | 代码优化

Self-test Checklist | 自测清单

  • Verified locally | 本地验证通过
  • No existing features affected | 无影响现有功能

Additional Notes | 其他说明

Reproduced against the real context-token-tracker.ts: seed the cache, strip a deleted tool_use block in place, snapshot again.

Before:

after strip, no inval : 537
ground truth (fresh)  : 18
drift                 : 519

After (invalidate following the splice): 18, drift 0.

npm test -> 47 passed. No new test added: the strip loops live in non-exported hook closures, so a unit test cannot drive them without a mocking harness the repo does not have. The runtime repro above is the evidence.

… blocks

The two hooks that strip deleted tool_use blocks in place never invalidated the
per-message token cache, so the next snapshot returned the stale pre-strip count
and could over-delete messages. Mirror the other in-place mutations and call
invalidateTokenCache after each strip loop.
@Maxwell-Code07

Copy link
Copy Markdown
Collaborator

@akhilesharora Clean fix! Adding invalidateTokenCache() at both splice sites is the right approach — targeted and thorough with regression tests. Nice work 👍

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] Stale per-message token cache after in-place tool_use strip over-deletes messages

2 participants