use futures::StreamExt;
use objectiveai_sdk::agent::completions::message::RichContent;
use objectiveai_sdk::cli::command::agents::message::{Request, RequestMessage, Response};
use objectiveai_sdk::cli::command::agents::selector::{AgentRef, AgentSelector};
use objectiveai_sdk::cli::command::agents::spawn as spawn_sdk;
use objectiveai_sdk::cli::command::{BinaryExecutor, CommandExecutor};
use crate::context::Context;
use crate::error::Error;
pub async fn execute(ctx: &Context, request: Request) -> Result<Response, Error> {
let Request {
agent,
message,
dangerous_advanced,
..
} = request;
let seed = dangerous_advanced.as_ref().and_then(|a| a.seed);
let content = resolve_message(ctx, message).await?;
let state_dir = ctx.filesystem.state_dir();
let route = match agent {
AgentSelector::Ref { agent } => {
let resolved = super::spawn::resolve_agent_ref(ctx, agent).await?;
Route::Ref {
child: AgentSelector::Ref {
agent: AgentRef::Resolved(resolved),
},
}
}
AgentSelector::Instance {
parent_agent_instance_hierarchy,
agent_instance,
} => {
let parent = parent_agent_instance_hierarchy
.as_deref()
.unwrap_or(&ctx.config.agent_instance_hierarchy);
instance_route(&state_dir, format!("{parent}/{agent_instance}"))
}
AgentSelector::Tag { agent_tag } => {
match crate::db::tags::lookup(ctx.db_client().await?, &agent_tag).await? {
crate::db::tags::LookupState::Bound {
agent_instance_hierarchy,
} => instance_route(&state_dir, agent_instance_hierarchy),
crate::db::tags::LookupState::Grouped { tag_group_id, .. } => {
let (dir, key) = super::locks::agent_tag_lock(&state_dir, &agent_tag);
Route::Locked {
dir,
key,
family: super::locks::Family::Group(tag_group_id),
hierarchy: None,
tag: Some(agent_tag.clone()),
child: AgentSelector::Tag { agent_tag },
}
}
crate::db::tags::LookupState::Absent => {
return Err(Error::TagNotFound(agent_tag));
}
}
}
};
match route {
Route::Ref { child } => spawn_child(child, content, seed, Vec::new()).await,
Route::Locked {
dir,
key,
family,
hierarchy,
tag,
child,
} => {
if let Some(fam) = super::locks::try_acquire_family(
ctx.agent_locks(),
ctx.db_client().await?,
&state_dir,
family.clone(),
)
.await?
{
return spawn_locked(child, content, seed, fam.into_locks()).await;
}
let queue_id = crate::db::message_queue::enqueue_with_content(
ctx.db_client().await?,
hierarchy,
tag,
&ctx.config.agent_instance_hierarchy,
None,
content.clone(),
)
.await?;
let pool = ctx.db_client().await?.clone();
let primary = (dir, key);
let mut child = Some(child);
let mut content = Some(content);
loop {
tokio::select! {
delivery = crate::db::message_queue::subscribe_delivered(&pool, queue_id) => {
delivery?;
return Ok(Response::Delivered);
}
released = objectiveai_sdk::lockfile::wait_released(&primary.0, &primary.1) => {
released.map_err(|e| Error::Lockfile {
key: primary.1.clone(),
source: e,
})?;
match super::locks::try_acquire_family(
ctx.agent_locks(),
ctx.db_client().await?,
&state_dir,
family.clone(),
)
.await?
{
Some(fam) => {
let _ = crate::db::message_queue::delete_by_id(
ctx.db_client().await?,
queue_id,
&ctx.config.agent_instance_hierarchy,
)
.await;
return spawn_locked(
child.take().expect("child consumed once"),
content.take().expect("content consumed once"),
seed,
fam.into_locks(),
)
.await;
}
None => continue,
}
}
}
}
}
}
}
enum Route {
Ref { child: AgentSelector },
Locked {
dir: std::path::PathBuf,
key: String,
family: super::locks::Family,
hierarchy: Option<String>,
tag: Option<String>,
child: AgentSelector,
},
}
fn instance_route(state_dir: &std::path::Path, hierarchy: String) -> Route {
let (dir, key) = super::locks::agent_instance_lock(state_dir, &hierarchy);
let child = AgentSelector::Instance {
parent_agent_instance_hierarchy: Some(
crate::db::tags::parent_of(&hierarchy).to_string(),
),
agent_instance: crate::db::tags::leaf_of(&hierarchy).to_string(),
};
Route::Locked {
dir,
key,
family: super::locks::Family::Hierarchy(hierarchy.clone()),
hierarchy: Some(hierarchy),
tag: None,
child,
}
}
async fn spawn_locked(
agent: AgentSelector,
content: RichContent,
seed: Option<i64>,
family: Vec<super::locks::AgentLock>,
) -> Result<Response, Error> {
spawn_child(agent, content, seed, family).await
}
async fn spawn_child(
agent: AgentSelector,
content: RichContent,
seed: Option<i64>,
transfer: Vec<super::locks::AgentLock>,
) -> Result<Response, Error> {
let child_request = spawn_sdk::Request {
path_type: spawn_sdk::Path::AgentsSpawn,
message: RequestMessage::Inline(content),
agent,
dangerous_advanced: Some(spawn_sdk::RequestDangerousAdvanced {
stream: Some(true),
seed,
}),
base: Default::default(),
};
let exe = std::env::current_exe()
.map_err(|e| Error::Spawn("current_exe".into(), e))?;
let mut executor = BinaryExecutor::from_path(exe).detach(true);
let mut transfer = transfer;
let claims: Vec<_> = transfer
.iter_mut()
.filter_map(|lock| lock.take_claim())
.collect();
if !claims.is_empty() {
executor = executor.transfer_locks(claims);
}
let mut stream = executor
.execute::<spawn_sdk::Request, spawn_sdk::ResponseItem>(child_request, None)
.await
.map_err(|e| Error::Instance(format!("exec agents spawn child: {e}")))?;
let first = stream
.next()
.await
.ok_or(Error::EmptyStream)?
.map_err(|e| Error::Instance(format!("exec agents spawn child: {e}")))?;
match first {
spawn_sdk::ResponseItem::Id(agent_instance_hierarchy) => Ok(Response::Id {
agent_instance_hierarchy,
}),
spawn_sdk::ResponseItem::Chunk(_) => Err(Error::Instance(
"agents spawn child emitted a chunk before its id".to_string(),
)),
}
}
pub async fn resolve_message(
ctx: &Context,
message: RequestMessage,
) -> Result<RichContent, Error> {
let (simple, inline, file, python_inline, python_file) = match message {
RequestMessage::Inline(rich) => return Ok(rich),
RequestMessage::Simple(s) => (Some(s), None, None, None, None),
RequestMessage::File(p) => (None, None, Some(p), None, None),
RequestMessage::PythonInline(code) => (None, None, None, Some(code), None),
RequestMessage::PythonFile(p) => (None, None, None, None, Some(p)),
};
crate::source_resolver::resolve_source(
ctx,
simple,
inline,
file,
python_inline,
python_file,
RichContent::Text,
)
.await
}
pub mod request_schema {
use objectiveai_sdk::cli::command::agents::message as sdk;
use objectiveai_sdk::cli::command::agents::message::request_schema::{Request, Response};
use crate::context::Context;
use crate::error::Error;
pub async fn execute(_ctx: &Context, _request: Request) -> Result<Response, Error> {
Ok(objectiveai_sdk::cli::command::ResponseSchema(schemars::schema_for!(sdk::Request)))
}
}
pub mod response_schema {
use objectiveai_sdk::cli::command::agents::message as sdk;
use objectiveai_sdk::cli::command::agents::message::response_schema::{Request, Response};
use crate::context::Context;
use crate::error::Error;
pub async fn execute(_ctx: &Context, _request: Request) -> Result<Response, Error> {
Ok(objectiveai_sdk::cli::command::ResponseSchema(schemars::schema_for!(sdk::Response)))
}
}