use std::pin::Pin;
use futures::Stream;
use futures::StreamExt;
use objectiveai_sdk::agent::InlineAgentBaseWithFallbacksOrRemoteCommitOptional;
use objectiveai_sdk::agent::completions::message::{Message, UserMessage};
use objectiveai_sdk::agent::completions::request::AgentCompletionCreateParams;
use objectiveai_sdk::cli::command::agents::selector::{AgentRef, AgentSelector};
use objectiveai_sdk::cli::command::agents::spawn::{
Request, RequestDangerousAdvanced, ResponseItem,
};
use objectiveai_sdk::cli::command::{BinaryExecutor, CommandExecutor};
use crate::context::Context;
use crate::error::Error;
use crate::websockets::agent_hierarchies::ChunkAgentHierarchies;
use crate::websockets::agent_registry::AgentInstanceRegistry;
type ItemStream = Pin<Box<dyn Stream<Item = Result<ResponseItem, Error>> + Send>>;
pub async fn execute(
ctx: &Context,
request: Request,
) -> Result<ItemStream, Error> {
let want_stream = request
.dangerous_advanced
.as_ref()
.and_then(|a| a.stream)
.unwrap_or(false);
if want_stream {
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 = Some(true),
None => {
child_request.dangerous_advanced = Some(RequestDangerousAdvanced {
stream: Some(true),
..Default::default()
})
}
}
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 spawn: {e}"
)))?;
let first = stream
.next()
.await
.ok_or(Error::EmptyStream)?
.map_err(|e| Error::Instance(format!(
"self-respawn for agents spawn: {e}"
)))?;
Ok(Box::pin(
objectiveai_sdk::cli::command::StreamOnce::new(Ok(first)),
))
}
enum Mode {
Fresh {
agent: InlineAgentBaseWithFallbacksOrRemoteCommitOptional,
tag: Option<String>,
},
Historic {
hierarchy: String,
},
}
async fn execute_streaming(
ctx: &Context,
request: Request,
) -> Result<ItemStream, Error> {
let content = super::message::resolve_message(ctx, request.message).await?;
let messages = if content.is_empty() {
Vec::new()
} else {
vec![Message::User(UserMessage {
content,
})]
};
let seed = request.dangerous_advanced.as_ref().and_then(|a| a.seed);
use crate::command::agents::locks::Family;
use crate::db::laboratory_attachments::Target;
let by_spec = matches!(request.agent, AgentSelector::Ref { .. });
let (mode, lab_targets, family): (Mode, Vec<Target>, Option<Family>) = match request.agent {
AgentSelector::Ref { agent } => (
Mode::Fresh {
agent: resolve_agent_ref(ctx, agent).await?,
tag: None,
},
Vec::new(),
None,
),
AgentSelector::Tag { agent_tag } => {
match crate::db::tags::lookup(ctx.db_client().await?, &agent_tag).await? {
crate::db::tags::LookupState::Bound { agent_instance_hierarchy } => {
let lab_targets = vec![
Target::Tag(agent_tag.clone()),
Target::Aih(agent_instance_hierarchy.clone()),
];
(
Mode::Historic {
hierarchy: agent_instance_hierarchy.clone(),
},
lab_targets,
Some(Family::Hierarchy(agent_instance_hierarchy)),
)
}
crate::db::tags::LookupState::Grouped { agent_spec, tag_group_id, .. } => {
let lab_targets = vec![Target::Tag(agent_tag.clone())];
(
Mode::Fresh {
agent: agent_spec,
tag: Some(agent_tag),
},
lab_targets,
Some(Family::Group(tag_group_id)),
)
}
crate::db::tags::LookupState::Absent => {
return Err(Error::TagNotFound(agent_tag));
}
}
}
AgentSelector::Instance {
parent_agent_instance_hierarchy,
agent_instance,
} => {
let parent = parent_agent_instance_hierarchy
.as_deref()
.unwrap_or(&ctx.config.agent_instance_hierarchy);
let hierarchy = format!("{parent}/{agent_instance}");
let lab_targets = vec![Target::Aih(hierarchy.clone())];
(
Mode::Historic {
hierarchy: hierarchy.clone(),
},
lab_targets,
Some(Family::Hierarchy(hierarchy)),
)
}
};
let state_dir = ctx.filesystem.state_dir();
let mut registry = AgentInstanceRegistry::new(state_dir.clone(), ctx.agent_locks_arc());
if let Some(family) = family {
let is_group = matches!(family, Family::Group(_));
match super::locks::try_acquire_family(
ctx.agent_locks(),
ctx.db_client().await?,
&state_dir,
family,
)
.await?
{
Some(fam) => {
if let Some((hierarchy, aih_lock)) = fam.aih {
registry.preseed(hierarchy, aih_lock);
}
registry.hold_tag_claims(fam.tags);
}
None if is_group => {
let tag = match &mode {
Mode::Fresh { tag: Some(tag), .. } => tag.clone(),
_ => String::new(),
};
return Err(Error::AgentTagActive { tag });
}
None => {
let agent_instance_hierarchy = match &mode {
Mode::Historic { hierarchy } => hierarchy.clone(),
_ => String::new(),
};
return Err(Error::AgentInstanceActive {
agent_instance_hierarchy,
});
}
}
}
let (agent, agent_tag, continuation) = match mode {
Mode::Fresh { agent, tag } => (agent, tag, None),
Mode::Historic { hierarchy } => {
let lookup = async {
crate::db::logs::lookup_session(ctx.db_client().await?, &hierarchy)
.await?
.ok_or(Error::AgentNoPriorRequest {
agent_instance_hierarchy: hierarchy.clone(),
})
}
.await;
let lookup = match lookup {
Ok(lookup) => lookup,
Err(e) => {
let tee =
crate::db::logs::ConversationTee::spawn(ctx.filesystem.state_dir());
note_error(ctx, &tee, Some(&hierarchy), None, &e).await;
return Err(e);
}
};
(lookup.agent, None, lookup.continuation)
}
};
let agent_ref = if by_spec {
match &agent {
InlineAgentBaseWithFallbacksOrRemoteCommitOptional::Remote(remote) => {
crate::db::agent_refs::AgentRefValue::remote(remote)
}
InlineAgentBaseWithFallbacksOrRemoteCommitOptional::AgentBase(spec) => {
crate::db::agent_refs::AgentRefValue::inline(spec)
}
}
} else {
None
};
let ctx_clone = ctx.clone();
Ok(Box::pin(run_multi_pass(
ctx_clone,
messages,
agent,
seed,
continuation,
agent_tag,
lab_targets,
agent_ref,
registry,
)))
}
async fn resolve_laboratories(
ctx: &Context,
lab_targets: &[crate::db::laboratory_attachments::Target],
) -> Result<Option<Vec<objectiveai_sdk::laboratories::Laboratory>>, Error> {
if lab_targets.is_empty() {
return Ok(None);
}
let pool = ctx.db_client().await?;
let lists = futures::future::try_join_all(
lab_targets
.iter()
.map(|target| crate::db::laboratory_attachments::list(pool, target)),
)
.await?;
let mut ids: Vec<String> = Vec::new();
for list in lists {
for id in list {
if !ids.contains(&id) {
ids.push(id);
}
}
}
if ids.is_empty() {
return Ok(None);
}
Ok(Some(
ids.into_iter()
.map(|id| {
objectiveai_sdk::laboratories::Laboratory::Client(
objectiveai_sdk::laboratories::ClientLaboratory {
r#type: objectiveai_sdk::laboratories::ClientLaboratoryType::Client,
id,
},
)
})
.collect(),
))
}
pub(crate) async fn note_error(
ctx: &Context,
tee: &crate::db::logs::ConversationTee,
aih: Option<&str>,
response_id: Option<&str>,
error: &Error,
) {
let Some(aih) = aih else { return };
let Ok(pool) = ctx.db_client().await else {
return;
};
let value = error.output_message();
let timestamp = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs() as i64)
.unwrap_or(0);
if crate::db::logs::insert_error(pool, aih, response_id, &value, timestamp)
.await
.is_ok()
{
tee.send(crate::db::logs::error_frame(
aih.to_string(),
response_id.map(str::to_string),
value,
timestamp,
));
}
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn run_multi_pass(
ctx: Context,
messages: Vec<Message>,
agent: InlineAgentBaseWithFallbacksOrRemoteCommitOptional,
seed: Option<i64>,
continuation: Option<String>,
agent_tag: Option<String>,
lab_targets: Vec<crate::db::laboratory_attachments::Target>,
agent_ref: Option<crate::db::agent_refs::AgentRefValue>,
mut registry: AgentInstanceRegistry,
) -> impl Stream<Item = Result<ResponseItem, Error>> + Send {
async_stream::try_stream! {
let mut agent_ref = agent_ref;
let conversation_tee =
crate::db::logs::ConversationTee::spawn(ctx.filesystem.state_dir());
let laboratories = match resolve_laboratories(&ctx, &lab_targets).await {
Ok(laboratories) => laboratories,
Err(e) => {
note_error(&ctx, &conversation_tee, registry.aih(), None, &e).await;
Err(e)?;
unreachable!("Err(e)? diverges");
}
};
let mut params = AgentCompletionCreateParams {
messages,
provider: None,
agent,
response_format: None,
seed,
stream: Some(true),
continuation,
laboratories,
};
let mut identity: Option<(String, String)> = None;
let mut id_emitted = false;
let tuning = async {
Ok::<_, Error>((
ctx.resolve_mcp_timeout_ms().await?,
ctx.resolve_backoff_max_elapsed_time_ms().await?,
))
}
.await;
let (mcp_timeout_ms, backoff_max_elapsed_time_ms) = match tuning {
Ok(tuning) => tuning,
Err(e) => {
note_error(&ctx, &conversation_tee, registry.aih(), None, &e).await;
Err(e)?;
unreachable!("Err(e)? diverges");
}
};
let mut last_response_id: Option<String> = None;
loop {
let mcp_server =
crate::websockets::mcp_server::spawn(ctx.clone());
let conduit =
crate::websockets::conduit::ConduitMcpHandler::new(
mcp_server,
ctx.clone(),
agent_tag.clone(),
mcp_timeout_ms,
backoff_max_elapsed_time_ms,
);
let pool = match ctx.db_client().await {
Ok(pool) => pool,
Err(e) => {
let e = Error::from(e);
note_error(
&ctx,
&conversation_tee,
identity.as_ref().map(|(h, _)| h.as_str()).or(registry.aih()),
last_response_id.as_deref(),
&e,
)
.await;
Err(e)?;
unreachable!("Err(e)? diverges");
}
};
let (log_writer, _ready_rx) = match crate::db::logs::write_agent_completion(
pool,
¶ms,
ctx.config.agent_instance_hierarchy.clone(),
Some(conversation_tee.clone()),
)
.map_err(|e| Error::Instance(format!(
"failed to build agent-completion log writer: {e}"
))) {
Ok(writer) => writer,
Err(e) => {
note_error(
&ctx,
&conversation_tee,
identity.as_ref().map(|(h, _)| h.as_str()).or(registry.aih()),
last_response_id.as_deref(),
&e,
)
.await;
Err(e)?;
unreachable!("Err(e)? diverges");
}
};
let stream_open = async {
objectiveai_sdk::agent::completions::create_agent_completion_streaming(
ctx.api_client().await?,
params.clone(),
conduit.clone(),
)
.await
.map_err(|e| Error::Instance(format!(
"failed to open agent-completion stream: {e}"
)))
}
.await;
let (sdk_stream, notifier) = match stream_open {
Ok(opened) => opened,
Err(e) => {
note_error(
&ctx,
&conversation_tee,
identity.as_ref().map(|(h, _)| h.as_str()).or(registry.aih()),
last_response_id.as_deref(),
&e,
)
.await;
Err(e)?;
unreachable!("Err(e)? diverges");
}
};
conduit.install_notifier(notifier);
let mut sdk_stream = Box::pin(sdk_stream);
let mut last_continuation: Option<String> = None;
let mut buffered: Vec<
objectiveai_sdk::agent::completions::response::streaming::AgentCompletionChunk,
> = Vec::new();
let mut stream_err: Option<String> = None;
while let Some(item) = sdk_stream.next().await {
let chunk = match item {
Ok(c) => c,
Err(e) => {
stream_err = Some(format!("agent stream item error: {e}"));
break;
}
};
last_response_id = Some(chunk.id.clone());
if identity.is_none() {
let hier = chunk.agent_instance_hierarchy.clone();
let full_id = chunk.agent_full_id.clone();
registry.observe(&hier).await;
if let Some(value) = agent_ref.take() {
let upsert = async {
crate::db::agent_refs::upsert(
ctx.db_client().await?,
&hier,
value,
)
.await?;
Ok::<_, Error>(())
}
.await;
if let Err(e) = upsert {
stream_err = Some(format!("agent_refs upsert: {e}"));
break;
}
}
identity = Some((hier, full_id));
}
if let Some(c) = chunk.continuation.as_deref() {
last_continuation = Some(c.to_string());
}
let mut continuation_upserts: Vec<_> = Vec::new();
for (hier, continuation) in chunk.agent_instance_hierarchies() {
if let Some(c) = continuation {
continuation_upserts.push(
crate::db::agent_continuations::upsert(ctx.db_client().await?, hier, c),
);
}
}
if let Err(e) =
futures::future::try_join_all(continuation_upserts).await
{
stream_err =
Some(format!("agent_continuations upsert: {e}"));
break;
}
if let Err(e) = log_writer.write(chunk.clone()) {
stream_err = Some(format!("log writer error: {e}"));
break;
}
if !id_emitted && log_writer.written_once() {
let (hier, _) = identity
.as_ref()
.expect("identity set above on the first chunk");
yield ResponseItem::Id(hier.clone());
for c in buffered.drain(..) {
yield ResponseItem::Chunk(c);
}
id_emitted = true;
}
if id_emitted {
yield ResponseItem::Chunk(chunk);
} else {
buffered.push(chunk);
}
}
if !id_emitted && !buffered.is_empty() {
if let Err(e) = log_writer.wait_written_once().await {
stream_err.get_or_insert_with(|| format!("log writer wait: {e}"));
} else {
let (hier, _) = identity
.as_ref()
.expect("identity set on the first chunk");
yield ResponseItem::Id(hier.clone());
for c in buffered.drain(..) {
yield ResponseItem::Chunk(c);
}
id_emitted = true;
}
}
if let Err(e) = log_writer.finalize().await {
stream_err.get_or_insert_with(|| format!("log writer finalize: {e}"));
}
drop(sdk_stream);
drop(conduit);
if let Some(e) = stream_err {
let e = Error::Instance(e);
note_error(
&ctx,
&conversation_tee,
identity.as_ref().map(|(h, _)| h.as_str()).or(registry.aih()),
last_response_id.as_deref(),
&e,
)
.await;
Err(e)?;
}
let Some((hier, _full_id)) = identity.as_ref() else {
break;
};
let pending = crate::db::message_queue::check_any_pending(
ctx.db_client().await?, hier,
)
.await
.unwrap_or(false);
if !pending {
break;
}
params.messages = Vec::new();
params.continuation = last_continuation;
params.laboratories = match resolve_laboratories(&ctx, &lab_targets).await {
Ok(laboratories) => laboratories,
Err(e) => {
note_error(
&ctx,
&conversation_tee,
identity.as_ref().map(|(h, _)| h.as_str()).or(registry.aih()),
last_response_id.as_deref(),
&e,
)
.await;
Err(e)?;
unreachable!("Err(e)? diverges");
}
};
}
}
}
pub(crate) async fn resolve_agent_ref(
ctx: &Context,
agent: AgentRef,
) -> Result<InlineAgentBaseWithFallbacksOrRemoteCommitOptional, Error> {
let (file, python_inline, python_file) = match agent {
AgentRef::Resolved(resolved) => return Ok(resolved),
AgentRef::File(p) => (Some(p), None, None),
AgentRef::PythonInline(code) => (None, Some(code), None),
AgentRef::PythonFile(p) => (None, None, Some(p)),
};
crate::source_resolver::resolve_source(
ctx,
None,
None,
file,
python_inline,
python_file,
|_| unreachable!("agent refs have no plain-text variant"),
)
.await
}
pub mod request_schema {
use objectiveai_sdk::cli::command::agents::spawn as sdk;
use objectiveai_sdk::cli::command::agents::spawn::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::spawn as sdk;
use objectiveai_sdk::cli::command::agents::spawn::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)))
}
}