Skip to main content

objectiveai_cli/command/agents/laboratories/
mod.rs

1//! `agents laboratories` — CLI-side dispatch for attach/detach/list of
2//! laboratory ids on an agent target (a tag or an instance hierarchy).
3
4use std::pin::Pin;
5
6use futures::{Stream, StreamExt};
7use objectiveai_sdk::cli::command::agents::laboratories::{Request, ResponseItem};
8use objectiveai_sdk::cli::command::agents::selector::AgentSelector;
9
10use crate::command::agents::locks;
11use crate::context::Context;
12use crate::db::laboratory_attachments::Target;
13use crate::error::Error;
14
15pub mod attach;
16pub mod detach;
17pub mod list;
18
19type ItemStream = Pin<Box<dyn Stream<Item = Result<ResponseItem, Error>> + Send>>;
20
21fn once<T: Send + 'static>(
22    item: Result<T, Error>,
23) -> Pin<Box<dyn Stream<Item = Result<T, Error>> + Send>> {
24    Box::pin(futures::stream::once(async move { item }))
25}
26
27pub async fn execute(ctx: &Context, request: Request) -> Result<ItemStream, Error> {
28    let stream: ItemStream = match request {
29        Request::Attach(req) => {
30            let value = attach::execute(ctx, req).await?;
31            once(Ok(ResponseItem::Attach(value)))
32        }
33        Request::AttachRequestSchema(req) => {
34            let value = attach::request_schema::execute(ctx, req).await?;
35            once(Ok(ResponseItem::AttachRequestSchema(value)))
36        }
37        Request::AttachResponseSchema(req) => {
38            let value = attach::response_schema::execute(ctx, req).await?;
39            once(Ok(ResponseItem::AttachResponseSchema(value)))
40        }
41        Request::Detach(req) => {
42            let value = detach::execute(ctx, req).await?;
43            once(Ok(ResponseItem::Detach(value)))
44        }
45        Request::DetachRequestSchema(req) => {
46            let value = detach::request_schema::execute(ctx, req).await?;
47            once(Ok(ResponseItem::DetachRequestSchema(value)))
48        }
49        Request::DetachResponseSchema(req) => {
50            let value = detach::response_schema::execute(ctx, req).await?;
51            once(Ok(ResponseItem::DetachResponseSchema(value)))
52        }
53        Request::List(req) => {
54            let inner = list::execute(ctx, req).await?;
55            Box::pin(inner.map(|r| r.map(ResponseItem::List)))
56        }
57        Request::ListRequestSchema(req) => {
58            let value = list::request_schema::execute(ctx, req).await?;
59            once(Ok(ResponseItem::ListRequestSchema(value)))
60        }
61        Request::ListResponseSchema(req) => {
62            let value = list::response_schema::execute(ctx, req).await?;
63            once(Ok(ResponseItem::ListResponseSchema(value)))
64        }
65    };
66    Ok(stream)
67}
68
69/// Resolve the agent target to its DB key AND acquire the lock(s) that
70/// guard mutation for it. Shared by `attach` + `detach`.
71///
72/// - **Instance** (PAIH + `--agent-instance`) → the AIH lock; keyed on
73///   the AIH.
74/// - **Tag, GROUPED** → the tag lock; keyed on the tag.
75/// - **Tag, BOUND** → the tag lock AND the bound AIH lock (tag order
76///   first, for deadlock-freedom); keyed on the **tag**.
77/// - **Ref** (a direct agent spec) → error (no tag/AIH to key on).
78///
79/// Locks are taken with `try_acquire`, so a live owner ⇒ an "active"
80/// error. On any error the partial claims are released here. On success
81/// the caller MUST [`release_all`] after its DB op — an [`locks::AgentLock`]
82/// does not release on drop of its lockfile claim implicitly via this path.
83pub(super) async fn lock_target(
84    ctx: &Context,
85    selector: &AgentSelector,
86) -> Result<(Target, Vec<locks::AgentLock>), Error> {
87    let state_dir = ctx.filesystem.state_dir();
88    match selector {
89        AgentSelector::Ref { .. } => Err(Error::LaboratoryRefTarget),
90        AgentSelector::Instance {
91            parent_agent_instance_hierarchy,
92            agent_instance,
93        } => {
94            let parent = parent_agent_instance_hierarchy
95                .as_deref()
96                .unwrap_or(&ctx.config.agent_instance_hierarchy);
97            let aih = format!("{parent}/{agent_instance}");
98            let (dir, key) = locks::agent_instance_lock(&state_dir, &aih);
99            let claim = locks::try_acquire(ctx.agent_locks(), &dir, &key)
100                .await
101                .ok_or(Error::AgentInstanceActive {
102                    agent_instance_hierarchy: aih.clone(),
103                })?;
104            Ok((Target::Aih(aih), vec![claim]))
105        }
106        AgentSelector::Tag { agent_tag } => {
107            let (tdir, tkey) = locks::agent_tag_lock(&state_dir, agent_tag);
108            let tag_claim = locks::try_acquire(ctx.agent_locks(), &tdir, &tkey)
109                .await
110                .ok_or_else(|| Error::AgentTagActive {
111                    tag: agent_tag.clone(),
112                })?;
113            let pool = ctx.db_client().await?;
114            match crate::db::tags::lookup(pool, agent_tag).await {
115                Ok(crate::db::tags::LookupState::Absent) => {
116                    let _ = tag_claim.release();
117                    Err(Error::TagNotFound(agent_tag.clone()))
118                }
119                Ok(crate::db::tags::LookupState::Grouped { .. }) => {
120                    Ok((Target::Tag(agent_tag.clone()), vec![tag_claim]))
121                }
122                Ok(crate::db::tags::LookupState::Bound {
123                    agent_instance_hierarchy,
124                }) => {
125                    let (idir, ikey) =
126                        locks::agent_instance_lock(&state_dir, &agent_instance_hierarchy);
127                    match locks::try_acquire(ctx.agent_locks(), &idir, &ikey).await {
128                        Some(aih_claim) => Ok((
129                            Target::Tag(agent_tag.clone()),
130                            vec![tag_claim, aih_claim],
131                        )),
132                        None => {
133                            let _ = tag_claim.release();
134                            Err(Error::AgentInstanceActive {
135                                agent_instance_hierarchy,
136                            })
137                        }
138                    }
139                }
140                Err(e) => {
141                    let _ = tag_claim.release();
142                    Err(e.into())
143                }
144            }
145        }
146    }
147}
148
149/// Release every claim (best-effort). Call after the DB op on both the
150/// success and error paths — claims do not release on drop.
151pub(super) fn release_all(claims: Vec<locks::AgentLock>) {
152    for claim in claims {
153        let _ = claim.release();
154    }
155}