Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

fix: synchronize live worker context metadata

Browse Source
This commit is contained in:
Hare
2026-10-02 16:04:55 +09:00
parent cf3cf4b766
commit 3d37f056ad
12 changed files with 352 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/protocol/src/lib.rs
+34
View File
@@ -1245,6 +1245,16 @@ pub enum Event {
#[serde(default, skip_serializing_if = "Option::is_none")]
cache_read_input_tokens: Option<u64>,
},
/// Authoritative current-context occupancy after a non-request boundary.
///
/// Unlike [`Event::Usage`], this is not cumulative request traffic. It lets
/// connected clients converge with reconnect snapshots after run completion,
/// compaction, rewind, and other history rewrites. `None` explicitly means
/// the current value is unavailable and stale occupancy must be cleared.
ContextUsage {
#[serde(default, skip_serializing_if = "Option::is_none")]
usage: Option<ContextUsage>,
},
RunEnd {
result: RunResult,
},
@@ -2440,6 +2450,30 @@ mod tests {
assert_eq!(parsed["data"]["state"]["state"]["state"]["state"], "paused");
}
#[test]
fn context_usage_event_roundtrips_typed_value_and_unknown() {
let event = Event::ContextUsage {
usage: Some(ContextUsage {
tokens: 25_000,
source: ContextTokenSource::Estimated,
}),
};
let json = serde_json::to_string(&event).unwrap();
assert_eq!(
serde_json::from_str::<serde_json::Value>(&json).unwrap(),
serde_json::json!({
"event": "context_usage",
"data": {
"usage": { "tokens": 25_000, "source": "estimated" }
}
})
);
assert!(matches!(
serde_json::from_str::<Event>(r#"{"event":"context_usage","data":{}}"#).unwrap(),
Event::ContextUsage { usage: None }
));
}
#[test]
fn event_snapshot_in_flight_roundtrip_and_default() {
let inbound = r#"{"event":"snapshot","data":{"session":{"entries":[]},"greeting":{"worker_name":"test","cwd":"/tmp","provider":"p","model":"m","scope_summary":"s","tools":[]},"state":{"last_command_id":0,"state":{"kind":"busy","state":{"kind":"run","state":"running"}}}}}"#;
crates/tui/src/app.rs
+62 -4
View File
@@ -247,6 +247,9 @@ pub struct App {
/// Latest session context tokens reported by the Worker. This is the raw
/// `input_tokens` value and is independent from per-run upload totals.
pub session_context_tokens: u64,
/// Availability and provenance for `session_context_tokens`. `None` means
/// the current value is unknown and clients must not render a fabricated 0%.
pub session_context_source: Option<protocol::ContextTokenSource>,
pub context_window: u64,
pub turn_index: usize,
pub current_tool: Option<String>,
@@ -352,6 +355,7 @@ impl App {
run_upload_tokens: 0,
run_output_tokens: 0,
session_context_tokens: 0,
session_context_source: None,
context_window: 0,
turn_index: 0,
current_tool: None,
@@ -1405,6 +1409,9 @@ impl App {
} => {
if let Some(input_tokens) = input_tokens {
self.session_context_tokens = input_tokens;
self.session_context_source = Some(protocol::ContextTokenSource::Measured);
} else {
self.session_context_source = None;
}
// Subtract the cache-hit portion so a tool loop that
// re-sends the same prefix on every request doesn't
@@ -1416,6 +1423,14 @@ impl App {
self.run_upload_tokens += net_input;
self.run_output_tokens += output_tokens.unwrap_or(0);
}
Event::ContextUsage { usage } => {
if let Some(usage) = usage {
self.session_context_tokens = usage.tokens;
self.session_context_source = Some(usage.source);
} else {
self.session_context_source = None;
}
}
Event::Error { code, message } => {
self.handle_error(code, message);
}
@@ -2429,10 +2444,14 @@ impl App {
) {
self.greeting = Some(greeting.clone());
self.context_window = greeting.context_window;
self.session_context_tokens = greeting
.context_usage
.map(|usage| usage.tokens)
.unwrap_or(greeting.context_tokens);
let context_usage = greeting.context_usage.or_else(|| {
(greeting.context_window > 0).then_some(protocol::ContextUsage {
tokens: greeting.context_tokens,
source: protocol::ContextTokenSource::Estimated,
})
});
self.session_context_tokens = context_usage.map(|usage| usage.tokens).unwrap_or_default();
self.session_context_source = context_usage.map(|usage| usage.source);
self.restore_session(session, Some(greeting));
self.apply_in_flight_snapshot(in_flight);
}
@@ -2458,6 +2477,7 @@ impl App {
self.greeting = Some(greeting.clone());
self.context_window = greeting.context_window;
self.session_context_tokens = 0;
self.session_context_source = None;
}
let missing_greeting = greeting.is_none();
self.restore_session(session, greeting);
@@ -2850,6 +2870,7 @@ mod rewind_refresh_tests {
assert_eq!(composer_text(&app), "selected rewind input");
assert!(blocks_contain(&app, "Rewound session"));
assert_eq!(app.session_context_tokens, 0);
assert_eq!(app.session_context_source, None);
assert_eq!(app.greeting.as_ref().and_then(|g| g.context_usage), None);
}
@@ -4575,6 +4596,10 @@ mod completion_flow_tests {
assert_eq!(app.context_window, 123_000);
assert_eq!(app.session_context_tokens, 45_000);
assert_eq!(
app.session_context_source,
Some(protocol::ContextTokenSource::Measured)
);
}
#[test]
@@ -4597,6 +4622,10 @@ mod completion_flow_tests {
assert_eq!(app.context_window, 100_000);
assert_eq!(app.session_context_tokens, 25_000);
assert_eq!(
app.session_context_source,
Some(protocol::ContextTokenSource::Estimated)
);
}
#[test]
@@ -4610,10 +4639,39 @@ mod completion_flow_tests {
});
assert_eq!(app.session_context_tokens, 42_000);
assert_eq!(
app.session_context_source,
Some(protocol::ContextTokenSource::Measured)
);
assert_eq!(app.run_upload_tokens, 2_000);
assert_eq!(app.run_output_tokens, 9);
}
#[test]
fn context_usage_event_converges_live_state_without_changing_run_traffic() {
let mut app = App::new("test".into());
app.run_upload_tokens = 2_000;
app.run_output_tokens = 9;
app.handle_worker_event(Event::ContextUsage {
usage: Some(protocol::ContextUsage {
tokens: 45_000,
source: protocol::ContextTokenSource::Estimated,
}),
});
assert_eq!(app.session_context_tokens, 45_000);
assert_eq!(
app.session_context_source,
Some(protocol::ContextTokenSource::Estimated)
);
assert_eq!(app.run_upload_tokens, 2_000);
assert_eq!(app.run_output_tokens, 9);
app.handle_worker_event(Event::ContextUsage { usage: None });
assert_eq!(app.session_context_source, None);
}
#[test]
fn memory_worker_event_updates_actionbar_state() {
let mut app = App::new("test".into());
crates/tui/src/ui.rs
+29 -5
View File
@@ -1830,13 +1830,20 @@ fn draw_separator(frame: &mut Frame, area: Rect) {
}
fn context_usage_text(app: &App) -> String {
let pct = if app.context_window == 0 {
0
} else {
((app.session_context_tokens as f64 / app.context_window as f64) * 100.0).round() as u64
let Some(source) = app.session_context_source else {
return "Context unavailable".to_string();
};
if app.context_window == 0 {
return "Context unavailable".to_string();
}
let pct =
((app.session_context_tokens as f64 / app.context_window as f64) * 100.0).round() as u64;
let estimate = matches!(source, protocol::ContextTokenSource::Estimated)
.then_some("~")
.unwrap_or_default();
format!(
"{} / {} ({}%)",
"{}{} / {} ({}%)",
estimate,
fmt_tokens(app.session_context_tokens),
fmt_tokens(app.context_window),
pct
@@ -2194,6 +2201,23 @@ mod tests {
);
}
#[test]
fn context_usage_distinguishes_unknown_measured_and_estimated_values() {
let mut app = App::new("worker".into());
app.context_window = 100_000;
app.session_context_tokens = 0;
app.session_context_source = None;
assert_eq!(context_usage_text(&app), "Context unavailable");
app.session_context_tokens = 25_000;
app.session_context_source = Some(protocol::ContextTokenSource::Measured);
assert_eq!(context_usage_text(&app), "25.0k / 100.0k (25%)");
app.session_context_tokens = 20_000;
app.session_context_source = Some(protocol::ContextTokenSource::Estimated);
assert_eq!(context_usage_text(&app), "~20.0k / 100.0k (20%)");
}
#[test]
fn run_elapsed_uses_console_style_units() {
assert_eq!(fmt_run_elapsed(9), "9s");
crates/worker/src/controller.rs
+68 -5
View File
@@ -269,6 +269,17 @@ async fn set_controller_state(
snapshot
}
fn publish_context_usage(
shared_state: &Arc<WorkerSharedState>,
working_event_tx: &broadcast::Sender<Event>,
tokens: Option<u64>,
source: protocol::ContextTokenSource,
) {
shared_state.update_context_usage(tokens, source);
let usage = tokens.map(|tokens| protocol::ContextUsage { tokens, source });
let _ = working_event_tx.send(Event::ContextUsage { usage });
}
async fn set_controller_status(
shared_state: &Arc<WorkerSharedState>,
runtime_dir: &RuntimeDir,
@@ -303,7 +314,9 @@ async fn finish_controller_run<C, St>(
// the terminal run boundary so reconnect snapshots cannot append stale
// partial text/tool arguments after newer entries.
worker.clear_in_flight_events();
shared_state.update_context_usage(
publish_context_usage(
shared_state,
working_event_tx,
Some(worker.total_tokens().tokens),
protocol::ContextTokenSource::Estimated,
);
@@ -797,9 +810,14 @@ impl WorkerController {
greeting,
));
let usage_state = shared_state.clone();
let usage_events = working_event_tx.clone();
worker.engine_mut().on_usage(move |event| {
usage_state
.update_context_usage(event.input_tokens, protocol::ContextTokenSource::Measured);
publish_context_usage(
&usage_state,
&usage_events,
event.input_tokens,
protocol::ContextTokenSource::Measured,
);
});
if let Some(fs_for_view) = fs_for_view {
shared_state.set_fs_view(crate::fs_view::WorkerFsView::new(fs_for_view));
@@ -2518,7 +2536,9 @@ async fn controller_loop<C, St>(
}
}
};
shared_state.update_context_usage(
publish_context_usage(
&shared_state,
&working_event_tx,
Some(worker.total_tokens().tokens),
protocol::ContextTokenSource::Estimated,
);
@@ -2589,7 +2609,9 @@ async fn controller_loop<C, St>(
.await
{
worker.clear_in_flight_events();
shared_state.update_context_usage(
publish_context_usage(
&shared_state,
&working_event_tx,
Some(worker.total_tokens().tokens),
protocol::ContextTokenSource::Estimated,
);
@@ -3566,6 +3588,47 @@ mod tests {
);
}
#[test]
fn context_usage_publication_updates_reconnect_state_and_live_clients() {
let shared_state = Arc::new(WorkerSharedState::new(
"worker".into(),
session_store::new_segment_id(),
String::new(),
protocol::Greeting {
worker_name: "worker".into(),
cwd: "/tmp".into(),
provider: "provider".into(),
model: "model".into(),
reasoning: None,
scope_summary: String::new(),
tools: Vec::new(),
context_window: 100_000,
context_tokens: 0,
context_usage: None,
},
));
let (events, mut receiver) = broadcast::channel(4);
publish_context_usage(
&shared_state,
&events,
Some(25_000),
protocol::ContextTokenSource::Estimated,
);
let expected = protocol::ContextUsage {
tokens: 25_000,
source: protocol::ContextTokenSource::Estimated,
};
assert_eq!(shared_state.greeting().context_usage, Some(expected));
assert!(matches!(
receiver.try_recv(),
Ok(Event::ContextUsage {
usage: Some(usage)
}) if usage == expected
));
}
#[test]
fn automatic_memory_extraction_destinations_are_mutually_exclusive() {
let mut memory = manifest::MemoryFeatureProfileConfig::default();
crates/worker/src/internal_worker.rs
+95 -46
View File
@@ -834,6 +834,23 @@ fn spawn_internal_log_event_bridge(sink: SegmentLogSink, event_tx: broadcast::Se
});
}
fn publish_internal_context_usage(
greeting: &RwLock<protocol::Greeting>,
event_tx: &broadcast::Sender<Event>,
tokens: Option<u64>,
source: protocol::ContextTokenSource,
) {
let usage = tokens.map(|tokens| protocol::ContextUsage { tokens, source });
{
let mut greeting = greeting
.write()
.expect("internal Worker greeting lock poisoned");
greeting.context_tokens = tokens.unwrap_or_default();
greeting.context_usage = usage;
}
let _ = event_tx.send(Event::ContextUsage { usage });
}
pub(crate) async fn prepare_internal_worker_session(
mut worker: Worker<Box<dyn LlmClient>, EphemeralSessionStore>,
store: EphemeralSessionStore,
@@ -864,17 +881,17 @@ pub(crate) async fn prepare_internal_worker_session(
let greeting = Arc::new(RwLock::new(build_greeting(&worker)));
wire_event_bridges_on_engine(&mut worker, &event_tx, &alerter, &in_flight);
let usage_greeting = greeting.clone();
let usage_events = event_tx.clone();
worker.engine_mut().on_usage(move |event| {
let mut greeting = usage_greeting
.write()
.expect("internal Worker greeting lock poisoned");
greeting.context_tokens = event.input_tokens.unwrap_or_default();
greeting.context_usage = event.input_tokens.map(|tokens| protocol::ContextUsage {
tokens,
source: protocol::ContextTokenSource::Measured,
});
publish_internal_context_usage(
&usage_greeting,
&usage_events,
event.input_tokens,
protocol::ContextTokenSource::Measured,
);
});
let actor_greeting = greeting.clone();
let session_id = worker.session_id();
let segment_id = worker.segment_id();
let (command_tx, mut command_rx) = tokio::sync::mpsc::channel(8);
@@ -909,50 +926,56 @@ pub(crate) async fn prepare_internal_worker_session(
match command {
InternalWorkerSessionCommand::Run(input) => {
actor_in_flight.clear();
let cancel_sender = worker.engine_mut().cancel_sender();
let mut run = std::pin::pin!(worker.run_text(&input));
loop {
tokio::select! {
result = &mut run => {
let (turn_status, error, outcome) =
classify_internal_turn_result(result);
actor_in_flight.clear();
*last_outcome.lock().unwrap() = Some(outcome);
status.store(turn_status.encode(), std::sync::atomic::Ordering::Release);
if let Some(message) = error {
*last_error.lock().unwrap() = Some(message.clone());
let _ = event_tx.send(Event::Error {
code: protocol::ErrorCode::Internal,
message,
});
let (turn_status, error, outcome) = {
let cancel_sender = worker.engine_mut().cancel_sender();
let mut run = std::pin::pin!(worker.run_text(&input));
loop {
tokio::select! {
result = &mut run => {
break classify_internal_turn_result(result);
}
send_internal_worker_state(&event_tx, turn_status);
if let Some(callback) = &on_turn_end {
callback(turn_status);
}
state_changed.notify_waiters();
break;
}
command = command_rx.recv() => {
match command {
Some(InternalWorkerSessionCommand::Stop(done)) => {
let _ = cancel_sender.send(()).await;
let _ = (&mut run).await;
stop_done = Some(done);
break 'actor;
}
Some(InternalWorkerSessionCommand::Run(_)) => {
// `send` reserves Running atomically, so a second Run cannot be enqueued.
}
None => {
let _ = cancel_sender.send(()).await;
let _ = (&mut run).await;
break 'actor;
command = command_rx.recv() => {
match command {
Some(InternalWorkerSessionCommand::Stop(done)) => {
let _ = cancel_sender.send(()).await;
let _ = (&mut run).await;
stop_done = Some(done);
break 'actor;
}
Some(InternalWorkerSessionCommand::Run(_)) => {
// `send` reserves Running atomically, so a second Run cannot be enqueued.
}
None => {
let _ = cancel_sender.send(()).await;
let _ = (&mut run).await;
break 'actor;
}
}
}
}
}
};
publish_internal_context_usage(
&actor_greeting,
&event_tx,
Some(worker.total_tokens().tokens),
protocol::ContextTokenSource::Estimated,
);
actor_in_flight.clear();
*last_outcome.lock().unwrap() = Some(outcome);
status.store(turn_status.encode(), std::sync::atomic::Ordering::Release);
if let Some(message) = error {
*last_error.lock().unwrap() = Some(message.clone());
let _ = event_tx.send(Event::Error {
code: protocol::ErrorCode::Internal,
message,
});
}
send_internal_worker_state(&event_tx, turn_status);
if let Some(callback) = &on_turn_end {
callback(turn_status);
}
state_changed.notify_waiters();
}
InternalWorkerSessionCommand::Stop(done) => {
stop_done = Some(done);
@@ -1509,11 +1532,37 @@ permission = "write"
assert_eq!(handle.protocol_snapshot().in_flight.blocks.len(), 1);
let entries_after_first = handle.entries().len();
assert!(entries_after_first >= 4);
assert!(matches!(
handle
.protocol_snapshot()
.greeting
.and_then(|greeting| greeting.context_usage),
Some(protocol::ContextUsage {
source: protocol::ContextTokenSource::Estimated,
..
})
));
let mut events = handle.subscribe_events();
handle.send("follow-up").await.expect("send follow-up turn");
assert_eq!(
handle.wait_until_idle().await,
InternalWorkerSessionStatus::Idle
);
let mut saw_estimated_context = false;
while let Ok(event) = events.try_recv() {
if matches!(
event,
Event::ContextUsage {
usage: Some(protocol::ContextUsage {
source: protocol::ContextTokenSource::Estimated,
..
})
}
) {
saw_estimated_context = true;
}
}
assert!(saw_estimated_context);
assert!(handle.protocol_snapshot().in_flight.blocks.is_empty());
assert_eq!(calls.load(Ordering::SeqCst), 2);
assert!(handle.entries().len() > entries_after_first);
web/workspace/src/lib/generated/protocol-validator.ts
+1 -1
View File
File diff suppressed because one or more lines are too long
web/workspace/src/lib/generated/protocol.ts
+1 -1
View File
@@ -381,7 +381,7 @@ summary: string,
* Full tool output. Absent when the tool chose to return
* summary-only, or when the result was pruned.
*/
output?: string | null, disposition?: ToolResultDisposition | null, is_error: boolean, } } | { "event": "usage", "data": { input_tokens: number | null, output_tokens: number | null, cache_read_input_tokens?: number | null, } } | { "event": "run_end", "data": { result: RunResult, } } | { "event": "error", "data": { code: ErrorCode, message: string, } } | { "event": "snapshot", "data": { session: SessionSnapshot, greeting: Greeting,
output?: string | null, disposition?: ToolResultDisposition | null, is_error: boolean, } } | { "event": "usage", "data": { input_tokens: number | null, output_tokens: number | null, cache_read_input_tokens?: number | null, } } | { "event": "context_usage", "data": { usage?: ContextUsage | null, } } | { "event": "run_end", "data": { result: RunResult, } } | { "event": "error", "data": { code: ErrorCode, message: string, } } | { "event": "snapshot", "data": { session: SessionSnapshot, greeting: Greeting,
/**
* Full revisioned live execution state. `Stopped` remains Runtime
* catalog authority and is deliberately not represented here.
web/workspace/src/lib/workspace/console/WorkerContextStatus.browser.spec.ts
+11
View File
@@ -39,4 +39,15 @@ test("renders honest unavailable state instead of a fabricated zero percent", as
});
expect(view.container.textContent).toContain("claude · 8.2k token budget");
expect(view.container.textContent).toContain("Context ~12.0k / 64.0k (19%)");
await view.rerender({
metadata: {
model: "claude",
reasoning: { kind: "budget_tokens", budget_tokens: -1 },
contextWindow: 64_000,
contextTokens: 12_000,
contextSource: "estimated",
},
});
expect(view.container.textContent).toContain("claude · dynamic token budget");
});
web/workspace/src/lib/workspace/console/model.test.ts
+19 -2
View File
@@ -224,14 +224,17 @@ Deno.test("snapshot and usage keep current context separate from run traffic", (
cache_read_input_tokens: 140_000,
output_tokens: 2_000,
} }, observedAtMs: 2 },
{ eventId: "context", event: { event: "context_usage", data: {
usage: { tokens: 158_000, source: "estimated" },
} }, observedAtMs: 3 },
]);
assertEquals(projection.workerMetadata, {
model: "gpt-6-astra",
reasoning: { kind: "effort", effort: "high" },
contextWindow: 272_000,
contextTokens: 150_000,
contextSource: "measured",
contextTokens: 158_000,
contextSource: "estimated",
});
assertEquals(projection.runActivity.uploadTokens, 10_000);
assertEquals(projection.runActivity.outputTokens, 2_000);
@@ -279,6 +282,20 @@ Deno.test("unknown and destructive boundaries never fabricate context percentage
]);
assertEquals(projection.workerMetadata?.contextTokens, null);
assertEquals(projection.workerMetadata?.contextSource, null);
projection = projectConsole([
{ eventId: "live", event: live },
{ eventId: "rewind", event: { event: "rewind_applied", data: {
session: canonicalSession([]),
input: [],
summary: { truncated_to_entries: 0, discarded_entries: 1, tool_side_effect_warning: false },
} } },
{ eventId: "post-rewind-context", event: { event: "context_usage", data: {
usage: { tokens: 12, source: "estimated" },
} } },
]);
assertEquals(projection.workerMetadata?.contextTokens, 12);
assertEquals(projection.workerMetadata?.contextSource, "estimated");
});
Deno.test("large paste segments project compact artifact metadata", () => {
web/workspace/src/lib/workspace/console/model.ts
+21
View File
@@ -3,6 +3,7 @@ import type {
CommandEvent,
CommandSnapshot,
CommandStreamSlice,
ContextUsage,
Event as ProtocolEvent,
Greeting,
InFlightBlock,
@@ -256,6 +257,20 @@ function metadataWithUsage(
};
}
function metadataWithContextUsage(
metadata: ConsoleWorkerMetadata | null,
usage: ContextUsage | null | undefined,
): ConsoleWorkerMetadata {
const tokens = safeTokenCount(usage?.tokens);
return {
model: metadata?.model ?? null,
reasoning: metadata?.reasoning ?? null,
contextWindow: metadata?.contextWindow ?? null,
contextTokens: tokens,
contextSource: tokens === null ? null : usage?.source ?? null,
};
}
function metadataWithoutContext(
metadata: ConsoleWorkerMetadata | null,
): ConsoleWorkerMetadata | null {
@@ -1127,6 +1142,12 @@ export function applyProtocolEvent(
event.data.input_tokens,
);
break;
case "context_usage":
next.workerMetadata = metadataWithContextUsage(
next.workerMetadata,
event.data.usage,
);
break;
case "error":
if (hasDurableRunFailure(next, event.data.message)) break;
next.lines.push(
web/workspace/src/lib/workspace/console/run-status.test.ts
+4
View File
@@ -58,6 +58,10 @@ Deno.test("worker metadata formats effort, budget, estimate, and unknown safely"
);
assertEquals(
formatReasoning({ kind: "budget_tokens", budget_tokens: -1 }),
"dynamic token budget",
);
assertEquals(
formatReasoning({ kind: "budget_tokens", budget_tokens: -2 }),
"reasoning unavailable",
);
assertEquals(formatContextSummary(null), "Context unavailable");
web/workspace/src/lib/workspace/console/worker-metadata.ts
+7 -4
View File
@@ -8,10 +8,13 @@ export function formatReasoning(reasoning: ReasoningConfig | null): string {
const effort = reasoning.effort.trim();
return effort.length > 0 ? effort : "reasoning unavailable";
}
if (
!Number.isSafeInteger(reasoning.budget_tokens) ||
reasoning.budget_tokens <= 0
) {
if (!Number.isSafeInteger(reasoning.budget_tokens)) {
return "reasoning unavailable";
}
if (reasoning.budget_tokens === -1) {
return "dynamic token budget";
}
if (reasoning.budget_tokens <= 0) {
return "reasoning unavailable";
}
return `${formatRunTokens(reasoning.budget_tokens)} token budget`;