Skip to content
This repository was archived by the owner on Jul 30, 2026. It is now read-only.

fix(core): make local source attribution structural #194

Merged
makoMakoGo merged 2 commits into from
Jul 25, 2026
Merged

fix(core): make local source attribution structural #194

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1b740e4
fix(core): make adapters authoritative for input attribution
makoMakoGo Jul 25, 2026
cb02f77
refactor(core): make client attribution structural
makoMakoGo Jul 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
108 changes: 90 additions & 18 deletions crates/tokscale-cli/src/tui/cache.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ use super::data::{

/// Cache staleness threshold: 5 minutes (matches TS implementation)
const CACHE_STALE_THRESHOLD_MS: u64 = 5 * 60 * 1000;
const CACHE_SCHEMA_VERSION: u32 = 50;
const CACHE_SCHEMA_VERSION: u32 = 51;

fn sha256_hex(bytes: &[u8]) -> String {
const HEX: &[u8; 16] = b"0123456789abcdef";
Expand Down Expand Up @@ -126,6 +126,15 @@ mod bundle_tests {
InputInventorySignature::from_bytes([0x39; 32])
}

fn health_for_client_space(
client_space: &BTreeMap<String, u64>,
) -> tokscale_core::input_health::HealthReport {
tokscale_core::input_health::HealthReport {
input_data_bytes: client_space.values().copied().sum(),
..Default::default()
}
}

fn nonempty_accumulator(home: &std::path::Path) -> TuiAcc {
for (project, workspace, model, input_tokens, hour) in [
("project-a", "/work/alpha", "claude-sonnet-4.6", 101, 10),
Expand Down Expand Up @@ -283,6 +292,49 @@ mod bundle_tests {
assert_eq!(loaded.data.health, health);
}

#[test]
#[serial]
fn bundle_rejects_data_size_outside_client_space_invariant() {
let (_temp, _guard, clients, scope, sessions, client_space) = fixture();
let invalid_health = tokscale_core::input_health::HealthReport::default();

let error = save_tui_bundle_cache(
&TuiAcc::new(),
&sessions,
&client_space,
&invalid_health,
&clients,
&scope,
signature(),
)
.unwrap_err();
assert!(error
.to_string()
.contains("TUI Data Size 0 does not match client-space total 4096"));

save_tui_bundle_cache(
&TuiAcc::new(),
&sessions,
&client_space,
&health_for_client_space(&client_space),
&clients,
&scope,
signature(),
)
.unwrap();
let path = cache_file().unwrap();
let mut value: serde_json::Value =
serde_json::from_reader(File::open(&path).unwrap()).unwrap();
value["health"]["inputDataBytes"] = serde_json::Value::from(1_u64);
tokscale_core::fs_atomic::write_atomic(&path, &serde_json::to_vec(&value).unwrap())
.unwrap();

assert!(matches!(
load_cache(&clients, &GroupBy::Model, &scope),
CacheResult::Miss
));
}

#[test]
#[serial]
fn bundle_keeps_same_named_agents_separate_across_clients() {
Expand All @@ -298,7 +350,7 @@ mod bundle_tests {
&accumulator,
&[],
&client_space,
&Default::default(),
&health_for_client_space(&client_space),
&clients,
&scope,
signature(),
Expand Down Expand Up @@ -340,7 +392,7 @@ mod bundle_tests {
&TuiAcc::new(),
&sessions,
&client_space,
&Default::default(),
&health_for_client_space(&client_space),
&clients,
&scope,
signature(),
Expand Down Expand Up @@ -584,7 +636,7 @@ mod bundle_tests {
&accumulator,
&[],
&client_space,
&Default::default(),
&health_for_client_space(&client_space),
&clients,
&scope,
signature(),
Expand Down Expand Up @@ -645,7 +697,7 @@ mod bundle_tests {
&accumulator,
&[],
&client_space,
&Default::default(),
&health_for_client_space(&client_space),
&clients,
&scope,
signature(),
Expand Down Expand Up @@ -684,7 +736,7 @@ mod bundle_tests {
&accumulator,
&[],
&client_space,
&Default::default(),
&health_for_client_space(&client_space),
&clients,
&scope,
signature(),
Expand Down Expand Up @@ -776,7 +828,7 @@ mod bundle_tests {
&accumulator,
&[],
&client_space,
&Default::default(),
&health_for_client_space(&client_space),
&clients,
&scope,
signature(),
Expand Down Expand Up @@ -813,7 +865,7 @@ mod bundle_tests {
&accumulator,
&[],
&client_space,
&Default::default(),
&health_for_client_space(&client_space),
&clients,
&scope,
signature(),
Expand Down Expand Up @@ -851,7 +903,7 @@ mod bundle_tests {
&accumulator,
&[],
&client_space,
&Default::default(),
&health_for_client_space(&client_space),
&clients,
&scope,
signature(),
Expand All @@ -878,7 +930,7 @@ mod bundle_tests {
&TuiAcc::new(),
&sessions,
&client_space,
&Default::default(),
&health_for_client_space(&client_space),
&clients,
&scope,
signature(),
Expand Down Expand Up @@ -909,7 +961,7 @@ mod bundle_tests {
&TuiAcc::new(),
&sessions,
&client_space,
&Default::default(),
&health_for_client_space(&client_space),
&clients,
&scope,
signature(),
Expand All @@ -936,7 +988,7 @@ mod bundle_tests {
&TuiAcc::new(),
&sessions,
&client_space,
&Default::default(),
&health_for_client_space(&client_space),
&clients,
&scope,
signature(),
Expand Down Expand Up @@ -964,7 +1016,7 @@ mod bundle_tests {
&TuiAcc::new(),
&sessions,
&client_space,
&Default::default(),
&health_for_client_space(&client_space),
&clients,
&scope,
signature(),
Expand All @@ -989,7 +1041,7 @@ mod bundle_tests {
&TuiAcc::new(),
&sessions,
&client_space,
&Default::default(),
&health_for_client_space(&client_space),
&clients,
&scope,
signature(),
Expand Down Expand Up @@ -1018,7 +1070,7 @@ mod bundle_tests {
&TuiAcc::new(),
&sessions,
&client_space,
&Default::default(),
&health_for_client_space(&client_space),
&clients,
&scope,
signature(),
Expand All @@ -1039,7 +1091,7 @@ mod bundle_tests {
&TuiAcc::new(),
&sessions,
&client_space,
&Default::default(),
&health_for_client_space(&client_space),
&clients,
&scope,
signature(),
Expand All @@ -1065,7 +1117,7 @@ mod bundle_tests {
&TuiAcc::new(),
&sessions,
&client_space,
&Default::default(),
&health_for_client_space(&client_space),
&clients,
&scope,
signature(),
Expand All @@ -1092,7 +1144,7 @@ mod bundle_tests {
&TuiAcc::new(),
&sessions,
&client_space,
&Default::default(),
&health_for_client_space(&client_space),
&clients,
&scope,
signature(),
Expand Down Expand Up @@ -2198,6 +2250,24 @@ fn cache_client_space_matches_exact(
.all(|client| client_space.contains_key(client.as_str()))
}

fn validate_data_size_matches_client_space(
client_space: &BTreeMap<String, u64>,
health: &tokscale_core::input_health::HealthReport,
) -> anyhow::Result<()> {
let total = client_space
.values()
.copied()
.try_fold(0_u64, u64::checked_add)
.ok_or_else(|| anyhow::anyhow!("TUI client-space total exceeds u64::MAX"))?;
if health.input_data_bytes != total {
anyhow::bail!(
"TUI Data Size {} does not match client-space total {total}",
health.input_data_bytes
);
}
Ok(())
}

fn cache_session_clients_are_enabled(
client_universe: &HashSet<ClientId>,
sessions: &[TuiSessionEntry],
Expand Down Expand Up @@ -3067,6 +3137,7 @@ fn load_bundle_from_file(
if !cache_client_space_matches_exact(client_universe, &parsed.client_space) {
anyhow::bail!("cached TUI client-space keys do not match the client universe");
}
validate_data_size_matches_client_space(&parsed.client_space, &parsed.health)?;
if !cache_session_clients_are_enabled(client_universe, &parsed.sessions) {
anyhow::bail!("cached TUI Sessions contain a client outside the client universe");
}
Expand Down Expand Up @@ -3151,6 +3222,7 @@ pub fn save_tui_bundle_cache(
if !cache_client_space_matches_exact(client_universe, client_space) {
anyhow::bail!("TUI client-space keys do not match the client universe");
}
validate_data_size_matches_client_space(client_space, health)?;
if !cache_session_clients_are_enabled(client_universe, sessions) {
anyhow::bail!("TUI Sessions contain a client outside the client universe");
}
Expand Down
8 changes: 6 additions & 2 deletions crates/tokscale-cli/src/tui/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1118,7 +1118,10 @@ mod tests {
&tokscale_core::TuiAcc::new(),
&[session("amp", "cached-session")],
&std::collections::BTreeMap::from([("amp".to_string(), 512)]),
&Default::default(),
&tokscale_core::input_health::HealthReport {
input_data_bytes: 512,
..Default::default()
},
&clients,
&scope,
signature,
Expand All @@ -1130,13 +1133,14 @@ mod tests {
}

fn loaded_snapshot(
data: UsageData,
mut data: UsageData,
sessions: Vec<tokscale_core::TuiSessionEntry>,
client_space: std::collections::BTreeMap<String, u64>,
accumulator: tokscale_core::TuiAcc,
group_by: tokscale_core::GroupBy,
signature: tokscale_core::InputInventorySignature,
) -> BackgroundLoad {
data.health.input_data_bytes = client_space.values().copied().sum();
BackgroundLoad::Loaded {
data: Box::new(data),
sessions,
Expand Down
18 changes: 10 additions & 8 deletions crates/tokscale-core/src/adapters/antigravity.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -131,11 +131,12 @@ fn fold_antigravity_units(
ctx.input_cache.remove(&path, unit.parser_version);
}
let cache_write_outcome = cache_write_outcome?;
sink.extend_messages(
adapter_cache::emit_messages(
unit.client,
messages
.into_iter()
.filter(|message| crate::should_keep_deduped_message(seen, message))
.collect(),
.filter(|message| crate::should_keep_deduped_message(seen, message)),
sink,
);

if cache_write_outcome == adapter_cache::CacheWriteOutcome::NotPlanned && invalidate_cache {
Expand Down Expand Up @@ -188,7 +189,8 @@ mod tests {
use super::*;
use crate::adapters::UnitMessagePayload;
use crate::scanner::ScannerSettings;
use crate::{message_cache, TokenBreakdown, UnifiedMessage};
use crate::sessions::ParsedMessage;
use crate::{message_cache, TokenBreakdown};
use std::collections::BTreeMap;
use std::path::Path;

Expand Down Expand Up @@ -278,7 +280,7 @@ mod tests {
assert!(units.is_empty());
}

fn parsed_unit(path: &Path, meta: InputUnitMeta, message: UnifiedMessage) -> ParsedUnit {
fn parsed_unit(path: &Path, meta: InputUnitMeta, message: ParsedMessage) -> ParsedUnit {
ParsedUnit::healthy(
InputUnit::plain_file(ClientId::Antigravity, path.to_path_buf()).with_meta(meta),
UnitMessagePayload::Fresh(vec![message]),
Expand All @@ -287,9 +289,8 @@ mod tests {
)
}

fn antigravity_message(session_id: &str, dedup_key: Option<u64>) -> UnifiedMessage {
UnifiedMessage::new_with_dedup(
"antigravity",
fn antigravity_message(session_id: &str, dedup_key: Option<u64>) -> ParsedMessage {
ParsedMessage::new_with_dedup(
"gemini-3.1-pro",
"google",
session_id,
Expand Down Expand Up @@ -332,6 +333,7 @@ mod tests {
.unwrap();

assert_eq!(messages.len(), 1);
assert_eq!(messages[0].client.as_ref(), ClientId::Antigravity.as_str());
assert_eq!(messages[0].session_id.as_ref(), "first-session");
}
}
Loading
Loading