use std::sync::Arc;
use aion_core::{
AssistantCommandInvocation, AssistantSessionEvent, AssistantSessionFrame, AssistantSessionId,
AssistantSessionProjection, AssistantSessionState, AssistantSessionSummary,
AssistantTurnContext, ContentType, Payload,
};
use aion_integration_acp::TurnHandle;
use aion_integration_acp::catalogue::{self, CatalogueHarness};
use aion_integrations::{ActivityId, AgentHarness, AgentRunSpec, RunId, WorkflowId};
use aion_store::assistant::AssistantSessionRecord;
use chrono::Utc;
use tokio::sync::broadcast;
use crate::config::ResolvedAssistantAccount;
use super::error::AssistantSessionError;
use super::launch;
use super::live::LiveSession;
use super::prompt;
use super::registry::{AssistantSessions, Availability};
use super::turn_driver;
use crate::assistant::grounding;
impl AssistantSessions {
pub async fn create(
&self,
subject: &str,
harness: Option<&str>,
account: Option<&str>,
title: Option<String>,
) -> Result<AssistantSessionSummary, AssistantSessionError> {
if let Availability::Unavailable { reason } = self.availability() {
tracing::warn!(%reason, "an assistant session was requested on a server with none");
return Err(AssistantSessionError::NotCommissioned { reason });
}
let harness = self.resolve_harness(subject, harness).await?;
let account = self.resolve_account(harness.id, account)?;
let session_id = AssistantSessionId::new_v4();
let plan = launch::plan(session_id, harness, account, self.aion_endpoint())?;
drop(plan);
let now = Utc::now();
let record = AssistantSessionRecord {
session_id,
subject: subject.to_owned(),
harness: harness.id.to_owned(),
account: account.map(|account| account.name.clone()),
title,
created_at: now,
updated_at: now,
turns: 0,
mcp_token_digest: None,
commands: Vec::new(),
};
self.store().put_assistant_session(record.clone()).await?;
self.store()
.put_assistant_default_harness(subject, harness.id)
.await?;
self.settle(
session_id,
AssistantSessionState::Dormant,
CREATED_AWAITING_FIRST_TURN,
)
.await?;
Ok(record.summary(
AssistantSessionState::Dormant,
Some(CREATED_AWAITING_FIRST_TURN.to_owned()),
))
}
pub async fn current(
&self,
subject: &str,
) -> Result<Option<AssistantSessionSummary>, AssistantSessionError> {
Ok(self
.list(subject)
.await?
.into_iter()
.find(|summary| summary.state.is_continuable()))
}
pub async fn read(
&self,
subject: &str,
session_id: AssistantSessionId,
) -> Result<(AssistantSessionSummary, Vec<AssistantSessionFrame>), AssistantSessionError> {
let record = self.owned_record(subject, session_id).await?;
let frames = self.transcript_from(session_id, None).await?;
let live = self.is_live(session_id).await;
let projection = AssistantSessionProjection::of(frames.iter().map(|frame| &frame.event));
let (state, reason) = projection.state(live.then_some(AssistantSessionState::Live));
Ok((record.summary(state, reason), frames))
}
pub async fn watch(
&self,
subject: &str,
session_id: AssistantSessionId,
after: Option<u64>,
) -> Result<
(
Vec<AssistantSessionFrame>,
broadcast::Receiver<AssistantSessionFrame>,
),
AssistantSessionError,
> {
let _record = self.owned_record(subject, session_id).await?;
let receiver = self.recorder(session_id).subscribe();
let replay = self.transcript_from(session_id, after).await?;
Ok((replay, receiver))
}
pub async fn push_context(
&self,
subject: &str,
session_id: AssistantSessionId,
context: AssistantTurnContext,
) -> Result<(), AssistantSessionError> {
let _record = self.owned_record(subject, session_id).await?;
self.recorder(session_id)
.record(AssistantSessionEvent::ContextShared {
context,
source: CONTEXT_SOURCE_PUSH.to_owned(),
})
.await
.map(drop)
}
pub async fn turn(
&self,
subject: &str,
session_id: AssistantSessionId,
text: String,
context: Option<AssistantTurnContext>,
command: Option<AssistantCommandInvocation>,
) -> Result<String, AssistantSessionError> {
let record = self.owned_record(subject, session_id).await?;
if let Some(command) = command.as_ref() {
self.require_advertised(&record, command).await?;
}
let asked = command
.as_ref()
.map_or_else(|| text.clone(), AssistantCommandInvocation::prompt_line);
let composed = prompt::compose(context.as_ref(), &asked);
let composed = if record.turns == 0 {
let dir = grounding::directory().map_err(AssistantSessionError::Internal)?;
format!("{}\n\n{composed}", grounding::preamble(&dir))
} else {
composed
};
let turn_id = uuid::Uuid::new_v4().to_string();
let (live, opening) = match self.ensure_live(&record, &composed).await {
Ok(started) => started,
Err(error) => {
self.record_turn_failure(session_id, &turn_id, &error).await;
return Err(error);
}
};
live.claim_turn()?;
let started = self
.record_turn_start(&live, &turn_id, context, &asked, command, &composed)
.await;
if let Err(error) = started {
live.release_turn();
return Err(error);
}
let handle = if let Some(handle) = opening {
handle
} else {
let submitted = live
.with_session(async |session| session.prompt(composed.clone()).await)
.await;
match submitted {
Some(Ok(handle)) => handle,
Some(Err(error)) => {
live.release_turn();
return Err(AssistantSessionError::HarnessFailed {
harness: record.harness.clone(),
reason: error.to_string(),
});
}
None => {
live.release_turn();
return Err(AssistantSessionError::Ended {
session_id,
reason: "the harness process has been shut down".to_owned(),
});
}
}
};
self.touch(session_id, Some(&asked)).await?;
let driver_turn = turn_id.clone();
tokio::spawn(turn_driver::drive(live, handle, driver_turn));
Ok(turn_id)
}
async fn record_turn_failure(
&self,
session_id: AssistantSessionId,
turn_id: &str,
error: &AssistantSessionError,
) {
if let Err(recording) = self
.recorder(session_id)
.record(AssistantSessionEvent::TurnFailed {
turn_id: turn_id.to_owned(),
code: error.code().to_owned(),
message: error.to_string(),
})
.await
{
tracing::warn!(
session = %session_id,
%recording,
original = %error,
"an assistant turn failed before it reached the agent, and the failure could not \
be recorded on its transcript"
);
}
}
pub async fn cancel(
&self,
subject: &str,
session_id: AssistantSessionId,
) -> Result<(), AssistantSessionError> {
let _record = self.owned_record(subject, session_id).await?;
let Some(live) = self.live(session_id) else {
return Err(AssistantSessionError::Ended {
session_id,
reason: "no harness process is running for this session".to_owned(),
});
};
turn_driver::cancel(&live).await
}
pub async fn resume(
&self,
subject: &str,
session_id: AssistantSessionId,
) -> Result<AssistantSessionSummary, AssistantSessionError> {
let record = self.owned_record(subject, session_id).await?;
if self.is_live(session_id).await {
return Ok(record.summary(AssistantSessionState::Live, None));
}
let projection = self.projection(session_id).await?;
if projection.acp_session_ref.is_none() || projection.is_resumable() {
let (state, reason) = projection.state(None);
return Ok(record.summary(state, reason));
}
self.settle(session_id, AssistantSessionState::Ended, RESUME_REFUSED)
.await?;
Ok(record.summary(
AssistantSessionState::Ended,
Some(RESUME_REFUSED.to_owned()),
))
}
pub async fn delete(
&self,
subject: &str,
session_id: AssistantSessionId,
) -> Result<(), AssistantSessionError> {
let _record = self.owned_record(subject, session_id).await?;
if let Some(live) = self.forget(session_id) {
live.close().await;
}
self.document_edit_locks().remove(&session_id);
self.fail_open_turn(session_id, "process_exited").await?;
self.settle(session_id, AssistantSessionState::Ended, DELETED)
.await?;
self.recorder(session_id)
.record(AssistantSessionEvent::Ended {
reason: DELETED.to_owned(),
})
.await
.map(drop)
}
pub async fn sweep_orphans(&self) -> Result<usize, AssistantSessionError> {
match self.sweep_orphans_inner().await {
Ok(settled) => {
self.clear_store_fault();
Ok(settled)
}
Err(error) => {
self.report_store_fault(&error);
Err(error)
}
}
}
async fn fail_open_turn(
&self,
session_id: AssistantSessionId,
code: &str,
) -> Result<(), AssistantSessionError> {
let projection = self.projection(session_id).await?;
let Some(turn_id) = projection.open_turn_id else {
return Ok(());
};
self.recorder(session_id)
.record(AssistantSessionEvent::TurnFailed {
turn_id,
code: code.to_owned(),
message: OPEN_TURN_ORPHANED.to_owned(),
})
.await
.map(drop)
}
async fn sweep_orphans_inner(&self) -> Result<usize, AssistantSessionError> {
let listing = self.store().list_assistant_sessions().await?;
let mut settled = 0_usize;
for record in listing.sessions {
let session_id = record.session_id;
if self.is_live(session_id).await {
continue;
}
let projection = self.projection(session_id).await?;
if projection.settled.is_some() {
continue;
}
let (state, reason) = if projection.is_resumable() {
(
AssistantSessionState::Dormant,
PROCESS_EXITED_RESUMABLE.to_owned(),
)
} else {
(
AssistantSessionState::Ended,
PROCESS_EXITED_ENDED.to_owned(),
)
};
self.fail_open_turn(session_id, "process_exited").await?;
self.settle(session_id, state, reason).await?;
settled = settled.saturating_add(1);
}
if settled > 0 {
tracing::info!(
settled,
"assistant sessions whose harness process is gone were settled at boot"
);
}
Ok(settled)
}
pub async fn shutdown(&self) {
let ids: Vec<AssistantSessionId> = self.live_ids();
for session_id in ids {
if let Some(live) = self.forget(session_id) {
live.close().await;
}
if let Err(error) = self.fail_open_turn(session_id, "server_stopped").await {
tracing::warn!(
session = %session_id,
%error,
"an assistant session's open turn could not be failed while the server \
stopped; the transcript keeps an unanswered question until the boot sweep"
);
}
if let Err(error) = self
.settle(session_id, AssistantSessionState::Dormant, SERVER_STOPPED)
.await
{
tracing::warn!(
session = %session_id,
%error,
"an assistant session could not be settled while the server stopped; the boot \
sweep will settle it on the next start"
);
}
}
}
async fn ensure_live(
&self,
record: &AssistantSessionRecord,
composed_prompt: &str,
) -> Result<(Arc<LiveSession>, Option<TurnHandle>), AssistantSessionError> {
let session_id = record.session_id;
let spawn_lock = self.spawn_lock(session_id);
let _held = spawn_lock.lock().await;
if let Some(live) = self.live(session_id)
&& live.is_alive().await
{
return Ok((live, None));
}
let projection = self.projection(session_id).await?;
let prior = if projection.acp_session_ref.is_none() {
None
} else if projection.is_resumable() {
projection.acp_session_ref.clone()
} else {
self.settle(session_id, AssistantSessionState::Ended, RESUME_REFUSED)
.await?;
return Err(AssistantSessionError::Ended {
session_id,
reason: RESUME_REFUSED.to_owned(),
});
};
let (live, handle) = self.spawn(record, composed_prompt, prior).await?;
Ok((live, Some(handle)))
}
async fn spawn(
&self,
record: &AssistantSessionRecord,
composed_prompt: &str,
prior: Option<String>,
) -> Result<(Arc<LiveSession>, TurnHandle), AssistantSessionError> {
let session_id = record.session_id;
let harness = catalogue::harness(&record.harness).ok_or_else(|| {
AssistantSessionError::UnknownHarness {
requested: record.harness.clone(),
declared: catalogue::ids(),
}
})?;
let account = match record.account.as_deref() {
Some(name) => Some(self.resolve_account_by_name(harness.id, name)?),
None => None,
};
let plan = launch::plan(session_id, harness, account, self.aion_endpoint())?;
if let Some(minted) = plan.token.as_ref() {
let mut updated = self
.store()
.get_assistant_session(&session_id)
.await?
.ok_or(AssistantSessionError::NotFound { session_id })?;
updated.mcp_token_digest = Some(minted.digest().to_owned());
updated.updated_at = Utc::now();
self.store().put_assistant_session(updated).await?;
}
let grounding_dir = grounding::directory().map_err(AssistantSessionError::Internal)?;
grounding::materialize(&grounding_dir).map_err(AssistantSessionError::Internal)?;
let mut built = plan.harness;
if prior.is_some() {
built = built.with_session_parameter(RESUME_PARAMETER);
}
let input = spawn_input(composed_prompt, prior.as_deref())?;
let spec = AgentRunSpec::new(
WorkflowId::new(session_id.as_uuid()),
RunId::new_v4(),
ActivityId::from_sequence_position(1),
1,
ACTIVITY_TYPE.to_owned(),
input,
);
let started =
built
.start(spec)
.await
.map_err(|error| AssistantSessionError::HarnessFailed {
harness: record.harness.clone(),
reason: error.to_string(),
})?;
let recorder = self.recorder(session_id);
let mut started = started;
let handle = started.take_first_turn().ok_or_else(|| {
AssistantSessionError::Internal(
"a freshly started assistant harness had no opening turn to drive".to_owned(),
)
})?;
let live = Arc::new(LiveSession::new(session_id, started, recorder));
self.adopt(Arc::clone(&live));
live.recorder()
.record(AssistantSessionEvent::SessionOpened {
acp_session_ref: live.acp_session_ref().to_owned(),
load_session: live.supports_load_session(),
at: Utc::now(),
resumed: prior.is_some(),
})
.await?;
live.recorder()
.record(AssistantSessionEvent::State {
state: AssistantSessionState::Live,
reason: None,
})
.await?;
Ok((live, handle))
}
async fn record_turn_start(
&self,
live: &Arc<LiveSession>,
turn_id: &str,
context: Option<AssistantTurnContext>,
text: &str,
command: Option<AssistantCommandInvocation>,
composed: &str,
) -> Result<(), AssistantSessionError> {
live.recorder()
.record(AssistantSessionEvent::Request {
turn_id: turn_id.to_owned(),
text: text.to_owned(),
context: context.filter(|context| !prompt::is_empty(context)),
command,
})
.await?;
live.recorder()
.record(AssistantSessionEvent::TurnStarted {
turn_id: turn_id.to_owned(),
at: Utc::now(),
prompt: composed.to_owned(),
})
.await
.map(drop)
}
async fn require_advertised(
&self,
record: &AssistantSessionRecord,
command: &AssistantCommandInvocation,
) -> Result<(), AssistantSessionError> {
if record
.commands
.iter()
.any(|advertised| advertised.name == command.name)
{
return Ok(());
}
let advertised = self.projection(record.session_id).await?.commands;
if advertised.iter().any(|entry| entry.name == command.name) {
return Ok(());
}
Err(AssistantSessionError::UnknownCommand {
session_id: record.session_id,
requested: command.name.clone(),
advertised: declared_names(advertised.iter().map(|entry| entry.name.as_str())),
})
}
async fn resolve_harness(
&self,
subject: &str,
requested: Option<&str>,
) -> Result<&'static CatalogueHarness, AssistantSessionError> {
if let Some(name) = requested {
return catalogue::harness(name).ok_or_else(|| AssistantSessionError::UnknownHarness {
requested: name.to_owned(),
declared: catalogue::ids(),
});
}
if let Some(remembered) = self.last_harness_pick(subject).await?
&& let Some(harness) = catalogue::harness(&remembered)
{
return Ok(harness);
}
catalogue::CATALOGUE
.iter()
.find(|entry| entry.available())
.or_else(|| catalogue::CATALOGUE.first())
.ok_or_else(|| AssistantSessionError::UnknownHarness {
requested: "(none named)".to_owned(),
declared: catalogue::ids(),
})
}
fn resolve_account(
&self,
harness: &str,
requested: Option<&str>,
) -> Result<Option<&ResolvedAssistantAccount>, AssistantSessionError> {
match requested {
None => Ok(None),
Some(name) => self.resolve_account_by_name(harness, name).map(Some),
}
}
fn resolve_account_by_name(
&self,
harness: &str,
name: &str,
) -> Result<&ResolvedAssistantAccount, AssistantSessionError> {
self.config()
.account(harness, name)
.ok_or_else(|| AssistantSessionError::UnknownAccount {
harness: harness.to_owned(),
requested: name.to_owned(),
declared: declared_names(
self.config()
.harness(harness)
.into_iter()
.flat_map(|declared| declared.accounts.iter())
.map(|account| account.name.as_str()),
),
})
}
}
fn spawn_input(prompt: &str, prior: Option<&str>) -> Result<Payload, AssistantSessionError> {
let mut input = serde_json::Map::new();
input.insert(
PROMPT_PARAMETER.to_owned(),
serde_json::Value::String(prompt.to_owned()),
);
if let Some(prior) = prior {
input.insert(
RESUME_PARAMETER.to_owned(),
serde_json::Value::String(prior.to_owned()),
);
}
let bytes = serde_json::to_vec(&serde_json::Value::Object(input)).map_err(|error| {
AssistantSessionError::Internal(format!("the harness job input is not encodable: {error}"))
})?;
Ok(Payload::new(ContentType::Json, bytes))
}
fn declared_names<'name>(names: impl Iterator<Item = &'name str>) -> String {
let rendered: Vec<String> = names.map(|name| format!("`{name}`")).collect();
if rendered.is_empty() {
return "none".to_owned();
}
rendered.join(", ")
}
const PROMPT_PARAMETER: &str = "prompt";
const RESUME_PARAMETER: &str = "acp_session";
const ACTIVITY_TYPE: &str = "assistant.session";
const CONTEXT_SOURCE_PUSH: &str = "push";
pub const CREATED_AWAITING_FIRST_TURN: &str = "created; the harness starts with the first turn, because ACP opens a conversation by asking \
something";
pub const PROCESS_EXITED_RESUMABLE: &str = "process_exited; the agent advertised `loadSession`, so the next turn reopens this \
conversation";
pub const PROCESS_EXITED_ENDED: &str = "process_exited; the agent did not advertise `loadSession`, so its conversation cannot be \
reopened";
pub const RESUME_REFUSED: &str = "resume_refused: loadSession not advertised — this agent cannot reload a prior conversation, \
and starting a fresh one would silently discard it";
pub const DELETED: &str = "deleted by the caller";
pub const SERVER_STOPPED: &str = "process_exited; the server that held this session stopped";
pub const OPEN_TURN_ORPHANED: &str =
"the agent's process ended before this turn finished; nothing further arrives for it";