a3s_code_core/agent_api/
capability_facade.rs

1use super::*;
2
3impl AgentSession {
4    // Advanced optional Queue API
5    // ========================================================================
6
7    /// Returns whether this session has an advanced lane queue configured.
8    pub fn has_queue(&self) -> bool {
9        QueueControl::from_session(self).has_queue()
10    }
11
12    /// Configure a lane's handler mode for explicit external/hybrid dispatch.
13    ///
14    /// Only effective when a queue is configured via `SessionOptions::with_queue_config`.
15    pub async fn set_lane_handler(
16        &self,
17        lane: SessionLane,
18        config: LaneHandlerConfig,
19    ) -> crate::error::Result<()> {
20        let _mutation = self.close_handle.extension_mutation.lock().await;
21        if self.is_closed() {
22            return Err(crate::error::CodeError::SessionClosed {
23                session_id: self.session_id.clone(),
24            });
25        }
26        QueueControl::from_session(self)
27            .set_lane_handler(lane, config)
28            .await;
29        if self.is_closed() {
30            return Err(crate::error::CodeError::SessionClosed {
31                session_id: self.session_id.clone(),
32            });
33        }
34        Ok(())
35    }
36
37    /// Complete an external queue task by ID.
38    ///
39    /// Returns `true` if the task was found and completed, `false` if not found.
40    pub async fn complete_external_task(&self, task_id: &str, result: ExternalTaskResult) -> bool {
41        QueueControl::from_session(self)
42            .complete_external_task(task_id, result)
43            .await
44    }
45
46    /// Get pending external queue tasks awaiting completion by an external handler.
47    pub async fn pending_external_tasks(&self) -> Vec<ExternalTask> {
48        QueueControl::from_session(self)
49            .pending_external_tasks()
50            .await
51    }
52
53    /// Get optional queue statistics (pending, active, external counts per lane).
54    pub async fn queue_stats(&self) -> SessionQueueStats {
55        QueueControl::from_session(self).stats().await
56    }
57
58    /// Get a metrics snapshot from the optional queue (if metrics are enabled).
59    pub async fn queue_metrics(&self) -> Option<MetricsSnapshot> {
60        QueueControl::from_session(self).metrics().await
61    }
62
63    /// Get dead letters from the optional queue's DLQ (if DLQ is enabled).
64    pub async fn dead_letters(&self) -> Vec<DeadLetter> {
65        QueueControl::from_session(self).dead_letters().await
66    }
67
68    // ========================================================================
69    // MCP API
70    // ========================================================================
71
72    /// Register all agents found in a directory with the live session.
73    ///
74    /// Scans `dir` for `*.yaml`, `*.yml`, and `*.md` agent definition files,
75    /// parses them, and adds each one to the shared `AgentRegistry` used by the
76    /// `task` tool. New agents are usable by the next admitted Run in the same
77    /// Session; an already admitted Run retains its frozen registry.
78    ///
79    /// Returns the number of agents successfully loaded from the directory.
80    pub fn register_agent_dir(&self, dir: &std::path::Path) -> crate::error::Result<usize> {
81        let agents = crate::subagent::load_agents_from_dir(dir);
82        self.close_handle.mutate_immediate(|| {
83            for agent in &agents {
84                self.ensure_compatibility_name_available(
85                    crate::capability::CapabilityKind::Agent,
86                    &agent.name,
87                )?;
88            }
89            let count = agents.len();
90            for agent in agents {
91                tracing::info!(
92                    session_id = %self.session_id,
93                    agent = agent.name,
94                    dir = %dir.display(),
95                    "Dynamically registered agent"
96                );
97                self.agent_registry.register(agent);
98            }
99            Ok(count)
100        })?
101    }
102
103    /// Register a disposable worker agent with the live session.
104    ///
105    /// The returned definition enters the `task` lookup and the model-facing
106    /// `task` definitions on the next admitted Run. Callers
107    /// can create discoverable reproducible workers without writing temporary
108    /// agent files or restarting the Session, while an active Run remains
109    /// generation-stable.
110    pub fn register_worker_agent(
111        &self,
112        spec: crate::subagent::WorkerAgentSpec,
113    ) -> crate::error::Result<crate::subagent::AgentDefinition> {
114        self.close_handle.mutate_immediate(|| {
115            self.ensure_compatibility_name_available(
116                crate::capability::CapabilityKind::Agent,
117                &spec.name,
118            )?;
119            Ok(SessionExtensionRuntime::from_session(self).register_worker_agent(spec))
120        })?
121    }
122
123    /// Register multiple disposable worker agents with the live session.
124    pub fn register_worker_agents<I>(
125        &self,
126        specs: I,
127    ) -> crate::error::Result<Vec<crate::subagent::AgentDefinition>>
128    where
129        I: IntoIterator<Item = crate::subagent::WorkerAgentSpec>,
130    {
131        let specs = specs.into_iter().collect::<Vec<_>>();
132        self.close_handle.mutate_immediate(|| {
133            for spec in &specs {
134                self.ensure_compatibility_name_available(
135                    crate::capability::CapabilityKind::Agent,
136                    &spec.name,
137                )?;
138            }
139            Ok(SessionExtensionRuntime::from_session(self).register_worker_agents(specs))
140        })?
141    }
142
143    /// Add or replace a skill in this live session.
144    ///
145    /// The Skill and `search_skills` tools observe the new definition
146    /// immediately, and the model-visible skills catalog observes it on the
147    /// next turn. Removing the live definition restores any session skill it
148    /// shadowed at installation time.
149    pub fn add_skill(&self, skill: Arc<crate::skills::Skill>) -> crate::error::Result<()> {
150        self.close_handle.mutate_immediate(|| {
151            self.ensure_compatibility_name_available(
152                crate::capability::CapabilityKind::Skill,
153                &skill.name,
154            )?;
155            SessionExtensionRuntime::from_session(self).add_skill(skill)
156        })?
157    }
158
159    /// Remove a skill previously installed through [`Self::add_skill`].
160    ///
161    /// This is a no-op when the name is not owned by the live session API; base
162    /// session skills and later host registrations are never removed.
163    pub fn remove_skill(&self, name: &str) -> crate::error::Result<()> {
164        self.close_handle
165            .mutate_immediate(|| SessionExtensionRuntime::from_session(self).remove_skill(name))
166    }
167
168    /// Return the names in the session's current live skill registry.
169    pub fn skill_names(&self) -> Vec<String> {
170        self.close_handle.skill_registry.list()
171    }
172
173    /// Add an MCP server to this session.
174    ///
175    /// Registers, connects, and makes all tools immediately available for the
176    /// agent to call. Tool names follow the convention `mcp__<name>__<tool>`.
177    ///
178    /// Returns the number of tools registered from the server.
179    pub async fn add_mcp_server(
180        &self,
181        config: crate::mcp::McpServerConfig,
182    ) -> crate::error::Result<usize> {
183        SessionExtensionRuntime::from_session(self)
184            .add_mcp_server(config)
185            .await
186    }
187
188    /// The session's tool executor, for installing agent-dir `tools/` entries
189    /// (e.g. a `kind = "script"` tool) into the live registry. Internal seam used
190    /// by [`serve::install_agent_dir_tools`](crate::serve::install_agent_dir_tools)
191    /// (the only caller, hence the `serve` gate).
192    #[cfg(feature = "serve")]
193    pub(crate) fn tool_executor(&self) -> &Arc<crate::tools::ToolExecutor> {
194        &self.tool_executor
195    }
196
197    /// Register a host-provided dynamic tool into the live session. Enables an
198    /// embedding app (e.g. the a3s-code CLI's login-gated `runtime` A3S Runtime
199    /// offload tool) to add a native tool at runtime; it enters the LLM's toolset
200    /// on the next run (`build_agent_loop` re-snapshots `definitions()` per run),
201    /// the same way MCP tools surface after `add_mcp_server`. Idempotent by name.
202    pub fn register_dynamic_tool(
203        &self,
204        tool: Arc<dyn crate::tools::Tool>,
205    ) -> crate::error::Result<()> {
206        self.close_handle.mutate_immediate(|| {
207            self.ensure_compatibility_name_available(
208                crate::capability::CapabilityKind::Tool,
209                tool.name(),
210            )?;
211            self.tool_executor.register_dynamic_tool(tool);
212            Ok(())
213        })?
214    }
215
216    /// Register the A3S Flow-backed dynamic workflow tool for this live session.
217    ///
218    /// The tool is named `dynamic_workflow`. It accepts a sandboxed JavaScript
219    /// PTC workflow script and executes it through
220    /// [`crate::DynamicWorkflowRuntime`], so A3S Flow owns workflow replay while
221    /// the script can still call A3S Code tools.
222    #[cfg(feature = "dynamic-workflow")]
223    pub fn register_dynamic_workflow_runtime(&self) -> crate::error::Result<()> {
224        self.close_handle.mutate_immediate(|| {
225            self.ensure_compatibility_name_available(
226                crate::capability::CapabilityKind::Tool,
227                "dynamic_workflow",
228            )?;
229            crate::tools::register_dynamic_workflow(self.tool_executor.registry());
230            Ok(())
231        })?
232    }
233
234    /// Coding-only builds omit the dynamic workflow module; keep the method so
235    /// hosts fail closed instead of losing the API surface at compile time.
236    #[cfg(not(feature = "dynamic-workflow"))]
237    pub fn register_dynamic_workflow_runtime(&self) -> crate::error::Result<()> {
238        self.close_handle.mutate_immediate(|| {
239            Err(crate::error::CodeError::Config(
240                "dynamic_workflow requires the advanced-harness / dynamic-workflow feature".into(),
241            ))
242        })?
243    }
244
245    /// Remove a previously host-registered dynamic tool by name (e.g. on logout).
246    /// No-op if no tool of that name is registered.
247    pub fn unregister_dynamic_tool(&self, name: &str) -> crate::error::Result<()> {
248        self.close_handle
249            .mutate_immediate(|| self.tool_executor.unregister_dynamic_tool(name))
250    }
251
252    /// Remove an MCP server from this session.
253    ///
254    /// Disconnects the server and unregisters all its tools from the executor.
255    /// No-op if the server was never added.
256    pub async fn remove_mcp_server(&self, server_name: &str) -> crate::error::Result<()> {
257        SessionExtensionRuntime::from_session(self)
258            .remove_mcp_server(server_name)
259            .await
260    }
261
262    /// Return current projected and compatibility MCP server status.
263    pub async fn mcp_status(
264        &self,
265    ) -> std::collections::HashMap<String, crate::mcp::McpServerStatus> {
266        SessionExtensionRuntime::from_session(self)
267            .mcp_status()
268            .await
269    }
270
271    /// Return the exact generation and digest currently visible to new Runs.
272    pub fn capability_catalog_stamp(&self) -> crate::capability::CapabilityCatalogStamp {
273        self.capability_catalog.current_stamp()
274    }
275
276    /// Verify that a new Run on this Session would receive the exact persisted
277    /// scoped capability catalog and authority ceiling.
278    pub fn ensure_recovery_capability_binding(
279        &self,
280        expected: &crate::capability::RunCapabilityBindingV1,
281    ) -> std::result::Result<(), crate::capability::RunCapabilityBindingError> {
282        super::agent_loop_runtime::validate_run_capability_binding(self, expected)
283    }
284
285    /// Prepare and atomically publish one complete host capability generation.
286    ///
287    /// Preparation may perform asynchronous work, but no value becomes visible
288    /// until every adapter has succeeded and the complete projection wins its
289    /// generation-and-digest compare-and-swap. A Use-backed batch publishes its
290    /// generation-specific lease provider in that same commit.
291    pub async fn apply_capability_batch(
292        &self,
293        batch: crate::capability::SessionCapabilityBatch,
294        cancellation: tokio_util::sync::CancellationToken,
295    ) -> std::result::Result<
296        crate::capability::CapabilityCommitReceipt,
297        crate::capability::CapabilityRuntimeError,
298    > {
299        let _mutation = self.close_handle.extension_mutation.lock().await;
300        if self.is_closed() {
301            return Err(crate::capability::CapabilityRuntimeError::SessionClosed);
302        }
303
304        let preparation_cancellation = tokio_util::sync::CancellationToken::new();
305        let prepared = tokio::select! {
306            biased;
307            _ = self.session_cancel.cancelled() => {
308                preparation_cancellation.cancel();
309                return Err(if self.is_closed() {
310                    crate::capability::CapabilityRuntimeError::SessionClosed
311                } else {
312                    crate::capability::CapabilityRuntimeError::Cancelled
313                });
314            }
315            _ = cancellation.cancelled() => {
316                preparation_cancellation.cancel();
317                return Err(crate::capability::CapabilityRuntimeError::Cancelled);
318            }
319            result = batch.prepare(
320                &self.capability_catalog,
321                preparation_cancellation.clone(),
322            ) => result?,
323        };
324        self.ensure_projected_mcp_server_names_available(prepared.projection()?)
325            .await?;
326
327        // The close boundary and Run pinning use this same short mutex. The
328        // prepared transaction holds no registry write lock while waiting.
329        let _publication = self
330            .close_handle
331            .immediate_extension_mutation
332            .lock()
333            .unwrap_or_else(std::sync::PoisonError::into_inner);
334        if self.is_closed() {
335            return Err(crate::capability::CapabilityRuntimeError::SessionClosed);
336        }
337        if cancellation.is_cancelled() || self.session_cancel.is_cancelled() {
338            return Err(crate::capability::CapabilityRuntimeError::Cancelled);
339        }
340        // The legacy public registry guard does not participate in the
341        // Session mutation gate. Keep its lock through catalog publication so
342        // direct mutation linearizes wholly before validation or after CAS.
343        let command_registry = self
344            .command_registry
345            .lock()
346            .unwrap_or_else(std::sync::PoisonError::into_inner);
347        let current_projection = self.capability_catalog.pin();
348        super::agent_loop_runtime::validate_capability_projection_runtime(
349            self,
350            prepared.projection()?,
351            &command_registry,
352        )?;
353        super::agent_loop_runtime::validate_capability_projection_transition(
354            self,
355            current_projection.projection(),
356            prepared.projection()?,
357        )?;
358        prepared.commit()
359    }
360
361    /// Apply one SDK-transported capability batch (SDK-CAP1 Skill slice).
362    ///
363    /// Cross-language hosts stage serializable Skill values. The session cancel
364    /// token bounds preparation; hosts that need cooperative abort should close
365    /// the session or cancel the owning Run.
366    pub async fn apply_sdk_capability_batch(
367        &self,
368        batch: crate::capability::SdkCapabilityBatchV1,
369    ) -> std::result::Result<
370        crate::capability::SdkCapabilityCommitReceiptV1,
371        crate::capability::SdkCapabilityBatchError,
372    > {
373        let session_batch = batch.into_session_batch()?;
374        let receipt = self
375            .apply_capability_batch(session_batch, self.session_cancel.child_token())
376            .await?;
377        Ok(crate::capability::SdkCapabilityCommitReceiptV1::from_receipt(&receipt))
378    }
379
380    /// Reconstruct one exact historical capability generation on an untouched
381    /// recovery Session.
382    ///
383    /// Unlike ordinary [`Self::apply_capability_batch`], this one-time path may
384    /// jump directly from the empty generation-zero catalog to the generation
385    /// named by `expected`. It never resolves packages or `latest`: the host
386    /// supplies every runtime adapter and any exact A3S Use lease provider in
387    /// `batch`, and Code verifies the resulting catalog plus authority ceiling
388    /// before publication.
389    pub async fn bootstrap_recovery_capability_batch(
390        &self,
391        expected: &crate::capability::RunCapabilityBindingV1,
392        batch: crate::capability::SessionCapabilityBatch,
393        cancellation: tokio_util::sync::CancellationToken,
394    ) -> std::result::Result<
395        crate::capability::CapabilityCommitReceipt,
396        crate::capability::CapabilityRuntimeError,
397    > {
398        let _mutation = self.close_handle.extension_mutation.lock().await;
399        if self.is_closed() {
400            return Err(crate::capability::CapabilityRuntimeError::SessionClosed);
401        }
402
403        let target_ceiling = self.capability_run_ceiling(batch.target())?;
404        expected
405            .ensure_matches(batch.target(), &target_ceiling)
406            .map_err(
407                |error| crate::capability::CapabilityRuntimeError::RecoveryBinding {
408                    message: error.to_string(),
409                },
410            )?;
411
412        let preparation_cancellation = tokio_util::sync::CancellationToken::new();
413        let prepared = tokio::select! {
414            biased;
415            _ = self.session_cancel.cancelled() => {
416                preparation_cancellation.cancel();
417                return Err(if self.is_closed() {
418                    crate::capability::CapabilityRuntimeError::SessionClosed
419                } else {
420                    crate::capability::CapabilityRuntimeError::Cancelled
421                });
422            }
423            _ = cancellation.cancelled() => {
424                preparation_cancellation.cancel();
425                return Err(crate::capability::CapabilityRuntimeError::Cancelled);
426            }
427            result = batch.prepare_recovery_bootstrap(
428                &self.capability_catalog,
429                preparation_cancellation.clone(),
430            ) => result?,
431        };
432        self.ensure_projected_mcp_server_names_available(prepared.projection()?)
433            .await?;
434
435        let _publication = self
436            .close_handle
437            .immediate_extension_mutation
438            .lock()
439            .unwrap_or_else(std::sync::PoisonError::into_inner);
440        if self.is_closed() {
441            return Err(crate::capability::CapabilityRuntimeError::SessionClosed);
442        }
443        if cancellation.is_cancelled() || self.session_cancel.is_cancelled() {
444            return Err(crate::capability::CapabilityRuntimeError::Cancelled);
445        }
446        let command_registry = self
447            .command_registry
448            .lock()
449            .unwrap_or_else(std::sync::PoisonError::into_inner);
450        let current_projection = self.capability_catalog.pin();
451        super::agent_loop_runtime::validate_capability_projection_runtime(
452            self,
453            prepared.projection()?,
454            &command_registry,
455        )?;
456        super::agent_loop_runtime::validate_capability_projection_transition(
457            self,
458            current_projection.projection(),
459            prepared.projection()?,
460        )?;
461        prepared.commit()
462    }
463
464    /// Drain prepared effects from failed or retired host generations.
465    pub async fn drain_capability_cleanup(&self) -> crate::capability::CapabilityCleanupReport {
466        self.capability_catalog.drain_cleanup().await
467    }
468
469    #[cfg(test)]
470    pub(crate) async fn admit_capability_run(
471        &self,
472    ) -> std::result::Result<
473        crate::capability::SessionCapabilityRun,
474        crate::capability::CapabilityRuntimeError,
475    > {
476        let projection = {
477            let _admission = self
478                .close_handle
479                .immediate_extension_mutation
480                .lock()
481                .unwrap_or_else(std::sync::PoisonError::into_inner);
482            if self.is_closed() {
483                return Err(crate::capability::CapabilityRuntimeError::SessionClosed);
484            }
485            self.capability_catalog.pin()
486        };
487        let ceiling = self.capability_run_ceiling(projection.projection().set())?;
488        crate::capability::SessionCapabilityRun::admit(
489            projection,
490            "active",
491            "active",
492            ceiling,
493            self.session_cancel.child_token(),
494        )
495        .await
496    }
497
498    pub(super) fn capability_run_ceiling(
499        &self,
500        set: &crate::capability::CapabilitySet,
501    ) -> std::result::Result<
502        crate::capability::CapabilityCeiling,
503        crate::capability::CapabilityRuntimeError,
504    > {
505        let mut governance = crate::capability::GovernanceCapabilityCeiling::none_required();
506        if self.config.permission_checker.is_some() || self.config.permission_policy.is_some() {
507            governance = governance.require_permission_guard();
508        }
509        if self.config.confirmation_manager.is_some() || self.config.confirmation_policy.is_some() {
510            governance = governance.require_confirmation_guard();
511        }
512        if self.config.security_provider.is_some() {
513            governance = governance.require_security_guard();
514        }
515        if self.config.budget_guard.is_some() || self.budget_guard().is_some() {
516            governance = governance.require_budget_guard();
517        }
518        if self.config.enforce_active_skill_tool_restrictions {
519            governance = governance.require_active_skill_restrictions();
520        }
521        let execution = crate::capability::CapabilityExecutionCeiling::new(
522            self.config.max_tool_rounds,
523            self.config.max_parallel_tasks,
524            self.config.tool_timeout_ms,
525            self.config.llm_api_timeout_ms,
526            self.config.max_execution_time_ms,
527        )?;
528        crate::capability::CapabilityCeiling::all(
529            set,
530            crate::capability::WorkspaceCapabilityCeiling::all(),
531            governance,
532            execution,
533        )
534        .map_err(Into::into)
535    }
536
537    pub(super) fn ensure_compatibility_name_available(
538        &self,
539        kind: crate::capability::CapabilityKind,
540        public_name: &str,
541    ) -> crate::error::Result<()> {
542        let projection = self.capability_catalog.pin();
543        if projection
544            .projection()
545            .iter()
546            .any(|(_, value)| match value {
547                crate::capability::CapabilityValue::Mcp(binding)
548                    if kind == crate::capability::CapabilityKind::Tool =>
549                {
550                    binding.contains_public_tool_name(public_name)
551                }
552                crate::capability::CapabilityValue::Agent(agent) => {
553                    kind == crate::capability::CapabilityKind::Agent
554                        && crate::subagent::agent_names_conflict(&agent.name, public_name)
555                }
556                _ => value.kind() == kind && value.public_name() == Some(public_name),
557            })
558        {
559            return Err(
560                crate::capability::CapabilityRuntimeError::RuntimeNameConflict {
561                    kind,
562                    public_name: public_name.to_owned(),
563                }
564                .into(),
565            );
566        }
567        Ok(())
568    }
569
570    async fn ensure_projected_mcp_server_names_available(
571        &self,
572        projection: &crate::capability::CapabilityProjection,
573    ) -> std::result::Result<(), crate::capability::CapabilityRuntimeError> {
574        let server_names = projection
575            .iter()
576            .filter_map(|(_, value)| match value {
577                crate::capability::CapabilityValue::Mcp(binding) => {
578                    Some(binding.server_name().to_owned())
579                }
580                _ => None,
581            })
582            .collect::<Vec<_>>();
583        for server_name in server_names {
584            for manager in &self.mcp_managers {
585                if manager.contains_server(&server_name).await {
586                    return Err(
587                        crate::capability::CapabilityRuntimeError::RuntimeNameConflict {
588                            kind: crate::capability::CapabilityKind::Mcp,
589                            public_name: server_name,
590                        },
591                    );
592                }
593            }
594        }
595        Ok(())
596    }
597}
a3s_code_core/agent_api/capability_facade.rs

a3s_code_core/agent_api/
capability_facade.rs
