a3s_code_core/
agent_protocol_host.rs1use crate::agent_api::{AgentRunSpawn, AgentSession};
8use crate::agent_protocol::{
9 validate_lower_sha256, AgentProtocolCommandReceiptV1, AgentProtocolCommandV1,
10 AgentProtocolError, AgentProtocolEventPageRequestV1, AgentProtocolEventPageV1,
11 AgentProtocolRunIdentityV1, AGENT_PROTOCOL_MAX_EVENTS_PER_PAGE,
12};
13use crate::error::CodeError;
14use crate::release::{AgentReleaseManifest, AGENT_PROTOCOL_V1};
15use crate::run::RunSnapshot;
16use std::sync::Arc;
17use thiserror::Error;
18
19#[derive(Debug, Error)]
21pub enum AgentProtocolHostError {
22 #[error(transparent)]
23 Protocol(#[from] AgentProtocolError),
24 #[error("A3S Code Agent command targets another release")]
25 ReleaseMismatch,
26 #[error("A3S Code Agent release declares another protocol")]
27 ReleaseProtocolMismatch,
28 #[error("A3S Code Agent command targets another session")]
29 SessionMismatch,
30 #[error("A3S Code Agent run was not found")]
31 RunNotFound,
32 #[error("A3S Code Agent run is not active; recover it from a durable checkpoint")]
33 RunUnavailable,
34 #[error("A3S Code Agent sequence cannot be represented on this host")]
35 SequenceOverflow,
36 #[error(transparent)]
37 Code(#[from] CodeError),
38}
39
40impl AgentProtocolHostError {
41 pub const fn code(&self) -> &'static str {
42 match self {
43 Self::Protocol(error) => error.code(),
44 Self::ReleaseMismatch => "a3s.code.agent_protocol.release_mismatch",
45 Self::ReleaseProtocolMismatch => "a3s.code.agent_protocol.release_protocol_mismatch",
46 Self::SessionMismatch => "a3s.code.agent_protocol.session_mismatch",
47 Self::RunNotFound => "a3s.code.agent_protocol.run_not_found",
48 Self::RunUnavailable => "a3s.code.agent_protocol.run_unavailable",
49 Self::SequenceOverflow => "a3s.code.agent_protocol.sequence_overflow",
50 Self::Code(error) => error.code(),
51 }
52 }
53}
54
55#[derive(Clone)]
60pub struct AgentProtocolHost {
61 agent_release_identity: String,
62 session: Arc<AgentSession>,
63}
64
65impl AgentProtocolHost {
66 pub fn new(
67 agent_release_identity: impl Into<String>,
68 session: Arc<AgentSession>,
69 ) -> Result<Self, AgentProtocolHostError> {
70 let agent_release_identity = agent_release_identity.into();
71 validate_lower_sha256("agent_release_identity", &agent_release_identity)?;
72 Ok(Self {
73 agent_release_identity,
74 session,
75 })
76 }
77
78 pub fn from_manifest(
83 manifest: &AgentReleaseManifest,
84 session: Arc<AgentSession>,
85 ) -> Result<Self, AgentProtocolHostError> {
86 if manifest.protocol() != AGENT_PROTOCOL_V1 {
87 return Err(AgentProtocolHostError::ReleaseProtocolMismatch);
88 }
89 Ok(Self {
90 agent_release_identity: manifest.artifact().digest().to_string(),
91 session,
92 })
93 }
94
95 pub fn agent_release_identity(&self) -> &str {
96 &self.agent_release_identity
97 }
98
99 pub fn session(&self) -> &Arc<AgentSession> {
100 &self.session
101 }
102
103 pub async fn execute(
107 &self,
108 command: &AgentProtocolCommandV1,
109 ) -> Result<AgentProtocolCommandReceiptV1, AgentProtocolHostError> {
110 command.validate()?;
111 self.validate_identity(command.identity())?;
112
113 let replayed = match command {
114 AgentProtocolCommandV1::Start { request } => {
115 let spawned = self
116 .session
117 .spawn_run_with_id(&request.identity.run_id, &request.prompt)
118 .await?;
119 detach(spawned)
120 }
121 AgentProtocolCommandV1::Recover { request } => {
122 let spawned = self
123 .session
124 .spawn_recovery_with_run_id(
125 &request.checkpoint_run_id,
126 &request.identity.run_id,
127 )
128 .await?;
129 detach(spawned)
130 }
131 AgentProtocolCommandV1::Cancel { request } => {
132 let snapshot = self.snapshot(&request.identity).await?;
133 if snapshot.status.is_terminal() {
134 true
135 } else if self.session.cancel_run(&request.identity.run_id).await {
136 false
137 } else if self.snapshot(&request.identity).await?.status.is_terminal() {
138 true
139 } else {
140 return Err(AgentProtocolHostError::RunUnavailable);
141 }
142 }
143 };
144
145 let snapshot = self.snapshot(command.identity()).await?;
146 let receipt = AgentProtocolCommandReceiptV1 {
147 schema: AgentProtocolCommandReceiptV1::SCHEMA.into(),
148 action: command.action(),
149 request_id: command.request_id().into(),
150 identity: command.identity().clone(),
151 command_digest: command.digest()?,
152 state: snapshot.status.into(),
153 latest_event_sequence_exclusive: u64::try_from(snapshot.event_count)
154 .map_err(|_| AgentProtocolHostError::SequenceOverflow)?,
155 observed_at_ms: now_ms().max(snapshot.updated_at_ms),
156 replayed,
157 };
158 receipt.validate_for(command)?;
159 Ok(receipt)
160 }
161
162 pub async fn event_page(
165 &self,
166 identity: &AgentProtocolRunIdentityV1,
167 after_event_sequence: Option<u64>,
168 limit: usize,
169 ) -> Result<AgentProtocolEventPageV1, AgentProtocolHostError> {
170 identity.validate()?;
171 self.validate_identity(identity)?;
172 if limit == 0 || limit > AGENT_PROTOCOL_MAX_EVENTS_PER_PAGE {
173 return Err(AgentProtocolError::InvalidField("limit").into());
174 }
175 let after_sequence = after_event_sequence
176 .map(|sequence| {
177 usize::try_from(sequence).map_err(|_| AgentProtocolHostError::SequenceOverflow)
178 })
179 .transpose()?;
180 let snapshot = self.snapshot(identity).await?;
181 let page = self
182 .session
183 .run_event_page(&identity.run_id, after_sequence, limit)
184 .await
185 .ok_or(AgentProtocolHostError::RunNotFound)?;
186 AgentProtocolEventPageV1::from_run_page(
187 identity.clone(),
188 snapshot.status,
189 now_ms().max(snapshot.updated_at_ms),
190 after_sequence,
191 &page,
192 )
193 .map_err(Into::into)
194 }
195
196 pub async fn event_page_for(
198 &self,
199 request: &AgentProtocolEventPageRequestV1,
200 ) -> Result<AgentProtocolEventPageV1, AgentProtocolHostError> {
201 request.validate()?;
202 self.event_page(
203 &request.identity,
204 request.after_event_sequence,
205 usize::from(request.limit),
206 )
207 .await
208 }
209
210 fn validate_identity(
211 &self,
212 identity: &AgentProtocolRunIdentityV1,
213 ) -> Result<(), AgentProtocolHostError> {
214 if identity.agent_release_identity != self.agent_release_identity {
215 return Err(AgentProtocolHostError::ReleaseMismatch);
216 }
217 if identity.session_id != self.session.session_id() {
218 return Err(AgentProtocolHostError::SessionMismatch);
219 }
220 Ok(())
221 }
222
223 async fn snapshot(
224 &self,
225 identity: &AgentProtocolRunIdentityV1,
226 ) -> Result<RunSnapshot, AgentProtocolHostError> {
227 let snapshot = self
228 .session
229 .run_snapshot(&identity.run_id)
230 .await
231 .ok_or(AgentProtocolHostError::RunNotFound)?;
232 if snapshot.session_id != identity.session_id {
233 return Err(AgentProtocolHostError::SessionMismatch);
234 }
235 Ok(snapshot)
236 }
237}
238
239fn detach(spawned: AgentRunSpawn) -> bool {
240 match spawned {
241 AgentRunSpawn::Started { worker, .. } => {
242 drop(worker);
245 false
246 }
247 AgentRunSpawn::Replayed { .. } => true,
248 }
249}
250
251fn now_ms() -> u64 {
252 std::time::SystemTime::now()
253 .duration_since(std::time::UNIX_EPOCH)
254 .map(|duration| duration.as_millis() as u64)
255 .unwrap_or_default()
256}