use aion_core::{InterventionCommand, InterventionOutcome};
use aion_integrations::contract::{AgentHarness, AgentSession, DynAgentHarness, DynAgentSession};
use aion_integrations::error::HarnessError;
use aion_integrations::spec::AgentRunSpec;
use futures::StreamExt;
use tokio::sync::{mpsc, oneshot};
use tracing::{debug, warn};
use crate::activity::ActivityFailure;
use crate::runtime::loop_::DispatchOutcome;
pub type ActivityEventSender = mpsc::UnboundedSender<aion_core::ActivityEvent>;
#[derive(Debug)]
pub struct ControlMessage {
pub command: InterventionCommand,
pub ack: Option<oneshot::Sender<InterventionOutcome>>,
}
impl ControlMessage {
#[must_use]
pub const fn new(command: InterventionCommand) -> Self {
Self { command, ack: None }
}
#[must_use]
pub const fn with_ack(
command: InterventionCommand,
ack: oneshot::Sender<InterventionOutcome>,
) -> Self {
Self {
command,
ack: Some(ack),
}
}
}
pub type ControlReceiver = mpsc::UnboundedReceiver<ControlMessage>;
pub async fn spawn_agent<H>(
harness: &H,
spec: AgentRunSpec,
event_sender: ActivityEventSender,
control_receiver: Option<ControlReceiver>,
) -> Result<DispatchOutcome, HarnessError>
where
H: AgentHarness,
{
let mut session = harness.start(spec).await?;
let mut events = AgentSession::events(&mut session);
let mut control = control_receiver;
loop {
tokio::select! {
biased;
maybe_message = recv_control(&mut control) => {
match maybe_message {
Some(message) => deliver_command(&session, message).await,
None => control = None,
}
}
event = events.next() => {
match event {
Some(event) => {
if event_sender.send(event).is_err() {
debug!("agent driver: event sink closed; stopping event forwarding");
break;
}
}
None => break,
}
}
}
}
let output = AgentSession::wait_result(session).await?;
Ok(DispatchOutcome::Completed { output })
}
pub async fn spawn_dyn_agent(
harness: &dyn DynAgentHarness,
spec: AgentRunSpec,
event_sender: ActivityEventSender,
control_receiver: Option<ControlReceiver>,
) -> Result<DispatchOutcome, HarnessError> {
let mut session = harness.start_dyn(spec).await?;
let mut events = session.events();
let mut control = control_receiver;
loop {
tokio::select! {
biased;
maybe_message = recv_control(&mut control) => {
match maybe_message {
Some(message) => deliver_dyn_command(session.as_ref(), message).await,
None => control = None,
}
}
event = events.next() => {
match event {
Some(event) => {
if event_sender.send(event).is_err() {
debug!("agent driver: event sink closed; stopping event forwarding");
break;
}
}
None => break,
}
}
}
}
let output = session.wait_result().await?;
Ok(DispatchOutcome::Completed { output })
}
async fn deliver_dyn_command(session: &dyn DynAgentSession, message: ControlMessage) {
let ControlMessage { command, ack } = message;
let primitive = command.kind.primitive();
let outcome = match session.intervene(command).await {
Ok(()) => InterventionOutcome::Applied,
Err(HarnessError::CapabilityNotSupported { .. }) => {
InterventionOutcome::capability_not_supported(primitive)
}
Err(HarnessError::StaleTarget { detail }) => InterventionOutcome::stale_target(detail),
Err(error) => InterventionOutcome::stale_target(error.to_string()),
};
match &outcome {
InterventionOutcome::Applied => debug!(?primitive, "agent driver: intervention delivered"),
InterventionOutcome::CapabilityNotSupported { primitive } => {
debug!(
?primitive,
"agent driver: intervention gated (capability not supported)"
);
}
InterventionOutcome::StaleTarget { detail } => {
warn!(?primitive, %detail, "agent driver: intervention delivery failed");
}
}
if let Some(ack) = ack {
drop(ack.send(outcome));
}
}
async fn recv_control(control: &mut Option<ControlReceiver>) -> Option<ControlMessage> {
match control {
Some(receiver) => receiver.recv().await,
None => std::future::pending().await,
}
}
async fn deliver_command<S>(session: &S, message: ControlMessage)
where
S: AgentSession,
{
let ControlMessage { command, ack } = message;
let primitive = command.kind.primitive();
let outcome = apply_command(session, command).await;
match &outcome {
InterventionOutcome::Applied => {
debug!(?primitive, "agent driver: intervention delivered");
}
InterventionOutcome::CapabilityNotSupported { primitive } => {
debug!(
?primitive,
"agent driver: intervention gated (capability not supported)"
);
}
InterventionOutcome::StaleTarget { detail } => {
warn!(?primitive, %detail, "agent driver: intervention delivery failed");
}
}
if let Some(ack) = ack {
drop(ack.send(outcome));
}
}
async fn apply_command<S>(session: &S, command: InterventionCommand) -> InterventionOutcome
where
S: AgentSession,
{
let primitive = command.kind.primitive();
match session.intervene(command).await {
Ok(()) => InterventionOutcome::Applied,
Err(HarnessError::CapabilityNotSupported { .. }) => {
InterventionOutcome::capability_not_supported(primitive)
}
Err(HarnessError::StaleTarget { detail }) => InterventionOutcome::stale_target(detail),
Err(error) => InterventionOutcome::stale_target(error.to_string()),
}
}
#[must_use]
pub fn harness_error_to_outcome(error: &HarnessError) -> DispatchOutcome {
DispatchOutcome::Failed {
failure: ActivityFailure::retryable(error.to_string()).into(),
}
}
#[cfg(test)]
#[path = "agent_tests.rs"]
mod tests;