a3s_code_core/
agent_protocol_harness.rs1use 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
23pub const AGENT_PROTOCOL_HARNESS_MAX_SESSIONS: usize = 1_024;
25
26#[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
59pub 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 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 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 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 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 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 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}