Synced from monorepo
Synced from monorepo Changes: - Shell: accept target response id on rewind execute - Shell: stamp response id on chat user message chunks - Worktree: optional rebuild and stale git registration cleanup in auto-GC - Worktree: kind-aware auto-GC TTLs and config knobs - Worktree: macOS process CWD scan and Unix PID liveness for GC guards - Worktree: automatic throttled GC on startup (Linux age-based; non-Linux dead-only) - Pager: add `[ui].combine_queued_prompts` to batch queued follow-ups - Shell: stop overwriting user skills - Tools: read markdown in `skills/` directories untruncated - `/usage` shows per-session token and dollar usage in the TUI - Security: prompt on environment-dumping `ps` variants - Security: always-safe `kubectl` no longer runs arbitrary kubeconfig credential plugins without permission - Tools: make scheduler deletion durable - Shell: add relocation storage primitives - Shell: give side model calls their own conversation ids - Fix five workflow-runtime bugs (budget, pause, cancel, reconnect) - Security: peel `env -S` / `--split-string` operands in the Bash permission gate (managed deny/ask) - Pager: expose doctor in the TUI - Security: block unauthorized RCE via abused safe commands - Pager idle watcher cue: "1 subagent still running" instead of "watching · 1 subagent" - Security: block `rg --pre` arbitrary code execution in auto-mode - Voice: diagnose silent-mic failures (macOS permission) and add doctor/terminal-setup Voice section - App builder deployer: `allow_forking` and `show_built_with_grok` - Pager: stop stacking duplicate "Worked for" markers on parked turns - Shell: support `max` as a distinct reasoning effort tier - Tools: serialize background `/loop` fires on the whole work unit - Shell: add working-directory relocation state primitives - Proto: `ClientToolResult` and `ChatConfig` client-side tools - Shell: model providers - Chat: select App Builder product on the Build path - Shell: attach author identity to feedback when the deployment opts in - Doctor: fix for SSH wrap setup - Workflow authoring skills: create-workflow and import-claude-workflow docs - Add read-only grok doctor - Sandbox: apply Landlock without a controlling TTY - Pager: recover image paste over grok wrap on headless remotes - Pager: make actions screen-mode aware - Shell: resume sessions when the working directory moves - Pager: centralize terminal diagnostics - Workspace: gate inline shell file access - Pager: centralize terminal probes - Pager: edit minimal prompts in an external editor - Pager: standardize backgrounding on Ctrl+B - Shell: recap rides the parent turn's prompt cache - Tools: add scheduler lifecycle version clock Source-Revision: 0f4d7c91b8b2b408333f6de1e8a76cb8eaa71899
This commit is contained in:
grokkybara[bot]
parent
a881e6703f
commit
3af4d5d398
556 changed files with 56609 additions and 21892 deletions
|
|
@ -241,31 +241,42 @@ fn internal_reload_request_line(id: &str, method: &str, params: serde_json::Valu
|
|||
/// Start a skills file watcher and wire it to inject `x.ai/internal/reload_skills`
|
||||
/// messages into the shared ACP incoming stream when SKILL.md files change on disk.
|
||||
///
|
||||
/// Returns the watcher guard (must be kept alive for the lifetime of the session)
|
||||
/// or `None` if no directories could be watched.
|
||||
fn spawn_skills_file_watcher<W>(
|
||||
acp_incoming_tx: &Arc<TokioMutex<W>>,
|
||||
skills_paths: &[String],
|
||||
) -> Option<crate::config::watcher::SkillsFileWatcher>
|
||||
) -> Option<tokio::task::JoinHandle<()>>
|
||||
where
|
||||
W: tokio::io::AsyncWrite + Unpin + Send + 'static,
|
||||
{
|
||||
let cwd = std::env::current_dir().unwrap_or_default();
|
||||
let workspace_user_dir = xai_grok_agent::prompt::workspace_user::optional_workspace_user_dir();
|
||||
let (watcher, mut skills_rx) = crate::config::watcher::SkillsFileWatcher::start(
|
||||
let (mut watcher, mut skills_rx) = crate::config::watcher::SkillsFileWatcher::start(
|
||||
Some(cwd.as_path()),
|
||||
workspace_user_dir.as_deref(),
|
||||
skills_paths,
|
||||
)?;
|
||||
let skills_tx = acp_incoming_tx.clone();
|
||||
tokio::spawn(async move {
|
||||
while skills_rx.recv().await.is_some() {
|
||||
info!("Skill directory changed on disk, reloading skills for all sessions");
|
||||
let line = internal_reload_request_line(
|
||||
"skills-reload",
|
||||
"x.ai/internal/reload_skills",
|
||||
serde_json::json!({}),
|
||||
);
|
||||
let task = tokio::spawn(async move {
|
||||
while let Some(change) = skills_rx.recv().await {
|
||||
let created_discovery_dir = watcher.refresh_new_discovery_dirs();
|
||||
let (id, method) = match change {
|
||||
crate::config::watcher::DiscoveryChange::Skills if !created_discovery_dir => {
|
||||
info!("Skill directory changed on disk, reloading skills for all sessions");
|
||||
("skills-reload", "x.ai/internal/reload_skills")
|
||||
}
|
||||
crate::config::watcher::DiscoveryChange::Skills => {
|
||||
info!("Discovery directory created on disk, reloading skills and workflows");
|
||||
("skills-reload", "x.ai/internal/reload_skills")
|
||||
}
|
||||
crate::config::watcher::DiscoveryChange::Workflows => {
|
||||
info!(
|
||||
"Workflow directory changed on disk, re-advertising commands for all sessions"
|
||||
);
|
||||
("workflows-reload", "x.ai/internal/reload_workflows")
|
||||
}
|
||||
};
|
||||
let line = internal_reload_request_line(id, method, serde_json::json!({}));
|
||||
let mut tx = skills_tx.lock().await;
|
||||
if let Err(e) = tx.write_all(line.as_bytes()).await {
|
||||
warn!(
|
||||
|
|
@ -275,7 +286,7 @@ where
|
|||
}
|
||||
}
|
||||
});
|
||||
Some(watcher)
|
||||
Some(task)
|
||||
}
|
||||
|
||||
/// Register the process-lifetime runtime so shared filesystem watchers
|
||||
|
|
|
|||
|
|
@ -1,4 +1,7 @@
|
|||
use crate::agent::auth_method::ModelByok;
|
||||
use crate::agent::model_providers::{
|
||||
ModelProviderConfig, auth_config_issues, model_provider_auth_name, parse_model_providers,
|
||||
};
|
||||
use crate::auth::{AuthManager, GrokComConfig, OidcAuthConfig};
|
||||
use crate::remote::DEFAULT_CONTEXT_WINDOW;
|
||||
use crate::{config::StorageMode, sampling::ApiBackend, tools::config::ShellToolsetConfig};
|
||||
|
|
@ -947,7 +950,30 @@ impl PluginsConfig {
|
|||
/// Feedback submission configuration (`[feedback]` in config.toml).
|
||||
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
|
||||
#[serde(default)]
|
||||
pub struct FeedbackConfig {}
|
||||
pub struct FeedbackConfig {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub user: Option<FeedbackUserConfig>,
|
||||
}
|
||||
/// Self-reported feedback author identity (never used for authorization).
|
||||
/// Merged only from trusted config tiers, so a cloned repo can't inject the
|
||||
/// `command` escape hatch.
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(default)]
|
||||
pub struct FeedbackUserConfig {
|
||||
/// Sources tried in order for the name. `os_user` yields the OS user name;
|
||||
/// any other entry is a literal (`$VAR` expanded at load).
|
||||
pub name: Vec<String>,
|
||||
/// Sources tried in order for the email. `git_email` yields the global git
|
||||
/// email; any other entry is a literal (`$VAR` expanded at load) needing `@`.
|
||||
pub email: Vec<String>,
|
||||
/// Fallback domain for `<name>@<domain>` when no `email` source resolves.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub email_domain: Option<String>,
|
||||
/// Optional `sh -c` script printing `{"name","email"}` JSON; its fields win
|
||||
/// over the lists above, with per-field fallback. Trusted config tiers only.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub command: Option<String>,
|
||||
}
|
||||
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
|
||||
#[serde(default)]
|
||||
pub struct CompactionConfig {
|
||||
|
|
@ -1121,6 +1147,13 @@ pub struct WorktreePoolConfig {
|
|||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub parallelism: Option<usize>,
|
||||
}
|
||||
/// `[worktree]` section from config.toml (auto-GC policy lives under `auto_gc`).
|
||||
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
|
||||
#[serde(default)]
|
||||
pub struct WorktreeConfigSection {
|
||||
#[serde(default)]
|
||||
pub auto_gc: crate::util::config::WorktreeAutoGcSettings,
|
||||
}
|
||||
/// `[sandbox]` section from config.toml.
|
||||
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
|
||||
#[serde(default)]
|
||||
|
|
@ -1288,19 +1321,22 @@ pub struct Config {
|
|||
/// `[goal]` section: canonical `/goal` configuration. See [`GoalConfig`].
|
||||
#[serde(default)]
|
||||
pub goal: GoalConfig,
|
||||
#[serde(default)]
|
||||
pub workflows: WorkflowsConfig,
|
||||
/// `[doom_loop_recovery]` section: the shared settings struct — ONE type
|
||||
/// serves this TOML table and the remote remote settings `doom_loop_recovery`
|
||||
/// object. See [`crate::util::config::DoomLoopRecoverySettings`].
|
||||
#[serde(default)]
|
||||
pub doom_loop_recovery: crate::util::config::DoomLoopRecoverySettings,
|
||||
/// `[worktree]` section (currently `[worktree.auto_gc]` only).
|
||||
#[serde(default)]
|
||||
pub worktree: WorktreeConfigSection,
|
||||
/// `[auto_mode]` section: Auto permission-mode configuration. See [`AutoModeConfig`].
|
||||
#[serde(default)]
|
||||
pub auto_mode: AutoModeConfig,
|
||||
/// `[model.*]` overrides from config.toml. Resolve via `resolve_model_list()`.
|
||||
#[serde(skip)]
|
||||
pub config_models: IndexMap<String, ConfigModelOverride>,
|
||||
/// Warnings from `[model.*]` and `[auth_provider.*]` parsing; surfaced by
|
||||
/// `grok inspect`.
|
||||
#[serde(skip)]
|
||||
pub config_warnings: Vec<super::config_model_override_parse::ConfigWarning>,
|
||||
pub grok_com_config: GrokComConfig,
|
||||
|
|
@ -1308,6 +1344,8 @@ pub struct Config {
|
|||
/// [`parse_auth_providers`] from trusted config layers only.
|
||||
#[serde(skip)]
|
||||
pub auth_providers: IndexMap<String, crate::auth::AuthProviderConfig>,
|
||||
#[serde(skip)]
|
||||
pub model_providers: IndexMap<String, ModelProviderConfig>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub shortcuts: Option<toml::Value>,
|
||||
/// Written by the client via `config_toml_edit`; absorbed so it isn't
|
||||
|
|
@ -1729,12 +1767,15 @@ impl Default for Config {
|
|||
let mut cfg = Self {
|
||||
features: Features::default(),
|
||||
goal: GoalConfig::default(),
|
||||
workflows: WorkflowsConfig::default(),
|
||||
doom_loop_recovery: crate::util::config::DoomLoopRecoverySettings::default(),
|
||||
worktree: WorktreeConfigSection::default(),
|
||||
auto_mode: AutoModeConfig::default(),
|
||||
config_models: IndexMap::new(),
|
||||
config_warnings: Vec::new(),
|
||||
grok_com_config: GrokComConfig::default(),
|
||||
auth_providers: IndexMap::new(),
|
||||
model_providers: IndexMap::new(),
|
||||
shortcuts: None,
|
||||
hints: None,
|
||||
ui: UiConfig::default(),
|
||||
|
|
@ -1858,41 +1899,12 @@ fn parse_auth_providers(
|
|||
"unrecognized key; field ignored".to_owned(),
|
||||
));
|
||||
}
|
||||
if !provider.is_usable() {
|
||||
for (field, kind, reason) in auth_config_issues(&provider) {
|
||||
warnings.push(ConfigWarning::auth_provider(
|
||||
name,
|
||||
Some("command"),
|
||||
ConfigWarningKind::InvalidValue,
|
||||
"missing or empty command; referencing models resolve \
|
||||
with no credential"
|
||||
.to_owned(),
|
||||
));
|
||||
}
|
||||
let skew = crate::auth::PROVIDER_TOKEN_EXPIRY_SKEW_SECS;
|
||||
if provider.token_ttl_secs.is_some_and(|ttl| ttl <= skew) {
|
||||
warnings.push(ConfigWarning::auth_provider(
|
||||
name,
|
||||
Some("token_ttl_secs"),
|
||||
ConfigWarningKind::InvalidValue,
|
||||
format!(
|
||||
"at or below the {skew}s refresh margin; the command will \
|
||||
run before every turn"
|
||||
),
|
||||
));
|
||||
}
|
||||
if let Some(timeout) = provider.timeout_secs
|
||||
&& !(1..=crate::auth::PROVIDER_TIMEOUT_CEILING_SECS).contains(&timeout)
|
||||
{
|
||||
let ceiling = crate::auth::PROVIDER_TIMEOUT_CEILING_SECS;
|
||||
warnings.push(ConfigWarning::auth_provider(
|
||||
name,
|
||||
Some("timeout_secs"),
|
||||
ConfigWarningKind::InvalidValue,
|
||||
if timeout == 0 {
|
||||
"below the 1 second minimum; clamped to 1".to_owned()
|
||||
} else {
|
||||
format!("above the {ceiling}s maximum; clamped to {ceiling}")
|
||||
},
|
||||
Some(field),
|
||||
kind,
|
||||
reason,
|
||||
));
|
||||
}
|
||||
providers.insert(name.clone(), provider);
|
||||
|
|
@ -1969,7 +1981,29 @@ impl Config {
|
|||
models: config_models,
|
||||
warnings: config_warnings,
|
||||
} = super::config_model_override_parse::parse_model_overrides(raw_config);
|
||||
let (auth_providers, auth_provider_warnings) = parse_auth_providers(raw_config);
|
||||
let (mut auth_providers, auth_provider_warnings) = parse_auth_providers(raw_config);
|
||||
let (model_providers, mut model_provider_warnings) = parse_model_providers(raw_config);
|
||||
for (id, provider) in &model_providers {
|
||||
if let Some(auth) = &provider.auth {
|
||||
let synthetic = model_provider_auth_name(id);
|
||||
if auth_providers.contains_key(&synthetic) {
|
||||
model_provider_warnings
|
||||
.push(
|
||||
super::config_model_override_parse::ConfigWarning::model_provider(
|
||||
id,
|
||||
Some("auth"),
|
||||
super::config_model_override_parse::ConfigWarningKind::ConflictingFields,
|
||||
format!(
|
||||
"inline auth overwrites a hand-written \
|
||||
[auth_provider.\"{synthetic}\"]; the `model_provider:` prefix is \
|
||||
a reserved namespace"
|
||||
),
|
||||
),
|
||||
);
|
||||
}
|
||||
auth_providers.insert(synthetic, auth.clone());
|
||||
}
|
||||
}
|
||||
let mut base = toml::Value::try_from(Self::default()).map_err(|e| e.to_string())?;
|
||||
if let toml::Value::Table(ref mut t) = base {
|
||||
t.remove("model");
|
||||
|
|
@ -1978,6 +2012,7 @@ impl Config {
|
|||
if let toml::Value::Table(ref mut t) = raw_without_model_sections {
|
||||
t.remove("model");
|
||||
t.remove("auth_provider");
|
||||
t.remove("model_providers");
|
||||
}
|
||||
crate::config::deep_merge_toml(&mut base, &raw_without_model_sections);
|
||||
let (mut config, user_unused) =
|
||||
|
|
@ -1991,12 +2026,19 @@ impl Config {
|
|||
config.config_models = config_models;
|
||||
config.config_warnings = config_warnings;
|
||||
config.auth_providers = auth_providers;
|
||||
config.model_providers = model_providers;
|
||||
config.config_warnings.extend(auth_provider_warnings);
|
||||
config.config_warnings.extend(model_provider_warnings);
|
||||
let declared_provider_names: std::collections::HashSet<&str> = raw_config
|
||||
.get("auth_provider")
|
||||
.and_then(toml::Value::as_table)
|
||||
.map(|t| t.keys().map(String::as_str).collect())
|
||||
.unwrap_or_default();
|
||||
let declared_model_provider_names: std::collections::HashSet<&str> = raw_config
|
||||
.get("model_providers")
|
||||
.and_then(toml::Value::as_table)
|
||||
.map(|t| t.keys().map(String::as_str).collect())
|
||||
.unwrap_or_default();
|
||||
for (model_key, model) in &config.config_models {
|
||||
if let Some(ref name) = model.auth_provider
|
||||
&& !config.auth_providers.contains_key(name)
|
||||
|
|
@ -2014,6 +2056,41 @@ impl Config {
|
|||
),
|
||||
);
|
||||
}
|
||||
if let Some(ref id) = model.model_provider
|
||||
&& !config.model_providers.contains_key(id)
|
||||
&& !declared_model_provider_names.contains(id.as_str())
|
||||
{
|
||||
config.config_warnings.push(
|
||||
super::config_model_override_parse::ConfigWarning::model(
|
||||
model_key,
|
||||
Some("model_provider"),
|
||||
super::config_model_override_parse::ConfigWarningKind::InvalidValue,
|
||||
format!(
|
||||
"references [model_providers.{id}], which is not defined; \
|
||||
provider defaults are not applied — the model uses its own \
|
||||
credential if set, otherwise fails closed on a custom endpoint"
|
||||
),
|
||||
),
|
||||
);
|
||||
}
|
||||
}
|
||||
for (id, provider) in &config.model_providers {
|
||||
if let Some(ref name) = provider.auth_provider
|
||||
&& !config.auth_providers.contains_key(name)
|
||||
&& !declared_provider_names.contains(name.as_str())
|
||||
{
|
||||
config.config_warnings.push(
|
||||
super::config_model_override_parse::ConfigWarning::model_provider(
|
||||
id,
|
||||
Some("auth_provider"),
|
||||
super::config_model_override_parse::ConfigWarningKind::InvalidValue,
|
||||
format!(
|
||||
"references [auth_provider.{name}], which is not defined; \
|
||||
inheriting models fail closed with no provider credential"
|
||||
),
|
||||
),
|
||||
);
|
||||
}
|
||||
}
|
||||
super::config_model_override_parse::log_config_warnings(&config.config_warnings);
|
||||
if config.grok_com_config.oidc.is_none() {
|
||||
|
|
@ -2345,6 +2422,18 @@ impl Config {
|
|||
.map_or(Policy::DEFAULT_MAX_RETRIES, Policy::clamp_max_retries),
|
||||
})
|
||||
}
|
||||
/// Automatic worktree GC policy. Precedence: env kill/dry-run >
|
||||
/// `[worktree.auto_gc]` TOML > remote `worktree_auto_gc` > defaults.
|
||||
/// Platform age-expiry (non-Linux dead-only) is enforced inside
|
||||
/// `xai_fast_worktree::maybe_auto_gc`, not here.
|
||||
pub fn resolve_worktree_auto_gc(&self) -> xai_fast_worktree::ResolvedWorktreeAutoGc {
|
||||
crate::util::config::resolve_worktree_auto_gc_from_settings(
|
||||
Some(&self.worktree.auto_gc),
|
||||
self.remote_settings
|
||||
.as_ref()
|
||||
.and_then(|s| s.worktree_auto_gc.as_ref()),
|
||||
)
|
||||
}
|
||||
/// Gate first-run auto-registration of the official xAI marketplace source.
|
||||
/// Precedence: env `GROK_OFFICIAL_MARKETPLACE_AUTO_REGISTER` > remote settings >
|
||||
/// default off (so only remote settings-targeted teams get it pre-public). No
|
||||
|
|
@ -2474,16 +2563,31 @@ impl Config {
|
|||
/// reach cli-chat-proxy `/v1/settings` (custom `models_base_url`, external
|
||||
/// `auth_provider_command`, air-gapped proxies) never receive the
|
||||
/// remote settings `goal_enabled` flag, so the default must not carve them out.
|
||||
/// Env, `[goal] enabled`, and the remote flag (`Some(false)` kill-switch)
|
||||
/// all still override.
|
||||
pub(crate) fn resolve_goal(&self) -> Resolved<bool> {
|
||||
let ff = self.remote_settings.as_ref().and_then(|s| s.goal_enabled);
|
||||
if ff == Some(false) {
|
||||
return Resolved::new(false, ConfigSource::Remote);
|
||||
}
|
||||
BoolFlag::env("GROK_GOAL")
|
||||
.config(self.goal.enabled)
|
||||
.feature_flag(ff)
|
||||
.default(true)
|
||||
.resolve()
|
||||
}
|
||||
pub(crate) fn resolve_workflows(&self) -> Resolved<bool> {
|
||||
let ff = self
|
||||
.remote_settings
|
||||
.as_ref()
|
||||
.and_then(|s| s.workflows_enabled);
|
||||
if ff == Some(false) {
|
||||
return Resolved::new(false, ConfigSource::Remote);
|
||||
}
|
||||
BoolFlag::env("GROK_WORKFLOWS")
|
||||
.config(self.workflows.enabled)
|
||||
.feature_flag(ff)
|
||||
.default(false)
|
||||
.resolve()
|
||||
}
|
||||
/// Classifier, planner, and summary all default to goal mode itself: when
|
||||
/// `/goal` is on they are on unless config/env/remote says otherwise.
|
||||
/// `goal_enabled` is the session's already-resolved master switch (the same
|
||||
|
|
@ -3340,22 +3444,46 @@ pub fn resolve_model_list(
|
|||
);
|
||||
}
|
||||
}
|
||||
let entry = model_override.apply(key, base, &cfg.endpoints);
|
||||
let with_provider = model_override.model_provider.as_deref().map(|pid| {
|
||||
match cfg.model_providers.get(pid) {
|
||||
Some(provider) => model_override.with_provider_defaults(provider, pid),
|
||||
None => model_override.with_missing_provider(),
|
||||
}
|
||||
});
|
||||
let effective = with_provider.as_ref().unwrap_or(model_override);
|
||||
let mut entry = effective.apply(key, base, &cfg.endpoints);
|
||||
let session_bearer_unsafe = !crate::util::is_xai_api_bearer_url(&entry.info.base_url)
|
||||
|| entry
|
||||
.api_base_url
|
||||
.as_deref()
|
||||
.is_some_and(|url| !crate::util::is_xai_api_bearer_url(url));
|
||||
if let Some(pid) = model_override.model_provider.as_deref()
|
||||
&& entry.auth_provider.is_none()
|
||||
&& session_bearer_unsafe
|
||||
{
|
||||
entry.auth_provider = Some(crate::auth::AuthProviderRef::fail_closed(format!(
|
||||
"model_provider:{pid} (fail-closed)"
|
||||
)));
|
||||
}
|
||||
tracing::debug!(
|
||||
model_key = % key, base_url = % entry.info.base_url, has_api_key = entry
|
||||
.api_key.is_some(), env_key = ? entry.env_key, auth_provider = entry
|
||||
.auth_provider.as_ref().map(| p | p.name.as_str()), had_base,
|
||||
.auth_provider.as_ref().map(| p | p.name.as_str()), model_provider =
|
||||
model_override.model_provider.as_deref(), had_base,
|
||||
"config model override applied"
|
||||
);
|
||||
resolved.insert(key.clone(), entry);
|
||||
}
|
||||
for (key, entry) in resolved.iter_mut() {
|
||||
if let Some(ref mut provider) = entry.auth_provider {
|
||||
if provider.is_fail_closed() {
|
||||
continue;
|
||||
}
|
||||
let config = cfg.auth_providers.get(&provider.name);
|
||||
if config.is_none() {
|
||||
tracing::debug!(
|
||||
model_key = % key, provider = % provider.name,
|
||||
"model references an undefined [auth_provider.*] table"
|
||||
"provider ref has no trusted config; failing closed with an empty command"
|
||||
);
|
||||
}
|
||||
provider.attach_trusted_config(config);
|
||||
|
|
@ -3746,6 +3874,7 @@ pub struct ConfigModelOverride {
|
|||
/// this model's bearer token. Static `api_key` / `env_key` win when both
|
||||
/// are set.
|
||||
pub auth_provider: Option<String>,
|
||||
pub model_provider: Option<String>,
|
||||
pub api_base_url: Option<String>,
|
||||
pub max_completion_tokens: Option<u32>,
|
||||
pub temperature: Option<f32>,
|
||||
|
|
@ -4270,6 +4399,12 @@ pub struct GoalConfig {
|
|||
)]
|
||||
pub skeptic_models: Vec<crate::util::config::GoalRoleModel>,
|
||||
}
|
||||
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
|
||||
#[serde(default)]
|
||||
pub struct WorkflowsConfig {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub enabled: Option<bool>,
|
||||
}
|
||||
/// `[auto_mode]` section: server-side configuration for Auto permission mode.
|
||||
/// ONE struct serves both the local `[auto_mode]` TOML table and the remote
|
||||
/// remote settings `auto_mode` JSON object (coerced via `serde_json::from_value`), so
|
||||
|
|
@ -5789,8 +5924,7 @@ if n == name && f.as_deref() == field
|
|||
cfg.config_warnings.iter().any(|w| {
|
||||
w.kind == ConfigWarningKind::InvalidValue
|
||||
&& matches!(& w.target, WarningTarget::Model
|
||||
{ field, .. }
|
||||
if field.as_deref() == Some("auth_provider"))
|
||||
{ field, .. } if field.as_deref() == Some("auth_provider"))
|
||||
}),
|
||||
"undefined reference warns at parse time: {:?}",
|
||||
cfg.config_warnings
|
||||
|
|
@ -6050,6 +6184,32 @@ if field.as_deref() == Some("auth_provider"))
|
|||
let provider = resolved["cached-model"].auth_provider.as_ref().unwrap();
|
||||
assert_eq!(provider.config.command, "printf local");
|
||||
}
|
||||
#[test]
|
||||
fn provider_model_fails_closed_on_prefetched_custom_base_url() {
|
||||
let mut cfg = Config::default();
|
||||
cfg.model_providers.insert(
|
||||
"gw".to_string(),
|
||||
crate::agent::model_providers::ModelProviderConfig::default(),
|
||||
);
|
||||
cfg.config_models.insert(
|
||||
"m".to_string(),
|
||||
ConfigModelOverride {
|
||||
model_provider: Some("gw".to_string()),
|
||||
..Default::default()
|
||||
},
|
||||
);
|
||||
let mut prefetched = IndexMap::new();
|
||||
prefetched.insert(
|
||||
"m".to_string(),
|
||||
test_model_entry("m", "https://evil.example/v1", None, None, None),
|
||||
);
|
||||
let resolved = resolve_model_list(&cfg, Some(prefetched));
|
||||
assert_eq!(
|
||||
resolve_credentials(&resolved["m"], Some("session-jwt")).api_key,
|
||||
None,
|
||||
"a prefetched custom base_url must fail closed, not leak the session token",
|
||||
);
|
||||
}
|
||||
fn test_model_entry(
|
||||
model: &str,
|
||||
base_url: &str,
|
||||
|
|
@ -8710,6 +8870,47 @@ if field.as_deref() == Some("auth_provider"))
|
|||
assert_eq!(p.max_threshold, 12);
|
||||
assert_eq!(p.max_retries, 1);
|
||||
}
|
||||
/// `[worktree.auto_gc]` deserializes through Config and resolve honors it.
|
||||
#[test]
|
||||
#[serial]
|
||||
fn worktree_auto_gc_section_parses_from_toml() {
|
||||
unsafe {
|
||||
std::env::remove_var(xai_fast_worktree::ENV_AUTO_GC);
|
||||
std::env::remove_var(xai_fast_worktree::ENV_AUTO_GC_DRY_RUN);
|
||||
std::env::remove_var(xai_fast_worktree::ENV_AUTO_GC_MAX_AGE);
|
||||
}
|
||||
let raw: toml::Value = toml::from_str(
|
||||
r#"
|
||||
[worktree.auto_gc]
|
||||
enabled = true
|
||||
max_age_secs = 7200
|
||||
min_interval_secs = 120
|
||||
dry_run = true
|
||||
[worktree.auto_gc.max_age_by_kind]
|
||||
subagent = 3600
|
||||
manual = "never"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
let cfg = Config::new_from_toml_cfg(&raw).unwrap();
|
||||
assert_eq!(cfg.worktree.auto_gc.enabled, Some(true));
|
||||
assert_eq!(cfg.worktree.auto_gc.max_age_secs, Some(7200));
|
||||
let p = cfg.resolve_worktree_auto_gc();
|
||||
assert!(p.enabled);
|
||||
assert_eq!(p.max_age_secs, 7200);
|
||||
assert_eq!(p.min_interval_secs, 120);
|
||||
assert!(p.dry_run);
|
||||
assert_eq!(
|
||||
p.max_age_by_kind
|
||||
.get(&xai_fast_worktree::WorktreeKind::Subagent),
|
||||
Some(&Some(3600))
|
||||
);
|
||||
assert_eq!(
|
||||
p.max_age_by_kind
|
||||
.get(&xai_fast_worktree::WorktreeKind::Manual),
|
||||
Some(&None)
|
||||
);
|
||||
}
|
||||
/// Out-of-range tunables clamp instead of being honored or dropped.
|
||||
#[test]
|
||||
#[serial]
|
||||
|
|
@ -8867,14 +9068,10 @@ if field.as_deref() == Some("auth_provider"))
|
|||
}
|
||||
#[test]
|
||||
#[serial]
|
||||
fn resolve_goal_env_overrides_config() {
|
||||
fn resolve_goal_env_overrides_config_without_remote_kill_switch() {
|
||||
unsafe { std::env::set_var("GROK_GOAL", "1") };
|
||||
let mut cfg = Config::default();
|
||||
cfg.goal.enabled = Some(false);
|
||||
cfg.remote_settings = Some(crate::util::config::RemoteSettings {
|
||||
goal_enabled: Some(false),
|
||||
..Default::default()
|
||||
});
|
||||
let r = cfg.resolve_goal();
|
||||
assert_eq!(r.source, ConfigSource::Env);
|
||||
assert!(r.value);
|
||||
|
|
@ -8882,8 +9079,8 @@ if field.as_deref() == Some("auth_provider"))
|
|||
}
|
||||
#[test]
|
||||
#[serial]
|
||||
fn resolve_goal_config_overrides_remote_settings() {
|
||||
unsafe { std::env::remove_var("GROK_GOAL") };
|
||||
fn resolve_goal_remote_false_kills_local_opt_in() {
|
||||
unsafe { std::env::set_var("GROK_GOAL", "1") };
|
||||
let mut cfg = Config::default();
|
||||
cfg.goal.enabled = Some(true);
|
||||
cfg.remote_settings = Some(crate::util::config::RemoteSettings {
|
||||
|
|
@ -8891,8 +9088,9 @@ if field.as_deref() == Some("auth_provider"))
|
|||
..Default::default()
|
||||
});
|
||||
let r = cfg.resolve_goal();
|
||||
assert_eq!(r.source, ConfigSource::Config);
|
||||
assert!(r.value);
|
||||
assert_eq!(r.source, ConfigSource::Remote);
|
||||
assert!(!r.value);
|
||||
unsafe { std::env::remove_var("GROK_GOAL") };
|
||||
}
|
||||
#[test]
|
||||
#[serial]
|
||||
|
|
@ -8928,6 +9126,54 @@ if field.as_deref() == Some("auth_provider"))
|
|||
}
|
||||
#[test]
|
||||
#[serial]
|
||||
fn background_workflows_default_off_without_affecting_goal() {
|
||||
unsafe { std::env::remove_var("GROK_WORKFLOWS") };
|
||||
let cfg = Config::default();
|
||||
assert!(!cfg.resolve_workflows().value);
|
||||
assert!(cfg.resolve_goal().value);
|
||||
}
|
||||
#[test]
|
||||
#[serial]
|
||||
fn resolve_workflows_remote_settings_opt_in() {
|
||||
unsafe { std::env::remove_var("GROK_WORKFLOWS") };
|
||||
let cfg = Config {
|
||||
remote_settings: Some(crate::util::config::RemoteSettings {
|
||||
workflows_enabled: Some(true),
|
||||
..Default::default()
|
||||
}),
|
||||
..Default::default()
|
||||
};
|
||||
let r = cfg.resolve_workflows();
|
||||
assert_eq!(r.source, ConfigSource::Remote);
|
||||
assert!(r.value);
|
||||
}
|
||||
#[test]
|
||||
#[serial]
|
||||
fn resolve_workflows_remote_false_kills_local_opt_in() {
|
||||
unsafe { std::env::set_var("GROK_WORKFLOWS", "1") };
|
||||
let mut cfg = Config::default();
|
||||
cfg.workflows.enabled = Some(true);
|
||||
cfg.remote_settings = Some(crate::util::config::RemoteSettings {
|
||||
workflows_enabled: Some(false),
|
||||
..Default::default()
|
||||
});
|
||||
let r = cfg.resolve_workflows();
|
||||
assert_eq!(r.source, ConfigSource::Remote);
|
||||
assert!(!r.value);
|
||||
unsafe { std::env::remove_var("GROK_WORKFLOWS") };
|
||||
}
|
||||
#[test]
|
||||
#[serial]
|
||||
fn resolve_workflows_env_wins() {
|
||||
unsafe { std::env::set_var("GROK_WORKFLOWS", "1") };
|
||||
let cfg = Config::default();
|
||||
let r = cfg.resolve_workflows();
|
||||
assert_eq!(r.source, ConfigSource::Env);
|
||||
assert!(r.value);
|
||||
unsafe { std::env::remove_var("GROK_WORKFLOWS") };
|
||||
}
|
||||
#[test]
|
||||
#[serial]
|
||||
fn resolve_ask_user_question_defaults_to_true_when_unset() {
|
||||
unsafe { std::env::remove_var("GROK_ASK_USER_QUESTION") };
|
||||
let cfg = Config::default();
|
||||
|
|
@ -9882,6 +10128,11 @@ agent_type = "cursor"
|
|||
enabled = true
|
||||
[harness]
|
||||
block_for_upload = true
|
||||
[feedback.user]
|
||||
name = ["os_user"]
|
||||
email = ["git_email", "team@example.com"]
|
||||
email_domain = "example.com"
|
||||
command = "/opt/bin/grok-identity"
|
||||
[repo_changes_dedup]
|
||||
enabled = false
|
||||
[relay]
|
||||
|
|
|
|||
|
|
@ -63,6 +63,12 @@ pub enum WarningTarget {
|
|||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
field: Option<String>,
|
||||
},
|
||||
ModelProviderSection,
|
||||
ModelProvider {
|
||||
id: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
field: Option<String>,
|
||||
},
|
||||
}
|
||||
|
||||
impl WarningTarget {
|
||||
|
|
@ -73,13 +79,17 @@ impl WarningTarget {
|
|||
Self::Model { key, .. } => format!("model.\"{key}\""),
|
||||
Self::AuthProviderSection => "auth_provider".to_owned(),
|
||||
Self::AuthProvider { name, .. } => format!("auth_provider.\"{name}\""),
|
||||
Self::ModelProviderSection => "model_providers".to_owned(),
|
||||
Self::ModelProvider { id, .. } => format!("model_providers.\"{id}\""),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn field(&self) -> Option<&str> {
|
||||
match self {
|
||||
Self::Model { field, .. } | Self::AuthProvider { field, .. } => field.as_deref(),
|
||||
Self::ModelSection | Self::AuthProviderSection => None,
|
||||
Self::Model { field, .. }
|
||||
| Self::AuthProvider { field, .. }
|
||||
| Self::ModelProvider { field, .. } => field.as_deref(),
|
||||
Self::ModelSection | Self::AuthProviderSection | Self::ModelProviderSection => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -145,6 +155,30 @@ impl ConfigWarning {
|
|||
}
|
||||
}
|
||||
|
||||
pub(crate) fn model_provider(
|
||||
id: &str,
|
||||
field: Option<&str>,
|
||||
kind: ConfigWarningKind,
|
||||
reason: String,
|
||||
) -> Self {
|
||||
Self {
|
||||
target: WarningTarget::ModelProvider {
|
||||
id: id.to_owned(),
|
||||
field: field.map(str::to_owned),
|
||||
},
|
||||
kind,
|
||||
reason,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn model_provider_section(kind: ConfigWarningKind, reason: String) -> Self {
|
||||
Self {
|
||||
target: WarningTarget::ModelProviderSection,
|
||||
kind,
|
||||
reason,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn field(&self) -> Option<&str> {
|
||||
self.target.field()
|
||||
}
|
||||
|
|
@ -623,10 +657,10 @@ mod tests {
|
|||
assert_eq!(warnings.len(), 1);
|
||||
assert_eq!(warnings[0].kind, ConfigWarningKind::NotATable);
|
||||
assert!(matches!(
|
||||
&warnings[0].target,
|
||||
WarningTarget::Model { key, field: None }
|
||||
if key == "oops"
|
||||
));
|
||||
&warnings[0].target,
|
||||
WarningTarget::Model { key, field: None }
|
||||
if key == "oops"
|
||||
));
|
||||
}
|
||||
|
||||
/// Exhaustive literal (no `..`): a new struct field is a compile error
|
||||
|
|
@ -640,6 +674,7 @@ mod tests {
|
|||
api_key: Some("key".into()),
|
||||
env_key: Some(crate::agent::config::EnvKeys::single("ENV_KEY")),
|
||||
auth_provider: Some("corp-gateway".into()),
|
||||
model_provider: Some("gateway".into()),
|
||||
api_base_url: Some("https://api.example.com".into()),
|
||||
max_completion_tokens: Some(1024),
|
||||
temperature: Some(0.5),
|
||||
|
|
|
|||
|
|
@ -80,6 +80,14 @@ pub(super) fn parse_queue_edit_command(
|
|||
editor: owner,
|
||||
})
|
||||
}
|
||||
"x.ai/queue/hold_edit" => {
|
||||
let id = params.get("id").and_then(|v| v.as_str())?.to_string();
|
||||
Some(SessionCommand::HoldCombineEdit { id })
|
||||
}
|
||||
"x.ai/queue/release_edit" => {
|
||||
let id = params.get("id").and_then(|v| v.as_str())?.to_string();
|
||||
Some(SessionCommand::ReleaseCombineEdit { id })
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -121,7 +121,7 @@ fn resolve_config(cfg: &AgentConfig, auth_manager: &AuthManager) -> AgentConfig
|
|||
cfg
|
||||
}
|
||||
|
||||
/// Initialize process-level singletons (deployment sync, bundled files,
|
||||
/// Initialize process-level singletons (deployment sync, built-in metadata,
|
||||
/// telemetry). `Once`-guarded: only the first call takes effect.
|
||||
/// Telemetry user ID is updated separately via [`update_telemetry_config`].
|
||||
fn init_process(cfg: &AgentConfig, auth_manager: &AuthManager) {
|
||||
|
|
@ -135,12 +135,12 @@ fn init_process(cfg: &AgentConfig, auth_manager: &AuthManager) {
|
|||
}
|
||||
|
||||
let grok_home = crate::util::grok_home::grok_home();
|
||||
crate::builtin::extract_bundled_files(&grok_home);
|
||||
crate::builtin::extract_builtin_files(&grok_home);
|
||||
|
||||
crate::extensions::marketplace::purge_default_skills_installs(&grok_home);
|
||||
|
||||
// Auto-register is gated (default off; env/remote settings enables). Kept out
|
||||
// of extract_bundled_files so the gate can read the resolved
|
||||
// of built-in extraction so the gate can read the resolved
|
||||
// remote_settings, which resolve_config has populated by now.
|
||||
if cfg.resolve_official_marketplace_auto_register().value {
|
||||
crate::extensions::marketplace::ensure_official_marketplace_source(&grok_home);
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ pub mod feedback_client;
|
|||
pub mod folder_trust;
|
||||
pub(crate) mod handlers;
|
||||
pub mod init;
|
||||
pub mod model_providers;
|
||||
pub mod models;
|
||||
pub mod mvp_agent;
|
||||
pub(crate) mod proxy;
|
||||
|
|
|
|||
917
crates/codegen/xai-grok-shell/src/agent/model_providers.rs
Normal file
917
crates/codegen/xai-grok-shell/src/agent/model_providers.rs
Normal file
|
|
@ -0,0 +1,917 @@
|
|||
use indexmap::IndexMap;
|
||||
|
||||
use super::config::{ConfigModelOverride, EnvKeys};
|
||||
use super::config_model_override_parse::{ConfigWarning, ConfigWarningKind};
|
||||
use crate::sampling::ApiBackend;
|
||||
|
||||
#[derive(Clone, Debug, Default, serde::Deserialize)]
|
||||
#[serde(default)]
|
||||
pub struct ModelProviderConfig {
|
||||
pub base_url: Option<String>,
|
||||
pub api_base_url: Option<String>,
|
||||
pub env_key: Option<EnvKeys>,
|
||||
pub api_key: Option<String>,
|
||||
pub api_backend: Option<ApiBackend>,
|
||||
pub extra_headers: IndexMap<String, String>,
|
||||
pub auth_provider: Option<String>,
|
||||
pub auth: Option<crate::auth::AuthProviderConfig>,
|
||||
pub context_window: Option<u64>,
|
||||
}
|
||||
|
||||
pub(crate) fn model_provider_auth_name(provider_id: &str) -> String {
|
||||
format!("model_provider:{provider_id}")
|
||||
}
|
||||
|
||||
pub(crate) fn auth_config_issues(
|
||||
config: &crate::auth::AuthProviderConfig,
|
||||
) -> Vec<(&'static str, ConfigWarningKind, String)> {
|
||||
let mut issues = Vec::new();
|
||||
if !config.is_usable() {
|
||||
issues.push((
|
||||
"command",
|
||||
ConfigWarningKind::InvalidValue,
|
||||
"missing or empty command; models resolve with no credential".to_owned(),
|
||||
));
|
||||
}
|
||||
let skew = crate::auth::PROVIDER_TOKEN_EXPIRY_SKEW_SECS;
|
||||
if config.token_ttl_secs.is_some_and(|ttl| ttl <= skew) {
|
||||
issues.push((
|
||||
"token_ttl_secs",
|
||||
ConfigWarningKind::InvalidValue,
|
||||
format!(
|
||||
"at or below the {skew}s refresh margin; the command will run before every turn"
|
||||
),
|
||||
));
|
||||
}
|
||||
if let Some(timeout) = config.timeout_secs
|
||||
&& !(1..=crate::auth::PROVIDER_TIMEOUT_CEILING_SECS).contains(&timeout)
|
||||
{
|
||||
let ceiling = crate::auth::PROVIDER_TIMEOUT_CEILING_SECS;
|
||||
issues.push((
|
||||
"timeout_secs",
|
||||
ConfigWarningKind::InvalidValue,
|
||||
if timeout == 0 {
|
||||
"below the 1 second minimum; clamped to 1".to_owned()
|
||||
} else {
|
||||
format!("above the {ceiling}s maximum; clamped to {ceiling}")
|
||||
},
|
||||
));
|
||||
}
|
||||
issues
|
||||
}
|
||||
|
||||
pub(crate) fn parse_model_providers(
|
||||
raw_config: &toml::Value,
|
||||
) -> (IndexMap<String, ModelProviderConfig>, Vec<ConfigWarning>) {
|
||||
let mut providers = IndexMap::new();
|
||||
let mut warnings = Vec::new();
|
||||
let Some(section) = raw_config.get("model_providers") else {
|
||||
return (providers, warnings);
|
||||
};
|
||||
let Some(table) = section.as_table() else {
|
||||
warnings.push(ConfigWarning::model_provider_section(
|
||||
ConfigWarningKind::NotATable,
|
||||
format!(
|
||||
"`model_providers` must be a table of [model_providers.<id>] entries, got {}; \
|
||||
all model providers ignored",
|
||||
section.type_str()
|
||||
),
|
||||
));
|
||||
return (providers, warnings);
|
||||
};
|
||||
for (id, value) in table {
|
||||
let mut unknown = Vec::new();
|
||||
match serde_ignored::deserialize::<_, _, ModelProviderConfig>(value.clone(), |path| {
|
||||
unknown.push(path.to_string());
|
||||
}) {
|
||||
Ok(provider) => {
|
||||
for key in unknown {
|
||||
warnings.push(ConfigWarning::model_provider(
|
||||
id,
|
||||
Some(key.as_str()),
|
||||
ConfigWarningKind::UnknownField,
|
||||
"unrecognized key; field ignored".to_owned(),
|
||||
));
|
||||
}
|
||||
if let Some(auth) = &provider.auth {
|
||||
for (field, kind, reason) in auth_config_issues(auth) {
|
||||
warnings.push(ConfigWarning::model_provider(
|
||||
id,
|
||||
Some(&format!("auth.{field}")),
|
||||
kind,
|
||||
reason,
|
||||
));
|
||||
}
|
||||
}
|
||||
let has_helper = provider.auth.is_some() || provider.auth_provider.is_some();
|
||||
let has_static_api_key = provider
|
||||
.api_key
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.is_some_and(|k| !k.is_empty());
|
||||
if has_helper && has_static_api_key {
|
||||
warnings.push(ConfigWarning::model_provider(
|
||||
id,
|
||||
Some("api_key"),
|
||||
ConfigWarningKind::ConflictingFields,
|
||||
"api_key shadows this provider's auth helper; the static key always \
|
||||
takes precedence, so the helper never runs for inheriting models"
|
||||
.to_owned(),
|
||||
));
|
||||
} else if has_helper
|
||||
&& provider
|
||||
.env_key
|
||||
.as_ref()
|
||||
.and_then(EnvKeys::primary)
|
||||
.is_some()
|
||||
{
|
||||
warnings.push(ConfigWarning::model_provider(
|
||||
id,
|
||||
Some("env_key"),
|
||||
ConfigWarningKind::ConflictingFields,
|
||||
"env_key may shadow this provider's auth helper; env_key takes precedence \
|
||||
when its variable resolves, otherwise the helper runs"
|
||||
.to_owned(),
|
||||
));
|
||||
}
|
||||
if provider.auth_provider.is_some() && provider.auth.is_some() {
|
||||
warnings.push(ConfigWarning::model_provider(
|
||||
id,
|
||||
Some("auth"),
|
||||
ConfigWarningKind::ConflictingFields,
|
||||
"inline auth is shadowed by auth_provider on this provider; the referenced \
|
||||
provider takes precedence, so the inline helper never runs"
|
||||
.to_owned(),
|
||||
));
|
||||
}
|
||||
providers.insert(id.clone(), provider);
|
||||
}
|
||||
Err(error) => {
|
||||
warnings.push(ConfigWarning::model_provider(
|
||||
id,
|
||||
None,
|
||||
ConfigWarningKind::InvalidValue,
|
||||
format!(
|
||||
"failed to parse ({error}); provider skipped, inheriting models \
|
||||
resolve with defaults"
|
||||
),
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
(providers, warnings)
|
||||
}
|
||||
|
||||
impl ConfigModelOverride {
|
||||
pub(crate) fn with_provider_defaults(
|
||||
&self,
|
||||
provider: &ModelProviderConfig,
|
||||
provider_id: &str,
|
||||
) -> Self {
|
||||
let ModelProviderConfig {
|
||||
base_url,
|
||||
api_base_url,
|
||||
env_key,
|
||||
api_key,
|
||||
api_backend,
|
||||
extra_headers,
|
||||
auth_provider,
|
||||
auth,
|
||||
context_window,
|
||||
} = provider;
|
||||
|
||||
let mut merged = self.clone();
|
||||
merged.model_provider = None;
|
||||
merged.base_url = merged.base_url.or_else(|| base_url.clone());
|
||||
merged.api_base_url = merged.api_base_url.or_else(|| api_base_url.clone());
|
||||
merged.api_backend = merged.api_backend.or_else(|| api_backend.clone());
|
||||
merged.context_window = merged.context_window.or(*context_window);
|
||||
if merged.extra_headers.is_empty() {
|
||||
merged.extra_headers = extra_headers.clone();
|
||||
}
|
||||
let model_sets_own_api_key = self
|
||||
.api_key
|
||||
.as_deref()
|
||||
.is_some_and(|k| !k.trim().is_empty());
|
||||
let model_sets_own_env_key = self.env_key.as_ref().and_then(EnvKeys::primary).is_some();
|
||||
let model_has_own_auth =
|
||||
model_sets_own_api_key || model_sets_own_env_key || self.auth_provider.is_some();
|
||||
if !model_has_own_auth {
|
||||
merged.api_key = api_key.clone();
|
||||
merged.env_key = env_key.clone();
|
||||
merged.auth_provider = auth_provider
|
||||
.clone()
|
||||
.or_else(|| auth.as_ref().map(|_| model_provider_auth_name(provider_id)));
|
||||
}
|
||||
merged
|
||||
}
|
||||
|
||||
pub(crate) fn with_missing_provider(&self) -> Self {
|
||||
let mut merged = self.clone();
|
||||
merged.model_provider = None;
|
||||
merged
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::agent::config::{Config, resolve_credentials, resolve_model_list};
|
||||
#[test]
|
||||
fn model_inherits_provider_connection_defaults() {
|
||||
let raw_config: toml::Value = toml::from_str(
|
||||
r#"
|
||||
[model_providers.gateway]
|
||||
base_url = "https://gateway.example/v1"
|
||||
context_window = 123456
|
||||
|
||||
[model_providers.gateway.extra_headers]
|
||||
X-Corp = "yes"
|
||||
|
||||
[model.via-gateway]
|
||||
model = "m"
|
||||
model_provider = "gateway"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let cfg = Config::new_from_toml_cfg(&raw_config).expect("config should parse");
|
||||
assert!(cfg.model_providers.contains_key("gateway"));
|
||||
let resolved = resolve_model_list(&cfg, None);
|
||||
let model = resolved.get("via-gateway").expect("model should exist");
|
||||
assert_eq!(model.info.base_url, "https://gateway.example/v1");
|
||||
assert_eq!(model.info.context_window.get(), 123456);
|
||||
assert_eq!(
|
||||
model.info.extra_headers.get("X-Corp").map(String::as_str),
|
||||
Some("yes")
|
||||
);
|
||||
assert!(
|
||||
model.has_own_credentials(),
|
||||
"a custom endpoint without a credential is BYOK, not session-authed"
|
||||
);
|
||||
assert_eq!(
|
||||
resolve_credentials(model, Some("session-jwt")).api_key,
|
||||
None,
|
||||
"the session token must not leak to the provider's custom endpoint"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn model_fields_override_provider_defaults() {
|
||||
let raw_config: toml::Value = toml::from_str(
|
||||
r#"
|
||||
[model_providers.gateway]
|
||||
base_url = "https://gateway.example/v1"
|
||||
context_window = 100000
|
||||
|
||||
[model.override-url]
|
||||
model = "m"
|
||||
model_provider = "gateway"
|
||||
base_url = "https://model-specific.example/v1"
|
||||
context_window = 200000
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let cfg = Config::new_from_toml_cfg(&raw_config).expect("config should parse");
|
||||
let resolved = resolve_model_list(&cfg, None);
|
||||
let model = resolved.get("override-url").expect("model should exist");
|
||||
assert_eq!(model.info.base_url, "https://model-specific.example/v1");
|
||||
assert_eq!(model.info.context_window.get(), 200000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn model_provider_inline_auth_registers_synthetic_provider() {
|
||||
let raw_config: toml::Value = toml::from_str(
|
||||
r#"
|
||||
[model_providers.gateway]
|
||||
base_url = "https://gateway.example/v1"
|
||||
context_window = 200000
|
||||
|
||||
[model_providers.gateway.auth]
|
||||
command = "printf gw-token"
|
||||
token_ttl_secs = 3600
|
||||
|
||||
[model.byok-via-gateway]
|
||||
model = "m"
|
||||
model_provider = "gateway"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let cfg = Config::new_from_toml_cfg(&raw_config).expect("config should parse");
|
||||
assert_eq!(
|
||||
cfg.auth_providers
|
||||
.get("model_provider:gateway")
|
||||
.map(|c| c.command.as_str()),
|
||||
Some("printf gw-token"),
|
||||
"inline auth registers a synthetic provider keyed by the id"
|
||||
);
|
||||
let resolved = resolve_model_list(&cfg, None);
|
||||
let model = resolved
|
||||
.get("byok-via-gateway")
|
||||
.expect("model should exist");
|
||||
let provider = model
|
||||
.auth_provider
|
||||
.as_ref()
|
||||
.expect("the model inherits the provider's auth");
|
||||
assert_eq!(provider.name, "model_provider:gateway");
|
||||
assert_eq!(provider.config.command, "printf gw-token");
|
||||
assert!(
|
||||
model.has_own_credentials(),
|
||||
"a provider-backed model is BYOK (session token must not leak)"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn model_with_own_key_ignores_provider_auth() {
|
||||
let raw_config: toml::Value = toml::from_str(
|
||||
r#"
|
||||
[model_providers.gateway]
|
||||
base_url = "https://gateway.example/v1"
|
||||
context_window = 200000
|
||||
|
||||
[model_providers.gateway.auth]
|
||||
command = "printf gw-token"
|
||||
|
||||
[model.own-key]
|
||||
model = "m"
|
||||
model_provider = "gateway"
|
||||
api_key = "sk-model-own"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let cfg = Config::new_from_toml_cfg(&raw_config).expect("config should parse");
|
||||
let resolved = resolve_model_list(&cfg, None);
|
||||
let model = resolved.get("own-key").expect("model should exist");
|
||||
assert_eq!(
|
||||
model.info.base_url, "https://gateway.example/v1",
|
||||
"non-auth connection fields are still inherited"
|
||||
);
|
||||
assert_eq!(
|
||||
model.effective_auth_provider().map(|p| p.name.as_str()),
|
||||
None,
|
||||
"the model's own key shadows the provider's auth"
|
||||
);
|
||||
let creds = resolve_credentials(model, Some("session-jwt"));
|
||||
assert_eq!(creds.api_key.as_deref(), Some("sk-model-own"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn undefined_model_provider_fails_closed() {
|
||||
use super::super::config_model_override_parse::{ConfigWarningKind, WarningTarget};
|
||||
|
||||
let raw_config: toml::Value = toml::from_str(
|
||||
r#"
|
||||
[model.dangling]
|
||||
model = "m"
|
||||
base_url = "https://third-party.example/v1"
|
||||
context_window = 200000
|
||||
model_provider = "ghost"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let cfg = Config::new_from_toml_cfg(&raw_config).expect("config should parse");
|
||||
assert!(
|
||||
cfg.config_warnings.iter().any(|w| {
|
||||
w.kind == ConfigWarningKind::InvalidValue
|
||||
&& matches!(
|
||||
&w.target,
|
||||
WarningTarget::Model { field, .. }
|
||||
if field.as_deref() == Some("model_provider")
|
||||
)
|
||||
}),
|
||||
"an undefined provider reference warns: {:?}",
|
||||
cfg.config_warnings
|
||||
);
|
||||
let resolved = resolve_model_list(&cfg, None);
|
||||
let model = resolved.get("dangling").expect("model should exist");
|
||||
assert_eq!(
|
||||
model.info.base_url, "https://third-party.example/v1",
|
||||
"the model keeps its own connection fields"
|
||||
);
|
||||
assert!(
|
||||
model.has_own_credentials(),
|
||||
"an undefined provider leaves the model BYOK, not session-authed"
|
||||
);
|
||||
let creds = resolve_credentials(model, Some("session-jwt"));
|
||||
assert_eq!(
|
||||
creds.api_key, None,
|
||||
"no credential resolves and the session token does not leak to the model's base_url"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn undefined_model_provider_keeps_model_own_key() {
|
||||
let raw_config: toml::Value = toml::from_str(
|
||||
r#"
|
||||
[model.own-key]
|
||||
model = "m"
|
||||
base_url = "https://third-party.example/v1"
|
||||
context_window = 200000
|
||||
api_key = "sk-model-own"
|
||||
model_provider = "ghost"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let cfg = Config::new_from_toml_cfg(&raw_config).expect("config should parse");
|
||||
let resolved = resolve_model_list(&cfg, None);
|
||||
let model = resolved.get("own-key").expect("model should exist");
|
||||
let creds = resolve_credentials(model, Some("session-jwt"));
|
||||
assert_eq!(creds.api_key.as_deref(), Some("sk-model-own"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn model_provider_parse_warnings_are_lenient_and_specific() {
|
||||
use super::super::config_model_override_parse::{ConfigWarningKind, WarningTarget};
|
||||
|
||||
let raw_config: toml::Value = toml::from_str(
|
||||
r#"
|
||||
[model_providers.good]
|
||||
base_url = "https://good.example/v1"
|
||||
|
||||
[model_providers.bad-type]
|
||||
context_window = "not-a-number"
|
||||
|
||||
[model_providers.typo]
|
||||
base_url = "https://typo.example/v1"
|
||||
unknown_field = 5
|
||||
|
||||
[model.on-broken-provider]
|
||||
model = "m"
|
||||
base_url = "https://x.example/v1"
|
||||
context_window = 200000
|
||||
model_provider = "bad-type"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let cfg = Config::new_from_toml_cfg(&raw_config)
|
||||
.expect("one bad provider must not fail the config");
|
||||
assert!(cfg.model_providers.contains_key("good"));
|
||||
assert!(
|
||||
!cfg.model_providers.contains_key("bad-type"),
|
||||
"a malformed provider is skipped"
|
||||
);
|
||||
|
||||
let has_provider = |id: &str, field: Option<&str>, kind: ConfigWarningKind| {
|
||||
cfg.config_warnings.iter().any(|w| {
|
||||
w.kind == kind
|
||||
&& matches!(
|
||||
&w.target,
|
||||
WarningTarget::ModelProvider { id: i, field: f }
|
||||
if i == id && f.as_deref() == field
|
||||
)
|
||||
})
|
||||
};
|
||||
assert!(has_provider(
|
||||
"bad-type",
|
||||
None,
|
||||
ConfigWarningKind::InvalidValue
|
||||
));
|
||||
assert!(has_provider(
|
||||
"typo",
|
||||
Some("unknown_field"),
|
||||
ConfigWarningKind::UnknownField
|
||||
));
|
||||
assert!(
|
||||
!cfg.config_warnings.iter().any(|w| {
|
||||
matches!(
|
||||
&w.target,
|
||||
WarningTarget::Model { field, .. }
|
||||
if field.as_deref() == Some("model_provider")
|
||||
)
|
||||
}),
|
||||
"a declared-but-malformed provider must not also warn as undefined: {:?}",
|
||||
cfg.config_warnings
|
||||
);
|
||||
|
||||
let raw_config: toml::Value = toml::from_str(r#"model_providers = "oops""#).unwrap();
|
||||
let cfg = Config::new_from_toml_cfg(&raw_config)
|
||||
.expect("a non-table model_providers must not fail the config");
|
||||
assert!(cfg.model_providers.is_empty());
|
||||
assert!(
|
||||
cfg.config_warnings.iter().any(|w| {
|
||||
matches!(w.target, WarningTarget::ModelProviderSection)
|
||||
&& w.kind == ConfigWarningKind::NotATable
|
||||
}),
|
||||
"non-table section warns: {:?}",
|
||||
cfg.config_warnings
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn model_provider_conflicting_credentials_warn() {
|
||||
use super::super::config_model_override_parse::{ConfigWarningKind, WarningTarget};
|
||||
|
||||
let raw_config: toml::Value = toml::from_str(
|
||||
r#"
|
||||
[model_providers.static-shadows]
|
||||
base_url = "https://a.example/v1"
|
||||
api_key = "sk-static"
|
||||
[model_providers.static-shadows.auth]
|
||||
command = "printf tok"
|
||||
|
||||
[model_providers.env-shadows]
|
||||
base_url = "https://b.example/v1"
|
||||
env_key = "SOME_VAR"
|
||||
[model_providers.env-shadows.auth]
|
||||
command = "printf tok"
|
||||
|
||||
[model_providers.two-helpers]
|
||||
base_url = "https://c.example/v1"
|
||||
auth_provider = "corp"
|
||||
[model_providers.two-helpers.auth]
|
||||
command = "printf tok"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let cfg = Config::new_from_toml_cfg(&raw_config).expect("config should parse");
|
||||
let has = |id: &str, field: &str| {
|
||||
cfg.config_warnings.iter().any(|w| {
|
||||
w.kind == ConfigWarningKind::ConflictingFields
|
||||
&& matches!(
|
||||
&w.target,
|
||||
WarningTarget::ModelProvider { id: i, field: f }
|
||||
if i == id && f.as_deref() == Some(field)
|
||||
)
|
||||
})
|
||||
};
|
||||
assert!(
|
||||
has("static-shadows", "api_key"),
|
||||
"a static api_key shadows the helper: {:?}",
|
||||
cfg.config_warnings
|
||||
);
|
||||
assert!(
|
||||
has("env-shadows", "env_key"),
|
||||
"an env_key may shadow the helper: {:?}",
|
||||
cfg.config_warnings
|
||||
);
|
||||
assert!(
|
||||
has("two-helpers", "auth"),
|
||||
"auth_provider shadows the inline auth helper: {:?}",
|
||||
cfg.config_warnings
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn model_provider_undefined_auth_provider_warns() {
|
||||
use super::super::config_model_override_parse::{ConfigWarningKind, WarningTarget};
|
||||
|
||||
let raw_config: toml::Value = toml::from_str(
|
||||
r#"
|
||||
[model_providers.gateway]
|
||||
base_url = "https://gateway.example/v1"
|
||||
auth_provider = "nonexistent"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let cfg = Config::new_from_toml_cfg(&raw_config).expect("config should parse");
|
||||
assert!(
|
||||
cfg.config_warnings.iter().any(|w| {
|
||||
w.kind == ConfigWarningKind::InvalidValue
|
||||
&& matches!(
|
||||
&w.target,
|
||||
WarningTarget::ModelProvider { id, field }
|
||||
if id == "gateway" && field.as_deref() == Some("auth_provider")
|
||||
)
|
||||
}),
|
||||
"an undefined provider auth_provider reference warns: {:?}",
|
||||
cfg.config_warnings
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn model_provider_inline_auth_namespace_collision_warns() {
|
||||
use super::super::config_model_override_parse::{ConfigWarningKind, WarningTarget};
|
||||
|
||||
let raw_config: toml::Value = toml::from_str(
|
||||
r#"
|
||||
[auth_provider."model_provider:gateway"]
|
||||
command = "printf hand-written"
|
||||
|
||||
[model_providers.gateway]
|
||||
base_url = "https://gateway.example/v1"
|
||||
|
||||
[model_providers.gateway.auth]
|
||||
command = "printf inline"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let cfg = Config::new_from_toml_cfg(&raw_config).expect("config should parse");
|
||||
assert!(
|
||||
cfg.config_warnings.iter().any(|w| {
|
||||
w.kind == ConfigWarningKind::ConflictingFields
|
||||
&& matches!(
|
||||
&w.target,
|
||||
WarningTarget::ModelProvider { id, field }
|
||||
if id == "gateway" && field.as_deref() == Some("auth")
|
||||
)
|
||||
}),
|
||||
"a reserved-namespace collision warns: {:?}",
|
||||
cfg.config_warnings
|
||||
);
|
||||
assert_eq!(
|
||||
cfg.auth_providers
|
||||
.get("model_provider:gateway")
|
||||
.map(|c| c.command.as_str()),
|
||||
Some("printf inline"),
|
||||
"inline auth wins the reserved name"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn model_inherits_provider_named_auth_provider() {
|
||||
let raw_config: toml::Value = toml::from_str(
|
||||
r#"
|
||||
[auth_provider.corp]
|
||||
command = "printf corp-token"
|
||||
token_ttl_secs = 3600
|
||||
|
||||
[model_providers.gateway]
|
||||
base_url = "https://gateway.example/v1"
|
||||
auth_provider = "corp"
|
||||
|
||||
[model.via-gateway]
|
||||
model = "m"
|
||||
model_provider = "gateway"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let cfg = Config::new_from_toml_cfg(&raw_config).expect("config should parse");
|
||||
let resolved = resolve_model_list(&cfg, None);
|
||||
let model = resolved.get("via-gateway").expect("model should exist");
|
||||
let provider = model
|
||||
.auth_provider
|
||||
.as_ref()
|
||||
.expect("the model inherits the provider's named auth_provider");
|
||||
assert_eq!(provider.name, "corp");
|
||||
assert_eq!(provider.config.command, "printf corp-token");
|
||||
assert!(model.has_own_credentials());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn model_inherits_provider_static_key() {
|
||||
let raw_config: toml::Value = toml::from_str(
|
||||
r#"
|
||||
[model_providers.gateway]
|
||||
base_url = "https://gateway.example/v1"
|
||||
api_key = "sk-provider"
|
||||
|
||||
[model.via-gateway]
|
||||
model = "m"
|
||||
model_provider = "gateway"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let cfg = Config::new_from_toml_cfg(&raw_config).expect("config should parse");
|
||||
let resolved = resolve_model_list(&cfg, None);
|
||||
let model = resolved.get("via-gateway").expect("model should exist");
|
||||
assert_eq!(
|
||||
resolve_credentials(model, Some("session-jwt"))
|
||||
.api_key
|
||||
.as_deref(),
|
||||
Some("sk-provider"),
|
||||
"the provider's static key resolves for the inheriting model"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declared_unresolved_credential_fails_closed_on_provider_endpoint() {
|
||||
let raw_config: toml::Value = toml::from_str(
|
||||
r#"
|
||||
[model_providers.gateway]
|
||||
base_url = "https://gateway.example/v1"
|
||||
|
||||
[model.via-gateway]
|
||||
model = "m"
|
||||
model_provider = "gateway"
|
||||
env_key = "DEFINITELY_UNSET_MODEL_PROVIDER_TEST_VAR"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let cfg = Config::new_from_toml_cfg(&raw_config).expect("config should parse");
|
||||
let resolved = resolve_model_list(&cfg, None);
|
||||
let model = resolved.get("via-gateway").expect("model should exist");
|
||||
assert_eq!(
|
||||
resolve_credentials(model, Some("session-jwt")).api_key,
|
||||
None,
|
||||
"an unresolved declared credential must not fall back to the session token"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn model_inherits_provider_api_backend_and_base_url() {
|
||||
let raw_config: toml::Value = toml::from_str(
|
||||
r#"
|
||||
[model_providers.gateway]
|
||||
base_url = "https://gateway.example/v1"
|
||||
api_base_url = "https://gateway.example/api"
|
||||
api_backend = "responses"
|
||||
api_key = "sk-provider"
|
||||
|
||||
[model.via-gateway]
|
||||
model = "m"
|
||||
model_provider = "gateway"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let cfg = Config::new_from_toml_cfg(&raw_config).expect("config should parse");
|
||||
let resolved = resolve_model_list(&cfg, None);
|
||||
let model = resolved.get("via-gateway").expect("model should exist");
|
||||
assert_eq!(
|
||||
model.info.api_backend,
|
||||
crate::sampling::ApiBackend::Responses
|
||||
);
|
||||
assert_eq!(
|
||||
model.api_base_url.as_deref(),
|
||||
Some("https://gateway.example/api")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn model_own_unresolved_key_ignores_provider_inline_auth() {
|
||||
let raw_config: toml::Value = toml::from_str(
|
||||
r#"
|
||||
[model_providers.gateway]
|
||||
base_url = "https://gateway.example/v1"
|
||||
|
||||
[model_providers.gateway.auth]
|
||||
command = "printf gw-token"
|
||||
|
||||
[model.own-env]
|
||||
model = "m"
|
||||
model_provider = "gateway"
|
||||
env_key = "DEFINITELY_UNSET_MODEL_PROVIDER_INLINE_VAR"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let cfg = Config::new_from_toml_cfg(&raw_config).expect("config should parse");
|
||||
let resolved = resolve_model_list(&cfg, None);
|
||||
let model = resolved.get("own-env").expect("model should exist");
|
||||
let effective = model
|
||||
.effective_auth_provider()
|
||||
.expect("an unresolved own credential fails closed via a provider ref");
|
||||
assert!(
|
||||
effective.name.contains("fail-closed"),
|
||||
"must pin the unusable fail-closed ref, not the live inline auth: {}",
|
||||
effective.name
|
||||
);
|
||||
assert!(
|
||||
effective.config.command.is_empty(),
|
||||
"the fail-closed ref is unusable"
|
||||
);
|
||||
assert_eq!(
|
||||
resolve_credentials(model, Some("session-jwt")).api_key,
|
||||
None,
|
||||
"must not fall back to the session token"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn fail_closed_ref_ignores_a_colliding_auth_provider_table() {
|
||||
let raw_config: toml::Value = toml::from_str(
|
||||
r#"
|
||||
[auth_provider."model_provider:gateway (fail-closed)"]
|
||||
command = "printf sneaky-token"
|
||||
|
||||
[model_providers.gateway]
|
||||
base_url = "https://gateway.example/v1"
|
||||
|
||||
[model.via-gateway]
|
||||
model = "m"
|
||||
context_window = 200000
|
||||
model_provider = "gateway"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let cfg = Config::new_from_toml_cfg(&raw_config).expect("config should parse");
|
||||
let resolved = resolve_model_list(&cfg, None);
|
||||
let model = resolved.get("via-gateway").expect("model should exist");
|
||||
assert_eq!(
|
||||
resolve_credentials(model, Some("session-jwt")).api_key,
|
||||
None,
|
||||
"a fail-closed ref must never resolve a colliding auth_provider table"
|
||||
);
|
||||
let effective = model
|
||||
.effective_auth_provider()
|
||||
.expect("fails closed via a provider ref");
|
||||
assert!(
|
||||
effective.config.command.is_empty(),
|
||||
"the fail-closed ref stays unusable despite the name collision"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn model_headers_shadow_provider_headers() {
|
||||
let raw_config: toml::Value = toml::from_str(
|
||||
r#"
|
||||
[model_providers.gateway]
|
||||
base_url = "https://gateway.example/v1"
|
||||
api_key = "sk-provider"
|
||||
|
||||
[model_providers.gateway.extra_headers]
|
||||
X-Corp = "yes"
|
||||
|
||||
[model.via-gateway]
|
||||
model = "m"
|
||||
context_window = 200000
|
||||
model_provider = "gateway"
|
||||
|
||||
[model.via-gateway.extra_headers]
|
||||
X-Model = "own"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let cfg = Config::new_from_toml_cfg(&raw_config).expect("config should parse");
|
||||
let resolved = resolve_model_list(&cfg, None);
|
||||
let model = resolved.get("via-gateway").expect("model should exist");
|
||||
assert_eq!(
|
||||
model.info.extra_headers.get("X-Model").map(String::as_str),
|
||||
Some("own")
|
||||
);
|
||||
assert!(
|
||||
model.info.extra_headers.get("X-Corp").is_none(),
|
||||
"a model that sets any header inherits none of the provider's"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn model_provider_inline_auth_ttl_and_timeout_warn() {
|
||||
use super::super::config_model_override_parse::{ConfigWarningKind, WarningTarget};
|
||||
|
||||
let raw_config: toml::Value = toml::from_str(
|
||||
r#"
|
||||
[model_providers.gateway]
|
||||
base_url = "https://gateway.example/v1"
|
||||
|
||||
[model_providers.gateway.auth]
|
||||
command = "printf tok"
|
||||
token_ttl_secs = 5
|
||||
timeout_secs = 0
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let cfg = Config::new_from_toml_cfg(&raw_config).expect("config should parse");
|
||||
let has = |field: &str| {
|
||||
cfg.config_warnings.iter().any(|w| {
|
||||
w.kind == ConfigWarningKind::InvalidValue
|
||||
&& matches!(
|
||||
&w.target,
|
||||
WarningTarget::ModelProvider { id, field: f }
|
||||
if id == "gateway" && f.as_deref() == Some(field)
|
||||
)
|
||||
})
|
||||
};
|
||||
assert!(
|
||||
has("auth.token_ttl_secs"),
|
||||
"inline auth ttl below the refresh margin warns: {:?}",
|
||||
cfg.config_warnings
|
||||
);
|
||||
assert!(
|
||||
has("auth.timeout_secs"),
|
||||
"inline auth timeout out of range warns: {:?}",
|
||||
cfg.config_warnings
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn blank_api_key_does_not_shadow_provider_auth() {
|
||||
let raw_config: toml::Value = toml::from_str(
|
||||
r#"
|
||||
[model_providers.gateway]
|
||||
base_url = "https://gateway.example/v1"
|
||||
|
||||
[model_providers.gateway.auth]
|
||||
command = "printf tok"
|
||||
|
||||
[model.m]
|
||||
model = "m"
|
||||
model_provider = "gateway"
|
||||
api_key = " "
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
let cfg = Config::new_from_toml_cfg(&raw_config).expect("config should parse");
|
||||
let resolved = resolve_model_list(&cfg, None);
|
||||
let provider = resolved["m"]
|
||||
.auth_provider
|
||||
.as_ref()
|
||||
.expect("blank api_key must not fail-close a working gateway");
|
||||
assert_eq!(provider.name.as_str(), "model_provider:gateway");
|
||||
assert!(!provider.is_fail_closed());
|
||||
}
|
||||
}
|
||||
|
|
@ -25,8 +25,28 @@ impl acp::Agent for MvpAgent {
|
|||
tracing::debug!(target : "sampling_log", "Received initialize request");
|
||||
xai_grok_telemetry::unified_log::info("agent initialized", None, None);
|
||||
self.start_subagent_coordinator();
|
||||
tokio::task::spawn_blocking(|| {
|
||||
let (auto_gc_policy, run_auto_gc) = {
|
||||
let cfg = self.cfg.borrow();
|
||||
let has_remote = cfg.remote_settings.is_some();
|
||||
let run = has_remote || !crate::util::config::resolve_remote_fetch_enabled();
|
||||
(cfg.resolve_worktree_auto_gc(), run)
|
||||
};
|
||||
if !run_auto_gc {
|
||||
tracing::debug!(
|
||||
"auto worktree gc deferred until remote_settings are available"
|
||||
);
|
||||
}
|
||||
tokio::task::spawn_blocking(move || {
|
||||
crate::session::worktree_pool::cleanup_stale_pool_worktrees(None);
|
||||
if !run_auto_gc {
|
||||
return;
|
||||
}
|
||||
let opts = xai_fast_worktree::AutoGcOptions::from_resolved(auto_gc_policy);
|
||||
if let Err(e) = xai_fast_worktree::WorktreeDb::open_default()
|
||||
.and_then(|db| xai_fast_worktree::maybe_auto_gc(&db, &opts))
|
||||
{
|
||||
tracing::warn!(error = % e, "auto worktree gc failed");
|
||||
}
|
||||
});
|
||||
tokio::task::spawn_blocking(|| {
|
||||
crate::session::persistence::cleanup_stale_sessions(None);
|
||||
|
|
@ -1072,6 +1092,7 @@ impl acp::Agent for MvpAgent {
|
|||
persisted_signals: None,
|
||||
persisted_plan_mode: None,
|
||||
persisted_goal_mode: None,
|
||||
persisted_workflow_runs: Vec::new(),
|
||||
persisted_announcement_state: None,
|
||||
session_meta: arguments.meta.as_ref(),
|
||||
managed_mcp_expires_at,
|
||||
|
|
@ -1361,6 +1382,7 @@ impl acp::Agent for MvpAgent {
|
|||
signals: persisted_signals,
|
||||
announcement_state: persisted_announcement_state,
|
||||
goal_mode_state: _persisted_goal_mode,
|
||||
workflow_runs: persisted_workflow_runs,
|
||||
} = persistence_info;
|
||||
let restored_compaction_count = persisted_signals
|
||||
.as_ref()
|
||||
|
|
@ -1612,6 +1634,7 @@ impl acp::Agent for MvpAgent {
|
|||
persisted_signals,
|
||||
persisted_plan_mode,
|
||||
persisted_goal_mode: _persisted_goal_mode,
|
||||
persisted_workflow_runs,
|
||||
persisted_announcement_state,
|
||||
session_meta: request_meta.as_ref(),
|
||||
managed_mcp_expires_at,
|
||||
|
|
@ -2169,7 +2192,6 @@ impl acp::Agent for MvpAgent {
|
|||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string())
|
||||
.or_else(|| self.cfg.borrow().client_version.clone());
|
||||
let agent_config = self.cfg.borrow().clone();
|
||||
let plugin_registry = self.plugin_registry_snapshot();
|
||||
let prompt_images: Vec<agent_client_protocol::ImageContent> = arguments
|
||||
.prompt
|
||||
|
|
@ -2251,11 +2273,9 @@ impl acp::Agent for MvpAgent {
|
|||
let before_workspace_fut = async {};
|
||||
futures::join!(
|
||||
upload_session_state(& ctx, "before", session_copy_rx,
|
||||
UploadWait::Confirm), before_workspace_fut, upload_config(& ctx,
|
||||
& agent_config), crate
|
||||
::upload::config_files::upload_config_files(& ctx),
|
||||
upload_images(& ctx, & prompt_images), upload_plugin_state(& ctx,
|
||||
plugin_registry.as_deref()),
|
||||
UploadWait::Confirm), before_workspace_fut, upload_images(& ctx,
|
||||
& prompt_images), upload_plugin_state(& ctx, plugin_registry
|
||||
.as_deref()),
|
||||
);
|
||||
},
|
||||
);
|
||||
|
|
@ -3205,12 +3225,13 @@ impl acp::Agent for MvpAgent {
|
|||
| "x.ai/session/update_mcp_servers" | "x.ai/session/fork"
|
||||
| "x.ai/internal/reload_all_mcp_servers"
|
||||
| "x.ai/internal/reload_project_mcp_servers" | "x.ai/internal/reload_skills"
|
||||
| "x.ai/internal/reload_models" | "x.ai/internal/reload_models_cache"
|
||||
| "x.ai/internal/auth_cleared" | "x.ai/plugins/reload"
|
||||
| "x.ai/commands/list" => {
|
||||
| "x.ai/internal/reload_workflows" | "x.ai/internal/reload_models"
|
||||
| "x.ai/internal/reload_models_cache" | "x.ai/internal/auth_cleared"
|
||||
| "x.ai/plugins/reload" | "x.ai/commands/list" => {
|
||||
crate::extensions::session_admin::handle(self, &args).await
|
||||
}
|
||||
"x.ai/session/repair" => crate::extensions::repair::handle(self, &args).await,
|
||||
"x.ai/session/usage" => crate::extensions::usage::handle(self, &args).await,
|
||||
"x.ai/memory/flush" | "x.ai/memory/rewrite" => {
|
||||
crate::extensions::memory::handle(self, &args).await
|
||||
}
|
||||
|
|
@ -3497,9 +3518,10 @@ impl acp::Agent for MvpAgent {
|
|||
let ops = self.resolve_workspace_ops()?;
|
||||
crate::extensions::code_nav::handle(self, &ops, &args).await
|
||||
}
|
||||
s if s.starts_with("x.ai/skills/") => {
|
||||
s if s.starts_with("x.ai/skills/") || s == "x.ai/workflows/list" => {
|
||||
let compat = self.cfg.borrow().compat_resolved;
|
||||
crate::extensions::skills::handle(
|
||||
self,
|
||||
&args,
|
||||
self.plugin_registry_handle.snapshot().as_deref(),
|
||||
compat,
|
||||
|
|
|
|||
|
|
@ -16,6 +16,34 @@ fn byok_from_models(
|
|||
.or_else(|| models.values().find_map(|m| m.own_credential()))
|
||||
}
|
||||
impl MvpAgent {
|
||||
pub fn reload_skills_all_sessions(&self) -> usize {
|
||||
let session_ids: Vec<agent_client_protocol::SessionId> = self
|
||||
.sessions
|
||||
.borrow()
|
||||
.keys()
|
||||
.cloned()
|
||||
.collect();
|
||||
for sid in &session_ids {
|
||||
if let Some(handle) = self.sessions.borrow().get(sid).cloned() {
|
||||
let _ = handle.cmd_tx.send(SessionCommand::ReloadSkills);
|
||||
}
|
||||
}
|
||||
session_ids.len()
|
||||
}
|
||||
pub fn advertise_commands_all_sessions(&self) -> usize {
|
||||
let session_ids: Vec<agent_client_protocol::SessionId> = self
|
||||
.sessions
|
||||
.borrow()
|
||||
.keys()
|
||||
.cloned()
|
||||
.collect();
|
||||
for session_id in &session_ids {
|
||||
if let Some(handle) = self.sessions.borrow().get(session_id).cloned() {
|
||||
let _ = handle.cmd_tx.send(SessionCommand::AdvertiseCommands);
|
||||
}
|
||||
}
|
||||
session_ids.len()
|
||||
}
|
||||
pub(super) fn resolve_image_description_model(&self) -> String {
|
||||
self.cfg
|
||||
.borrow()
|
||||
|
|
@ -550,17 +578,13 @@ impl MvpAgent {
|
|||
/// `acp_session.rs` fill in the real per-model gating as soon as a
|
||||
/// session starts.
|
||||
///
|
||||
/// Exception: `/goal` is gated on the `resolve_goal()` feature flag
|
||||
/// (a config/managed-settings switch known at initialize time) plus
|
||||
/// the `update_goal` tool, which is part of the default coding-agent
|
||||
/// toolset. So when the flag is on we advertise `/goal` pre-session;
|
||||
/// otherwise it wouldn't appear in the slash menu until after the
|
||||
/// first user turn created a session.
|
||||
pub(crate) fn command_availability(
|
||||
&self,
|
||||
) -> crate::session::slash_commands::CommandAvailability {
|
||||
crate::session::slash_commands::CommandAvailability {
|
||||
goal: self.cfg.borrow().resolve_goal().value,
|
||||
workflows: self.cfg.borrow().resolve_workflows().value,
|
||||
..crate::session::slash_commands::CommandAvailability::default()
|
||||
}
|
||||
}
|
||||
|
|
@ -800,6 +824,7 @@ impl MvpAgent {
|
|||
let is_xai = auth.is_xai_auth();
|
||||
let user_id = auth.user_id.clone();
|
||||
let team_id = auth.team_id.clone();
|
||||
let remote_was_absent = self.cfg.borrow().remote_settings.is_none();
|
||||
let Some(settings) = self.fetch_remote_settings(auth.clone()).await else {
|
||||
tracing::warn!("post-auth settings refresh failed (HTTP or parse error)");
|
||||
return;
|
||||
|
|
@ -861,6 +886,9 @@ impl MvpAgent {
|
|||
crate::auth::credential_provider::sync_external_otel_identity();
|
||||
self.emit_announcements(AnnouncementsPushMode::IfChanged);
|
||||
self.reconfigure_heap_profile_monitor();
|
||||
if remote_was_absent {
|
||||
self.spawn_auto_worktree_gc();
|
||||
}
|
||||
}
|
||||
/// Refresh remote settings settings and re-resolve eagerly-resolved config fields.
|
||||
///
|
||||
|
|
@ -1471,9 +1499,9 @@ impl MvpAgent {
|
|||
raw,
|
||||
cfg.remote_settings.as_ref(),
|
||||
);
|
||||
let session_registry_local = config_root
|
||||
.as_ref()
|
||||
.and_then(crate::util::config::session_registry_from_toml_opt);
|
||||
let session_registry_local = crate::util::config::session_registry_local_override(
|
||||
config_root.as_ref(),
|
||||
);
|
||||
tracing::info!(
|
||||
worktree_type = ? worktree_type, source = wt_source,
|
||||
"WORKTREE_CONFIG_SHELL: resolved worktree type at agent startup"
|
||||
|
|
@ -2892,6 +2920,7 @@ impl MvpAgent {
|
|||
persisted_signals,
|
||||
persisted_plan_mode,
|
||||
persisted_goal_mode,
|
||||
persisted_workflow_runs,
|
||||
persisted_announcement_state,
|
||||
session_meta,
|
||||
managed_mcp_expires_at,
|
||||
|
|
@ -3039,6 +3068,7 @@ impl MvpAgent {
|
|||
let resolved = cfg.resolve_feedback();
|
||||
let flags = crate::session::feedback_manager::FeedbackFlags {
|
||||
enabled: resolved.value,
|
||||
user: cfg.feedback.user.clone(),
|
||||
};
|
||||
(resolved, flags)
|
||||
};
|
||||
|
|
@ -3353,6 +3383,7 @@ impl MvpAgent {
|
|||
let web_fetch_config = self.prepare_web_fetch_config();
|
||||
let write_file_enabled = self.cfg.borrow().resolve_write_file().value;
|
||||
let goal_enabled = self.cfg.borrow().resolve_goal().value;
|
||||
let background_workflows_enabled = self.cfg.borrow().resolve_workflows().value;
|
||||
let subagents_enabled = self.cfg.borrow().subagents_enabled;
|
||||
let ask_user_question_enabled = crate::upload::turn::parse_ask_user_question_from_meta(
|
||||
session_meta,
|
||||
|
|
@ -3566,6 +3597,7 @@ impl MvpAgent {
|
|||
persisted_signals,
|
||||
persisted_plan_mode,
|
||||
persisted_goal_mode,
|
||||
persisted_workflow_runs,
|
||||
persisted_announcement_state,
|
||||
self.memory_config.clone(),
|
||||
loc_tracking_enabled,
|
||||
|
|
@ -3586,6 +3618,7 @@ impl MvpAgent {
|
|||
app_builder_deployer_config,
|
||||
write_file_enabled,
|
||||
goal_enabled,
|
||||
background_workflows_enabled,
|
||||
subagents_enabled,
|
||||
ask_user_question_enabled,
|
||||
client_hooks,
|
||||
|
|
|
|||
|
|
@ -84,10 +84,9 @@ use crate::tools::ToolContext;
|
|||
use crate::upload::manifest::write_error_manifest;
|
||||
use crate::upload::trace::{
|
||||
GCS_SCHEMA_VERSION, PromptMetadata, TurnResultMetadata,
|
||||
build_chat_history_session_state, local_sandbox_telemetry, upload_config,
|
||||
upload_full_prompt_txt, upload_harness_session_archive, upload_images,
|
||||
upload_metadata, upload_plugin_state, upload_session_state, upload_turn_messages,
|
||||
upload_turn_result, upload_unified_log,
|
||||
build_chat_history_session_state, local_sandbox_telemetry, upload_full_prompt_txt,
|
||||
upload_harness_session_archive, upload_images, upload_metadata, upload_plugin_state,
|
||||
upload_session_state, upload_turn_messages, upload_turn_result, upload_unified_log,
|
||||
};
|
||||
use crate::upload::turn::{
|
||||
PromptTraceContext, UploadWait, complete_prompt_trace, spawn_upload_task,
|
||||
|
|
@ -204,6 +203,9 @@ pub(crate) struct SessionSpawnOptions<'a> {
|
|||
pub persisted_signals: Option<crate::session::signals::SessionSignals>,
|
||||
pub persisted_plan_mode: Option<crate::session::plan_mode::PlanModeSnapshot>,
|
||||
pub persisted_goal_mode: Option<crate::session::goal_tracker::GoalOrchestration>,
|
||||
pub persisted_workflow_runs: Vec<
|
||||
crate::session::workflow::store::RestoredWorkflowRun,
|
||||
>,
|
||||
pub persisted_announcement_state: Option<
|
||||
crate::session::announcement_state::AnnouncementState,
|
||||
>,
|
||||
|
|
@ -342,6 +344,7 @@ pub(crate) fn chat_session_spawn_options<'a>(
|
|||
persisted_signals: None,
|
||||
persisted_plan_mode: None,
|
||||
persisted_goal_mode: None,
|
||||
persisted_workflow_runs: Vec::new(),
|
||||
persisted_announcement_state: None,
|
||||
session_meta,
|
||||
managed_mcp_expires_at: None,
|
||||
|
|
@ -753,7 +756,8 @@ pub struct MvpAgent {
|
|||
pub(crate) worktree_type: crate::util::config::WorktreeType,
|
||||
/// Restore codebase state on worktree resume (resolved: local config > remote > default false).
|
||||
pub(crate) restore_code: bool,
|
||||
/// Local config.toml override for session registry (`[cli] session_registry`).
|
||||
/// Local session-registry override: `GROK_SESSION_REGISTRY` env, else
|
||||
/// `[cli] session_registry`.
|
||||
/// `Some(true)` enables, `Some(false)` disables, `None` defers to remote settings.
|
||||
session_registry_local: Option<bool>,
|
||||
/// Managed MCP configs and gateway tool catalog; lazily fetched.
|
||||
|
|
@ -1813,6 +1817,7 @@ impl MvpAgent {
|
|||
),
|
||||
);
|
||||
if let Some(settings) = unblocked.settings {
|
||||
let remote_was_absent = self.cfg.borrow().remote_settings.is_none();
|
||||
{
|
||||
let mut cfg = self.cfg.borrow_mut();
|
||||
cfg.remote_settings = Some(settings);
|
||||
|
|
@ -1823,6 +1828,9 @@ impl MvpAgent {
|
|||
self.sync_collection_config_gate();
|
||||
self.emit_announcements(AnnouncementsPushMode::IfChanged);
|
||||
self.reconfigure_heap_profile_monitor();
|
||||
if remote_was_absent {
|
||||
self.spawn_auto_worktree_gc();
|
||||
}
|
||||
}
|
||||
if crate::util::config::resolve_remote_fetch_enabled()
|
||||
&& !settings_allow_access(self.cfg.borrow().remote_settings.as_ref())
|
||||
|
|
@ -2035,6 +2043,19 @@ impl MvpAgent {
|
|||
self.emit_settings_update_notification();
|
||||
self.emit_announcements(AnnouncementsPushMode::IfChanged);
|
||||
self.reconfigure_heap_profile_monitor();
|
||||
self.spawn_auto_worktree_gc();
|
||||
}
|
||||
/// Resolve current auto-GC policy and run it on the blocking pool.
|
||||
pub(super) fn spawn_auto_worktree_gc(&self) {
|
||||
let auto_gc_policy = self.cfg.borrow().resolve_worktree_auto_gc();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let opts = xai_fast_worktree::AutoGcOptions::from_resolved(auto_gc_policy);
|
||||
if let Err(e) = xai_fast_worktree::WorktreeDb::open_default()
|
||||
.and_then(|db| xai_fast_worktree::maybe_auto_gc(&db, &opts))
|
||||
{
|
||||
tracing::warn!(error = % e, "auto worktree gc failed");
|
||||
}
|
||||
});
|
||||
}
|
||||
/// Fire-and-forget `x.ai/settings/update` from the current remote snapshot.
|
||||
pub(super) fn emit_settings_update_notification(&self) {
|
||||
|
|
@ -2217,16 +2238,7 @@ async fn handle_synthetic_turn_trace(
|
|||
use crate::session::SessionCommand;
|
||||
use crate::upload::turn::{UploadWait, complete_prompt_trace, spawn_upload_task};
|
||||
let turn_started_at = chrono::Utc::now().to_rfc3339();
|
||||
let (
|
||||
info,
|
||||
turn_number,
|
||||
agent_config,
|
||||
user_id,
|
||||
user_email,
|
||||
client_source,
|
||||
client_version,
|
||||
model,
|
||||
) = {
|
||||
let (info, turn_number, user_id, user_email, client_source, client_version, model) = {
|
||||
let this = agent_ref.get();
|
||||
let session_info = {
|
||||
let sessions = this.sessions.borrow();
|
||||
|
|
@ -2265,17 +2277,7 @@ async fn handle_synthetic_turn_trace(
|
|||
.map(|h| h.model_id.0.to_string())
|
||||
.unwrap_or_else(|| this.models_manager.current_model_id().0.to_string())
|
||||
};
|
||||
let agent_config = this.cfg.borrow().clone();
|
||||
(
|
||||
info,
|
||||
turn_number,
|
||||
agent_config,
|
||||
user_id,
|
||||
user_email,
|
||||
client_source,
|
||||
client_version,
|
||||
model,
|
||||
)
|
||||
(info, turn_number, user_id, user_email, client_source, client_version, model)
|
||||
};
|
||||
let this = agent_ref.get();
|
||||
let trace_context = this.get_trace_context(&info, turn_number).await;
|
||||
|
|
@ -2323,8 +2325,7 @@ async fn handle_synthetic_turn_trace(
|
|||
futures::join!(
|
||||
upload_session_state(& before_ctx, "before", request
|
||||
.before_session_copy_rx, UploadWait::Confirm,), upload_metadata(&
|
||||
before_ctx, metadata), upload_config(& before_ctx, & agent_config), crate
|
||||
::upload::config_files::upload_config_files(& before_ctx),
|
||||
before_ctx, metadata),
|
||||
);
|
||||
},
|
||||
);
|
||||
|
|
|
|||
|
|
@ -26,26 +26,43 @@ impl MvpAgent {
|
|||
match event {
|
||||
SubagentEvent::Spawn(boxed) => {
|
||||
let mut request = *boxed;
|
||||
let agent_ref = agent_ref.clone();
|
||||
tokio::task::spawn_local(async move {
|
||||
{
|
||||
let this = agent_ref.get();
|
||||
let parent_is_session = this.sessions.borrow().contains_key(
|
||||
&acp::SessionId::new(request.parent_session_id.clone()),
|
||||
);
|
||||
if !parent_is_session
|
||||
&& let Some(root) = this
|
||||
.subagent_coordinator
|
||||
.borrow()
|
||||
.parent_of_child_session(&request.parent_session_id)
|
||||
{
|
||||
tracing::info!(
|
||||
child_session_id = % request.parent_session_id,
|
||||
root_session_id = % root, subagent_id = % request.id,
|
||||
"Re-parenting child-session spawn to root session"
|
||||
);
|
||||
request.parent_session_id = root;
|
||||
request.surface_completion = false;
|
||||
if !parent_is_session {
|
||||
let child_sess = request.parent_session_id.clone();
|
||||
let reparent = {
|
||||
let coord = this.subagent_coordinator.borrow();
|
||||
coord.parent_of_child_session(&child_sess).map(|root| {
|
||||
(root, coord.loop_task_id_of_child_session(&child_sess))
|
||||
})
|
||||
};
|
||||
if let Some((root, inherited_loop)) = reparent {
|
||||
tracing::info!(
|
||||
child_session_id = % child_sess, root_session_id = % root,
|
||||
subagent_id = % request.id,
|
||||
"Re-parenting child-session spawn to root session"
|
||||
);
|
||||
request.parent_session_id = root;
|
||||
request.surface_completion = false;
|
||||
if request.runtime_overrides.loop_task_id.is_none() {
|
||||
request.runtime_overrides.loop_task_id = inherited_loop;
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(task_id) =
|
||||
request.runtime_overrides.loop_task_id.clone()
|
||||
{
|
||||
this.subagent_coordinator
|
||||
.borrow_mut()
|
||||
.record_loop_owner(&request.id, &task_id);
|
||||
}
|
||||
}
|
||||
let agent_ref = agent_ref.clone();
|
||||
tokio::task::spawn_local(async move {
|
||||
let this = agent_ref.get();
|
||||
let parent_sid = request.parent_session_id.clone();
|
||||
let Some(mut ctx) =
|
||||
this.try_build_subagent_spawn_context(&parent_sid)
|
||||
|
|
@ -55,6 +72,9 @@ impl MvpAgent {
|
|||
.id,
|
||||
"Spawn for unknown/evicted parent session, failing request"
|
||||
);
|
||||
this.subagent_coordinator
|
||||
.borrow_mut()
|
||||
.remove_loop_owner(&request.id);
|
||||
crate::agent::subagent::send_failure(
|
||||
request,
|
||||
"Parent session not found (evicted or torn down); cannot spawn subagent.",
|
||||
|
|
@ -164,23 +184,57 @@ impl MvpAgent {
|
|||
}
|
||||
});
|
||||
}
|
||||
SubagentEvent::Cancel(request) => {
|
||||
let this = agent_ref.get();
|
||||
let outcome = {
|
||||
let mut coord = this.subagent_coordinator.borrow_mut();
|
||||
match request.target {
|
||||
SubagentCancelTarget::SubagentId(ref subagent_id) => {
|
||||
coord.mark_explicitly_killed(subagent_id);
|
||||
coord.cancel_with_outcome(subagent_id)
|
||||
SubagentEvent::Cancel(request) => match request.target {
|
||||
SubagentCancelTarget::WorkflowRunId(run_id) => {
|
||||
let agent_ref = agent_ref.clone();
|
||||
tokio::task::spawn_local(async move {
|
||||
let notify = {
|
||||
let this = agent_ref.get();
|
||||
let mut coord = this.subagent_coordinator.borrow_mut();
|
||||
coord.cancel_workflow_children(&run_id);
|
||||
coord.completion_notify()
|
||||
};
|
||||
loop {
|
||||
let notified = notify.notified();
|
||||
let outstanding = {
|
||||
let this = agent_ref.get();
|
||||
this.subagent_coordinator
|
||||
.borrow()
|
||||
.outstanding_for_workflow(&run_id)
|
||||
};
|
||||
if outstanding == 0 {
|
||||
let _ = request
|
||||
.respond_to
|
||||
.send(SubagentCancelOutcome::Cancelled);
|
||||
break;
|
||||
}
|
||||
notified.await;
|
||||
}
|
||||
SubagentCancelTarget::ParentPromptId(ref parent_prompt_id) => {
|
||||
coord.cancel_by_parent_prompt_id(parent_prompt_id);
|
||||
SubagentCancelOutcome::Cancelled
|
||||
});
|
||||
}
|
||||
target => {
|
||||
let this = agent_ref.get();
|
||||
let outcome = {
|
||||
let mut coord = this.subagent_coordinator.borrow_mut();
|
||||
match target {
|
||||
SubagentCancelTarget::SubagentId(ref subagent_id) => {
|
||||
coord.mark_explicitly_killed(subagent_id);
|
||||
coord.cancel_with_outcome(subagent_id)
|
||||
}
|
||||
SubagentCancelTarget::ParentPromptId(
|
||||
ref parent_prompt_id,
|
||||
) => {
|
||||
coord.cancel_by_parent_prompt_id(parent_prompt_id);
|
||||
SubagentCancelOutcome::Cancelled
|
||||
}
|
||||
SubagentCancelTarget::WorkflowRunId(_) => {
|
||||
unreachable!("handled above")
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
let _ = request.respond_to.send(outcome);
|
||||
}
|
||||
};
|
||||
let _ = request.respond_to.send(outcome);
|
||||
}
|
||||
},
|
||||
SubagentEvent::ListActive(request) => {
|
||||
let this = agent_ref.get();
|
||||
let summaries = this
|
||||
|
|
@ -257,6 +311,14 @@ impl MvpAgent {
|
|||
let _ = request.respond_to.send(outcome);
|
||||
});
|
||||
}
|
||||
SubagentEvent::LoopUnitActive(request) => {
|
||||
let this = agent_ref.get();
|
||||
let active = this
|
||||
.subagent_coordinator
|
||||
.borrow()
|
||||
.loop_unit_active(&request.task_id);
|
||||
let _ = request.respond_to.send(active);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -519,6 +581,7 @@ impl MvpAgent {
|
|||
app_builder_deployer_config: self.prepare_app_builder_deployer_config(),
|
||||
write_file_enabled: self.cfg.borrow().resolve_write_file().value,
|
||||
goal_enabled: self.cfg.borrow().resolve_goal().value,
|
||||
background_workflows_enabled: self.cfg.borrow().resolve_workflows().value,
|
||||
ask_user_question_enabled,
|
||||
parent_cmd_tx: parent_cmd_tx.clone(),
|
||||
parent_session_info: {
|
||||
|
|
@ -634,7 +697,9 @@ impl MvpAgent {
|
|||
sessions
|
||||
.get(&parent_sid)
|
||||
.map(|h| h.tool_context.blocking_wait_depth.clone())
|
||||
.unwrap_or_else(|| std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)))
|
||||
.unwrap_or_else(|| {
|
||||
std::sync::Arc::new(crate::tools::tool_context::BlockingWaitState::new())
|
||||
})
|
||||
},
|
||||
parent_terminal_backend: parent_terminal_backend.clone(),
|
||||
parent_notification_handle: parent_notification_handle.clone(),
|
||||
|
|
|
|||
|
|
@ -1611,6 +1611,38 @@ fn build_minimal_agent_for_tests() -> MvpAgent {
|
|||
let cfg = AgentConfig::default();
|
||||
MvpAgent::new(gateway, &cfg, auth_manager, None).expect("valid test config")
|
||||
}
|
||||
fn session_usage_request(session_id: &str) -> acp::ExtRequest {
|
||||
acp::ExtRequest::new(
|
||||
"x.ai/session/usage",
|
||||
serde_json::value::to_raw_value(&serde_json::json!({ "sessionId" : session_id }))
|
||||
.unwrap()
|
||||
.into(),
|
||||
)
|
||||
}
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn session_usage_unknown_session_is_resource_not_found() {
|
||||
let agent = build_minimal_agent_for_tests();
|
||||
let err = crate::extensions::usage::handle(&agent, &session_usage_request("no-such-session"))
|
||||
.await
|
||||
.expect_err("unknown session");
|
||||
assert_eq!(
|
||||
err.code,
|
||||
acp::Error::resource_not_found(None::<String>).code
|
||||
);
|
||||
}
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn session_usage_dead_chat_state_actor_fails_closed() {
|
||||
let agent = build_minimal_agent_for_tests();
|
||||
let sid = acp::SessionId::new("usage-dead-actor-sess");
|
||||
let mut handle = make_test_handle("test-model", false, None);
|
||||
handle.info.id = sid.clone();
|
||||
agent.sessions.borrow_mut().insert(sid, handle);
|
||||
let err =
|
||||
crate::extensions::usage::handle(&agent, &session_usage_request("usage-dead-actor-sess"))
|
||||
.await
|
||||
.expect_err("dead chat-state actor");
|
||||
assert_eq!(err.code, acp::Error::internal_error().code);
|
||||
}
|
||||
/// Build a minimal MvpAgent with pre-loaded auth for gate tests.
|
||||
fn build_agent_with_auth(auth: crate::auth::GrokAuth) -> MvpAgent {
|
||||
use crate::agent::config::Config as AgentConfig;
|
||||
|
|
|
|||
|
|
@ -509,25 +509,25 @@ where
|
|||
let mut keepalive = tokio::time::interval(Duration::from_secs(KEEPALIVE_INTERVAL_SECS));
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = cancel_write.cancelled() => break, msg_opt = from_agent_rx.recv() =>
|
||||
{ match msg_opt { Some(msg) => { if
|
||||
tracing::enabled!(tracing::Level::DEBUG) { if let Ok(json_val) =
|
||||
serde_json::from_str::< serde_json::Value > (& msg) { let method =
|
||||
json_val.get("method").and_then(| m | m.as_str()); let line_to_print =
|
||||
match method { Some("session/update") => { let params = json_val
|
||||
.get("params").unwrap_or(& serde_json::Value::Null);
|
||||
format!("acp_outbound::session/update::{params}") } Some(m) =>
|
||||
format!("acp_outbound::{m}"), None => "acp_outbound::response"
|
||||
.to_string(), }; debug!("{line_to_print}"); } else {
|
||||
debug!("acp_outbound::response"); } }
|
||||
if ! msg.is_empty() && let Err(e) =
|
||||
ws_outbound.send(Message::Text(Utf8Bytes::from(msg))). await {
|
||||
warn!(error = ? e, "failed to send to WS"); break; } } None => {
|
||||
info!("Agent outbound channel closed"); break; } } } _ = keepalive.tick()
|
||||
=> { tprintln!("ws::keep_alive_tick"); if let Err(e) = ws_outbound
|
||||
.send(Message::Ping(Vec::new().into())). await {
|
||||
tprintln!("ws::keep_alive::error::{:?}", & e); break; } }
|
||||
}
|
||||
_ = cancel_write.cancelled() => break, msg_opt = from_agent_rx.recv() =>
|
||||
{ match msg_opt { Some(msg) => { if
|
||||
tracing::enabled!(tracing::Level::DEBUG) { if let Ok(json_val) =
|
||||
serde_json::from_str::< serde_json::Value > (& msg) { let method =
|
||||
json_val.get("method").and_then(| m | m.as_str()); let line_to_print =
|
||||
match method { Some("session/update") => { let params = json_val
|
||||
.get("params").unwrap_or(& serde_json::Value::Null);
|
||||
format!("acp_outbound::session/update::{params}") } Some(m) =>
|
||||
format!("acp_outbound::{m}"), None => "acp_outbound::response"
|
||||
.to_string(), }; debug!("{line_to_print}"); } else {
|
||||
debug!("acp_outbound::response"); } }
|
||||
if ! msg.is_empty() && let Err(e) =
|
||||
ws_outbound.send(Message::Text(Utf8Bytes::from(msg))). await {
|
||||
warn!(error = ? e, "failed to send to WS"); break; } } None => {
|
||||
info!("Agent outbound channel closed"); break; } } } _ = keepalive.tick()
|
||||
=> { tprintln!("ws::keep_alive_tick"); if let Err(e) = ws_outbound
|
||||
.send(Message::Ping(Vec::new().into())). await {
|
||||
tprintln!("ws::keep_alive::error::{:?}", & e); break; } }
|
||||
}
|
||||
}
|
||||
anyhow::Ok(())
|
||||
};
|
||||
|
|
|
|||
|
|
@ -59,6 +59,7 @@ fn effort_label(effort: ReasoningEffort) -> String {
|
|||
ReasoningEffort::Medium => "Medium",
|
||||
ReasoningEffort::High => "High",
|
||||
ReasoningEffort::Xhigh => "X-High",
|
||||
ReasoningEffort::Max => "Max",
|
||||
}
|
||||
.to_string()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -71,7 +71,7 @@ pub struct UpdateRequest {
|
|||
// Response types
|
||||
// ============================================================================
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SessionRecord {
|
||||
pub session_id: String,
|
||||
|
|
|
|||
|
|
@ -16,7 +16,7 @@ use crate::terminal::AsyncTerminalRunner;
|
|||
use crate::tools::ToolContext;
|
||||
use crate::upload::trace::{
|
||||
GCS_SCHEMA_VERSION, PromptMetadata, SubagentSpawnedRef, TurnResultMetadata,
|
||||
local_sandbox_telemetry, upload_config, upload_metadata, upload_session_state,
|
||||
local_sandbox_telemetry, upload_metadata, upload_session_state,
|
||||
upload_subagent_metadata, upload_turn_result,
|
||||
};
|
||||
use crate::upload::turn::{PromptTraceContext, complete_prompt_trace};
|
||||
|
|
@ -38,6 +38,7 @@ impl SubagentCoordinator {
|
|||
running_gauge: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
|
||||
block_wait_slots: HashMap::new(),
|
||||
subagent_usage_not_applied_prompts: std::collections::HashSet::new(),
|
||||
loop_owned: HashMap::new(),
|
||||
}
|
||||
}
|
||||
pub fn mark_subagent_usage_not_applied(&mut self, prompt_id: &str) {
|
||||
|
|
@ -76,14 +77,6 @@ impl SubagentCoordinator {
|
|||
std::sync::atomic::Ordering::Relaxed,
|
||||
);
|
||||
}
|
||||
/// Returns a handle to the completion [`Notify`].
|
||||
#[cfg_attr(
|
||||
not(test),
|
||||
expect(
|
||||
dead_code,
|
||||
reason = "used from tests only; remove expect when wired in production"
|
||||
)
|
||||
)]
|
||||
pub fn completion_notify(&self) -> Arc<Notify> {
|
||||
Arc::clone(&self.completion_notify)
|
||||
}
|
||||
|
|
@ -236,6 +229,7 @@ impl SubagentCoordinator {
|
|||
description: pending.description,
|
||||
parent_prompt_id: pending.parent_prompt_id,
|
||||
parent_session_id: pending.parent_session_id,
|
||||
owner: pending.owner,
|
||||
persona: pending.persona,
|
||||
started_at: pending.started_at,
|
||||
error,
|
||||
|
|
@ -261,6 +255,7 @@ impl SubagentCoordinator {
|
|||
description: String,
|
||||
parent_prompt_id: Option<String>,
|
||||
parent_session_id: String,
|
||||
owner: SubagentOwner,
|
||||
error: &str,
|
||||
surface_completion: bool,
|
||||
) {
|
||||
|
|
@ -270,6 +265,7 @@ impl SubagentCoordinator {
|
|||
description,
|
||||
parent_prompt_id,
|
||||
parent_session_id,
|
||||
owner,
|
||||
persona: None,
|
||||
started_at: std::time::Instant::now(),
|
||||
error,
|
||||
|
|
@ -281,6 +277,7 @@ impl SubagentCoordinator {
|
|||
/// Clears any stale pending entry for the same id.
|
||||
fn record_failure_completion(&mut self, c: FailureCompletion<'_>) {
|
||||
self.pending.remove(&c.subagent_id);
|
||||
self.loop_owned.remove(&c.subagent_id);
|
||||
self.sync_running_gauge();
|
||||
let FailureCompletion {
|
||||
subagent_id,
|
||||
|
|
@ -288,6 +285,7 @@ impl SubagentCoordinator {
|
|||
description,
|
||||
parent_prompt_id,
|
||||
parent_session_id,
|
||||
owner,
|
||||
persona,
|
||||
started_at,
|
||||
error,
|
||||
|
|
@ -309,6 +307,7 @@ impl SubagentCoordinator {
|
|||
subagent_id: subagent_id.clone(),
|
||||
parent_session_id,
|
||||
parent_prompt_id,
|
||||
owner,
|
||||
child_session_id: String::new(),
|
||||
description: description.clone(),
|
||||
subagent_type: subagent_type.clone(),
|
||||
|
|
@ -359,6 +358,7 @@ impl SubagentCoordinator {
|
|||
persisted_output_dir: Option<PathBuf>,
|
||||
) -> Option<SubagentTracker> {
|
||||
let tracker = self.active.remove(id);
|
||||
self.loop_owned.remove(id);
|
||||
self.sync_running_gauge();
|
||||
let started_at = tracker
|
||||
.as_ref()
|
||||
|
|
@ -373,6 +373,7 @@ impl SubagentCoordinator {
|
|||
.map(|t| t.child_session_id.0.to_string())
|
||||
.unwrap_or_default();
|
||||
let parent_prompt_id = tracker.as_ref().and_then(|t| t.parent_prompt_id.clone());
|
||||
let owner = tracker.as_ref().map(|t| t.owner.clone()).unwrap_or_default();
|
||||
let persona = tracker.as_ref().and_then(|t| t.persona.clone());
|
||||
let child_cwd = tracker
|
||||
.as_ref()
|
||||
|
|
@ -394,6 +395,7 @@ impl SubagentCoordinator {
|
|||
subagent_id: id.to_string(),
|
||||
parent_session_id,
|
||||
parent_prompt_id,
|
||||
owner,
|
||||
child_session_id,
|
||||
description,
|
||||
subagent_type,
|
||||
|
|
@ -481,6 +483,31 @@ impl SubagentCoordinator {
|
|||
}
|
||||
}
|
||||
}
|
||||
pub fn cancel_workflow_children(&mut self, run_id: &str) -> usize {
|
||||
for tracker in self.active.values() {
|
||||
if tracker.owner.workflow_run_id() == Some(run_id) {
|
||||
Self::cancel_tracker(tracker);
|
||||
}
|
||||
}
|
||||
for pending in self.pending.values() {
|
||||
if pending.owner.workflow_run_id() == Some(run_id) {
|
||||
pending.cancel_token.cancel();
|
||||
}
|
||||
}
|
||||
self.outstanding_for_workflow(run_id)
|
||||
}
|
||||
pub fn outstanding_for_workflow(&self, run_id: &str) -> usize {
|
||||
self
|
||||
.pending
|
||||
.values()
|
||||
.filter(|entry| entry.owner.workflow_run_id() == Some(run_id))
|
||||
.count()
|
||||
+ self
|
||||
.active
|
||||
.values()
|
||||
.filter(|entry| entry.owner.workflow_run_id() == Some(run_id))
|
||||
.count()
|
||||
}
|
||||
/// Attempt to cancel a subagent. Returns a typed outcome covering all cases:
|
||||
/// - Active → cancel it, return Cancelled
|
||||
/// - Pending (initializing) → fire its spawn token, return Cancelled
|
||||
|
|
|
|||
|
|
@ -16,7 +16,7 @@ use crate::terminal::AsyncTerminalRunner;
|
|||
use crate::tools::ToolContext;
|
||||
use crate::upload::trace::{
|
||||
GCS_SCHEMA_VERSION, PromptMetadata, SubagentSpawnedRef, TurnResultMetadata,
|
||||
local_sandbox_telemetry, upload_config, upload_metadata, upload_session_state,
|
||||
local_sandbox_telemetry, upload_metadata, upload_session_state,
|
||||
upload_subagent_metadata, upload_turn_result,
|
||||
};
|
||||
use crate::upload::turn::{PromptTraceContext, complete_prompt_trace};
|
||||
|
|
@ -37,6 +37,9 @@ impl SubagentCoordinator {
|
|||
/// - `None` — ID not found in active, completed, or pending maps.
|
||||
pub(crate) fn lookup(&self, id: &str) -> Option<SnapshotLookup> {
|
||||
if let Some(tracker) = self.active.get(id) {
|
||||
if tracker.owner.is_workflow() {
|
||||
return None;
|
||||
}
|
||||
return Some(
|
||||
SnapshotLookup::NeedsSignals(RunningSnapshotSeed {
|
||||
subagent_id: tracker.subagent_id.clone(),
|
||||
|
|
@ -50,6 +53,9 @@ impl SubagentCoordinator {
|
|||
);
|
||||
}
|
||||
if let Some(completed) = self.completed.get(id) {
|
||||
if completed.owner.is_workflow() {
|
||||
return None;
|
||||
}
|
||||
let status = if completed.result.cancelled {
|
||||
SubagentSnapshotStatus::Cancelled {
|
||||
reason: completed.result.error.clone(),
|
||||
|
|
@ -92,6 +98,9 @@ impl SubagentCoordinator {
|
|||
);
|
||||
}
|
||||
if let Some(pending) = self.pending.get(id) {
|
||||
if pending.owner.is_workflow() {
|
||||
return None;
|
||||
}
|
||||
return Some(
|
||||
SnapshotLookup::Ready(SubagentSnapshot {
|
||||
subagent_id: pending.subagent_id.clone(),
|
||||
|
|
@ -291,6 +300,27 @@ impl SubagentCoordinator {
|
|||
pub(crate) fn is_active_or_pending(&self, id: &str) -> bool {
|
||||
self.active.contains_key(id) || self.pending.contains_key(id)
|
||||
}
|
||||
pub(crate) fn record_loop_owner(&mut self, subagent_id: &str, task_id: &str) {
|
||||
self.loop_owned.insert(subagent_id.to_string(), task_id.to_string());
|
||||
}
|
||||
pub(crate) fn remove_loop_owner(&mut self, subagent_id: &str) {
|
||||
self.loop_owned.remove(subagent_id);
|
||||
}
|
||||
pub(crate) fn loop_task_id_of_child_session(
|
||||
&self,
|
||||
child_session_id: &str,
|
||||
) -> Option<String> {
|
||||
let subagent_id = self
|
||||
.active
|
||||
.values()
|
||||
.find(|t| t.child_session_id.0.as_ref() == child_session_id)?
|
||||
.subagent_id
|
||||
.clone();
|
||||
self.loop_owned.get(&subagent_id).cloned()
|
||||
}
|
||||
pub(crate) fn loop_unit_active(&self, task_id: &str) -> bool {
|
||||
self.loop_owned.values().any(|t| t == task_id)
|
||||
}
|
||||
/// The terminal `SubagentFinished` for an id the coordinator already holds in
|
||||
/// `completed`, else `None`. Lets orphan reconcile re-emit a subagent's real
|
||||
/// outcome when only its terminal meta write was lost (reconnect race: entry
|
||||
|
|
@ -337,7 +367,6 @@ impl SubagentCoordinator {
|
|||
}
|
||||
/// Snapshot all currently-running subagents for compaction state context.
|
||||
///
|
||||
/// Returns one `ActiveSubagentSummary` per entry in the `active` map.
|
||||
/// Completed/failed/cancelled subagents are NOT included — they live in
|
||||
/// the `completed` map and are irrelevant for post-compaction reminders
|
||||
/// (the model already saw their tool results before compaction).
|
||||
|
|
@ -347,7 +376,11 @@ impl SubagentCoordinator {
|
|||
/// compaction since it happens once and the reminder is static.
|
||||
#[cfg(test)]
|
||||
pub fn active_summaries(&self) -> Vec<ActiveSubagentSummary> {
|
||||
self.active.values().map(tracker_to_summary).collect()
|
||||
self.active
|
||||
.values()
|
||||
.filter(|t| !t.owner.is_workflow())
|
||||
.map(tracker_to_summary)
|
||||
.collect()
|
||||
}
|
||||
pub fn active_summaries_for(
|
||||
&self,
|
||||
|
|
@ -355,11 +388,12 @@ impl SubagentCoordinator {
|
|||
) -> Vec<ActiveSubagentSummary> {
|
||||
self.active
|
||||
.values()
|
||||
.filter(|t| t.parent_session_id == parent_session_id)
|
||||
.filter(|t| {
|
||||
t.parent_session_id == parent_session_id && !t.owner.is_workflow()
|
||||
})
|
||||
.map(tracker_to_summary)
|
||||
.collect()
|
||||
}
|
||||
/// Return seeds for all running subagents belonging to `parent_session_id`.
|
||||
///
|
||||
/// Each seed carries copied identity metadata plus a cloned
|
||||
/// `SessionSignalsHandle` so the caller can resolve live progress
|
||||
|
|
@ -375,7 +409,9 @@ impl SubagentCoordinator {
|
|||
) -> Vec<RunningSubagentListSeed> {
|
||||
self.active
|
||||
.values()
|
||||
.filter(|t| t.parent_session_id == parent_session_id)
|
||||
.filter(|t| {
|
||||
t.parent_session_id == parent_session_id && !t.owner.is_workflow()
|
||||
})
|
||||
.map(|t| RunningSubagentListSeed {
|
||||
subagent_id: t.subagent_id.clone(),
|
||||
parent_session_id: t.parent_session_id.clone(),
|
||||
|
|
|
|||
|
|
@ -16,7 +16,7 @@ use crate::terminal::AsyncTerminalRunner;
|
|||
use crate::tools::ToolContext;
|
||||
use crate::upload::trace::{
|
||||
GCS_SCHEMA_VERSION, PromptMetadata, SubagentSpawnedRef, TurnResultMetadata,
|
||||
local_sandbox_telemetry, upload_config, upload_metadata, upload_session_state,
|
||||
local_sandbox_telemetry, upload_metadata, upload_session_state,
|
||||
upload_subagent_metadata, upload_turn_result,
|
||||
};
|
||||
use crate::upload::turn::{PromptTraceContext, complete_prompt_trace};
|
||||
|
|
@ -43,6 +43,17 @@ pub(super) fn strip_task_tools_at_max_depth(
|
|||
prune_orphaned_background_task_tools(tool_config);
|
||||
stripped
|
||||
}
|
||||
pub(super) fn canonical_total_tokens(totals: &xai_chat_state::UsageTotals) -> u64 {
|
||||
totals.total_tokens()
|
||||
}
|
||||
pub(super) fn usage_is_incomplete(
|
||||
ledger_incomplete: bool,
|
||||
cancellation_may_hide_usage: bool,
|
||||
_known_total_tokens: u64,
|
||||
_has_usage_entries: bool,
|
||||
) -> bool {
|
||||
ledger_incomplete || cancellation_may_hide_usage
|
||||
}
|
||||
pub(super) fn task_model_override_error(
|
||||
requested: Option<&str>,
|
||||
provenance: ModelOverrideProvenance,
|
||||
|
|
@ -81,10 +92,15 @@ pub(crate) async fn handle_subagent_request(
|
|||
gateway: &GatewaySender,
|
||||
) {
|
||||
let start = std::time::Instant::now();
|
||||
let mut parent_wait_guard = (!request.run_in_background)
|
||||
let mut parent_wait_guard = subagent_blocks_parent_turn(&request)
|
||||
.then(|| crate::tools::tool_context::BlockingWaitGuard::enter(
|
||||
ctx.parent_blocking_wait_depth.clone(),
|
||||
));
|
||||
if request.owner.is_workflow() && request.cancel_token.is_cancelled() {
|
||||
parent_wait_guard.take();
|
||||
send_pre_spawn_cancelled(request, "Subagent was cancelled");
|
||||
return;
|
||||
}
|
||||
let Some(mut definition) = resolve_agent_definition(&request.subagent_type, &ctx)
|
||||
else {
|
||||
let msg = format!("Unknown subagent type: {}", request.subagent_type);
|
||||
|
|
@ -112,7 +128,7 @@ pub(crate) async fn handle_subagent_request(
|
|||
}
|
||||
let run_in_background = request.run_in_background
|
||||
|| definition.background.unwrap_or(false);
|
||||
let cancel_token = CancellationToken::new();
|
||||
let cancel_token = request.cancel_token.clone();
|
||||
coordinator
|
||||
.borrow_mut()
|
||||
.insert_pending(PendingSubagent {
|
||||
|
|
@ -122,6 +138,7 @@ pub(crate) async fn handle_subagent_request(
|
|||
persona: request.runtime_overrides.persona.clone(),
|
||||
parent_prompt_id: request.parent_prompt_id.clone(),
|
||||
parent_session_id: ctx.parent_session_id.clone(),
|
||||
owner: request.owner.clone(),
|
||||
started_at: start,
|
||||
run_in_background,
|
||||
surface_completion: request.surface_completion,
|
||||
|
|
@ -411,6 +428,10 @@ pub(crate) async fn handle_subagent_request(
|
|||
"Resolved runtime overrides for subagent"
|
||||
);
|
||||
}
|
||||
effective_runtime.capability_mode = xai_grok_subagent_resolution::intersect_capability_modes(
|
||||
effective_runtime.capability_mode,
|
||||
definition.capability_mode,
|
||||
);
|
||||
if let Some(mode) = effective_runtime.capability_mode {
|
||||
mode.filter_tool_config(&mut definition.tool_config);
|
||||
tracing::info!(
|
||||
|
|
@ -429,6 +450,17 @@ pub(crate) async fn handle_subagent_request(
|
|||
"Stripped task tool from child at max depth"
|
||||
);
|
||||
}
|
||||
if request.owner.is_workflow() {
|
||||
definition
|
||||
.tool_config
|
||||
.tools
|
||||
.retain(|tool| {
|
||||
!matches!(
|
||||
tool.id.rsplit(':').next(), Some("scheduler_create" |
|
||||
"scheduler_list" | "scheduler_delete")
|
||||
)
|
||||
});
|
||||
}
|
||||
if request.fork_context {
|
||||
effective_runtime.model = Some(ctx.model_id.0.to_string());
|
||||
}
|
||||
|
|
@ -661,6 +693,7 @@ pub(crate) async fn handle_subagent_request(
|
|||
role: effective_runtime.role_name.clone(),
|
||||
model: Some(effective_model_id.0.to_string()),
|
||||
resumed_from: request.resume_from.clone(),
|
||||
workflow_run_id: request.owner.workflow_run_id().map(str::to_string),
|
||||
},
|
||||
ctx.parent_cmd_tx.as_ref(),
|
||||
);
|
||||
|
|
@ -759,6 +792,12 @@ pub(crate) async fn handle_subagent_request(
|
|||
)
|
||||
.with_hunk_tracking_enabled(ctx.hunk_tracking_enabled);
|
||||
tool_ctx.subagent_event_tx = Some(ctx.subagent_event_tx.clone());
|
||||
let task_output_budget = request
|
||||
.runtime_overrides
|
||||
.output_token_budget
|
||||
.map(crate::tools::tool_context::TaskOutputTokenBudget::limited);
|
||||
tool_ctx.task_output_token_budget = task_output_budget.clone();
|
||||
tool_ctx.sampler_retry_only_before_output = task_output_budget.is_some();
|
||||
tool_ctx.monitor_event_buffer = Some(MonitorEventBuffer::default());
|
||||
tool_ctx.subagent_depth = child_depth;
|
||||
tool_ctx.lsp = ctx.lsp.clone();
|
||||
|
|
@ -859,7 +898,7 @@ pub(crate) async fn handle_subagent_request(
|
|||
memory_dir.display()
|
||||
);
|
||||
definition.prompt_body = Some(
|
||||
definition.prompt_body.unwrap_or_default() + &injection,
|
||||
definition.prompt_body.unwrap_or_default() + injection.as_str(),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -1051,7 +1090,7 @@ pub(crate) async fn handle_subagent_request(
|
|||
agent_name: Some(definition.name.clone()),
|
||||
reasoning_effort: Some(effective_sampling_config.reasoning_effort),
|
||||
});
|
||||
let forked_tool_override = if verbatim_mirror_fork {
|
||||
let forked_tool_override = if verbatim_mirror_fork && !request.owner.is_workflow() {
|
||||
ctx.parent_tool_snapshot.clone()
|
||||
} else {
|
||||
None
|
||||
|
|
@ -1124,6 +1163,7 @@ pub(crate) async fn handle_subagent_request(
|
|||
None,
|
||||
None,
|
||||
None,
|
||||
Vec::new(),
|
||||
None,
|
||||
if verbatim_mirror_fork {
|
||||
None
|
||||
|
|
@ -1164,6 +1204,7 @@ pub(crate) async fn handle_subagent_request(
|
|||
ctx.app_builder_deployer_config.clone(),
|
||||
ctx.write_file_enabled,
|
||||
ctx.goal_enabled,
|
||||
ctx.background_workflows_enabled,
|
||||
true,
|
||||
ctx.ask_user_question_enabled,
|
||||
ctx.client_hooks.clone(),
|
||||
|
|
@ -1192,7 +1233,11 @@ pub(crate) async fn handle_subagent_request(
|
|||
std::mem::take(&mut ctx.remote_settings),
|
||||
std::mem::take(&mut ctx.laziness_debug_log),
|
||||
ctx.parent_terminal_backend.clone(),
|
||||
ctx.parent_scheduler_handle.clone(),
|
||||
if request.owner.is_workflow() {
|
||||
None
|
||||
} else {
|
||||
ctx.parent_scheduler_handle.clone()
|
||||
},
|
||||
subagent_max_turns,
|
||||
forked_tool_override,
|
||||
)
|
||||
|
|
@ -1245,6 +1290,7 @@ pub(crate) async fn handle_subagent_request(
|
|||
subagent_id: request.id.clone(),
|
||||
parent_session_id: ctx.parent_session_id.clone(),
|
||||
parent_prompt_id: request.parent_prompt_id.clone(),
|
||||
owner: request.owner.clone(),
|
||||
child_session_id: child_session_id.clone(),
|
||||
subagent_type: request.subagent_type.clone(),
|
||||
persona: effective_runtime.persona.clone(),
|
||||
|
|
@ -1316,7 +1362,7 @@ pub(crate) async fn handle_subagent_request(
|
|||
screen_mode: None,
|
||||
verbatim: true,
|
||||
traceparent: xai_file_utils::trace_context::current_traceparent(),
|
||||
json_schema: None,
|
||||
json_schema: request.runtime_overrides.output_schema.clone(),
|
||||
send_now: false,
|
||||
admission: None,
|
||||
respond_to: prompt_tx,
|
||||
|
|
@ -1347,10 +1393,16 @@ pub(crate) async fn handle_subagent_request(
|
|||
None => std::future::pending::<()>().await,
|
||||
}
|
||||
};
|
||||
let budget = async {
|
||||
if request.await_to_completion {
|
||||
std::future::pending::<()>().await
|
||||
} else {
|
||||
tokio::time::sleep(subagent_await_budget()).await
|
||||
}
|
||||
};
|
||||
tokio::select! {
|
||||
biased; outcome = & mut fut => ForegroundWait::Done(outcome), _ =
|
||||
parent_await_dropped => ForegroundWait::ParentGone, _ =
|
||||
tokio::time::sleep(subagent_await_budget()) =>
|
||||
parent_await_dropped => ForegroundWait::ParentGone, _ = budget =>
|
||||
ForegroundWait::Budget,
|
||||
}
|
||||
};
|
||||
|
|
@ -1363,13 +1415,22 @@ pub(crate) async fn handle_subagent_request(
|
|||
}
|
||||
ForegroundWait::ParentGone => {
|
||||
parent_wait_guard.take();
|
||||
tracing::info!(
|
||||
subagent_id = % request.id,
|
||||
"foreground subagent await abandoned by its parent turn; detaching child to background (child keeps running)",
|
||||
);
|
||||
if !cancel_token.is_cancelled() {
|
||||
request.run_in_background = true;
|
||||
coordinator.borrow_mut().mark_backgrounded(&request.id);
|
||||
if request.owner.is_workflow() {
|
||||
tracing::info!(
|
||||
subagent_id = % request.id, workflow_run_id = ? request.owner
|
||||
.workflow_run_id(),
|
||||
"workflow subagent result receiver dropped; cancelling child",
|
||||
);
|
||||
cancel_token.cancel();
|
||||
} else {
|
||||
tracing::info!(
|
||||
subagent_id = % request.id,
|
||||
"foreground subagent await abandoned by its parent turn; detaching child to background (child keeps running)",
|
||||
);
|
||||
if !cancel_token.is_cancelled() {
|
||||
request.run_in_background = true;
|
||||
coordinator.borrow_mut().mark_backgrounded(&request.id);
|
||||
}
|
||||
}
|
||||
fut.await
|
||||
}
|
||||
|
|
@ -1400,9 +1461,11 @@ pub(crate) async fn handle_subagent_request(
|
|||
};
|
||||
let duration_ms = start.elapsed().as_millis() as u64;
|
||||
let mut turn_token_totals: Option<(u64, u64, u64)> = None;
|
||||
let mut cancellation_may_hide_usage = false;
|
||||
let mut result = match wait_outcome {
|
||||
SubagentWaitOutcome::Cancelled => {
|
||||
let (tool_calls, turns) = signals_snapshot_counts(&child_handle).await;
|
||||
cancellation_may_hide_usage = turns > 0 || tool_calls > 0;
|
||||
SubagentResult {
|
||||
success: false,
|
||||
cancelled: true,
|
||||
|
|
@ -1456,6 +1519,7 @@ pub(crate) async fn handle_subagent_request(
|
|||
},
|
||||
),
|
||||
) => {
|
||||
cancellation_may_hide_usage = true;
|
||||
let reason = cancellation_error_message(category, context.as_ref());
|
||||
SubagentResult {
|
||||
success: false,
|
||||
|
|
@ -1478,6 +1542,9 @@ pub(crate) async fn handle_subagent_request(
|
|||
turns,
|
||||
duration_ms,
|
||||
tokens_used: result_tokens,
|
||||
output_tokens_used: 0,
|
||||
output_usage_incomplete: true,
|
||||
total_tokens_used: 0,
|
||||
worktree_path: worktree_path
|
||||
.as_ref()
|
||||
.map(|p| p.to_string_lossy().to_string()),
|
||||
|
|
@ -1515,32 +1582,76 @@ pub(crate) async fn handle_subagent_request(
|
|||
turns,
|
||||
duration_ms,
|
||||
tokens_used: result_tokens,
|
||||
output_tokens_used: 0,
|
||||
output_usage_incomplete: true,
|
||||
total_tokens_used: 0,
|
||||
worktree_path: worktree_path
|
||||
.as_ref()
|
||||
.map(|p| p.to_string_lossy().to_string()),
|
||||
backgrounded: false,
|
||||
}
|
||||
}
|
||||
Ok(Ok(_)) => {
|
||||
SubagentResult {
|
||||
success: true,
|
||||
output: if final_text.is_empty() {
|
||||
std::sync::Arc::from(
|
||||
format!(
|
||||
"Subagent '{}' ({}) completed successfully. {} tool calls, {} turns.",
|
||||
request.description, request.subagent_type, tool_calls,
|
||||
turns
|
||||
),
|
||||
Ok(
|
||||
Ok(crate::session::commands::PromptTurnOk { structured_output, .. }),
|
||||
) => {
|
||||
let wanted_schema = request
|
||||
.runtime_overrides
|
||||
.output_schema
|
||||
.is_some();
|
||||
let (success, error, output) = match (
|
||||
wanted_schema,
|
||||
structured_output,
|
||||
) {
|
||||
(true, Some(Ok(value))) => {
|
||||
(true, None, std::sync::Arc::from(value.to_string()))
|
||||
}
|
||||
(true, Some(Err(e))) => {
|
||||
(
|
||||
false,
|
||||
Some(format!("structured output validation failed: {e}")),
|
||||
std::sync::Arc::from(final_text),
|
||||
)
|
||||
} else {
|
||||
std::sync::Arc::from(final_text)
|
||||
},
|
||||
}
|
||||
(true, None) => {
|
||||
(
|
||||
false,
|
||||
Some(
|
||||
"structured output requested but none produced".to_string(),
|
||||
),
|
||||
std::sync::Arc::from(final_text),
|
||||
)
|
||||
}
|
||||
(false, _) => {
|
||||
(
|
||||
true,
|
||||
None,
|
||||
if final_text.is_empty() {
|
||||
std::sync::Arc::from(
|
||||
format!(
|
||||
"Subagent '{}' ({}) completed successfully. {} tool calls, {} turns.",
|
||||
request.description, request.subagent_type, tool_calls,
|
||||
turns
|
||||
),
|
||||
)
|
||||
} else {
|
||||
std::sync::Arc::from(final_text)
|
||||
},
|
||||
)
|
||||
}
|
||||
};
|
||||
SubagentResult {
|
||||
success,
|
||||
error,
|
||||
output,
|
||||
subagent_id: request.id.clone(),
|
||||
child_session_id: child_session_id.0.to_string(),
|
||||
tool_calls,
|
||||
turns,
|
||||
duration_ms,
|
||||
tokens_used: result_tokens,
|
||||
output_tokens_used: 0,
|
||||
output_usage_incomplete: true,
|
||||
total_tokens_used: 0,
|
||||
worktree_path: worktree_path
|
||||
.as_ref()
|
||||
.map(|p| p.to_string_lossy().to_string()),
|
||||
|
|
@ -1548,6 +1659,7 @@ pub(crate) async fn handle_subagent_request(
|
|||
}
|
||||
}
|
||||
Ok(Err(e)) => {
|
||||
cancellation_may_hide_usage = was_cancelled;
|
||||
SubagentResult {
|
||||
success: false,
|
||||
cancelled: was_cancelled,
|
||||
|
|
@ -1570,6 +1682,7 @@ pub(crate) async fn handle_subagent_request(
|
|||
}
|
||||
}
|
||||
Err(_) => {
|
||||
cancellation_may_hide_usage = was_cancelled;
|
||||
SubagentResult {
|
||||
success: false,
|
||||
cancelled: was_cancelled,
|
||||
|
|
@ -1719,10 +1832,6 @@ pub(crate) async fn handle_subagent_request(
|
|||
sandbox: local_sandbox_telemetry(),
|
||||
};
|
||||
upload_metadata(&trace_ctx, metadata).await;
|
||||
if let Some(ref agent_config) = ctx.agent_config {
|
||||
upload_config(&trace_ctx, agent_config).await;
|
||||
}
|
||||
crate::upload::config_files::upload_config_files(&trace_ctx).await;
|
||||
let resolved_model = child_handle
|
||||
.get_model_metadata()
|
||||
.await
|
||||
|
|
@ -1789,14 +1898,42 @@ pub(crate) async fn handle_subagent_request(
|
|||
} else {
|
||||
0
|
||||
};
|
||||
let (subagent_usage_by_model, subagent_usage_incomplete) = match child_handle
|
||||
.chat_state_handle
|
||||
.try_get_session_usage()
|
||||
.await
|
||||
{
|
||||
Ok(u) => (Some(u.by_model.into_iter().collect::<Vec<_>>()), u.incomplete),
|
||||
Err(()) => (None, true),
|
||||
let task_budget_usage = task_output_budget.as_ref().map(|budget| budget.usage());
|
||||
let (
|
||||
subagent_usage_by_model,
|
||||
subagent_usage_incomplete,
|
||||
output_tokens_used,
|
||||
total_tokens_used,
|
||||
) = match child_handle.chat_state_handle.try_get_session_usage().await {
|
||||
Ok(u) => {
|
||||
let output_tokens = u.totals.output_tokens;
|
||||
let total_tokens = canonical_total_tokens(&u.totals);
|
||||
let has_usage_entries = !u.by_model.is_empty();
|
||||
let usage_incomplete = usage_is_incomplete(
|
||||
u.incomplete,
|
||||
cancellation_may_hide_usage,
|
||||
total_tokens,
|
||||
has_usage_entries,
|
||||
);
|
||||
(
|
||||
Some(u.by_model.into_iter().collect::<Vec<_>>()),
|
||||
usage_incomplete,
|
||||
(!usage_incomplete).then_some(output_tokens),
|
||||
Some(total_tokens),
|
||||
)
|
||||
}
|
||||
Err(()) => (None, true, None, None),
|
||||
};
|
||||
result.total_tokens_used = total_tokens_used.unwrap_or(0);
|
||||
if let Some((task_spent, task_incomplete)) = task_budget_usage {
|
||||
result.output_tokens_used = output_tokens_used.unwrap_or(task_spent);
|
||||
result.output_usage_incomplete = task_incomplete || subagent_usage_incomplete
|
||||
|| output_tokens_used.is_none();
|
||||
} else {
|
||||
result.output_tokens_used = output_tokens_used.unwrap_or(0);
|
||||
result.output_usage_incomplete = subagent_usage_incomplete
|
||||
|| output_tokens_used.is_none();
|
||||
}
|
||||
let fold_acked = match subagent_usage_by_model {
|
||||
None => false,
|
||||
Some(ref by_model) if by_model.is_empty() && !subagent_usage_incomplete => true,
|
||||
|
|
|
|||
|
|
@ -22,8 +22,8 @@ use crate::terminal::AsyncTerminalRunner;
|
|||
use crate::tools::ToolContext;
|
||||
use crate::upload::trace::{
|
||||
GCS_SCHEMA_VERSION, PromptMetadata, SubagentSpawnedRef, TurnResultMetadata,
|
||||
local_sandbox_telemetry, upload_config, upload_metadata, upload_session_state,
|
||||
upload_subagent_metadata, upload_turn_result,
|
||||
local_sandbox_telemetry, upload_metadata, upload_session_state, upload_subagent_metadata,
|
||||
upload_turn_result,
|
||||
};
|
||||
use crate::upload::turn::{PromptTraceContext, complete_prompt_trace};
|
||||
use agent_client_protocol as acp;
|
||||
|
|
@ -59,6 +59,7 @@ pub(crate) struct SubagentTracker {
|
|||
pub subagent_id: String,
|
||||
pub parent_session_id: String,
|
||||
pub parent_prompt_id: Option<String>,
|
||||
pub owner: SubagentOwner,
|
||||
pub child_session_id: acp::SessionId,
|
||||
pub subagent_type: String,
|
||||
pub persona: Option<String>,
|
||||
|
|
@ -227,6 +228,7 @@ pub(crate) struct SubagentSpawnContext {
|
|||
pub write_file_enabled: bool,
|
||||
/// Whether goal mode (`/goal`) is enabled.
|
||||
pub goal_enabled: bool,
|
||||
pub background_workflows_enabled: bool,
|
||||
/// Whether the `ask_user_question` tool is exposed to this subagent,
|
||||
/// inherited from the parent session (see `build_subagent_spawn_context`).
|
||||
pub ask_user_question_enabled: bool,
|
||||
|
|
@ -363,7 +365,7 @@ pub(crate) struct SubagentSpawnContext {
|
|||
/// Parent's `blocking_wait_depth` (same `Arc`). A foreground spawn holds a
|
||||
/// `BlockingWaitGuard` on it for the blocking await so `queue_input` routes
|
||||
/// a prompt sent during the wait onto send-now; never for background spawns.
|
||||
pub parent_blocking_wait_depth: Arc<std::sync::atomic::AtomicUsize>,
|
||||
pub parent_blocking_wait_depth: Arc<crate::tools::tool_context::BlockingWaitState>,
|
||||
}
|
||||
impl SubagentSpawnContext {
|
||||
/// Check if a subagent is enabled via the toggle config.
|
||||
|
|
@ -478,6 +480,7 @@ pub(crate) struct CompletedSubagent {
|
|||
pub subagent_id: String,
|
||||
pub parent_session_id: String,
|
||||
pub parent_prompt_id: Option<String>,
|
||||
pub owner: SubagentOwner,
|
||||
pub child_session_id: String,
|
||||
pub description: String,
|
||||
pub subagent_type: String,
|
||||
|
|
@ -539,6 +542,7 @@ pub(crate) struct PendingSubagent {
|
|||
pub persona: Option<String>,
|
||||
pub parent_prompt_id: Option<String>,
|
||||
pub parent_session_id: String,
|
||||
pub owner: SubagentOwner,
|
||||
pub started_at: std::time::Instant,
|
||||
pub run_in_background: bool,
|
||||
/// Mirrors `SubagentRequest::surface_completion`.
|
||||
|
|
@ -559,6 +563,7 @@ struct FailureCompletion<'a> {
|
|||
description: String,
|
||||
parent_prompt_id: Option<String>,
|
||||
parent_session_id: String,
|
||||
owner: SubagentOwner,
|
||||
persona: Option<String>,
|
||||
started_at: std::time::Instant,
|
||||
error: &'a str,
|
||||
|
|
@ -618,6 +623,7 @@ pub(crate) struct SubagentCoordinator {
|
|||
/// marks ledgers by itself (a true apply-miss marks them at fold time).
|
||||
/// Cleared on freeze/cancel. See AGENTS.md rule 3 for the completeness model.
|
||||
subagent_usage_not_applied_prompts: std::collections::HashSet<String>,
|
||||
loop_owned: HashMap<String, String>,
|
||||
}
|
||||
/// Cap on the completed map (entries are small: identity, counts, and an
|
||||
/// error string; successful output text lives in `output.json`).
|
||||
|
|
@ -768,6 +774,9 @@ pub(crate) struct SubagentProvenance {
|
|||
/// ID of the source subagent this session was resumed from.
|
||||
pub(crate) resumed_from: Option<String>,
|
||||
}
|
||||
fn subagent_blocks_parent_turn(request: &SubagentRequest) -> bool {
|
||||
!request.run_in_background && !request.owner.is_workflow()
|
||||
}
|
||||
/// Convert a `std::time::Instant` to approximate epoch milliseconds.
|
||||
///
|
||||
/// `Instant` has no absolute epoch, so we compute the offset from
|
||||
|
|
@ -2110,6 +2119,15 @@ pub(crate) fn send_failure(request: SubagentRequest, error: &str) {
|
|||
..Default::default()
|
||||
});
|
||||
}
|
||||
fn send_pre_spawn_cancelled(request: SubagentRequest, error: &str) {
|
||||
let _ = request.result_tx.send(SubagentResult {
|
||||
success: false,
|
||||
cancelled: true,
|
||||
error: Some(error.to_string()),
|
||||
subagent_id: request.id,
|
||||
..Default::default()
|
||||
});
|
||||
}
|
||||
/// Fail BEFORE `insert_pending`. Sends via oneshot; for background-mode
|
||||
/// requests also records a synthetic `CompletedSubagent` + emits a
|
||||
/// `SubagentFinished` notification (persisted + live).
|
||||
|
|
@ -2125,6 +2143,7 @@ fn send_pre_spawn_failure(
|
|||
subagent_type,
|
||||
description,
|
||||
parent_prompt_id,
|
||||
owner,
|
||||
result_tx,
|
||||
run_in_background,
|
||||
surface_completion,
|
||||
|
|
@ -2138,6 +2157,7 @@ fn send_pre_spawn_failure(
|
|||
description,
|
||||
parent_prompt_id,
|
||||
ctx.parent_session_id.clone(),
|
||||
owner,
|
||||
error,
|
||||
surface_completion,
|
||||
);
|
||||
|
|
|
|||
|
|
@ -1,8 +1,26 @@
|
|||
#![cfg_attr(rustfmt, rustfmt::skip)]
|
||||
use super::*;
|
||||
use super::handle_request::{canonical_total_tokens, usage_is_incomplete};
|
||||
use crate::test_support::lsp_runtime::{
|
||||
DummyLspDispatch, ctx_with_toggle, make_request, test_gateway,
|
||||
};
|
||||
#[test]
|
||||
fn canonical_total_tokens_does_not_double_count_reasoning() {
|
||||
let totals = xai_chat_state::UsageTotals {
|
||||
input_tokens: 100,
|
||||
output_tokens: 40,
|
||||
reasoning_tokens: 25,
|
||||
..Default::default()
|
||||
};
|
||||
assert_eq!(canonical_total_tokens(& totals), 140);
|
||||
}
|
||||
#[test]
|
||||
fn cancellation_makes_an_otherwise_complete_usage_snapshot_incomplete() {
|
||||
assert!(usage_is_incomplete(false, true, 0, false));
|
||||
assert!(usage_is_incomplete(false, true, 10, false));
|
||||
assert!(! usage_is_incomplete(false, false, 0, false));
|
||||
assert!(usage_is_incomplete(true, false, 0, false));
|
||||
}
|
||||
/// Invariant: resolving a subagent applies the parent session's
|
||||
/// `--tools`/`--disallowed-tools`/`--permission-mode` — driven through
|
||||
/// `resolve_agent_definition` so the spawn path can't skip them.
|
||||
|
|
@ -336,6 +354,7 @@ fn lookup_returns_initializing_for_pending_subagent() {
|
|||
persona: None,
|
||||
parent_prompt_id: None,
|
||||
parent_session_id: String::new(),
|
||||
owner: SubagentOwner::Task,
|
||||
started_at: std::time::Instant::now(),
|
||||
run_in_background: false,
|
||||
surface_completion: true,
|
||||
|
|
@ -367,6 +386,7 @@ async fn running_gauge_tracks_pending_and_active() {
|
|||
persona: None,
|
||||
parent_prompt_id: None,
|
||||
parent_session_id: String::new(),
|
||||
owner: SubagentOwner::Task,
|
||||
started_at: std::time::Instant::now(),
|
||||
run_in_background: true,
|
||||
surface_completion: true,
|
||||
|
|
@ -396,6 +416,7 @@ async fn running_gauge_tracks_pending_and_active() {
|
|||
persona: None,
|
||||
parent_prompt_id: None,
|
||||
parent_session_id: String::new(),
|
||||
owner: SubagentOwner::Task,
|
||||
started_at: std::time::Instant::now(),
|
||||
run_in_background: true,
|
||||
surface_completion: true,
|
||||
|
|
@ -414,6 +435,7 @@ async fn running_gauge_tracks_pending_and_active() {
|
|||
persona: None,
|
||||
parent_prompt_id: None,
|
||||
parent_session_id: String::new(),
|
||||
owner: SubagentOwner::Task,
|
||||
started_at: std::time::Instant::now(),
|
||||
run_in_background: true,
|
||||
surface_completion: true,
|
||||
|
|
@ -581,7 +603,10 @@ fn auto_wake_test_request(id: &str) -> SubagentRequest {
|
|||
runtime_overrides: Default::default(),
|
||||
run_in_background: true,
|
||||
surface_completion: true,
|
||||
await_to_completion: false,
|
||||
fork_context: false,
|
||||
owner: SubagentOwner::Task,
|
||||
cancel_token: CancellationToken::new(),
|
||||
result_tx,
|
||||
}
|
||||
}
|
||||
|
|
@ -738,6 +763,7 @@ fn fail_pending(coordinator: &mut SubagentCoordinator, id: &str, surface: bool)
|
|||
persona: None,
|
||||
parent_prompt_id: None,
|
||||
parent_session_id: String::new(),
|
||||
owner: SubagentOwner::Task,
|
||||
started_at: std::time::Instant::now(),
|
||||
run_in_background: false,
|
||||
surface_completion: surface,
|
||||
|
|
@ -786,6 +812,7 @@ fn remove_pending_clears_entry() {
|
|||
persona: None,
|
||||
parent_prompt_id: None,
|
||||
parent_session_id: String::new(),
|
||||
owner: SubagentOwner::Task,
|
||||
started_at: std::time::Instant::now(),
|
||||
run_in_background: false,
|
||||
surface_completion: true,
|
||||
|
|
@ -810,6 +837,7 @@ fn move_pending_to_failed_creates_completed_entry() {
|
|||
persona: None,
|
||||
parent_prompt_id: None,
|
||||
parent_session_id: String::new(),
|
||||
owner: SubagentOwner::Task,
|
||||
started_at: std::time::Instant::now(),
|
||||
run_in_background: true,
|
||||
surface_completion: true,
|
||||
|
|
@ -848,6 +876,7 @@ fn move_pending_to_failed_fires_completion_notify() {
|
|||
persona: None,
|
||||
parent_prompt_id: None,
|
||||
parent_session_id: String::new(),
|
||||
owner: SubagentOwner::Task,
|
||||
started_at: std::time::Instant::now(),
|
||||
run_in_background: true,
|
||||
surface_completion: true,
|
||||
|
|
@ -877,6 +906,7 @@ fn move_pending_to_cancelled_creates_cancelled_entry() {
|
|||
persona: None,
|
||||
parent_prompt_id: None,
|
||||
parent_session_id: String::new(),
|
||||
owner: SubagentOwner::Task,
|
||||
started_at: std::time::Instant::now(),
|
||||
run_in_background: true,
|
||||
surface_completion: true,
|
||||
|
|
@ -905,6 +935,7 @@ fn completed_with_output(
|
|||
CompletedSubagent {
|
||||
subagent_id: id.into(),
|
||||
parent_session_id: String::new(),
|
||||
owner: SubagentOwner::Task,
|
||||
parent_prompt_id: None,
|
||||
child_session_id: String::new(),
|
||||
description: "task".into(),
|
||||
|
|
@ -1076,6 +1107,7 @@ fn cancel_with_outcome_fires_pending_token() {
|
|||
persona: None,
|
||||
parent_prompt_id: None,
|
||||
parent_session_id: String::new(),
|
||||
owner: SubagentOwner::Task,
|
||||
started_at: std::time::Instant::now(),
|
||||
run_in_background: false,
|
||||
surface_completion: true,
|
||||
|
|
@ -1115,8 +1147,7 @@ async fn cancel_with_outcome_returns_variant_for_active_finished_unknown() {
|
|||
);
|
||||
assert!(
|
||||
matches!(coordinator.cancel_with_outcome("sub-done"),
|
||||
SubagentCancelOutcome::AlreadyFinished { status }
|
||||
if status == "completed")
|
||||
SubagentCancelOutcome::AlreadyFinished { status } if status == "completed")
|
||||
);
|
||||
assert!(
|
||||
matches!(coordinator.cancel_with_outcome("nonexistent"),
|
||||
|
|
@ -1136,6 +1167,7 @@ fn cancel_by_parent_prompt_id_fires_matching_pending_token() {
|
|||
persona: None,
|
||||
parent_prompt_id: Some("prompt-A".to_string()),
|
||||
parent_session_id: String::new(),
|
||||
owner: SubagentOwner::Task,
|
||||
started_at: std::time::Instant::now(),
|
||||
run_in_background: false,
|
||||
surface_completion: true,
|
||||
|
|
@ -1150,6 +1182,7 @@ fn cancel_by_parent_prompt_id_fires_matching_pending_token() {
|
|||
persona: None,
|
||||
parent_prompt_id: Some("prompt-B".to_string()),
|
||||
parent_session_id: String::new(),
|
||||
owner: SubagentOwner::Task,
|
||||
started_at: std::time::Instant::now(),
|
||||
run_in_background: false,
|
||||
surface_completion: true,
|
||||
|
|
@ -1176,6 +1209,7 @@ fn completed_takes_precedence_over_pending_in_lookup() {
|
|||
persona: None,
|
||||
parent_prompt_id: None,
|
||||
parent_session_id: String::new(),
|
||||
owner: SubagentOwner::Task,
|
||||
started_at: std::time::Instant::now(),
|
||||
run_in_background: false,
|
||||
surface_completion: true,
|
||||
|
|
@ -1294,6 +1328,7 @@ fn dummy_tracker(
|
|||
SubagentTracker {
|
||||
subagent_id: subagent_id.into(),
|
||||
parent_session_id: parent_session_id.into(),
|
||||
owner: SubagentOwner::Task,
|
||||
parent_prompt_id: None,
|
||||
child_session_id: acp::SessionId::new(subagent_id),
|
||||
subagent_type: subagent_type.into(),
|
||||
|
|
@ -1419,7 +1454,7 @@ fn explicit_override_takes_precedence_over_role() {
|
|||
);
|
||||
assert_eq!(resolved.model.as_deref(), Some("explicit-model"));
|
||||
assert_eq!(
|
||||
resolved.capability_mode, Some(xai_tool_types::SubagentCapabilityMode::All)
|
||||
resolved.capability_mode, Some(xai_tool_types::SubagentCapabilityMode::ReadOnly)
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
|
|
@ -1914,8 +1949,7 @@ fn verbatim_fork_keeps_items_byte_for_byte_when_small() {
|
|||
.any(|i| {
|
||||
matches!(
|
||||
i, ConversationItem::User(u) if u.content.iter().any(| p |
|
||||
matches!(p, ContentPart::Text { text }
|
||||
if text.contains(needle)))
|
||||
matches!(p, ContentPart::Text { text } if text.contains(needle)))
|
||||
)
|
||||
})
|
||||
};
|
||||
|
|
@ -1999,8 +2033,7 @@ fn verbatim_fork_falls_back_to_summary_when_oversize() {
|
|||
.any(|i| {
|
||||
matches!(
|
||||
i, ConversationItem::User(u) if u.content.iter().any(| p | matches!(p,
|
||||
ContentPart::Text { text }
|
||||
if text.contains("<background_context>")))
|
||||
ContentPart::Text { text } if text.contains("<background_context>")))
|
||||
)
|
||||
});
|
||||
assert!(has_blob, "oversize fallback must produce a background_context blob");
|
||||
|
|
@ -2079,7 +2112,10 @@ fn bootstrap_test_request(fork_context: bool) -> SubagentRequest {
|
|||
runtime_overrides: Default::default(),
|
||||
run_in_background: false,
|
||||
surface_completion: false,
|
||||
await_to_completion: false,
|
||||
fork_context,
|
||||
owner: SubagentOwner::Task,
|
||||
cancel_token: CancellationToken::new(),
|
||||
result_tx,
|
||||
}
|
||||
}
|
||||
|
|
@ -2940,6 +2976,7 @@ async fn cancel_pending_subagent_at_promote_emits_exactly_one_cancelled_finish()
|
|||
persona: None,
|
||||
parent_prompt_id: None,
|
||||
parent_session_id: ctx.parent_session_id.clone(),
|
||||
owner: SubagentOwner::Task,
|
||||
started_at: std::time::Instant::now(),
|
||||
run_in_background: true,
|
||||
surface_completion: true,
|
||||
|
|
@ -3041,6 +3078,7 @@ async fn run_promote_cancel_with_worktree(
|
|||
persona: None,
|
||||
parent_prompt_id: None,
|
||||
parent_session_id: ctx.parent_session_id.clone(),
|
||||
owner: SubagentOwner::Task,
|
||||
started_at: std::time::Instant::now(),
|
||||
run_in_background: true,
|
||||
surface_completion: true,
|
||||
|
|
@ -3156,6 +3194,7 @@ fn record_pre_spawn_failure_populates_completed_and_summary() {
|
|||
"bg job".to_string(),
|
||||
Some("prompt-1".to_string()),
|
||||
"parent-1".to_string(),
|
||||
SubagentOwner::Task,
|
||||
"Unknown subagent type: invented",
|
||||
true,
|
||||
);
|
||||
|
|
@ -3189,6 +3228,7 @@ fn record_pre_spawn_failure_skips_buffer_when_flag_false() {
|
|||
"bg job".to_string(),
|
||||
None,
|
||||
"parent-1".to_string(),
|
||||
SubagentOwner::Task,
|
||||
"Unknown subagent type: invented",
|
||||
false,
|
||||
);
|
||||
|
|
@ -3207,6 +3247,7 @@ async fn record_pre_spawn_failure_notifies_waiters() {
|
|||
"bg job".to_string(),
|
||||
None,
|
||||
"parent-1".to_string(),
|
||||
SubagentOwner::Task,
|
||||
"error",
|
||||
true,
|
||||
);
|
||||
|
|
@ -3227,6 +3268,7 @@ async fn record_pre_spawn_failure_notifies_all_waiters() {
|
|||
"bg job".to_string(),
|
||||
None,
|
||||
"parent-1".to_string(),
|
||||
SubagentOwner::Task,
|
||||
"error",
|
||||
true,
|
||||
);
|
||||
|
|
@ -3245,6 +3287,7 @@ fn record_pre_spawn_failure_clears_stale_pending_entry() {
|
|||
persona: None,
|
||||
parent_prompt_id: Some("prompt-X".to_string()),
|
||||
parent_session_id: "parent-1".to_string(),
|
||||
owner: SubagentOwner::Task,
|
||||
started_at: std::time::Instant::now(),
|
||||
run_in_background: true,
|
||||
surface_completion: true,
|
||||
|
|
@ -3259,6 +3302,7 @@ fn record_pre_spawn_failure_clears_stale_pending_entry() {
|
|||
"bg job".to_string(),
|
||||
Some("prompt-X".to_string()),
|
||||
"parent-1".to_string(),
|
||||
SubagentOwner::Task,
|
||||
"Unknown subagent type: invented",
|
||||
true,
|
||||
);
|
||||
|
|
|
|||
|
|
@ -327,6 +327,7 @@ fn resumable_source_returns_info_for_completed_subagent() {
|
|||
subagent_id: "sub-resume".into(),
|
||||
parent_session_id: "parent-1".into(),
|
||||
parent_prompt_id: Some("prompt-1".into()),
|
||||
owner: SubagentOwner::Task,
|
||||
child_session_id: "child-resume".into(),
|
||||
description: "resumable task".into(),
|
||||
subagent_type: "general-purpose".into(),
|
||||
|
|
@ -713,6 +714,35 @@ fn coordinator_with_completed(id: &str) -> SubagentCoordinator {
|
|||
);
|
||||
coordinator
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn loop_unit_active_tracks_and_prunes_owned_subagents() {
|
||||
let mut coordinator = SubagentCoordinator::new();
|
||||
coordinator
|
||||
.insert(
|
||||
dummy_tracker("iter-1", "root-sess", "general-purpose", "loop: watch ci"),
|
||||
);
|
||||
coordinator.record_loop_owner("iter-1", "task-42");
|
||||
assert!(coordinator.loop_unit_active("task-42"));
|
||||
assert!(! coordinator.loop_unit_active("other-task"));
|
||||
assert_eq!(
|
||||
coordinator.loop_task_id_of_child_session("iter-1"), Some("task-42".to_string())
|
||||
);
|
||||
assert_eq!(coordinator.loop_task_id_of_child_session("unknown"), None);
|
||||
coordinator
|
||||
.move_to_completed(
|
||||
"iter-1",
|
||||
"loop: watch ci".into(),
|
||||
"general-purpose".into(),
|
||||
SubagentResult {
|
||||
success: true,
|
||||
subagent_id: "iter-1".into(),
|
||||
child_session_id: "iter-1".into(),
|
||||
..Default::default()
|
||||
},
|
||||
None,
|
||||
);
|
||||
assert!(! coordinator.loop_unit_active("task-42"));
|
||||
}
|
||||
/// End-to-end glue: gate ON + a worktree present runs the completion
|
||||
/// sequence (snapshot → persist ref to meta.json AND in-memory → remove)
|
||||
/// and asserts all three post-conditions hold together.
|
||||
|
|
@ -1343,6 +1373,7 @@ fn resumable_source_rejects_cross_session_lookup() {
|
|||
subagent_id: "sub-other".into(),
|
||||
parent_session_id: "session-A".into(),
|
||||
parent_prompt_id: None,
|
||||
owner: SubagentOwner::Task,
|
||||
child_session_id: "child-other".into(),
|
||||
description: "other task".into(),
|
||||
subagent_type: "explore".into(),
|
||||
|
|
@ -1640,6 +1671,7 @@ fn reconcile_orphan_skips_pending_ids_in_live_registry() {
|
|||
persona: None,
|
||||
parent_prompt_id: None,
|
||||
parent_session_id: "parent-x".to_string(),
|
||||
owner: SubagentOwner::Task,
|
||||
started_at: std::time::Instant::now(),
|
||||
run_in_background: false,
|
||||
surface_completion: true,
|
||||
|
|
@ -2098,6 +2130,7 @@ fn notification_subagent_spawned_includes_resumed_from() {
|
|||
role: None,
|
||||
model: None,
|
||||
resumed_from: Some("prev-agent-id".into()),
|
||||
workflow_run_id: None,
|
||||
};
|
||||
let json = serde_json::to_value(¬ification).unwrap();
|
||||
assert_eq!(json["resumed_from"], "prev-agent-id");
|
||||
|
|
@ -2118,6 +2151,7 @@ fn notification_subagent_spawned_includes_resumed_from() {
|
|||
role: None,
|
||||
model: None,
|
||||
resumed_from: None,
|
||||
workflow_run_id: None,
|
||||
};
|
||||
let json = serde_json::to_value(&fresh).unwrap();
|
||||
assert!(json.get("resumed_from").is_none());
|
||||
|
|
@ -2162,6 +2196,7 @@ fn completed_subagent_propagates_resumed_from() {
|
|||
subagent_id: "sub-prov".into(),
|
||||
parent_session_id: "parent".into(),
|
||||
parent_prompt_id: Some("prompt-1".into()),
|
||||
owner: SubagentOwner::Task,
|
||||
child_session_id: "child-prov".into(),
|
||||
description: "provenance test".into(),
|
||||
subagent_type: "general-purpose".into(),
|
||||
|
|
@ -2299,6 +2334,7 @@ async fn outstanding_for_prompt_includes_pending_and_active() {
|
|||
persona: None,
|
||||
parent_prompt_id: Some("prompt-X".to_string()),
|
||||
parent_session_id: String::new(),
|
||||
owner: SubagentOwner::Task,
|
||||
started_at: std::time::Instant::now(),
|
||||
run_in_background: false,
|
||||
surface_completion: true,
|
||||
|
|
@ -2411,6 +2447,7 @@ fn outstanding_for_prompt_returns_sorted_ids() {
|
|||
persona: None,
|
||||
parent_prompt_id: Some("p".to_string()),
|
||||
parent_session_id: String::new(),
|
||||
owner: SubagentOwner::Task,
|
||||
started_at: std::time::Instant::now(),
|
||||
run_in_background: false,
|
||||
surface_completion: true,
|
||||
|
|
@ -2425,6 +2462,7 @@ fn outstanding_for_prompt_returns_sorted_ids() {
|
|||
persona: None,
|
||||
parent_prompt_id: Some("p".to_string()),
|
||||
parent_session_id: String::new(),
|
||||
owner: SubagentOwner::Task,
|
||||
started_at: std::time::Instant::now(),
|
||||
run_in_background: false,
|
||||
surface_completion: true,
|
||||
|
|
|
|||
Loading…
Reference in a new issue