agents_kit/builder.rs
1use std::collections::BTreeSet;
2use std::sync::{Arc, Mutex};
3use std::time::Duration;
4use tools_kit::ToolsetBackend;
5
6use providers_kit::ProviderBackend;
7use tokio::sync::oneshot;
8use tokio::task::JoinHandle;
9
10use crate::agent_storage::AgentStorageClaim;
11use crate::backends::{codex, custom, opencode};
12use crate::{
13 Agent, AgentBackend, AgentConfig, AgentId, AgentInfo, AgentStorage, BackendConfig, KitError,
14 KnowledgeOwnerId, ManagedKnowledge, ToolSelection,
15};
16use storage_kit::{Database, DatabaseBackend, StateStore, StateStoreBackend};
17
18/// Configures creation or reopening of one durable [`Agent`].
19///
20/// The builder accepts already-built provider, tool, and storage values without
21/// exposing their implementation details through `agents-kit`.
22///
23/// The hierarchy preserves concrete `S`, `D`, `P` and `T` types, with no Clone
24/// requirement on caller backends or `T::Handler`. New identity creation and
25/// durable reopen use separate constructors; reopening resupplies integrations
26/// instead of reconstructing credentials or lower configuration from records.
27///
28/// Native Codex/OpenCode can select `P` as static composition context without
29/// owning a provider value. Custom requires a value through [`Self::with_provider`].
30/// The builder owns supplied values, while [`Self::build`] accepts construction
31/// and cleanup work. Await consuming [`AgentBackend::shutdown`] on the built Agent
32/// while its runtime remains driven.
33///
34/// The local example uses memory Lance and a passively constructed Provider. It
35/// selects an exact model without verifying a running service or generating output.
36///
37/// See the [construction example](crate) and [`Self::reopen`] for durable reopen.
38pub struct AgentBuilder<
39 S: StateStoreBackend = StateStore,
40 D: DatabaseBackend = Database,
41 P: ProviderBackend = providers_kit::Provider,
42 T: ToolsetBackend = tools_kit::Toolset,
43> {
44 config: AgentConfig,
45 storage: AgentStorage<S>,
46 mode: AgentBuildMode,
47 provider: Option<P>,
48 toolsets: Vec<Arc<T>>,
49 knowledge_database: Option<D>,
50}
51
52enum AgentBuildMode {
53 Create,
54 Reopen(AgentId),
55}
56
57impl<S: StateStoreBackend> AgentBuilder<S, Database, providers_kit::Provider> {
58 /// Configures a new durable agent over the required control-plane store.
59 ///
60 /// This records creation intent without opening a backend. [`Self::build`] creates
61 /// a fresh durable identity. Storage is required up front and retains its exact
62 /// associated transaction type. For an existing identity use [`Self::reopen`].
63 ///
64 /// See the [construction example](crate). To choose caller-owned lower types, see [`Self::with_provider`] and [`Self::with_toolset`].
65 pub fn create(config: AgentConfig, storage: AgentStorage<S>) -> Self {
66 Self::create_for_backends::<providers_kit::Provider, tools_kit::Toolset>(config, storage)
67 }
68
69 /// Configures a new agent with fixed Provider and Toolset implementation types.
70 ///
71 /// Codex and OpenCode retain their native providers, so this selects only
72 /// the static provider type shared with other agents in a caller-owned
73 /// composition. Custom agents receive their provider value through
74 /// [`Self::with_provider`]. Repeated [`Self::with_toolset`] calls append T2
75 /// values to one homogeneous inventory without normalization or a Clone bound.
76 ///
77 /// See [`Self::reopen_for_backends`] for typed creation and reopening.
78 /// [`Self::with_provider`] and [`Self::with_toolset`] accept caller-owned lower types.
79 pub fn create_for_backends<P2: ProviderBackend, T2: ToolsetBackend>(
80 config: AgentConfig,
81 storage: AgentStorage<S>,
82 ) -> AgentBuilder<S, Database, P2, T2> {
83 AgentBuilder {
84 config,
85 storage,
86 mode: AgentBuildMode::Create,
87 provider: None,
88 toolsets: Vec::new(),
89 knowledge_database: None,
90 }
91 }
92
93 /// Configures reopening of an existing durable agent.
94 ///
95 /// Selects an existing identity; it never falls back to creating one. Build checks
96 /// stored compatibility and the storage clone family's live-ownership claim.
97 /// Resupply fresh provider/tool/database values required by that identity.
98 /// An AgentStorage clone shares coordination without requiring `S: Clone`.
99 ///
100 /// # Examples
101 ///
102 /// This owns one Lance directory and reopens it through a fresh connection. Memory
103 /// stores instead require the same retained storage family. Shutdown is awaited
104 /// before inspection errors are propagated; unsuccessful runs retain their files.
105 ///
106 /// ```
107 /// use std::{num::NonZeroUsize, time::Duration};
108 /// use agents_kit::{AgentBackend, AgentBuilder, AgentStorage};
109 /// use agents_kit::config::{AgentConfig, CustomAgentConfig, CustomLoopLimits, CustomToolApproval};
110 /// use agents_kit::{model::ModelSelector, tool::ToolSelection};
111 /// use providers_kit::{Provider, ProviderConfig};
112 /// use providers_kit::config::{OllamaConfig, ProviderKind, TransportConfig};
113 /// use storage_kit::{StateStore, state::config::{LanceStateStoreConfig, StateStoreConfig}};
114 /// let directory = std::env::temp_dir().join(format!("agent-reopen-doc-{}", uuid::Uuid::new_v4()));
115 /// std::fs::create_dir(&directory)?;
116 /// let store_config = StateStoreConfig::lance(LanceStateStoreConfig::directory(&directory));
117 /// # tokio::runtime::Builder::new_current_thread().enable_all().build()?.block_on(async {
118 /// let custom = CustomAgentConfig::new(ModelSelector::new("exact-model")?, ToolSelection::None,
119 /// CustomLoopLimits::new(4, 0, 1)?, CustomToolApproval::Deny)?;
120 /// let config = AgentConfig::new(custom.into(), Duration::from_secs(10), NonZeroUsize::MIN)?;
121 /// let provider_config = ProviderConfig::new(
122 /// ProviderKind::Ollama(OllamaConfig::local()), TransportConfig::default(),
123 /// );
124 /// let store = StateStore::connect(store_config.clone()).await?;
125 /// let agent = AgentBuilder::create(config.clone(), AgentStorage::new(store))
126 /// .with_provider(Provider::from_config(provider_config.clone())?).build().await?;
127 /// let id = agent.id();
128 /// agent.shutdown().await?;
129 ///
130 /// let store = StateStore::connect(store_config).await?;
131 /// let reopened = AgentBuilder::reopen(config, AgentStorage::new(store), id.clone())
132 /// .with_provider(Provider::from_config(provider_config)?).build().await?;
133 /// let result = reopened.info().await;
134 /// let cleanup = reopened.shutdown().await;
135 /// let info = match (result, cleanup) {
136 /// (Ok(info), Ok(())) => Ok(info),
137 /// (Err(primary), Err(cleanup)) => Err(primary.with_supplementary(cleanup)),
138 /// (Err(error), _) | (_, Err(error)) => Err(error),
139 /// }?;
140 /// assert_eq!(info.id(), &id);
141 /// # Ok::<(), Box<dyn std::error::Error>>(())
142 /// # })?;
143 /// // Only successful shutdown and readback permit deleting this owned directory.
144 /// std::fs::remove_dir_all(directory)?;
145 /// # Ok::<(), Box<dyn std::error::Error>>(())
146 /// ```
147 pub fn reopen(config: AgentConfig, storage: AgentStorage<S>, id: AgentId) -> Self {
148 Self::reopen_for_backends::<providers_kit::Provider, tools_kit::Toolset>(
149 config, storage, id,
150 )
151 }
152
153 /// Configures reopening with fixed Provider and Toolset implementation types.
154 ///
155 /// This is the reopening counterpart to [`Self::create_for_backends`]. It
156 /// selects no provider value and performs no provider operation.
157 ///
158 /// # Examples
159 ///
160 /// ```
161 /// use std::{num::NonZeroUsize, time::Duration};
162 /// use agents_kit::{AgentBackend, AgentBuilder, AgentStorage};
163 /// use agents_kit::config::{AgentConfig, CustomAgentConfig, CustomLoopLimits, CustomToolApproval};
164 /// use agents_kit::{model::ModelSelector, tool::ToolSelection};
165 /// use providers_kit::{Provider, ProviderConfig, config::{OllamaConfig, ProviderKind, TransportConfig}};
166 /// use storage_kit::{StateStore, state::config::{LanceStateStoreConfig, StateStoreConfig}};
167 /// # let runtime = tokio::runtime::Builder::new_current_thread().enable_all().build()?;
168 /// # runtime.block_on(async {
169 /// let store = StateStore::connect(StateStoreConfig::Lance(LanceStateStoreConfig::memory())).await?;
170 /// let storage = AgentStorage::new(store);
171 /// let custom = CustomAgentConfig::new(ModelSelector::new("exact-model")?, ToolSelection::None,
172 /// CustomLoopLimits::new(4, 0, 1)?, CustomToolApproval::Deny)?;
173 /// let config = AgentConfig::new(custom.into(), Duration::from_secs(10), NonZeroUsize::MIN)?;
174 /// let provider_config = ProviderConfig::new(ProviderKind::Ollama(OllamaConfig::local()), TransportConfig::default());
175 /// let provider = Provider::from_config(provider_config.clone())?;
176 /// let agent = AgentBuilder::create_for_backends::<Provider, tools_kit::Toolset>(
177 /// config.clone(), storage.clone(),
178 /// )
179 /// .with_provider(provider).build().await?;
180 /// let id = agent.id();
181 /// agent.shutdown().await?;
182 /// let builder = AgentBuilder::reopen_for_backends::<Provider, tools_kit::Toolset>(
183 /// config, storage, id.clone(),
184 /// );
185 /// let reopened = builder.with_provider(Provider::from_config(provider_config)?).build().await?;
186 /// let reopened_id = reopened.id();
187 /// reopened.shutdown().await?;
188 /// assert_eq!(reopened_id, id);
189 /// # Ok::<(), Box<dyn std::error::Error>>(())
190 /// # })?;
191 /// # Ok::<(), Box<dyn std::error::Error>>(())
192 /// ```
193 pub fn reopen_for_backends<P2: ProviderBackend, T2: ToolsetBackend>(
194 config: AgentConfig,
195 storage: AgentStorage<S>,
196 id: AgentId,
197 ) -> AgentBuilder<S, Database, P2, T2> {
198 let mut builder = Self::create_for_backends::<P2, T2>(config, storage);
199 builder.mode = AgentBuildMode::Reopen(id);
200 builder
201 }
202}
203
204impl<S: StateStoreBackend, D: DatabaseBackend, P: ProviderBackend, T: ToolsetBackend>
205 AgentBuilder<S, D, P, T>
206{
207 /// Supplies the already-built provider required by the Custom backend.
208 ///
209 /// Consumes the provider and changes the builder's concrete provider parameter;
210 /// accepted Toolsets and the Database remain attached. Replacing an earlier value
211 /// drops that value, so complete any required lower lifecycle work before replacing
212 /// it. Codex/OpenCode reject a supplied provider value during build; use typed
213 /// construction to select their static provider parameter without a value.
214 ///
215 /// See the [construction example](crate).
216 ///
217 /// Direct implementations retain `P::Model` and `T::Handler`; neither `T` nor
218 /// its handler requires Clone. See [`providers_kit::ProviderBackend`] and
219 /// [`tools_kit::ToolsetBackend`] for runnable implementations, and the
220 /// [Custom composition examples](https://github.com/jecsand838/ai-kit/blob/main/agents-kit/tests/custom_mock_composition.rs)
221 /// for their use together with the real agent loop. Other lower operations
222 /// keep their owning trait's explicit unsupported behavior.
223 pub fn with_provider<P2: ProviderBackend>(self, provider: P2) -> AgentBuilder<S, D, P2, T> {
224 AgentBuilder {
225 config: self.config,
226 storage: self.storage,
227 mode: self.mode,
228 provider: Some(provider),
229 toolsets: self.toolsets,
230 knowledge_database: self.knowledge_database,
231 }
232 }
233
234 /// Appends one already-built executable Toolset implementation.
235 ///
236 /// T remains fixed by construction. Changing the Provider or Database keeps
237 /// every accepted Toolset; its unique Handler remains associated as `T::Handler`.
238 ///
239 /// Each call appends a value of the construction-selected `T`; use caller-owned
240 /// exhaustive Toolset/Handler enums for heterogeneous sources. Custom executes the
241 /// associated Handler directly. Build validates unique effective names and selected
242 /// tools before publication. Native support is backend-specific: OpenCode rejects
243 /// supplied executable toolsets. Builder transfer alone never invokes a handler.
244 ///
245 /// # Examples
246 ///
247 /// ```
248 /// use std::{num::NonZeroUsize, time::Duration};
249 /// use agents_kit::{AgentBackend, AgentBuilder, AgentStorage};
250 /// use agents_kit::config::{AgentConfig, CustomAgentConfig, CustomLoopLimits, CustomToolApproval};
251 /// use agents_kit::{model::ModelSelector, tool::ToolSelection};
252 /// use providers_kit::{Provider, ProviderConfig, config::{OllamaConfig, ProviderKind, TransportConfig}};
253 /// use storage_kit::{StateStore, state::config::{LanceStateStoreConfig, StateStoreConfig}};
254 /// # let runtime = tokio::runtime::Builder::new_current_thread().enable_all().build()?;
255 /// # runtime.block_on(async {
256 /// let store = StateStore::connect(StateStoreConfig::Lance(LanceStateStoreConfig::memory())).await?;
257 /// let storage = AgentStorage::new(store);
258 /// let custom = CustomAgentConfig::new(ModelSelector::new("exact-model")?, ToolSelection::None,
259 /// CustomLoopLimits::new(4, 0, 1)?, CustomToolApproval::Deny)?;
260 /// let config = AgentConfig::new(custom.into(), Duration::from_secs(10), NonZeroUsize::MIN)?;
261 /// let provider_config = ProviderConfig::new(ProviderKind::Ollama(OllamaConfig::local()), TransportConfig::default());
262 /// let provider = Provider::from_config(provider_config.clone())?;
263 /// use tools_kit::{Toolset, declaration::{ToolInfo, ToolName}, execution::ToolResult,
264 /// local::ToolRegistration, schema::JsonSchema};
265 /// let tools = Toolset::local(vec![ToolRegistration::function(
266 /// ToolInfo::new(ToolName::new("ping")?, JsonSchema::strict(serde_json::json!({"type":"object"}))),
267 /// |_| async { Ok(ToolResult::text("pong")) },
268 /// )])?;
269 /// let agent = AgentBuilder::create(config, storage)
270 /// .with_toolset(tools).with_provider(provider).build().await?;
271 /// let result = agent.list_workspaces().await;
272 /// // Shutdown settles the transferred Toolset and Provider even if listing failed.
273 /// let cleanup = agent.shutdown().await;
274 /// let workspaces = match (result, cleanup) {
275 /// (Ok(value), Ok(())) => Ok(value),
276 /// (Err(primary), Err(cleanup)) => Err(primary.with_supplementary(cleanup)),
277 /// (Err(error), _) | (_, Err(error)) => Err(error),
278 /// }?;
279 /// assert!(workspaces.is_empty());
280 /// # Ok::<(), Box<dyn std::error::Error>>(())
281 /// # })?;
282 /// # Ok::<(), Box<dyn std::error::Error>>(())
283 /// ```
284 pub fn with_toolset(mut self, toolset: T) -> Self {
285 self.toolsets.push(Arc::new(toolset));
286 self
287 }
288
289 /// Supplies the Database used by the private managed-knowledge owner.
290 ///
291 /// Consumes a direct Database implementation and changes only the builder's
292 /// concrete `D` parameter. Existing Toolsets and Provider remain attached. This
293 /// Database is separate from AgentStorage's control-plane records and exposes no
294 /// public knowledge runtime. Reopen must resupply a compatible Database when the
295 /// durable identity requires managed knowledge.
296 ///
297 /// See [`IngestKnowledgeRequest`](crate::knowledge::IngestKnowledgeRequest) for request construction and the [Custom composition examples](https://github.com/jecsand838/ai-kit/blob/main/agents-kit/tests/custom_mock_composition.rs) for caller-owned knowledge storage.
298 pub fn with_knowledge_database<D2: DatabaseBackend>(
299 self,
300 database: D2,
301 ) -> AgentBuilder<S, D2, P, T> {
302 AgentBuilder {
303 config: self.config,
304 storage: self.storage,
305 mode: self.mode,
306 provider: self.provider,
307 toolsets: self.toolsets,
308 knowledge_database: Some(database),
309 }
310 }
311
312 /// Validates the complete configuration and creates or reopens the agent.
313 ///
314 /// Takes ownership of construction after acceptance. On failure, prepared resources
315 /// are cleaned in reverse ownership order within the lifecycle deadline; the first
316 /// failure remains primary and later failures remain supplementary. Losing the
317 /// caller observation does not detach accepted build/cleanup ownership.
318 ///
319 /// On success the unique Agent owns its runtime and transferred integrations.
320 /// Keep Tokio and lower runtimes driven through consuming shutdown. Custom model
321 /// selection alone does not verify remote model availability; later provider
322 /// operations own service validation.
323 ///
324 /// # Errors
325 ///
326 /// Returns InvalidConfig for invalid backend/integration combinations (including a
327 /// missing Custom provider), Unsupported for a supplied integration the selected
328 /// backend cannot represent, and Conflict/Busy for incompatible or already-owned
329 /// reopen state. Missing identities use NotFound; malformed durable evidence uses
330 /// InvalidState. Process, transport, protocol, storage and deadline failures retain
331 /// their owning portable classification. Inspect [`KitError::kind`] and
332 /// [`KitError::context`]; build and reopen carry distinct operation context.
333 ///
334 /// This example requires an installed, configured `codex` executable and its real
335 /// App Server process. It does not submit a model prompt.
336 ///
337 /// # Examples
338 ///
339 /// ```no_run
340 /// use std::{num::NonZeroUsize, time::Duration};
341 /// use agents_kit::{AgentBackend, AgentBuilder, AgentStorage};
342 /// use agents_kit::config::{AgentConfig, CodexConfig};
343 /// use storage_kit::{StateStore, state::config::{LanceStateStoreConfig, StateStoreConfig}};
344 /// # let runtime = tokio::runtime::Builder::new_current_thread().enable_all().build()?;
345 /// # runtime.block_on(async {
346 /// let config = AgentConfig::new(CodexConfig::spawn(Duration::from_secs(15),
347 /// Duration::from_secs(60))?.into(), Duration::from_secs(30), NonZeroUsize::MIN)?;
348 /// let store = StateStore::connect(StateStoreConfig::Lance(LanceStateStoreConfig::memory())).await?;
349 /// let agent = AgentBuilder::create(config, AgentStorage::new(store)).build().await?;
350 /// let observed = agent.info().await;
351 /// let cleanup = agent.shutdown().await;
352 /// let info = match (observed, cleanup) {
353 /// (Ok(value), Ok(())) => Ok(value),
354 /// (Err(primary), Err(cleanup)) => Err(primary.with_supplementary(cleanup)),
355 /// (Err(error), _) | (_, Err(error)) => Err(error),
356 /// }?;
357 /// assert!(!info.id().as_str().is_empty());
358 /// # Ok::<(), Box<dyn std::error::Error>>(())
359 /// # })?;
360 /// # Ok::<(), Box<dyn std::error::Error>>(())
361 /// ```
362 pub async fn build(self) -> Result<Agent<S, D, P, T>, KitError> {
363 let operation = match self.mode {
364 AgentBuildMode::Create => crate::error::Operation::Build,
365 AgentBuildMode::Reopen(_) => crate::error::Operation::Reopen,
366 };
367 AgentBuildPublication::<Agent<S, D, P, T>>::run(AgentBuildOwner::new(self))
368 .await
369 .map_err(|error| {
370 error
371 .at_operation(operation)
372 .in_phase(crate::error::ErrorPhase::Admission)
373 })
374 }
375}
376
377struct AgentBuildOwner<
378 S: StateStoreBackend,
379 D: DatabaseBackend,
380 P: ProviderBackend,
381 T: ToolsetBackend,
382> {
383 config: Option<AgentConfig>,
384 mode: AgentBuildMode,
385 provider: Option<P>,
386 toolsets: Vec<Arc<T>>,
387 storage: Option<AgentStorage<S>>,
388 knowledge_database: Option<D>,
389 claim: Option<AgentStorageClaim>,
390 agent: Option<Agent<S, D, P, T>>,
391 cleanup_timeout: Duration,
392 tracing_enabled: bool,
393}
394
395impl<S: StateStoreBackend, D: DatabaseBackend, P: ProviderBackend, T: ToolsetBackend>
396 AgentBuildOwner<S, D, P, T>
397{
398 fn new(builder: AgentBuilder<S, D, P, T>) -> Self {
399 let cleanup_timeout = builder.config.lifecycle_timeout();
400 let tracing_enabled = builder.config.tracing_enabled();
401 let AgentBuilder {
402 config,
403 storage,
404 mode,
405 provider,
406 toolsets,
407 knowledge_database,
408 } = builder;
409 Self {
410 config: Some(config),
411 mode,
412 provider,
413 toolsets,
414 storage: Some(storage),
415 knowledge_database,
416 claim: None,
417 agent: None,
418 cleanup_timeout,
419 tracing_enabled,
420 }
421 }
422
423 async fn build(mut self) -> Result<Agent<S, D, P, T>, KitError> {
424 if let Err(error) = self.validate().await {
425 return Err(self.cleanup_after_error(error).await);
426 }
427 let lifecycle_deadline = match tokio::time::Instant::now().checked_add(self.cleanup_timeout)
428 {
429 Some(deadline) => deadline,
430 None => {
431 return Err(self
432 .cleanup_after_error(KitError::invalid_config(
433 "lifecycle timeout cannot be represented",
434 ))
435 .await);
436 }
437 };
438
439 let storage = self
440 .storage
441 .as_ref()
442 .expect("validated AgentStorage must be present")
443 .clone();
444 let reopen = matches!(&self.mode, AgentBuildMode::Reopen(_));
445 let info = match &self.mode {
446 AgentBuildMode::Reopen(id) => match storage.agent_info(id.clone()).await {
447 Ok(info)
448 if self
449 .config
450 .as_ref()
451 .expect("build configuration must remain owned")
452 .display_name()
453 .is_some_and(|name| Some(name) != info.display_name()) =>
454 {
455 return Err(self
456 .cleanup_after_error(KitError::conflict(
457 "configured agent display name differs from durable identity",
458 ))
459 .await);
460 }
461 Ok(info) => info,
462 Err(error) => return Err(self.cleanup_after_error(error).await),
463 },
464 AgentBuildMode::Create => match AgentInfo::new(
465 AgentId::new(),
466 self.config
467 .as_ref()
468 .expect("build configuration must remain owned")
469 .display_name()
470 .map(str::to_owned),
471 ) {
472 Ok(info) => info,
473 Err(error) => return Err(self.cleanup_after_error(error).await),
474 },
475 };
476
477 self.claim = match storage.claim_agent(info.id()) {
478 Ok(claim) => Some(claim),
479 Err(error) => return Err(self.cleanup_after_error(error).await),
480 };
481 let knowledge = self
482 .knowledge_database
483 .take()
484 .map(|database| ManagedKnowledge::new(storage.clone(), database));
485 if reopen {
486 if let Err(error) = storage
487 .settle_knowledge_forgets(knowledge.as_ref(), info.id(), lifecycle_deadline)
488 .await
489 {
490 return Err(self.cleanup_after_error(error).await);
491 }
492 if let Some(knowledge) = &knowledge {
493 let owner = KnowledgeOwnerId::Agent {
494 agent_id: info.id().clone(),
495 };
496 let verification = match tokio::time::timeout_at(
497 lifecycle_deadline,
498 knowledge.verify(owner),
499 )
500 .await
501 {
502 Ok(result) => result,
503 Err(_) => Err(KitError::timeout(
504 "knowledge reconciliation deadline elapsed",
505 )),
506 };
507 if let Err(error) = verification {
508 return Err(self.cleanup_after_error(error).await);
509 }
510 }
511 }
512
513 let config = self
514 .config
515 .take()
516 .expect("build configuration must transfer exactly once");
517 let backend = config.backend().clone();
518 let claim = self
519 .claim
520 .take()
521 .expect("validated live-agent claim must transfer exactly once");
522 let built = match backend {
523 BackendConfig::Codex(_) => {
524 let toolsets = std::mem::take(&mut self.toolsets);
525 codex::agent::CodexAgent::build(
526 info, config, storage, knowledge, claim, toolsets, reopen,
527 )
528 .await
529 .map(Agent::codex)
530 }
531 BackendConfig::OpenCode(_) => opencode::agent::OpenCodeAgent::build(
532 info, config, storage, knowledge, claim, reopen,
533 )
534 .await
535 .map(Agent::opencode),
536 BackendConfig::Custom(_) => {
537 let composition = custom::CustomAgent::<S, D, P, T>::compose(
538 self.provider
539 .take()
540 .expect("validated provider must transfer exactly once"),
541 std::mem::take(&mut self.toolsets),
542 );
543 custom::CustomAgent::build(
544 info,
545 config,
546 storage,
547 knowledge,
548 claim,
549 composition,
550 reopen,
551 )
552 .await
553 .map(Agent::custom)
554 }
555 };
556 let agent = match built {
557 Ok(agent) => agent,
558 Err(error) => return Err(self.cleanup_after_error(error).await),
559 };
560 self.agent = Some(agent);
561
562 Ok(self
563 .agent
564 .take()
565 .expect("constructed Agent must publish exactly once"))
566 }
567
568 async fn validate(&mut self) -> Result<(), KitError> {
569 let agent_config = self
570 .config
571 .as_ref()
572 .expect("build configuration must remain owned");
573 agent_config.validate()?;
574 if tokio::time::Instant::now()
575 .checked_add(agent_config.lifecycle_timeout())
576 .is_none()
577 {
578 return Err(KitError::invalid_config(
579 "lifecycle timeout cannot be represented",
580 ));
581 }
582 let has_provider = self.provider.is_some();
583 let has_tools = !self.toolsets.is_empty();
584 match agent_config.backend() {
585 BackendConfig::Codex(_) => {
586 if has_provider {
587 return Err(KitError::invalid_config(
588 "selected runtime does not accept a model runtime",
589 ));
590 }
591 codex::tools::CodexTools::<S, T>::validate_registrations(&self.toolsets).await?;
592 }
593 BackendConfig::OpenCode(_) => {
594 if has_provider {
595 return Err(KitError::invalid_config(
596 "selected runtime does not accept a model runtime",
597 ));
598 }
599 if has_tools {
600 return Err(KitError::unsupported(
601 "selected runtime does not accept executable tools",
602 ));
603 }
604 }
605 BackendConfig::Custom(config) => {
606 if !has_provider {
607 return Err(KitError::invalid_config(
608 "selected runtime requires a configured model runtime",
609 ));
610 }
611 let catalog = self.validate_catalog().await?;
612 Self::validate_custom_selection(config.default_tool_selection(), &catalog)?;
613 }
614 }
615 Ok(())
616 }
617
618 async fn validate_catalog(&self) -> Result<BTreeSet<String>, KitError> {
619 let mut executable_names = BTreeSet::new();
620 for toolset in &self.toolsets {
621 for tool in toolset.list_tools().await? {
622 let name = tool.name.as_str().to_owned();
623 if !executable_names.insert(name.clone()) {
624 return Err(KitError::invalid_config(
625 "configured tool names must be unique",
626 ));
627 }
628 }
629 }
630
631 Ok(executable_names)
632 }
633
634 fn validate_custom_selection(
635 selection: &ToolSelection,
636 executable_names: &BTreeSet<String>,
637 ) -> Result<(), KitError> {
638 selection.validate()?;
639 match selection {
640 ToolSelection::Auto | ToolSelection::None => Ok(()),
641 ToolSelection::Required if executable_names.is_empty() => {
642 Err(KitError::invalid_config(
643 "required tool selection needs at least one configured tool",
644 ))
645 }
646 ToolSelection::Required => Ok(()),
647 ToolSelection::Allowed(names) | ToolSelection::RequiredOneOf(names) => {
648 Self::require_known_names(names, executable_names, "default selected tool")
649 }
650 }
651 }
652
653 fn require_known_names(
654 requested: &[String],
655 available: &BTreeSet<String>,
656 label: &str,
657 ) -> Result<(), KitError> {
658 let missing = requested
659 .iter()
660 .filter(|name| !available.contains(*name))
661 .cloned()
662 .collect::<Vec<_>>();
663 if missing.is_empty() {
664 Ok(())
665 } else {
666 Err(KitError::invalid_config(format!(
667 "{label} names must exist in the supplied Toolsets"
668 )))
669 }
670 }
671
672 async fn cleanup_after_error(mut self, mut primary: KitError) -> KitError {
673 let cleanup_deadline = tokio::time::Instant::now()
674 .checked_add(self.cleanup_timeout)
675 .unwrap_or_else(tokio::time::Instant::now);
676 if let Some(agent) = self.agent.take()
677 && let Err(cleanup_error) = agent.shutdown().await
678 {
679 if self.tracing_enabled {
680 tracing::warn!(
681 cleanup_target = "agent",
682 "agent build cleanup did not complete successfully"
683 );
684 }
685 primary = primary
686 .with_supplementary(cleanup_error.at_phase(crate::error::ErrorPhase::Cleanup));
687 }
688
689 drop(self.claim.take());
690 while let Some(toolset) = self.toolsets.pop() {
691 let cleanup =
692 tokio::time::timeout_at(cleanup_deadline, async { toolset.shutdown().await }).await;
693 if !matches!(cleanup, Ok(Ok(()))) && self.tracing_enabled {
694 tracing::warn!(
695 cleanup_target = "toolset",
696 "agent build cleanup did not complete successfully"
697 );
698 }
699 let error = match cleanup {
700 Ok(Ok(())) => None,
701 Ok(Err(error)) => Some(KitError::from(error)),
702 Err(_) => Some(
703 KitError::timeout("tool cleanup deadline elapsed").with_detail(
704 crate::error::ErrorDetail::DeadlineElapsed {
705 kind: crate::error::DeadlineKind::Lifecycle,
706 timeout: Some(self.cleanup_timeout),
707 },
708 ),
709 ),
710 };
711 if let Some(error) = error {
712 primary =
713 primary.with_supplementary(error.at_phase(crate::error::ErrorPhase::Cleanup));
714 }
715 }
716 if let Some(provider) = self.provider.take() {
717 let cleanup =
718 tokio::time::timeout_at(cleanup_deadline, async { provider.shutdown().await })
719 .await;
720 if !matches!(cleanup, Ok(Ok(()))) && self.tracing_enabled {
721 tracing::warn!(
722 cleanup_target = "provider",
723 "agent build cleanup did not complete successfully"
724 );
725 }
726 let error = match cleanup {
727 Ok(Ok(())) => None,
728 Ok(Err(error)) => Some(KitError::from_provider(&error)),
729 Err(_) => Some(
730 KitError::timeout("provider cleanup deadline elapsed").with_detail(
731 crate::error::ErrorDetail::DeadlineElapsed {
732 kind: crate::error::DeadlineKind::Lifecycle,
733 timeout: Some(self.cleanup_timeout),
734 },
735 ),
736 ),
737 };
738 if let Some(error) = error {
739 primary =
740 primary.with_supplementary(error.at_phase(crate::error::ErrorPhase::Cleanup));
741 }
742 }
743 primary
744 }
745}
746
747struct AgentBuildPublication<A> {
748 result: Mutex<Option<Result<A, KitError>>>,
749 handle: Mutex<Option<JoinHandle<()>>>,
750}
751
752impl<S: StateStoreBackend, D: DatabaseBackend, P: ProviderBackend, T: ToolsetBackend>
753 AgentBuildPublication<Agent<S, D, P, T>>
754{
755 fn new() -> Self {
756 Self {
757 result: Mutex::new(None),
758 handle: Mutex::new(None),
759 }
760 }
761
762 async fn run(owner: AgentBuildOwner<S, D, P, T>) -> Result<Agent<S, D, P, T>, KitError> {
763 let publication = Arc::new(Self::new());
764 let task_publication = Arc::clone(&publication);
765 let (start_sender, start_receiver) = oneshot::channel();
766 let (ready_sender, ready_receiver) = oneshot::channel();
767 let (claimed_sender, claimed_receiver) = oneshot::channel();
768 let handle = tokio::spawn(async move {
769 if start_receiver.await.is_err() {
770 let _ = owner
771 .cleanup_after_error(KitError::invalid_state(
772 "agent build owner could not start",
773 ))
774 .await;
775 task_publication.finish();
776 return;
777 }
778
779 let result = owner.build().await;
780 let successful = result.is_ok();
781 task_publication.publish(result);
782 if ready_sender.send(()).is_err() {
783 if successful {
784 task_publication.cleanup_unclaimed().await;
785 }
786 task_publication.finish();
787 return;
788 }
789 if successful && claimed_receiver.await.is_err() {
790 task_publication.cleanup_unclaimed().await;
791 }
792 task_publication.finish();
793 });
794 publication.install(handle);
795 start_sender
796 .send(())
797 .map_err(|_| KitError::invalid_state("agent build owner could not start"))?;
798 ready_receiver
799 .await
800 .map_err(|_| KitError::invalid_state("agent build owner terminated unexpectedly"))?;
801
802 let mut agent = publication.take()?;
803 agent.retain_storage_claim();
804 if claimed_sender.send(()).is_err() {
805 let primary =
806 KitError::invalid_state("agent build publication terminated unexpectedly")
807 .with_detail(crate::error::ErrorDetail::TaskTerminated)
808 .in_phase(crate::error::ErrorPhase::Publication);
809 return Err(match agent.shutdown().await {
810 Ok(()) => primary,
811 Err(cleanup) => {
812 primary.with_supplementary(cleanup.at_phase(crate::error::ErrorPhase::Cleanup))
813 }
814 });
815 }
816 Ok(agent)
817 }
818
819 fn install(&self, handle: JoinHandle<()>) {
820 *self
821 .handle
822 .lock()
823 .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(handle);
824 }
825
826 fn publish(&self, result: Result<Agent<S, D, P, T>, KitError>) {
827 *self
828 .result
829 .lock()
830 .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(result);
831 }
832
833 fn take(&self) -> Result<Agent<S, D, P, T>, KitError> {
834 self.result
835 .lock()
836 .unwrap_or_else(|poisoned| poisoned.into_inner())
837 .take()
838 .ok_or_else(|| KitError::invalid_state("agent build result was not published"))?
839 }
840
841 async fn cleanup_unclaimed(&self) {
842 let result = self
843 .result
844 .lock()
845 .unwrap_or_else(|poisoned| poisoned.into_inner())
846 .take();
847 if let Some(Ok(agent)) = result {
848 let _ = agent.shutdown().await;
849 }
850 }
851
852 fn finish(&self) {
853 self.handle
854 .lock()
855 .unwrap_or_else(|poisoned| poisoned.into_inner())
856 .take();
857 }
858}