bingshuo/grok-build-upstream-mirror
Watch
1
0
Fork
You've already forked grok-build-upstream-mirror
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Synced from monorepo

Browse source 
Synced from monorepo

Changes:
- Workspace server: report `/ready` as failed with dwell on hub connect failure
- Refresh OIDC token for the Grok agent in the shell
- ACP terminal output recorder
- Cross-platform provider auth commands in the shell
- Default `/resume` to Grok sessions with a hint for hidden external sessions
- Resume sessions by title with `--resume`
- Limit app-builder archive size
- Data-driven tag labels for slash commands
- Doctor fixes for tmux
- Custom provider gateways and subprocess environment policy in the shell
- `/tutorial` — opt-in onboarding tour of Grok Build
- Soft and required CLI version checks in the shell
- Privacy banner env overrides survive live settings updates
- Add remote flag to override the image-edit model
- Return profile fields from auth info even when the access token is expired
- Add edit control on queued prompt rows
- Keep fail-closed policy when clearing orphans with no team
- Setting to disable the Ctrl+Space/F8 voice shortcut
- Pass `--raw` to pw-record so Linux dictation works on older PipeWire
- Validate git URLs when adding marketplace entries
- Stop shipping stale tool-doc parameter and tool names
- Re-point dashboard attach after `/fork` only when the parent was attached
- Surface Grok Computer media-generation results as file-path chunks
- Clear web background-task tray on kill and keep the task description
- Show privacy upsell banner in agent view until acted on
- Add tools-server client callback surface
- Protect persistent global hook sources

Source-Revision: 95d84f443eddcbed6cbfd6eed22e2eafe6b3939d
This commit is contained in:
grokkybara[bot] 2026-07-23 17:12:33 +00:00
commit 69f0ba880a
286 changed files with 22939 additions and 9624 deletions
Download patch file Download diff file Expand all files Collapse all files

8
crates/codegen/xai-grok-tools/src/implementations/grok_build/bash/mod.rs
View file

