perf(core): reduce source cache memory peaks - #78
Conversation
概览本 PR 将源消息缓存从即时写入的 变更说明缓存版本化与写入重构
时序图sequenceDiagram
participant Adapter as 适配器
participant CachePipeline as cache.rs
participant SourceMessageCache as SourceMessageCache
participant FoldUnits as fold_units
participant WriteCache as write_cache
Adapter->>CachePipeline: SourceUnit (含 parser_version)
CachePipeline->>SourceMessageCache: get_meta(path, parser_version)
alt 命中
SourceMessageCache-->>CachePipeline: CachedSourceMeta
CachePipeline-->>Adapter: ParsedUnit { cache_write: None }
else 未命中
CachePipeline->>CachePipeline: 解析 messages
CachePipeline-->>Adapter: ParsedUnit { cache_write: Some(...) }
end
Adapter->>FoldUnits: ParsedUnit
FoldUnits->>FoldUnits: resolve_messages
FoldUnits->>WriteCache: write_cache(cache_write, ctx, messages)
WriteCache->>SourceMessageCache: write_messages / insert
预估代码审查工作量🎯 4 (Complex) | ⏱️ ~60 minutes 相关 PR
小诗
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Merging this PR will improve performance by 10.85%
|
| Benchmark | BASE |
HEAD |
Efficiency | |
|---|---|---|---|---|
| ⚡ | normalize_model_for_grouping[claude_anthropic_date] |
3.2 µs | 2.8 µs | +11.28% |
| ⚡ | normalize_model_for_grouping[kimi_free_tier] |
3.2 µs | 2.9 µs | +11.14% |
| ⚡ | normalize_model_for_grouping[longcat_quantized] |
3.2 µs | 2.9 µs | +10.13% |
Tip
Curious why this is faster? Comment @codspeedbot explain why this is faster on this PR, or directly use the CodSpeed MCP with your agent.
Comparing fix/source-cache-memory-revisions (114ba5a) with personal/local-clients (812c926)
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
crates/tokscale-core/src/adapters/junie.rs (1)
155-191: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win用
Drop守卫恢复TOKSCALE_CONFIG_DIR。这里直到 Line 191 才恢复环境变量;前面的
unwrap()或断言一旦失败,进程级 env 就会残留在临时目录,后续测试会被污染。更稳妥的做法是像crates/tokscale-core/src/lib.rs里的HomeEnvGuard一样用 RAII 包住这次覆盖。🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/tokscale-core/src/adapters/junie.rs` around lines 155 - 191, The test in JUNIE adapter setup restores TOKSCALE_CONFIG_DIR manually at the end, which can leak the temporary env value if any unwrap or assertion fails. Update the adapter_cache_hit_matches_fresh_parse test to use an RAII guard, similar to HomeEnvGuard in lib.rs, around the temporary TOKSCALE_CONFIG_DIR override so the original value is always restored automatically.
🧹 Nitpick comments (2)
crates/tokscale-core/src/adapters/discover.rs (1)
124-130: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win不要在
PrimaryWithSiblings分支硬编码 parser revision。Line 129 绕过了
SourceUnitMeta::parser_revision()的集中映射;以后 bumpNone的 revision 时,这个分支会继续命中旧缓存。改成同一来源可以避免版本漂移。建议修改
FingerprintPolicy::PrimaryWithSiblings { sibling_names } => SourceUnit { client, path, fingerprint_policy: FingerprintPolicy::PrimaryWithSiblings { sibling_names }, meta: crate::adapters::SourceUnitMeta::None, - parser_revision: 1, + parser_revision: crate::adapters::SourceUnitMeta::None.parser_revision(), },🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/tokscale-core/src/adapters/discover.rs` around lines 124 - 130, The PrimaryWithSiblings branch is hardcoding parser_revision instead of using the centralized mapping, which can cause revision drift and stale cache hits. Update the SourceUnit construction in discover.rs so the parser revision comes from SourceUnitMeta::parser_revision() (or the same shared source used by the other branches) rather than a literal value, keeping the revision logic consistent with the rest of SourceUnit creation.crates/tokscale-core/src/message_cache.rs (1)
1476-1502: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win补上旧分片头缺少
parser_revision的回归用例。这里覆盖了
format_version不匹配,但 PR 目标还要求旧 shard(只有旧schema_version、没有parser_revision)必须 stale。建议用一个本地 legacy header struct 按旧字段顺序序列化,断言get_meta(..., 1)仍返回None,避免旧schema_version == CACHE_FORMAT_VERSION时被误收。建议补充测试
+ #[test] + #[serial_test::serial] + fn test_get_meta_ignores_legacy_shard_without_parser_revision() { + #[derive(Serialize)] + struct LegacyCachedShardHeader { + schema_version: u32, + path: CachedPath, + fingerprint: SourceFingerprint, + fallback_timestamp_indices: Vec<usize>, + codex_incremental: Option<CodexIncrementalCache>, + message_count: usize, + } + + let temp_home = TempDir::new().unwrap(); + let prev_env = sandbox_cache_env(temp_home.path()); + + let source = write_temp_file(b"source\n"); + let shard = shard_path(source.path()).unwrap(); + ensure_cache_dir(shard.parent().unwrap()).unwrap(); + let header = LegacyCachedShardHeader { + schema_version: CACHE_FORMAT_VERSION, + path: CachedPath::from_path(source.path()), + fingerprint: SourceFingerprint::from_path(source.path()).unwrap(), + fallback_timestamp_indices: Vec::new(), + codex_incremental: None, + message_count: 0, + }; + let header_bytes = bincode::options().serialize(&header).unwrap(); + let mut file = File::create(&shard).unwrap(); + file.write_all(&(header_bytes.len() as u64).to_le_bytes()).unwrap(); + file.write_all(&header_bytes).unwrap(); + file.flush().unwrap(); + + let loaded = SourceMessageCache::load(); + assert!(loaded.get_meta(source.path(), 1).is_none()); + + restore_cache_env(prev_env); + }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/tokscale-core/src/message_cache.rs` around lines 1476 - 1502, Add a regression test for legacy shard headers that are missing parser_revision, not just mismatched format_version. Update the existing test around test_get_meta_ignores_stale_shard_format_version or add a sibling case that serializes an old header layout with only schema_version/CACHE_FORMAT_VERSION-equivalent fields, then verify SourceMessageCache::load().get_meta(source.path(), 1) returns None. Use a local legacy header struct with the old field order so the test covers shards that match the current format version but are still stale because parser_revision is absent.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/tokscale-cli/tests/copilot_memory.rs`:
- Around line 69-77: In the copilot_memory test flow, remove the hardcoded
no-write-cache behavior from the CLI args so the warm run can actually persist
and hit source-cache; update the test setup around the repeated CLI invocation
in copilot_memory.rs to allow cache writes for the warm path while keeping the
cold-path setup intact. Use the existing copilot_memory test helpers and the CLI
argument builder in the affected test cases so the cold/warm RSS assertions
reflect a real cache hit.
In `@crates/tokscale-core/src/lib.rs`:
- Line 3502: The get_meta test assertions are hardcoding parser revision 1,
which can hide cache persistence regressions when revisions change. Update the
affected tests around get_meta to derive the revision from the source/unit under
test or use a shared test helper/constant, and replace the literal 1 in the
is_none() checks so the revision matches the adapter/parser being exercised
consistently.
---
Outside diff comments:
In `@crates/tokscale-core/src/adapters/junie.rs`:
- Around line 155-191: The test in JUNIE adapter setup restores
TOKSCALE_CONFIG_DIR manually at the end, which can leak the temporary env value
if any unwrap or assertion fails. Update the
adapter_cache_hit_matches_fresh_parse test to use an RAII guard, similar to
HomeEnvGuard in lib.rs, around the temporary TOKSCALE_CONFIG_DIR override so the
original value is always restored automatically.
---
Nitpick comments:
In `@crates/tokscale-core/src/adapters/discover.rs`:
- Around line 124-130: The PrimaryWithSiblings branch is hardcoding
parser_revision instead of using the centralized mapping, which can cause
revision drift and stale cache hits. Update the SourceUnit construction in
discover.rs so the parser revision comes from SourceUnitMeta::parser_revision()
(or the same shared source used by the other branches) rather than a literal
value, keeping the revision logic consistent with the rest of SourceUnit
creation.
In `@crates/tokscale-core/src/message_cache.rs`:
- Around line 1476-1502: Add a regression test for legacy shard headers that are
missing parser_revision, not just mismatched format_version. Update the existing
test around test_get_meta_ignores_stale_shard_format_version or add a sibling
case that serializes an old header layout with only
schema_version/CACHE_FORMAT_VERSION-equivalent fields, then verify
SourceMessageCache::load().get_meta(source.path(), 1) returns None. Use a local
legacy header struct with the old field order so the test covers shards that
match the current format version but are still stale because parser_revision is
absent.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 1fa05b1f-d62e-442e-96ed-64a783c6dba1
📒 Files selected for processing (19)
crates/tokscale-cli/tests/copilot_memory.rscrates/tokscale-core/src/adapters/antigravity.rscrates/tokscale-core/src/adapters/cache.rscrates/tokscale-core/src/adapters/claude.rscrates/tokscale-core/src/adapters/codex.rscrates/tokscale-core/src/adapters/discover.rscrates/tokscale-core/src/adapters/file.rscrates/tokscale-core/src/adapters/goose.rscrates/tokscale-core/src/adapters/hermes.rscrates/tokscale-core/src/adapters/junie.rscrates/tokscale-core/src/adapters/kilo.rscrates/tokscale-core/src/adapters/kiro.rscrates/tokscale-core/src/adapters/mod.rscrates/tokscale-core/src/adapters/opencode.rscrates/tokscale-core/src/adapters/pi.rscrates/tokscale-core/src/adapters/trae.rscrates/tokscale-core/src/lib.rscrates/tokscale-core/src/message_cache.rscrates/tokscale-core/src/sessions/copilot.rs
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
crates/tokscale-core/src/message_cache.rs (1)
501-525: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift让 cache hit 携带并复验
parser_version。Line 520 只在
get_meta阶段校验版本,但后续 cache-hit 路径只保留path,take_messages*会重新读取当前 shard;如果同一路径被另一解析器或并发进程在 parse 与 fold 之间重写,可能返回错误解析器的消息。建议让命中结果携带期望的ParserVersion(最好也携带 fingerprint),并在读取 body 前再次校验 header。Also applies to: 747-765
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/tokscale-core/src/message_cache.rs` around lines 501 - 525, The cache-hit path in get_meta currently validates parser_version only when locating metadata, but the later take_messages* flow still re-reads the shard without rechecking that the body matches the same parser. Update the CachedSourceMeta/lookup result to carry the expected ParserVersion, and in the relevant take_messages* read path re-validate the shard header before returning messages, using the existing get_meta, read_shard_header, and meta_from_header/meta_from_entry flow to ensure the hit still belongs to the same parser version.
🧹 Nitpick comments (1)
crates/tokscale-core/src/adapters/file.rs (1)
242-243: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win避免在测试里硬编码
/tmp路径。这个测试只需要一个稳定的
PathBuf,用TempDir下的路径即可避免平台/环境假设。建议修改
- use std::path::{Path, PathBuf}; + use std::path::Path; @@ - let path = PathBuf::from("/tmp/shared-source.jsonl"); + let dir = tempfile::TempDir::new().unwrap(); + let path = dir.path().join("shared-source.jsonl");As per coding guidelines,
crates/**/{tests,**/}*.rs: “Use temporary directories or fixtures rather than developer-local paths.”Also applies to: 341-345
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/tokscale-core/src/adapters/file.rs` around lines 242 - 243, 测试中不要硬编码 /tmp 路径;在 file.rs 里的相关测试应改为使用 TempDir 生成的稳定 PathBuf,并通过 PathBuf/Path 相关辅助构造替代任何开发机本地假设。请定位测试代码中引用 PathBuf、Path 以及相关断言/初始化逻辑,确保所有临时文件路径都来自 TempDir,而不是固定的 /tmp。Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@crates/tokscale-core/src/message_cache.rs`:
- Around line 501-525: The cache-hit path in get_meta currently validates
parser_version only when locating metadata, but the later take_messages* flow
still re-reads the shard without rechecking that the body matches the same
parser. Update the CachedSourceMeta/lookup result to carry the expected
ParserVersion, and in the relevant take_messages* read path re-validate the
shard header before returning messages, using the existing get_meta,
read_shard_header, and meta_from_header/meta_from_entry flow to ensure the hit
still belongs to the same parser version.
---
Nitpick comments:
In `@crates/tokscale-core/src/adapters/file.rs`:
- Around line 242-243: 测试中不要硬编码 /tmp 路径;在 file.rs 里的相关测试应改为使用 TempDir 生成的稳定
PathBuf,并通过 PathBuf/Path 相关辅助构造替代任何开发机本地假设。请定位测试代码中引用 PathBuf、Path
以及相关断言/初始化逻辑,确保所有临时文件路径都来自 TempDir,而不是固定的 /tmp。
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 3470a449-87ef-4a3c-a3e3-1b882052d9b9
📒 Files selected for processing (9)
crates/tokscale-cli/tests/copilot_memory.rscrates/tokscale-core/src/adapters/cache.rscrates/tokscale-core/src/adapters/codex.rscrates/tokscale-core/src/adapters/discover.rscrates/tokscale-core/src/adapters/file.rscrates/tokscale-core/src/adapters/mod.rscrates/tokscale-core/src/lib.rscrates/tokscale-core/src/message_cache.rscrates/tokscale-core/src/sessions/copilot.rs
🚧 Files skipped from review as they are similar to previous changes (4)
- crates/tokscale-core/src/sessions/copilot.rs
- crates/tokscale-core/src/lib.rs
- crates/tokscale-core/src/adapters/codex.rs
- crates/tokscale-core/src/adapters/cache.rs
There was a problem hiding this comment.
1 issue found across 2 files (changes from recent commits).
Tip: Review your code locally with the cubic CLI to iterate faster.
Re-trigger cubic
makoMakoGo
left a comment
There was a problem hiding this comment.
Code review complete: no blocking findings. The stale read-plan path no longer poisons taken_paths; exact-hit and append-hit cache reads recover by using a verified replacement shard or by reparsing the full Codex file; and append recovery only writes after the full parse can produce a cache entry tied to the final fingerprint. The double-writer, missing-base, and stale-plan regressions cover the previously identified data-loss windows. I would approve this revision; GitHub does not permit approving one's own pull request.
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
crates/tokscale-core/src/adapters/codex.rs (1)
653-657: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win从被测 unit 派生 Codex
ParserVersion。这里手写
ParserVersion::new(ParserId::Codex, 1)会在 Codex parser revision bump 后查询/删除错误 shard;尤其 Line 807-810 的“base cache disappears”测试会因为没有删到当前版本缓存而假阳性通过。请复用codex_unit(...).parser_version或测试 helper 中的当前版本常量。Also applies to: 807-810
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/tokscale-core/src/adapters/codex.rs` around lines 653 - 657, The tests are hardcoding the Codex ParserVersion instead of deriving it from the unit under test, which can cause stale shard lookups and false positives after parser revision bumps. Update the assertions and cache operations in codex.rs to use codex_unit(...).parser_version or the shared current-version test helper constant rather than ParserVersion::new(ParserId::Codex, 1), especially in the base cache disappears coverage.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@crates/tokscale-core/src/adapters/codex.rs`:
- Around line 653-657: The tests are hardcoding the Codex ParserVersion instead
of deriving it from the unit under test, which can cause stale shard lookups and
false positives after parser revision bumps. Update the assertions and cache
operations in codex.rs to use codex_unit(...).parser_version or the shared
current-version test helper constant rather than
ParserVersion::new(ParserId::Codex, 1), especially in the base cache disappears
coverage.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 0e25113d-a48a-491a-8a10-d71d426f18c1
📒 Files selected for processing (9)
crates/tokscale-core/src/adapters/antigravity.rscrates/tokscale-core/src/adapters/cache.rscrates/tokscale-core/src/adapters/claude.rscrates/tokscale-core/src/adapters/codex.rscrates/tokscale-core/src/adapters/mod.rscrates/tokscale-core/src/adapters/opencode.rscrates/tokscale-core/src/adapters/pi.rscrates/tokscale-core/src/lib.rscrates/tokscale-core/src/message_cache.rs
🚧 Files skipped from review as they are similar to previous changes (7)
- crates/tokscale-core/src/adapters/claude.rs
- crates/tokscale-core/src/adapters/pi.rs
- crates/tokscale-core/src/adapters/opencode.rs
- crates/tokscale-core/src/adapters/mod.rs
- crates/tokscale-core/src/adapters/antigravity.rs
- crates/tokscale-core/src/adapters/cache.rs
- crates/tokscale-core/src/message_cache.rs
What changed
SourceUnitMeta.messages.clone()path with a borrowedCacheWritePlanthat writes the shard before moving the same message vector into the sink.This keeps old shards without the new metadata stale with no compatibility migration. Antigravity parser, discovery, model decoding, and dedup semantics are unchanged; it only participates in the generic cache-write plumbing.
Closes #76.
Validation
cargo fmt --all -- --checkcargo test -p tokscale-core sessions::copilotcargo test -p tokscale-core message_cachecargo test -p tokscale-core adapters::antigravitycargo test -p tokscale-core adapters::codexcargo test -p tokscale-core adapters::opencodecargo test -p tokscale-cli --test copilot_memory -- --nocapturecargo test --workspacecargo clippy --workspace --all-targets -- -D warningscargo build --release -p tokscale-cliMemory checks
Generated Copilot OTEL fixture:
52,429,642bytes.17,676 KBmax RSS17,680 KBmax RSSSingle-client real-data smoke checks with release binary:
10,720 KBmax RSS16,860 KBmax RSS37,716 KBmax RSSSummary by cubic
Reduce source-cache memory peaks by writing shards from borrowed message buffers and making cache hits parser-aware by per-source parser identity (id + revision). Cache reads now verify path, parser, and fingerprint; Codex append cache reads are restored; stale or mismatched shards are skipped. Closes #76.
Refactors
parser_version(id + revision); stored in shard headers and keys.parser_versiontoSourceUnit/CachedSourceEntry; adapters set explicit versions, with sensible client defaults.CacheReadPlanfor hits andCacheWritePlan/CacheWritefor writes;write_messageswrites borrowed buffers, andtake_messages[_with_fallback]require a read plan and revalidate before returning.get_meta(path, parser_version)andremove(path, parser_version); Codex/file/policy adapters updated to use plans and borrowed writes while preserving incremental metadata.Bug Fixes
CacheReadPlan, ensuring tails append correctly on cache hits.Written for commit 114ba5a. Summary will update on new commits.
Summary by CodeRabbit
New Features
Bug Fixes
Tests
stdout完全一致,并校验峰值内存与源缓存分片复用。