Skip to main content

a3s_code_core/
agent_protocol_harness.rs

1//! Code-owned multi-session kernel for the native `a3s code harness` process.
2//!
3//! The executable supplies HTTP and health transport. This kernel owns only
4//! admission into existing `Agent`/`AgentSession` state and deliberately has
5//! no parallel run store, scheduler, event journal, or recovery semantics.
6
7use crate::agent_api::{Agent, SessionOptions};
8use crate::agent_protocol::{
9    AgentProtocolCommandReceiptV1, AgentProtocolCommandV1, AgentProtocolError,
10    AgentProtocolEventPageRequestV1, AgentProtocolEventPageV1, AgentProtocolRunIdentityV1,
11};
12use crate::agent_protocol_host::{AgentProtocolHost, AgentProtocolHostError};
13use crate::error::CodeError;
14use crate::release::{
15    agent_harness_compatibility_v1, AgentReleaseError, AgentReleaseManifest, AGENT_PROTOCOL_V1,
16};
17use std::collections::HashMap;
18use std::sync::atomic::{AtomicBool, Ordering};
19use std::sync::Arc;
20use thiserror::Error;
21use tokio::sync::{Mutex, RwLock};
22
23/// Finite default number of conversation sessions retained by one Harness.
24pub const AGENT_PROTOCOL_HARNESS_MAX_SESSIONS: usize = 1_024;
25
26/// Stable failures returned by the Code-owned multi-session Harness kernel.
27#[derive(Debug, Error)]
28pub enum AgentProtocolHarnessError {
29    #[error(transparent)]
30    Protocol(#[from] AgentProtocolError),
31    #[error(transparent)]
32    Release(#[from] AgentReleaseError),
33    #[error(transparent)]
34    Host(#[from] AgentProtocolHostError),
35    #[error(transparent)]
36    Code(#[from] CodeError),
37    #[error("A3S Code Harness session was not found")]
38    SessionNotFound,
39    #[error("A3S Code Harness session capacity is exhausted")]
40    SessionCapacity,
41    #[error("A3S Code Harness is draining or stopped")]
42    Closed,
43}
44
45impl AgentProtocolHarnessError {
46    pub const fn code(&self) -> &'static str {
47        match self {
48            Self::Protocol(error) => error.code(),
49            Self::Release(error) => error.code(),
50            Self::Host(error) => error.code(),
51            Self::Code(error) => error.code(),
52            Self::SessionNotFound => "a3s.code.agent_protocol.session_not_found",
53            Self::SessionCapacity => "a3s.code.agent_protocol.session_capacity",
54            Self::Closed => "a3s.code.agent_protocol.harness_closed",
55        }
56    }
57}
58
59/// Release-bound, multi-session kernel used by the sole native Harness.
60///
61/// Each entry is an [`AgentProtocolHost`] over one ordinary [`AgentSession`](crate::AgentSession).
62/// The map only retains those Code-owned sessions for conversation reuse; it
63/// never mirrors their runs or events. Miss admission is serialized so two
64/// concurrent commands cannot construct the same session twice, while work on
65/// already admitted sessions remains concurrent.
66pub struct AgentProtocolHarness {
67    manifest: Arc<AgentReleaseManifest>,
68    agent: Arc<Agent>,
69    workspace: String,
70    session_options: SessionOptions,
71    max_sessions: usize,
72    sessions: RwLock<HashMap<String, Arc<AgentProtocolHost>>>,
73    admission: Mutex<()>,
74    closed: AtomicBool,
75}
76
77impl std::fmt::Debug for AgentProtocolHarness {
78    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
79        formatter
80            .debug_struct("AgentProtocolHarness")
81            .field("agent_release_identity", &self.manifest.artifact().digest())
82            .field("manifest_identity", &self.manifest.identity())
83            .field("workspace", &self.workspace)
84            .field("max_sessions", &self.max_sessions)
85            .field("closed", &self.closed.load(Ordering::Acquire))
86            .finish_non_exhaustive()
87    }
88}
89
90impl AgentProtocolHarness {
91    /// Admit one release into the native Harness compatibility surface.
92    pub fn new(
93        manifest: AgentReleaseManifest,
94        agent: Arc<Agent>,
95        workspace: impl Into<String>,
96    ) -> Result<Self, AgentProtocolHarnessError> {
97        manifest.verify_compatibility(&agent_harness_compatibility_v1())?;
98        if manifest.protocol() != AGENT_PROTOCOL_V1 {
99            return Err(AgentProtocolHostError::ReleaseProtocolMismatch.into());
100        }
101        Ok(Self {
102            manifest: Arc::new(manifest),
103            agent,
104            workspace: workspace.into(),
105            session_options: SessionOptions::new(),
106            max_sessions: AGENT_PROTOCOL_HARNESS_MAX_SESSIONS,
107            sessions: RwLock::new(HashMap::new()),
108            admission: Mutex::new(()),
109            closed: AtomicBool::new(false),
110        })
111    }
112
113    /// Apply common options to every Code session created by this Harness.
114    ///
115    /// A caller-provided session ID is ignored. The exact protocol identity is
116    /// authoritative, and auto-save is always enabled when a store is present.
117    pub fn with_session_options(mut self, options: SessionOptions) -> Self {
118        self.session_options = options;
119        self.session_options.session_id = None;
120        self.session_options.auto_save = true;
121        self
122    }
123
124    /// Override the finite retained-session limit.
125    pub fn with_max_sessions(
126        mut self,
127        max_sessions: usize,
128    ) -> Result<Self, AgentProtocolHarnessError> {
129        if max_sessions == 0 {
130            return Err(AgentProtocolHarnessError::SessionCapacity);
131        }
132        self.max_sessions = max_sessions;
133        Ok(self)
134    }
135
136    pub fn manifest(&self) -> &AgentReleaseManifest {
137        &self.manifest
138    }
139
140    pub fn agent_release_identity(&self) -> &str {
141        self.manifest.artifact().digest()
142    }
143
144    pub fn max_sessions(&self) -> usize {
145        self.max_sessions
146    }
147
148    pub fn is_closed(&self) -> bool {
149        self.closed.load(Ordering::Acquire)
150    }
151
152    pub async fn session_count(&self) -> usize {
153        self.sessions.read().await.len()
154    }
155
156    /// Route an exact command into its Code-owned conversation session.
157    pub async fn execute(
158        &self,
159        command: &AgentProtocolCommandV1,
160    ) -> Result<AgentProtocolCommandReceiptV1, AgentProtocolHarnessError> {
161        command.validate()?;
162        let create_if_missing = matches!(
163            command,
164            AgentProtocolCommandV1::Start { .. } | AgentProtocolCommandV1::Recover { .. }
165        );
166        let host = self.host_for(command.identity(), create_if_missing).await?;
167        host.execute(command).await.map_err(Into::into)
168    }
169
170    /// Route a bounded event query into the same authoritative Code session.
171    pub async fn event_page(
172        &self,
173        request: &AgentProtocolEventPageRequestV1,
174    ) -> Result<AgentProtocolEventPageV1, AgentProtocolHarnessError> {
175        request.validate()?;
176        let host = self.host_for(&request.identity, false).await?;
177        host.event_page_for(request).await.map_err(Into::into)
178    }
179
180    /// Stop admission and close every Code-owned session and Agent resource.
181    pub async fn close(&self) {
182        if self.closed.swap(true, Ordering::AcqRel) {
183            return;
184        }
185        let _admission = self.admission.lock().await;
186        self.agent.close().await;
187        self.sessions.write().await.clear();
188    }
189
190    async fn host_for(
191        &self,
192        identity: &AgentProtocolRunIdentityV1,
193        create_if_missing: bool,
194    ) -> Result<Arc<AgentProtocolHost>, AgentProtocolHarnessError> {
195        identity.validate()?;
196        if identity.agent_release_identity != self.manifest.artifact().digest() {
197            return Err(AgentProtocolHostError::ReleaseMismatch.into());
198        }
199        if self.is_closed() {
200            return Err(AgentProtocolHarnessError::Closed);
201        }
202        if let Some(host) = self
203            .sessions
204            .read()
205            .await
206            .get(&identity.session_id)
207            .cloned()
208        {
209            return Ok(host);
210        }
211
212        let _admission = self.admission.lock().await;
213        if self.is_closed() {
214            return Err(AgentProtocolHarnessError::Closed);
215        }
216        if let Some(host) = self
217            .sessions
218            .read()
219            .await
220            .get(&identity.session_id)
221            .cloned()
222        {
223            return Ok(host);
224        }
225        if self.sessions.read().await.len() >= self.max_sessions {
226            return Err(AgentProtocolHarnessError::SessionCapacity);
227        }
228
229        let options = self
230            .session_options
231            .clone()
232            .with_session_id(&identity.session_id)
233            .with_auto_save(true);
234        let session = self
235            .agent
236            .open_protocol_session_async(&self.workspace, options, create_if_missing)
237            .await?
238            .ok_or(AgentProtocolHarnessError::SessionNotFound)?;
239        let host = Arc::new(AgentProtocolHost::from_manifest(
240            &self.manifest,
241            Arc::new(session),
242        )?);
243        self.sessions
244            .write()
245            .await
246            .insert(identity.session_id.clone(), Arc::clone(&host));
247        Ok(host)
248    }
249}