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

worker: make coder review routing durable

Browse Source
This commit is contained in:
Hare
2026-08-10 02:01:08 +09:00
parent b9dadb6a08
commit 64ced7dbad
18 changed files with 925 additions and 136 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/controller.rs
+47 -6
View File
@@ -5,7 +5,7 @@ use std::sync::atomic::Ordering;
use llm_engine::EngineError;
use llm_engine::llm_client::client::LlmClient;
use session_store::WorkerMetadataStore;
use session_store::{LogEntry, Store};
use session_store::{LogEntry, SessionExtension, Store};
use tokio::sync::{broadcast, mpsc, oneshot};
use crate::discovery::WorkerDiscovery;
@@ -24,7 +24,10 @@ use crate::shutdown_after_idle::{
use crate::spawn::comm_tools::{sub_worker_list_tool, sub_worker_send_tool, sub_worker_stop_tool};
use crate::spawn::registry::SpawnedWorkerRegistry;
use crate::spawn::tool::sub_worker_spawn_tool;
use crate::worker::{SystemItemCommitter, Worker, WorkerError, WorkerRunResult};
use crate::worker::{
SystemItemCommitter, WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN, Worker, WorkerError,
WorkerRunResult,
};
use protocol::{
AlertLevel, AlertSource, ErrorCode, Event, Method, RewindTargetId, RunResult, Segment,
TurnResult, WorkerStatus,
@@ -160,6 +163,10 @@ async fn finish_controller_run<C, St>(
/// `worker.run_for_notification()` drains the NotifyBuffer on its own.
enum PendingRun {
Run(Vec<Segment>),
RunTracked {
input: Vec<Segment>,
extension: SessionExtension,
},
/// Self-initiated turn kicked from the notify buffer. The carried
/// `InvokeKind` is the trigger that flipped the Worker from IDLE
/// (Notify or WorkerEvent) and is recorded by the Invoke marker
@@ -177,7 +184,7 @@ impl PendingRun {
/// notify buffer (Notify / inbound WorkerEvent) and stays silent.
fn is_parent_originated(&self) -> bool {
match self {
PendingRun::Run(_) | PendingRun::Resume => true,
PendingRun::Run(_) | PendingRun::RunTracked { .. } | PendingRun::Resume => true,
PendingRun::RunForNotification(_) => false,
}
}
@@ -340,6 +347,7 @@ impl WorkerController {
bash_output_dir,
runtime_base.to_path_buf(),
spawned_registry.clone(),
Some(method_tx.downgrade()),
)
.await?;
@@ -591,6 +599,7 @@ pub(crate) async fn register_worker_tools<C, St>(
bash_output_dir: PathBuf,
runtime_base: PathBuf,
spawned_registry: Arc<SpawnedWorkerRegistry>,
parent_method_tx: Option<mpsc::WeakSender<Method>>,
) -> std::io::Result<Option<workdir::WorkdirSessionHandle>>
where
C: LlmClient + Clone + 'static,
@@ -621,7 +630,11 @@ where
let spawner_name = worker.manifest().worker.name.clone();
let spawner_manifest = worker.manifest().clone();
let spawner_workspace_context = worker.workspace_context_handle();
let parent_notifies = worker.notify_buffer_handle();
let parent_notifications = parent_method_tx
.map(crate::spawn::tool::ParentNotificationTarget::Controller)
.unwrap_or_else(|| {
crate::spawn::tool::ParentNotificationTarget::Buffer(worker.notify_buffer_handle())
});
let prompts = worker.prompts().clone();
// Resolve the existing Worker–Workdir binding into the domain provider.
// Tools only consume the provider handle; they do not own its root, cwd,
@@ -810,7 +823,7 @@ where
engine.register_tool(sub_worker_spawn_tool(
spawner_name.clone(),
spawner_workspace_context,
parent_notifies,
parent_notifications,
runtime_base.clone(),
spawner_workspace_root,
spawner_cwd.clone(),
@@ -933,6 +946,21 @@ async fn controller_loop<C, St>(
)
.await
}
PendingRun::RunTracked { input, extension } => {
drive_turn(
worker.run_with_input_extensions(input, vec![extension]),
&mut method_rx,
&event_tx,
&cancel_tx,
&shared_state,
&notify_buffer,
self_parent_socket.as_ref(),
&spawner_name,
&spawned_registry,
parent_originated,
)
.await
}
PendingRun::RunForNotification(kind) => {
drive_turn(
worker.run_for_notification(kind),
@@ -1018,6 +1046,19 @@ async fn controller_loop<C, St>(
pending = Some(PendingRun::Run(input));
}
Method::RunTracked {
input,
submission_id,
} => {
// Runtime-correlated submissions retain their opaque id in the
// same durable UserInput record used for Flow state.
let extension = SessionExtension::new(
WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN,
serde_json::json!({ "submission_id": submission_id }),
);
pending = Some(PendingRun::RunTracked { input, extension });
}
Method::Notify { message, auto_run } => {
// Client-side live echo is delivered as `Event::SystemItem`
// once the interceptor commits the corresponding
@@ -1411,7 +1452,7 @@ where
shutdown_requested = true;
let _ = cancel_tx.try_send(());
}
Some(Method::Run { .. } | Method::Resume) => {
Some(Method::Run { .. } | Method::RunTracked { .. } | Method::Resume) => {
let _ = event_tx.send(Event::Error {
code: ErrorCode::AlreadyRunning,
message: "Worker is already executing a turn".into(),
crates/worker/src/feature/builtin/manage_worker.rs
+36 -3
View File
@@ -96,6 +96,9 @@ struct WorkerSpawnInput {
runtime_id: String,
working_directory_id: String,
profile: String,
/// Optional queued Ticket to assign atomically to the new Coder Worker.
#[serde(default)]
ticket_id: Option<String>,
#[serde(default)]
display_name: Option<String>,
/// Normal typed initial user submission delivered after spawn. An empty
@@ -106,11 +109,19 @@ struct WorkerSpawnInput {
relative_cwd: Option<String>,
}
#[derive(Debug, Serialize)]
struct WorkerSpawnTicketAssignmentRequest {
ticket_id: String,
operation_id: String,
}
#[derive(Debug, Serialize)]
struct WorkerSpawnRequest {
runtime_id: String,
display_name: String,
profile: String,
#[serde(skip_serializing_if = "Option::is_none")]
ticket_assignment: Option<WorkerSpawnTicketAssignmentRequest>,
initial_submit: Vec<Segment>,
working_directory: WorkerWorkingDirectorySelection,
}
@@ -170,7 +181,7 @@ impl WorkerOperation {
"List Backend/Runtime Worker sessions in the current Workspace. SubWorkers are excluded."
}
Self::Spawn => {
"Spawn a Backend/Runtime Worker session in an existing Workspace Workdir. The Workdir id is authority; filesystem paths and Runtime URLs are not accepted."
"Spawn a Backend/Runtime Worker session in an existing Workspace Workdir. The Workdir id is authority; filesystem paths and Runtime URLs are not accepted. `initial_submit` carries the normal typed user submission. Set `ticket_id` with a Flow segment in `initial_submit` to atomically assign a queued Ticket to the new Coder Worker; the operation id is derived from the durable tool call rather than model input."
}
Self::Stop => "Stop a Backend/Runtime Worker session in the current Workspace.",
Self::Restore => {
@@ -185,7 +196,7 @@ impl Tool for WorkspaceWorkerTool {
async fn execute(
&self,
input_json: &str,
_ctx: ToolExecutionContext,
ctx: ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
let request = match self.operation {
WorkerOperation::List => {
@@ -194,6 +205,17 @@ impl Tool for WorkspaceWorkerTool {
}
WorkerOperation::Spawn => {
let input = parse::<WorkerSpawnInput>(input_json, "WorkerSpawn")?;
let ticket_assignment = input
.ticket_id
.map(|ticket_id| {
let ticket_id = authority_id(&ticket_id, "ticket_id")?;
let call_id = non_empty(ctx.call_id.clone(), "tool call_id")?;
Ok::<_, ToolError>(WorkerSpawnTicketAssignmentRequest {
operation_id: format!("worker-spawn:{ticket_id}:{call_id}"),
ticket_id,
})
})
.transpose()?;
let request = WorkerSpawnRequest {
runtime_id: authority_id(&input.runtime_id, "runtime_id")?,
display_name: input
@@ -201,6 +223,7 @@ impl Tool for WorkspaceWorkerTool {
.filter(|value| !value.trim().is_empty())
.unwrap_or_else(|| "Workspace Worker".to_string()),
profile: non_empty(input.profile, "profile")?,
ticket_assignment,
initial_submit: input.initial_submit,
working_directory: WorkerWorkingDirectorySelection {
working_directory_id: authority_id(
@@ -372,13 +395,14 @@ mod tests {
"runtime_id": "runtime-1",
"working_directory_id": "workdir-1",
"profile": "builtin:coder",
"ticket_id": "00001KZ9E0DBS",
"initial_submit": [
{ "kind": "flow", "selector": "builtin:coder-review" },
{ "kind": "text", "content": "Implement Ticket 00001" }
]
})
.to_string(),
ToolExecutionContext::direct(),
ToolExecutionContext::new("call-1", "batch-1", 0),
)
.await
.unwrap();
@@ -394,6 +418,13 @@ mod tests {
"builtin:coder-review"
);
assert_eq!(body["initial_submit"][1]["kind"], "text");
assert_eq!(
body["ticket_assignment"],
serde_json::json!({
"ticket_id": "00001KZ9E0DBS",
"operation_id": "worker-spawn:00001KZ9E0DBS:call-1"
})
);
assert!(body.get("initial_text").is_none());
}
@@ -410,6 +441,7 @@ mod tests {
let schema = serde_json::to_value(schemars::schema_for!(WorkerSpawnInput)).unwrap();
let text = serde_json::to_string(&schema).unwrap();
assert!(text.contains("initial_submit"));
assert!(text.contains("ticket_id"));
assert!(text.contains("selector"));
assert!(text.contains("flow"));
assert!(!text.contains("initial_text"));
@@ -421,6 +453,7 @@ mod tests {
runtime_id: "runtime-1".to_string(),
display_name: "Coder".to_string(),
profile: "builtin:coder".to_string(),
ticket_assignment: None,
initial_submit: vec![
Segment::Flow {
selector: "builtin:coder-review".to_string(),
crates/worker/src/lib.rs
+4 -4
View File
@@ -40,9 +40,9 @@ pub use runtime::dir::RuntimeDir;
pub use segment_log_sink::SegmentLogSink;
pub use shared_state::WorkerSharedState;
pub use worker::{
LocalWorkingDirectory, RuntimeWorkspaceHttpClient, Worker, WorkerError,
WorkerFilesystemAuthority, WorkerRunResult, WorkerWorkspaceContext, WorkspaceClient,
WorkspaceClientError, WorkspaceId, WorkspaceIdError, WorkspaceRequest, WorkspaceRequestMethod,
WorkspaceResponse, apply_worker_manifest, marker_workspace_client,
LocalWorkingDirectory, RuntimeWorkspaceHttpClient, WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN,
Worker, WorkerError, WorkerFilesystemAuthority, WorkerRunResult, WorkerWorkspaceContext,
WorkspaceClient, WorkspaceClientError, WorkspaceId, WorkspaceIdError, WorkspaceRequest,
WorkspaceRequestMethod, WorkspaceResponse, apply_worker_manifest, marker_workspace_client,
unavailable_workspace_client,
};
crates/worker/src/spawn/tool.rs
+78 -21
View File
@@ -18,6 +18,7 @@ use manifest::{
WorkerManifest, WorkerManifestConfig, WorkerMetaConfig,
};
use serde::Deserialize;
use tokio::sync::mpsc;
use crate::PromptLoader;
use crate::controller::register_worker_tools;
@@ -27,6 +28,7 @@ use crate::internal_worker::{
use crate::prompt::catalog::PromptCatalog;
use crate::spawn::registry::SpawnedWorkerRegistry;
use crate::worker::{Worker, WorkerFilesystemAuthority};
use protocol::Method;
#[derive(Debug, Deserialize, schemars::JsonSchema)]
struct SubWorkerSpawnInput {
@@ -205,6 +207,39 @@ fn parse_spawn_profile_selector(raw: Option<&str>) -> Result<SpawnProfileSelecto
Ok(SpawnProfileSelector::Registry(ProfileSelector::named(raw)))
}
#[derive(Clone)]
pub(crate) enum ParentNotificationTarget {
Controller(mpsc::WeakSender<Method>),
Buffer(crate::ipc::notify_buffer::NotifyBuffer),
}
impl ParentNotificationTarget {
fn notify(&self, message: String, auto_run: bool) {
match self {
Self::Controller(parent_method_tx) => {
let Some(parent_method_tx) = parent_method_tx.upgrade() else {
tracing::warn!(
"parent Worker controller closed before Internal SubWorker completion notification"
);
return;
};
tokio::spawn(async move {
if let Err(error) = parent_method_tx
.send(Method::Notify { message, auto_run })
.await
{
tracing::warn!(
%error,
"failed to notify parent Worker about Internal SubWorker completion"
);
}
});
}
Self::Buffer(parent_notifies) => parent_notifies.push_notify(message, auto_run),
}
}
}
/// Runtime dependencies the `SubWorkerSpawn` tool needs in order to launch a
/// child SubWorker and record the handoff locally. Constructed by the Worker
/// controller once per Worker lifetime.
@@ -212,7 +247,7 @@ pub struct SubWorkerSpawnTool {
/// Spawner's own Worker name, used for direct-child identity collision checks.
spawner_name: String,
workspace_context: crate::worker::WorkerWorkspaceContext,
parent_notifies: crate::ipc::notify_buffer::NotifyBuffer,
parent_notifications: ParentNotificationTarget,
/// Runtime-owned root used only for bounded Internal Worker tool artifacts such as Bash spill
/// output. It is not an Internal Worker identity or catalog location.
runtime_base: PathBuf,
@@ -256,7 +291,7 @@ impl SubWorkerSpawnTool {
fn new(
spawner_name: String,
workspace_context: crate::worker::WorkerWorkspaceContext,
parent_notifies: crate::ipc::notify_buffer::NotifyBuffer,
parent_notifications: ParentNotificationTarget,
runtime_base: PathBuf,
workspace_root: PathBuf,
spawner_cwd: PathBuf,
@@ -270,7 +305,7 @@ impl SubWorkerSpawnTool {
Self {
spawner_name,
workspace_context,
parent_notifies,
parent_notifications,
runtime_base,
workspace_root,
spawner_cwd,
@@ -368,6 +403,7 @@ impl Tool for SubWorkerSpawnTool {
.join("bash-output"),
self.runtime_base.clone(),
child_registry,
None,
)
.await
.map_err(|error| {
@@ -392,7 +428,7 @@ impl Tool for SubWorkerSpawnTool {
let child_name = input.name.clone();
let registry = Arc::downgrade(&self.registry);
let parent_notifies = self.parent_notifies.clone();
let parent_notifications = self.parent_notifications.clone();
let session_result = prepare_internal_worker_session(
child,
store,
@@ -408,10 +444,10 @@ impl Tool for SubWorkerSpawnTool {
}
}
}
parent_notifies.push_notify(
format!("SubWorker `{child_name}` turn ended with status {status:?}. Inspect its committed session with worker-observation tools before making completion decisions."),
true,
let message = format!(
"SubWorker `{child_name}` turn ended with status {status:?}. Inspect its committed session with worker-observation tools before making completion decisions."
);
parent_notifications.notify(message, true);
})),
)
.await;
@@ -764,10 +800,10 @@ fn manifest_to_reusable_config(manifest: &WorkerManifest) -> WorkerManifestConfi
/// tool-result budget — debugging beyond this should read the file
/// directly.
/// Factory for the `SubWorkerSpawn` tool.
pub fn sub_worker_spawn_tool(
pub(crate) fn sub_worker_spawn_tool(
spawner_name: String,
workspace_context: crate::worker::WorkerWorkspaceContext,
parent_notifies: crate::ipc::notify_buffer::NotifyBuffer,
parent_notifications: ParentNotificationTarget,
runtime_base: PathBuf,
workspace_root: PathBuf,
spawner_cwd: PathBuf,
@@ -779,7 +815,7 @@ pub fn sub_worker_spawn_tool(
sub_worker_spawn_tool_impl(
spawner_name,
workspace_context,
parent_notifies,
parent_notifications,
runtime_base,
workspace_root,
spawner_cwd,
@@ -793,7 +829,7 @@ pub fn sub_worker_spawn_tool(
fn sub_worker_spawn_tool_impl(
spawner_name: String,
workspace_context: crate::worker::WorkerWorkspaceContext,
parent_notifies: crate::ipc::notify_buffer::NotifyBuffer,
parent_notifications: ParentNotificationTarget,
runtime_base: PathBuf,
workspace_root: PathBuf,
spawner_cwd: PathBuf,
@@ -824,7 +860,7 @@ fn sub_worker_spawn_tool_impl(
let tool: Arc<dyn Tool> = Arc::new(SubWorkerSpawnTool::new(
spawner_name.clone(),
workspace_context.clone(),
parent_notifies.clone(),
parent_notifications.clone(),
runtime_base.clone(),
workspace_root.clone(),
spawner_cwd.clone(),
@@ -845,6 +881,7 @@ mod tests {
use super::*;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::time::Duration;
use crate::WorkspaceId;
use async_trait::async_trait;
@@ -896,7 +933,18 @@ extract_threshold = 4000
"#;
#[tokio::test]
async fn reviewer_profile_spawns_as_workspace_aware_internal_session() {
async fn parent_controller_notification_target_does_not_keep_channel_open() {
let (parent_method_tx, mut parent_method_rx) = mpsc::channel(1);
let target = ParentNotificationTarget::Controller(parent_method_tx.downgrade());
drop(parent_method_tx);
assert!(parent_method_rx.recv().await.is_none());
target.notify("late completion".to_string(), true);
}
#[tokio::test]
async fn reviewer_profile_spawns_and_notifies_parent_controller() {
let runtime = TempDir::new().unwrap();
let workspace_root = runtime.path().join("project");
let available_profiles = write_project_profile_registry(
@@ -927,11 +975,11 @@ extract_threshold = 4000
)
.unwrap();
let prompt_loader = PromptLoader::new(None, Some(workspace_prompts));
let parent_notifies = crate::ipc::notify_buffer::NotifyBuffer::new();
let (parent_method_tx, mut parent_method_rx) = mpsc::channel(8);
let tool = SubWorkerSpawnTool::new(
"parent".into(),
workspace_context,
parent_notifies.clone(),
ParentNotificationTarget::Controller(parent_method_tx.downgrade()),
runtime.path().to_path_buf(),
workspace_root.clone(),
workspace_root.clone(),
@@ -995,11 +1043,17 @@ extract_threshold = 4000
assert_eq!(calls.load(Ordering::SeqCst), 1);
assert!(observed_parent_write_revoked.load(Ordering::SeqCst));
assert!(observed_instruction_override.load(Ordering::SeqCst));
assert_eq!(parent_notifies.len(), 1);
assert!(
parent_notifies.has_auto_run_pending(),
"SubWorker completion must auto-invoke the parent"
);
let completion = tokio::time::timeout(Duration::from_secs(1), parent_method_rx.recv())
.await
.expect("SubWorker completion must wake the parent method channel")
.expect("parent method channel remains open");
assert!(matches!(
completion,
Method::Notify {
message,
auto_run: true,
} if message.contains("SubWorker `reviewer-child` turn ended with status Idle")
));
assert!(!runtime.path().join("reviewer-child/sock").exists());
let duplicate_error = tool
@@ -1018,7 +1072,10 @@ extract_threshold = 4000
1,
"duplicate rejection must not invoke the child provider"
);
assert_eq!(parent_notifies.len(), 1);
assert!(matches!(
parent_method_rx.try_recv(),
Err(tokio::sync::mpsc::error::TryRecvError::Empty)
));
let context = llm_engine::tool::ToolExecutionContext::direct();
let list = (crate::spawn::comm_tools::sub_worker_list_tool(registry.clone()))().1;
crates/worker/src/worker.rs
+19 -12
View File
@@ -650,6 +650,8 @@ where
}
}
pub const WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN: &str = "worker.input-submission.v1";
/// An independent agent execution unit.
///
/// Holds a [`Engine`] directly and persists session state via
@@ -2182,19 +2184,24 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
/// the Engine is aborted, history is compacted, and execution resumes
/// automatically.
pub async fn run(&mut self, input: Vec<Segment>) -> Result<WorkerRunResult, WorkerError> {
self.run_with_input_extensions(input, Vec::new()).await
}
pub(crate) async fn run_with_input_extensions(
&mut self,
input: Vec<Segment>,
mut input_extensions: Vec<SessionExtension>,
) -> Result<WorkerRunResult, WorkerError> {
let (input, pending_flow_state) = self.prepare_flow_input(input)?;
let input_extensions = pending_flow_state
.as_ref()
.map(|state| {
serde_json::to_value(state)
.map(|payload| SessionExtension::new(FLOW_RUNTIME_EXTENSION_DOMAIN, payload))
.map_err(|error| {
WorkerError::FlowInput(format!("serialize Flow runtime state: {error}"))
})
})
.transpose()?
.into_iter()
.collect();
if let Some(state) = pending_flow_state.as_ref() {
let payload = serde_json::to_value(state).map_err(|error| {
WorkerError::FlowInput(format!("serialize Flow runtime state: {error}"))
})?;
input_extensions.push(SessionExtension::new(
FLOW_RUNTIME_EXTENSION_DOMAIN,
payload,
));
}
// Paused→Run transition: if the previous turn was cut short,
// any `Item::ToolCall` whose tool never produced a matching