@ -2037,6 +2037,7 @@ impl xai_tool_runtime::Tool for BashTool {
foreground_block_budget: None,
kind: crate::computer::types::TaskKind::Bash,
owner_session_id: owner_session_id.clone(),
description: Some(input.description.clone()).filter(|d| !d.trim().is_empty()),
};
let handle = match backend.run_background(request).await {
@ -2072,7 +2073,7 @@ impl xai_tool_runtime::Tool for BashTool {
output_file: bg_output_file.clone(),
task_id: task_id.clone(),
monitor_description: None,
description: Some(input.description.clone()),
description: Some(input.description.clone()).filter(|d| !d.trim().is_empty()),
});
let retrieval_hint = Self::background_retrieval_hint(&resources, &task_id).await?;
@ -2133,6 +2134,7 @@ impl xai_tool_runtime::Tool for BashTool {
foreground_block_budget: Self::effective_foreground_block_budget(&params),
kind: crate::computer::types::TaskKind::Bash,
owner_session_id: owner_session_id.clone(),
description: Some(input.description.clone()).filter(|d| !d.trim().is_empty()),
};
let result = match backend.run(request).await {
@ -2166,7 +2168,7 @@ impl xai_tool_runtime::Tool for BashTool {
output_file: output_file.clone(),
task_id: tool_call_id.as_str().to_owned(),
monitor_description: None,
description: Some(input.description.clone()),
description: Some(input.description.clone()).filter(|d| !d.trim().is_empty()),
});
let retrieval_hint =
@ -2233,7 +2235,7 @@ impl xai_tool_runtime::Tool for BashTool {
truncated: result.truncated,
signal: result.signal,
timed_out: result.timed_out,
description: Some(input.description),
description: Some(input.description).filter(|d| !d.trim().is_empty()),
current_dir: cwd.to_string_lossy().to_string(),
output_file: output_file.to_string_lossy().to_string(),
total_bytes: result.total_bytes,

4
crates/codegen/xai-grok-tools/src/implementations/grok_build/image_edit/mod.rs
View file

@ -26,7 +26,7 @@ use crate::types::resources::SessionFolder;
use crate::types::tool::{ToolKind, ToolNamespace};
use crate::util::image_compress::{FilterType, ReEncodeParams, re_encode_under_limit};
const XAI_IMAGINE_MODEL: &str = "grok-imagine-image-quality";
pub(crate) const XAI_IMAGINE_EDIT_MODEL: &str = "grok-imagine-image-quality";
/// Size/dimension limits for reference images sent to the Imagine API.
/// Tighter than the vision path; the backend returns 400 when exceeded.
@ -353,7 +353,7 @@ impl xai_tool_runtime::Tool for ImageEditTool {
let url = format!("{base}/images/edits");
let mut payload = serde_json::json!({
"model": XAI_IMAGINE_MODEL,
"model": client.edit_model(),
"prompt": input.prompt,
"n": 1,
"resolution": "1k",

42
crates/codegen/xai-grok-tools/src/implementations/grok_build/image_gen/mod.rs
View file

@ -58,6 +58,7 @@ pub struct ImageGenClient {
/// [`XAI_IMAGINE_MODEL`]). `image_edit` uses its own model and is
/// unaffected.
model: String,
edit_model: String,
writer: super::storage::SessionFileWriter,
api_key_provider: Option<SharedApiKeyProvider>,
/// Optional 401-attribution hook. Hosts wire this so a 401 from the
@ -81,6 +82,7 @@ impl ImageGenClient {
base_url,
extra_headers,
model_override,
edit_model_override,
tier_restricted,
..
} = config
@ -93,6 +95,10 @@ impl ImageGenClient {
.clone()
.filter(|m| !m.trim().is_empty())
.unwrap_or_else(|| XAI_IMAGINE_MODEL.to_owned());
let edit_model = edit_model_override
.clone()
.filter(|m| !m.trim().is_empty())
.unwrap_or_else(|| super::image_edit::XAI_IMAGINE_EDIT_MODEL.to_owned());
let mut headers = reqwest::header::HeaderMap::new();
headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
@ -138,6 +144,7 @@ impl ImageGenClient {
http,
base_url: base_url.clone(),
model,
edit_model,
writer: super::storage::SessionFileWriter::new(DEFAULT_IMAGE_DIR, "jpg"),
api_key_provider,
attribution_callback: None,
@ -183,6 +190,10 @@ impl ImageGenClient {
&self.writer
}
pub(crate) fn edit_model(&self) -> &str {
&self.edit_model
}
pub async fn generate(
&self,
prompt: &str,
@ -277,6 +288,7 @@ pub enum ImageGenConfig {
/// ([`XAI_IMAGINE_MODEL`]). Driven by the remote
/// `image_gen_model_override` config flag. `image_edit` is unaffected.
model_override: Option<String>,
edit_model_override: Option<String>,
/// `true` when the user is on a tier the Imagine server zero-limits
/// (free / X Basic). The tools stay advertised to the model, but
/// `image_gen` / `image_edit` short-circuit at call time with the
@ -483,6 +495,7 @@ mod tests {
image_gen_enabled: false,
image_edit_enabled: true,
model_override: Some("grok-imagine-image".into()),
edit_model_override: None,
tier_restricted: false,
};
assert!(cfg.has_credentials());
@ -502,6 +515,7 @@ mod tests {
image_gen_enabled: true,
image_edit_enabled: true,
model_override: model_override.map(String::from),
edit_model_override: None,
tier_restricted: false,
};
// No override → default quality model.
@ -523,6 +537,33 @@ mod tests {
);
}
#[test]
fn client_selects_edit_model_from_override() {
let mk = |edit_model_override: Option<&str>| ImageGenConfig::Enabled {
api_key: "k".into(),
base_url: "https://api.x.ai/v1".into(),
extra_headers: indexmap::IndexMap::new(),
image_gen_enabled: true,
image_edit_enabled: true,
model_override: None,
edit_model_override: edit_model_override.map(String::from),
tier_restricted: false,
};
assert_eq!(
ImageGenClient::new(&mk(None), None).unwrap().edit_model(),
super::super::image_edit::XAI_IMAGINE_EDIT_MODEL
);
assert_eq!(
ImageGenClient::new(&mk(Some(" ")), None)
.unwrap()
.edit_model(),
super::super::image_edit::XAI_IMAGINE_EDIT_MODEL
);
let client = ImageGenClient::new(&mk(Some("grok-imagine-image-v2")), None).unwrap();
assert_eq!(client.edit_model(), "grok-imagine-image-v2");
assert_eq!(client.model, XAI_IMAGINE_MODEL);
}
#[tokio::test]
async fn errors_when_client_missing() {
let tool = ImageGenTool;
@ -558,6 +599,7 @@ mod tests {
image_gen_enabled: true,
image_edit_enabled: true,
model_override: None,
edit_model_override: None,
tier_restricted: true,
};
let mut resources = crate::types::resources::Resources::new();

15
crates/codegen/xai-grok-tools/src/implementations/grok_build/monitor/tool.rs
View file

@ -83,7 +83,7 @@ impl xai_tool_runtime::Tool for MonitorTool {
.map_err(|e| xai_tool_runtime::ToolError::invalid_arguments(e.to_string()))?;
let resolved_timeout = input.resolved_timeout_ms();
let description = input.description.clone();
let description = input.description;
let (terminal, notification_handle, cwd, session_folder, owner_session_id) = {
let res = resources.lock().await;
@ -127,16 +127,18 @@ impl xai_tool_runtime::Tool for MonitorTool {
output_file,
notification_handle: notification_handle.clone(),
tool_call_id: ctx.call_id.as_str().to_owned(),
display_command: Some(format!("[monitor] {}", input.description)),
display_command: Some(format!("[monitor] {description}")),
auto_background_on_timeout: false,
foreground_block_budget: None,
kind: crate::computer::types::TaskKind::Monitor,
owner_session_id,
description: Some(description.clone()).filter(|d| !d.trim().is_empty()),
})
.await
.map_err(|e| xai_tool_runtime::ToolError::custom("process_manager", e.to_string()))?;
let task_id = bg_handle.task_id.clone();
let tray_description = Some(description.clone()).filter(|d| !d.trim().is_empty());
// Notify the pager so the monitor appears in the tasks pane
// (same notification that bash background tasks send).
@ -155,15 +157,15 @@ impl xai_tool_runtime::Tool for MonitorTool {
},
output_file: bg_handle.output_file.clone(),
task_id: task_id.clone(),
monitor_description: Some(input.description.clone()),
description: None,
monitor_description: tray_description.clone(),
description: tray_description,
});
// Spawn the stdout processing pipeline.
// Reads the output file, processes lines through the rate limiter,
// and emits MonitorEvent notifications.
let pipeline_task_id = task_id.clone();
let pipeline_description = description.clone();
let pipeline_description = description;
// Weak handle: the pipeline must not keep the session's terminal backend
// (and the monitored process) alive past session end. See
// `run_monitor_pipeline`.
@ -449,6 +451,7 @@ mod tests {
foreground_block_budget: None,
kind: TaskKind::Monitor,
owner_session_id: Some("session-A".to_string()),
description: None,
})
.await
.expect("spawn monitor");
@ -525,6 +528,7 @@ mod tests {
foreground_block_budget: None,
kind: TaskKind::Monitor,
owner_session_id: Some("session-A".to_string()),
description: None,
})
.await
.expect("spawn monitor");
@ -594,6 +598,7 @@ mod tests {
foreground_block_budget: None,
kind: TaskKind::Monitor,
owner_session_id: Some("child-session".to_string()),
description: None,
})
.await
.expect("spawn monitor");

11
crates/codegen/xai-grok-tools/src/implementations/grok_build/scheduler/actor.rs
View file

@ -8,7 +8,7 @@ use tokio_util::sync::CancellationToken;
use crate::implementations::grok_build::task::types::{
SessionIdResource, SubagentEvent, SubagentEventSender, SubagentLoopUnitActiveRequest,
SubagentOwner, SubagentQueryRequest, SubagentRequest, SubagentRuntimeOverrides,
SubagentSnapshotStatus,
SubagentSnapshotStatus, SubagentSpawnRequest,
};
use crate::notification::types::ToolNotificationHandle;
use crate::notification::{
@ -526,6 +526,7 @@ impl SchedulerActor {
.0
.send(SubagentEvent::Query(SubagentQueryRequest {
subagent_id: prev_id.clone(),
parent_session_id: Some(parent_session_id.clone()),
block: false,
timeout_ms: None,
respond_to,
@ -675,12 +676,14 @@ impl SchedulerActor {
fork_context: false,
owner: SubagentOwner::Task,
cancel_token: CancellationToken::new(),
result_tx,
};
if events
.0
.send(SubagentEvent::Spawn(Box::new(request)))
.send(SubagentEvent::Spawn(SubagentSpawnRequest {
request: Box::new(request),
result_tx,
}))
.is_err()
{
let mut res = self.resources.lock().await;
@ -1818,7 +1821,7 @@ mod tests {
let SubagentEvent::Spawn(spawn) = next_event(rx).await else {
panic!("expected subagent spawn");
};
spawn
spawn.request
}
async fn answer_loop_unit_active(

783
crates/codegen/xai-grok-tools/src/implementations/grok_build/task/backend.rs
View file

@ -4,21 +4,21 @@
//! `TaskOutputTool`, `KillTaskTool`) from the transport mechanism used to
//! communicate with the subagent coordinator.
//!
//! Two implementations are planned:
//!
//! - [`ChannelBackend`] — wraps in-process `tokio::mpsc` channels used by
//! the local host shell. This is the only implementation today.
//! - `RemoteBackend` (future) — dispatches over a remote transport to an
//! out-of-process spawner.
//! All hosts use [`ChannelBackend`].
//! The receiver is owned by the shared single-writer coordinator actor; only
//! the child runner plugged into that actor differs by host.
use std::sync::Arc;
use tokio::sync::{mpsc, oneshot};
use super::types::{
SubagentCancelOutcome, SubagentCancelRequest, SubagentCancelTarget, SubagentDescribeOutcome,
SubagentDescribeRequest, SubagentEvent, SubagentQueryRequest, SubagentRequest, SubagentResult,
SubagentSnapshot, SubagentValidateTypeOutcome, SubagentValidateTypeRequest,
SpawnedSubagentRef, SubagentCancelOutcome, SubagentCancelRequest, SubagentCancelTarget,
SubagentDescribeOutcome, SubagentDescribeRequest, SubagentEvent, SubagentInspectRequest,
SubagentInspection, SubagentListRunningRequest, SubagentQueryRequest, SubagentRegistryCounts,
SubagentRegistryCountsRequest, SubagentRequest, SubagentResult, SubagentSnapshot,
SubagentSpawnRequest, SubagentSpawnedRefsRequest, SubagentValidateTypeOutcome,
SubagentValidateTypeRequest,
};
use crate::register_resource;
use xai_tool_runtime::ToolError;
@ -110,13 +110,132 @@ register_resource!(
/// Wraps a single `mpsc::UnboundedSender<SubagentEvent>` that carries
/// spawn, query, and cancel messages to the coordinator. The oneshot for
/// `spawn` is created inside the backend so callers never manage it.
#[derive(Clone)]
pub struct ChannelBackend {
tx: mpsc::UnboundedSender<SubagentEvent>,
parent_session_id: Option<Arc<str>>,
}
impl ChannelBackend {
pub fn new(tx: mpsc::UnboundedSender<SubagentEvent>) -> Self {
Self { tx }
Self {
tx,
parent_session_id: None,
}
}
/// Bind model-facing operations to one parent session.
pub fn for_session(
tx: mpsc::UnboundedSender<SubagentEvent>,
parent_session_id: impl Into<Arc<str>>,
) -> Self {
Self {
tx,
parent_session_id: Some(parent_session_id.into()),
}
}
fn parent_session_id(&self) -> Option<String> {
self.parent_session_id.as_deref().map(str::to_owned)
}
pub fn sender(&self) -> mpsc::UnboundedSender<SubagentEvent> {
self.tx.clone()
}
pub fn into_resource(self) -> SubagentBackendResource {
SubagentBackendResource(Arc::new(self))
}
pub async fn cancel_parent_prompt(&self, parent_prompt_id: &str) -> SubagentCancelOutcome {
let (respond_to, response_rx) = oneshot::channel();
if self
.tx
.send(SubagentEvent::Cancel(SubagentCancelRequest {
parent_session_id: self.parent_session_id(),
target: SubagentCancelTarget::ParentPromptId(parent_prompt_id.to_owned()),
respond_to,
}))
.is_err()
{
return SubagentCancelOutcome::NotFound;
}
response_rx.await.unwrap_or(SubagentCancelOutcome::NotFound)
}
pub async fn inspect(&self, id: &str) -> Option<SubagentInspection> {
let (respond_to, response_rx) = oneshot::channel();
self.tx
.send(SubagentEvent::Inspect(SubagentInspectRequest {
subagent_id: id.to_owned(),
parent_session_id: self.parent_session_id(),
respond_to,
}))
.ok()?;
response_rx.await.ok().flatten()
}
pub async fn list_running(&self, parent_session_id: &str) -> Vec<SubagentInspection> {
let (respond_to, response_rx) = oneshot::channel();
if self
.tx
.send(SubagentEvent::ListRunning(SubagentListRunningRequest {
parent_session_id: parent_session_id.to_owned(),
respond_to,
}))
.is_err()
{
return Vec::new();
}
response_rx.await.unwrap_or_default()
}
pub async fn spawned_refs_for_prompt(
&self,
parent_session_id: &str,
prompt_id: &str,
) -> Vec<SpawnedSubagentRef> {
let (respond_to, response_rx) = oneshot::channel();
if self
.tx
.send(SubagentEvent::SpawnedRefs(SubagentSpawnedRefsRequest {
parent_session_id: self
.parent_session_id
.as_deref()
.unwrap_or(parent_session_id)
.to_owned(),
prompt_id: prompt_id.to_owned(),
respond_to,
}))
.is_err()
{
return Vec::new();
}
response_rx.await.unwrap_or_default()
}
pub async fn registry_counts(&self) -> SubagentRegistryCounts {
let (respond_to, response_rx) = oneshot::channel();
if self
.tx
.send(SubagentEvent::RegistryCounts(
SubagentRegistryCountsRequest { respond_to },
))
.is_err()
{
return SubagentRegistryCounts::default();
}
response_rx.await.unwrap_or_default()
}
/// Spawn while holding the host's interruptible foreground-wait token.
pub async fn spawn_with_foreground_wait(
&self,
request: SubagentRequest,
wait: Option<&super::types::SubagentForegroundWait>,
) -> Result<SubagentResult, ToolError> {
let _wait = wait.map(super::types::SubagentForegroundWait::enter);
self.spawn(request).await
}
}
@ -135,20 +254,18 @@ impl Drop for CancelResultReceiverOnDrop {
#[async_trait::async_trait]
impl SubagentBackend for ChannelBackend {
async fn spawn(&self, request: SubagentRequest) -> Result<SubagentResult, ToolError> {
let (result_tx, result_rx) = oneshot::channel();
async fn spawn(&self, mut request: SubagentRequest) -> Result<SubagentResult, ToolError> {
if let Some(parent_session_id) = self.parent_session_id.as_deref() {
request.parent_session_id = parent_session_id.to_owned();
}
let (respond_to, response_rx) = oneshot::channel();
let cancel_on_receiver_drop = request.owner.is_workflow();
let cancel_token = request.cancel_token.clone();
// Replace the dummy oneshot with our fresh one. Using struct update
// syntax (`..request`) ensures new fields added to `SubagentRequest`
// are forwarded automatically — a field-by-field copy would silently
// drop them.
self.tx
.send(SubagentEvent::Spawn(Box::new(SubagentRequest {
result_tx,
..request
})))
.send(SubagentEvent::Spawn(SubagentSpawnRequest {
request: Box::new(request),
result_tx: respond_to,
}))
.map_err(|_| {
ToolError::custom(
"channel_closed",
@ -160,7 +277,7 @@ impl SubagentBackend for ChannelBackend {
cancel_token: cancel_token.clone(),
armed: true,
});
let result = result_rx.await;
let result = response_rx.await;
if result.is_ok() {
if let Some(guard) = receiver_guard.as_mut() {
guard.armed = false;
@ -185,6 +302,7 @@ impl SubagentBackend for ChannelBackend {
let (respond_to, response_rx) = oneshot::channel();
let sent = self.tx.send(SubagentEvent::Query(SubagentQueryRequest {
subagent_id: id.to_string(),
parent_session_id: self.parent_session_id(),
block,
timeout_ms,
respond_to,
@ -198,6 +316,7 @@ impl SubagentBackend for ChannelBackend {
async fn cancel(&self, id: &str) -> SubagentCancelOutcome {
let (respond_to, response_rx) = oneshot::channel();
let sent = self.tx.send(SubagentEvent::Cancel(SubagentCancelRequest {
parent_session_id: self.parent_session_id(),
target: SubagentCancelTarget::SubagentId(id.to_string()),
respond_to,
}));
@ -212,6 +331,10 @@ impl SubagentBackend for ChannelBackend {
subagent_type: &str,
parent_session_id: &str,
) -> SubagentValidateTypeOutcome {
let parent_session_id = self
.parent_session_id
.as_deref()
.unwrap_or(parent_session_id);
let (respond_to, response_rx) = oneshot::channel();
if self
.tx
@ -255,6 +378,10 @@ impl SubagentBackend for ChannelBackend {
harness_agent_type: Option<&str>,
parent_session_id: &str,
) -> SubagentDescribeOutcome {
let parent_session_id = self
.parent_session_id
.as_deref()
.unwrap_or(parent_session_id);
let (respond_to, response_rx) = oneshot::channel();
if self
.tx
@ -301,7 +428,7 @@ pub const VALIDATE_TYPE_TIMEOUT: std::time::Duration = std::time::Duration::from
pub const VALIDATE_TYPE_TIMEOUT_ENV_VAR: &str = "XAI_VALIDATE_TYPE_TIMEOUT_MS";
/// Validation timeout, honoring the env-var override.
pub(crate) fn validate_type_timeout() -> std::time::Duration {
pub fn validate_type_timeout() -> std::time::Duration {
let raw = std::env::var(VALIDATE_TYPE_TIMEOUT_ENV_VAR).ok();
parse_timeout_ms(raw.as_deref())
.map(std::time::Duration::from_millis)
@ -313,604 +440,14 @@ pub(crate) fn parse_timeout_ms(value: Option<&str>) -> Option<u64> {
value?.parse::<u64>().ok().filter(|&ms| ms > 0)
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
use tokio::sync::mpsc;
/// Helper: receive the next event, match the expected variant, or panic.
macro_rules! recv_event {
($rx:expr, Spawn) => {{
let event = $rx.recv().await.unwrap();
match event {
SubagentEvent::Spawn(inner) => *inner,
_ => panic!("Expected SubagentEvent::Spawn, got different variant"),
}
}};
($rx:expr, $variant:ident) => {{
let event = $rx.recv().await.unwrap();
match event {
SubagentEvent::$variant(inner) => inner,
_ => panic!(
"Expected SubagentEvent::{}, got different variant",
stringify!($variant)
),
}
}};
}
#[tokio::test]
async fn channel_backend_spawn_success() {
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
let handle = tokio::spawn(async move {
let req = recv_event!(rx, Spawn);
assert_eq!(req.id, "test-id");
assert_eq!(req.prompt, "do something");
req.result_tx
.send(SubagentResult {
success: true,
output: Arc::from("done"),
subagent_id: "test-id".to_string(),
child_session_id: "test-id".to_string(),
tool_calls: 3,
turns: 1,
duration_ms: 500,
..Default::default()
})
.unwrap();
});
let (dummy_tx, _dummy_rx) = oneshot::channel();
let request = SubagentRequest {
id: "test-id".to_string(),
prompt: "do something".to_string(),
description: "test".to_string(),
subagent_type: "general-purpose".to_string(),
parent_session_id: "parent".to_string(),
parent_prompt_id: None,
resume_from: None,
cwd: None,
runtime_overrides: Default::default(),
run_in_background: false,
surface_completion: true,
await_to_completion: false,
fork_context: false,
owner: super::super::types::SubagentOwner::Task,
cancel_token: tokio_util::sync::CancellationToken::new(),
result_tx: dummy_tx,
};
let result = backend.spawn(request).await.unwrap();
assert!(result.success);
assert_eq!(result.subagent_id, "test-id");
assert_eq!(result.tool_calls, 3);
handle.await.unwrap();
}
#[tokio::test]
async fn channel_backend_spawn_closed_channel() {
let (tx, rx) = mpsc::unbounded_channel::<SubagentEvent>();
drop(rx);
let backend = ChannelBackend::new(tx);
let (dummy_tx, _dummy_rx) = oneshot::channel();
let request = SubagentRequest {
id: "test-id".to_string(),
prompt: "do something".to_string(),
description: "test".to_string(),
subagent_type: "general-purpose".to_string(),
parent_session_id: "parent".to_string(),
parent_prompt_id: None,
resume_from: None,
cwd: None,
runtime_overrides: Default::default(),
run_in_background: false,
surface_completion: true,
await_to_completion: false,
fork_context: false,
owner: super::super::types::SubagentOwner::Task,
cancel_token: tokio_util::sync::CancellationToken::new(),
result_tx: dummy_tx,
};
let err = backend.spawn(request).await.unwrap_err();
assert!(err.to_string().contains("channel closed"));
}
#[tokio::test]
async fn channel_backend_query_found() {
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
let handle = tokio::spawn(async move {
let req = recv_event!(rx, Query);
assert_eq!(req.subagent_id, "sub-1");
assert!(req.block);
assert_eq!(req.timeout_ms, Some(5000));
req.respond_to
.send(Some(SubagentSnapshot {
subagent_id: "sub-1".to_string(),
description: "find bugs".to_string(),
subagent_type: "explore".to_string(),
status: super::super::types::SubagentSnapshotStatus::Completed {
output: "result".to_string(),
tool_calls: 2,
turns: 1,
worktree_path: None,
},
started_at_epoch_ms: 1000,
duration_ms: 200,
persona: Some("reviewer".to_string()),
}))
.unwrap();
});
let snap = backend.query("sub-1", true, Some(5000)).await;
let snap = snap.expect("snapshot should be present");
assert_eq!(snap.subagent_id, "sub-1");
assert_eq!(snap.description, "find bugs");
assert_eq!(snap.subagent_type, "explore");
assert_eq!(snap.started_at_epoch_ms, 1000);
assert_eq!(snap.duration_ms, 200);
assert_eq!(snap.persona.as_deref(), Some("reviewer"));
match &snap.status {
super::super::types::SubagentSnapshotStatus::Completed {
output,
tool_calls,
turns,
worktree_path,
} => {
assert_eq!(output, "result");
assert_eq!(*tool_calls, 2);
assert_eq!(*turns, 1);
assert!(worktree_path.is_none());
}
other => panic!("Expected Completed, got {:?}", other),
}
handle.await.unwrap();
}
#[tokio::test]
async fn channel_backend_query_non_blocking_passes_through() {
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
let handle = tokio::spawn(async move {
let req = recv_event!(rx, Query);
assert_eq!(req.subagent_id, "sub-nb");
assert!(!req.block, "block should be false");
assert_eq!(req.timeout_ms, None, "timeout_ms should be None");
req.respond_to.send(None).unwrap();
});
let snap = backend.query("sub-nb", false, None).await;
assert!(snap.is_none());
handle.await.unwrap();
}
#[tokio::test]
async fn channel_backend_query_not_found() {
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
let handle = tokio::spawn(async move {
let req = recv_event!(rx, Query);
req.respond_to.send(None).unwrap();
});
let snap = backend.query("nonexistent", false, None).await;
assert!(snap.is_none());
handle.await.unwrap();
}
#[tokio::test]
async fn channel_backend_cancel_success() {
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
let handle = tokio::spawn(async move {
let req = recv_event!(rx, Cancel);
match &req.target {
SubagentCancelTarget::SubagentId(id) => assert_eq!(id, "sub-cancel"),
other => panic!("Expected SubagentId, got {:?}", other),
}
req.respond_to
.send(SubagentCancelOutcome::Cancelled)
.unwrap();
});
let outcome = backend.cancel("sub-cancel").await;
assert!(matches!(outcome, SubagentCancelOutcome::Cancelled));
handle.await.unwrap();
}
#[tokio::test]
async fn channel_backend_cancel_closed_channel() {
let (tx, rx) = mpsc::unbounded_channel::<SubagentEvent>();
drop(rx);
let backend = ChannelBackend::new(tx);
let outcome = backend.cancel("sub-cancel").await;
assert!(matches!(outcome, SubagentCancelOutcome::NotFound));
}
#[tokio::test]
async fn workflow_spawn_future_drop_cancels_but_task_drop_does_not() {
fn request_for(owner: super::super::types::SubagentOwner) -> SubagentRequest {
let (dummy_tx, _dummy_rx) = oneshot::channel();
SubagentRequest {
id: "drop-owner-test".to_string(),
prompt: "test".to_string(),
description: "test".to_string(),
subagent_type: "general-purpose".to_string(),
parent_session_id: "parent".to_string(),
parent_prompt_id: None,
resume_from: None,
cwd: None,
runtime_overrides: Default::default(),
run_in_background: false,
surface_completion: false,
await_to_completion: true,
fork_context: false,
owner,
cancel_token: tokio_util::sync::CancellationToken::new(),
result_tx: dummy_tx,
}
}
for (owner, should_cancel) in [
(super::super::types::SubagentOwner::Task, false),
(super::super::types::SubagentOwner::workflow("wf-1"), true),
] {
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = Arc::new(ChannelBackend::new(tx));
let request = request_for(owner);
let cancel_token = request.cancel_token.clone();
let task = tokio::spawn({
let backend = backend.clone();
async move { backend.spawn(request).await }
});
let spawned = recv_event!(rx, Spawn);
task.abort();
let _ = task.await;
assert_eq!(
cancel_token.is_cancelled(),
should_cancel,
"only workflow receiver drop owns cancellation"
);
drop(spawned.result_tx);
}
}
#[tokio::test]
async fn channel_backend_spawn_result_dropped() {
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
let handle = tokio::spawn(async move {
let req = recv_event!(rx, Spawn);
drop(req.result_tx);
});
let (dummy_tx, _dummy_rx) = oneshot::channel();
let request = SubagentRequest {
id: "drop-test".to_string(),
prompt: "test".to_string(),
description: "test".to_string(),
subagent_type: "general-purpose".to_string(),
parent_session_id: "parent".to_string(),
parent_prompt_id: None,
resume_from: None,
cwd: None,
runtime_overrides: Default::default(),
run_in_background: false,
surface_completion: true,
await_to_completion: false,
fork_context: false,
owner: super::super::types::SubagentOwner::Task,
cancel_token: tokio_util::sync::CancellationToken::new(),
result_tx: dummy_tx,
};
let err = backend.spawn(request).await.unwrap_err();
assert!(
err.to_string().contains("result channel dropped"),
"error: {err}"
);
handle.await.unwrap();
}
#[tokio::test]
async fn channel_backend_query_closed_channel() {
let (tx, rx) = mpsc::unbounded_channel::<SubagentEvent>();
drop(rx);
let backend = ChannelBackend::new(tx);
let snap = backend.query("sub-1", false, None).await;
assert!(snap.is_none());
}
// ── validate_type ────────────────────────────────────────────────
#[tokio::test]
async fn channel_backend_validate_type_round_trips_outcome() {
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
let handle = tokio::spawn(async move {
let event = rx.recv().await.unwrap();
match event {
SubagentEvent::ValidateType(req) => {
assert_eq!(req.subagent_type, "explore");
assert_eq!(req.parent_session_id, "parent-1");
req.respond_to
.send(SubagentValidateTypeOutcome::Ok)
.unwrap();
}
_ => panic!("Expected ValidateType event"),
}
});
let outcome = backend.validate_type("explore", "parent-1").await;
assert!(matches!(outcome, SubagentValidateTypeOutcome::Ok));
handle.await.unwrap();
}
#[tokio::test]
async fn channel_backend_validate_type_propagates_unknown_outcome() {
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
let handle = tokio::spawn(async move {
if let Some(SubagentEvent::ValidateType(req)) = rx.recv().await {
req.respond_to
.send(SubagentValidateTypeOutcome::Unknown {
available: vec!["explore".into(), "plan".into()],
})
.unwrap();
}
});
let outcome = backend.validate_type("invented", "p").await;
match outcome {
SubagentValidateTypeOutcome::Unknown { available } => {
assert_eq!(available, vec!["explore".to_string(), "plan".to_string()]);
}
other => panic!("expected Unknown, got {other:?}"),
}
handle.await.unwrap();
}
#[tokio::test]
async fn channel_backend_validate_type_returns_validation_unavailable_when_channel_closed() {
let (tx, rx) = mpsc::unbounded_channel::<SubagentEvent>();
drop(rx);
let backend = ChannelBackend::new(tx);
let outcome = backend.validate_type("explore", "p").await;
assert!(matches!(
outcome,
SubagentValidateTypeOutcome::ValidationUnavailable
));
}
#[tokio::test]
async fn channel_backend_validate_type_returns_validation_unavailable_when_responder_dropped() {
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
let handle = tokio::spawn(async move {
if let Some(SubagentEvent::ValidateType(req)) = rx.recv().await {
drop(req.respond_to);
}
});
let outcome = backend.validate_type("explore", "p").await;
assert!(matches!(
outcome,
SubagentValidateTypeOutcome::ValidationUnavailable,
));
handle.await.unwrap();
}
use super::super::types::test_capture;
#[tokio::test(start_paused = true)]
async fn channel_backend_validate_type_logs_warn_on_timeout() {
let captured = test_capture::capture();
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
// Coordinator receives but never replies; keeps the responder
// alive so the timeout arm fires (not responder-dropped).
let holder = tokio::spawn(async move {
if let Some(SubagentEvent::ValidateType(req)) = rx.recv().await {
std::mem::forget(req.respond_to);
std::future::pending::<()>().await;
}
});
let validate = tokio::spawn(async move { backend.validate_type("explore", "p").await });
tokio::time::advance(VALIDATE_TYPE_TIMEOUT + std::time::Duration::from_millis(1)).await;
let outcome = validate.await.unwrap();
assert!(matches!(
outcome,
SubagentValidateTypeOutcome::ValidationUnavailable
));
let mut events_rx = captured.events_rx;
let mut saw_timeout_warn = false;
while let Ok(event) = events_rx.try_recv() {
if event.level == tracing::Level::WARN
&& event.fields.contains("coordinator validation timed out")
&& event.fields.contains("subagent_type=explore")
&& event.fields.contains("timeout_ms=")
{
saw_timeout_warn = true;
break;
}
}
assert!(saw_timeout_warn, "must emit WARN with timeout_ms field");
holder.abort();
}
// ── describe_subagent_type ───────────────────────────────────────
#[tokio::test]
async fn channel_backend_describe_round_trips_summary() {
use super::super::types::{SubagentDescribeOutcome, SubagentTypeSummary};
use crate::types::tool::ToolKind;
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
let handle = tokio::spawn(async move {
match rx.recv().await.unwrap() {
SubagentEvent::DescribeType(req) => {
assert_eq!(req.subagent_type, "explore");
assert_eq!(req.harness_agent_type.as_deref(), Some("cursor"));
assert_eq!(req.parent_session_id, "parent-1");
let mut summary = SubagentTypeSummary {
can_read: true,
can_search: true,
..Default::default()
};
summary
.tool_names
.insert(ToolKind::Read, "read_file".to_string());
req.respond_to
.send(SubagentDescribeOutcome::Ok(summary))
.unwrap();
}
_ => panic!("Expected DescribeType event"),
}
});
let outcome = backend
.describe_subagent_type("explore", Some("cursor"), "parent-1")
.await;
match outcome {
SubagentDescribeOutcome::Ok(summary) => {
assert!(summary.can_read && summary.can_search && !summary.can_execute);
assert_eq!(
summary.tool_names.get(&ToolKind::Read).unwrap(),
"read_file"
);
}
other => panic!("expected Ok, got {other:?}"),
}
handle.await.unwrap();
}
#[tokio::test]
async fn channel_backend_describe_propagates_not_allowed_outcome() {
use super::super::types::SubagentDescribeOutcome;
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
let handle = tokio::spawn(async move {
if let Some(SubagentEvent::DescribeType(req)) = rx.recv().await {
req.respond_to
.send(SubagentDescribeOutcome::NotAllowed {
allowed: vec!["explore".into()],
})
.unwrap();
}
});
match backend.describe_subagent_type("plan", None, "p").await {
SubagentDescribeOutcome::NotAllowed { allowed } => {
assert_eq!(allowed, vec!["explore".to_string()]);
}
other => panic!("expected NotAllowed, got {other:?}"),
}
handle.await.unwrap();
}
#[tokio::test]
async fn channel_backend_describe_returns_unavailable_when_channel_closed() {
use super::super::types::SubagentDescribeOutcome;
let (tx, rx) = mpsc::unbounded_channel::<SubagentEvent>();
drop(rx);
let backend = ChannelBackend::new(tx);
assert!(matches!(
backend.describe_subagent_type("explore", None, "p").await,
SubagentDescribeOutcome::Unavailable
));
}
#[tokio::test]
async fn channel_backend_describe_returns_unavailable_when_responder_dropped() {
use super::super::types::SubagentDescribeOutcome;
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
let handle = tokio::spawn(async move {
if let Some(SubagentEvent::DescribeType(req)) = rx.recv().await {
drop(req.respond_to);
}
});
assert!(matches!(
backend.describe_subagent_type("explore", None, "p").await,
SubagentDescribeOutcome::Unavailable
));
handle.await.unwrap();
}
#[tokio::test(start_paused = true)]
async fn channel_backend_describe_returns_unavailable_on_timeout() {
use super::super::types::SubagentDescribeOutcome;
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
let holder = tokio::spawn(async move {
if let Some(SubagentEvent::DescribeType(req)) = rx.recv().await {
std::mem::forget(req.respond_to);
std::future::pending::<()>().await;
}
});
let describe =
tokio::spawn(async move { backend.describe_subagent_type("explore", None, "p").await });
tokio::time::advance(VALIDATE_TYPE_TIMEOUT + std::time::Duration::from_millis(1)).await;
assert!(matches!(
describe.await.unwrap(),
SubagentDescribeOutcome::Unavailable
));
holder.abort();
}
#[test]
fn parse_timeout_ms_returns_none_for_unset() {
assert_eq!(parse_timeout_ms(None), None);
}
#[test]
fn parse_timeout_ms_returns_none_for_unparseable() {
assert_eq!(parse_timeout_ms(Some("not-a-number")), None);
assert_eq!(parse_timeout_ms(Some("")), None);
assert_eq!(parse_timeout_ms(Some("3.14")), None);
assert_eq!(parse_timeout_ms(Some("-100")), None);
}
#[test]
fn parse_timeout_ms_returns_none_for_zero() {
assert_eq!(parse_timeout_ms(Some("0")), None);
}
#[test]
fn parse_timeout_ms_returns_value_for_positive_integer() {
assert_eq!(parse_timeout_ms(Some("5000")), Some(5000));
assert_eq!(parse_timeout_ms(Some("1")), Some(1));
}
/// Resolve a `Duration` from a positive-millisecond env override, falling back
/// to `default` when the var is unset / non-numeric / zero.
pub fn env_duration_or(env_var: &str, default: std::time::Duration) -> std::time::Duration {
parse_timeout_ms(std::env::var(env_var).ok().as_deref())
.map(std::time::Duration::from_millis)
.unwrap_or(default)
}
#[cfg(test)]
#[path = "backend_tests.rs"]
mod tests;

590
crates/codegen/xai-grok-tools/src/implementations/grok_build/task/backend_tests.rs Normal file
View file

@ -0,0 +1,590 @@
use super::*;
use std::sync::Arc;
use tokio::sync::mpsc;
/// Helper: receive the next event, match the expected variant, or panic.
macro_rules! recv_event {
($rx:expr, Spawn) => {{
let event = $rx.recv().await.unwrap();
match event {
SubagentEvent::Spawn(inner) => inner,
_ => panic!("Expected SubagentEvent::Spawn, got different variant"),
}
}};
($rx:expr, $variant:ident) => {{
let event = $rx.recv().await.unwrap();
match event {
SubagentEvent::$variant(inner) => inner,
_ => panic!(
"Expected SubagentEvent::{}, got different variant",
stringify!($variant)
),
}
}};
}
#[tokio::test]
async fn channel_backend_spawn_success() {
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
let handle = tokio::spawn(async move {
let req = recv_event!(rx, Spawn);
assert_eq!(req.request.id, "test-id");
assert_eq!(req.request.prompt, "do something");
req.result_tx
.send(SubagentResult {
success: true,
output: Arc::from("done"),
subagent_id: "test-id".to_string(),
child_session_id: "test-id".to_string(),
tool_calls: 3,
turns: 1,
duration_ms: 500,
..Default::default()
})
.unwrap();
});
let request = SubagentRequest {
id: "test-id".to_string(),
prompt: "do something".to_string(),
description: "test".to_string(),
subagent_type: "general-purpose".to_string(),
parent_session_id: "parent".to_string(),
parent_prompt_id: None,
resume_from: None,
cwd: None,
runtime_overrides: Default::default(),
run_in_background: false,
surface_completion: true,
await_to_completion: false,
fork_context: false,
owner: super::super::types::SubagentOwner::Task,
cancel_token: tokio_util::sync::CancellationToken::new(),
};
let result = backend.spawn(request).await.unwrap();
assert!(result.success);
assert_eq!(result.subagent_id, "test-id");
assert_eq!(result.tool_calls, 3);
handle.await.unwrap();
}
#[tokio::test]
async fn channel_backend_spawn_closed_channel() {
let (tx, rx) = mpsc::unbounded_channel::<SubagentEvent>();
drop(rx);
let backend = ChannelBackend::new(tx);
let request = SubagentRequest {
id: "test-id".to_string(),
prompt: "do something".to_string(),
description: "test".to_string(),
subagent_type: "general-purpose".to_string(),
parent_session_id: "parent".to_string(),
parent_prompt_id: None,
resume_from: None,
cwd: None,
runtime_overrides: Default::default(),
run_in_background: false,
surface_completion: true,
await_to_completion: false,
fork_context: false,
owner: super::super::types::SubagentOwner::Task,
cancel_token: tokio_util::sync::CancellationToken::new(),
};
let err = backend.spawn(request).await.unwrap_err();
assert!(err.to_string().contains("channel closed"));
}
#[tokio::test]
async fn channel_backend_query_found() {
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
let handle = tokio::spawn(async move {
let req = recv_event!(rx, Query);
assert_eq!(req.subagent_id, "sub-1");
assert!(req.block);
assert_eq!(req.timeout_ms, Some(5000));
req.respond_to
.send(Some(SubagentSnapshot {
subagent_id: "sub-1".to_string(),
description: "find bugs".to_string(),
subagent_type: "explore".to_string(),
status: super::super::types::SubagentSnapshotStatus::Completed {
output: "result".to_string(),
tool_calls: 2,
turns: 1,
worktree_path: None,
},
started_at_epoch_ms: 1000,
duration_ms: 200,
persona: Some("reviewer".to_string()),
}))
.unwrap();
});
let snap = backend.query("sub-1", true, Some(5000)).await;
let snap = snap.expect("snapshot should be present");
assert_eq!(snap.subagent_id, "sub-1");
assert_eq!(snap.description, "find bugs");
assert_eq!(snap.subagent_type, "explore");
assert_eq!(snap.started_at_epoch_ms, 1000);
assert_eq!(snap.duration_ms, 200);
assert_eq!(snap.persona.as_deref(), Some("reviewer"));
match &snap.status {
super::super::types::SubagentSnapshotStatus::Completed {
output,
tool_calls,
turns,
worktree_path,
} => {
assert_eq!(output, "result");
assert_eq!(*tool_calls, 2);
assert_eq!(*turns, 1);
assert!(worktree_path.is_none());
}
other => panic!("Expected Completed, got {:?}", other),
}
handle.await.unwrap();
}
#[tokio::test]
async fn channel_backend_query_non_blocking_passes_through() {
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
let handle = tokio::spawn(async move {
let req = recv_event!(rx, Query);
assert_eq!(req.subagent_id, "sub-nb");
assert!(!req.block, "block should be false");
assert_eq!(req.timeout_ms, None, "timeout_ms should be None");
req.respond_to.send(None).unwrap();
});
let snap = backend.query("sub-nb", false, None).await;
assert!(snap.is_none());
handle.await.unwrap();
}
#[tokio::test]
async fn channel_backend_query_not_found() {
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
let handle = tokio::spawn(async move {
let req = recv_event!(rx, Query);
req.respond_to.send(None).unwrap();
});
let snap = backend.query("nonexistent", false, None).await;
assert!(snap.is_none());
handle.await.unwrap();
}
#[tokio::test]
async fn channel_backend_cancel_success() {
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
let handle = tokio::spawn(async move {
let req = recv_event!(rx, Cancel);
match &req.target {
SubagentCancelTarget::SubagentId(id) => assert_eq!(id, "sub-cancel"),
other => panic!("Expected SubagentId, got {:?}", other),
}
req.respond_to
.send(SubagentCancelOutcome::Cancelled)
.unwrap();
});
let outcome = backend.cancel("sub-cancel").await;
assert!(matches!(outcome, SubagentCancelOutcome::Cancelled));
handle.await.unwrap();
}
#[tokio::test]
async fn channel_backend_cancel_closed_channel() {
let (tx, rx) = mpsc::unbounded_channel::<SubagentEvent>();
drop(rx);
let backend = ChannelBackend::new(tx);
let outcome = backend.cancel("sub-cancel").await;
assert!(matches!(outcome, SubagentCancelOutcome::NotFound));
}
#[tokio::test]
async fn workflow_spawn_future_drop_cancels_but_task_drop_does_not() {
fn request_for(owner: super::super::types::SubagentOwner) -> SubagentRequest {
SubagentRequest {
id: "drop-owner-test".to_string(),
prompt: "test".to_string(),
description: "test".to_string(),
subagent_type: "general-purpose".to_string(),
parent_session_id: "parent".to_string(),
parent_prompt_id: None,
resume_from: None,
cwd: None,
runtime_overrides: Default::default(),
run_in_background: false,
surface_completion: false,
await_to_completion: true,
fork_context: false,
owner,
cancel_token: tokio_util::sync::CancellationToken::new(),
}
}
for (owner, should_cancel) in [
(super::super::types::SubagentOwner::Task, false),
(super::super::types::SubagentOwner::workflow("wf-1"), true),
] {
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = Arc::new(ChannelBackend::new(tx));
let request = request_for(owner);
let cancel_token = request.cancel_token.clone();
let task = tokio::spawn({
let backend = backend.clone();
async move { backend.spawn(request).await }
});
let spawned = recv_event!(rx, Spawn);
task.abort();
let _ = task.await;
assert_eq!(
cancel_token.is_cancelled(),
should_cancel,
"only workflow receiver drop owns cancellation"
);
drop(spawned.result_tx);
}
}
#[tokio::test]
async fn channel_backend_spawn_result_dropped() {
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
let handle = tokio::spawn(async move {
let req = recv_event!(rx, Spawn);
drop(req.result_tx);
});
let request = SubagentRequest {
id: "drop-test".to_string(),
prompt: "test".to_string(),
description: "test".to_string(),
subagent_type: "general-purpose".to_string(),
parent_session_id: "parent".to_string(),
parent_prompt_id: None,
resume_from: None,
cwd: None,
runtime_overrides: Default::default(),
run_in_background: false,
surface_completion: true,
await_to_completion: false,
fork_context: false,
owner: super::super::types::SubagentOwner::Task,
cancel_token: tokio_util::sync::CancellationToken::new(),
};
let err = backend.spawn(request).await.unwrap_err();
assert!(
err.to_string().contains("result channel dropped"),
"error: {err}"
);
handle.await.unwrap();
}
#[tokio::test]
async fn channel_backend_query_closed_channel() {
let (tx, rx) = mpsc::unbounded_channel::<SubagentEvent>();
drop(rx);
let backend = ChannelBackend::new(tx);
let snap = backend.query("sub-1", false, None).await;
assert!(snap.is_none());
}
// ── validate_type ────────────────────────────────────────────────
#[tokio::test]
async fn channel_backend_validate_type_round_trips_outcome() {
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
let handle = tokio::spawn(async move {
let event = rx.recv().await.unwrap();
match event {
SubagentEvent::ValidateType(req) => {
assert_eq!(req.subagent_type, "explore");
assert_eq!(req.parent_session_id, "parent-1");
req.respond_to
.send(SubagentValidateTypeOutcome::Ok)
.unwrap();
}
_ => panic!("Expected ValidateType event"),
}
});
let outcome = backend.validate_type("explore", "parent-1").await;
assert!(matches!(outcome, SubagentValidateTypeOutcome::Ok));
handle.await.unwrap();
}
#[tokio::test]
async fn channel_backend_validate_type_propagates_unknown_outcome() {
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
let handle = tokio::spawn(async move {
if let Some(SubagentEvent::ValidateType(req)) = rx.recv().await {
req.respond_to
.send(SubagentValidateTypeOutcome::Unknown {
available: vec!["explore".into(), "plan".into()],
})
.unwrap();
}
});
let outcome = backend.validate_type("invented", "p").await;
match outcome {
SubagentValidateTypeOutcome::Unknown { available } => {
assert_eq!(available, vec!["explore".to_string(), "plan".to_string()]);
}
other => panic!("expected Unknown, got {other:?}"),
}
handle.await.unwrap();
}
#[tokio::test]
async fn channel_backend_validate_type_returns_validation_unavailable_when_channel_closed() {
let (tx, rx) = mpsc::unbounded_channel::<SubagentEvent>();
drop(rx);
let backend = ChannelBackend::new(tx);
let outcome = backend.validate_type("explore", "p").await;
assert!(matches!(
outcome,
SubagentValidateTypeOutcome::ValidationUnavailable
));
}
#[tokio::test]
async fn channel_backend_validate_type_returns_validation_unavailable_when_responder_dropped() {
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
let handle = tokio::spawn(async move {
if let Some(SubagentEvent::ValidateType(req)) = rx.recv().await {
drop(req.respond_to);
}
});
let outcome = backend.validate_type("explore", "p").await;
assert!(matches!(
outcome,
SubagentValidateTypeOutcome::ValidationUnavailable,
));
handle.await.unwrap();
}
use super::super::types::test_capture;
#[tokio::test(start_paused = true)]
async fn channel_backend_validate_type_logs_warn_on_timeout() {
let captured = test_capture::capture();
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
// Coordinator receives but never replies; keeps the responder
// alive so the timeout arm fires (not responder-dropped).
let holder = tokio::spawn(async move {
if let Some(SubagentEvent::ValidateType(req)) = rx.recv().await {
std::mem::forget(req.respond_to);
std::future::pending::<()>().await;
}
});
let validate = tokio::spawn(async move { backend.validate_type("explore", "p").await });
tokio::time::advance(VALIDATE_TYPE_TIMEOUT + std::time::Duration::from_millis(1)).await;
let outcome = validate.await.unwrap();
assert!(matches!(
outcome,
SubagentValidateTypeOutcome::ValidationUnavailable
));
let mut events_rx = captured.events_rx;
let mut saw_timeout_warn = false;
while let Ok(event) = events_rx.try_recv() {
if event.level == tracing::Level::WARN
&& event.fields.contains("coordinator validation timed out")
&& event.fields.contains("subagent_type=explore")
&& event.fields.contains("timeout_ms=")
{
saw_timeout_warn = true;
break;
}
}
assert!(saw_timeout_warn, "must emit WARN with timeout_ms field");
holder.abort();
}
// ── describe_subagent_type ───────────────────────────────────────
#[tokio::test]
async fn channel_backend_describe_round_trips_summary() {
use super::super::types::{SubagentDescribeOutcome, SubagentTypeSummary};
use crate::types::tool::ToolKind;
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
let handle = tokio::spawn(async move {
match rx.recv().await.unwrap() {
SubagentEvent::DescribeType(req) => {
assert_eq!(req.subagent_type, "explore");
assert_eq!(req.harness_agent_type.as_deref(), Some("cursor"));
assert_eq!(req.parent_session_id, "parent-1");
let mut summary = SubagentTypeSummary {
can_read: true,
can_search: true,
..Default::default()
};
summary
.tool_names
.insert(ToolKind::Read, "read_file".to_string());
req.respond_to
.send(SubagentDescribeOutcome::Ok(summary))
.unwrap();
}
_ => panic!("Expected DescribeType event"),
}
});
let outcome = backend
.describe_subagent_type("explore", Some("cursor"), "parent-1")
.await;
match outcome {
SubagentDescribeOutcome::Ok(summary) => {
assert!(summary.can_read && summary.can_search && !summary.can_execute);
assert_eq!(
summary.tool_names.get(&ToolKind::Read).unwrap(),
"read_file"
);
}
other => panic!("expected Ok, got {other:?}"),
}
handle.await.unwrap();
}
#[tokio::test]
async fn channel_backend_describe_propagates_not_allowed_outcome() {
use super::super::types::SubagentDescribeOutcome;
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
let handle = tokio::spawn(async move {
if let Some(SubagentEvent::DescribeType(req)) = rx.recv().await {
req.respond_to
.send(SubagentDescribeOutcome::NotAllowed {
allowed: vec!["explore".into()],
})
.unwrap();
}
});
match backend.describe_subagent_type("plan", None, "p").await {
SubagentDescribeOutcome::NotAllowed { allowed } => {
assert_eq!(allowed, vec!["explore".to_string()]);
}
other => panic!("expected NotAllowed, got {other:?}"),
}
handle.await.unwrap();
}
#[tokio::test]
async fn channel_backend_describe_returns_unavailable_when_channel_closed() {
use super::super::types::SubagentDescribeOutcome;
let (tx, rx) = mpsc::unbounded_channel::<SubagentEvent>();
drop(rx);
let backend = ChannelBackend::new(tx);
assert!(matches!(
backend.describe_subagent_type("explore", None, "p").await,
SubagentDescribeOutcome::Unavailable
));
}
#[tokio::test]
async fn channel_backend_describe_returns_unavailable_when_responder_dropped() {
use super::super::types::SubagentDescribeOutcome;
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
let handle = tokio::spawn(async move {
if let Some(SubagentEvent::DescribeType(req)) = rx.recv().await {
drop(req.respond_to);
}
});
assert!(matches!(
backend.describe_subagent_type("explore", None, "p").await,
SubagentDescribeOutcome::Unavailable
));
handle.await.unwrap();
}
#[tokio::test(start_paused = true)]
async fn channel_backend_describe_returns_unavailable_on_timeout() {
use super::super::types::SubagentDescribeOutcome;
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
let backend = ChannelBackend::new(tx);
let holder = tokio::spawn(async move {
if let Some(SubagentEvent::DescribeType(req)) = rx.recv().await {
std::mem::forget(req.respond_to);
std::future::pending::<()>().await;
}
});
let describe =
tokio::spawn(async move { backend.describe_subagent_type("explore", None, "p").await });
tokio::time::advance(VALIDATE_TYPE_TIMEOUT + std::time::Duration::from_millis(1)).await;
assert!(matches!(
describe.await.unwrap(),
SubagentDescribeOutcome::Unavailable
));
holder.abort();
}
#[test]
fn parse_timeout_ms_returns_none_for_unset() {
assert_eq!(parse_timeout_ms(None), None);
}
#[test]
fn parse_timeout_ms_returns_none_for_unparseable() {
assert_eq!(parse_timeout_ms(Some("not-a-number")), None);
assert_eq!(parse_timeout_ms(Some("")), None);
assert_eq!(parse_timeout_ms(Some("3.14")), None);
assert_eq!(parse_timeout_ms(Some("-100")), None);
}
#[test]
fn parse_timeout_ms_returns_none_for_zero() {
assert_eq!(parse_timeout_ms(Some("0")), None);
}
#[test]
fn parse_timeout_ms_returns_value_for_positive_integer() {
assert_eq!(parse_timeout_ms(Some("5000")), Some(5000));
assert_eq!(parse_timeout_ms(Some("1")), Some(1));
}

840
crates/codegen/xai-grok-tools/src/implementations/grok_build/task/coordinator.rs Normal file
View file

@ -0,0 +1,840 @@
//! Single-writer subagent coordinator actor.
//!
//! The actor owns the command receiver, pending/active/completed state,
//! concrete blocking waiters, foreground deadlines, cancellation, and the
//! terminal delivery disposition. All hosts drive it through `ChannelBackend`;
//! only their `ChildRunner` implementations differ.
//!
//! There is intentionally no shared mutable state in this module. A runner's
//! associated futures may be `Send` or non-`Send`; the resulting actor future
//! inherits that property naturally on stable Rust.
mod query;
use std::collections::{HashMap, HashSet, VecDeque};
use std::sync::Arc;
use futures::FutureExt;
use futures::stream::{FuturesUnordered, StreamExt};
use tokio::sync::{mpsc, oneshot};
use super::coordinator_state::{
ActiveChild, BlockingWaiter, BufferedCompletion, ChildRecord, CompletedChild, InternalEvent,
ListRequest, MAX_COMPLETED_ENTRIES, PendingChild, ProgressFuture, ProgressTarget, ReplyFuture,
TaggedFuture, active_summary, background_at_deadline, background_if_caller_gone,
completed_snapshot, completion_summary, sleep_until, workflow_outstanding,
};
use super::types::{
SpawnedSubagentRef, SubagentCancelOutcome, SubagentCancelTarget, SubagentDescribeOutcome,
SubagentEvent, SubagentOutstandingReply, SubagentRegistryCounts, SubagentRequest,
SubagentResult, SubagentResumeLookup, SubagentResumeSource, SubagentValidateTypeOutcome,
};
pub use super::coordinator_state::{
ChildCompletion, ChildControl, ChildReporter, ChildRunOutput, ChildRunRequest, ChildRunner,
CompletionDisposition, CoordinatorConfig, LocalBoxFuture, SendBoxFuture, StartedChild,
SubagentProgress,
};
/// Channel-owned subagent lifecycle actor.
pub struct SubagentCoordinator<R: ChildRunner> {
commands: mpsc::UnboundedReceiver<SubagentEvent>,
internal_tx: mpsc::UnboundedSender<InternalEvent<R::Control>>,
internal_rx: mpsc::UnboundedReceiver<InternalEvent<R::Control>>,
runner: R,
config: CoordinatorConfig,
pending: HashMap<String, PendingChild>,
active: HashMap<String, ActiveChild<R::Control>>,
completed: HashMap<String, CompletedChild>,
completed_order: VecDeque<String>,
waiters: HashMap<String, Vec<BlockingWaiter>>,
workflow_cancel_waiters: HashMap<String, Vec<oneshot::Sender<SubagentCancelOutcome>>>,
usage_not_applied_prompts: HashSet<PromptScope>,
pending_completions: Vec<BufferedCompletion>,
runs: FuturesUnordered<
TaggedFuture<futures::future::CatchUnwind<std::panic::AssertUnwindSafe<R::RunFuture>>>,
>,
validations: FuturesUnordered<ReplyFuture<R::ValidateFuture, SubagentValidateTypeOutcome>>,
descriptions: FuturesUnordered<ReplyFuture<R::DescribeFuture, SubagentDescribeOutcome>>,
progress: FuturesUnordered<ProgressFuture<<R::Control as ChildControl>::ProgressFuture>>,
list_requests: HashMap<u64, ListRequest>,
next_list_request_id: u64,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct PromptScope {
parent_session_id: String,
prompt_id: String,
}
impl PromptScope {
fn new(parent_session_id: String, prompt_id: String) -> Self {
Self {
parent_session_id,
prompt_id,
}
}
}
impl<R: ChildRunner> SubagentCoordinator<R> {
pub fn new(
commands: mpsc::UnboundedReceiver<SubagentEvent>,
runner: R,
config: CoordinatorConfig,
) -> Self {
let (internal_tx, internal_rx) = mpsc::unbounded_channel();
Self {
commands,
internal_tx,
internal_rx,
runner,
config,
pending: HashMap::new(),
active: HashMap::new(),
completed: HashMap::new(),
completed_order: VecDeque::new(),
waiters: HashMap::new(),
workflow_cancel_waiters: HashMap::new(),
usage_not_applied_prompts: HashSet::new(),
pending_completions: Vec::new(),
runs: FuturesUnordered::new(),
validations: FuturesUnordered::new(),
descriptions: FuturesUnordered::new(),
progress: FuturesUnordered::new(),
list_requests: HashMap::new(),
next_list_request_id: 0,
}
}
pub async fn run(mut self) {
let mut commands_open = true;
loop {
if !commands_open
&& self.runs.is_empty()
&& self.validations.is_empty()
&& self.descriptions.is_empty()
&& self.progress.is_empty()
{
break;
}
let deadline = self.next_deadline();
tokio::select! {
biased;
Some(event) = self.internal_rx.recv() => self.handle_internal(event),
Some((id, output)) = self.runs.next(), if !self.runs.is_empty() => {
match output {
Ok(output) => self.finish_child(&id, output),
Err(_) => self.finish_panicked_child(&id),
}
}
Some((respond_to, outcome)) = self.validations.next(), if !self.validations.is_empty() => {
let _ = respond_to.send(outcome);
}
Some((respond_to, outcome)) = self.descriptions.next(), if !self.descriptions.is_empty() => {
let _ = respond_to.send(outcome);
}
Some((seed, target, progress)) = self.progress.next(), if !self.progress.is_empty() => {
self.finish_progress(seed, target, progress);
}
command = self.commands.recv(), if commands_open => {
match command {
Some(command) => {
self.reap_abandoned_callers();
self.handle_command(command);
}
None => commands_open = false,
}
}
_ = sleep_until(deadline), if deadline.is_some() => self.process_deadlines(),
}
while self.completed.len() > MAX_COMPLETED_ENTRIES {
let Some(id) = self.completed_order.pop_front() else {
break;
};
self.completed.remove(&id);
}
}
self.cancel_all_children();
}
fn handle_command(&mut self, command: SubagentEvent) {
match command {
SubagentEvent::Spawn(command) => {
let mut request = *command.request;
if let Some((root_parent, loop_task_id)) = self
.active
.values()
.find(|child| child.child_session_id == request.parent_session_id)
.map(|child| {
(
child.request.parent_session_id.clone(),
child.request.runtime_overrides.loop_task_id.clone(),
)
})
{
request.parent_session_id = root_parent;
request.surface_completion = false;
if request.runtime_overrides.loop_task_id.is_none() {
request.runtime_overrides.loop_task_id = loop_task_id;
}
}
let id = request.id.clone();
if self.pending.contains_key(&id)
|| self.active.contains_key(&id)
|| self.completed.contains_key(&id)
{
let _ = command.result_tx.send(SubagentResult {
success: false,
error: Some(format!("Subagent id '{id}' already exists")),
subagent_id: id.clone(),
child_session_id: id,
..Default::default()
});
return;
}
let cancellation = request.cancel_token.clone();
let handle_only = request.run_in_background;
let foreground_deadline = (!request.run_in_background
&& !request.await_to_completion)
.then(|| tokio::time::Instant::now() + self.config.foreground_budget);
self.pending.insert(
id.clone(),
PendingChild {
request: request.clone(),
started_at: std::time::Instant::now(),
cancellation: cancellation.clone(),
spawn_reply: Some(command.result_tx),
foreground_deadline,
handle_only,
explicitly_killed: false,
},
);
self.running_count_changed();
let reporter = ChildReporter {
subagent_id: id.clone(),
tx: self.internal_tx.clone(),
};
self.runs.push(TaggedFuture {
subagent_id: id,
future: Box::pin(
std::panic::AssertUnwindSafe(self.runner.run(ChildRunRequest {
request,
cancellation,
reporter,
}))
.catch_unwind(),
),
});
}
SubagentEvent::Query(query) => {
self.handle_query(
query.subagent_id,
query.parent_session_id,
query.block,
query.timeout_ms,
query.respond_to,
);
}
SubagentEvent::Cancel(request) => match request.target {
SubagentCancelTarget::SubagentId(id) => {
let outcome = self.cancel_one(&id, request.parent_session_id.as_deref(), true);
let _ = request.respond_to.send(outcome);
}
SubagentCancelTarget::ParentPromptId(prompt_id) => {
self.cancel_parent_prompt(&prompt_id, request.parent_session_id.as_deref());
let _ = request.respond_to.send(SubagentCancelOutcome::Cancelled);
}
SubagentCancelTarget::WorkflowRunId(run_id) => {
self.cancel_workflow_children(&run_id, request.parent_session_id.as_deref());
if workflow_outstanding(&self.pending, &self.active, &run_id) == 0 {
let _ = request.respond_to.send(SubagentCancelOutcome::Cancelled);
} else {
self.workflow_cancel_waiters
.entry(run_id)
.or_default()
.push(request.respond_to);
}
}
},
SubagentEvent::ListActive(request) => {
let summaries = self
.active
.values()
.filter(|child| {
child.request.parent_session_id == request.parent_session_id
&& !child.request.owner.is_workflow()
})
.map(active_summary)
.collect();
let _ = request.respond_to.send(summaries);
}
SubagentEvent::ListRunning(request) => {
self.handle_list_running(request.parent_session_id, request.respond_to);
}
SubagentEvent::Completions(request) => {
let (owned, foreign): (Vec<_>, Vec<_>) =
std::mem::take(&mut self.pending_completions)
.into_iter()
.partition(|completion| {
request
.parent_session_id
.as_ref()
.is_none_or(|id| completion.parent_session_id == *id)
});
self.pending_completions = foreign;
let completions = owned
.into_iter()
.map(|completion| completion.summary)
.filter(|summary| !request.suppress_ids.contains(&summary.subagent_id))
.collect();
let _ = request.respond_to.send(completions);
}
SubagentEvent::DiscardSessionCompletions { parent_session_id } => {
self.pending_completions
.retain(|completion| completion.parent_session_id != parent_session_id);
}
SubagentEvent::Outstanding(request) => {
// Reap again here so turn-freeze / Outstanding polls see
// ParentGone even if no other command woke the actor first.
self.reap_abandoned_callers();
let mut live_ids: Vec<_> = self
.pending
.values()
.filter(|child| {
child.request.parent_session_id == request.parent_session_id
&& child.request.parent_prompt_id.as_deref() == Some(&request.prompt_id)
&& !child.request.owner.is_workflow()
&& !child.handle_only
})
.map(|child| child.request.id.clone())
.chain(
self.active
.values()
.filter(|child| {
child.request.parent_session_id == request.parent_session_id
&& child.request.parent_prompt_id.as_deref()
== Some(&request.prompt_id)
&& !child.request.owner.is_workflow()
// Definition-declared background children are
// background for accounting even while the
// spawning tool block-awaits them.
&& !child.handle_only
&& !child.definition_background
})
.map(|child| child.request.id.clone()),
)
.collect();
live_ids.sort();
let background_live = self.pending.values().any(|child| {
child.request.parent_session_id == request.parent_session_id
&& child.request.parent_prompt_id.as_deref() == Some(&request.prompt_id)
&& !child.request.owner.is_workflow()
&& child.handle_only
}) || self.active.values().any(|child| {
child.request.parent_session_id == request.parent_session_id
&& child.request.parent_prompt_id.as_deref() == Some(&request.prompt_id)
&& !child.request.owner.is_workflow()
&& (child.handle_only || child.definition_background)
});
let scope =
PromptScope::new(request.parent_session_id.clone(), request.prompt_id.clone());
let _ = request.respond_to.send(SubagentOutstandingReply {
live_ids,
background_live,
subagent_usage_not_applied: self.usage_not_applied_prompts.contains(&scope),
});
}
SubagentEvent::ClearUsageNotApplied(request) => {
self.usage_not_applied_prompts.remove(&PromptScope::new(
request.parent_session_id,
request.prompt_id,
));
}
SubagentEvent::MarkUsageNotApplied(request) => {
self.usage_not_applied_prompts.insert(PromptScope::new(
request.parent_session_id,
request.prompt_id,
));
let _ = request.respond_to.send(());
}
SubagentEvent::RegistryCounts(request) => {
let _ = request.respond_to.send(SubagentRegistryCounts {
pending: self.pending.len(),
active: self.active.len(),
completed: self.completed.len(),
});
}
SubagentEvent::Inspect(request) => {
self.handle_inspect(
request.subagent_id,
request.parent_session_id,
request.respond_to,
);
}
SubagentEvent::SpawnedRefs(request) => {
let mut refs: Vec<_> = self
.active
.values()
.filter(|child| {
child.request.parent_session_id == request.parent_session_id
&& child.request.parent_prompt_id.as_deref() == Some(&request.prompt_id)
})
.map(|child| SpawnedSubagentRef {
subagent_id: child.request.id.clone(),
child_session_id: child.child_session_id.clone(),
subagent_type: child.request.subagent_type.clone(),
description: child.request.description.clone(),
persona: child.persona.clone(),
resumed_from: child.resumed_from.clone(),
})
.chain(
self.completed
.values()
.filter(|child| {
child.request.parent_session_id == request.parent_session_id
&& child.request.parent_prompt_id.as_deref()
== Some(&request.prompt_id)
})
.map(|child| SpawnedSubagentRef {
subagent_id: child.request.id.clone(),
child_session_id: child.child_session_id.clone(),
subagent_type: child.request.subagent_type.clone(),
description: child.request.description.clone(),
persona: child.persona.clone(),
resumed_from: child.resumed_from.clone(),
}),
)
.collect();
refs.sort_by(|a, b| a.subagent_id.cmp(&b.subagent_id));
let _ = request.respond_to.send(refs);
}
SubagentEvent::ValidateType(request) => {
self.validations.push(ReplyFuture {
future: Box::pin(
self.runner
.validate_type(request.subagent_type, request.parent_session_id),
),
respond_to: Some(request.respond_to),
});
}
SubagentEvent::DescribeType(request) => {
self.descriptions.push(ReplyFuture {
future: Box::pin(self.runner.describe_type(
request.subagent_type,
request.harness_agent_type,
request.parent_session_id,
)),
respond_to: Some(request.respond_to),
});
}
SubagentEvent::LoopUnitActive(request) => {
let is_active = self.pending.values().any(|child| {
child.request.runtime_overrides.loop_task_id.as_deref()
== Some(&request.task_id)
}) || self.active.values().any(|child| {
child.request.runtime_overrides.loop_task_id.as_deref()
== Some(&request.task_id)
});
let _ = request.respond_to.send(is_active);
}
}
}
fn handle_internal(&mut self, event: InternalEvent<R::Control>) {
match event {
InternalEvent::Started {
subagent_id,
child,
respond_to,
} => {
let Some(pending) = self.pending.remove(&subagent_id) else {
let _ = respond_to.send(false);
return;
};
if pending.cancellation.is_cancelled() {
self.pending.insert(subagent_id, pending);
let _ = respond_to.send(false);
return;
}
self.active.insert(
subagent_id,
ActiveChild {
request: pending.request,
started_at: pending.started_at,
cancellation: pending.cancellation,
spawn_reply: pending.spawn_reply,
foreground_deadline: pending.foreground_deadline,
handle_only: pending.handle_only,
definition_background: child.definition_background,
explicitly_killed: pending.explicitly_killed,
child_session_id: child.child_session_id,
persona: child.persona,
resumed_from: child.resumed_from,
child_cwd: child.child_cwd,
worktree_path: child.worktree_path,
effective_model_id: child.effective_model_id,
control: child.control,
},
);
let _ = respond_to.send(true);
}
InternalEvent::ResumeSource {
source_id,
parent_session_id,
respond_to,
} => {
let source_is_active =
self.pending
.get(&source_id)
.is_some_and(|child| child.request.parent_session_id == parent_session_id)
|| self.active.get(&source_id).is_some_and(|child| {
child.request.parent_session_id == parent_session_id
});
let lookup = if source_is_active {
SubagentResumeLookup::Active
} else if let Some(child) = self.completed.get(&source_id)
&& child.request.parent_session_id == parent_session_id
{
SubagentResumeLookup::Completed(SubagentResumeSource {
subagent_id: child.request.id.clone(),
child_session_id: child.child_session_id.clone(),
child_cwd: child.child_cwd.clone(),
worktree_path: child.worktree_path.clone(),
snapshot_ref: child.snapshot_ref.clone(),
subagent_type: child.request.subagent_type.clone(),
persona: child.persona.clone(),
model_id: Some(child.effective_model_id.clone()),
})
} else {
SubagentResumeLookup::Missing
};
let _ = respond_to.send(lookup);
}
}
}
fn finish_child(&mut self, id: &str, output: ChildRunOutput<R::CompletionData>) {
let record = if let Some(child) = self.active.remove(id) {
ChildRecord::Active(child)
} else if let Some(child) = self.pending.remove(id) {
ChildRecord::Pending(child)
} else {
return;
};
let request = record.request().clone();
let explicitly_killed = record.explicitly_killed();
let (
started_at,
child_session_id,
persona,
resumed_from,
child_cwd,
worktree_path,
effective_model_id,
mut spawn_reply,
mut handle_only,
) = match record {
ChildRecord::Pending(child) => (
child.started_at,
output.result.child_session_id.clone(),
child.request.runtime_overrides.persona.clone(),
child.request.resume_from.clone(),
child.request.cwd.clone().unwrap_or_default(),
output.result.worktree_path.clone(),
String::new(),
child.spawn_reply,
child.handle_only,
),
ChildRecord::Active(child) => (
child.started_at,
child.child_session_id,
child.persona,
child.resumed_from,
child.child_cwd,
child.worktree_path,
child.effective_model_id,
child.spawn_reply,
child.handle_only,
),
};
let persisted_output_ref = self.runner.persisted_output_ref(&output.completion_data);
let mut completed = CompletedChild {
request: request.clone(),
started_at,
child_session_id,
persona,
resumed_from,
child_cwd,
worktree_path,
snapshot_ref: output.snapshot_ref,
persisted_output_ref,
effective_model_id,
result: output.result.clone(),
};
let snapshot = completed_snapshot(&completed, None);
let mut waiter_delivered = false;
for waiter in self.waiters.remove(id).unwrap_or_default() {
waiter_delivered |= waiter.respond_to.send(Some(snapshot.clone())).is_ok();
}
let mut foreground_delivered = false;
if let Some(respond_to) = spawn_reply.take() {
let sent = respond_to.send(output.result.clone()).is_ok();
if !handle_only {
foreground_delivered = sent;
handle_only = !sent;
}
} else if !handle_only {
handle_only = true;
}
if self.config.buffer_completions
&& request.surface_completion
&& !request.owner.is_workflow()
{
let mut summary = completion_summary(&request, &output.result);
if let Some(cap) = self.config.buffered_completion_output_cap {
summary.output = super::cap_completion_output(&summary.output, cap);
}
self.pending_completions.push(BufferedCompletion {
parent_session_id: request.parent_session_id.clone(),
summary,
});
// Bound the buffer (drop oldest): sessions unloaded without a
// DiscardSessionCompletions cannot grow it unboundedly.
const MAX_PENDING_COMPLETIONS: usize = 256;
if self.pending_completions.len() > MAX_PENDING_COMPLETIONS {
let excess = self.pending_completions.len() - MAX_PENDING_COMPLETIONS;
self.pending_completions.drain(..excess);
}
}
if completed.persisted_output_ref.is_some() {
completed.result.output = Arc::from("");
}
let should_surface = request.surface_completion
&& handle_only
&& !output.result.cancelled
&& !waiter_delivered
&& !explicitly_killed;
let disposition = CompletionDisposition {
foreground_delivered,
backgrounded: handle_only,
waiter_delivered,
explicitly_killed,
should_surface,
};
self.completed.insert(id.to_owned(), completed);
self.completed_order.push_back(id.to_owned());
self.running_count_changed();
let workflow_run_id = request.owner.workflow_run_id().map(str::to_owned);
self.runner.on_completed(ChildCompletion {
request,
result: output.result,
completion_data: output.completion_data,
disposition,
});
if let Some(run_id) = workflow_run_id {
self.resolve_workflow_cancel_waiters(&run_id);
}
}
fn finish_panicked_child(&mut self, id: &str) {
let request = self
.active
.get(id)
.map(|child| child.request.clone())
.or_else(|| self.pending.get(id).map(|child| child.request.clone()));
let Some(request) = request else {
return;
};
tracing::error!(subagent_id = id, "subagent child runner panicked");
self.finish_child(
id,
ChildRunOutput {
result: SubagentResult {
success: false,
error: Some("Subagent runtime panicked".to_owned()),
subagent_id: request.id.clone(),
child_session_id: request.id,
..Default::default()
},
completion_data: R::CompletionData::default(),
snapshot_ref: None,
},
);
}
fn cancel_one(
&mut self,
id: &str,
parent_session_id: Option<&str>,
explicit: bool,
) -> SubagentCancelOutcome {
if let Some(child) = self.active.get_mut(id)
&& belongs_to_session(&child.request, parent_session_id)
{
child.explicitly_killed |= explicit;
child.cancellation.cancel();
child.control.cancel();
return SubagentCancelOutcome::Cancelled;
}
if let Some(child) = self.pending.get_mut(id)
&& belongs_to_session(&child.request, parent_session_id)
{
child.explicitly_killed |= explicit;
child.cancellation.cancel();
return SubagentCancelOutcome::Cancelled;
}
if let Some(child) = self.completed.get(id)
&& belongs_to_session(&child.request, parent_session_id)
{
return SubagentCancelOutcome::AlreadyFinished {
status: child.result.status().to_owned(),
};
}
SubagentCancelOutcome::NotFound
}
fn cancel_parent_prompt(&mut self, parent_prompt_id: &str, parent_session_id: Option<&str>) {
for child in self.active.values() {
if child.request.parent_prompt_id.as_deref() == Some(parent_prompt_id)
&& belongs_to_session(&child.request, parent_session_id)
{
child.cancellation.cancel();
child.control.cancel();
}
}
for child in self.pending.values() {
if child.request.parent_prompt_id.as_deref() == Some(parent_prompt_id)
&& belongs_to_session(&child.request, parent_session_id)
{
child.cancellation.cancel();
}
}
}
fn cancel_workflow_children(&mut self, run_id: &str, parent_session_id: Option<&str>) {
for child in self.active.values() {
if child.request.owner.workflow_run_id() == Some(run_id)
&& belongs_to_session(&child.request, parent_session_id)
{
child.cancellation.cancel();
child.control.cancel();
}
}
for child in self.pending.values() {
if child.request.owner.workflow_run_id() == Some(run_id)
&& belongs_to_session(&child.request, parent_session_id)
{
child.cancellation.cancel();
}
}
}
fn resolve_workflow_cancel_waiters(&mut self, run_id: &str) {
if workflow_outstanding(&self.pending, &self.active, run_id) != 0 {
return;
}
for respond_to in self
.workflow_cancel_waiters
.remove(run_id)
.unwrap_or_default()
{
let _ = respond_to.send(SubagentCancelOutcome::Cancelled);
}
}
fn next_deadline(&self) -> Option<tokio::time::Instant> {
self.pending
.values()
.filter_map(|child| child.foreground_deadline)
.chain(
self.active
.values()
.filter_map(|child| child.foreground_deadline),
)
.chain(
self.waiters
.values()
.flatten()
.map(|waiter| waiter.deadline),
)
.min()
}
fn reap_abandoned_callers(&mut self) {
for child in self.pending.values_mut() {
background_if_caller_gone(child);
}
for child in self.active.values_mut() {
background_if_caller_gone(child);
}
}
fn process_deadlines(&mut self) {
self.reap_abandoned_callers();
let now = tokio::time::Instant::now();
for child in self.pending.values_mut() {
background_at_deadline(child, now, self.config.foreground_budget);
}
for child in self.active.values_mut() {
background_at_deadline(child, now, self.config.foreground_budget);
}
let ids: Vec<_> = self.waiters.keys().cloned().collect();
for id in ids {
let waiters = self.waiters.remove(&id).unwrap_or_default();
let (due, live): (Vec<_>, Vec<_>) = waiters
.into_iter()
.partition(|waiter| waiter.deadline <= now);
if !live.is_empty() {
self.waiters.insert(id.clone(), live);
}
for waiter in due {
if waiter.respond_to.is_closed() {
continue;
}
if self.active.contains_key(&id) {
self.queue_active_progress(&id, ProgressTarget::Query(waiter.respond_to));
} else {
let _ = waiter.respond_to.send(self.ready_snapshot(&id));
}
}
}
}
fn running_count_changed(&self) {
self.runner
.running_count_changed(self.pending.len() + self.active.len());
}
fn cancel_all_children(&self) {
for child in self.active.values() {
child.cancellation.cancel();
child.control.cancel();
}
for child in self.pending.values() {
child.cancellation.cancel();
}
}
}
fn belongs_to_session(request: &SubagentRequest, parent_session_id: Option<&str>) -> bool {
parent_session_id.is_none_or(|id| request.parent_session_id == id)
}
impl<R: ChildRunner> Drop for SubagentCoordinator<R> {
fn drop(&mut self) {
self.cancel_all_children();
}
}
#[cfg(test)]
#[path = "coordinator_tests.rs"]
mod tests;

256
crates/codegen/xai-grok-tools/src/implementations/grok_build/task/coordinator/query.rs Normal file
View file

@ -0,0 +1,256 @@
//! Session-scoped query, inspection, and progress delivery.
use std::sync::Arc;
use tokio::sync::oneshot;
use super::super::coordinator_state::{
BlockingWaiter, CompletedChild, ListRequest, OUTPUT_UNAVAILABLE_PLACEHOLDER, ProgressFuture,
ProgressTarget, RunningSeed, completed_inspection, completed_snapshot, pending_inspection,
pending_snapshot, running_inspection, running_seed,
};
use super::super::types::{SubagentInspection, SubagentSnapshot};
use super::{ChildControl, ChildRunner, SubagentCoordinator, SubagentProgress, belongs_to_session};
impl<R: ChildRunner> SubagentCoordinator<R> {
pub(super) fn handle_query(
&mut self,
id: String,
parent_session_id: Option<String>,
block: bool,
timeout_ms: Option<u64>,
respond_to: oneshot::Sender<Option<SubagentSnapshot>>,
) {
if let Some(child) = self
.completed
.get(&id)
.filter(|child| belongs_to_session(&child.request, parent_session_id.as_deref()))
{
let snapshot = (!child.request.owner.is_workflow())
.then(|| self.completed_snapshot_for_query(child));
let _ = respond_to.send(snapshot);
return;
}
if let Some(child) = self
.active
.get(&id)
.filter(|child| belongs_to_session(&child.request, parent_session_id.as_deref()))
{
if child.request.owner.is_workflow() {
let _ = respond_to.send(None);
return;
}
if block {
self.waiters.entry(id).or_default().push(BlockingWaiter {
deadline: tokio::time::Instant::now()
+ std::time::Duration::from_millis(timeout_ms.unwrap_or(30_000)),
respond_to,
});
} else {
self.queue_active_progress(&id, ProgressTarget::Query(respond_to));
}
return;
}
if let Some(child) = self
.pending
.get(&id)
.filter(|child| belongs_to_session(&child.request, parent_session_id.as_deref()))
{
if child.request.owner.is_workflow() {
let _ = respond_to.send(None);
return;
}
if block {
self.waiters.entry(id).or_default().push(BlockingWaiter {
deadline: tokio::time::Instant::now()
+ std::time::Duration::from_millis(timeout_ms.unwrap_or(30_000)),
respond_to,
});
} else {
let _ = respond_to.send(Some(pending_snapshot(child)));
}
return;
}
let _ = respond_to.send(None);
}
pub(super) fn handle_inspect(
&mut self,
id: String,
parent_session_id: Option<String>,
respond_to: oneshot::Sender<Option<SubagentInspection>>,
) {
if let Some(child) = self
.completed
.get(&id)
.filter(|child| belongs_to_session(&child.request, parent_session_id.as_deref()))
{
let _ = respond_to.send(Some(self.completed_inspection_for_query(child)));
} else if let Some(child) = self
.pending
.get(&id)
.filter(|child| belongs_to_session(&child.request, parent_session_id.as_deref()))
{
let _ = respond_to.send(Some(pending_inspection(child)));
} else if self
.active
.get(&id)
.is_some_and(|child| belongs_to_session(&child.request, parent_session_id.as_deref()))
{
self.queue_active_progress(&id, ProgressTarget::Inspect(respond_to));
} else {
let _ = respond_to.send(None);
}
}
fn persisted_output(&self, child: &CompletedChild) -> Option<Arc<str>> {
child.persisted_output_ref.as_deref().map(|reference| {
self.runner
.load_persisted_output(reference)
.unwrap_or_else(|| Arc::from(OUTPUT_UNAVAILABLE_PLACEHOLDER))
})
}
fn completed_snapshot_for_query(&self, child: &CompletedChild) -> SubagentSnapshot {
let output = self.persisted_output(child);
completed_snapshot(child, output.as_deref())
}
fn completed_inspection_for_query(&self, child: &CompletedChild) -> SubagentInspection {
let output = self.persisted_output(child);
completed_inspection(child, output.as_deref())
}
pub(super) fn ready_snapshot(&self, id: &str) -> Option<SubagentSnapshot> {
self.completed
.get(id)
.filter(|child| !child.request.owner.is_workflow())
.map(|child| self.completed_snapshot_for_query(child))
.or_else(|| {
self.pending
.get(id)
.filter(|child| !child.request.owner.is_workflow())
.map(pending_snapshot)
})
}
pub(super) fn handle_list_running(
&mut self,
parent_session_id: String,
respond_to: oneshot::Sender<Vec<SubagentInspection>>,
) {
let ids: Vec<_> = self
.active
.values()
.filter(|child| {
child.request.parent_session_id == parent_session_id
&& !child.request.owner.is_workflow()
})
.map(|child| child.request.id.clone())
.collect();
if ids.is_empty() {
let _ = respond_to.send(Vec::new());
return;
}
let request_id = self.next_list_request_id;
self.next_list_request_id = self.next_list_request_id.wrapping_add(1);
self.list_requests.insert(
request_id,
ListRequest {
slots: vec![None; ids.len()],
remaining: ids.len(),
respond_to,
},
);
for (index, id) in ids.into_iter().enumerate() {
self.queue_active_progress(&id, ProgressTarget::List { request_id, index });
}
}
pub(super) fn queue_active_progress(&mut self, id: &str, target: ProgressTarget) {
let Some(child) = self.active.get(id) else {
match target {
ProgressTarget::Query(tx) => {
let _ = tx.send(self.ready_snapshot(id));
}
ProgressTarget::Inspect(tx) => {
let value = self
.completed
.get(id)
.map(|child| self.completed_inspection_for_query(child));
let _ = tx.send(value);
}
ProgressTarget::List { request_id, index } => {
self.finish_list_slot(request_id, index, None);
}
}
return;
};
self.progress.push(ProgressFuture {
future: Box::pin(child.control.progress()),
seed: Some(running_seed(child)),
target: Some(target),
});
}
pub(super) fn finish_progress(
&mut self,
seed: RunningSeed,
target: ProgressTarget,
progress: SubagentProgress,
) {
let still_active = self.active.contains_key(&seed.subagent_id);
if !still_active {
match target {
ProgressTarget::Query(respond_to) => {
let _ = respond_to.send(self.ready_snapshot(&seed.subagent_id));
}
ProgressTarget::Inspect(respond_to) => {
let value = self
.completed
.get(&seed.subagent_id)
.map(|child| self.completed_inspection_for_query(child));
let _ = respond_to.send(value);
}
ProgressTarget::List { request_id, index } => {
self.finish_list_slot(request_id, index, None);
}
}
return;
}
let inspection = running_inspection(seed, progress);
match target {
ProgressTarget::Query(respond_to) => {
let _ = respond_to.send(Some(inspection.snapshot));
}
ProgressTarget::Inspect(respond_to) => {
let _ = respond_to.send(Some(inspection));
}
ProgressTarget::List { request_id, index } => {
self.finish_list_slot(request_id, index, Some(inspection));
}
}
}
fn finish_list_slot(
&mut self,
request_id: u64,
index: usize,
inspection: Option<SubagentInspection>,
) {
let Some(request) = self.list_requests.get_mut(&request_id) else {
return;
};
request.slots[index] = inspection;
request.remaining = request.remaining.saturating_sub(1);
if request.remaining != 0 {
return;
}
let Some(request) = self.list_requests.remove(&request_id) else {
return;
};
let values = request.slots.into_iter().flatten().collect();
let _ = request.respond_to.send(values);
}
}

731
crates/codegen/xai-grok-tools/src/implementations/grok_build/task/coordinator_state.rs Normal file
View file

@ -0,0 +1,731 @@
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use tokio::sync::{mpsc, oneshot};
use tokio_util::sync::CancellationToken;
use super::types::{
ActiveSubagentSummary, SubagentCompletionSummary, SubagentDescribeOutcome, SubagentInspection,
SubagentRequest, SubagentResult, SubagentResumeLookup, SubagentSnapshot,
SubagentSnapshotStatus, SubagentValidateTypeOutcome,
};
pub(super) const MAX_COMPLETED_ENTRIES: usize = 1024;
pub(super) const OUTPUT_UNAVAILABLE_PLACEHOLDER: &str = "[subagent output no longer available]";
pub type LocalBoxFuture<T> = Pin<Box<dyn Future<Output = T> + 'static>>;
pub type SendBoxFuture<T> = Pin<Box<dyn Future<Output = T> + Send + 'static>>;
/// Runtime-specific live progress for one active child.
#[derive(Debug, Clone, Default)]
pub struct SubagentProgress {
pub turn_count: u32,
pub tool_call_count: u32,
pub tokens_used: u64,
pub context_window_tokens: u64,
pub context_usage_pct: u8,
pub tools_used: Vec<String>,
pub error_count: u32,
}
/// Runtime handle retained while a child is active.
pub trait ChildControl: 'static {
type ProgressFuture: Future<Output = SubagentProgress> + 'static;
fn progress(&self) -> Self::ProgressFuture;
fn cancel(&self);
}
/// Data reported when runtime initialization has produced a live child.
pub struct StartedChild<C> {
pub child_session_id: String,
pub persona: Option<String>,
pub resumed_from: Option<String>,
pub child_cwd: String,
pub worktree_path: Option<String>,
pub effective_model_id: String,
/// The resolved agent definition declares `background: true`. Folded into
/// `Outstanding` accounting (background, never turn-blocking) while the
/// foreground await budget stays gated on the tool's own
/// `run_in_background` flag.
pub definition_background: bool,
pub control: C,
}
/// Input to one runtime-specific child run.
pub struct ChildRunRequest<C> {
pub request: SubagentRequest,
pub cancellation: CancellationToken,
pub reporter: ChildReporter<C>,
}
/// Terminal output from one runtime-specific child run.
pub struct ChildRunOutput<D> {
pub result: SubagentResult,
pub completion_data: D,
pub snapshot_ref: Option<String>,
}
/// Coordinator-owned delivery decision passed to host presentation.
#[derive(Debug, Clone)]
pub struct CompletionDisposition {
pub foreground_delivered: bool,
pub backgrounded: bool,
pub waiter_delivered: bool,
pub explicitly_killed: bool,
pub should_surface: bool,
}
/// Terminal event delivered to the runtime adapter after state is committed.
pub struct ChildCompletion<D> {
pub request: SubagentRequest,
pub result: SubagentResult,
pub completion_data: D,
pub disposition: CompletionDisposition,
}
/// The only host-specific seam.
///
/// Associated future types intentionally carry no unconditional `Send` bound.
/// A local runner may return non-`Send` futures, while a multithreaded runner
/// may return `Send` futures.
pub trait ChildRunner: 'static {
type Control: ChildControl;
type CompletionData: Default + 'static;
type RunFuture: Future<Output = ChildRunOutput<Self::CompletionData>> + 'static;
type ValidateFuture: Future<Output = SubagentValidateTypeOutcome> + 'static;
type DescribeFuture: Future<Output = SubagentDescribeOutcome> + 'static;
fn run(&self, request: ChildRunRequest<Self::Control>) -> Self::RunFuture;
fn validate_type(
&self,
subagent_type: String,
parent_session_id: String,
) -> Self::ValidateFuture;
fn describe_type(
&self,
subagent_type: String,
harness_agent_type: Option<String>,
parent_session_id: String,
) -> Self::DescribeFuture;
fn on_completed(&self, completion: ChildCompletion<Self::CompletionData>);
fn running_count_changed(&self, _running: usize) {}
fn persisted_output_ref(&self, _completion_data: &Self::CompletionData) -> Option<String> {
None
}
fn load_persisted_output(&self, _reference: &str) -> Option<Arc<str>> {
None
}
}
/// Host-configurable lifecycle policy. The transition logic remains shared.
#[derive(Debug, Clone)]
pub struct CoordinatorConfig {
pub foreground_budget: std::time::Duration,
/// Whether the host drains completion summaries between turns.
pub buffer_completions: bool,
/// Extra cap applied to BUFFERED summary outputs only (the request's own
/// `completion_output_cap` still applies first). Buffered entries pin the
/// child's output `Arc` until drained; hosts whose reminder rendering
/// never inlines the output (a polling tool exists, e.g. the callback
/// tools-server) should bound it. `None` keeps outputs verbatim — the
/// shell needs this for toolsets with no polling tool, where the inline
/// reminder is the model's only chance to see the output.
pub buffered_completion_output_cap: Option<usize>,
}
impl Default for CoordinatorConfig {
fn default() -> Self {
Self {
foreground_budget: std::time::Duration::from_secs(45),
buffer_completions: false,
buffered_completion_output_cap: None,
}
}
}
/// Runner-side channel back into the actor.
pub struct ChildReporter<C> {
pub(super) subagent_id: String,
pub(super) tx: mpsc::UnboundedSender<InternalEvent<C>>,
}
impl<C> Clone for ChildReporter<C> {
fn clone(&self) -> Self {
Self {
subagent_id: self.subagent_id.clone(),
tx: self.tx.clone(),
}
}
}
impl<C: 'static> ChildReporter<C> {
/// Promote the pending child to active. The acknowledgement closes the
/// cancel-at-promote race: `false` means cancellation won and the adapter
/// must tear down the half-initialized runtime.
pub async fn started(&self, child: StartedChild<C>) -> bool {
let (respond_to, response_rx) = oneshot::channel();
if self
.tx
.send(InternalEvent::Started {
subagent_id: self.subagent_id.clone(),
child,
respond_to,
})
.is_err()
{
return false;
}
response_rx.await.unwrap_or(false)
}
/// Resolve an in-memory resume source without sharing coordinator state.
pub async fn resume_source(
&self,
source_id: &str,
parent_session_id: &str,
) -> SubagentResumeLookup {
let (respond_to, response_rx) = oneshot::channel();
if self
.tx
.send(InternalEvent::ResumeSource {
source_id: source_id.to_owned(),
parent_session_id: parent_session_id.to_owned(),
respond_to,
})
.is_err()
{
return SubagentResumeLookup::Missing;
}
response_rx.await.unwrap_or(SubagentResumeLookup::Missing)
}
}
pub(super) enum InternalEvent<C> {
Started {
subagent_id: String,
child: StartedChild<C>,
respond_to: oneshot::Sender<bool>,
},
ResumeSource {
source_id: String,
parent_session_id: String,
respond_to: oneshot::Sender<SubagentResumeLookup>,
},
}
pub(super) struct PendingChild {
pub(super) request: SubagentRequest,
pub(super) started_at: std::time::Instant,
pub(super) cancellation: CancellationToken,
pub(super) spawn_reply: Option<oneshot::Sender<SubagentResult>>,
pub(super) foreground_deadline: Option<tokio::time::Instant>,
pub(super) handle_only: bool,
pub(super) explicitly_killed: bool,
}
pub(super) struct ActiveChild<C> {
pub(super) request: SubagentRequest,
pub(super) started_at: std::time::Instant,
pub(super) cancellation: CancellationToken,
pub(super) spawn_reply: Option<oneshot::Sender<SubagentResult>>,
pub(super) foreground_deadline: Option<tokio::time::Instant>,
pub(super) handle_only: bool,
/// Definition-declared background (see [`StartedChild`]): background for
/// `Outstanding` accounting even while the spawn caller block-awaits.
pub(super) definition_background: bool,
pub(super) explicitly_killed: bool,
pub(super) child_session_id: String,
pub(super) persona: Option<String>,
pub(super) resumed_from: Option<String>,
pub(super) child_cwd: String,
pub(super) worktree_path: Option<String>,
pub(super) effective_model_id: String,
pub(super) control: C,
}
pub(super) struct CompletedChild {
pub(super) request: SubagentRequest,
pub(super) started_at: std::time::Instant,
pub(super) child_session_id: String,
pub(super) persona: Option<String>,
pub(super) resumed_from: Option<String>,
pub(super) child_cwd: String,
pub(super) worktree_path: Option<String>,
pub(super) snapshot_ref: Option<String>,
pub(super) persisted_output_ref: Option<String>,
pub(super) effective_model_id: String,
pub(super) result: SubagentResult,
}
pub(super) struct BlockingWaiter {
pub(super) deadline: tokio::time::Instant,
pub(super) respond_to: oneshot::Sender<Option<SubagentSnapshot>>,
}
pub(super) struct BufferedCompletion {
pub(super) parent_session_id: String,
pub(super) summary: SubagentCompletionSummary,
}
pub(super) struct TaggedFuture<F> {
pub(super) subagent_id: String,
pub(super) future: Pin<Box<F>>,
}
impl<F: Future> Future for TaggedFuture<F> {
type Output = (String, F::Output);
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.get_mut();
this.future
.as_mut()
.poll(cx)
.map(|output| (this.subagent_id.clone(), output))
}
}
pub(super) struct ReplyFuture<F, T> {
pub(super) future: Pin<Box<F>>,
pub(super) respond_to: Option<oneshot::Sender<T>>,
}
impl<F, T> Future for ReplyFuture<F, T>
where
F: Future<Output = T>,
{
type Output = (oneshot::Sender<T>, T);
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.get_mut();
this.future.as_mut().poll(cx).map(|output| {
let respond_to = match this.respond_to.take() {
Some(respond_to) => respond_to,
None => unreachable!("reply future polled after completion"),
};
(respond_to, output)
})
}
}
#[derive(Clone)]
pub(super) struct RunningSeed {
pub(super) subagent_id: String,
pub(super) description: String,
pub(super) subagent_type: String,
pub(super) started_at_epoch_ms: u64,
pub(super) duration_ms: u64,
pub(super) persona: Option<String>,
pub(super) parent_session_id: String,
pub(super) child_session_id: String,
pub(super) fork_parent_prompt_id: Option<String>,
pub(super) resumed_from: Option<String>,
}
pub(super) enum ProgressTarget {
Query(oneshot::Sender<Option<SubagentSnapshot>>),
Inspect(oneshot::Sender<Option<SubagentInspection>>),
List { request_id: u64, index: usize },
}
pub(super) struct ProgressFuture<F> {
pub(super) future: Pin<Box<F>>,
pub(super) seed: Option<RunningSeed>,
pub(super) target: Option<ProgressTarget>,
}
impl<F> Future for ProgressFuture<F>
where
F: Future<Output = SubagentProgress>,
{
type Output = (RunningSeed, ProgressTarget, SubagentProgress);
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.get_mut();
this.future.as_mut().poll(cx).map(|progress| {
let seed = match this.seed.take() {
Some(seed) => seed,
None => unreachable!("progress future polled without a seed"),
};
let target = match this.target.take() {
Some(target) => target,
None => unreachable!("progress future polled without a target"),
};
(seed, target, progress)
})
}
}
pub(super) struct ListRequest {
pub(super) slots: Vec<Option<SubagentInspection>>,
pub(super) remaining: usize,
pub(super) respond_to: oneshot::Sender<Vec<SubagentInspection>>,
}
pub(super) enum ChildRecord<C> {
Pending(PendingChild),
Active(ActiveChild<C>),
}
impl<C> ChildRecord<C> {
pub(super) fn request(&self) -> &SubagentRequest {
match self {
Self::Pending(child) => &child.request,
Self::Active(child) => &child.request,
}
}
pub(super) fn explicitly_killed(&self) -> bool {
match self {
Self::Pending(child) => child.explicitly_killed,
Self::Active(child) => child.explicitly_killed,
}
}
}
pub(super) trait ForegroundChild {
fn id(&self) -> &str;
fn child_session_id(&self) -> &str;
fn deadline(&self) -> Option<tokio::time::Instant>;
/// True when the spawn caller dropped its result receiver while this
/// child was still treated as turn-blocking (old shell `ParentGone`).
fn caller_gone(&self) -> bool;
fn is_workflow(&self) -> bool;
fn take_reply(&mut self) -> Option<oneshot::Sender<SubagentResult>>;
fn mark_backgrounded(&mut self);
/// Cancel the child's execution (token + active control where present).
fn cancel(&mut self);
}
impl ForegroundChild for PendingChild {
fn id(&self) -> &str {
&self.request.id
}
fn child_session_id(&self) -> &str {
&self.request.id
}
fn deadline(&self) -> Option<tokio::time::Instant> {
self.foreground_deadline
}
fn caller_gone(&self) -> bool {
!self.handle_only && self.spawn_reply.as_ref().is_some_and(|tx| tx.is_closed())
}
fn is_workflow(&self) -> bool {
self.request.owner.is_workflow()
}
fn take_reply(&mut self) -> Option<oneshot::Sender<SubagentResult>> {
self.spawn_reply.take()
}
fn mark_backgrounded(&mut self) {
self.handle_only = true;
self.foreground_deadline = None;
}
fn cancel(&mut self) {
self.cancellation.cancel();
}
}
impl<C: ChildControl> ForegroundChild for ActiveChild<C> {
fn id(&self) -> &str {
&self.request.id
}
fn child_session_id(&self) -> &str {
&self.child_session_id
}
fn deadline(&self) -> Option<tokio::time::Instant> {
self.foreground_deadline
}
fn caller_gone(&self) -> bool {
!self.handle_only && self.spawn_reply.as_ref().is_some_and(|tx| tx.is_closed())
}
fn is_workflow(&self) -> bool {
self.request.owner.is_workflow()
}
fn take_reply(&mut self) -> Option<oneshot::Sender<SubagentResult>> {
self.spawn_reply.take()
}
fn mark_backgrounded(&mut self) {
self.handle_only = true;
self.foreground_deadline = None;
}
fn cancel(&mut self) {
self.cancellation.cancel();
self.control.cancel();
}
}
pub(super) fn background_at_deadline(
child: &mut impl ForegroundChild,
now: tokio::time::Instant,
budget: std::time::Duration,
) {
if child.deadline().is_none_or(|deadline| deadline > now) {
return;
}
tracing::warn!(
subagent_id = child.id(),
budget_ms = budget.as_millis() as u64,
"foreground subagent exceeded await budget; auto-backgrounding (child keeps running)",
);
if let Some(respond_to) = child.take_reply() {
// Interim handoff, not a completion: keep `success: false` (default)
// so `SubagentResult::status()` consumers cannot record a completed
// status for a still-running child. Callers branch on `backgrounded`.
let _ = respond_to.send(SubagentResult {
backgrounded: true,
subagent_id: child.id().to_owned(),
child_session_id: child.child_session_id().to_owned(),
..Default::default()
});
}
child.mark_backgrounded();
}
/// Handle a foreground child whose spawn caller dropped the result channel
/// (parent turn stop / cancelled await). Task-owned children keep running and
/// just leave the turn-blocking `Outstanding` set — shell `ParentGone` parity.
/// Workflow-owned children are CANCELLED instead (old shell `ParentGone`
/// cancelled workflow children); `ChannelBackend`'s drop-cancel arming remains
/// defense in depth for hosts that go through it.
pub(super) fn background_if_caller_gone(child: &mut impl ForegroundChild) {
if !child.caller_gone() {
return;
}
let _ = child.take_reply();
if child.is_workflow() {
tracing::debug!(
subagent_id = child.id(),
"workflow subagent caller gone; cancelling child",
);
child.cancel();
return;
}
tracing::debug!(
subagent_id = child.id(),
"foreground subagent caller gone; auto-backgrounding (child keeps running)",
);
child.mark_backgrounded();
}
pub(super) async fn sleep_until(deadline: Option<tokio::time::Instant>) {
match deadline {
Some(deadline) => tokio::time::sleep_until(deadline).await,
None => std::future::pending().await,
}
}
fn instant_to_epoch_ms(instant: std::time::Instant) -> u64 {
let now_instant = std::time::Instant::now();
let now_system = std::time::SystemTime::now();
let elapsed = now_instant.saturating_duration_since(instant);
now_system
.checked_sub(elapsed)
.unwrap_or(now_system)
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64
}
pub(super) fn active_summary<C>(child: &ActiveChild<C>) -> ActiveSubagentSummary {
ActiveSubagentSummary {
subagent_id: child.request.id.clone(),
subagent_type: child.request.subagent_type.clone(),
description: child.request.description.clone(),
elapsed_ms: child.started_at.elapsed().as_millis() as u64,
}
}
pub(super) fn running_seed<C>(child: &ActiveChild<C>) -> RunningSeed {
RunningSeed {
subagent_id: child.request.id.clone(),
description: child.request.description.clone(),
subagent_type: child.request.subagent_type.clone(),
started_at_epoch_ms: instant_to_epoch_ms(child.started_at),
duration_ms: child.started_at.elapsed().as_millis() as u64,
persona: child.persona.clone(),
parent_session_id: child.request.parent_session_id.clone(),
child_session_id: child.child_session_id.clone(),
fork_parent_prompt_id: child.request.parent_prompt_id.clone(),
resumed_from: child.resumed_from.clone(),
}
}
pub(super) fn running_inspection(
seed: RunningSeed,
progress: SubagentProgress,
) -> SubagentInspection {
SubagentInspection {
snapshot: SubagentSnapshot {
subagent_id: seed.subagent_id,
description: seed.description,
subagent_type: seed.subagent_type,
status: SubagentSnapshotStatus::Running {
turn_count: progress.turn_count,
tool_call_count: progress.tool_call_count,
tokens_used: progress.tokens_used,
context_window_tokens: progress.context_window_tokens,
context_usage_pct: progress.context_usage_pct,
tools_used: progress.tools_used,
error_count: progress.error_count,
},
started_at_epoch_ms: seed.started_at_epoch_ms,
duration_ms: seed.duration_ms,
persona: seed.persona,
},
parent_session_id: seed.parent_session_id,
child_session_id: seed.child_session_id,
fork_parent_prompt_id: seed.fork_parent_prompt_id,
resumed_from: seed.resumed_from,
}
}
pub(super) fn pending_snapshot(child: &PendingChild) -> SubagentSnapshot {
SubagentSnapshot {
subagent_id: child.request.id.clone(),
description: child.request.description.clone(),
subagent_type: child.request.subagent_type.clone(),
status: SubagentSnapshotStatus::Initializing,
started_at_epoch_ms: instant_to_epoch_ms(child.started_at),
duration_ms: child.started_at.elapsed().as_millis() as u64,
persona: child.request.runtime_overrides.persona.clone(),
}
}
pub(super) fn pending_inspection(child: &PendingChild) -> SubagentInspection {
SubagentInspection {
snapshot: pending_snapshot(child),
parent_session_id: child.request.parent_session_id.clone(),
child_session_id: String::new(),
fork_parent_prompt_id: child.request.parent_prompt_id.clone(),
resumed_from: child.request.resume_from.clone(),
}
}
pub(super) fn completed_snapshot(
child: &CompletedChild,
persisted_output: Option<&str>,
) -> SubagentSnapshot {
let status = if child.result.cancelled {
SubagentSnapshotStatus::Cancelled {
reason: child.result.error.clone(),
}
} else if child.result.success {
SubagentSnapshotStatus::Completed {
output: persisted_output
.map(str::to_owned)
.unwrap_or_else(|| child.result.output.to_string()),
tool_calls: child.result.tool_calls,
turns: child.result.turns,
worktree_path: child.result.worktree_path.clone(),
}
} else {
SubagentSnapshotStatus::Failed {
error: child
.result
.error
.clone()
.unwrap_or_else(|| "Unknown error".to_owned()),
}
};
SubagentSnapshot {
subagent_id: child.request.id.clone(),
description: child.request.description.clone(),
subagent_type: child.request.subagent_type.clone(),
status,
started_at_epoch_ms: instant_to_epoch_ms(child.started_at),
duration_ms: child.result.duration_ms,
persona: child.persona.clone(),
}
}
pub(super) fn completed_inspection(
child: &CompletedChild,
persisted_output: Option<&str>,
) -> SubagentInspection {
SubagentInspection {
snapshot: completed_snapshot(child, persisted_output),
parent_session_id: child.request.parent_session_id.clone(),
child_session_id: child.child_session_id.clone(),
fork_parent_prompt_id: child.request.parent_prompt_id.clone(),
resumed_from: child.resumed_from.clone(),
}
}
/// Truncate `output` to `cap` bytes (UTF-8 safe) with a truncation footer.
/// Returns a refcount clone when already within the cap.
pub fn cap_completion_output(output: &Arc<str>, cap: usize) -> Arc<str> {
if output.len() <= cap {
return output.clone();
}
let mut end = cap;
while end > 0 && !output.is_char_boundary(end) {
end -= 1;
}
Arc::from(format!(
"{}\n[output truncated: {} of {} bytes shown]",
&output[..end],
end,
output.len()
))
}
/// Model-facing summary for a finished child, honoring the request's
/// `completion_output_cap`. Shared by the coordinator's buffered reminder
/// path and the shell's auto-wake synthetic prompt.
pub fn completion_summary(
request: &SubagentRequest,
result: &SubagentResult,
) -> SubagentCompletionSummary {
let output = match request.runtime_overrides.completion_output_cap {
Some(cap) => cap_completion_output(&result.output, cap),
None => result.output.clone(),
};
SubagentCompletionSummary {
subagent_id: request.id.clone(),
subagent_type: request.subagent_type.clone(),
description: request.description.clone(),
success: result.success && !result.cancelled,
duration_ms: result.duration_ms,
tool_calls: result.tool_calls,
turns: result.turns,
output,
}
}
pub(super) fn workflow_outstanding<C>(
pending: &HashMap<String, PendingChild>,
active: &HashMap<String, ActiveChild<C>>,
run_id: &str,
) -> usize {
pending
.values()
.filter(|child| child.request.owner.workflow_run_id() == Some(run_id))
.count()
+ active
.values()
.filter(|child| child.request.owner.workflow_run_id() == Some(run_id))
.count()
}

1191
crates/codegen/xai-grok-tools/src/implementations/grok_build/task/coordinator_tests.rs Normal file
View file

File diff suppressed because it is too large Load diff

141
crates/codegen/xai-grok-tools/src/implementations/grok_build/task/mod.rs
View file

@ -2,17 +2,21 @@
//!
//! The TaskTool delegates subagent operations to a [`SubagentBackend`]
//! (injected as [`SubagentBackendResource`]). The backend abstracts over the
//! transport mechanism (in-process channels for the local host, remote
//! backends, etc.).
//! coordinator mailbox. All hosts use the same backend and coordinator actor;
//! only their child runners differ.
//!
//! ## Resources
//!
//! - `SubagentBackendResource` — backend for spawn/query/cancel (required)
//! - `SubagentDepthCounter` — current nesting depth (optional, defaults to 0)
//! - `SessionIdResource` — current session ID for parent scoping (optional)
//! - `SubagentForegroundWait` — host wait-window guard factory (optional)
//! - `TaskModelValidator` — validates explicit model slugs before spawn
pub mod backend;
pub mod coordinator;
mod coordinator_state;
pub use coordinator_state::{cap_completion_output, completion_summary};
pub mod types;
use self::backend::SubagentBackendResource;
@ -116,9 +120,12 @@ impl xai_tool_runtime::Tool for TaskTool {
) -> Result<ToolOutput, xai_tool_runtime::ToolError> {
use crate::types::tool_metadata::shared_resources;
let resources = shared_resources(&ctx)?;
let tool_cancellation = ctx
.get::<xai_tool_runtime::Cancellation>()
.map(|cancellation| cancellation.0.clone());
// 1. Depth check
let (depth, backend, model_validator, parent_session_id, parent_prompt_id) = {
let (depth, backend, model_validator, parent_session_id, parent_prompt_id, foreground_wait) = {
let res = resources.lock().await;
let depth = res.get::<SubagentDepthCounter>().map(|d| d.0).unwrap_or(0);
@ -144,6 +151,7 @@ impl xai_tool_runtime::Tool for TaskTool {
.get::<CurrentPromptIdResource>()
.map(|p| p.0.clone())
.filter(|prompt_id| !prompt_id.is_empty());
let foreground_wait = res.get::<SubagentForegroundWait>().cloned();
(
depth,
@ -151,6 +159,7 @@ impl xai_tool_runtime::Tool for TaskTool {
model_validator,
parent_session_id,
parent_prompt_id,
foreground_wait,
)
};
@ -289,9 +298,18 @@ impl xai_tool_runtime::Tool for TaskTool {
.task_id
.clone()
.unwrap_or_else(|| uuid::Uuid::now_v7().to_string());
// Placeholder; `ChannelBackend::spawn` replaces it with a fresh one.
let (result_tx, _) = tokio::sync::oneshot::channel();
let child_cancellation = tokio_util::sync::CancellationToken::new();
let cancellation_forwarder = (!input.run_in_background)
.then(|| {
tool_cancellation.map(|tool_cancellation| {
let child_cancellation = child_cancellation.clone();
tokio::spawn(async move {
tool_cancellation.cancelled().await;
child_cancellation.cancel();
})
})
})
.flatten();
let request = SubagentRequest {
id: id.clone(),
@ -325,8 +343,7 @@ impl xai_tool_runtime::Tool for TaskTool {
await_to_completion: false,
fork_context: false,
owner: SubagentOwner::Task,
cancel_token: tokio_util::sync::CancellationToken::new(),
result_tx,
cancel_token: child_cancellation,
};
// 4. Background mode: fire-and-forget via backend.spawn().
@ -377,7 +394,12 @@ impl xai_tool_runtime::Tool for TaskTool {
}
// 5. Blocking mode (default): spawn via backend and await result
let result = backend.backend().spawn(request).await?;
let _foreground_wait = foreground_wait.map(|wait| wait.enter());
let result = backend.backend().spawn(request).await;
if let Some(forwarder) = cancellation_forwarder {
forwarder.abort();
}
let result = result?;
// 5b. The await budget expired and the coordinator auto-backgrounded the
// still-running child — return a task_id to poll, like the background
@ -495,10 +517,10 @@ mod tests {
(backend, proxy_rx)
}
/// Extract a `SubagentRequest` from a `SubagentEvent`, panicking on wrong variant.
fn unwrap_spawn(event: SubagentEvent) -> SubagentRequest {
/// Extract a spawn envelope from a `SubagentEvent`.
fn unwrap_spawn(event: SubagentEvent) -> SubagentSpawnRequest {
match event {
SubagentEvent::Spawn(r) => *r,
SubagentEvent::Spawn(r) => r,
_ => panic!("Expected SubagentEvent::Spawn"),
}
}
@ -621,8 +643,7 @@ mod tests {
assert_eq!(request.parent_session_id, "parent-session");
assert_eq!(request.parent_prompt_id.as_deref(), Some("prompt-123"));
request
.result_tx
.send(SubagentResult {
.respond_with(|request| SubagentResult {
success: true,
output: std::sync::Arc::from("Found 3 auth middleware files"),
subagent_id: request.id.clone(),
@ -683,8 +704,7 @@ mod tests {
let handle = tokio::spawn(async move {
let request = unwrap_spawn(rx.recv().await.unwrap());
request
.result_tx
.send(SubagentResult {
.respond_with(|_| SubagentResult {
success: false,
error: Some("Child session crashed".to_string()),
..Default::default()
@ -765,11 +785,22 @@ mod tests {
#[tokio::test]
async fn auto_backgrounded_result_returns_task_id_text() {
let (backend, mut rx) = make_backend();
let resources = resources_for_task(backend);
let mut resources = resources_for_task(backend);
let wait_closed = Arc::new(std::sync::atomic::AtomicBool::new(false));
struct WaitProbe(Arc<std::sync::atomic::AtomicBool>);
impl Drop for WaitProbe {
fn drop(&mut self) {
self.0.store(true, std::sync::atomic::Ordering::Relaxed);
}
}
let wait_closed_for_factory = Arc::clone(&wait_closed);
resources.insert(SubagentForegroundWait::new(move || {
Box::new(WaitProbe(Arc::clone(&wait_closed_for_factory)))
}));
let drain = tokio::spawn(async move {
if let Some(SubagentEvent::Spawn(boxed)) = rx.recv().await {
let _ = boxed.result_tx.send(SubagentResult {
let _ = boxed.respond_with(|boxed| SubagentResult {
backgrounded: true,
subagent_id: boxed.id.clone(),
child_session_id: boxed.id.clone(),
@ -785,6 +816,10 @@ mod tests {
)
.await
.expect("auto-backgrounded blocking spawn returns Ok");
assert!(
wait_closed.load(std::sync::atomic::Ordering::Relaxed),
"auto-backgrounding must close the foreground wait window"
);
match result {
ToolOutput::Text(text) => {
@ -982,7 +1017,7 @@ mod tests {
let drain = tokio::spawn(async move {
if let Some(SubagentEvent::Spawn(boxed)) = rx.recv().await {
let _ = boxed.result_tx.send(SubagentResult {
let _ = boxed.respond_with(|boxed| SubagentResult {
success: true,
output: std::sync::Arc::from(""),
subagent_id: boxed.id.clone(),
@ -1023,7 +1058,7 @@ mod tests {
let (done_tx, done_rx) = tokio::sync::oneshot::channel::<()>();
let drain = tokio::spawn(async move {
if let Some(SubagentEvent::Spawn(boxed)) = rx.recv().await {
let _ = boxed.result_tx.send(SubagentResult {
let _ = boxed.respond_with(|boxed| SubagentResult {
success: false,
error: Some("worktree creation failed".to_string()),
subagent_id: boxed.id.clone(),
@ -1502,8 +1537,7 @@ mod tests {
"model-spawned task must not set fork_context"
);
request
.result_tx
.send(SubagentResult {
.respond_with(|request| SubagentResult {
success: true,
output: "ok".into(),
subagent_id: request.id.clone(),
@ -1585,8 +1619,7 @@ mod tests {
let request = unwrap_spawn(rx.recv().await.unwrap());
assert_eq!(request.resume_from.as_deref(), Some("prev-id"));
request
.result_tx
.send(SubagentResult {
.respond_with(|request| SubagentResult {
success: true,
output: "resumed".into(),
subagent_id: request.id.clone(),
@ -1652,8 +1685,7 @@ mod tests {
request.resume_from
);
request
.result_tx
.send(SubagentResult {
.respond_with(|request| SubagentResult {
success: true,
output: "fresh".into(),
subagent_id: request.id.clone(),
@ -1781,8 +1813,7 @@ mod tests {
request.cwd
);
request
.result_tx
.send(SubagentResult {
.respond_with(|request| SubagentResult {
success: true,
output: "ok".into(),
subagent_id: request.id.clone(),
@ -1832,8 +1863,7 @@ mod tests {
request.cwd
);
request
.result_tx
.send(SubagentResult {
.respond_with(|request| SubagentResult {
success: true,
output: "ok".into(),
subagent_id: request.id.clone(),
@ -1883,8 +1913,7 @@ mod tests {
request.cwd
);
request
.result_tx
.send(SubagentResult {
.respond_with(|request| SubagentResult {
success: true,
output: "ok".into(),
subagent_id: request.id.clone(),
@ -1937,8 +1966,7 @@ mod tests {
request.cwd
);
request
.result_tx
.send(SubagentResult {
.respond_with(|request| SubagentResult {
success: true,
output: "ok".into(),
subagent_id: request.id.clone(),
@ -2027,8 +2055,7 @@ mod tests {
request.cwd
);
request
.result_tx
.send(SubagentResult {
.respond_with(|request| SubagentResult {
success: true,
output: "ok".into(),
subagent_id: request.id.clone(),
@ -2081,8 +2108,7 @@ mod tests {
let request = unwrap_spawn(rx.recv().await.unwrap());
assert_eq!(request.cwd.as_deref(), Some("/tmp"));
request
.result_tx
.send(SubagentResult {
.respond_with(|request| SubagentResult {
success: true,
output: "done".into(),
subagent_id: request.id.clone(),
@ -2139,8 +2165,7 @@ mod tests {
"stray leading quote should be stripped before reaching the backend",
);
request
.result_tx
.send(SubagentResult {
.respond_with(|request| SubagentResult {
success: true,
output: "ok".into(),
subagent_id: request.id.clone(),
@ -2192,8 +2217,7 @@ mod tests {
let request = unwrap_spawn(rx.recv().await.unwrap());
assert_eq!(request.cwd.as_deref(), Some("/tmp"));
request
.result_tx
.send(SubagentResult {
.respond_with(|request| SubagentResult {
success: true,
output: "ok".into(),
subagent_id: request.id.clone(),
@ -2241,8 +2265,7 @@ mod tests {
assert_eq!(request.cwd.as_deref(), Some("/tmp/some-dir"));
assert_eq!(request.resume_from.as_deref(), Some("prev-id"));
request
.result_tx
.send(SubagentResult {
.respond_with(|request| SubagentResult {
success: true,
output: "resumed".into(),
subagent_id: request.id.clone(),
@ -2301,13 +2324,14 @@ mod tests {
);
assert!(request.runtime_overrides.reasoning_effort.is_none());
assert!(request.runtime_overrides.persona.is_none());
let id = request.id.clone();
request
.result_tx
.send(SubagentResult {
success: true,
output: "ok".into(),
subagent_id: request.id.clone(),
child_session_id: request.id.clone(),
subagent_id: id.clone(),
child_session_id: id,
..Default::default()
})
.unwrap();
@ -2338,13 +2362,14 @@ mod tests {
"omitted model must stay None, got {:?}",
request.runtime_overrides.model
);
let id = request.id.clone();
request
.result_tx
.send(SubagentResult {
success: true,
output: "ok".into(),
subagent_id: request.id.clone(),
child_session_id: request.id.clone(),
subagent_id: id.clone(),
child_session_id: id,
..Default::default()
})
.unwrap();
@ -2387,13 +2412,14 @@ mod tests {
"sentinel {sentinel:?} must normalize to None, got {:?}",
request.runtime_overrides.model
);
let id = request.id.clone();
request
.result_tx
.send(SubagentResult {
success: true,
output: "ok".into(),
subagent_id: request.id.clone(),
child_session_id: request.id.clone(),
subagent_id: id.clone(),
child_session_id: id,
..Default::default()
})
.unwrap();
@ -2427,13 +2453,14 @@ mod tests {
Some("test-model"),
"leading/trailing whitespace should be trimmed"
);
let id = request.id.clone();
request
.result_tx
.send(SubagentResult {
success: true,
output: "ok".into(),
subagent_id: request.id.clone(),
child_session_id: request.id.clone(),
subagent_id: id.clone(),
child_session_id: id,
..Default::default()
})
.unwrap();
@ -2467,13 +2494,14 @@ mod tests {
);
assert!(request.runtime_overrides.reasoning_effort.is_none());
assert!(request.runtime_overrides.persona.is_none());
let id = request.id.clone();
request
.result_tx
.send(SubagentResult {
success: true,
output: "resumed".into(),
subagent_id: request.id.clone(),
child_session_id: request.id.clone(),
subagent_id: id.clone(),
child_session_id: id,
..Default::default()
})
.unwrap();
@ -2502,13 +2530,14 @@ mod tests {
let request = unwrap_spawn(rx.recv().await.unwrap());
assert_eq!(request.resume_from.as_deref(), Some("prev-id"));
assert!(request.runtime_overrides.model.is_none());
let id = request.id.clone();
request
.result_tx
.send(SubagentResult {
success: true,
output: "resumed".into(),
subagent_id: request.id.clone(),
child_session_id: request.id.clone(),
subagent_id: id.clone(),
child_session_id: id,
..Default::default()
})
.unwrap();

241
crates/codegen/xai-grok-tools/src/implementations/grok_build/task/types.rs
View file

@ -1,7 +1,8 @@
//! Channel types for subagent communication (TaskTool ↔ MvpAgent coordinator).
//! Data and channel types for subagent coordination.
//!
//! These types define the request/response protocol between the `TaskTool`
//! (in `xai-grok-tools`) and the subagent coordinator (in `xai-grok-shell`).
//! Request data is deliberately separate from command reply envelopes. The
//! shared coordinator actor owns every reply sender and every lifecycle
//! transition; child runners receive only plain request data.
//!
//! ## Resource types
//!
@ -23,6 +24,8 @@ use tokio::sync::{mpsc, oneshot};
use tokio_util::sync::CancellationToken;
use xai_tool_types::{SubagentCapabilityMode, SubagentIsolationMode, WaitMode};
use crate::register_resource;
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub enum SubagentOwner {
#[default]
@ -51,13 +54,10 @@ impl SubagentOwner {
}
}
use crate::register_resource;
// Request / Response
/// Request emitted by TaskTool, received by MvpAgent coordinator.
#[derive(Educe)]
#[educe(Debug)]
/// Plain spawn request emitted by `TaskTool`.
#[derive(Debug, Clone)]
pub struct SubagentRequest {
/// Subagent ID (UUID v7). Same as `TaskToolInput.task_id`; becomes the child session ID.
pub id: String,
@ -75,15 +75,17 @@ pub struct SubagentRequest {
/// freshly rendered.
pub resume_from: Option<String>,
/// Explicit working directory for the child session.
/// Validated at spawn time in `handle_subagent_request()`.
/// Validated at spawn time by the injected child runner.
pub cwd: Option<String>,
/// Runtime overrides for the child agent.
pub runtime_overrides: SubagentRuntimeOverrides,
/// Whether this subagent was launched with `run_in_background: true`.
///
/// Background subagents survive parent-turn cancellation — they are
/// excluded from `cancel_by_parent_prompt_id` so the user can poll
/// results later via `get_task_output`.
/// Controls immediate handle delivery and completion surfacing. A
/// background child still auto-surfaces its completion to the model
/// (buffered reminder / auto-wake) when `surface_completion` is set —
/// background does not mean fire-and-forget. Prompt cancellation still
/// cancels every child owned by that prompt.
pub run_in_background: bool,
/// When false, the subagent's completion is NOT buffered for the
/// between-turn "idle completion" reminder — used by harness-internal
@ -95,11 +97,39 @@ pub struct SubagentRequest {
pub fork_context: bool,
pub owner: SubagentOwner,
pub cancel_token: CancellationToken,
/// Oneshot channel for the coordinator to send back the result.
}
/// Spawn command envelope owned by the coordinator mailbox.
#[derive(Educe)]
#[educe(Debug)]
pub struct SubagentSpawnRequest {
pub request: Box<SubagentRequest>,
#[educe(Debug(ignore))]
pub result_tx: oneshot::Sender<SubagentResult>,
}
impl std::ops::Deref for SubagentSpawnRequest {
type Target = SubagentRequest;
fn deref(&self) -> &Self::Target {
&self.request
}
}
impl SubagentSpawnRequest {
/// Build and send a reply while the plain request remains borrowable.
///
/// Primarily useful for channel adapters and deterministic test harnesses;
/// production lifecycle replies are owned by `SubagentCoordinator`.
pub fn respond_with(
self,
build: impl FnOnce(&SubagentRequest) -> SubagentResult,
) -> Result<(), SubagentResult> {
let result = build(&self.request);
self.result_tx.send(result)
}
}
/// Per-spawn dynamic runtime overrides for a subagent.
///
/// Optional values inherit from the parent or role default. Explicit values take
@ -410,12 +440,14 @@ impl SubagentResult {
// Query protocol
/// Query sent by TaskOutputTool, received by MvpAgent coordinator.
/// Query sent by `TaskOutputTool` to the shared coordinator actor.
#[derive(Educe)]
#[educe(Debug)]
pub struct SubagentQueryRequest {
/// The subagent ID to look up.
pub subagent_id: String,
/// Restrict the lookup to children owned by this parent session.
pub parent_session_id: Option<String>,
/// If true, coordinator waits for completion (up to timeout) before responding.
pub block: bool,
/// Max wait time in ms when blocking. Default 30s.
@ -449,6 +481,27 @@ pub struct SubagentSnapshot {
pub persona: Option<String>,
}
/// Lifecycle metadata returned to shell presentation and extension callers.
#[derive(Debug, Clone)]
pub struct SubagentInspection {
pub snapshot: SubagentSnapshot,
pub parent_session_id: String,
pub child_session_id: String,
pub fork_parent_prompt_id: Option<String>,
pub resumed_from: Option<String>,
}
impl SubagentSnapshot {
/// Whether the child is still in flight (initializing or running) — the
/// shared liveness rule every driver's blocking query loops on.
pub fn is_running(&self) -> bool {
matches!(
self.status,
SubagentSnapshotStatus::Running { .. } | SubagentSnapshotStatus::Initializing
)
}
}
/// Status of a subagent snapshot.
#[derive(Debug, Clone)]
pub enum SubagentSnapshotStatus {
@ -506,11 +559,11 @@ pub enum SubagentCancelTarget {
WorkflowRunId(String),
}
/// Cancel request sent by KillTaskTool or session cancellation paths,
/// received by MvpAgent coordinator.
/// Cancel request sent by `KillTaskTool` or session cancellation paths.
#[derive(Educe)]
#[educe(Debug)]
pub struct SubagentCancelRequest {
pub parent_session_id: Option<String>,
pub target: SubagentCancelTarget,
#[educe(Debug(ignore))]
pub respond_to: oneshot::Sender<SubagentCancelOutcome>,
@ -524,10 +577,11 @@ pub enum SubagentCancelOutcome {
}
/// Summary of a completed subagent, used for between-turn delivery.
/// Session ownership lives on the coordinator's `BufferedCompletion` wrapper;
/// drains are scoped there, so delivered summaries carry no owner field.
#[derive(Debug, Clone)]
pub struct SubagentCompletionSummary {
pub subagent_id: String,
pub owner_session_id: String,
pub subagent_type: String,
pub description: String,
pub success: bool,
@ -560,7 +614,7 @@ pub struct SubagentMultiWaitRequest {
#[derive(Educe)]
#[educe(Debug)]
pub struct SubagentCompletionsRequest {
pub session_id: String,
pub parent_session_id: Option<String>,
pub suppress_ids: Vec<String>,
#[educe(Debug(ignore))]
pub respond_to: oneshot::Sender<Vec<SubagentCompletionSummary>>,
@ -580,6 +634,7 @@ pub struct SubagentOutstandingReply {
#[derive(Educe)]
#[educe(Debug)]
pub struct SubagentOutstandingRequest {
pub parent_session_id: String,
pub prompt_id: String,
#[educe(Debug(ignore))]
pub respond_to: oneshot::Sender<SubagentOutstandingReply>,
@ -588,6 +643,7 @@ pub struct SubagentOutstandingRequest {
/// Clear sticky incomplete after freeze/cancel has snapshotted the bill.
#[derive(Debug)]
pub struct SubagentClearUsageNotAppliedRequest {
pub parent_session_id: String,
pub prompt_id: String,
}
@ -595,11 +651,94 @@ pub struct SubagentClearUsageNotAppliedRequest {
#[derive(Educe)]
#[educe(Debug)]
pub struct SubagentMarkUsageNotAppliedRequest {
pub parent_session_id: String,
pub prompt_id: String,
#[educe(Debug(ignore))]
pub respond_to: oneshot::Sender<()>,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct SubagentRegistryCounts {
pub pending: usize,
pub active: usize,
pub completed: usize,
}
#[derive(Educe)]
#[educe(Debug)]
pub struct SubagentRegistryCountsRequest {
#[educe(Debug(ignore))]
pub respond_to: oneshot::Sender<SubagentRegistryCounts>,
}
/// Request for full metadata plus a resolved progress snapshot.
#[derive(Educe)]
#[educe(Debug)]
pub struct SubagentInspectRequest {
pub subagent_id: String,
pub parent_session_id: Option<String>,
#[educe(Debug(ignore))]
pub respond_to: oneshot::Sender<Option<SubagentInspection>>,
}
/// Request for all running children owned by one parent session.
#[derive(Educe)]
#[educe(Debug)]
pub struct SubagentListRunningRequest {
pub parent_session_id: String,
#[educe(Debug(ignore))]
pub respond_to: oneshot::Sender<Vec<SubagentInspection>>,
}
/// Fork/resume provenance retained by the shared coordinator.
#[derive(Debug, Clone, Default)]
pub struct SubagentProvenance {
pub fork_parent_prompt_id: Option<String>,
pub resumed_from: Option<String>,
}
/// Reference to a child spawned during one parent prompt.
#[derive(Debug, Clone)]
pub struct SpawnedSubagentRef {
pub subagent_id: String,
pub child_session_id: String,
pub subagent_type: String,
pub description: String,
pub persona: Option<String>,
pub resumed_from: Option<String>,
}
/// Request for prompt-scoped spawned-child references.
#[derive(Educe)]
#[educe(Debug)]
pub struct SubagentSpawnedRefsRequest {
pub parent_session_id: String,
pub prompt_id: String,
#[educe(Debug(ignore))]
pub respond_to: oneshot::Sender<Vec<SpawnedSubagentRef>>,
}
/// In-memory source data used by a runtime adapter to resume a child.
#[derive(Debug, Clone)]
pub struct SubagentResumeSource {
pub subagent_id: String,
pub child_session_id: String,
pub child_cwd: String,
pub worktree_path: Option<String>,
pub snapshot_ref: Option<String>,
pub subagent_type: String,
pub persona: Option<String>,
pub model_id: Option<String>,
}
/// Result of a resume-source lookup.
#[derive(Debug, Clone)]
pub enum SubagentResumeLookup {
Active,
Completed(SubagentResumeSource),
Missing,
}
// Validate-type protocol
#[derive(Debug, Clone)]
@ -700,18 +839,25 @@ pub struct SubagentDescribeRequest {
pub respond_to: oneshot::Sender<SubagentDescribeOutcome>,
}
/// Coordinator message enum. Intentionally NOT `#[non_exhaustive]` —
/// the cross-crate drain loop in `xai-grok-shell` relies on
/// compile-time exhaustiveness.
/// Coordinator message enum. Kept exhaustive so every actor command is handled.
pub enum SubagentEvent {
Spawn(Box<SubagentRequest>),
Spawn(SubagentSpawnRequest),
Query(SubagentQueryRequest),
Cancel(SubagentCancelRequest),
ListActive(SubagentListActiveRequest),
ListRunning(SubagentListRunningRequest),
Completions(SubagentCompletionsRequest),
/// Fire-and-forget: drop buffered completions owned by a removed session
/// so unloaded sessions cannot leak entries into the shared buffer.
DiscardSessionCompletions {
parent_session_id: String,
},
Outstanding(SubagentOutstandingRequest),
ClearUsageNotApplied(SubagentClearUsageNotAppliedRequest),
MarkUsageNotApplied(SubagentMarkUsageNotAppliedRequest),
RegistryCounts(SubagentRegistryCountsRequest),
Inspect(SubagentInspectRequest),
SpawnedRefs(SubagentSpawnedRefsRequest),
ValidateType(SubagentValidateTypeRequest),
DescribeType(SubagentDescribeRequest),
LoopUnitActive(SubagentLoopUnitActiveRequest),
@ -780,10 +926,8 @@ pub fn drain_owned(
/// Lightweight summary of a running subagent.
///
/// This is the single shared definition of this type. The coordinator in
/// xai-grok-shell produces it, the channel protocol carries it, and the
/// compaction pipeline in xai-chat-state (via `RunningSubagentSummary`)
/// consumes it. Do not duplicate this type in other crates.
/// The shared coordinator produces this through the channel protocol, and the
/// compaction pipeline consumes it through `RunningSubagentSummary`.
#[derive(Debug, Clone)]
pub struct ActiveSubagentSummary {
/// The subagent's unique ID (same ID used by `get_task_output` / `kill_task`).
@ -799,8 +943,7 @@ pub struct ActiveSubagentSummary {
/// Request to list currently-running subagents for a specific parent session.
///
/// Sent by the compaction pipeline in `SessionActor::run_compact_inner()`.
/// Handled by `MvpAgent::start_subagent_coordinator()` which borrows the
/// coordinator and calls `active_summaries_for()`.
/// Handled by the shared coordinator actor.
#[derive(Educe)]
#[educe(Debug)]
pub struct SubagentListActiveRequest {
@ -853,6 +996,39 @@ pub struct SessionIdResource(pub String);
register_resource!("grok_build", "SessionIdResource", SessionIdResource);
/// Host-owned RAII token for an interruptible foreground wait.
pub trait ForegroundWaitGuard: Send {}
impl<T: Send> ForegroundWaitGuard for T {}
type ForegroundWaitFactory = dyn Fn() -> Box<dyn ForegroundWaitGuard> + Send + Sync;
/// Factory injected by hosts that expose a send-now wait window.
#[derive(Clone)]
pub struct SubagentForegroundWait(Arc<ForegroundWaitFactory>);
impl SubagentForegroundWait {
pub fn new(factory: impl Fn() -> Box<dyn ForegroundWaitGuard> + Send + Sync + 'static) -> Self {
Self(Arc::new(factory))
}
pub fn enter(&self) -> Box<dyn ForegroundWaitGuard> {
(self.0)()
}
}
impl std::fmt::Debug for SubagentForegroundWait {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SubagentForegroundWait").finish()
}
}
register_resource!(
"grok_build",
"SubagentForegroundWait",
SubagentForegroundWait
);
/// Carries the current parent prompt/turn ID for TaskTool subagent scoping.
///
/// Set by xai-grok-shell immediately before a prompt turn begins executing so
@ -1282,19 +1458,18 @@ mod tests {
let (respond_to, mut response_rx) = oneshot::channel();
tx.send(super::SubagentCompletionsRequest {
session_id: "session-1".into(),
parent_session_id: Some("parent".into()),
suppress_ids: vec!["id-1".into(), "id-2".into()],
respond_to,
})
.unwrap();
let req = rx.try_recv().unwrap();
assert_eq!(req.session_id, "session-1");
assert_eq!(req.parent_session_id.as_deref(), Some("parent"));
assert_eq!(req.suppress_ids, vec!["id-1", "id-2"]);
let summaries = vec![super::SubagentCompletionSummary {
subagent_id: "sub-1".into(),
owner_session_id: "session-1".into(),
subagent_type: "general-purpose".into(),
description: "test task".into(),
success: true,
@ -1373,7 +1548,7 @@ mod tests {
.0
.send(super::SubagentEvent::Completions(
super::SubagentCompletionsRequest {
session_id: String::new(),
parent_session_id: None,
suppress_ids: vec![],
respond_to,
},
@ -1405,7 +1580,7 @@ mod tests {
.0
.send(super::SubagentEvent::Completions(
super::SubagentCompletionsRequest {
session_id: String::new(),
parent_session_id: None,
suppress_ids: vec![],
respond_to,
},

1
crates/codegen/xai-grok-tools/src/implementations/grok_build/task_output/mod.rs
View file

@ -916,6 +916,7 @@ pub(crate) mod test_helpers {
block_waited: false,
explicitly_killed: false,
owner_session_id: None,
description: None,
}
}

2
crates/codegen/xai-grok-tools/src/implementations/grok_build_hashline/grep.rs
View file

@ -147,7 +147,7 @@ Content output format:
Usage:
- ${{ params.search.pattern }} is a regex: `log.*Error`, `function\s+\w+`, `TODO`
- Output modes: "content" (default, with anchors), "files_with_matches", "count"
- Default output is anchored content matches (no output-mode selector)
- Use -A, -B, -C for context lines around matches
- Only use '${{ params.search.type }}' or '${{ params.search.glob }}' when certain of the file type
- Results are capped; truncated results show "at least" counts"#;

4
crates/codegen/xai-grok-tools/src/implementations/opencode/bash/mod.rs
View file

@ -394,7 +394,9 @@ impl xai_tool_runtime::Tool for BashTool {
auto_background_on_timeout: false, // OpenCode doesn't support auto-backgrounding
foreground_block_budget: None,
kind: crate::computer::types::TaskKind::Bash,
owner_session_id: None, // OpenCode doesn't use shared terminal backends
// OpenCode doesn't use shared terminal backends.
owner_session_id: None,
description: None,
};
let result = match backend.run(request).await {

1
crates/codegen/xai-grok-tools/src/implementations/task_output/tool.rs
View file

@ -107,6 +107,7 @@ mod tests {
block_waited: false,
explicitly_killed: false,
owner_session_id: None,
description: None,
}
}