fix(kiro): de-duplicate snapshot traversal and align model-id key-sets - #752
Conversation
Kiro IDE globalStorage snapshots are parsed by collect_kiro_snapshot_text,
which recursed into every present key across three overlapping key-sets.
When a single object stored the same payload under aliased keys (e.g. both
`content` and `text`, or both `messages` and `entries`), the text was
counted once per alias, inflating estimated token totals.
(a) Treat each key-set as an ordered list of aliases for the same logical
payload and descend into only the first present key per group, so each
node's text is collected exactly once.
(b) find_kiro_snapshot_model_id omitted `prompt`/`response`/`parts`/`items`/
`nodes`, keys that collect_kiro_snapshot_text descends into, so a model
id nested under those keys was missed and fell back to `unknown`. Align
the key-sets.
(c) Snapshots have no per-turn timestamps, so the whole blob is emitted as a
single message at the file mtime, mis-bucketing historical usage. No such
timestamps exist in the schema, so mtime is kept and the limitation is
documented explicitly rather than synthesizing timestamps.
Adds regression tests for (a) (aliased text keys and aliased containers) and
(b) (model id under `parts` and `prompt`).
Confidence: high
Scope-risk: narrow
Not-tested: real-world snapshot schemas with deeply nested mixed aliases
|
The latest updates on your projects. Learn more about Vercel for GitHub. 1 Skipped Deployment
|
![cubic-dev-ai[bot]](https://avatars.githubusercontent.com/in/1082092?s=60&v=4){kind=small}
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0788385194
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
|
||
| for key in ["parts", "items", "nodes"] { | ||
| if let Some(item) = map.get(key) { | ||
| if let Some(item) = group.iter().find_map(|key| map.get(*key)) { |
There was a problem hiding this comment.
When a snapshot object contains more than one key from a group with different subtrees (for example a turn with both prompt and response, or a chat object with both messages and history), this find_map only visits the first present key and silently drops the rest, undercounting tokens that the previous traversal would have included. The double-count fix should only skip later keys when they are actually duplicate payloads, rather than treating every key in the group as mutually exclusive.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 86b0883. You're right — the find_map made every key in a group mutually exclusive, so an object holding distinct subtrees under multiple group keys (e.g. both prompt and response, or both messages and history) silently dropped all but the first, undercounting tokens.
collect_kiro_snapshot_text now descends into every present key in each alias group but de-duplicates by VALUE: a subtree structurally equal to one already visited in the same group is skipped. Distinct payloads are all counted; genuine aliases (identical text under content+text, or the same list under messages+entries) are still counted once, preserving this PR's original double-count fix.
Added two regression tests: test_collect_kiro_snapshot_text_counts_distinct_alias_subtrees (distinct prompt/response bodies both counted) and test_collect_kiro_snapshot_text_counts_distinct_container_subtrees (distinct messages/history lists both counted). The existing identical-alias dedup tests still pass. cargo test -p tokscale-core (988 passed) and cargo clippy -p tokscale-core --tests are green.
collect_kiro_snapshot_text used find_map to visit only the FIRST present key in each alias group, which silently dropped distinct subtrees stored under other keys of the same group (e.g. a turn with both `prompt` and `response`, or a chat with both `messages` and `history`), undercounting tokens. Now descend into every present key but skip subtrees structurally equal to one already visited in the same group, so distinct payloads are all counted while genuine aliases are still counted once. Constraint: Must preserve the PR's original double-count fix for identical aliased text Rejected: Visit only first present key | drops distinct sibling payloads Confidence: high Scope-risk: narrow
ported from upstream junhoyeo#735 ported from upstream junhoyeo#737 ported from upstream junhoyeo#747 ported from upstream junhoyeo#750 ported from upstream junhoyeo#752 ported from upstream junhoyeo#760 ported from upstream junhoyeo#766
junhoyeo#752) * fix(kiro): de-duplicate snapshot traversal and align model-id key-sets Kiro IDE globalStorage snapshots are parsed by collect_kiro_snapshot_text, which recursed into every present key across three overlapping key-sets. When a single object stored the same payload under aliased keys (e.g. both `content` and `text`, or both `messages` and `entries`), the text was counted once per alias, inflating estimated token totals. (a) Treat each key-set as an ordered list of aliases for the same logical payload and descend into only the first present key per group, so each node's text is collected exactly once. (b) find_kiro_snapshot_model_id omitted `prompt`/`response`/`parts`/`items`/ `nodes`, keys that collect_kiro_snapshot_text descends into, so a model id nested under those keys was missed and fell back to `unknown`. Align the key-sets. (c) Snapshots have no per-turn timestamps, so the whole blob is emitted as a single message at the file mtime, mis-bucketing historical usage. No such timestamps exist in the schema, so mtime is kept and the limitation is documented explicitly rather than synthesizing timestamps. Adds regression tests for (a) (aliased text keys and aliased containers) and (b) (model id under `parts` and `prompt`). Confidence: high Scope-risk: narrow Not-tested: real-world snapshot schemas with deeply nested mixed aliases * fix(kiro): de-duplicate snapshot aliases by value, not first-key-only collect_kiro_snapshot_text used find_map to visit only the FIRST present key in each alias group, which silently dropped distinct subtrees stored under other keys of the same group (e.g. a turn with both `prompt` and `response`, or a chat with both `messages` and `history`), undercounting tokens. Now descend into every present key but skip subtrees structurally equal to one already visited in the same group, so distinct payloads are all counted while genuine aliases are still counted once. Constraint: Must preserve the PR's original double-count fix for identical aliased text Rejected: Visit only first present key | drops distinct sibling payloads Confidence: high Scope-risk: narrow
1 participant
Problem
crates/tokscale-core/src/sessions/kiro.rsparses Kiro IDE globalStorage snapshots. Three confirmed bugs (#715):collect_kiro_snapshot_textrecursed into every present key across three overlapping key-sets ([prompt,response,content,text,message], the container set, and[parts,items,nodes]). When a single object stored the same payload under aliased keys — e.g. bothcontentandtext, or bothmessagesandentries— the text was counted once per alias, inflating estimated token totals.find_kiro_snapshot_model_idomittedprompt/response/parts/items/nodes— keys thatcollect_kiro_snapshot_textdescends into — so a model id nested under those keys was missed and fell back tounknown.Fix
find_kiro_snapshot_model_id's key-set withcollect_kiro_snapshot_text(addsprompt,response,parts,items,nodes)..jsonland sqliterequest_metadatasources). No timestamps are synthesized; mtime is kept and the limitation is documented with a clear comment so a future schema with per-turn times can split into per-turn messages.Tests
Added regression tests:
test_collect_kiro_snapshot_text_does_not_double_count_aliased_keys— same body undercontent+textcounted once.test_collect_kiro_snapshot_text_does_not_double_count_aliased_containers— same list undermessages+entriescounted once.test_find_kiro_snapshot_model_id_descends_into_aliased_text_keys— model id underpartsandpromptis discovered.cargo test -p tokscale-core(13 kiro tests pass) andcargo clippy -p tokscale-core --testsare green.Residual concern
(c) is documented, not fixed — it requires a snapshot schema that exposes per-turn timestamps, which does not currently exist.
🤖 Generated with Claude Code
Summary by cubic
Fixes double-counting in Kiro snapshot parsing by de-duplicating alias keys by value while still counting distinct subtrees, and aligns model-ID discovery. Fixes #715; date bucketing remains at file mtime and is documented.
collect_kiro_snapshot_textby visiting all keys in each alias group but skipping structurally equal subtrees, so identical aliased text is counted once while distinctprompt/responseormessages/historybodies are both counted (regression tests added).find_kiro_snapshot_model_idwith the same key-sets so IDs underprompt,response,content,text,message, container aliases, andparts/items/nodesare found (regression tests added).Written for commit 86b0883. Summary will update on new commits.