use crate::agent_api::{Agent, SessionOptions};
use crate::agent_protocol::{
AgentProtocolChangeSetRequestV1, AgentProtocolChangeSetV1, 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::path::{Path, PathBuf};
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,
#[error("A3S Code Harness workspace isolation failed: {0}")]
Workspace(String),
}
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",
Self::Workspace(_) => "a3s.code.agent_protocol.workspace_isolation",
}
}
}
struct HarnessSessionEntry {
host: Arc<AgentProtocolHost>,
_workspace: HarnessSessionWorkspace,
}
enum HarnessSessionWorkspace {
Shared(PathBuf),
Isolated {
source: PathBuf,
path: PathBuf,
_temporary_root: tempfile::TempDir,
},
}
impl HarnessSessionWorkspace {
async fn prepare(source: PathBuf) -> Result<Self, AgentProtocolHarnessError> {
tokio::task::spawn_blocking(move || {
if !crate::git::is_git_repo(&source) {
return Ok(Self::Shared(source));
}
let temporary_root = tempfile::Builder::new()
.prefix("a3s-code-harness-session-")
.tempdir()
.map_err(|error| AgentProtocolHarnessError::Workspace(error.to_string()))?;
let path = temporary_root.path().join("workspace");
crate::git::create_isolated_worktree(&source, &path)
.map_err(|error| AgentProtocolHarnessError::Workspace(error.to_string()))?;
Ok(Self::Isolated {
source,
path,
_temporary_root: temporary_root,
})
})
.await
.map_err(|error| AgentProtocolHarnessError::Workspace(error.to_string()))?
}
fn path(&self) -> &Path {
match self {
Self::Shared(path) | Self::Isolated { path, .. } => path,
}
}
}
impl Drop for HarnessSessionWorkspace {
fn drop(&mut self) {
if let Self::Isolated { source, path, .. } = self {
if let Err(error) = crate::git::remove_isolated_worktree(source, path) {
tracing::warn!(%error, workspace = %path.display(), "could not remove Agent Harness session worktree");
}
}
}
}
pub struct AgentProtocolHarness {
manifest: Arc<AgentReleaseManifest>,
agent: Arc<Agent>,
workspace: String,
session_options: SessionOptions,
max_sessions: usize,
sessions: RwLock<HashMap<String, Arc<HarnessSessionEntry>>>,
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 change_set(
&self,
request: &AgentProtocolChangeSetRequestV1,
) -> Result<AgentProtocolChangeSetV1, AgentProtocolHarnessError> {
request.validate()?;
let host = self.host_for(&request.identity, false).await?;
host.change_set_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)
.map(|entry| Arc::clone(&entry.host))
{
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)
.map(|entry| Arc::clone(&entry.host))
{
return Ok(host);
}
if self.sessions.read().await.len() >= self.max_sessions {
return Err(AgentProtocolHarnessError::SessionCapacity);
}
let workspace = HarnessSessionWorkspace::prepare(PathBuf::from(&self.workspace)).await?;
let options = self
.session_options
.clone()
.with_session_id(&identity.session_id)
.with_auto_save(true);
let session = self
.agent
.open_protocol_session_async(
workspace.path().to_string_lossy().into_owned(),
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::new(HarnessSessionEntry {
host: Arc::clone(&host),
_workspace: workspace,
}),
);
Ok(host)
}
}