objectiveai_cli/command/agents/laboratories/
mod.rs1use 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
69pub(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
149pub(super) fn release_all(claims: Vec<locks::AgentLock>) {
152 for claim in claims {
153 let _ = claim.release();
154 }
155}