use crate::agent_api::{Agent, SessionOptions};
use crate::agent_protocol::{
AgentProtocolCommandReceiptV1, AgentProtocolCommandV1, AgentProtocolError,
AgentProtocolEventPageRequestV1, AgentProtocolEventPageV1, AgentProtocolRunIdentityV1,
};
use crate::agent_protocol_host::{AgentProtocolHost, AgentProtocolHostError};
use crate::error::CodeError;
use crate::release::{
agent_harness_compatibility_v1, AgentReleaseError, AgentReleaseManifest, AGENT_PROTOCOL_V1,
};
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use thiserror::Error;
use tokio::sync::{Mutex, RwLock};
pub const AGENT_PROTOCOL_HARNESS_MAX_SESSIONS: usize = 1_024;
#[derive(Debug, Error)]
pub enum AgentProtocolHarnessError {
#[error(transparent)]
Protocol(#[from] AgentProtocolError),
#[error(transparent)]
Release(#[from] AgentReleaseError),
#[error(transparent)]
Host(#[from] AgentProtocolHostError),
#[error(transparent)]
Code(#[from] CodeError),
#[error("A3S Code Harness session was not found")]
SessionNotFound,
#[error("A3S Code Harness session capacity is exhausted")]
SessionCapacity,
#[error("A3S Code Harness is draining or stopped")]
Closed,
}
impl AgentProtocolHarnessError {
pub const fn code(&self) -> &'static str {
match self {
Self::Protocol(error) => error.code(),
Self::Release(error) => error.code(),
Self::Host(error) => error.code(),
Self::Code(error) => error.code(),
Self::SessionNotFound => "a3s.code.agent_protocol.session_not_found",
Self::SessionCapacity => "a3s.code.agent_protocol.session_capacity",
Self::Closed => "a3s.code.agent_protocol.harness_closed",
}
}
}
pub struct AgentProtocolHarness {
manifest: Arc<AgentReleaseManifest>,
agent: Arc<Agent>,
workspace: String,
session_options: SessionOptions,
max_sessions: usize,
sessions: RwLock<HashMap<String, Arc<AgentProtocolHost>>>,
admission: Mutex<()>,
closed: AtomicBool,
}
impl std::fmt::Debug for AgentProtocolHarness {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("AgentProtocolHarness")
.field("agent_release_identity", &self.manifest.artifact().digest())
.field("manifest_identity", &self.manifest.identity())
.field("workspace", &self.workspace)
.field("max_sessions", &self.max_sessions)
.field("closed", &self.closed.load(Ordering::Acquire))
.finish_non_exhaustive()
}
}
impl AgentProtocolHarness {
pub fn new(
manifest: AgentReleaseManifest,
agent: Arc<Agent>,
workspace: impl Into<String>,
) -> Result<Self, AgentProtocolHarnessError> {
manifest.verify_compatibility(&agent_harness_compatibility_v1())?;
if manifest.protocol() != AGENT_PROTOCOL_V1 {
return Err(AgentProtocolHostError::ReleaseProtocolMismatch.into());
}
Ok(Self {
manifest: Arc::new(manifest),
agent,
workspace: workspace.into(),
session_options: SessionOptions::new(),
max_sessions: AGENT_PROTOCOL_HARNESS_MAX_SESSIONS,
sessions: RwLock::new(HashMap::new()),
admission: Mutex::new(()),
closed: AtomicBool::new(false),
})
}
pub fn with_session_options(mut self, options: SessionOptions) -> Self {
self.session_options = options;
self.session_options.session_id = None;
self.session_options.auto_save = true;
self
}
pub fn with_max_sessions(
mut self,
max_sessions: usize,
) -> Result<Self, AgentProtocolHarnessError> {
if max_sessions == 0 {
return Err(AgentProtocolHarnessError::SessionCapacity);
}
self.max_sessions = max_sessions;
Ok(self)
}
pub fn manifest(&self) -> &AgentReleaseManifest {
&self.manifest
}
pub fn agent_release_identity(&self) -> &str {
self.manifest.artifact().digest()
}
pub fn max_sessions(&self) -> usize {
self.max_sessions
}
pub fn is_closed(&self) -> bool {
self.closed.load(Ordering::Acquire)
}
pub async fn session_count(&self) -> usize {
self.sessions.read().await.len()
}
pub async fn execute(
&self,
command: &AgentProtocolCommandV1,
) -> Result<AgentProtocolCommandReceiptV1, AgentProtocolHarnessError> {
command.validate()?;
let create_if_missing = matches!(
command,
AgentProtocolCommandV1::Start { .. } | AgentProtocolCommandV1::Recover { .. }
);
let host = self.host_for(command.identity(), create_if_missing).await?;
host.execute(command).await.map_err(Into::into)
}
pub async fn event_page(
&self,
request: &AgentProtocolEventPageRequestV1,
) -> Result<AgentProtocolEventPageV1, AgentProtocolHarnessError> {
request.validate()?;
let host = self.host_for(&request.identity, false).await?;
host.event_page_for(request).await.map_err(Into::into)
}
pub async fn close(&self) {
if self.closed.swap(true, Ordering::AcqRel) {
return;
}
let _admission = self.admission.lock().await;
self.agent.close().await;
self.sessions.write().await.clear();
}
async fn host_for(
&self,
identity: &AgentProtocolRunIdentityV1,
create_if_missing: bool,
) -> Result<Arc<AgentProtocolHost>, AgentProtocolHarnessError> {
identity.validate()?;
if identity.agent_release_identity != self.manifest.artifact().digest() {
return Err(AgentProtocolHostError::ReleaseMismatch.into());
}
if self.is_closed() {
return Err(AgentProtocolHarnessError::Closed);
}
if let Some(host) = self
.sessions
.read()
.await
.get(&identity.session_id)
.cloned()
{
return Ok(host);
}
let _admission = self.admission.lock().await;
if self.is_closed() {
return Err(AgentProtocolHarnessError::Closed);
}
if let Some(host) = self
.sessions
.read()
.await
.get(&identity.session_id)
.cloned()
{
return Ok(host);
}
if self.sessions.read().await.len() >= self.max_sessions {
return Err(AgentProtocolHarnessError::SessionCapacity);
}
let options = self
.session_options
.clone()
.with_session_id(&identity.session_id)
.with_auto_save(true);
let session = self
.agent
.open_protocol_session_async(&self.workspace, options, create_if_missing)
.await?
.ok_or(AgentProtocolHarnessError::SessionNotFound)?;
let host = Arc::new(AgentProtocolHost::from_manifest(
&self.manifest,
Arc::new(session),
)?);
self.sessions
.write()
.await
.insert(identity.session_id.clone(), Arc::clone(&host));
Ok(host)
}
}