feat: expose worker model and context metadata

This commit is contained in:
2026-10-02 15:29:45 +09:00
parent 43e3b42055
commit cf3cf4b766
25 changed files with 834 additions and 57 deletions
+79 -10
View File
@@ -1403,7 +1403,9 @@ impl App {
output_tokens,
cache_read_input_tokens,
} => {
self.session_context_tokens = input_tokens.unwrap_or(0);
if let Some(input_tokens) = input_tokens {
self.session_context_tokens = input_tokens;
}
// Subtract the cache-hit portion so a tool loop that
// re-sends the same prefix on every request doesn't
// re-count it. cache_creation stays in (it is full
@@ -1475,7 +1477,6 @@ impl App {
return None;
}
self.active_compaction = None;
self.session_context_tokens = 0;
let new_segment_id = lifecycle
.new_segment_id
.as_deref()
@@ -2337,8 +2338,12 @@ impl App {
) -> InternalWorkerView {
let mut app = App::new(snapshot.worker.name.clone());
app.mode = mode;
app.restore_session(&snapshot.session, None);
app.apply_in_flight_snapshot(snapshot.in_flight);
if let Some(greeting) = snapshot.greeting {
app.restore_snapshot(&snapshot.session, greeting, snapshot.in_flight);
} else {
app.restore_session(&snapshot.session, None);
app.apply_in_flight_snapshot(snapshot.in_flight);
}
app.set_worker_status(snapshot.status);
if let Some(error) = snapshot.error {
let _ = app.handle_worker_event(Event::Error {
@@ -2424,7 +2429,10 @@ impl App {
) {
self.greeting = Some(greeting.clone());
self.context_window = greeting.context_window;
self.session_context_tokens = greeting.context_tokens;
self.session_context_tokens = greeting
.context_usage
.map(|usage| usage.tokens)
.unwrap_or(greeting.context_tokens);
self.restore_session(session, Some(greeting));
self.apply_in_flight_snapshot(in_flight);
}
@@ -2441,10 +2449,15 @@ impl App {
_ => None,
})
});
let greeting = greeting.map(|mut greeting| {
greeting.context_tokens = 0;
greeting.context_usage = None;
greeting
});
if let Some(greeting) = greeting.clone() {
self.greeting = Some(greeting.clone());
self.context_window = greeting.context_window;
self.session_context_tokens = greeting.context_tokens;
self.session_context_tokens = 0;
}
let missing_greeting = greeting.is_none();
self.restore_session(session, greeting);
@@ -2836,6 +2849,8 @@ mod rewind_refresh_tests {
assert!(!blocks_contain(&app, "old post-target output"));
assert_eq!(composer_text(&app), "selected rewind input");
assert!(blocks_contain(&app, "Rewound session"));
assert_eq!(app.session_context_tokens, 0);
assert_eq!(app.greeting.as_ref().and_then(|g| g.context_usage), None);
}
#[test]
@@ -2946,6 +2961,8 @@ mod rewind_refresh_tests {
tools: vec![],
context_window: 100,
context_tokens: 10,
reasoning: None,
context_usage: None,
}
}
@@ -4021,6 +4038,7 @@ mod completion_flow_tests {
},
revision,
status: WorkerStatus::Idle,
greeting: None,
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
@@ -4031,6 +4049,28 @@ mod completion_flow_tests {
}
}
#[test]
fn internal_worker_snapshot_uses_child_greeting_metadata() {
let mut app = App::new("parent".into());
let mut child = test_internal_worker_snapshot("child", "research", 1);
let mut greeting = test_greeting();
greeting.worker_name = "research".into();
greeting.model = "child-model".into();
greeting.context_window = 64_000;
greeting.context_usage = Some(protocol::ContextUsage {
tokens: 12_000,
source: protocol::ContextTokenSource::Measured,
});
child.greeting = Some(greeting);
app.replace_internal_worker_snapshots(vec![child]);
let child_app = &app.internal_workers[0].app;
assert_eq!(child_app.greeting.as_ref().unwrap().model, "child-model");
assert_eq!(child_app.context_window, 64_000);
assert_eq!(child_app.session_context_tokens, 12_000);
}
#[test]
fn worker_view_cycle_uses_stable_session_identity_and_wraps_to_main() {
let mut app = App::new("parent".into());
@@ -4309,6 +4349,7 @@ mod completion_flow_tests {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
greeting: None,
status: WorkerStatus::Running,
error: None,
in_flight: Default::default(),
@@ -4505,15 +4546,21 @@ mod completion_flow_tests {
tools: Vec::new(),
context_window: 200_000,
context_tokens: 0,
reasoning: None,
context_usage: None,
}
}
#[test]
fn snapshot_initializes_context_usage() {
fn snapshot_prefers_typed_context_usage_over_legacy_mirror() {
let mut app = App::new("test".into());
let mut greeting = test_greeting();
greeting.context_window = 123_000;
greeting.context_tokens = 45_000;
greeting.context_tokens = 1;
greeting.context_usage = Some(protocol::ContextUsage {
tokens: 45_000,
source: protocol::ContextTokenSource::Measured,
});
app.handle_worker_event(Event::Snapshot {
session: protocol::SessionSnapshot {
@@ -4530,6 +4577,28 @@ mod completion_flow_tests {
assert_eq!(app.session_context_tokens, 45_000);
}
#[test]
fn snapshot_falls_back_to_legacy_context_tokens() {
let mut app = App::new("test".into());
let mut greeting = test_greeting();
greeting.context_window = 100_000;
greeting.context_tokens = 25_000;
app.handle_worker_event(Event::Snapshot {
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
greeting,
state: test_worker_state(WorkerStatus::Idle),
in_flight: Default::default(),
internal_workers: Vec::new(),
});
assert_eq!(app.context_window, 100_000);
assert_eq!(app.session_context_tokens, 25_000);
}
#[test]
fn usage_updates_session_context_tokens_without_cache_discount() {
let mut app = App::new("test".into());
@@ -4566,7 +4635,7 @@ mod completion_flow_tests {
}
#[test]
fn compact_done_resets_session_context_tokens() {
fn legacy_compact_done_does_not_fabricate_zero_context_usage() {
let mut app = App::new("test".into());
app.session_context_tokens = 42_000;
@@ -4574,7 +4643,7 @@ mod completion_flow_tests {
lifecycle.new_segment_id = Some(uuid::Uuid::nil().to_string());
app.handle_worker_event(Event::CompactDone { lifecycle });
assert_eq!(app.session_context_tokens, 0);
assert_eq!(app.session_context_tokens, 42_000);
}
#[test]
+6
View File
@@ -499,10 +499,12 @@ impl RetainedSnapshotSocket {
cwd: String::new(),
provider: "retained session".to_string(),
model: String::new(),
reasoning: None,
scope_summary: "read-only retained session".to_string(),
tools: Vec::new(),
context_window: 0,
context_tokens: 0,
context_usage: None,
},
state: protocol::WorkerStateSnapshot {
last_command_id: 0,
@@ -654,10 +656,12 @@ async fn run_e2e_rewind_fixture(
cwd: workspace_root.display().to_string(),
provider: "e2e-fixture".to_string(),
model: "canned".to_string(),
reasoning: None,
scope_summary: "isolated e2e rewind fixture".to_string(),
tools: Vec::new(),
context_window: 0,
context_tokens: 0,
context_usage: None,
},
});
@@ -2784,10 +2788,12 @@ mod tests {
cwd: "/tmp".into(),
provider: "test".into(),
model: "test".into(),
reasoning: None,
scope_summary: "".into(),
tools: vec![],
context_window: 0,
context_tokens: 0,
context_usage: None,
}
}