Synced from monorepo
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]
parent
a5727c5960
commit
69f0ba880a
286 changed files with 22939 additions and 9624 deletions
|
|
@ -141,17 +141,22 @@ impl ToolBridge {
|
|||
template: &str,
|
||||
placeholders: &serde_json::Value,
|
||||
) -> Option<String> {
|
||||
let registry = &*self.registry;
|
||||
let result;
|
||||
{
|
||||
result = registry
|
||||
.resources
|
||||
.lock()
|
||||
.await
|
||||
.get::<TemplateRenderer>()
|
||||
.and_then(|r| r.render_with_extra(template, placeholders).ok());
|
||||
}
|
||||
result
|
||||
self.registry
|
||||
.resources
|
||||
.lock()
|
||||
.await
|
||||
.get::<TemplateRenderer>()
|
||||
.and_then(|renderer| renderer.render_with_extra(template, placeholders).ok())
|
||||
}
|
||||
|
||||
/// Return the finalized template renderer for multi-part prompt assembly.
|
||||
pub async fn template_renderer_snapshot(&self) -> Option<TemplateRenderer> {
|
||||
self.registry
|
||||
.resources
|
||||
.lock()
|
||||
.await
|
||||
.get::<TemplateRenderer>()
|
||||
.cloned()
|
||||
}
|
||||
|
||||
pub async fn register_mcp_tools<T>(
|
||||
|
|
@ -829,6 +834,7 @@ mod tests {
|
|||
block_waited: false,
|
||||
explicitly_killed: false,
|
||||
owner_session_id: owner.map(|s| s.to_string()),
|
||||
description: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -272,6 +272,7 @@ impl ShellState {
|
|||
pub async fn init(
|
||||
shell: ShellKind,
|
||||
cwd: &Path,
|
||||
shell_env_policy: Option<&crate::util::ShellEnvironmentPolicy>,
|
||||
) -> Result<Self, crate::computer::types::ComputerError> {
|
||||
let dump_script = shell.dump_script();
|
||||
let dump_fn = shell.dump_function_name();
|
||||
|
|
@ -297,6 +298,21 @@ impl ShellState {
|
|||
.stderr(Stdio::null())
|
||||
.kill_on_drop(true);
|
||||
crate::util::detach_command(&mut cmd);
|
||||
// Apply the policy before the `export -p` snapshot so the replayed state
|
||||
// is already filtered; otherwise the restore would undo it. No-op unless set.
|
||||
//
|
||||
// SECURITY: this filters the base env only. Variables an rc file exports
|
||||
// during login are captured in the replay snapshot and are not
|
||||
// re-filtered by `exclude`/`include_only` on the persistent backend, so
|
||||
// warn when a policy is active. The non-persistent backend has no such
|
||||
// gap (it filters login capture directly).
|
||||
if shell_env_policy.is_some_and(|p| !p.is_noop()) {
|
||||
tracing::warn!(
|
||||
"shell_environment_policy filters the persistent shell's base env only; \
|
||||
variables exported by rc files enter the replay snapshot unfiltered"
|
||||
);
|
||||
}
|
||||
crate::util::apply_shell_environment_policy(&mut cmd, shell_env_policy);
|
||||
cmd.envs(crate::util::pager_env());
|
||||
let mut child = cmd.spawn().map_err(|e| {
|
||||
crate::computer::types::ComputerError::io(format!(
|
||||
|
|
@ -839,7 +855,7 @@ mod tests {
|
|||
return;
|
||||
}
|
||||
let cwd = std::env::current_dir().unwrap();
|
||||
let state = ShellState::init(ShellKind::Bash, &cwd).await.unwrap();
|
||||
let state = ShellState::init(ShellKind::Bash, &cwd, None).await.unwrap();
|
||||
assert!(state.cwd.is_absolute());
|
||||
// The snapshot should contain at least some env var exports
|
||||
assert!(
|
||||
|
|
@ -858,7 +874,7 @@ mod tests {
|
|||
return;
|
||||
}
|
||||
let cwd = std::env::current_dir().unwrap();
|
||||
let mut state = ShellState::init(ShellKind::Bash, &cwd).await.unwrap();
|
||||
let mut state = ShellState::init(ShellKind::Bash, &cwd, None).await.unwrap();
|
||||
|
||||
// Run "export GROK_TEST_VAR=hello" and capture the new state
|
||||
let prep = state
|
||||
|
|
@ -964,7 +980,7 @@ mod tests {
|
|||
return;
|
||||
}
|
||||
let cwd = std::env::current_dir().unwrap();
|
||||
let mut state = ShellState::init(ShellKind::Bash, &cwd).await.unwrap();
|
||||
let mut state = ShellState::init(ShellKind::Bash, &cwd, None).await.unwrap();
|
||||
|
||||
// cd to /tmp (macOS resolves to /private/tmp via symlink)
|
||||
let (code, _) = run_command(&mut state, "cd /tmp").await;
|
||||
|
|
@ -987,7 +1003,7 @@ mod tests {
|
|||
return;
|
||||
}
|
||||
let cwd = std::env::current_dir().unwrap();
|
||||
let mut state = ShellState::init(ShellKind::Bash, &cwd).await.unwrap();
|
||||
let mut state = ShellState::init(ShellKind::Bash, &cwd, None).await.unwrap();
|
||||
|
||||
// Export a variable
|
||||
let (code, _) = run_command(&mut state, "export MY_TEST_VAR=persistent_value").await;
|
||||
|
|
@ -1006,7 +1022,7 @@ mod tests {
|
|||
return;
|
||||
}
|
||||
let cwd = std::env::current_dir().unwrap();
|
||||
let mut state = ShellState::init(ShellKind::Bash, &cwd).await.unwrap();
|
||||
let mut state = ShellState::init(ShellKind::Bash, &cwd, None).await.unwrap();
|
||||
|
||||
let (code, _) = run_command(&mut state, "export GPG_TTY=/grok-sentinel-tty").await;
|
||||
assert_eq!(code, 0);
|
||||
|
|
@ -1026,7 +1042,7 @@ mod tests {
|
|||
return;
|
||||
}
|
||||
let cwd = std::env::current_dir().unwrap();
|
||||
let mut state = ShellState::init(ShellKind::Zsh, &cwd).await.unwrap();
|
||||
let mut state = ShellState::init(ShellKind::Zsh, &cwd, None).await.unwrap();
|
||||
|
||||
let (code, _) = run_command(&mut state, "export GPG_TTY=/grok-sentinel-tty").await;
|
||||
assert_eq!(code, 0);
|
||||
|
|
@ -1046,7 +1062,7 @@ mod tests {
|
|||
return;
|
||||
}
|
||||
let cwd = std::env::current_dir().unwrap();
|
||||
let mut state = ShellState::init(ShellKind::Zsh, &cwd).await.unwrap();
|
||||
let mut state = ShellState::init(ShellKind::Zsh, &cwd, None).await.unwrap();
|
||||
|
||||
let prep = state
|
||||
.prepare_command(
|
||||
|
|
@ -1091,7 +1107,7 @@ mod tests {
|
|||
return;
|
||||
}
|
||||
let cwd = std::env::current_dir().unwrap();
|
||||
let mut state = ShellState::init(ShellKind::Bash, &cwd).await.unwrap();
|
||||
let mut state = ShellState::init(ShellKind::Bash, &cwd, None).await.unwrap();
|
||||
|
||||
// Define a function
|
||||
let (code, _) = run_command(&mut state, "greet() { echo \"hello $1\"; }").await;
|
||||
|
|
@ -1109,7 +1125,7 @@ mod tests {
|
|||
return;
|
||||
}
|
||||
let cwd = std::env::current_dir().unwrap();
|
||||
let mut state = ShellState::init(ShellKind::Bash, &cwd).await.unwrap();
|
||||
let mut state = ShellState::init(ShellKind::Bash, &cwd, None).await.unwrap();
|
||||
|
||||
// Define an alias
|
||||
let (code, _) = run_command(&mut state, "alias ll='ls -la'").await;
|
||||
|
|
@ -1136,7 +1152,7 @@ mod tests {
|
|||
return;
|
||||
}
|
||||
let cwd = std::env::current_dir().unwrap();
|
||||
let mut state = ShellState::init(ShellKind::Bash, &cwd).await.unwrap();
|
||||
let mut state = ShellState::init(ShellKind::Bash, &cwd, None).await.unwrap();
|
||||
|
||||
let prep = state.prepare_command("true", None, shadows, None).unwrap();
|
||||
// Shadows enabled → the self-resolving find/grep functions are always
|
||||
|
|
@ -1175,7 +1191,7 @@ mod tests {
|
|||
return;
|
||||
}
|
||||
let cwd = std::env::current_dir().unwrap();
|
||||
let mut state = ShellState::init(ShellKind::Bash, &cwd).await.unwrap();
|
||||
let mut state = ShellState::init(ShellKind::Bash, &cwd, None).await.unwrap();
|
||||
|
||||
// Set up some state
|
||||
let (_, _) = run_command(&mut state, "export SURVIVE_TEST=yes").await;
|
||||
|
|
|
|||
|
|
@ -303,6 +303,7 @@ struct ProcessState {
|
|||
/// Session that owns this process. Used to scope kill operations so
|
||||
/// subagent teardown only kills the subagent's own tasks.
|
||||
owner_session_id: Option<String>,
|
||||
description: Option<String>,
|
||||
}
|
||||
|
||||
impl ProcessState {
|
||||
|
|
@ -442,6 +443,7 @@ impl ProcessState {
|
|||
explicitly_killed: self.explicitly_killed,
|
||||
kind: self.kind,
|
||||
owner_session_id: self.owner_session_id.clone(),
|
||||
description: self.description.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -516,6 +518,10 @@ struct LocalTerminalActor {
|
|||
/// a subagent reusing this backend can't clobber the parent's shadows.
|
||||
search_shadows: SearchShadowConfig,
|
||||
|
||||
/// Shell-environment policy baked in at construction (like `search_shadows`);
|
||||
/// `None` inherits the full environment.
|
||||
shell_env_policy: Option<crate::util::ShellEnvironmentPolicy>,
|
||||
|
||||
/// Persistent shell state (env vars, cwd, functions, aliases).
|
||||
/// Lazily initialized on first command when `persistent_shell` is true.
|
||||
#[cfg(unix)]
|
||||
|
|
@ -542,11 +548,13 @@ impl LocalTerminalActor {
|
|||
foreground_block_budget: Duration,
|
||||
output_file_cap: u64,
|
||||
scope: crate::util::ProcessScope,
|
||||
shell_env_policy: Option<crate::util::ShellEnvironmentPolicy>,
|
||||
) -> Self {
|
||||
Self {
|
||||
cmd_rx,
|
||||
cancel_token,
|
||||
scope,
|
||||
shell_env_policy,
|
||||
processes: HashMap::new(),
|
||||
completion_waiters: HashMap::new(),
|
||||
completed_task_snapshots: HashMap::new(),
|
||||
|
|
@ -598,8 +606,14 @@ impl LocalTerminalActor {
|
|||
#[cfg(not(unix))]
|
||||
let login_env: Option<&HashMap<String, String>> = None;
|
||||
|
||||
let (child, process_group) =
|
||||
spawn_shell_command(command, cwd, env, login_env, self.search_shadows)?;
|
||||
let (child, process_group) = spawn_shell_command(
|
||||
command,
|
||||
cwd,
|
||||
env,
|
||||
login_env,
|
||||
self.search_shadows,
|
||||
self.shell_env_policy.as_ref(),
|
||||
)?;
|
||||
Ok(SpawnResult {
|
||||
child,
|
||||
process_group,
|
||||
|
|
@ -658,22 +672,12 @@ impl LocalTerminalActor {
|
|||
.stderr(Stdio::piped())
|
||||
.kill_on_drop(true);
|
||||
|
||||
if let Some(login) = self.login_env.as_ref() {
|
||||
for (key, value) in login {
|
||||
if key != "PATH" && std::env::var_os(key).is_none() {
|
||||
cmd.env(key, value);
|
||||
}
|
||||
}
|
||||
}
|
||||
cmd.envs(shell_state::shell_env_overrides());
|
||||
for (key, value) in env {
|
||||
cmd.env(key, value);
|
||||
}
|
||||
cmd.envs(crate::util::pager_env());
|
||||
if let Some(path) = self.login_env.as_ref().and_then(|l| l.get("PATH")) {
|
||||
cmd.env("PATH", path);
|
||||
}
|
||||
crate::util::apply_grok_agent_marker(&mut cmd);
|
||||
apply_child_env(
|
||||
&mut cmd,
|
||||
self.shell_env_policy.as_ref(),
|
||||
self.login_env.as_ref(),
|
||||
env,
|
||||
);
|
||||
|
||||
cmd.fd_mappings(prep.fd_mappings)
|
||||
.map_err(|e| ComputerError::io(format!("fd mapping: {e}")))?;
|
||||
|
|
@ -722,7 +726,7 @@ impl LocalTerminalActor {
|
|||
return;
|
||||
}
|
||||
let shell = shell_state::ShellKind::detect();
|
||||
match shell_state::ShellState::init(shell, cwd).await {
|
||||
match shell_state::ShellState::init(shell, cwd, self.shell_env_policy.as_ref()).await {
|
||||
Ok(state) => self.shell_state = Some(state),
|
||||
Err(e) => {
|
||||
tracing::warn!("persistent shell init failed, using empty state: {e}");
|
||||
|
|
@ -791,17 +795,9 @@ impl LocalTerminalActor {
|
|||
.stderr(Stdio::piped())
|
||||
.kill_on_drop(true);
|
||||
|
||||
// Apply SHELL_ENV_OVERRIDES (TERM=dumb, NO_COLOR, GROK_AGENT=1, etc.)
|
||||
// + request env + pager env. Agent marker is re-applied last so request
|
||||
// env cannot clear it.
|
||||
cmd.envs(shell_state::shell_env_overrides());
|
||||
|
||||
for (key, value) in env {
|
||||
cmd.env(key, value);
|
||||
}
|
||||
|
||||
cmd.envs(crate::util::pager_env());
|
||||
crate::util::apply_grok_agent_marker(&mut cmd);
|
||||
// The persistent backend restores login state from its snapshot, so no
|
||||
// login-env layering here.
|
||||
apply_child_env(&mut cmd, self.shell_env_policy.as_ref(), None, env);
|
||||
|
||||
cmd.fd_mappings(prep.fd_mappings)
|
||||
.map_err(|e| ComputerError::io(format!("fd mapping: {e}")))?;
|
||||
|
|
@ -1095,6 +1091,7 @@ impl LocalTerminalActor {
|
|||
explicitly_killed: false,
|
||||
state_dump_handle,
|
||||
owner_session_id: request.owner_session_id.clone(),
|
||||
description: request.description.filter(|d| !d.trim().is_empty()),
|
||||
};
|
||||
|
||||
// Send an initial empty notification so the TUI shows the execution
|
||||
|
|
@ -1238,6 +1235,7 @@ impl LocalTerminalActor {
|
|||
None
|
||||
},
|
||||
owner_session_id: request.owner_session_id.clone(),
|
||||
description: request.description.filter(|d| !d.trim().is_empty()),
|
||||
};
|
||||
|
||||
// Store under task_id — this is the key that get_task/kill_task will use
|
||||
|
|
@ -1618,6 +1616,7 @@ impl LocalTerminalActor {
|
|||
block_waited: p.block_waited,
|
||||
explicitly_killed: p.explicitly_killed,
|
||||
owner_session_id: p.owner_session_id.clone(),
|
||||
description: p.description.clone(),
|
||||
};
|
||||
self.completed_task_snapshots.insert(id.clone(), snapshot);
|
||||
}
|
||||
|
|
@ -2044,15 +2043,25 @@ impl LocalTerminalActor {
|
|||
// "Monitor" row (matching the original-spawn path) rather than a
|
||||
// bash-highlighted "[monitor] …".
|
||||
let is_monitor = process.kind == crate::computer::types::TaskKind::Monitor;
|
||||
let monitor_description = if is_monitor {
|
||||
// Recover monitor label once; reuse for backgrounded notify + pipeline.
|
||||
// Filter empty/whitespace the same way as spawn so `[monitor] `
|
||||
// / blank recovery does not stick as Some("") and block the
|
||||
// command fallback for the re-spawned pipeline label.
|
||||
let recovered_monitor_description = if is_monitor {
|
||||
process
|
||||
.display_command
|
||||
.as_deref()
|
||||
.and_then(|d| d.strip_prefix("[monitor] "))
|
||||
.map(str::to_string)
|
||||
.filter(|d| !d.trim().is_empty())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let effective_description = process
|
||||
.description
|
||||
.clone()
|
||||
.filter(|d| !d.trim().is_empty())
|
||||
.or_else(|| recovered_monitor_description.clone());
|
||||
let reparent_command = if is_monitor {
|
||||
process.command.clone()
|
||||
} else {
|
||||
|
|
@ -2072,20 +2081,16 @@ impl LocalTerminalActor {
|
|||
},
|
||||
output_file: process.output_file.clone(),
|
||||
task_id: task_id.clone(),
|
||||
monitor_description,
|
||||
// Reparent path has no model tool description; monitors use
|
||||
// `monitor_description` above.
|
||||
description: None,
|
||||
monitor_description: recovered_monitor_description,
|
||||
description: effective_description.clone(),
|
||||
});
|
||||
|
||||
// Re-spawn the monitor pipeline so events continue streaming.
|
||||
// The old pipeline died with the child's runtime.
|
||||
if process.kind == crate::computer::types::TaskKind::Monitor {
|
||||
let pipeline_task_id = task_id.clone();
|
||||
let pipeline_description = process
|
||||
.display_command
|
||||
.clone()
|
||||
.unwrap_or_else(|| process.command.clone());
|
||||
let pipeline_description =
|
||||
effective_description.unwrap_or_else(|| process.command.clone());
|
||||
// Weak so the reparented monitor doesn't pin the backend.
|
||||
let pipeline_terminal = backend_weak.clone();
|
||||
let pipeline_notif = new_handle.clone();
|
||||
|
|
@ -2127,6 +2132,31 @@ pub struct LocalTerminalBackend {
|
|||
cancel_token: CancellationToken,
|
||||
}
|
||||
|
||||
/// Grouped inputs for [`LocalTerminalBackend::new_inner`], so call sites read as
|
||||
/// named fields instead of a telescoping list of positional `bool`s. Constructors
|
||||
/// override only the fields they vary via `..Default::default()`.
|
||||
struct LocalTerminalConfig {
|
||||
memory_config: Option<CgroupMemoryConfig>,
|
||||
use_spawn_local: bool,
|
||||
persistent_shell: bool,
|
||||
login_shell_capture: bool,
|
||||
search_shadows: SearchShadowConfig,
|
||||
shell_env_policy: Option<crate::util::ShellEnvironmentPolicy>,
|
||||
}
|
||||
|
||||
impl Default for LocalTerminalConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
memory_config: None,
|
||||
use_spawn_local: false,
|
||||
persistent_shell: false,
|
||||
login_shell_capture: true,
|
||||
search_shadows: SearchShadowConfig::default(),
|
||||
shell_env_policy: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl LocalTerminalBackend {
|
||||
/// Create a new LocalTerminalBackend and spawn the actor task.
|
||||
///
|
||||
|
|
@ -2134,7 +2164,7 @@ impl LocalTerminalBackend {
|
|||
/// If `memory_config` is provided, a cgroupv2 memory limit is enforced on
|
||||
/// all spawned commands (Linux only; silently degrades to no-op elsewhere).
|
||||
pub fn new() -> Self {
|
||||
Self::new_inner(None, false, false, true, SearchShadowConfig::default())
|
||||
Self::new_inner(LocalTerminalConfig::default())
|
||||
}
|
||||
|
||||
/// Create a new LocalTerminalBackend with persistent shell state.
|
||||
|
|
@ -2143,31 +2173,29 @@ impl LocalTerminalBackend {
|
|||
/// and shell options persist across command invocations. The user's login shell
|
||||
/// (bash or zsh) is detected and its rc files are loaded once on first command.
|
||||
pub fn with_persistent_shell() -> Self {
|
||||
Self::new_inner(None, false, true, true, SearchShadowConfig::default())
|
||||
Self::new_inner(LocalTerminalConfig {
|
||||
persistent_shell: true,
|
||||
..Default::default()
|
||||
})
|
||||
}
|
||||
|
||||
/// Create a new LocalTerminalBackend with cgroup memory limits.
|
||||
///
|
||||
/// See [`CgroupMemoryConfig`] for details on the soft/hard limit model.
|
||||
pub fn with_memory_limit(config: CgroupMemoryConfig) -> Self {
|
||||
Self::new_inner(
|
||||
Some(config),
|
||||
false,
|
||||
false,
|
||||
true,
|
||||
SearchShadowConfig::default(),
|
||||
)
|
||||
Self::new_inner(LocalTerminalConfig {
|
||||
memory_config: Some(config),
|
||||
..Default::default()
|
||||
})
|
||||
}
|
||||
|
||||
/// Create a new LocalTerminalBackend with both memory limits and persistent shell.
|
||||
pub fn with_memory_limit_and_persistent_shell(config: CgroupMemoryConfig) -> Self {
|
||||
Self::new_inner(
|
||||
Some(config),
|
||||
false,
|
||||
true,
|
||||
true,
|
||||
SearchShadowConfig::default(),
|
||||
)
|
||||
Self::new_inner(LocalTerminalConfig {
|
||||
memory_config: Some(config),
|
||||
persistent_shell: true,
|
||||
..Default::default()
|
||||
})
|
||||
}
|
||||
|
||||
/// Create a new LocalTerminalBackend using spawn_local (for single-threaded runtimes).
|
||||
|
|
@ -2175,33 +2203,51 @@ impl LocalTerminalBackend {
|
|||
/// `search_shadows` is the host-resolved `find`→`bfs` / `grep`→`ugrep` enable
|
||||
/// state, baked into this backend (see [`SearchShadowConfig`]).
|
||||
pub fn new_local(search_shadows: SearchShadowConfig) -> Self {
|
||||
Self::new_inner(None, true, false, true, search_shadows)
|
||||
Self::new_inner(LocalTerminalConfig {
|
||||
use_spawn_local: true,
|
||||
search_shadows,
|
||||
..Default::default()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn new_local_with_login_shell_capture(
|
||||
search_shadows: SearchShadowConfig,
|
||||
login_shell_capture: bool,
|
||||
shell_env_policy: Option<crate::util::ShellEnvironmentPolicy>,
|
||||
) -> Self {
|
||||
Self::new_inner(None, true, false, login_shell_capture, search_shadows)
|
||||
Self::new_inner(LocalTerminalConfig {
|
||||
use_spawn_local: true,
|
||||
login_shell_capture,
|
||||
search_shadows,
|
||||
shell_env_policy,
|
||||
..Default::default()
|
||||
})
|
||||
}
|
||||
|
||||
/// Create a new LocalTerminalBackend using spawn_local with persistent shell.
|
||||
///
|
||||
/// `search_shadows` is the host-resolved `find`→`bfs` / `grep`→`ugrep` enable
|
||||
/// state, baked into this backend (see [`SearchShadowConfig`]).
|
||||
pub fn new_local_with_persistent_shell(search_shadows: SearchShadowConfig) -> Self {
|
||||
Self::new_inner(None, true, true, true, search_shadows)
|
||||
pub fn new_local_with_persistent_shell(
|
||||
search_shadows: SearchShadowConfig,
|
||||
shell_env_policy: Option<crate::util::ShellEnvironmentPolicy>,
|
||||
) -> Self {
|
||||
Self::new_inner(LocalTerminalConfig {
|
||||
use_spawn_local: true,
|
||||
persistent_shell: true,
|
||||
search_shadows,
|
||||
shell_env_policy,
|
||||
..Default::default()
|
||||
})
|
||||
}
|
||||
|
||||
/// Create a new LocalTerminalBackend using spawn_local with memory limits.
|
||||
pub fn new_local_with_memory_limit(config: CgroupMemoryConfig) -> Self {
|
||||
Self::new_inner(
|
||||
Some(config),
|
||||
true,
|
||||
false,
|
||||
true,
|
||||
SearchShadowConfig::default(),
|
||||
)
|
||||
Self::new_inner(LocalTerminalConfig {
|
||||
memory_config: Some(config),
|
||||
use_spawn_local: true,
|
||||
..Default::default()
|
||||
})
|
||||
}
|
||||
|
||||
/// Test-only: a spawn_local backend that enrolls spawned children into
|
||||
|
|
@ -2222,6 +2268,7 @@ impl LocalTerminalBackend {
|
|||
FOREGROUND_BLOCK_BUDGET,
|
||||
MAX_OUTPUT_FILE_BYTES,
|
||||
scope,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
|
|
@ -2238,6 +2285,7 @@ impl LocalTerminalBackend {
|
|||
FOREGROUND_BLOCK_BUDGET,
|
||||
MAX_OUTPUT_FILE_BYTES,
|
||||
crate::util::global_process_scope().clone(),
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
|
|
@ -2254,6 +2302,7 @@ impl LocalTerminalBackend {
|
|||
budget,
|
||||
MAX_OUTPUT_FILE_BYTES,
|
||||
crate::util::global_process_scope().clone(),
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
|
|
@ -2270,16 +2319,19 @@ impl LocalTerminalBackend {
|
|||
FOREGROUND_BLOCK_BUDGET,
|
||||
output_file_cap,
|
||||
crate::util::global_process_scope().clone(),
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
fn new_inner(
|
||||
memory_config: Option<CgroupMemoryConfig>,
|
||||
use_spawn_local: bool,
|
||||
persistent_shell: bool,
|
||||
login_shell_capture: bool,
|
||||
search_shadows: SearchShadowConfig,
|
||||
) -> Self {
|
||||
fn new_inner(config: LocalTerminalConfig) -> Self {
|
||||
let LocalTerminalConfig {
|
||||
memory_config,
|
||||
use_spawn_local,
|
||||
persistent_shell,
|
||||
login_shell_capture,
|
||||
search_shadows,
|
||||
shell_env_policy,
|
||||
} = config;
|
||||
Self::new_with_ttl(
|
||||
memory_config,
|
||||
use_spawn_local,
|
||||
|
|
@ -2290,6 +2342,7 @@ impl LocalTerminalBackend {
|
|||
foreground_block_budget_from_env(),
|
||||
output_file_cap_from_env(),
|
||||
crate::util::global_process_scope().clone(),
|
||||
shell_env_policy,
|
||||
)
|
||||
}
|
||||
|
||||
|
|
@ -2303,6 +2356,7 @@ impl LocalTerminalBackend {
|
|||
foreground_block_budget: Duration,
|
||||
output_file_cap: u64,
|
||||
scope: crate::util::ProcessScope,
|
||||
shell_env_policy: Option<crate::util::ShellEnvironmentPolicy>,
|
||||
) -> Self {
|
||||
let (cmd_tx, cmd_rx) = mpsc::channel(COMMAND_CHANNEL_SIZE);
|
||||
let cancel_token = CancellationToken::new();
|
||||
|
|
@ -2329,6 +2383,7 @@ impl LocalTerminalBackend {
|
|||
foreground_block_budget,
|
||||
output_file_cap,
|
||||
scope,
|
||||
shell_env_policy,
|
||||
);
|
||||
actor.run().await;
|
||||
};
|
||||
|
|
@ -2920,19 +2975,107 @@ async fn capture_login_env() -> HashMap<String, String> {
|
|||
}
|
||||
}
|
||||
|
||||
/// Spawn the shell command and attach the child to a [`ProcessGroup`].
|
||||
/// Layer login-shell captured vars (except `PATH`) onto `cmd`, dropping those the
|
||||
/// active policy filters out and those already set in grok's own environment.
|
||||
#[cfg(unix)]
|
||||
fn layer_login_env_vars(
|
||||
cmd: &mut tokio::process::Command,
|
||||
login_env: Option<&HashMap<String, String>>,
|
||||
active_policy: Option<&crate::util::ShellEnvironmentPolicy>,
|
||||
) {
|
||||
if let Some(login) = login_env {
|
||||
for (key, value) in login {
|
||||
// `var_os` reads grok's own process env (not the possibly cleared
|
||||
// child env): a login var already present in grok's environment is
|
||||
// left alone. Capture is filtered through the policy so an rc export
|
||||
// cannot bypass it.
|
||||
if key != "PATH"
|
||||
&& std::env::var_os(key).is_none()
|
||||
&& active_policy.is_none_or(|p| p.allows_with_inherit(key))
|
||||
{
|
||||
cmd.env(key, value);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Layer per-request env (`.envrc`, ACP, session settings) onto `cmd`, dropping
|
||||
/// names the active policy excludes so a request-supplied secret cannot bypass
|
||||
/// it. Honors `exclude`/`include_only`/default excludes, not `inherit`, since
|
||||
/// request env is provided explicitly rather than inherited.
|
||||
fn layer_request_env(
|
||||
cmd: &mut tokio::process::Command,
|
||||
env: &HashMap<String, String>,
|
||||
active_policy: Option<&crate::util::ShellEnvironmentPolicy>,
|
||||
) {
|
||||
for (key, value) in env {
|
||||
if active_policy.is_none_or(|p| p.allows(key)) {
|
||||
cmd.env(key, value);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Re-inject the login-shell `PATH` last (so rc-file additions win), unless the
|
||||
/// active policy filters `PATH` out.
|
||||
#[cfg(unix)]
|
||||
fn layer_login_path(
|
||||
cmd: &mut tokio::process::Command,
|
||||
login_env: Option<&HashMap<String, String>>,
|
||||
active_policy: Option<&crate::util::ShellEnvironmentPolicy>,
|
||||
) {
|
||||
if let Some(path) = login_env.and_then(|l| l.get("PATH"))
|
||||
&& active_policy.is_none_or(|p| p.allows_with_inherit("PATH"))
|
||||
{
|
||||
cmd.env("PATH", path);
|
||||
}
|
||||
}
|
||||
|
||||
/// Compose the child environment on `cmd` in one place, in a fixed order:
|
||||
/// policy base, login-shell capture, grok control vars, request env, pager
|
||||
/// vars, login `PATH` last, then the agent marker. Untrusted layers (login
|
||||
/// capture and request env) pass through the policy name filter so an excluded
|
||||
/// name cannot re-enter; grok's own control vars, login `PATH`, and the marker
|
||||
/// are applied unfiltered and last. `login_env` is `None` for the persistent
|
||||
/// backend, which restores login state from its own snapshot.
|
||||
///
|
||||
/// The returned `ProcessGroup` is what the teardown helpers
|
||||
/// ([`send_sigterm_to_group`], [`send_sigkill_to_group`]) dispatch to:
|
||||
/// `killpg` on Unix; `TerminateJobObject` on Windows. This gives
|
||||
/// grandchild teardown for fan-out workloads (npm install, git clone,
|
||||
/// cargo build) on both platforms.
|
||||
/// Layers are applied incrementally rather than composed into one map and
|
||||
/// installed via `env_clear`: the default policy is a no-op, and the common
|
||||
/// path must inherit grok's environment untouched (including non-UTF-8 vars).
|
||||
/// A base env is cleared and rebuilt only when a policy is active. Request env
|
||||
/// is filtered by name only, so `inherit = none` still admits explicitly
|
||||
/// provided `.envrc`/ACP vars.
|
||||
///
|
||||
/// Unix only: the Windows spawn path applies the policy inline (it has no
|
||||
/// login-shell capture and uses the shell-invocation env instead of overrides).
|
||||
#[cfg(unix)]
|
||||
fn apply_child_env(
|
||||
cmd: &mut tokio::process::Command,
|
||||
policy: Option<&crate::util::ShellEnvironmentPolicy>,
|
||||
login_env: Option<&HashMap<String, String>>,
|
||||
request_env: &HashMap<String, String>,
|
||||
) {
|
||||
let active_policy = policy.filter(|p| !p.is_noop());
|
||||
// 1. Base env: cleared and rebuilt from the policy only when one is active.
|
||||
crate::util::shell_env_policy::install_policy_base_env(cmd, active_policy);
|
||||
// 2. Login-shell capture (filtered). 3. Grok control vars. 4. Request env
|
||||
// (filtered). 5. Pager vars. 6. Login PATH last. 7. Agent marker wins.
|
||||
layer_login_env_vars(cmd, login_env, active_policy);
|
||||
cmd.envs(shell_state::shell_env_overrides());
|
||||
layer_request_env(cmd, request_env, active_policy);
|
||||
cmd.envs(crate::util::pager_env());
|
||||
layer_login_path(cmd, login_env, active_policy);
|
||||
crate::util::apply_grok_agent_marker(cmd);
|
||||
}
|
||||
|
||||
/// Spawn the shell command and attach the child to a [`ProcessGroup`] for
|
||||
/// grandchild teardown (`killpg` on Unix, `TerminateJobObject` on Windows).
|
||||
fn spawn_shell_command(
|
||||
command: &str,
|
||||
cwd: &std::path::Path,
|
||||
env: &HashMap<String, String>,
|
||||
login_env: Option<&HashMap<String, String>>,
|
||||
search_shadows: SearchShadowConfig,
|
||||
shell_env_policy: Option<&crate::util::ShellEnvironmentPolicy>,
|
||||
) -> std::io::Result<(tokio::process::Child, crate::util::ProcessGroup)> {
|
||||
// `login_env` and `search_shadows` are only consumed by the `#[cfg(unix)]`
|
||||
// shell wrapper below; keep them live on Windows to avoid unused-arg warnings.
|
||||
|
|
@ -2966,30 +3109,7 @@ fn spawn_shell_command(
|
|||
// detach_from_tty() handles both session and process group creation.
|
||||
.kill_on_drop(true);
|
||||
|
||||
if let Some(login) = login_env {
|
||||
for (key, value) in login {
|
||||
if key != "PATH" && std::env::var_os(key).is_none() {
|
||||
cmd.env(key, value);
|
||||
}
|
||||
}
|
||||
}
|
||||
// Apply env vars from the request (e.g., .envrc, color vars, ACP-provided vars).
|
||||
cmd.envs(shell_state::shell_env_overrides());
|
||||
for (key, value) in env {
|
||||
cmd.env(key, value);
|
||||
}
|
||||
cmd.envs(crate::util::pager_env());
|
||||
|
||||
// Inject the user's login-shell PATH LAST so tools installed via rc
|
||||
// files (.bashrc, .zshrc, virtualenvs) are always discoverable. The
|
||||
// request env often carries a copy of the parent process's PATH which
|
||||
// doesn't include rc-file additions — applying login PATH after the
|
||||
// request env ensures those additions aren't clobbered.
|
||||
if let Some(path) = login_env.and_then(|l| l.get("PATH")) {
|
||||
cmd.env("PATH", path);
|
||||
}
|
||||
// Agent marker must win over request/login env.
|
||||
crate::util::apply_grok_agent_marker(&mut cmd);
|
||||
apply_child_env(&mut cmd, shell_env_policy, login_env, env);
|
||||
|
||||
// Detach from the controlling terminal so subprocesses cannot open
|
||||
// /dev/tty and compete with the TUI for terminal input.
|
||||
|
|
@ -3018,18 +3138,21 @@ fn spawn_shell_command(
|
|||
let inv = xai_grok_config::shell::shell_command_argv(command);
|
||||
let mut cmd = tokio::process::Command::new(&inv.program);
|
||||
cmd.args(&inv.args)
|
||||
.envs(inv.env)
|
||||
.current_dir(cwd)
|
||||
.stdin(Stdio::null())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped())
|
||||
.kill_on_drop(true);
|
||||
|
||||
for (key, value) in env {
|
||||
cmd.env(key, value);
|
||||
}
|
||||
// Policy base first (cleared + rebuilt only when a policy is active), then
|
||||
// the shell-invocation env, the filtered request env, pager vars, and the
|
||||
// agent marker last. Mirrors the unix ordering in `apply_child_env`;
|
||||
// `inv.env` is grok's trusted shell setup, so it is not filtered.
|
||||
let active_policy = shell_env_policy.filter(|p| !p.is_noop());
|
||||
crate::util::shell_env_policy::install_policy_base_env(&mut cmd, active_policy);
|
||||
cmd.envs(inv.env);
|
||||
layer_request_env(&mut cmd, env, active_policy);
|
||||
cmd.envs(crate::util::pager_env());
|
||||
// Agent marker must win over request env.
|
||||
crate::util::apply_grok_agent_marker(&mut cmd);
|
||||
|
||||
// Set creation flags inline rather than via crate::util::detach_command
|
||||
|
|
@ -3142,9 +3265,42 @@ mod tests {
|
|||
foreground_block_budget: None,
|
||||
kind: TaskKind::Bash,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn run_background_preserves_description_on_snapshot() {
|
||||
let backend = LocalTerminalBackend::new();
|
||||
let mut with_desc = make_request("sleep 30");
|
||||
with_desc.description = Some("build frontend".to_string());
|
||||
let handle = backend.run_background(with_desc).await.unwrap();
|
||||
let snap = backend
|
||||
.get_task(&handle.task_id)
|
||||
.await
|
||||
.expect("running task snapshot");
|
||||
assert_eq!(snap.description.as_deref(), Some("build frontend"));
|
||||
let listed = backend.list_tasks().await;
|
||||
let listed_snap = listed
|
||||
.iter()
|
||||
.find(|t| t.task_id == handle.task_id)
|
||||
.expect("task listed");
|
||||
assert_eq!(listed_snap.description.as_deref(), Some("build frontend"));
|
||||
let _ = backend.kill_task(&handle.task_id).await;
|
||||
|
||||
let without = make_request("sleep 30");
|
||||
let handle = backend.run_background(without).await.unwrap();
|
||||
let snap = backend
|
||||
.get_task(&handle.task_id)
|
||||
.await
|
||||
.expect("running task snapshot");
|
||||
assert!(
|
||||
snap.description.is_none(),
|
||||
"absent description must stay None"
|
||||
);
|
||||
let _ = backend.kill_task(&handle.task_id).await;
|
||||
}
|
||||
|
||||
/// Poll `get_task` every 25ms until the task reports `completed`, returning
|
||||
/// `false` if `timeout` elapses first. Lets callers keep a bespoke assert
|
||||
/// message while sharing the poll-until-reaped boilerplate.
|
||||
|
|
@ -3170,6 +3326,79 @@ mod tests {
|
|||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn layer_request_env_drops_names_the_policy_excludes() {
|
||||
use crate::util::{EnvironmentVariablePattern, ShellEnvironmentPolicy};
|
||||
|
||||
let glob = EnvironmentVariablePattern::new_case_insensitive;
|
||||
let policy = ShellEnvironmentPolicy {
|
||||
exclude: vec![glob("AWS_*")],
|
||||
include_only: vec![glob("PATH"), glob("SAFE_*")],
|
||||
..Default::default()
|
||||
};
|
||||
let env = HashMap::from([
|
||||
("PATH".to_string(), "/bin".to_string()),
|
||||
("SAFE_FLAG".to_string(), "1".to_string()),
|
||||
("AWS_SECRET".to_string(), "leak".to_string()),
|
||||
("OTHER".to_string(), "x".to_string()),
|
||||
]);
|
||||
|
||||
let mut cmd = tokio::process::Command::new("true");
|
||||
layer_request_env(&mut cmd, &env, Some(&policy));
|
||||
let applied: HashMap<String, String> = cmd
|
||||
.as_std()
|
||||
.get_envs()
|
||||
.filter_map(|(k, v)| Some((k.to_str()?.to_string(), v?.to_str()?.to_string())))
|
||||
.collect();
|
||||
|
||||
assert_eq!(applied.get("PATH").map(String::as_str), Some("/bin"));
|
||||
assert_eq!(applied.get("SAFE_FLAG").map(String::as_str), Some("1"));
|
||||
assert!(!applied.contains_key("AWS_SECRET"));
|
||||
assert!(!applied.contains_key("OTHER"));
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn apply_child_env_layers_in_fixed_order() {
|
||||
use crate::util::{EnvironmentVariablePattern, ShellEnvironmentPolicy};
|
||||
|
||||
let policy = ShellEnvironmentPolicy {
|
||||
exclude: vec![EnvironmentVariablePattern::new_case_insensitive("*SECRET*")],
|
||||
set: HashMap::from([("GROK_TEST_BASE".to_string(), "1".to_string())]),
|
||||
..Default::default()
|
||||
};
|
||||
let login = HashMap::from([
|
||||
("GROK_TEST_LOGIN".to_string(), "l".to_string()),
|
||||
("PATH".to_string(), "/login/bin".to_string()),
|
||||
]);
|
||||
let request = HashMap::from([
|
||||
("GROK_TEST_REQ".to_string(), "r".to_string()),
|
||||
("PATH".to_string(), "/req/bin".to_string()),
|
||||
("GROK_TEST_SECRET".to_string(), "s".to_string()),
|
||||
]);
|
||||
|
||||
let mut cmd = tokio::process::Command::new("true");
|
||||
apply_child_env(&mut cmd, Some(&policy), Some(&login), &request);
|
||||
let env: HashMap<String, String> = cmd
|
||||
.as_std()
|
||||
.get_envs()
|
||||
.filter_map(|(k, v)| Some((k.to_str()?.to_string(), v?.to_str()?.to_string())))
|
||||
.collect();
|
||||
|
||||
assert_eq!(env.get("GROK_TEST_BASE").map(String::as_str), Some("1"));
|
||||
assert_eq!(env.get("GROK_TEST_LOGIN").map(String::as_str), Some("l"));
|
||||
assert_eq!(env.get("GROK_TEST_REQ").map(String::as_str), Some("r"));
|
||||
// Request env is filtered by the policy.
|
||||
assert!(!env.contains_key("GROK_TEST_SECRET"));
|
||||
// Login PATH is applied last and wins over the request PATH.
|
||||
assert_eq!(env.get("PATH").map(String::as_str), Some("/login/bin"));
|
||||
// The agent marker wins over every layer.
|
||||
assert_eq!(
|
||||
env.get(crate::util::GROK_AGENT_ENV).map(String::as_str),
|
||||
Some(crate::util::GROK_AGENT_ENV_VALUE)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore = "flaky: combined_output is sometimes empty in CI"]
|
||||
async fn test_simple_command() {
|
||||
|
|
@ -3220,6 +3449,7 @@ mod tests {
|
|||
foreground_block_budget: None,
|
||||
kind: TaskKind::Bash,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
|
||||
let result = backend.run(request).await.unwrap();
|
||||
|
|
@ -3251,6 +3481,7 @@ mod tests {
|
|||
foreground_block_budget: None,
|
||||
kind: TaskKind::Bash,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
|
||||
let result = backend.run(request).await.unwrap();
|
||||
|
|
@ -3314,6 +3545,7 @@ mod tests {
|
|||
foreground_block_budget: None,
|
||||
kind: TaskKind::Bash,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
|
||||
let start = Instant::now();
|
||||
|
|
@ -3385,6 +3617,7 @@ mod tests {
|
|||
foreground_block_budget: None,
|
||||
kind: TaskKind::Bash,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
|
||||
let result = backend.run(request).await.unwrap();
|
||||
|
|
@ -3431,6 +3664,7 @@ mod tests {
|
|||
foreground_block_budget: Some(Duration::from_millis(300)),
|
||||
kind: TaskKind::Bash,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
|
||||
let start = Instant::now();
|
||||
|
|
@ -3482,6 +3716,7 @@ mod tests {
|
|||
foreground_block_budget: Some(Duration::MAX),
|
||||
kind: TaskKind::Bash,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
|
||||
let start = Instant::now();
|
||||
|
|
@ -3533,6 +3768,7 @@ mod tests {
|
|||
foreground_block_budget: None,
|
||||
kind: TaskKind::Bash,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
|
||||
let result = backend.run(request).await.unwrap();
|
||||
|
|
@ -3591,6 +3827,7 @@ mod tests {
|
|||
foreground_block_budget: None,
|
||||
kind: TaskKind::Bash,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
|
||||
let result = backend.run(request).await.unwrap();
|
||||
|
|
@ -3627,6 +3864,7 @@ mod tests {
|
|||
foreground_block_budget: None,
|
||||
kind: TaskKind::Bash,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
|
||||
// Start background task
|
||||
|
|
@ -3667,6 +3905,7 @@ mod tests {
|
|||
foreground_block_budget: None,
|
||||
kind: TaskKind::Bash,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
|
||||
let handle = backend.run_background(request).await.unwrap();
|
||||
|
|
@ -3707,6 +3946,7 @@ mod tests {
|
|||
foreground_block_budget: None,
|
||||
kind: TaskKind::Bash,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
|
||||
let result = backend.run(request).await.unwrap();
|
||||
|
|
@ -3782,6 +4022,7 @@ mod tests {
|
|||
foreground_block_budget: None,
|
||||
kind: TaskKind::Bash,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
|
||||
let result = backend.run(request).await.unwrap();
|
||||
|
|
@ -3848,6 +4089,7 @@ mod tests {
|
|||
foreground_block_budget: None,
|
||||
kind: TaskKind::Bash,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
|
||||
let result = backend.run(request).await.unwrap();
|
||||
|
|
@ -3883,6 +4125,7 @@ mod tests {
|
|||
foreground_block_budget: None,
|
||||
kind: TaskKind::Bash,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
|
||||
let result = backend.run(request).await.unwrap();
|
||||
|
|
@ -3917,6 +4160,7 @@ mod tests {
|
|||
foreground_block_budget: None,
|
||||
kind: TaskKind::Bash,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
|
||||
let result = backend.run(request).await.unwrap();
|
||||
|
|
@ -3947,6 +4191,7 @@ mod tests {
|
|||
foreground_block_budget: None,
|
||||
kind: TaskKind::Bash,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
|
||||
backend.run(request).await.unwrap();
|
||||
|
|
@ -3986,6 +4231,7 @@ mod tests {
|
|||
foreground_block_budget: None,
|
||||
kind: TaskKind::Bash,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
|
||||
backend.run(request).await.unwrap();
|
||||
|
|
@ -4034,6 +4280,7 @@ mod tests {
|
|||
foreground_block_budget: None,
|
||||
kind: TaskKind::Bash,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
|
||||
let result = backend.run(request).await.unwrap();
|
||||
|
|
@ -4067,6 +4314,7 @@ mod tests {
|
|||
foreground_block_budget: None,
|
||||
kind: TaskKind::Bash,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
|
||||
let handle = backend.run_background(request).await.unwrap();
|
||||
|
|
@ -4111,6 +4359,7 @@ mod tests {
|
|||
foreground_block_budget: None,
|
||||
kind: TaskKind::Bash,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
|
||||
let result = backend.run(request).await.unwrap();
|
||||
|
|
@ -4143,6 +4392,7 @@ mod tests {
|
|||
foreground_block_budget: None,
|
||||
kind: TaskKind::Bash,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
|
||||
let start = Instant::now();
|
||||
|
|
|
|||
|
|
@ -110,6 +110,8 @@ pub struct TerminalRunRequest {
|
|||
/// `kill_all_background_tasks_by_owner` only targets the requesting
|
||||
/// session's processes — not the parent's or sibling's.
|
||||
pub owner_session_id: Option<String>,
|
||||
/// Model-supplied label for task UI / snapshots.
|
||||
pub description: Option<String>,
|
||||
}
|
||||
|
||||
/// Distinguishes different types of background tasks.
|
||||
|
|
@ -214,6 +216,9 @@ pub struct TaskSnapshot {
|
|||
/// the parent's or sibling's.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub owner_session_id: Option<String>,
|
||||
/// Model-supplied label for task UI / snapshots.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub description: Option<String>,
|
||||
}
|
||||
|
||||
impl TaskSnapshot {
|
||||
|
|
|
|||
|
|
@ -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(¶ms),
|
||||
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,
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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));
|
||||
}
|
||||
|
|
@ -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;
|
||||
|
|
@ -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);
|
||||
}
|
||||
}
|
||||
|
|
@ -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()
|
||||
}
|
||||
File diff suppressed because it is too large
Load diff
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
},
|
||||
|
|
|
|||
|
|
@ -916,6 +916,7 @@ pub(crate) mod test_helpers {
|
|||
block_waited: false,
|
||||
explicitly_killed: false,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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"#;
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -107,6 +107,7 @@ mod tests {
|
|||
block_waited: false,
|
||||
explicitly_killed: false,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,11 +1,13 @@
|
|||
use std::collections::VecDeque;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
|
||||
use super::types::{
|
||||
BashExecutionBackgrounded, BashExecutionComplete, BashExecutionFailed, BashExecutionTimeout,
|
||||
BashOutputChunk, FileWritten, LspServerCrashed, LspServerFailed, LspServerReady,
|
||||
LspServerRetrying, LspServerStarting, MonitorEvent, PlanModeEntered, PlanModeExited,
|
||||
ScheduledTaskCreated, ScheduledTaskFired, ScheduledTaskRemoved, ToolNotification,
|
||||
UserQuestionAsked,
|
||||
ScheduledTaskCreated, ScheduledTaskFired, ScheduledTaskRemoved, SubagentCompleted,
|
||||
ToolNotification, UserQuestionAsked,
|
||||
};
|
||||
use crate::types::TaskSnapshot;
|
||||
|
||||
|
|
@ -90,9 +92,87 @@ impl NotificationAcknowledgementBatch {
|
|||
#[derive(Clone)]
|
||||
enum ToolNotificationTarget {
|
||||
Plain(tokio::sync::mpsc::UnboundedSender<ToolNotification>),
|
||||
Bounded(tokio::sync::mpsc::Sender<ToolNotification>),
|
||||
Capped(Arc<CappedNotificationQueue>),
|
||||
Acknowledged(tokio::sync::mpsc::UnboundedSender<AcknowledgedToolNotification>),
|
||||
}
|
||||
|
||||
struct CappedNotificationQueue {
|
||||
queue: parking_lot::Mutex<VecDeque<ToolNotification>>,
|
||||
capacity: usize,
|
||||
closed: AtomicBool,
|
||||
ready: tokio::sync::Notify,
|
||||
}
|
||||
|
||||
impl CappedNotificationQueue {
|
||||
fn push(&self, notification: ToolNotification) {
|
||||
if self.closed.load(Ordering::Relaxed) {
|
||||
return;
|
||||
}
|
||||
let mut queue = self.queue.lock();
|
||||
if queue.len() >= self.capacity {
|
||||
if !is_critical_notification(¬ification) {
|
||||
tracing::warn!("tool notification queue full; dropping newest lossy event");
|
||||
return;
|
||||
}
|
||||
let evict = queue
|
||||
.iter()
|
||||
.position(|queued| !is_critical_notification(queued))
|
||||
.unwrap_or(0);
|
||||
queue.remove(evict);
|
||||
tracing::warn!("tool notification queue full; evicting older event for terminal event");
|
||||
}
|
||||
queue.push_back(notification);
|
||||
drop(queue);
|
||||
self.ready.notify_one();
|
||||
}
|
||||
}
|
||||
|
||||
fn is_critical_notification(notification: &ToolNotification) -> bool {
|
||||
matches!(
|
||||
notification,
|
||||
ToolNotification::BashExecutionComplete(_)
|
||||
| ToolNotification::BashExecutionTimeout(_)
|
||||
| ToolNotification::BashExecutionFailed(_)
|
||||
| ToolNotification::TaskCompleted(_)
|
||||
| ToolNotification::SubagentCompleted(_)
|
||||
| ToolNotification::PlanModeEntered(_)
|
||||
| ToolNotification::PlanModeExited(_)
|
||||
| ToolNotification::UserQuestionAsked(_)
|
||||
| ToolNotification::LspServerCrashed(_)
|
||||
| ToolNotification::LspServerFailed(_)
|
||||
| ToolNotification::ScheduledTaskFired(_)
|
||||
| ToolNotification::ScheduledTaskRemoved(_)
|
||||
)
|
||||
}
|
||||
|
||||
/// Receiver for a capped queue that preserves terminal notifications.
|
||||
pub struct CappedToolNotificationReceiver {
|
||||
queue: Arc<CappedNotificationQueue>,
|
||||
}
|
||||
|
||||
impl CappedToolNotificationReceiver {
|
||||
pub async fn recv(&mut self) -> Option<ToolNotification> {
|
||||
loop {
|
||||
let ready = self.queue.ready.notified();
|
||||
if let Some(notification) = self.queue.queue.lock().pop_front() {
|
||||
return Some(notification);
|
||||
}
|
||||
if self.queue.closed.load(Ordering::Relaxed) {
|
||||
return None;
|
||||
}
|
||||
ready.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for CappedToolNotificationReceiver {
|
||||
fn drop(&mut self) {
|
||||
self.queue.closed.store(true, Ordering::Relaxed);
|
||||
self.queue.ready.notify_waiters();
|
||||
}
|
||||
}
|
||||
|
||||
/// Cloneable notification fan-out with per-target FIFO ordering.
|
||||
#[derive(Clone)]
|
||||
pub struct ToolNotificationHandle {
|
||||
|
|
@ -127,6 +207,35 @@ impl ToolNotificationHandle {
|
|||
(Self::new(sender), receiver)
|
||||
}
|
||||
|
||||
/// Create a capped target that drops the newest event when full.
|
||||
pub fn bounded_channel(
|
||||
capacity: usize,
|
||||
) -> (Self, tokio::sync::mpsc::Receiver<ToolNotification>) {
|
||||
let (sender, receiver) = tokio::sync::mpsc::channel(capacity);
|
||||
(
|
||||
Self {
|
||||
targets: Arc::from([ToolNotificationTarget::Bounded(sender)]),
|
||||
},
|
||||
receiver,
|
||||
)
|
||||
}
|
||||
|
||||
/// Create a capped queue that evicts lossy events before terminal events.
|
||||
pub fn capped_channel(capacity: usize) -> (Self, CappedToolNotificationReceiver) {
|
||||
let queue = Arc::new(CappedNotificationQueue {
|
||||
queue: parking_lot::Mutex::new(VecDeque::new()),
|
||||
capacity: capacity.max(1),
|
||||
closed: AtomicBool::new(false),
|
||||
ready: tokio::sync::Notify::new(),
|
||||
});
|
||||
(
|
||||
Self {
|
||||
targets: Arc::from([ToolNotificationTarget::Capped(Arc::clone(&queue))]),
|
||||
},
|
||||
CappedToolNotificationReceiver { queue },
|
||||
)
|
||||
}
|
||||
|
||||
pub fn acknowledged_channel() -> (
|
||||
Self,
|
||||
tokio::sync::mpsc::UnboundedReceiver<AcknowledgedToolNotification>,
|
||||
|
|
@ -185,6 +294,12 @@ impl ToolNotificationHandle {
|
|||
ToolNotificationTarget::Plain(target) => {
|
||||
let _ = target.send(notification);
|
||||
}
|
||||
ToolNotificationTarget::Bounded(target) => {
|
||||
if target.try_send(notification).is_err() {
|
||||
tracing::warn!("tool notification queue full; dropping newest event");
|
||||
}
|
||||
}
|
||||
ToolNotificationTarget::Capped(target) => target.push(notification),
|
||||
ToolNotificationTarget::Acknowledged(target) => {
|
||||
let _ = target.send(AcknowledgedToolNotification {
|
||||
notification,
|
||||
|
|
@ -210,6 +325,12 @@ impl ToolNotificationHandle {
|
|||
ToolNotificationTarget::Plain(target) => {
|
||||
let _ = target.send(notification.clone());
|
||||
}
|
||||
ToolNotificationTarget::Bounded(target) => {
|
||||
if target.try_send(notification.clone()).is_err() {
|
||||
tracing::warn!("tool notification queue full; dropping newest event");
|
||||
}
|
||||
}
|
||||
ToolNotificationTarget::Capped(target) => target.push(notification.clone()),
|
||||
ToolNotificationTarget::Acknowledged(target) => {
|
||||
let (acknowledgement, receipt) = tokio::sync::oneshot::channel();
|
||||
if target
|
||||
|
|
@ -237,6 +358,7 @@ impl ToolNotificationHandle {
|
|||
send_failed, BashExecutionFailed, BashExecutionFailed;
|
||||
send_file_written, FileWritten, FileWritten;
|
||||
send_task_complete, TaskSnapshot, TaskCompleted;
|
||||
send_subagent_completed, SubagentCompleted, SubagentCompleted;
|
||||
send_plan_mode_entered, PlanModeEntered, PlanModeEntered;
|
||||
send_plan_mode_exited, PlanModeExited, PlanModeExited;
|
||||
send_user_question_asked, UserQuestionAsked, UserQuestionAsked;
|
||||
|
|
|
|||
|
|
@ -106,3 +106,22 @@ async fn batch_distinguishes_dropped_and_rejected_acknowledgements() {
|
|||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bounded_channel_drops_newest_when_full() {
|
||||
let (handle, mut receiver) = ToolNotificationHandle::bounded_channel(1);
|
||||
handle.send_scheduled_task_created(created("kept"));
|
||||
handle.send_scheduled_task_created(created("dropped"));
|
||||
|
||||
assert_eq!(task_id(&receiver.recv().await.unwrap()), "kept");
|
||||
assert!(receiver.try_recv().is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn capped_channel_evicts_lossy_event_for_terminal_event() {
|
||||
let (handle, mut receiver) = ToolNotificationHandle::capped_channel(1);
|
||||
handle.send_scheduled_task_created(created("lossy"));
|
||||
handle.send(ToolNotification::ScheduledTaskRemoved(removed("terminal")));
|
||||
|
||||
assert_eq!(task_id(&receiver.recv().await.unwrap()), "terminal");
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ pub mod handle;
|
|||
pub mod types;
|
||||
|
||||
pub use handle::AcknowledgedToolNotification;
|
||||
pub use handle::CappedToolNotificationReceiver;
|
||||
pub use handle::DurableNotificationTargets;
|
||||
pub use handle::NotificationAcknowledgementBatch;
|
||||
pub use handle::NotificationAcknowledgementError;
|
||||
|
|
|
|||
|
|
@ -368,6 +368,21 @@ pub struct MonitorEvent {
|
|||
pub owner_session_id: Option<String>,
|
||||
}
|
||||
|
||||
/// A background subagent reached a terminal state while the parent held a handle.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, schemars::JsonSchema)]
|
||||
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
|
||||
pub struct SubagentCompleted {
|
||||
pub subagent_id: String,
|
||||
pub subagent_type: String,
|
||||
pub description: String,
|
||||
pub status: String,
|
||||
#[cfg_attr(feature = "serde", serde(default))]
|
||||
pub error: Option<String>,
|
||||
pub duration_ms: u64,
|
||||
#[cfg_attr(feature = "serde", serde(default))]
|
||||
pub owner_session_id: Option<String>,
|
||||
}
|
||||
|
||||
/// A notification emitted by a tool during or after execution.
|
||||
/// These are sent to external consumers (TUI, logging, etc.) to provide
|
||||
/// real-time visibility into tool execution.
|
||||
|
|
@ -398,6 +413,9 @@ pub enum ToolNotification {
|
|||
/// about the task being finished status
|
||||
TaskCompleted(TaskSnapshot),
|
||||
|
||||
/// A background subagent reached a terminal state.
|
||||
SubagentCompleted(SubagentCompleted),
|
||||
|
||||
/// The agent requested to enter plan mode.
|
||||
/// Consumers (gateway, TUI) use this to transition the client into
|
||||
/// plan-mode UI state (e.g., enforce read-only, inject plan-mode
|
||||
|
|
@ -480,6 +498,7 @@ notification_variants! {
|
|||
BashExecutionFailed => BashExecutionFailed,
|
||||
FileWritten => FileWritten,
|
||||
TaskCompleted => TaskSnapshot,
|
||||
SubagentCompleted => SubagentCompleted,
|
||||
PlanModeEntered => PlanModeEntered,
|
||||
PlanModeExited => PlanModeExited,
|
||||
UserQuestionAsked => UserQuestionAsked,
|
||||
|
|
|
|||
|
|
@ -210,6 +210,15 @@ pub struct ToolServerConfig {
|
|||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub behavior_preset: Option<String>,
|
||||
}
|
||||
#[derive(Clone)]
|
||||
pub struct SubagentSessionResources {
|
||||
pub backend: crate::implementations::grok_build::task::backend::SubagentBackendResource,
|
||||
/// Same channel as [`Self::backend`]; required so `TaskCompletionReminder`
|
||||
/// can drain completions onto the next tool result (shell parity).
|
||||
pub event_sender: crate::implementations::grok_build::task::types::SubagentEventSender,
|
||||
pub depth: crate::implementations::grok_build::task::types::SubagentDepthCounter,
|
||||
pub session_id: crate::implementations::grok_build::task::types::SessionIdResource,
|
||||
}
|
||||
/// Everything a session provides at finalization time.
|
||||
///
|
||||
/// This is the **public API boundary** — callers pass concrete, strongly-typed
|
||||
|
|
@ -231,6 +240,8 @@ pub struct SessionContext {
|
|||
/// Session ID that owns processes spawned by this session's tools.
|
||||
/// Used to scope kill operations on a shared terminal backend.
|
||||
pub owner_session_id: Option<String>,
|
||||
/// Complete subagent capability for this session.
|
||||
pub subagent: Option<SubagentSessionResources>,
|
||||
/// Parent's scheduler handle. When `Some`, the session reuses the parent's
|
||||
/// scheduler actor instead of spawning its own, so scheduled tasks survive
|
||||
/// subagent exit.
|
||||
|
|
@ -975,6 +986,12 @@ impl ToolRegistryBuilder {
|
|||
if let Some(owner_session_id) = ctx.owner_session_id {
|
||||
resources.insert(crate::types::resources::OwnerSessionId(owner_session_id));
|
||||
}
|
||||
if let Some(subagent) = ctx.subagent {
|
||||
resources.insert(subagent.backend);
|
||||
resources.insert(subagent.event_sender);
|
||||
resources.insert(subagent.depth);
|
||||
resources.insert(subagent.session_id);
|
||||
}
|
||||
let scheduler_notification_handle = ctx.notification_handle.clone();
|
||||
resources.insert(crate::types::resources::NotificationHandle(
|
||||
ctx.notification_handle,
|
||||
|
|
@ -1303,6 +1320,23 @@ impl FinalizedToolset {
|
|||
.map(|t| (t.client_name.clone(), t.metadata.kind().as_key().to_owned()))
|
||||
.collect()
|
||||
}
|
||||
/// Map of client-facing tool name → typed [`ToolKind`].
|
||||
///
|
||||
/// Unlike the finalize-request `ToolConfig`s (whose `kind` is `None` when
|
||||
/// built from raw IDs over gRPC), the finalized tools always know their
|
||||
/// real kind from the registry metadata — use this for kind-derived
|
||||
/// metadata in server responses (e.g. capability-mode classification).
|
||||
pub fn tool_kind_map(&self) -> HashMap<String, ToolKind> {
|
||||
self.tools
|
||||
.read()
|
||||
.iter()
|
||||
.map(|t| (t.client_name.clone(), t.metadata.kind()))
|
||||
.collect()
|
||||
}
|
||||
/// Finalized canonical-to-client parameter names by tool kind.
|
||||
pub fn template_param_names(&self) -> HashMap<ToolKind, HashMap<String, String>> {
|
||||
self.renderer.param_names()
|
||||
}
|
||||
pub async fn update_resource<T: Send + Sync + 'static>(&self, resource: T) {
|
||||
self.resources.lock().await.insert(resource);
|
||||
}
|
||||
|
|
@ -1420,6 +1454,9 @@ impl FinalizedToolset {
|
|||
let mut ctx = xai_tool_runtime::ToolCallContext::new(parent_ctx.call_id.clone());
|
||||
ctx.extensions.insert(self.resources.clone());
|
||||
ctx.extensions.insert_arc(Arc::clone(&self.renderer));
|
||||
if let Some(cancellation) = parent_ctx.get::<xai_tool_runtime::Cancellation>() {
|
||||
ctx.extensions.insert((*cancellation).clone());
|
||||
}
|
||||
ctx.extensions.insert(
|
||||
crate::types::resources::InvokingToolParamNames::from_reverse_params(&reverse_params),
|
||||
);
|
||||
|
|
@ -1451,9 +1488,27 @@ impl FinalizedToolset {
|
|||
tool_args: serde_json::Value,
|
||||
tool_call_id: &str,
|
||||
cwd_override: Option<std::path::PathBuf>,
|
||||
) -> Result<ToolRunResult, xai_tool_runtime::ToolError> {
|
||||
self.call_with_cancellation(tool_name, tool_args, tool_call_id, cwd_override, None)
|
||||
.await
|
||||
}
|
||||
/// Dispatch with cooperative cancellation exposed to the tool.
|
||||
pub async fn call_with_cancellation(
|
||||
self: &Arc<Self>,
|
||||
tool_name: &str,
|
||||
tool_args: serde_json::Value,
|
||||
tool_call_id: &str,
|
||||
cwd_override: Option<std::path::PathBuf>,
|
||||
cancellation: Option<tokio_util::sync::CancellationToken>,
|
||||
) -> Result<ToolRunResult, xai_tool_runtime::ToolError> {
|
||||
use futures::StreamExt;
|
||||
let mut stream = self.call_streaming(tool_name, tool_args, tool_call_id, cwd_override);
|
||||
let mut stream = self.call_streaming_with_cancellation(
|
||||
tool_name,
|
||||
tool_args,
|
||||
tool_call_id,
|
||||
cwd_override,
|
||||
cancellation,
|
||||
);
|
||||
while let Some(item) = stream.next().await {
|
||||
match item {
|
||||
xai_tool_runtime::ToolStreamItem::Progress(_) => continue,
|
||||
|
|
@ -1482,6 +1537,23 @@ impl FinalizedToolset {
|
|||
tool_args: serde_json::Value,
|
||||
tool_call_id: &str,
|
||||
cwd_override: Option<std::path::PathBuf>,
|
||||
) -> xai_tool_runtime::ToolStream<ToolRunResult> {
|
||||
self.call_streaming_with_cancellation(
|
||||
tool_name,
|
||||
tool_args,
|
||||
tool_call_id,
|
||||
cwd_override,
|
||||
None,
|
||||
)
|
||||
}
|
||||
/// Streaming dispatch with cooperative cancellation exposed to the tool.
|
||||
pub fn call_streaming_with_cancellation(
|
||||
self: &Arc<Self>,
|
||||
tool_name: &str,
|
||||
tool_args: serde_json::Value,
|
||||
tool_call_id: &str,
|
||||
cwd_override: Option<std::path::PathBuf>,
|
||||
cancellation: Option<tokio_util::sync::CancellationToken>,
|
||||
) -> xai_tool_runtime::ToolStream<ToolRunResult> {
|
||||
use futures::StreamExt;
|
||||
let this = Arc::clone(self);
|
||||
|
|
@ -1493,6 +1565,7 @@ impl FinalizedToolset {
|
|||
tool_args,
|
||||
&tool_call_id,
|
||||
cwd_override,
|
||||
cancellation,
|
||||
) {
|
||||
Ok(parts) => parts,
|
||||
Err(e) => {
|
||||
|
|
@ -1542,6 +1615,7 @@ impl FinalizedToolset {
|
|||
tool_args: serde_json::Value,
|
||||
tool_call_id: &str,
|
||||
cwd_override: Option<std::path::PathBuf>,
|
||||
cancellation: Option<tokio_util::sync::CancellationToken>,
|
||||
) -> Result<DispatchParts, xai_tool_runtime::ToolError> {
|
||||
let (registry_id, output_converter, reverse_params) = {
|
||||
let tools = self.tools.read();
|
||||
|
|
@ -1581,6 +1655,10 @@ impl FinalizedToolset {
|
|||
if let Some(cwd) = cwd_override {
|
||||
ctx.extensions.insert(xai_tool_runtime::Cwd(cwd));
|
||||
}
|
||||
if let Some(cancellation) = cancellation {
|
||||
ctx.extensions
|
||||
.insert(xai_tool_runtime::Cancellation(cancellation));
|
||||
}
|
||||
if let Some(ref version) = contract_version {
|
||||
ctx.extensions
|
||||
.insert(xai_tool_runtime::BehaviorVersion(version.clone()));
|
||||
|
|
@ -1807,6 +1885,10 @@ pub fn generate_schema<T: schemars::JsonSchema>() -> serde_json::Value {
|
|||
let generator = settings.into_generator();
|
||||
let schema = generator.into_root_schema_for::<T>();
|
||||
let mut value = serde_json::to_value(&schema).unwrap_or_default();
|
||||
if let Some(obj) = value.as_object_mut() {
|
||||
obj.remove("title");
|
||||
obj.remove("description");
|
||||
}
|
||||
if let Some(obj) = value.as_object_mut()
|
||||
&& obj.get("type").and_then(|v| v.as_str()) == Some("object")
|
||||
{
|
||||
|
|
@ -2036,6 +2118,7 @@ mod tests {
|
|||
session_env: Arc::new(HashMap::new()),
|
||||
notification_handle: crate::notification::ToolNotificationHandle::noop(),
|
||||
owner_session_id: None,
|
||||
subagent: None,
|
||||
parent_scheduler_handle: None,
|
||||
skills: vec![],
|
||||
state_path: tmp.path().join("state.json"),
|
||||
|
|
@ -4415,6 +4498,29 @@ mod tests {
|
|||
assert_eq!(skills.0.len(), 2, "should have exactly 2 skills");
|
||||
}
|
||||
}
|
||||
/// generate_schema strips the boilerplate root `title` (struct name) and
|
||||
/// root `description` (struct doc "Input for the <canonical> tool") so the
|
||||
/// canonical name can't leak via parameters.description after randomization
|
||||
/// renames the tool. $schema and per-property descriptions are retained.
|
||||
#[test]
|
||||
fn generate_schema_strips_root_title_and_description() {
|
||||
let schema = generate_schema::<crate::implementations::grok_build::bash::BashToolInput>();
|
||||
assert!(
|
||||
schema.get("title").is_none(),
|
||||
"root title (struct name) must be stripped: {schema}"
|
||||
);
|
||||
assert!(
|
||||
schema.get("description").is_none(),
|
||||
"root description (leaks canonical tool name) must be stripped: {schema}"
|
||||
);
|
||||
assert!(schema.get("$schema").is_some(), "$schema must be retained");
|
||||
assert!(
|
||||
schema["properties"]
|
||||
.as_object()
|
||||
.is_some_and(|p| !p.is_empty()),
|
||||
"per-property schema must be retained: {schema}"
|
||||
);
|
||||
}
|
||||
fn toolset_with_viewer_ctx(
|
||||
viewer_ctx: Option<xai_tool_runtime::WorkspaceViewerContext>,
|
||||
) -> (Arc<FinalizedToolset>, TempDir) {
|
||||
|
|
@ -4455,6 +4561,7 @@ mod tests {
|
|||
serde_json::json!({"target_file": "noop"}),
|
||||
"test-call",
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("prepare_dispatch succeeds");
|
||||
let wvc = parts
|
||||
|
|
@ -4473,6 +4580,7 @@ mod tests {
|
|||
serde_json::json!({"target_file": "noop"}),
|
||||
"test-call",
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("prepare_dispatch succeeds");
|
||||
assert!(
|
||||
|
|
|
|||
|
|
@ -643,14 +643,13 @@ impl Reminder for TaskCompletionReminder {
|
|||
.chain(&reserved_ids)
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
let (terminal, event_sender, session_id) = {
|
||||
let (terminal, event_sender, parent_session_id) = {
|
||||
let res = resources.lock().await;
|
||||
(
|
||||
res.get::<Terminal>().map(|t| t.0.clone()),
|
||||
res.get::<SubagentEventSender>().cloned(),
|
||||
res.get::<crate::implementations::grok_build::task::types::SessionIdResource>()
|
||||
.map(|s| s.0.clone())
|
||||
.unwrap_or_default(),
|
||||
res.get::<crate::types::resources::OwnerSessionId>()
|
||||
.map(|owner| owner.0.clone()),
|
||||
)
|
||||
};
|
||||
let mut reminders = Vec::new();
|
||||
|
|
@ -733,7 +732,7 @@ impl Reminder for TaskCompletionReminder {
|
|||
if sender
|
||||
.0
|
||||
.send(SubagentEvent::Completions(SubagentCompletionsRequest {
|
||||
session_id,
|
||||
parent_session_id,
|
||||
suppress_ids,
|
||||
respond_to: tx,
|
||||
}))
|
||||
|
|
@ -802,6 +801,7 @@ mod tests {
|
|||
block_waited: false,
|
||||
explicitly_killed: false,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
let msg = format_bash_completion(&task, Some("get_command_or_subagent_output"), None);
|
||||
assert!(msg.contains("abc-123"));
|
||||
|
|
@ -828,6 +828,7 @@ mod tests {
|
|||
block_waited: false,
|
||||
explicitly_killed: false,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
let msg = format_monitor_completion(&task, Some("get_command_or_subagent_output"));
|
||||
assert!(
|
||||
|
|
@ -860,6 +861,7 @@ mod tests {
|
|||
block_waited: false,
|
||||
explicitly_killed: false,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
let msg = format_monitor_completion(&task, None);
|
||||
assert!(
|
||||
|
|
@ -887,6 +889,7 @@ mod tests {
|
|||
block_waited: false,
|
||||
explicitly_killed: false,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
let msg = format_bash_completion(&task, Some("get_command_or_subagent_output"), None);
|
||||
assert!(msg.contains("cargo test"));
|
||||
|
|
@ -911,6 +914,7 @@ mod tests {
|
|||
block_waited: false,
|
||||
explicitly_killed: false,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
let msg = format_bash_completion(&task, Some("get_command_or_subagent_output"), None);
|
||||
assert!(msg.contains("exit code: unknown"));
|
||||
|
|
@ -938,6 +942,7 @@ mod tests {
|
|||
block_waited: false,
|
||||
explicitly_killed: false,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
let msg = format_bash_completion(&task, Some("get_command_or_subagent_output"), None);
|
||||
assert!(
|
||||
|
|
@ -976,6 +981,7 @@ mod tests {
|
|||
block_waited: false,
|
||||
explicitly_killed: false,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
let msg = format_bash_completion(&task, Some("get_command_or_subagent_output"), None);
|
||||
assert!(
|
||||
|
|
@ -1013,6 +1019,7 @@ mod tests {
|
|||
block_waited: false,
|
||||
explicitly_killed: false,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
};
|
||||
let msg = format_bash_completion(&task, Some("get_command_or_subagent_output"), None);
|
||||
assert!(msg.contains("exit code: 0"));
|
||||
|
|
@ -1173,6 +1180,7 @@ mod tests {
|
|||
block_waited: false,
|
||||
explicitly_killed: false,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
}
|
||||
}
|
||||
fn make_running(id: &str) -> TaskSnapshot {
|
||||
|
|
@ -1193,6 +1201,7 @@ mod tests {
|
|||
block_waited: false,
|
||||
explicitly_killed: false,
|
||||
owner_session_id: None,
|
||||
description: None,
|
||||
}
|
||||
}
|
||||
fn make_bg_started(id: &str) -> crate::types::output::BackgroundTaskStarted {
|
||||
|
|
@ -1438,7 +1447,6 @@ mod tests {
|
|||
fn make_subagent_completion(id: &str, success: bool) -> SubagentCompletionSummary {
|
||||
SubagentCompletionSummary {
|
||||
subagent_id: id.into(),
|
||||
owner_session_id: String::new(),
|
||||
subagent_type: "general-purpose".into(),
|
||||
description: "test task".into(),
|
||||
success,
|
||||
|
|
|
|||
|
|
@ -175,6 +175,19 @@ impl TemplateRenderer {
|
|||
render_with_env(template, &self.ctx)
|
||||
}
|
||||
|
||||
/// Return the finalized canonical-to-client parameter names by tool kind.
|
||||
pub fn param_names(&self) -> HashMap<ToolKind, HashMap<String, String>> {
|
||||
self.ctx
|
||||
.params
|
||||
.iter()
|
||||
.filter_map(|(kind, names)| {
|
||||
serde_json::from_value(serde_json::Value::String(kind.clone()))
|
||||
.ok()
|
||||
.map(|kind| (kind, names.clone()))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Render `${{ ... }}` placeholders in every `description` string within a
|
||||
/// JSON Schema, in place — recursing into nested objects, array `items`, and
|
||||
/// `$defs`. Property keys are remapped separately; this resolves
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ pub mod path_suggestions;
|
|||
pub(crate) mod query_tools;
|
||||
pub mod remap;
|
||||
pub mod serde_base64;
|
||||
pub mod shell_env_policy;
|
||||
pub mod spawn;
|
||||
pub mod truncate;
|
||||
pub mod unicode_confusables;
|
||||
|
|
@ -26,6 +27,10 @@ pub use fs::{UnicodePathMatch, canonicalize_with_timeout, try_resolve_unicode_fi
|
|||
pub use grok_home::{grok_application, grok_home};
|
||||
pub use path_suggestions::format_not_found_error;
|
||||
pub use remap::{remap_json_keys, remap_schema_properties, reverse_map};
|
||||
pub use shell_env_policy::{
|
||||
EnvironmentVariablePattern, ShellEnvironmentPolicy, ShellEnvironmentPolicyInherit,
|
||||
apply_shell_environment_policy,
|
||||
};
|
||||
pub use spawn::{
|
||||
ProcessGroup, ProcessScope, detach_command, global_process_scope, new_process_group,
|
||||
};
|
||||
|
|
|
|||
237
crates/codegen/xai-grok-tools/src/util/shell_env_policy.rs
Normal file
237
crates/codegen/xai-grok-tools/src/util/shell_env_policy.rs
Normal file
|
|
@ -0,0 +1,237 @@
|
|||
//! Controls which environment variables agent subprocesses (bash tool,
|
||||
//! terminals) inherit. Default is a no-op (inherit everything); enforced at the
|
||||
//! shell spawn sites on macOS, Linux, and Windows.
|
||||
|
||||
use serde::Deserialize;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::LazyLock;
|
||||
use wildmatch::WildMatchPattern;
|
||||
|
||||
/// Case-insensitive environment-variable-name glob (`*`, `?`).
|
||||
pub type EnvironmentVariablePattern = WildMatchPattern<'*', '?'>;
|
||||
|
||||
fn deserialize_patterns<'de, D>(
|
||||
deserializer: D,
|
||||
) -> Result<Vec<EnvironmentVariablePattern>, D::Error>
|
||||
where
|
||||
D: serde::Deserializer<'de>,
|
||||
{
|
||||
let globs = Vec::<String>::deserialize(deserializer)?;
|
||||
Ok(globs
|
||||
.iter()
|
||||
.map(|s| EnvironmentVariablePattern::new_case_insensitive(s))
|
||||
.collect())
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Deserialize)]
|
||||
#[serde(rename_all = "kebab-case")]
|
||||
pub enum ShellEnvironmentPolicyInherit {
|
||||
/// Core platform variables only (PATH, HOME, SHELL, ...).
|
||||
Core,
|
||||
#[default]
|
||||
All,
|
||||
None,
|
||||
}
|
||||
|
||||
/// How to build the environment for agent subprocesses. Applied in order: start
|
||||
/// from `inherit`; if `ignore_default_excludes` is false, drop the secret
|
||||
/// patterns `*KEY*`/`*SECRET*`/`*TOKEN*`; drop `exclude`; insert `set`; if
|
||||
/// `include_only` is non-empty, keep only those. Patterns are case-insensitive
|
||||
/// globs (`*`, `?`).
|
||||
#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
|
||||
#[serde(default)]
|
||||
pub struct ShellEnvironmentPolicy {
|
||||
pub inherit: ShellEnvironmentPolicyInherit,
|
||||
/// Skip the built-in secret excludes (default `true`).
|
||||
pub ignore_default_excludes: bool,
|
||||
#[serde(deserialize_with = "deserialize_patterns")]
|
||||
pub exclude: Vec<EnvironmentVariablePattern>,
|
||||
/// Values inserted into the base environment before `include_only` filtering
|
||||
/// (an unmatched name is then dropped). These seed the base; request env
|
||||
/// layered at spawn can still override them.
|
||||
pub set: HashMap<String, String>,
|
||||
#[serde(deserialize_with = "deserialize_patterns")]
|
||||
pub include_only: Vec<EnvironmentVariablePattern>,
|
||||
}
|
||||
|
||||
impl Default for ShellEnvironmentPolicy {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
inherit: ShellEnvironmentPolicyInherit::All,
|
||||
ignore_default_excludes: true,
|
||||
exclude: Vec::new(),
|
||||
set: HashMap::new(),
|
||||
include_only: Vec::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ShellEnvironmentPolicy {
|
||||
/// True when the policy leaves the inherited environment untouched.
|
||||
pub fn is_noop(&self) -> bool {
|
||||
self.inherit == ShellEnvironmentPolicyInherit::All
|
||||
&& self.ignore_default_excludes
|
||||
&& self.exclude.is_empty()
|
||||
&& self.set.is_empty()
|
||||
&& self.include_only.is_empty()
|
||||
}
|
||||
|
||||
/// True if `name` matches a built-in secret exclude and those are enabled.
|
||||
fn matches_default_exclude(&self, name: &str) -> bool {
|
||||
!self.ignore_default_excludes && DEFAULT_SECRET_EXCLUDES.iter().any(|p| p.matches(name))
|
||||
}
|
||||
|
||||
fn matches_exclude(&self, name: &str) -> bool {
|
||||
self.exclude.iter().any(|p| p.matches(name))
|
||||
}
|
||||
|
||||
/// True if `include_only` is empty (all admitted) or `name` matches it.
|
||||
fn matches_include_only(&self, name: &str) -> bool {
|
||||
self.include_only.is_empty() || self.include_only.iter().any(|p| p.matches(name))
|
||||
}
|
||||
|
||||
/// Whether `name` survives the name filters (default excludes, `exclude`,
|
||||
/// `include_only`), ignoring `inherit`/`set`. Used to filter variables layered
|
||||
/// in after the policy base, e.g. login-shell capture. Shares its matchers
|
||||
/// with [`create_env_from_vars`] so the two cannot drift.
|
||||
pub fn allows(&self, name: &str) -> bool {
|
||||
!self.matches_default_exclude(name)
|
||||
&& !self.matches_exclude(name)
|
||||
&& self.matches_include_only(name)
|
||||
}
|
||||
|
||||
/// Like [`allows`](Self::allows) but also honors `inherit`: `none` admits
|
||||
/// nothing, `core` admits only core names, `all` defers to `allows`.
|
||||
pub fn allows_with_inherit(&self, name: &str) -> bool {
|
||||
match self.inherit {
|
||||
ShellEnvironmentPolicyInherit::None => return false,
|
||||
ShellEnvironmentPolicyInherit::Core => {
|
||||
if !CORE_ENV_VARS
|
||||
.iter()
|
||||
.any(|core| core.eq_ignore_ascii_case(name))
|
||||
{
|
||||
return false;
|
||||
}
|
||||
}
|
||||
ShellEnvironmentPolicyInherit::All => {}
|
||||
}
|
||||
self.allows(name)
|
||||
}
|
||||
}
|
||||
|
||||
/// Built-in secret excludes applied when `ignore_default_excludes` is false.
|
||||
/// Shared by the base-env build and the login-capture filter so they can't drift.
|
||||
static DEFAULT_SECRET_EXCLUDES: LazyLock<[EnvironmentVariablePattern; 3]> = LazyLock::new(|| {
|
||||
[
|
||||
EnvironmentVariablePattern::new_case_insensitive("*KEY*"),
|
||||
EnvironmentVariablePattern::new_case_insensitive("*SECRET*"),
|
||||
EnvironmentVariablePattern::new_case_insensitive("*TOKEN*"),
|
||||
]
|
||||
});
|
||||
|
||||
/// "Core" variables retained under [`ShellEnvironmentPolicyInherit::Core`].
|
||||
#[cfg(not(target_os = "windows"))]
|
||||
const CORE_ENV_VARS: &[&str] = &[
|
||||
"PATH", "SHELL", "TMPDIR", "TEMP", "TMP", "HOME", "LANG", "LC_ALL", "LC_CTYPE", "LOGNAME",
|
||||
"USER",
|
||||
];
|
||||
#[cfg(target_os = "windows")]
|
||||
const CORE_ENV_VARS: &[&str] = &[
|
||||
"PATH",
|
||||
"PATHEXT",
|
||||
"SHELL",
|
||||
"COMSPEC",
|
||||
"SYSTEMROOT",
|
||||
"SYSTEMDRIVE",
|
||||
"USERNAME",
|
||||
"USERDOMAIN",
|
||||
"USERPROFILE",
|
||||
"HOMEDRIVE",
|
||||
"HOMEPATH",
|
||||
"PROGRAMFILES",
|
||||
"PROGRAMFILES(X86)",
|
||||
"PROGRAMW6432",
|
||||
"PROGRAMDATA",
|
||||
"LOCALAPPDATA",
|
||||
"APPDATA",
|
||||
"TEMP",
|
||||
"TMP",
|
||||
"TMPDIR",
|
||||
"POWERSHELL",
|
||||
"PWSH",
|
||||
];
|
||||
|
||||
/// Build the child environment from `policy` and the process env. Uses `vars_os`
|
||||
/// and skips non-UTF-8 entries so a hostile variable cannot panic at spawn time.
|
||||
pub(crate) fn create_env(policy: &ShellEnvironmentPolicy) -> HashMap<String, String> {
|
||||
let vars = std::env::vars_os()
|
||||
.filter_map(|(k, v)| Some((k.into_string().ok()?, v.into_string().ok()?)));
|
||||
create_env_from_vars(vars, policy)
|
||||
}
|
||||
|
||||
pub(crate) fn create_env_from_vars<I>(
|
||||
vars: I,
|
||||
policy: &ShellEnvironmentPolicy,
|
||||
) -> HashMap<String, String>
|
||||
where
|
||||
I: IntoIterator<Item = (String, String)>,
|
||||
{
|
||||
let mut env: HashMap<String, String> = match policy.inherit {
|
||||
ShellEnvironmentPolicyInherit::All => vars.into_iter().collect(),
|
||||
ShellEnvironmentPolicyInherit::None => HashMap::new(),
|
||||
ShellEnvironmentPolicyInherit::Core => vars
|
||||
.into_iter()
|
||||
.filter(|(k, _)| {
|
||||
CORE_ENV_VARS
|
||||
.iter()
|
||||
.any(|allowed| allowed.eq_ignore_ascii_case(k))
|
||||
})
|
||||
.collect(),
|
||||
};
|
||||
|
||||
// Order matters: default excludes, then `exclude`, then `set`, then
|
||||
// `include_only`. `set` lands before `include_only` so an unmatched set name
|
||||
// is still dropped. The matchers are shared with `allows`.
|
||||
env.retain(|k, _| !policy.matches_default_exclude(k));
|
||||
env.retain(|k, _| !policy.matches_exclude(k));
|
||||
for (k, v) in &policy.set {
|
||||
env.insert(k.clone(), v.clone());
|
||||
}
|
||||
env.retain(|k, _| policy.matches_include_only(k));
|
||||
|
||||
// Windows resolves executables via PATHEXT; keep it present even under a
|
||||
// restrictive policy so commands stay runnable.
|
||||
if cfg!(target_os = "windows") && !env.keys().any(|k| k.eq_ignore_ascii_case("PATHEXT")) {
|
||||
env.insert("PATHEXT".to_string(), ".COM;.EXE;.BAT;.CMD".to_string());
|
||||
}
|
||||
|
||||
env
|
||||
}
|
||||
|
||||
/// Clear the command's inherited env and install the policy-derived base env.
|
||||
/// `active` must already be noop-filtered; `None` leaves the command untouched.
|
||||
/// The one base-env code path, shared by the public entry point and the spawn
|
||||
/// sites.
|
||||
pub(crate) fn install_policy_base_env(
|
||||
cmd: &mut tokio::process::Command,
|
||||
active: Option<&ShellEnvironmentPolicy>,
|
||||
) {
|
||||
if let Some(policy) = active {
|
||||
cmd.env_clear();
|
||||
cmd.envs(create_env(policy));
|
||||
}
|
||||
}
|
||||
|
||||
/// Install the policy-derived base env on `cmd` (clearing inherited env first);
|
||||
/// a `None` or no-op policy leaves it untouched. Call before any other
|
||||
/// `.env`/`.envs`.
|
||||
pub fn apply_shell_environment_policy(
|
||||
cmd: &mut tokio::process::Command,
|
||||
policy: Option<&ShellEnvironmentPolicy>,
|
||||
) {
|
||||
install_policy_base_env(cmd, policy.filter(|p| !p.is_noop()));
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "shell_env_policy_tests.rs"]
|
||||
mod tests;
|
||||
166
crates/codegen/xai-grok-tools/src/util/shell_env_policy_tests.rs
Normal file
166
crates/codegen/xai-grok-tools/src/util/shell_env_policy_tests.rs
Normal file
|
|
@ -0,0 +1,166 @@
|
|||
use super::{
|
||||
EnvironmentVariablePattern, ShellEnvironmentPolicy, ShellEnvironmentPolicyInherit,
|
||||
apply_shell_environment_policy, create_env_from_vars,
|
||||
};
|
||||
use std::collections::HashMap;
|
||||
|
||||
fn vars(pairs: &[(&str, &str)]) -> Vec<(String, String)> {
|
||||
pairs
|
||||
.iter()
|
||||
.map(|(k, v)| (k.to_string(), v.to_string()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn patterns(globs: &[&str]) -> Vec<EnvironmentVariablePattern> {
|
||||
globs
|
||||
.iter()
|
||||
.map(|g| EnvironmentVariablePattern::new_case_insensitive(g))
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn apply_policy_reshapes_command_env() {
|
||||
let mut set = HashMap::new();
|
||||
set.insert("MY_FLAG".to_string(), "1".to_string());
|
||||
let policy = ShellEnvironmentPolicy {
|
||||
inherit: ShellEnvironmentPolicyInherit::None,
|
||||
set,
|
||||
..Default::default()
|
||||
};
|
||||
let mut cmd = tokio::process::Command::new("true");
|
||||
apply_shell_environment_policy(&mut cmd, Some(&policy));
|
||||
let envs: HashMap<String, String> = cmd
|
||||
.as_std()
|
||||
.get_envs()
|
||||
.filter_map(|(k, v)| Some((k.to_str()?.to_string(), v?.to_str()?.to_string())))
|
||||
.collect();
|
||||
assert_eq!(envs.get("MY_FLAG").map(String::as_str), Some("1"));
|
||||
// inherit=None cleared the env, so no inherited PATH leaks through.
|
||||
assert!(!envs.contains_key("PATH"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn apply_noop_or_absent_policy_leaves_command_untouched() {
|
||||
let mut cmd = tokio::process::Command::new("true");
|
||||
apply_shell_environment_policy(&mut cmd, None);
|
||||
apply_shell_environment_policy(&mut cmd, Some(&ShellEnvironmentPolicy::default()));
|
||||
// No env_clear and no sets: the command carries no explicit env entries.
|
||||
assert_eq!(cmd.as_std().get_envs().count(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_excludes_drop_secrets_when_enabled() {
|
||||
let policy = ShellEnvironmentPolicy {
|
||||
ignore_default_excludes: false,
|
||||
..Default::default()
|
||||
};
|
||||
assert!(!policy.is_noop());
|
||||
let env = create_env_from_vars(
|
||||
vars(&[
|
||||
("PATH", "/bin"),
|
||||
("MY_API_KEY", "x"),
|
||||
("MY_SECRET", "y"),
|
||||
("GH_TOKEN", "z"),
|
||||
]),
|
||||
&policy,
|
||||
);
|
||||
assert_eq!(env.get("PATH").map(String::as_str), Some("/bin"));
|
||||
assert!(!env.contains_key("MY_API_KEY"));
|
||||
assert!(!env.contains_key("MY_SECRET"));
|
||||
assert!(!env.contains_key("GH_TOKEN"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn inherit_none_starts_empty_then_set_applies() {
|
||||
let mut set = HashMap::new();
|
||||
set.insert("PATH".to_string(), "/usr/bin".to_string());
|
||||
set.insert("MY_FLAG".to_string(), "1".to_string());
|
||||
let policy = ShellEnvironmentPolicy {
|
||||
inherit: ShellEnvironmentPolicyInherit::None,
|
||||
set,
|
||||
..Default::default()
|
||||
};
|
||||
let env = create_env_from_vars(vars(&[("PATH", "/bin"), ("HOME", "/root")]), &policy);
|
||||
assert_eq!(env.get("PATH").map(String::as_str), Some("/usr/bin"));
|
||||
assert_eq!(env.get("MY_FLAG").map(String::as_str), Some("1"));
|
||||
assert!(!env.contains_key("HOME"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn inherit_core_keeps_only_core_vars() {
|
||||
let policy = ShellEnvironmentPolicy {
|
||||
inherit: ShellEnvironmentPolicyInherit::Core,
|
||||
..Default::default()
|
||||
};
|
||||
let env = create_env_from_vars(vars(&[("PATH", "/bin"), ("RANDOM_VAR", "v")]), &policy);
|
||||
assert_eq!(env.get("PATH").map(String::as_str), Some("/bin"));
|
||||
assert!(!env.contains_key("RANDOM_VAR"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn exclude_and_include_only_filter() {
|
||||
let policy = ShellEnvironmentPolicy {
|
||||
exclude: patterns(&["AWS_*"]),
|
||||
include_only: patterns(&["PATH", "HOME"]),
|
||||
..Default::default()
|
||||
};
|
||||
let env = create_env_from_vars(
|
||||
vars(&[
|
||||
("PATH", "/bin"),
|
||||
("HOME", "/root"),
|
||||
("AWS_SECRET", "s"),
|
||||
("OTHER", "o"),
|
||||
]),
|
||||
&policy,
|
||||
);
|
||||
assert_eq!(env.get("PATH").map(String::as_str), Some("/bin"));
|
||||
assert_eq!(env.get("HOME").map(String::as_str), Some("/root"));
|
||||
assert!(!env.contains_key("AWS_SECRET"));
|
||||
assert!(!env.contains_key("OTHER"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn allows_filters_by_name_case_insensitively() {
|
||||
let policy = ShellEnvironmentPolicy {
|
||||
exclude: patterns(&["aws_*"]), // lowercase pattern, uppercase var
|
||||
include_only: patterns(&["PATH", "HOME"]),
|
||||
..Default::default()
|
||||
};
|
||||
assert!(policy.allows("PATH"));
|
||||
assert!(!policy.allows("AWS_SECRET")); // excluded (case-insensitive)
|
||||
assert!(!policy.allows("OTHER")); // not in include_only
|
||||
|
||||
let scrub = ShellEnvironmentPolicy {
|
||||
ignore_default_excludes: false,
|
||||
..Default::default()
|
||||
};
|
||||
assert!(!scrub.allows("my_api_key")); // `*KEY*` matches case-insensitively
|
||||
assert!(ShellEnvironmentPolicy::default().allows("MY_API_KEY")); // default allows all
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn allows_with_inherit_honors_inherit() {
|
||||
// inherit = none admits nothing.
|
||||
let none = ShellEnvironmentPolicy {
|
||||
inherit: ShellEnvironmentPolicyInherit::None,
|
||||
..Default::default()
|
||||
};
|
||||
assert!(!none.allows_with_inherit("PATH"));
|
||||
assert!(!none.allows_with_inherit("FOO"));
|
||||
|
||||
// inherit = core admits only core names.
|
||||
let core = ShellEnvironmentPolicy {
|
||||
inherit: ShellEnvironmentPolicyInherit::Core,
|
||||
..Default::default()
|
||||
};
|
||||
assert!(core.allows_with_inherit("PATH"));
|
||||
assert!(!core.allows_with_inherit("RANDOM_VAR"));
|
||||
|
||||
// inherit = all defers to `allows` (exclude still applies).
|
||||
let all = ShellEnvironmentPolicy {
|
||||
exclude: patterns(&["AWS_*"]),
|
||||
..Default::default()
|
||||
};
|
||||
assert!(all.allows_with_inherit("RANDOM_VAR"));
|
||||
assert!(!all.allows_with_inherit("AWS_SECRET"));
|
||||
}
|
||||
Loading…
Reference in a new issue