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

feat: expose worker model and context metadata

Browse Source
This commit is contained in:
Hare
2026-10-02 15:29:45 +09:00
parent 43e3b42055
commit cf3cf4b766
25 changed files with 834 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/controller.rs
+75 -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(),
@@ -303,6 +303,10 @@ 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(
Some(worker.total_tokens().tokens),
protocol::ContextTokenSource::Estimated,
);
set_controller_status(shared_state, runtime_dir, working_event_tx, new_status).await;
}
@@ -792,6 +796,11 @@ impl WorkerController {
manifest_toml.clone(),
greeting,
));
let usage_state = shared_state.clone();
worker.engine_mut().on_usage(move |event| {
usage_state
.update_context_usage(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));
}
@@ -2509,6 +2518,10 @@ async fn controller_loop<C, St>(
}
}
};
shared_state.update_context_usage(
Some(worker.total_tokens().tokens),
protocol::ContextTokenSource::Estimated,
);
if !matches!(
result,
Err(WorkerError::Store(_))
@@ -2576,6 +2589,10 @@ async fn controller_loop<C, St>(
.await
{
worker.clear_in_flight_events();
shared_state.update_context_usage(
Some(worker.total_tokens().tokens),
protocol::ContextTokenSource::Estimated,
);
let snapshot = shared_state.transition(WorkerState::Idle);
let _ = working_event_tx.send(Event::WorkerState { snapshot });
}
@@ -3383,7 +3400,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,
@@ -3426,6 +3460,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
@@ -3434,10 +3469,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,
}),
}
}
@@ -3500,6 +3544,28 @@ 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 automatic_memory_extraction_destinations_are_mutually_exclusive() {
let mut memory = manifest::MemoryFeatureProfileConfig::default();
@@ -3658,6 +3724,8 @@ mod tests {
tools: Vec::new(),
context_window: 200_000,
context_tokens: 0,
reasoning: None,
context_usage: None,
},
));
let notify_buffer = NotifyBuffer::new();
@@ -3707,6 +3775,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(),
@@ -4415,6 +4485,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
+26 -2
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,
@@ -851,7 +861,19 @@ 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();
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,
});
});
let session_id = worker.session_id();
let segment_id = worker.segment_id();
@@ -874,6 +896,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)]
@@ -1206,6 +1229,7 @@ pub(crate) fn test_internal_worker_session(
visibility,
last_error: Arc::new(Mutex::new(None)),
last_outcome: Arc::new(Mutex::new(None)),
greeting: None,
child_registry: None,
sink,
fail_stop: Arc::new(std::sync::atomic::AtomicBool::new(false)),
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
+1
View File
@@ -995,6 +995,7 @@ fn internal_worker_snapshot(
worker,
revision,
session: snapshot.session,
greeting: snapshot.greeting,
status: snapshot.status,
error: snapshot.error,
in_flight: snapshot.in_flight,