use std::sync::Arc;
use aion_core::AssistantSessionEvent;
use aion_integration_acp::{PromptOutcome, TurnHandle};
use aion_integrations::HarnessError;
use aion_integrations::contract::AgentSession as _;
use futures::StreamExt;
use super::error::AssistantSessionError;
use super::frames::TurnFrames;
use super::live::{LiveSession, Recorder};
pub(crate) const AUTH_REQUIRED_CODE: &str = "auth_required";
const ACP_AUTH_REQUIRED: i64 = -32000;
pub(crate) async fn drive(live: Arc<LiveSession>, handle: TurnHandle, turn_id: String) {
let recorder = live.recorder().clone();
if let Err(error) = pump(&recorder, handle, &turn_id).await {
record_failure(&recorder, &turn_id, "record_failed", &error).await;
}
live.release_turn();
}
async fn pump(
recorder: &Recorder,
mut handle: TurnHandle,
turn_id: &str,
) -> Result<(), AssistantSessionError> {
let mut frames = TurnFrames::new(turn_id);
let mut events = handle.events();
while let Some(event) = events.next().await {
for frame in frames.translate(event) {
recorder.record(frame).await?;
}
}
for frame in frames.flush() {
recorder.record(frame).await?;
}
let terminal = match handle.wait().await {
Ok(outcome) => completed_or_failed(turn_id, &outcome),
Err(error) => AssistantSessionEvent::TurnFailed {
turn_id: turn_id.to_owned(),
code: harness_error_code(&error).to_owned(),
message: error.to_string(),
},
};
recorder.record(terminal).await.map(drop)
}
fn completed_or_failed(turn_id: &str, outcome: &PromptOutcome) -> AssistantSessionEvent {
if let Some(error) = outcome.response.error.as_ref() {
let code = if error.code == ACP_AUTH_REQUIRED {
AUTH_REQUIRED_CODE.to_owned()
} else {
format!("agent_error_{}", error.code.abs())
};
return AssistantSessionEvent::TurnFailed {
turn_id: turn_id.to_owned(),
code,
message: error.message.clone(),
};
}
let stop_reason = outcome
.response
.result
.as_ref()
.and_then(|result| result.get("stopReason"))
.and_then(serde_json::Value::as_str)
.unwrap_or(UNSTATED_STOP_REASON)
.to_owned();
AssistantSessionEvent::TurnCompleted {
turn_id: turn_id.to_owned(),
final_message: outcome.final_message.clone(),
stop_reason,
session_ref: None,
}
}
pub(crate) const UNSTATED_STOP_REASON: &str = "unstated";
fn harness_error_code(error: &HarnessError) -> &'static str {
match error {
HarnessError::CapabilityNotSupported { .. } => "capability_not_supported",
HarnessError::StaleTarget { .. } => "stale_target",
HarnessError::Occupied { .. } => "occupied",
HarnessError::Transport { .. } => "transport",
HarnessError::Protocol { .. } => "protocol",
HarnessError::Harness { .. } => "harness",
HarnessError::PolicyRefused { .. } => "policy_refused",
HarnessError::Configuration { .. } => "configuration",
HarnessError::Contract { .. } => "contract",
_ => "harness_error",
}
}
pub(crate) async fn cancel(live: &Arc<LiveSession>) -> Result<(), AssistantSessionError> {
let session_id = live.session_id();
let delivered = live
.with_session(async |session| {
session
.intervene(aion_core::InterventionCommand {
workflow_id: aion_core::WorkflowId::new(session_id.as_uuid()),
run_id: aion_core::RunId::new(session_id.as_uuid()),
activity_id: aion_core::ActivityId::from_sequence_position(1),
attempt: 1,
issued_by: None,
issued_at: chrono::Utc::now(),
kind: aion_core::InterventionKind::Cancel {
reason: "the operator stopped this turn".to_owned(),
},
})
.await
})
.await;
match delivered {
None => Err(AssistantSessionError::Ended {
session_id,
reason: "the harness process has already been shut down".to_owned(),
}),
Some(Err(error)) => Err(AssistantSessionError::HarnessFailed {
harness: session_id.to_string(),
reason: error.to_string(),
}),
Some(Ok(())) => Ok(()),
}
}
async fn record_terminal(recorder: &Recorder, event: AssistantSessionEvent) {
if let Err(error) = recorder.record(event).await {
tracing::error!(
session = %recorder.session_id(),
%error,
"an assistant turn's terminal frame could not be recorded; the transcript is \
incomplete for this turn"
);
}
}
async fn record_failure(
recorder: &Recorder,
turn_id: &str,
code: &str,
error: &AssistantSessionError,
) {
tracing::error!(
session = %recorder.session_id(),
turn = %turn_id,
%error,
"an assistant turn could not be recorded"
);
record_terminal(
recorder,
AssistantSessionEvent::TurnFailed {
turn_id: turn_id.to_owned(),
code: code.to_owned(),
message: error.to_string(),
},
)
.await;
}