use std::collections::HashSet;
use std::pin::Pin;
use futures::{Stream, StreamExt};
use objectiveai_sdk::agent::completions::request::AgentCompletionCreateParams;
use objectiveai_sdk::cli::command::ResponseItem as RootResponseItem;
use objectiveai_sdk::cli::command::agents::ResponseItem as AgentsResponseItem;
use objectiveai_sdk::cli::command::agents::queue::deliver::{
AgentActiveResponseItem, AgentActiveType, AgentSpawnedResponseItem, AgentSpawnedType,
AllAgentsActive, Request, RequestDangerousAdvanced, ResponseItem, TagActiveResponseItem,
TagActiveType, TagSpawnedResponseItem, TagSpawnedType, ValueResponseItem,
};
use objectiveai_sdk::cli::command::agents::spawn::ResponseItem as SpawnResponseItem;
use objectiveai_sdk::cli::command::{BinaryExecutor, CommandExecutor};
use crate::context::Context;
use crate::db;
use crate::error::Error;
use crate::websockets::agent_registry::AgentInstanceRegistry;
type ItemStream = Pin<Box<dyn Stream<Item = Result<ResponseItem, Error>> + Send>>;
type TaggedStream =
Pin<Box<dyn Stream<Item = (usize, Result<ResponseItem, Error>)> + Send>>;
pub async fn execute(ctx: &Context, request: Request) -> Result<ItemStream, Error> {
if request
.dangerous_advanced
.as_ref()
.and_then(|adv| adv.stream_spawns)
== Some(true)
{
execute_streaming(ctx, request).await
} else {
execute_detached(request).await
}
}
async fn execute_detached(request: Request) -> Result<ItemStream, Error> {
let mut child_request = request;
match child_request.dangerous_advanced.as_mut() {
Some(adv) => adv.stream_spawns = Some(true),
None => {
child_request.dangerous_advanced = Some(RequestDangerousAdvanced {
stream_spawns: Some(true),
})
}
}
crate::command::reexec::strip_inherited(&mut child_request.base);
let exe = std::env::current_exe()
.map_err(|e| Error::Spawn("current_exe".into(), e))?;
let executor = BinaryExecutor::from_path(exe).detach(true);
let mut stream = executor
.execute::<Request, ResponseItem>(child_request, None)
.await
.map_err(|e| Error::Instance(format!(
"self-respawn for agents queue deliver: {e}"
)))?;
let out = async_stream::stream! {
while let Some(item) = stream.next().await {
match item {
Ok(ResponseItem::Value(_)) => {}
Ok(item) => {
let done = matches!(item, ResponseItem::AllAgentsActive(_));
yield Ok(item);
if done {
break;
}
}
Err(e) => yield Err(Error::Instance(format!(
"self-respawn for agents queue deliver: {e}"
))),
}
}
};
Ok(Box::pin(out))
}
async fn execute_streaming(ctx: &Context, _request: Request) -> Result<ItemStream, Error> {
let caller = ctx.config.agent_instance_hierarchy.clone();
let targets = db::message_queue::list_delivery_targets(ctx.db_client().await?, &caller).await?;
let mut hierarchies: Vec<String> = Vec::new();
let mut tags: Vec<String> = Vec::new();
for target in targets {
match target {
db::message_queue::DeliveryTarget::Hierarchy { agent_instance_hierarchy } => {
if agent_instance_hierarchy != caller
&& !hierarchies.contains(&agent_instance_hierarchy)
{
hierarchies.push(agent_instance_hierarchy);
}
}
db::message_queue::DeliveryTarget::GroupedTag { agent_tag } => {
if !tags.contains(&agent_tag) {
tags.push(agent_tag);
}
}
}
}
let n = hierarchies.len() + tags.len();
let mut select_all = futures::stream::SelectAll::new();
let mut idx = 0usize;
for hierarchy in hierarchies {
let i = idx;
idx += 1;
let tagged = deliver_one_hierarchy(ctx.clone(), hierarchy).map(move |item| (i, item));
select_all.push(Box::pin(tagged) as TaggedStream);
}
for tag in tags {
let i = idx;
idx += 1;
let tagged = deliver_one_tag(ctx.clone(), tag).map(move |item| (i, item));
select_all.push(Box::pin(tagged) as TaggedStream);
}
let out = async_stream::stream! {
if n == 0 {
yield Ok(ResponseItem::AllAgentsActive(AllAgentsActive::AllAgentsActive));
return;
}
let mut seen: HashSet<usize> = HashSet::new();
let mut merged = select_all;
while let Some((idx, item)) = merged.next().await {
let first = seen.insert(idx);
yield item;
if first && seen.len() == n {
yield Ok(ResponseItem::AllAgentsActive(AllAgentsActive::AllAgentsActive));
}
}
};
Ok(Box::pin(out))
}
fn deliver_one_hierarchy(
ctx: Context,
hierarchy: String,
) -> impl Stream<Item = Result<ResponseItem, Error>> + Send {
async_stream::stream! {
let state_dir = ctx.filesystem.state_dir();
let (dir, key) =
crate::command::agents::locks::agent_instance_lock(&state_dir, &hierarchy);
let Some(claim) = objectiveai_sdk::lockfile::try_acquire(&dir, &key, "").await else {
yield Ok(ResponseItem::AgentActive(AgentActiveResponseItem {
r#type: AgentActiveType::AgentActive,
agent_instance_hierarchy: hierarchy,
}));
return;
};
let pool = match ctx.db_client().await {
Ok(pool) => pool,
Err(e) => {
let _ = claim.release();
yield Err(e);
return;
}
};
let lookup = match crate::db::logs::lookup_session(pool, &hierarchy).await {
Ok(Some(lookup)) => lookup,
Ok(None) => {
let _ = claim.release();
yield Err(Error::AgentNoPriorRequest {
agent_instance_hierarchy: hierarchy,
});
return;
}
Err(e) => {
let _ = claim.release();
yield Err(e.into());
return;
}
};
let mut registry = AgentInstanceRegistry::new(state_dir);
registry.preseed(hierarchy.clone(), claim);
yield Ok(ResponseItem::AgentSpawned(AgentSpawnedResponseItem {
r#type: AgentSpawnedType::AgentSpawned,
agent_instance_hierarchy: hierarchy.clone(),
}));
let params = AgentCompletionCreateParams {
messages: Vec::new(),
provider: None,
agent: lookup.agent,
response_format: None,
seed: None,
stream: Some(true),
continuation: lookup.continuation,
};
let inner = crate::command::agents::spawn::run_multi_pass(
ctx.clone(),
params,
None,
registry,
);
let mut inner = Box::pin(inner);
while let Some(item) = inner.next().await {
match item {
Ok(spawn_item) => {
yield Ok(ResponseItem::Value(ValueResponseItem {
agent_instance_hierarchy: hierarchy.clone(),
value: Box::new(RootResponseItem::Agents(
AgentsResponseItem::Spawn(spawn_item),
)),
}));
}
Err(e) => yield Err(e),
}
}
}
}
fn deliver_one_tag(
ctx: Context,
agent_tag: String,
) -> impl Stream<Item = Result<ResponseItem, Error>> + Send {
async_stream::stream! {
let state_dir = ctx.filesystem.state_dir();
let (dir, key) =
crate::command::agents::locks::agent_tag_lock(&state_dir, &agent_tag);
let Some(claim) = objectiveai_sdk::lockfile::try_acquire(&dir, &key, "").await else {
yield Ok(ResponseItem::TagActive(TagActiveResponseItem {
r#type: TagActiveType::TagActive,
agent_tag,
}));
return;
};
let pool = match ctx.db_client().await {
Ok(pool) => pool,
Err(e) => {
let _ = claim.release();
yield Err(e);
return;
}
};
let agent = match crate::db::tags::lookup(pool, &agent_tag).await {
Ok(crate::db::tags::LookupState::Grouped { agent_spec, .. }) => {
agent_spec
}
Ok(crate::db::tags::LookupState::Bound { agent_instance_hierarchy }) => {
let _ = claim.release();
let mut inner = Box::pin(deliver_one_hierarchy(
ctx.clone(),
agent_instance_hierarchy,
));
while let Some(item) = inner.next().await {
yield item;
}
return;
}
Ok(crate::db::tags::LookupState::Absent) => {
let _ = claim.release();
yield Err(Error::TagNotFound(agent_tag));
return;
}
Err(e) => {
let _ = claim.release();
yield Err(e.into());
return;
}
};
let mut registry = AgentInstanceRegistry::new(state_dir);
registry.hold_tag_claim(claim);
yield Ok(ResponseItem::TagSpawned(TagSpawnedResponseItem {
r#type: TagSpawnedType::TagSpawned,
agent_tag: agent_tag.clone(),
}));
let params = AgentCompletionCreateParams {
messages: Vec::new(),
provider: None,
agent,
response_format: None,
seed: None,
stream: Some(true),
continuation: None,
};
let inner = crate::command::agents::spawn::run_multi_pass(
ctx.clone(),
params,
Some(agent_tag),
registry,
);
let mut inner = Box::pin(inner);
let mut minted: Option<String> = None;
while let Some(item) = inner.next().await {
match item {
Ok(spawn_item) => {
if let SpawnResponseItem::Id(id) = &spawn_item {
minted = Some(id.clone());
}
let aih = minted.clone().unwrap_or_default();
yield Ok(ResponseItem::Value(ValueResponseItem {
agent_instance_hierarchy: aih,
value: Box::new(RootResponseItem::Agents(
AgentsResponseItem::Spawn(spawn_item),
)),
}));
}
Err(e) => yield Err(e),
}
}
}
}
pub mod request_schema {
use objectiveai_sdk::cli::command::agents::queue::deliver as sdk;
use objectiveai_sdk::cli::command::agents::queue::deliver::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::queue::deliver as sdk;
use objectiveai_sdk::cli::command::agents::queue::deliver::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::ResponseItem)))
}
}