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

Merge remote-tracking branch 'origin/develop' into work/T-673-subjektiv-session-history

Browse Source
This commit is contained in:
Hare
2026-10-02 16:52:02 +09:00
parent ee6aa9de41 3ae9b1c6d1
commit ce841cfee0
28 changed files with 1470 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/protocol/src/lib.rs
+117 -7
View File
@@ -766,6 +766,10 @@ pub struct InternalWorkerSnapshot {
pub worker: InternalWorkerRef,
pub revision: u64,
pub session: SessionSnapshot,
/// Public execution metadata for this child. Older snapshots and service
/// workers may omit it; clients must never substitute the parent's values.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub greeting: Option<Greeting>,
#[serde(default)]
pub status: WorkerStatus,
#[serde(default, skip_serializing_if = "Option::is_none")]
@@ -1241,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,
},
@@ -1667,11 +1681,45 @@ impl InFlightToolCallState {
}
}
/// Worker self-description rendered by the TUI when a session starts empty.
/// Public reasoning configuration resolved for the model client.
///
/// Built once in the Worker controller from the resolved manifest and
/// transmitted alongside `Event::Snapshot` so clients don't need
/// their own view of the manifest.
/// This deliberately exposes only the effective provider control, not the
/// Worker manifest or provider credentials. The tagged shape prevents token
/// budgets from being presented as effort labels.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[cfg_attr(feature = "json-schema", derive(schemars::JsonSchema))]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum ReasoningConfig {
Effort { effort: String },
BudgetTokens { budget_tokens: i32 },
}
/// Provenance of the current-context token count in a [`Greeting`].
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[cfg_attr(feature = "json-schema", derive(schemars::JsonSchema))]
#[serde(rename_all = "snake_case")]
pub enum ContextTokenSource {
/// Provider-reported input occupancy from the most recent LLM request.
Measured,
/// Current history projected from the Worker's token accounting records.
Estimated,
}
/// Current prompt-context occupancy, distinct from cumulative run traffic.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[cfg_attr(feature = "json-schema", derive(schemars::JsonSchema))]
pub struct ContextUsage {
pub tokens: u64,
pub source: ContextTokenSource,
}
/// Worker self-description rendered by clients when a session starts empty.
///
/// Snapshot metadata contains only resolved public execution information so
/// clients do not need their own copy of the manifest.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[cfg_attr(feature = "json-schema", derive(schemars::JsonSchema))]
@@ -1680,14 +1728,22 @@ pub struct Greeting {
pub cwd: String,
pub provider: String,
pub model: String,
/// Effective reasoning control. Absent means unspecified or unavailable.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub reasoning: Option<ReasoningConfig>,
pub scope_summary: String,
pub tools: Vec<String>,
/// Model context window in tokens. Always filled by the Worker greeting.
/// Resolved effective model context window in tokens. Legacy producers use
/// zero when this is unavailable; clients must treat that as unknown.
#[serde(default)]
pub context_window: u64,
/// Estimated current session context tokens at connect time.
/// Legacy numeric mirror retained for mixed-version clients.
#[serde(default)]
pub context_tokens: u64,
/// Typed current-context value. Older snapshots omit it; clients may fall
/// back to `context_tokens` only when `context_window` is known.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub context_usage: Option<ContextUsage>,
}
// ---------------------------------------------------------------------------
@@ -2354,6 +2410,11 @@ mod tests {
tools: vec!["Read".into()],
context_window: 200_000,
context_tokens: 42_000,
reasoning: Some(ReasoningConfig::BudgetTokens { budget_tokens: -1 }),
context_usage: Some(ContextUsage {
tokens: 42_000,
source: ContextTokenSource::Measured,
}),
},
state: WorkerStatus::Paused.into(),
in_flight: InFlightSnapshot::default(),
@@ -2372,16 +2433,63 @@ mod tests {
assert_eq!(parsed["data"]["greeting"]["tools"][0], "Read");
assert_eq!(parsed["data"]["greeting"]["context_window"], 200_000);
assert_eq!(parsed["data"]["greeting"]["context_tokens"], 42_000);
assert_eq!(
parsed["data"]["greeting"]["reasoning"]["kind"],
"budget_tokens"
);
assert_eq!(parsed["data"]["greeting"]["reasoning"]["budget_tokens"], -1);
assert_eq!(
parsed["data"]["greeting"]["context_usage"]["tokens"],
42_000
);
assert_eq!(
parsed["data"]["greeting"]["context_usage"]["source"],
"measured"
);
assert_eq!(parsed["data"]["state"]["state"]["kind"], "busy");
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"}}}}}"#;
let decoded: Event = serde_json::from_str(inbound).unwrap();
match decoded {
Event::Snapshot { in_flight, .. } => assert!(in_flight.is_empty()),
Event::Snapshot {
greeting,
in_flight,
..
} => {
assert!(in_flight.is_empty());
assert_eq!(greeting.reasoning, None);
assert_eq!(greeting.context_usage, None);
assert_eq!(greeting.context_window, 0);
assert_eq!(greeting.context_tokens, 0);
}
other => panic!("expected Snapshot, got {other:?}"),
}
@@ -2399,6 +2507,8 @@ mod tests {
tools: Vec::new(),
context_window: 0,
context_tokens: 0,
reasoning: None,
context_usage: None,
},
state: WorkerStatus::Running.into(),
in_flight: InFlightSnapshot {
crates/protocol/src/typescript.rs
+71 -16
View File
@@ -6,17 +6,18 @@ use ts_rs::{Config, TS};
use crate::{
Alert, AlertLevel, AlertSource, CommandEvent, CommandSnapshot, CommandStatus, CommandStream,
CommandStreamSlice, CompactionLifecycle, CompactionLifecycleState, CompactionPhase,
CompactionTrigger, CompletionEntry, CompletionKind, ErrorCode, Event, Greeting, InFlightBlock,
InFlightCompaction, InFlightSnapshot, InFlightToolCallState, InternalWorkerKind,
InternalWorkerRef, InternalWorkerSnapshot, InvokeKind, MemoryWorkerEvent, Method,
PasteArtifactAvailability, PasteArtifactMediaType, PasteArtifactRef, PendingSubmissionSummary,
PendingSubmissionsSnapshot, Permission, RewindSummary, RewindTarget, RewindTargetId,
RunFailureKind, RunResult, RunResumeSource, RunYieldReason, ScopeRule, Segment,
SessionContentPart, SessionConversationTurn, SessionEntryProvenance,
SessionHistoryLineageBoundary, SessionHistoryPage, SessionMessageRole, SessionSnapshot,
SessionSnapshotEntry, SessionSnapshotEntryData, SessionToolAttachment, SubmissionDisposition,
SymlinkPolicy, ToolResultDisposition, TurnResult, UploadedFileAvailability, UploadedFileRef,
WorkerBusyState, WorkerCommandAcknowledgement, WorkerCommandDisposition, WorkerCommandEnvelope,
CompactionTrigger, CompletionEntry, CompletionKind, ContextTokenSource, ContextUsage,
ErrorCode, Event, Greeting, InFlightBlock, InFlightCompaction, InFlightSnapshot,
InFlightToolCallState, InternalWorkerKind, InternalWorkerRef, InternalWorkerSnapshot,
InvokeKind, MemoryWorkerEvent, Method, PasteArtifactAvailability, PasteArtifactMediaType,
PasteArtifactRef, PendingSubmissionSummary, PendingSubmissionsSnapshot, Permission,
ReasoningConfig, RewindSummary, RewindTarget, RewindTargetId, RunFailureKind, RunResult,
RunResumeSource, RunYieldReason, ScopeRule, Segment, SessionContentPart,
SessionConversationTurn, SessionEntryProvenance, SessionHistoryLineageBoundary,
SessionHistoryPage, SessionMessageRole, SessionSnapshot, SessionSnapshotEntry,
SessionSnapshotEntryData, SessionToolAttachment, SubmissionDisposition, SymlinkPolicy,
ToolResultDisposition, TurnResult, UploadedFileAvailability, UploadedFileRef, WorkerBusyState,
WorkerCommandAcknowledgement, WorkerCommandDisposition, WorkerCommandEnvelope,
WorkerCommandKind, WorkerEvent, WorkerMaintenanceState, WorkerRunState, WorkerState,
WorkerStateSnapshot, WorkerStatus,
subscription::{
@@ -116,10 +117,13 @@ pub fn generated_protocol_types() -> String {
push_decl::<PendingSubmissionsSnapshot>(&cfg, &mut output);
push_decl::<SubmissionDisposition>(&cfg, &mut output);
push_decl::<SessionSnapshot>(&cfg, &mut output);
push_decl::<ReasoningConfig>(&cfg, &mut output);
push_decl::<ContextTokenSource>(&cfg, &mut output);
push_decl::<ContextUsage>(&cfg, &mut output);
push_decl::<Greeting>(&cfg, &mut output);
push_decl::<InternalWorkerKind>(&cfg, &mut output);
push_decl::<InternalWorkerRef>(&cfg, &mut output);
push_decl::<InternalWorkerSnapshot>(&cfg, &mut output);
push_decl::<Greeting>(&cfg, &mut output);
push_decl::<Alert>(&cfg, &mut output);
push_decl::<MemoryWorkerEvent>(&cfg, &mut output);
push_decl::<PasteArtifactMediaType>(&cfg, &mut output);
@@ -256,20 +260,71 @@ fn rust_serialized_subscription_frame_fixtures() -> Vec<SubscriptionFrame> {
));
let pending =
SubscriptionFrame::new(SubscriptionFramePayload::Event(SubscriptionEvent::Event {
subscription_id,
subscription_id: subscription_id.clone(),
payload: SubscriptionEventPayload::WorkerProtocol {
worker_id,
worker_id: worker_id.clone(),
event: Event::PendingSubmissionsChanged {
pending: PendingSubmissionsSnapshot::default(),
},
},
}));
for fixture in [&subscribed, &pending] {
let greeting = Greeting {
worker_name: "fixture-worker".into(),
cwd: "/fixture".into(),
provider: "openai_responses".into(),
model: "gpt-6-astra".into(),
reasoning: Some(ReasoningConfig::Effort {
effort: "high".into(),
}),
scope_summary: "fixture scope".into(),
tools: vec!["Read".into()],
context_window: 272_000,
context_tokens: 142_000,
context_usage: Some(ContextUsage {
tokens: 142_000,
source: ContextTokenSource::Measured,
}),
};
let snapshot =
SubscriptionFrame::new(SubscriptionFramePayload::Event(SubscriptionEvent::Event {
subscription_id,
payload: SubscriptionEventPayload::WorkerProtocol {
worker_id,
event: Event::Snapshot {
session: SessionSnapshot {
pending_submissions: PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
greeting: greeting.clone(),
state: WorkerStateSnapshot::initial(),
in_flight: InFlightSnapshot::default(),
internal_workers: vec![InternalWorkerSnapshot {
worker: InternalWorkerRef {
session_id: "fixture-child-session".into(),
name: "fixture-child".into(),
parent_session_id: Some("fixture-parent-session".into()),
kind: InternalWorkerKind::SubWorker,
},
revision: 1,
session: SessionSnapshot {
pending_submissions: PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
greeting: Some(greeting),
status: WorkerStatus::Idle,
error: None,
in_flight: InFlightSnapshot::default(),
internal_workers: Vec::new(),
}],
},
},
}));
for fixture in [&subscribed, &pending, &snapshot] {
fixture
.validate()
.expect("Rust-serialized Browser compatibility fixture must validate");
}
vec![subscribed, pending]
vec![subscribed, pending, snapshot]
}
fn close_object_shapes(schema: &mut Value) {
crates/tui/src/app.rs
+185 -10
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,
@@ -1403,7 +1407,12 @@ 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;
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
// re-count it. cache_creation stays in (it is full
@@ -1414,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);
}
@@ -1475,7 +1492,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 +2353,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 +2444,14 @@ impl App {
) {
self.greeting = Some(greeting.clone());
self.context_window = greeting.context_window;
self.session_context_tokens = 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);
}
@@ -2441,10 +2468,16 @@ 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;
self.session_context_source = None;
}
let missing_greeting = greeting.is_none();
self.restore_session(session, greeting);
@@ -2836,6 +2869,9 @@ 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.session_context_source, None);
assert_eq!(app.greeting.as_ref().and_then(|g| g.context_usage), None);
}
#[test]
@@ -2946,6 +2982,8 @@ mod rewind_refresh_tests {
tools: vec![],
context_window: 100,
context_tokens: 10,
reasoning: None,
context_usage: None,
}
}
@@ -4021,6 +4059,7 @@ mod completion_flow_tests {
},
revision,
status: WorkerStatus::Idle,
greeting: None,
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
@@ -4031,6 +4070,76 @@ 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 live_internal_worker_snapshot_installs_child_metadata_before_context_updates() {
let mut app = App::new("parent".into());
let worker = test_internal_worker_snapshot("child", "research", 1).worker;
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,
});
app.handle_worker_event(Event::InternalWorker {
worker: worker.clone(),
revision: 1,
event: Box::new(Event::Snapshot {
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
},
greeting,
state: WorkerStatus::Running.into(),
in_flight: Default::default(),
internal_workers: Vec::new(),
}),
});
app.handle_worker_event(Event::InternalWorker {
worker,
revision: 2,
event: Box::new(Event::ContextUsage {
usage: Some(protocol::ContextUsage {
tokens: 14_000,
source: protocol::ContextTokenSource::Estimated,
}),
}),
});
let child = &app.internal_workers[0].app;
assert_eq!(child.greeting.as_ref().unwrap().model, "child-model");
assert_eq!(child.context_window, 64_000);
assert_eq!(child.session_context_tokens, 14_000);
assert_eq!(
child.session_context_source,
Some(protocol::ContextTokenSource::Estimated)
);
}
#[test]
fn worker_view_cycle_uses_stable_session_identity_and_wraps_to_main() {
let mut app = App::new("parent".into());
@@ -4309,6 +4418,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 +4615,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 {
@@ -4528,6 +4644,36 @@ 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]
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);
assert_eq!(
app.session_context_source,
Some(protocol::ContextTokenSource::Estimated)
);
}
#[test]
@@ -4541,10 +4687,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());
@@ -4566,7 +4741,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 +4749,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]
crates/tui/src/console/mod.rs
+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,
}
}
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-runtime/src/http_server.rs
+2
View File
@@ -3995,6 +3995,8 @@ mod ws_tests {
tools: Vec::new(),
context_window: 0,
context_tokens: 0,
reasoning: None,
context_usage: None,
},
state: protocol::WorkerStateSnapshot::initial(),
in_flight: protocol::InFlightSnapshot {
crates/worker-runtime/src/runtime.rs
+5
View File
@@ -5626,6 +5626,7 @@ mod tests {
worker,
revision,
status,
greeting: None,
session: protocol::SessionSnapshot {
pending_submissions: protocol::PendingSubmissionsSnapshot::default(),
entries: Vec::new(),
@@ -5651,6 +5652,8 @@ mod tests {
tools: Vec::new(),
context_window: 0,
context_tokens: 0,
reasoning: None,
context_usage: None,
},
state: protocol::WorkerStatus::Idle.into(),
in_flight: protocol::InFlightSnapshot::default(),
@@ -8263,6 +8266,8 @@ mod tests {
tools: Vec::new(),
context_window: 128,
context_tokens: 64,
reasoning: None,
context_usage: None,
},
state: protocol::WorkerStatus::Running.into(),
in_flight: protocol::InFlightSnapshot {
crates/worker-runtime/src/worker_backend.rs
+2
View File
@@ -4520,6 +4520,8 @@ mod tests {
tools: Vec::new(),
context_window: 1_000,
context_tokens: 0,
reasoning: None,
context_usage: None,
},
));
hub.workers.lock().unwrap().insert(
crates/worker/src/controller.rs
+138 -3
View File
@@ -151,7 +151,7 @@ impl WorkerHandle {
.snapshot();
let event = Event::Snapshot {
session,
greeting: self.shared_state.greeting.clone(),
greeting: self.shared_state.greeting(),
state: self.shared_state.snapshot(),
in_flight,
internal_workers: self.spawned_registry.internal_worker_snapshots(),
@@ -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,6 +314,12 @@ 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();
publish_context_usage(
shared_state,
working_event_tx,
Some(worker.total_tokens().tokens),
protocol::ContextTokenSource::Estimated,
);
set_controller_status(shared_state, runtime_dir, working_event_tx, new_status).await;
}
@@ -792,6 +809,16 @@ impl WorkerController {
manifest_toml.clone(),
greeting,
));
let usage_state = shared_state.clone();
let usage_events = working_event_tx.clone();
worker.engine_mut().on_usage(move |event| {
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,6 +2545,12 @@ async fn controller_loop<C, St>(
}
}
};
publish_context_usage(
&shared_state,
&working_event_tx,
Some(worker.total_tokens().tokens),
protocol::ContextTokenSource::Estimated,
);
if !matches!(
result,
Err(WorkerError::Store(_))
@@ -2585,6 +2618,12 @@ async fn controller_loop<C, St>(
.await
{
worker.clear_in_flight_events();
publish_context_usage(
&shared_state,
&working_event_tx,
Some(worker.total_tokens().tokens),
protocol::ContextTokenSource::Estimated,
);
let snapshot = shared_state.transition(WorkerState::Idle);
let _ = working_event_tx.send(Event::WorkerState { snapshot });
}
@@ -3392,7 +3431,24 @@ fn model_supports_image_attachments(model: &manifest::ModelManifest) -> bool {
})
}
fn build_greeting<C, St>(worker: &Worker<C, St>) -> protocol::Greeting
fn public_reasoning_config(
reasoning: &agen::llm_client::capability::ReasoningControl,
) -> protocol::ReasoningConfig {
match reasoning {
agen::llm_client::capability::ReasoningControl::Effort(effort) => {
protocol::ReasoningConfig::Effort {
effort: effort.as_str().to_owned(),
}
}
agen::llm_client::capability::ReasoningControl::BudgetTokens(budget_tokens) => {
protocol::ReasoningConfig::BudgetTokens {
budget_tokens: *budget_tokens,
}
}
}
}
pub(crate) fn build_greeting<C, St>(worker: &Worker<C, St>) -> protocol::Greeting
where
C: LlmClient + 'static,
St: Store,
@@ -3435,6 +3491,7 @@ where
.into_iter()
.map(|def| def.name)
.collect();
let context_tokens = worker.total_tokens().tokens;
protocol::Greeting {
worker_name: manifest.worker.name.clone(),
cwd: worker
@@ -3443,10 +3500,19 @@ where
.unwrap_or_default(),
provider: provider_name,
model: model_id,
reasoning: manifest
.engine
.reasoning
.as_ref()
.map(public_reasoning_config),
scope_summary: worker.scope_snapshot().summary(),
tools: tool_names,
context_window,
context_tokens: worker.total_tokens().tokens,
context_tokens,
context_usage: Some(protocol::ContextUsage {
tokens: context_tokens,
source: protocol::ContextTokenSource::Estimated,
}),
}
}
@@ -3509,6 +3575,69 @@ mod tests {
assert!(materialize < exposure);
}
#[test]
fn public_reasoning_config_preserves_effort_and_signed_budget() {
let effort =
public_reasoning_config(&agen::llm_client::capability::ReasoningControl::Effort(
agen::llm_client::capability::ReasoningEffort::Other("provider-native".into()),
));
assert_eq!(
effort,
protocol::ReasoningConfig::Effort {
effort: "provider-native".into(),
}
);
let budget = public_reasoning_config(
&agen::llm_client::capability::ReasoningControl::BudgetTokens(-1),
);
assert_eq!(
budget,
protocol::ReasoningConfig::BudgetTokens { budget_tokens: -1 }
);
}
#[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();
@@ -3667,6 +3796,8 @@ mod tests {
tools: Vec::new(),
context_window: 200_000,
context_tokens: 0,
reasoning: None,
context_usage: None,
},
));
let notify_buffer = NotifyBuffer::new();
@@ -3716,6 +3847,8 @@ mod tests {
tools: Vec::new(),
context_window: 200_000,
context_tokens: 0,
reasoning: None,
context_usage: None,
},
state: WorkerStatus::Idle.into(),
in_flight: Default::default(),
@@ -4424,6 +4557,8 @@ mod tests {
tools: Vec::new(),
context_window: 1,
context_tokens: 0,
reasoning: None,
context_usage: None,
},
);
assert!(
crates/worker/src/discovery.rs
+12
View File
@@ -1507,6 +1507,8 @@ mod tests {
tools: Vec::new(),
context_window: 0,
context_tokens: 0,
reasoning: None,
context_usage: None,
},
state: WorkerStatus::Idle.into(),
in_flight: Default::default(),
@@ -1543,6 +1545,8 @@ mod tests {
tools: Vec::new(),
context_window: 0,
context_tokens: 0,
reasoning: None,
context_usage: None,
},
state: WorkerStatus::Idle.into(),
in_flight: Default::default(),
@@ -1634,6 +1638,8 @@ mod tests {
tools: Vec::new(),
context_window: 0,
context_tokens: 0,
reasoning: None,
context_usage: None,
},
state: WorkerStatus::Idle.into(),
in_flight: Default::default(),
@@ -1661,6 +1667,8 @@ mod tests {
tools: Vec::new(),
context_window: 0,
context_tokens: 0,
reasoning: None,
context_usage: None,
},
state: WorkerStatus::Idle.into(),
in_flight: Default::default(),
@@ -1765,6 +1773,8 @@ mod tests {
tools: Vec::new(),
context_window: 0,
context_tokens: 0,
reasoning: None,
context_usage: None,
},
state: WorkerStatus::Paused.into(),
in_flight: Default::default(),
@@ -1820,6 +1830,8 @@ mod tests {
tools: Vec::new(),
context_window: 0,
context_tokens: 0,
reasoning: None,
context_usage: None,
},
state: WorkerStatus::Idle.into(),
in_flight: Default::default(),
crates/worker/src/internal_worker.rs
+121 -40
View File
@@ -7,7 +7,7 @@
//! need durable domain audit must keep using their domain authority (for example Memory audit).
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::sync::{Arc, Mutex, RwLock};
use agen::timeline::event::UsageEvent;
use agen::{Engine, EngineError, llm_client::LlmClient};
@@ -17,7 +17,9 @@ use session_store::{LogEntry, SegmentId, SessionId, Store, StoreError, TraceEntr
use tokio::sync::broadcast;
use uuid::Uuid;
use crate::controller::{wire_event_bridges_on_engine, wire_workdir_command_events};
use crate::controller::{
build_greeting, wire_event_bridges_on_engine, wire_workdir_command_events,
};
use crate::feature::FeatureRegistryBuilder;
use crate::in_flight::{InFlightEvents, snapshot_from_guard};
use crate::ipc::alerter::Alerter;
@@ -409,6 +411,7 @@ enum InternalWorkerSessionCommand {
#[derive(Debug, Clone)]
pub(crate) struct InternalWorkerSessionSnapshot {
pub session: protocol::SessionSnapshot,
pub greeting: Option<protocol::Greeting>,
pub status: WorkerStatus,
pub error: Option<String>,
pub in_flight: InFlightSnapshot,
@@ -428,6 +431,7 @@ pub(crate) struct InternalWorkerSessionHandle {
visibility: InternalWorkerVisibility,
last_error: Arc<Mutex<Option<String>>>,
last_outcome: Arc<Mutex<Option<InternalWorkerTurnOutcome>>>,
greeting: Option<Arc<RwLock<protocol::Greeting>>>,
child_registry: Option<Arc<SpawnedWorkerRegistry>>,
sink: SegmentLogSink,
#[cfg(test)]
@@ -481,6 +485,12 @@ impl InternalWorkerSessionHandle {
};
InternalWorkerSessionSnapshot {
session: session_store::public_snapshot::project_current_session_snapshot(&entries),
greeting: self.greeting.as_ref().map(|greeting| {
greeting
.read()
.expect("internal Worker greeting lock poisoned")
.clone()
}),
status: match self.status() {
InternalWorkerSessionStatus::Running => WorkerStatus::Running,
InternalWorkerSessionStatus::Paused => WorkerStatus::Paused,
@@ -824,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,
@@ -851,8 +878,20 @@ pub(crate) async fn prepare_internal_worker_session(
worker.attach_alerter(alerter.clone());
worker.attach_working_event_tx(event_tx.clone());
worker.attach_in_flight_events(in_flight.clone());
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| {
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);
@@ -874,6 +913,7 @@ pub(crate) async fn prepare_internal_worker_session(
visibility,
last_error: last_error.clone(),
last_outcome: last_outcome.clone(),
greeting: Some(greeting),
child_registry,
sink,
#[cfg(test)]
@@ -886,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);
@@ -1174,6 +1220,14 @@ impl session_store::WorkerMetadataStore for EphemeralSessionStore {
#[cfg(test)]
pub(crate) fn test_internal_worker_session(
visibility: InternalWorkerVisibility,
) -> (InternalWorkerSessionHandle, broadcast::Sender<Event>) {
test_internal_worker_session_with_greeting(visibility, None)
}
#[cfg(test)]
pub(crate) fn test_internal_worker_session_with_greeting(
visibility: InternalWorkerVisibility,
greeting: Option<protocol::Greeting>,
) -> (InternalWorkerSessionHandle, broadcast::Sender<Event>) {
let store = EphemeralSessionStore::default();
let session_id = session_store::new_session_id();
@@ -1206,6 +1260,7 @@ pub(crate) fn test_internal_worker_session(
visibility,
last_error: Arc::new(Mutex::new(None)),
last_outcome: Arc::new(Mutex::new(None)),
greeting: greeting.map(|greeting| Arc::new(RwLock::new(greeting))),
child_registry: None,
sink,
fail_stop: Arc::new(std::sync::atomic::AtomicBool::new(false)),
@@ -1485,11 +1540,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);
crates/worker/src/ipc/protocol_session.rs
+46 -2
View File
@@ -16,9 +16,23 @@ pub struct WorkerProtocolSessionStreams {
pub events: broadcast::Receiver<Event>,
}
fn subscribe_before_snapshot<Subscription, Snapshot>(
subscribe: impl FnOnce() -> Subscription,
snapshot: impl FnOnce() -> Snapshot,
) -> (Snapshot, Subscription) {
// Subscription must exist before any snapshot field is read. A concurrent
// state update can then appear in both lanes, but it can never be absent
// from both the snapshot and the queued live stream.
let subscription = subscribe();
let snapshot = snapshot();
(snapshot, subscription)
}
pub fn subscribe_worker_protocol_session(handle: &WorkerHandle) -> WorkerProtocolSessionStreams {
let (snapshot_event, log_entries) = handle.snapshot_event_with_entry_subscription();
let (alert_snapshot, events) = handle.alerter.subscribe_with_snapshot();
let ((snapshot_event, log_entries), (alert_snapshot, events)) = subscribe_before_snapshot(
|| handle.alerter.subscribe_with_snapshot(),
|| handle.snapshot_event_with_entry_subscription(),
);
WorkerProtocolSessionStreams {
snapshot_event,
alert_snapshot,
@@ -129,6 +143,36 @@ mod tests {
}
}
#[test]
fn context_publication_between_subscription_and_snapshot_cannot_be_missed() {
let (events, _) = broadcast::channel(4);
let context_tokens = std::cell::Cell::new(20_u64);
let (snapshot_tokens, mut receiver) = subscribe_before_snapshot(
|| events.subscribe(),
|| {
context_tokens.set(25);
events
.send(Event::ContextUsage {
usage: Some(protocol::ContextUsage {
tokens: 25,
source: protocol::ContextTokenSource::Estimated,
}),
})
.unwrap();
context_tokens.get()
},
);
assert_eq!(snapshot_tokens, 25);
assert!(matches!(
receiver.try_recv(),
Ok(Event::ContextUsage {
usage: Some(protocol::ContextUsage { tokens: 25, .. })
})
));
}
#[test]
fn persisted_run_transitions_keep_identity_across_live_and_snapshot_projection() {
let temp = tempfile::tempdir().unwrap();
crates/worker/src/runtime/dir.rs
+2
View File
@@ -212,6 +212,8 @@ mod tests {
tools: Vec::new(),
context_window: 200_000,
context_tokens: 0,
reasoning: None,
context_usage: None,
},
)
}
crates/worker/src/shared_state.rs
+41 -2
View File
@@ -39,7 +39,7 @@ pub struct WorkerSharedState {
pub worker_name: String,
pub segment_id: SegmentId,
pub manifest_toml: String,
pub greeting: protocol::Greeting,
greeting: RwLock<protocol::Greeting>,
state: RwLock<WorkerStateSnapshot>,
accepted_commands: RwLock<VecDeque<AcceptedWorkerCommand>>,
/// Worker-from-the-inside view of the filesystem. Set once in
@@ -62,7 +62,7 @@ impl WorkerSharedState {
worker_name,
segment_id,
manifest_toml,
greeting,
greeting: RwLock::new(greeting),
state: RwLock::new(WorkerStateSnapshot::initial()),
accepted_commands: RwLock::new(VecDeque::new()),
fs_view: OnceLock::new(),
@@ -90,6 +90,22 @@ impl WorkerSharedState {
self.flow_transition_enabled.load(Ordering::Acquire)
}
pub fn greeting(&self) -> protocol::Greeting {
self.greeting
.read()
.expect("worker greeting lock poisoned")
.clone()
}
pub fn update_context_usage(&self, tokens: Option<u64>, source: protocol::ContextTokenSource) {
let mut greeting = self
.greeting
.write()
.expect("worker greeting lock poisoned");
greeting.context_tokens = tokens.unwrap_or_default();
greeting.context_usage = tokens.map(|tokens| protocol::ContextUsage { tokens, source });
}
pub fn transition(&self, state: WorkerState) -> WorkerStateSnapshot {
let mut snapshot = self
.state
@@ -245,9 +261,32 @@ mod tests {
tools: Vec::new(),
context_window: 200_000,
context_tokens: 0,
reasoning: None,
context_usage: None,
}
}
#[test]
fn current_context_usage_is_reflected_in_reconnect_greeting() {
let state = test_state();
state.update_context_usage(Some(42_000), protocol::ContextTokenSource::Measured);
let greeting = state.greeting();
assert_eq!(greeting.context_tokens, 42_000);
assert_eq!(
greeting.context_usage,
Some(protocol::ContextUsage {
tokens: 42_000,
source: protocol::ContextTokenSource::Measured,
})
);
state.update_context_usage(None, protocol::ContextTokenSource::Measured);
let greeting = state.greeting();
assert_eq!(greeting.context_tokens, 0);
assert_eq!(greeting.context_usage, None);
}
#[test]
fn initial_snapshot_is_idle() {
let state = test_state();
crates/worker/src/spawn/comm_tools.rs
+2
View File
@@ -97,6 +97,8 @@ mod tests {
tools: Vec::new(),
context_window: 200_000,
context_tokens: 0,
reasoning: None,
context_usage: None,
},
state: WorkerStatus::Idle.into(),
in_flight: Default::default(),
crates/worker/src/spawn/registry.rs
+157 -3
View File
@@ -29,7 +29,9 @@ use tokio::sync::{Notify, broadcast};
use tracing::warn;
use workdir::WorkdirScopeLeaseSet;
use crate::internal_worker::{InternalWorkerSessionHandle, InternalWorkerVisibility};
use crate::internal_worker::{
InternalWorkerSessionHandle, InternalWorkerSessionSnapshot, InternalWorkerVisibility,
};
use crate::runtime::dir::{RuntimeDir, SpawnedWorkerRecord};
use crate::runtime::worker_allocation;
@@ -648,6 +650,14 @@ impl SpawnedWorkerRegistry {
let protocol_emit_lock = record.protocol_emit_lock.clone();
let protocol_terminal = record.protocol_terminal.clone();
let mut child_rx = record.session.subscribe_events();
publish_initial_internal_snapshot(
&parent_tx,
&worker,
&protocol_revision,
&protocol_emit_lock,
&protocol_terminal,
&record.session,
);
tokio::spawn(async move {
loop {
match child_rx.recv().await {
@@ -710,6 +720,14 @@ impl SpawnedWorkerRegistry {
let protocol_emit_lock = record.protocol_emit_lock.clone();
let protocol_terminal = record.protocol_terminal.clone();
let mut child_rx = record.session.subscribe_events();
publish_initial_internal_snapshot(
&parent_tx,
&worker,
&protocol_revision,
&protocol_emit_lock,
&protocol_terminal,
&record.session,
);
tokio::spawn(async move {
loop {
match child_rx.recv().await {
@@ -985,6 +1003,41 @@ impl SpawnedWorkerRegistry {
}
}
fn publish_initial_internal_snapshot(
parent_tx: &broadcast::Sender<Event>,
worker: &InternalWorkerRef,
protocol_revision: &AtomicU64,
protocol_emit_lock: &Mutex<()>,
protocol_terminal: &AtomicBool,
session: &InternalWorkerSessionHandle,
) {
let Some(event) = internal_worker_live_snapshot(session.protocol_snapshot()) else {
return;
};
let _emit_guard = protocol_emit_lock
.lock()
.unwrap_or_else(|error| error.into_inner());
if protocol_terminal.load(Ordering::Acquire) {
return;
}
let revision = protocol_revision.fetch_add(1, Ordering::AcqRel) + 1;
let _ = parent_tx.send(Event::InternalWorker {
worker: worker.clone(),
revision,
event: Box::new(event),
});
}
fn internal_worker_live_snapshot(snapshot: InternalWorkerSessionSnapshot) -> Option<Event> {
Some(Event::Snapshot {
session: snapshot.session,
greeting: snapshot.greeting?,
state: snapshot.status.into(),
in_flight: snapshot.in_flight,
internal_workers: snapshot.internal_workers,
})
}
fn internal_worker_snapshot(
worker: InternalWorkerRef,
revision: u64,
@@ -995,6 +1048,7 @@ fn internal_worker_snapshot(
worker,
revision,
session: snapshot.session,
greeting: snapshot.greeting,
status: snapshot.status,
error: snapshot.error,
in_flight: snapshot.in_flight,
@@ -1121,7 +1175,10 @@ mod tests {
use session_store::LogEntry;
use super::*;
use crate::internal_worker::{InternalWorkerSessionStatus, test_internal_worker_session};
use crate::internal_worker::{
InternalWorkerSessionStatus, test_internal_worker_session,
test_internal_worker_session_with_greeting,
};
fn registry() -> Arc<SpawnedWorkerRegistry> {
let scope = Scope::from_config(&ScopeConfig {
@@ -1141,7 +1198,15 @@ mod tests {
name: &str,
visibility: InternalWorkerVisibility,
) -> (InternalSpawnedWorkerRecord, broadcast::Sender<Event>) {
let (session, sender) = test_internal_worker_session(visibility);
record_with_greeting(name, visibility, None).await
}
async fn record_with_greeting(
name: &str,
visibility: InternalWorkerVisibility,
greeting: Option<protocol::Greeting>,
) -> (InternalSpawnedWorkerRecord, broadcast::Sender<Event>) {
let (session, sender) = test_internal_worker_session_with_greeting(visibility, greeting);
let root = std::path::PathBuf::from("/tmp");
let scope = Scope::from_config(&ScopeConfig {
allow: vec![ScopeRule {
@@ -1193,6 +1258,95 @@ mod tests {
)
}
#[tokio::test]
async fn visible_internal_worker_starts_with_own_live_snapshot_metadata() {
let registry = registry();
let (parent_tx, mut parent_rx) = broadcast::channel(16);
registry.attach_parent_protocol(parent_tx, "parent-session".into());
let greeting = protocol::Greeting {
worker_name: "research".into(),
cwd: "/child".into(),
provider: "provider".into(),
model: "child-model".into(),
reasoning: Some(protocol::ReasoningConfig::Effort {
effort: "high".into(),
}),
scope_summary: "child scope".into(),
tools: Vec::new(),
context_window: 64_000,
context_tokens: 12_000,
context_usage: Some(protocol::ContextUsage {
tokens: 12_000,
source: protocol::ContextTokenSource::Measured,
}),
};
let (record, child_tx) = record_with_greeting(
"research",
InternalWorkerVisibility::ParentClient,
Some(greeting),
)
.await;
registry
.internal_records
.lock()
.unwrap()
.push(record.clone());
registry.start_protocol_forwarding(record);
let initial = tokio::time::timeout(Duration::from_secs(1), parent_rx.recv())
.await
.unwrap()
.unwrap();
assert!(matches!(
initial,
Event::InternalWorker {
revision: 1,
event,
..
} if matches!(
*event,
Event::Snapshot {
greeting: protocol::Greeting {
ref model,
context_window: 64_000,
..
},
..
} if model == "child-model"
)
));
child_tx
.send(Event::ContextUsage {
usage: Some(protocol::ContextUsage {
tokens: 14_000,
source: protocol::ContextTokenSource::Estimated,
}),
})
.unwrap();
let update = tokio::time::timeout(Duration::from_secs(1), parent_rx.recv())
.await
.unwrap()
.unwrap();
assert!(matches!(
update,
Event::InternalWorker {
revision: 2,
event,
..
} if matches!(
*event,
Event::ContextUsage {
usage: Some(protocol::ContextUsage {
tokens: 14_000,
source: protocol::ContextTokenSource::Estimated,
})
}
)
));
}
#[tokio::test]
async fn visible_internal_output_is_wrapped_after_registry_insertion() {
let registry = registry();
crates/workspace-server/src/server.rs
+2
View File
@@ -39652,10 +39652,12 @@ mod tests {
cwd: String::new(),
provider: "deterministic-workspace-server-test".to_string(),
model: "deterministic-workspace-server-test".to_string(),
reasoning: None,
scope_summary: "test execution snapshot".to_string(),
tools: Vec::new(),
context_window: 0,
context_tokens: 0,
context_usage: None,
},
state: protocol::WorkerStateSnapshot::initial(),
in_flight: protocol::InFlightSnapshot {