Skip to main content

agent_engine/
orchestration.rs

1use agent_core::orchestration::{
2    CatalogSnapshot, CompletionGate, DelegationPolicy, WorkerRegistry, WorkerRole, WorkerTerminal,
3    WorkerWritePolicy,
4};
5use agent_core::prompt::QualifiedModelId;
6use serde::Serialize;
7use std::collections::{BTreeMap, HashMap, HashSet};
8use std::sync::Mutex;
9
10pub struct DelegationTreeBudget {
11    pub max_depth: u16,
12    pub max_children_per_worker: usize,
13    pub max_total_descendants: usize,
14}
15
16#[derive(Debug, Clone, Eq, PartialEq)]
17pub enum DelegationTreeDenied {
18    DepthLimit,
19    ChildLimit,
20    DescendantLimit,
21    UnknownParent,
22}
23
24#[derive(Debug, Default)]
25struct DelegationTreeState {
26    depth_by_worker: BTreeMap<String, u16>,
27    parent_by_worker: BTreeMap<String, String>,
28    child_count: BTreeMap<String, usize>,
29    descendants: usize,
30}
31
32/// Session-scoped runtime enforcement shared by every subagent tool path.
33pub struct OrchestrationRuntime {
34    inner: Mutex<Inner>,
35    tree_budget: DelegationTreeBudget,
36    tree: Mutex<DelegationTreeState>,
37}
38struct Inner {
39    registry: WorkerRegistry,
40    handles: HashMap<String, agent_core::orchestration::WorkerHandle>,
41}
42
43#[derive(Clone, Debug, Serialize)]
44pub struct AuthorizedWorkerModel {
45    pub model: QualifiedModelId,
46    pub selection_source: SelectionSource,
47    pub network_attempted: bool,
48    pub policy_digest: String,
49    pub catalog_snapshot_id: String,
50    pub catalog_digest: String,
51    pub correlation_id: String,
52    pub cross_provider_grant_id: Option<String>,
53}
54#[derive(Clone, Copy, Debug, Serialize)]
55#[serde(rename_all = "snake_case")]
56pub enum SelectionSource {
57    ForegroundInheritance,
58    ExplicitRequest,
59}
60impl SelectionSource {
61    pub fn as_str(self) -> &'static str {
62        match self {
63            Self::ForegroundInheritance => "foreground_inheritance",
64            Self::ExplicitRequest => "explicit_request",
65        }
66    }
67}
68#[derive(Clone, Debug, Serialize)]
69pub struct DispatchDenial {
70    pub kind: &'static str,
71    pub code: &'static str,
72    pub requested_model: Option<String>,
73    pub foreground_model: String,
74    pub selection_source: SelectionSource,
75    pub network_attempted: bool,
76    pub remediation: &'static str,
77}
78impl std::fmt::Display for DispatchDenial {
79    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
80        write!(
81            f,
82            "{}",
83            serde_json::to_string(self).unwrap_or_else(|_| self.code.into())
84        )
85    }
86}
87impl std::error::Error for DispatchDenial {}
88
89/// Canonicalize a runtime foreground model string into its exact qualified
90/// identity. Legacy configs persist bare Anthropic IDs (e.g. `claude-fable-5`)
91/// which the runtime router still deliberately routes to `anthropic`; the
92/// canonical identity is derived from that single routing entry point — never
93/// from substring heuristics. Unroutable values fail closed.
94pub fn canonical_foreground_identity(raw: &str) -> Result<QualifiedModelId, String> {
95    if let Ok(model) = QualifiedModelId::parse(raw) {
96        return Ok(model);
97    }
98    let route = crate::runtime::openai::resolve_route(raw)
99        .ok_or_else(|| format!("foreground model is unresolved: {raw}"))?;
100    QualifiedModelId::parse(format!("{}/{}", route.provider, route.model))
101        .map_err(|e| e.to_string())
102}
103
104/// Validate an exact identity against source-controlled runtime routing/catalog
105/// descriptors before allowing a model-driven session grant. This operation is
106/// credential- and network-blind: availability is checked later by normal
107/// provider execution, while invented identities fail before policy mutation.
108pub fn validate_user_authorizable_model(raw: &str) -> Result<QualifiedModelId, String> {
109    let model = QualifiedModelId::parse(raw)
110        .map_err(|_| "authorization denied: invalid qualified model".to_string())?;
111    let known = if model.provider() == "anthropic" {
112        agent_core::models::KNOWN_MODELS
113            .iter()
114            .any(|(id, _)| *id == model.model())
115    } else if model.provider() == "openai-codex" {
116        crate::runtime::openai::catalog::codex_static_catalog_models()
117            .iter()
118            .any(|entry| entry.runtime_id() == model.as_str())
119    } else if model.provider() == "xai-auth" {
120        crate::runtime::openai::catalog::xai_model(model.model()).is_some()
121    } else if model.provider() == "github-copilot" {
122        crate::runtime::openai::catalog::github_copilot_runtime_model(model.model()).is_some()
123    } else if model.provider() == "google-gemini" {
124        crate::runtime::openai::catalog::google_gemini_model(model.model()).is_some()
125    } else {
126        crate::runtime::openai::registry::providers()
127            .iter()
128            .find(|provider| provider.key == model.provider())
129            .is_some_and(|provider| {
130                provider
131                    .models
132                    .iter()
133                    .any(|(id, _, _)| *id == model.model())
134            })
135    };
136    if !known || crate::runtime::openai::resolve_route(model.as_str()).is_none() {
137        return Err(format!(
138            "authorization denied: '{}' is not a known routable model",
139            model.as_str()
140        ));
141    }
142    Ok(model)
143}
144
145/// Exact source-controlled OpenRouter worker stack. These identities are
146/// trusted local descriptors (not live catalog results or historical logs).
147const OPENROUTER_WORKER_MODELS: &[&str] = &[
148    "openrouter/deepseek/deepseek-v4-pro",
149    "openrouter/moonshotai/kimi-k2.7-code",
150    "openrouter/z-ai/glm-5.2",
151];
152
153/// Exact manifestless worker choices for a foreground provider.
154fn manifestless_worker_choices(foreground: &QualifiedModelId) -> Vec<QualifiedModelId> {
155    let source = if foreground.provider() == "openrouter" {
156        OPENROUTER_WORKER_MODELS
157    } else {
158        &[]
159    };
160    let mut choices: Vec<_> = source
161        .iter()
162        .filter_map(|value| QualifiedModelId::parse(*value).ok())
163        .collect();
164    choices.push(foreground.clone());
165    choices.sort();
166    choices.dedup();
167    choices
168}
169
170impl OrchestrationRuntime {
171    /// Build a snapshot exclusively from runtime-owned routing/catalog descriptors.
172    /// `manifest_references` are deliberately ignored here: manifests may narrow the
173    /// snapshot, but can never create catalog membership. The foreground is admitted
174    /// only after the credential-free runtime router resolves its exact identity.
175    pub fn trusted_catalog(
176        foreground: &QualifiedModelId,
177        _manifest_references: impl IntoIterator<Item = QualifiedModelId>,
178    ) -> Result<CatalogSnapshot, &'static str> {
179        use crate::runtime::openai::catalog;
180        use agent_core::orchestration::CatalogEntry;
181
182        let route = crate::runtime::openai::resolve_route(foreground.as_str())
183            .ok_or("foreground model is unresolved")?;
184        if route.provider != foreground.provider() {
185            return Err("foreground provider did not resolve exactly");
186        }
187
188        let mut descriptors = catalog::codex_static_catalog_models();
189        descriptors.extend(catalog::xai_static_catalog_models());
190        descriptors.extend(catalog::copilot_static_catalog_models());
191        descriptors.extend(catalog::google_gemini_static_catalog_models());
192        let mut entries: Vec<_> = descriptors
193            .into_iter()
194            .filter_map(|descriptor| QualifiedModelId::parse(descriptor.runtime_id()).ok())
195            .map(|model| CatalogEntry {
196                model,
197                available: true,
198                worker_eligible: true,
199            })
200            .collect();
201        entries.extend(
202            manifestless_worker_choices(foreground)
203                .into_iter()
204                .map(|model| CatalogEntry {
205                    model,
206                    available: true,
207                    worker_eligible: true,
208                }),
209        );
210        // `resolve_route` is itself a trusted local provider descriptor. This is
211        // independent of manifest content and permits configured generic/Anthropic
212        // foreground routes whose live catalogs are not fetched during startup.
213        entries.push(CatalogEntry {
214            model: foreground.clone(),
215            available: true,
216            worker_eligible: true,
217        });
218        Ok(CatalogSnapshot::from_entries(entries))
219    }
220
221    /// Secure manifestless baseline: the exact foreground identity is the only
222    /// worker choice in a deterministic runtime-controlled catalog.
223    pub fn baseline(
224        foreground: QualifiedModelId,
225        concurrent: usize,
226        total: usize,
227    ) -> Result<Self, &'static str> {
228        let choices = manifestless_worker_choices(&foreground);
229        let catalog = Self::trusted_catalog(&foreground, std::iter::empty())?;
230        let mut policy =
231            DelegationPolicy::provider_baseline(foreground, catalog, choices, concurrent, total)?;
232        policy.recycle_reconciled_workers = true;
233        Ok(Self::new(policy))
234    }
235
236    pub fn new(policy: DelegationPolicy) -> Self {
237        let max_total_descendants = policy.max_total_workers;
238        Self {
239            inner: Mutex::new(Inner {
240                registry: WorkerRegistry::new(policy),
241                handles: HashMap::new(),
242            }),
243            tree_budget: DelegationTreeBudget {
244                max_depth: 4,
245                max_children_per_worker: 8,
246                max_total_descendants,
247            },
248            tree: Mutex::new(DelegationTreeState::default()),
249        }
250    }
251
252    pub fn with_tree_budget(mut self, budget: DelegationTreeBudget) -> Result<Self, &'static str> {
253        if budget.max_depth == 0
254            || budget.max_children_per_worker == 0
255            || budget.max_total_descendants == 0
256        {
257            return Err("invalid delegation tree budget");
258        }
259        self.tree_budget = budget;
260        Ok(self)
261    }
262
263    /// Reserve one tree edge before allocating channels/threads/provider
264    /// runtimes. `parent=None` is a foreground-root child. Every denial is
265    /// fail-closed and leaves the counters unchanged.
266    pub fn reserve_delegation(
267        &self,
268        worker_id: &str,
269        parent: Option<&str>,
270    ) -> Result<u16, DelegationTreeDenied> {
271        let mut tree = self.tree.lock().unwrap();
272        if tree.descendants >= self.tree_budget.max_total_descendants {
273            return Err(DelegationTreeDenied::DescendantLimit);
274        }
275        let depth = match parent {
276            None => 1,
277            Some(parent_id) => tree
278                .depth_by_worker
279                .get(parent_id)
280                .copied()
281                .ok_or(DelegationTreeDenied::UnknownParent)?
282                .saturating_add(1),
283        };
284        if depth > self.tree_budget.max_depth {
285            return Err(DelegationTreeDenied::DepthLimit);
286        }
287        let parent_key = parent.unwrap_or("<root>");
288        if tree.child_count.get(parent_key).copied().unwrap_or(0)
289            >= self.tree_budget.max_children_per_worker
290        {
291            return Err(DelegationTreeDenied::ChildLimit);
292        }
293        if tree.depth_by_worker.contains_key(worker_id) {
294            return Err(DelegationTreeDenied::ChildLimit);
295        }
296        tree.depth_by_worker.insert(worker_id.to_string(), depth);
297        tree.parent_by_worker
298            .insert(worker_id.to_string(), parent_key.to_string());
299        *tree.child_count.entry(parent_key.to_string()).or_default() += 1;
300        tree.descendants += 1;
301        Ok(depth)
302    }
303
304    /// Release one reserved tree edge. Descendants are recursively removed
305    /// because a live child cannot outlive its parent lineage. The caller-
306    /// supplied `parent` is intentionally ignored: ownership is taken from the
307    /// reservation recorded under the tree lock, so a stale or hostile caller
308    /// cannot decrement another parent's fanout.
309    pub fn release_delegation(&self, worker_id: &str, _parent: Option<&str>) {
310        let mut tree = self.tree.lock().unwrap();
311        let mut stack = vec![worker_id.to_string()];
312        while let Some(current) = stack.pop() {
313            let children = tree
314                .parent_by_worker
315                .iter()
316                .filter_map(|(child, parent)| (parent == &current).then_some(child.clone()))
317                .collect::<Vec<_>>();
318            stack.extend(children);
319            if tree.depth_by_worker.remove(&current).is_some() {
320                tree.descendants = tree.descendants.saturating_sub(1);
321                if let Some(parent_key) = tree.parent_by_worker.remove(&current) {
322                    if let Some(count) = tree.child_count.get_mut(&parent_key) {
323                        *count = count.saturating_sub(1);
324                    }
325                }
326                tree.child_count.remove(&current);
327            }
328        }
329    }
330
331    pub fn delegation_descendants(&self) -> usize {
332        self.tree.lock().unwrap().descendants
333    }
334    pub fn preflight(&self, model: &str) -> Result<(), String> {
335        let model = QualifiedModelId::parse(model)
336            .map_err(|_| "delegation denied: invalid qualified model".to_string())?;
337        self.inner
338            .lock()
339            .unwrap()
340            .registry
341            .validate_dispatch(&model)
342            .map_err(|e| format!("delegation denied: {}", e.code()))
343    }
344
345    /// Honors an explicit mid-session user trust grant for one exact worker
346    /// model. Trusted models were never meant to be pinned at session start:
347    /// the live policy and catalog are extended in place while worker
348    /// lifecycle state and concurrency limits are preserved. The grant takes
349    /// effect for the next dispatch decision.
350    pub fn grant_worker_model(&self, model: &str) -> Result<(), String> {
351        let model = QualifiedModelId::parse(model)
352            .map_err(|_| "grant denied: invalid qualified model".to_string())?;
353        self.inner
354            .lock()
355            .unwrap()
356            .registry
357            .grant_worker_model(model)
358            .map_err(|error| format!("grant denied: {error}"))
359    }
360
361    /// Read-only UltraCode authorization snapshot. The exact foreground identity,
362    /// policy authorization, and configured limits are checked under one lock; no
363    /// worker reservation or lifecycle state is created.
364    pub fn ultracode_readiness(&self, model: &str) -> Result<(usize, usize), String> {
365        let model = QualifiedModelId::parse(model)
366            .map_err(|_| "delegation denied: invalid qualified model".to_string())?;
367        let inner = self.inner.lock().unwrap();
368        if inner.registry.foreground_model() != &model {
369            return Err("delegation denied: model is not exact foreground".into());
370        }
371        inner
372            .registry
373            .validate_dispatch(&model)
374            .map_err(|e| format!("delegation denied: {}", e.code()))?;
375        let policy = inner.registry.policy();
376        if policy.max_concurrent_workers == 0 || policy.max_total_workers == 0 {
377            return Err("delegation denied: invalid worker limits".into());
378        }
379        Ok((policy.max_concurrent_workers, policy.max_total_workers))
380    }
381
382    /// Single parse/resolve/catalog/policy/limits decision point used by every
383    /// public spawn path. No credential or provider object is reachable here.
384    pub fn resolve_and_authorize(
385        &self,
386        runtime_handle: &str,
387        requested: Option<&str>,
388    ) -> Result<AuthorizedWorkerModel, DispatchDenial> {
389        let mut inner = self.inner.lock().unwrap();
390        let foreground = inner.registry.foreground_model().clone();
391        let selection_source = if requested.is_some() {
392            SelectionSource::ExplicitRequest
393        } else {
394            SelectionSource::ForegroundInheritance
395        };
396        let model = match requested {
397            Some(value) => QualifiedModelId::parse(value).map_err(|_| DispatchDenial {
398                kind: "dispatch_denied",
399                code: "invalid_qualified_model",
400                requested_model: None,
401                foreground_model: foreground.as_str().into(),
402                selection_source,
403                network_attempted: false,
404                remediation: "Omit model to inherit foreground or select an exact session choice.",
405            })?,
406            None => foreground.clone(),
407        };
408        if inner.handles.contains_key(runtime_handle) {
409            return Err(DispatchDenial {
410                kind: "dispatch_denied",
411                code: "duplicate_runtime_handle",
412                requested_model: Some(model.as_str().into()),
413                foreground_model: foreground.as_str().into(),
414                selection_source,
415                network_attempted: false,
416                remediation: "Use a new worker handle.",
417            });
418        }
419        let policy_digest = inner.registry.policy().digest();
420        let catalog_snapshot_id = inner.registry.policy().catalog_snapshot_id().to_owned();
421        let catalog_digest = inner.registry.policy().catalog_digest().to_owned();
422        let cross_provider_grant_id = inner
423            .registry
424            .policy()
425            .authorize(&model)
426            .ok()
427            .flatten()
428            .map(str::to_owned);
429        let handle = inner
430            .registry
431            .authorize_dispatch(&model, WorkerRole::Implementer, WorkerWritePolicy::ReadOnly)
432            .map_err(|error| DispatchDenial {
433                kind: "dispatch_denied",
434                code: error.code(),
435                requested_model: Some(model.as_str().into()),
436                foreground_model: foreground.as_str().into(),
437                selection_source,
438                network_attempted: false,
439                remediation: match error.typed_code() {
440                    agent_core::orchestration::DispatchFailureCode::ConcurrencyLimit =>
441                        "Wait for running workers to finish; collect and reconcile their results. Changing model does not release capacity.",
442                    agent_core::orchestration::DispatchFailureCode::TotalWorkerLimit =>
443                        if inner.registry.policy().recycle_reconciled_workers {
444                            "Collect finished workers with reconciled=true to release outstanding-worker capacity. Changing model does not release capacity."
445                        } else {
446                            "The explicit delegation policy cumulative worker budget is exhausted. Changing model or collecting workers does not renew it."
447                        },
448                    _ => "Omit model to inherit foreground, select an exact session choice, or trust the model mid-session (favorite it in the models picker).",
449                },
450            })?;
451        inner.handles.insert(runtime_handle.to_owned(), handle);
452        Ok(AuthorizedWorkerModel {
453            model,
454            selection_source,
455            network_attempted: false,
456            policy_digest,
457            catalog_snapshot_id,
458            catalog_digest,
459            correlation_id: runtime_handle.to_owned(),
460            cross_provider_grant_id,
461        })
462    }
463
464    /// Marks the persisted handle Starting immediately before the runtime/thread
465    /// factory is invoked. Starting remains rollback-safe until `mark_running`.
466    pub fn mark_starting(&self, runtime_handle: &str) -> Result<(), String> {
467        let mut inner = self.inner.lock().unwrap();
468        let handle = inner
469            .handles
470            .get(runtime_handle)
471            .cloned()
472            .ok_or_else(|| "worker lifecycle unavailable".to_string())?;
473        inner
474            .registry
475            .mark_starting(&handle)
476            .map_err(|_| "worker lifecycle unavailable".into())
477    }
478    pub fn mark_running(&self, runtime_handle: &str) -> Result<(), String> {
479        let mut inner = self.inner.lock().unwrap();
480        let handle = inner
481            .handles
482            .get(runtime_handle)
483            .cloned()
484            .ok_or_else(|| "worker lifecycle unavailable".to_string())?;
485        inner
486            .registry
487            .mark_running(&handle)
488            .map_err(|_| "worker lifecycle unavailable".into())
489    }
490
491    pub fn rollback(&self, runtime_handle: &str) {
492        let mut inner = self.inner.lock().unwrap();
493        if let Some(handle) = inner.handles.get(runtime_handle).cloned() {
494            let _ = inner.registry.rollback_dispatch(&handle);
495            inner.handles.remove(runtime_handle);
496        }
497    }
498
499    pub fn effective_choices(&self) -> Vec<String> {
500        self.inner
501            .lock()
502            .unwrap()
503            .registry
504            .policy()
505            .effective_choices()
506            .iter()
507            .map(|model| model.as_str().to_owned())
508            .collect()
509    }
510
511    pub fn foreground_model(&self) -> String {
512        self.inner
513            .lock()
514            .unwrap()
515            .registry
516            .foreground_model()
517            .as_str()
518            .to_owned()
519    }
520    pub fn authorize_with_policy(
521        &self,
522        runtime_handle: &str,
523        model: &str,
524        role: WorkerRole,
525        writes: WorkerWritePolicy,
526    ) -> Result<(), String> {
527        let model = QualifiedModelId::parse(model)
528            .map_err(|_| "delegation denied: invalid qualified model".to_string())?;
529        let mut inner = self.inner.lock().unwrap();
530        if inner.handles.contains_key(runtime_handle) {
531            return Err("delegation denied: duplicate runtime handle".to_string());
532        }
533        let handle = inner
534            .registry
535            .authorize_dispatch(&model, role, writes)
536            .map_err(|e| format!("delegation denied: {}", e.code()))?;
537        inner.registry.mark_starting(&handle).map_err(|error| {
538            let _ = inner.registry.rollback_dispatch(&handle);
539            error.to_string()
540        })?;
541        if let Err(error) = inner.registry.mark_running(&handle) {
542            let _ = inner.registry.rollback_dispatch(&handle);
543            return Err(error.to_string());
544        }
545        inner.handles.insert(runtime_handle.to_string(), handle);
546        Ok(())
547    }
548    pub fn authorize(&self, runtime_handle: &str, model: &str) -> Result<(), String> {
549        self.authorize_with_policy(
550            runtime_handle,
551            model,
552            WorkerRole::Implementer,
553            WorkerWritePolicy::ReadOnly,
554        )
555    }
556    pub fn poll(&self, id: &str, fingerprint: &str) -> Result<(), String> {
557        let mut inner = self.inner.lock().unwrap();
558        let h = inner
559            .handles
560            .get(id)
561            .cloned()
562            .ok_or_else(|| "unknown worker".to_string())?;
563        inner.registry.poll(&h, fingerprint).map_err(str::to_string)
564    }
565    pub fn steer(&self, id: &str) -> Result<(), String> {
566        let mut inner = self.inner.lock().unwrap();
567        let h = inner
568            .handles
569            .get(id)
570            .cloned()
571            .ok_or_else(|| "unknown worker".to_string())?;
572        inner.registry.steer(&h).map_err(str::to_string)
573    }
574    pub fn terminal_and_collect(&self, id: &str, terminal: WorkerTerminal) -> Result<(), String> {
575        let mut inner = self.inner.lock().unwrap();
576        let h = inner
577            .handles
578            .get(id)
579            .cloned()
580            .ok_or_else(|| "unknown worker".to_string())?;
581        inner
582            .registry
583            .mark_terminal(&h, terminal)
584            .map_err(str::to_string)?;
585        inner.registry.collect(&h).map_err(str::to_string)
586    }
587    pub fn finish_one_shot(&self, id: &str, terminal: WorkerTerminal) -> Result<(), String> {
588        self.terminal_and_collect(id, terminal)?;
589        self.reconcile(id)?;
590        // One-shot results have no retained reactive handle to be reaped later.
591        self.retire_reconciled(&[id.to_owned()].into_iter().collect());
592        Ok(())
593    }
594    pub fn reconcile(&self, id: &str) -> Result<(), String> {
595        let mut inner = self.inner.lock().unwrap();
596        let h = inner
597            .handles
598            .get(id)
599            .cloned()
600            .ok_or_else(|| "unknown worker".to_string())?;
601        inner.registry.reconcile(&h).map_err(str::to_string)
602    }
603    /// Reaper handoff: use exact IDs removed under the runtime registry lock,
604    /// never a complement snapshot that could include newly registered workers.
605    pub fn retire_reconciled(&self, removed_handles: &HashSet<String>) {
606        let mut inner = self.inner.lock().unwrap();
607        for id in removed_handles {
608            if let Some(handle) = inner.handles.get(id).cloned() {
609                if inner.registry.retire_reconciled(&handle) {
610                    inner.handles.remove(id);
611                }
612            }
613        }
614    }
615
616    /// Snapshot runtime handle IDs still named by the completion gate. Taking this
617    /// snapshot before locking the subagent registry avoids cross-registry lock order.
618    pub fn unreconciled_runtime_handles(&self) -> HashSet<String> {
619        let inner = self.inner.lock().unwrap();
620        // Use all_unreconciled_ids (includes Running workers) — NOT
621        // completion_gate (which only reports Terminal|Collected). The reaper
622        // must retain handles the orchestration still tracks, even if the
623        // gate lets the foreground turn through.
624        let unreconciled = inner.registry.all_unreconciled_ids();
625        if unreconciled.is_empty() {
626            return HashSet::new();
627        }
628        inner
629            .handles
630            .iter()
631            .filter(|(_, handle)| unreconciled.iter().any(|id| id == handle.id()))
632            .map(|(runtime_id, _)| runtime_id.clone())
633            .collect()
634    }
635
636    /// Whether a runtime handle is still named by the policy completion gate.
637    pub fn is_unreconciled(&self, runtime_handle: &str) -> bool {
638        self.unreconciled_runtime_handles().contains(runtime_handle)
639    }
640
641    pub fn completion_gate(&self) -> CompletionGate {
642        // Core WorkerRegistry reports internal policy IDs (`worker-N`). Map them
643        // back through the runtime handle table so tool-facing remediation can
644        // cite `sa_*` handles accepted by `subagent_collect`.
645        let inner = self.inner.lock().unwrap();
646        match inner.registry.completion_gate() {
647            CompletionGate::Allowed => CompletionGate::Allowed,
648            CompletionGate::Warning { workers } => CompletionGate::Warning {
649                workers: Self::map_policy_ids_to_runtime(&inner.handles, workers),
650            },
651            CompletionGate::Blocked { workers } => CompletionGate::Blocked {
652                workers: Self::map_policy_ids_to_runtime(&inner.handles, workers),
653            },
654        }
655    }
656
657    /// Map policy `WorkerHandle` IDs (`worker-N`) to runtime handles (`sa_*`).
658    /// Preserves the deterministic order of the policy gate's worker list.
659    fn map_policy_ids_to_runtime(
660        handles: &HashMap<String, agent_core::orchestration::WorkerHandle>,
661        policy_ids: Vec<String>,
662    ) -> Vec<String> {
663        let reverse: HashMap<&str, &str> = handles
664            .iter()
665            .map(|(runtime_handle, policy_handle)| (policy_handle.id(), runtime_handle.as_str()))
666            .collect();
667        policy_ids
668            .into_iter()
669            .map(|policy_id| {
670                reverse
671                    .get(policy_id.as_str())
672                    .map(|runtime| (*runtime).to_owned())
673                    .unwrap_or(policy_id)
674            })
675            .collect()
676    }
677    pub fn check_foreground_write(&self, path: &str) -> agent_core::orchestration::ScopeDecision {
678        self.inner
679            .lock()
680            .unwrap()
681            .registry
682            .check_foreground_write(path)
683    }
684    pub fn telemetry_json(&self) -> String {
685        serde_json::to_string(self.inner.lock().unwrap().registry.telemetry())
686            .unwrap_or_else(|_| "[]".into())
687    }
688    pub fn enforcement_mode(&self) -> agent_core::orchestration::EnforcementMode {
689        self.inner.lock().unwrap().registry.mode()
690    }
691}
692
693#[cfg(test)]
694mod tests {
695    use super::*;
696    use std::sync::Arc;
697    fn model(s: &str) -> QualifiedModelId {
698        QualifiedModelId::parse(s).unwrap()
699    }
700    /// Regression: a legacy config (`model = claude-fable-5`, no provider
701    /// segment) must not brick startup. The router deliberately routes bare
702    /// Claude IDs to `anthropic`, so the baseline foreground identity must be
703    /// the router's canonical `anthropic/<id>` — not a startup Config error.
704    #[test]
705    fn legacy_bare_claude_foreground_canonicalizes_via_router() {
706        let id = canonical_foreground_identity("claude-fable-5").unwrap();
707        assert_eq!(id.as_str(), "anthropic/claude-fable-5");
708        // The canonical identity must satisfy the manifestless baseline.
709        OrchestrationRuntime::baseline(id, 8, 64).unwrap();
710    }
711
712    #[test]
713    fn baseline_openrouter_catalog_authorizes_exact_chinese_worker_stack() {
714        let deepseek = model("openrouter/deepseek/deepseek-v4-pro");
715        let glm = model("openrouter/z-ai/glm-5.2");
716        let kimi = model("openrouter/moonshotai/kimi-k2.7-code");
717        let rt = OrchestrationRuntime::baseline(deepseek.clone(), 3, 8).unwrap();
718
719        assert_eq!(
720            rt.effective_choices(),
721            vec![
722                deepseek.as_str().to_owned(),
723                kimi.as_str().to_owned(),
724                glm.as_str().to_owned(),
725            ]
726        );
727        assert_eq!(
728            rt.resolve_and_authorize("sa_glm", Some(glm.as_str()))
729                .unwrap()
730                .model,
731            glm
732        );
733        assert_eq!(
734            rt.resolve_and_authorize("sa_kimi", Some(kimi.as_str()))
735                .unwrap()
736                .model,
737            kimi
738        );
739    }
740
741    #[test]
742    fn manifestless_non_openrouter_baseline_remains_foreground_only() {
743        let foreground = model("anthropic/claude-fable-5");
744        let rt = OrchestrationRuntime::baseline(foreground.clone(), 3, 8).unwrap();
745        assert_eq!(rt.effective_choices(), vec![foreground.as_str().to_owned()]);
746        let denied = rt
747            .resolve_and_authorize(
748                "sa_openrouter",
749                Some("openrouter/moonshotai/kimi-k2.7-code"),
750            )
751            .unwrap_err();
752        assert!(matches!(
753            denied.code,
754            "provider_not_allowed" | "catalog_model_unknown"
755        ));
756    }
757
758    /// Regression: an explicit mid-session user trust grant must flip an
759    /// exact-model dispatch from `provider_not_allowed` to authorized without
760    /// restarting the session. Trusted models were never meant to be pinned
761    /// at session start.
762    #[test]
763    fn baseline_honors_mid_session_user_grant() {
764        let foreground = model("anthropic/claude-fable-5");
765        let requested = "openai-codex/gpt-5.6-sol";
766        let rt = OrchestrationRuntime::baseline(foreground.clone(), 3, 8).unwrap();
767        assert_eq!(rt.effective_choices(), vec![foreground.as_str().to_owned()]);
768        let denied = rt
769            .resolve_and_authorize("sa_denied", Some(requested))
770            .unwrap_err();
771        assert!(matches!(
772            denied.code,
773            "provider_not_allowed" | "catalog_model_unknown"
774        ));
775        assert!(!denied.network_attempted);
776
777        rt.grant_worker_model(requested).unwrap();
778
779        assert!(rt.effective_choices().contains(&requested.to_owned()));
780        let authorized = rt
781            .resolve_and_authorize("sa_granted", Some(requested))
782            .unwrap();
783        assert_eq!(authorized.model.as_str(), requested);
784        assert_eq!(
785            authorized.cross_provider_grant_id.as_deref(),
786            Some("session-user-grant-openai-codex/gpt-5.6-sol")
787        );
788        assert!(!authorized.network_attempted);
789        rt.rollback("sa_granted");
790        assert_eq!(rt.completion_gate(), CompletionGate::Allowed);
791    }
792
793    /// A grant for a malformed identity fails closed and changes nothing.
794    #[test]
795    fn mid_session_grant_rejects_invalid_qualified_model() {
796        let foreground = model("anthropic/claude-fable-5");
797        let rt = OrchestrationRuntime::baseline(foreground.clone(), 3, 8).unwrap();
798        assert!(rt.grant_worker_model("").is_err());
799        assert!(rt.grant_worker_model("not-qualified").is_err());
800        assert_eq!(rt.effective_choices(), vec![foreground.as_str().to_owned()]);
801    }
802
803    #[test]
804    fn qualified_foreground_identity_passes_through_unchanged() {
805        for raw in [
806            "openai-codex/gpt-5.6-sol",
807            "anthropic/claude-fable-5",
808            "openrouter/z-ai/glm-5.1",
809        ] {
810            assert_eq!(canonical_foreground_identity(raw).unwrap().as_str(), raw);
811        }
812    }
813
814    #[test]
815    fn unroutable_foreground_identity_fails_closed() {
816        // Neither a valid qualified ID nor routable: empty model segment
817        // under an unknown provider prefix.
818        assert!(canonical_foreground_identity("foo/").is_err());
819        assert!(canonical_foreground_identity("").is_err());
820    }
821
822    #[test]
823    fn trusted_catalog_rejects_unresolved_foreground_and_ignores_manifest_candidates() {
824        let unresolved = model("invented/foreground");
825        assert!(OrchestrationRuntime::trusted_catalog(
826            &unresolved,
827            [model("anthropic/claude-manifest-injected")]
828        )
829        .is_err());
830
831        let foreground = model("openai-codex/gpt-5.1-codex-mini");
832        let catalog = OrchestrationRuntime::trusted_catalog(
833            &foreground,
834            [model("anthropic/claude-manifest-injected")],
835        )
836        .unwrap();
837        assert!(catalog.contains(&foreground));
838        assert!(!catalog.contains(&model("anthropic/claude-manifest-injected")));
839    }
840
841    #[test]
842    fn authorization_reserves_without_claiming_worker_is_running() {
843        let foreground = model("anthropic/claude-foreground");
844        let rt = OrchestrationRuntime::new(DelegationPolicy::enforced(
845            foreground.clone(),
846            [foreground.clone()],
847            1,
848            1,
849        ));
850        rt.resolve_and_authorize("sa_reserved", None).unwrap();
851        assert!(rt.poll("sa_reserved", "not-started").is_err());
852        rt.mark_starting("sa_reserved").unwrap();
853        assert!(rt.poll("sa_reserved", "not-running").is_err());
854        rt.rollback("sa_reserved");
855        assert_eq!(rt.completion_gate(), CompletionGate::Allowed);
856    }
857
858    #[test]
859    fn one_shot_uses_the_full_lifecycle_before_terminal_collection() {
860        let foreground = model("anthropic/foreground");
861        let rt = OrchestrationRuntime::new(DelegationPolicy::enforced(
862            foreground.clone(),
863            [foreground],
864            1,
865            1,
866        ));
867        rt.resolve_and_authorize("sa_one_shot", None).unwrap();
868        rt.mark_starting("sa_one_shot").unwrap();
869        rt.mark_running("sa_one_shot").unwrap();
870        rt.finish_one_shot("sa_one_shot", WorkerTerminal::Completed)
871            .unwrap();
872        assert_eq!(rt.completion_gate(), CompletionGate::Allowed);
873        let telemetry = rt.telemetry_json();
874        for event in [
875            "worker.dispatch_allowed",
876            "worker.starting",
877            "worker.running",
878            "worker.terminal",
879            "worker.collected",
880            "worker.reconciled",
881        ] {
882            assert!(telemetry.contains(event), "missing {event}: {telemetry}");
883        }
884    }
885
886    #[test]
887    fn authorization_happens_before_a_worker_is_registered() {
888        let rt = OrchestrationRuntime::new(DelegationPolicy::enforced(
889            model("anthropic/foreground"),
890            [model("anthropic/worker")],
891            1,
892            1,
893        ));
894        assert!(rt.authorize("sa_1", "openrouter/worker").is_err());
895        assert_eq!(rt.completion_gate(), CompletionGate::Allowed);
896        rt.authorize("sa_2", "anthropic/worker").unwrap();
897        // Running workers pass through the gate (reactive pattern).
898        assert_eq!(rt.completion_gate(), CompletionGate::Allowed);
899        rt.poll("sa_2", "same").unwrap();
900        rt.steer("sa_2").unwrap();
901        rt.terminal_and_collect("sa_2", WorkerTerminal::Completed)
902            .unwrap();
903        // Terminal+Collected: gate blocks until reconciled.
904        assert!(matches!(
905            rt.completion_gate(),
906            CompletionGate::Blocked { .. }
907        ));
908        rt.reconcile("sa_2").unwrap();
909        assert_eq!(rt.completion_gate(), CompletionGate::Allowed);
910    }
911
912    /// Completion remediation must cite tool-facing `sa_*` handles, never the
913    /// internal policy IDs (`worker-N`) that `WorkerRegistry` allocates.
914    #[test]
915    fn delegation_tree_depth_children_and_total_are_independently_bounded() {
916        let foreground = model("anthropic/foreground");
917        let rt = OrchestrationRuntime::new(DelegationPolicy::enforced(
918            foreground.clone(),
919            [foreground],
920            8,
921            16,
922        ))
923        .with_tree_budget(DelegationTreeBudget {
924            max_depth: 2,
925            max_children_per_worker: 2,
926            max_total_descendants: 3,
927        })
928        .unwrap();
929
930        assert_eq!(rt.reserve_delegation("a", None), Ok(1));
931        assert_eq!(
932            rt.reserve_delegation("a", None),
933            Err(DelegationTreeDenied::ChildLimit),
934            "registration/spawn retry cannot double-reserve one handle"
935        );
936        assert_eq!(rt.reserve_delegation("b", Some("a")), Ok(2));
937        assert_eq!(
938            rt.reserve_delegation("too-deep", Some("b")),
939            Err(DelegationTreeDenied::DepthLimit)
940        );
941        assert_eq!(rt.reserve_delegation("c", Some("a")), Ok(2));
942        assert_eq!(
943            rt.reserve_delegation("too-many", Some("a")),
944            Err(DelegationTreeDenied::DescendantLimit)
945        );
946        assert_eq!(rt.delegation_descendants(), 3);
947        rt.release_delegation("c", Some("a"));
948        assert_eq!(rt.delegation_descendants(), 2);
949    }
950
951    /// Active reservation count is exact across reserve/release races.
952    #[test]
953    fn delegation_release_is_idempotent_under_completion_cancel_race() {
954        let foreground = model("anthropic/foreground");
955        let rt = Arc::new(OrchestrationRuntime::new(DelegationPolicy::enforced(
956            foreground.clone(),
957            [foreground],
958            4,
959            8,
960        )));
961        rt.reserve_delegation("child", None).unwrap();
962        let a = Arc::clone(&rt);
963        let b = Arc::clone(&rt);
964        let t1 = std::thread::spawn(move || a.release_delegation("child", None));
965        let t2 = std::thread::spawn(move || b.release_delegation("child", None));
966        t1.join().unwrap();
967        t2.join().unwrap();
968        assert_eq!(rt.delegation_descendants(), 0);
969    }
970
971    #[test]
972    fn parent_exit_releases_entire_live_descendant_subtree() {
973        let foreground = model("anthropic/foreground");
974        let rt = OrchestrationRuntime::new(DelegationPolicy::enforced(
975            foreground.clone(),
976            [foreground],
977            8,
978            16,
979        ));
980        rt.reserve_delegation("parent", None).unwrap();
981        rt.reserve_delegation("child", Some("parent")).unwrap();
982        rt.reserve_delegation("grandchild", Some("child")).unwrap();
983        assert_eq!(rt.delegation_descendants(), 3);
984        rt.release_delegation("parent", None);
985        assert_eq!(rt.delegation_descendants(), 0);
986        assert_eq!(
987            rt.reserve_delegation("orphan", Some("child")),
988            Err(DelegationTreeDenied::UnknownParent)
989        );
990    }
991
992    /// Active reservation count is exact when a live worker completes while
993    /// its parent is cancelled and recursively releases the same subtree.
994    #[test]
995    fn parent_cancel_racing_child_completion_releases_every_lease() {
996        let foreground = model("anthropic/foreground");
997        let rt = Arc::new(OrchestrationRuntime::new(DelegationPolicy::enforced(
998            foreground.clone(),
999            [foreground],
1000            8,
1001            16,
1002        )));
1003        rt.reserve_delegation("parent", None).unwrap();
1004        rt.reserve_delegation("child", Some("parent")).unwrap();
1005        rt.reserve_delegation("grandchild", Some("child")).unwrap();
1006        let cancel = Arc::clone(&rt);
1007        let complete = Arc::clone(&rt);
1008        let t1 = std::thread::spawn(move || cancel.release_delegation("parent", None));
1009        let t2 = std::thread::spawn(move || complete.release_delegation("child", Some("parent")));
1010        t1.join().unwrap();
1011        t2.join().unwrap();
1012
1013        assert_eq!(rt.delegation_descendants(), 0);
1014        assert_eq!(
1015            rt.reserve_delegation("orphan", Some("grandchild")),
1016            Err(DelegationTreeDenied::UnknownParent)
1017        );
1018        assert_eq!(rt.reserve_delegation("replacement", None), Ok(1));
1019        rt.release_delegation("replacement", None);
1020        assert_eq!(rt.delegation_descendants(), 0);
1021    }
1022
1023    /// A child receives tree lineage only, never a new policy/grant mutation:
1024    /// even when the session already has one explicit cross-provider grant,
1025    /// child lineage cannot use that fact to authorize a different identity.
1026    #[test]
1027    fn descendant_lineage_cannot_mint_or_broaden_session_model_grants() {
1028        let foreground = model("anthropic/foreground");
1029        let rt = OrchestrationRuntime::new(DelegationPolicy::enforced(
1030            foreground.clone(),
1031            [foreground],
1032            4,
1033            8,
1034        ));
1035        rt.grant_worker_model("openai-codex/gpt-5.6-sol").unwrap();
1036        rt.reserve_delegation("parent", None).unwrap();
1037        rt.reserve_delegation("child", Some("parent")).unwrap();
1038
1039        assert!(rt.preflight("openai-codex/gpt-5.6-sol").is_ok());
1040        assert!(
1041            rt.preflight("openai-codex/gpt-5.6-luna").is_err(),
1042            "a child lineage must not broaden one exact session grant to a sibling model"
1043        );
1044        assert!(
1045            rt.preflight("openrouter/ungranted").is_err(),
1046            "tree lineage must never mint a cross-provider grant"
1047        );
1048        rt.release_delegation("parent", None);
1049        assert_eq!(rt.delegation_descendants(), 0);
1050    }
1051
1052    #[test]
1053    fn completion_gate_reports_runtime_handles_not_policy_ids() {
1054        let rt = OrchestrationRuntime::new(DelegationPolicy::enforced(
1055            model("anthropic/foreground"),
1056            [model("anthropic/worker")],
1057            2,
1058            2,
1059        ));
1060        rt.authorize("sa_alpha", "anthropic/worker").unwrap();
1061        rt.authorize("sa_beta", "anthropic/worker").unwrap();
1062        // Transition both to Terminal so the gate reports them.
1063        rt.terminal_and_collect("sa_alpha", WorkerTerminal::Completed)
1064            .unwrap();
1065        rt.terminal_and_collect("sa_beta", WorkerTerminal::Completed)
1066            .unwrap();
1067
1068        match rt.completion_gate() {
1069            CompletionGate::Blocked { workers } => {
1070                assert_eq!(
1071                    workers,
1072                    vec!["sa_alpha".to_string(), "sa_beta".to_string()],
1073                    "blocked IDs must be runtime handles in deterministic policy order"
1074                );
1075                for id in &workers {
1076                    assert!(
1077                        id.starts_with("sa_"),
1078                        "blocked id must be a subagent_collect handle, got {id}"
1079                    );
1080                    assert!(
1081                        !id.starts_with("worker-"),
1082                        "blocked id must never leak policy WorkerHandle, got {id}"
1083                    );
1084                }
1085            }
1086            other => panic!("expected Blocked with two sa_* handles, got {other:?}"),
1087        }
1088
1089        // After one reconcile, only the remaining runtime handle is reported.
1090        rt.reconcile("sa_alpha").unwrap();
1091        match rt.completion_gate() {
1092            CompletionGate::Blocked { workers } => {
1093                assert_eq!(workers, vec!["sa_beta".to_string()]);
1094            }
1095            other => panic!("expected only sa_beta still blocked, got {other:?}"),
1096        }
1097    }
1098}
1099
1100#[cfg(test)]
1101mod capacity_recovery_tests {
1102    use super::*;
1103    fn baseline(concurrent: usize, total: usize) -> OrchestrationRuntime {
1104        OrchestrationRuntime::baseline(
1105            QualifiedModelId::parse("anthropic/claude-fable-5").unwrap(),
1106            concurrent,
1107            total,
1108        )
1109        .unwrap()
1110    }
1111    #[test]
1112    fn manifestless_long_session_reuses_capacity_without_replacing_policy() {
1113        let rt = baseline(8, 64);
1114        let choices = rt.effective_choices();
1115        for n in 0..150 {
1116            let id = format!("sa_{n}");
1117            rt.resolve_and_authorize(&id, None).unwrap();
1118            rt.mark_starting(&id).unwrap();
1119            rt.mark_running(&id).unwrap();
1120            rt.terminal_and_collect(&id, WorkerTerminal::Completed)
1121                .unwrap();
1122            rt.reconcile(&id).unwrap();
1123            rt.reconcile(&id).unwrap();
1124        }
1125        assert_eq!(rt.effective_choices(), choices);
1126        assert_eq!(rt.completion_gate(), CompletionGate::Allowed);
1127    }
1128    #[test]
1129    fn retirement_only_touches_explicit_reconciled_ids() {
1130        let rt = baseline(8, 64);
1131        for id in ["old", "new", "uncollected"] {
1132            rt.authorize(id, &rt.foreground_model()).unwrap();
1133            rt.terminal_and_collect(id, WorkerTerminal::Completed)
1134                .unwrap();
1135        }
1136        rt.reconcile("old").unwrap();
1137        rt.reconcile("new").unwrap();
1138        rt.retire_reconciled(&["old".into(), "uncollected".into()].into_iter().collect());
1139        let inner = rt.inner.lock().unwrap();
1140        assert!(!inner.handles.contains_key("old"));
1141        assert!(inner.handles.contains_key("new"));
1142        assert!(inner.handles.contains_key("uncollected"));
1143        assert_eq!(inner.registry.total_dispatched(), 3);
1144    }
1145
1146    #[test]
1147    fn capacity_returns_only_after_reconcile_and_concurrency_still_applies() {
1148        let rt = baseline(1, 2);
1149        rt.authorize("a", &rt.foreground_model()).unwrap();
1150        let denial = rt.resolve_and_authorize("b", None).unwrap_err();
1151        assert_eq!(denial.code, "concurrency_limit");
1152        assert!(denial.remediation.contains("Changing model does not"));
1153        rt.terminal_and_collect("a", WorkerTerminal::Failed)
1154            .unwrap();
1155        rt.authorize("b", &rt.foreground_model()).unwrap();
1156        rt.terminal_and_collect("b", WorkerTerminal::TimedOut)
1157            .unwrap();
1158        assert_eq!(
1159            rt.resolve_and_authorize("c", None).unwrap_err().code,
1160            "total_limit"
1161        );
1162        rt.reconcile("a").unwrap();
1163        rt.resolve_and_authorize("c", None).unwrap();
1164        assert!(matches!(
1165            rt.completion_gate(),
1166            CompletionGate::Blocked { .. }
1167        ));
1168    }
1169}