a3s_code_core/
agent_protocol_host.rs

1//! Code-owned adapter from the versioned headless protocol to `AgentSession`.
2//!
3//! This adapter deliberately stores no parallel run state or event journal.
4//! Exact command replay is resolved by the session's authoritative run store,
5//! and event pages are projected directly from that same store.
6
7use crate::agent_api::{AgentRunSpawn, AgentSession};
8use crate::agent_protocol::{
9    validate_lower_sha256, AgentProtocolChangeSetRequestV1, AgentProtocolChangeSetV1,
10    AgentProtocolCommandReceiptV1, AgentProtocolCommandV1, AgentProtocolError,
11    AgentProtocolEventPageRequestV1, AgentProtocolEventPageV1, AgentProtocolRunIdentityV1,
12    AGENT_PROTOCOL_CHANGE_SET_ENCODING_V1, AGENT_PROTOCOL_CHANGE_SET_FORMAT_V1,
13    AGENT_PROTOCOL_MAX_CHANGE_SET_BYTES, AGENT_PROTOCOL_MAX_EVENTS_PER_PAGE,
14};
15use crate::error::CodeError;
16use crate::release::{AgentReleaseManifest, AGENT_PROTOCOL_V1};
17use crate::run::{RunSnapshot, RunWorkspaceChangeSet};
18use base64::Engine as _;
19use sha2::{Digest, Sha256};
20use std::collections::HashMap;
21use std::sync::Arc;
22use thiserror::Error;
23use tokio::sync::{Mutex, RwLock};
24
25/// Stable failures returned by the Code-owned headless protocol adapter.
26#[derive(Debug, Error)]
27pub enum AgentProtocolHostError {
28    #[error(transparent)]
29    Protocol(#[from] AgentProtocolError),
30    #[error("A3S Code Agent command targets another release")]
31    ReleaseMismatch,
32    #[error("A3S Code Agent release declares another protocol")]
33    ReleaseProtocolMismatch,
34    #[error("A3S Code Agent command targets another session")]
35    SessionMismatch,
36    #[error("A3S Code Agent run was not found")]
37    RunNotFound,
38    #[error("A3S Code Agent run is not active; recover it from a durable checkpoint")]
39    RunUnavailable,
40    #[error("A3S Code Agent sequence cannot be represented on this host")]
41    SequenceOverflow,
42    #[error("A3S Code Agent change set is still being captured")]
43    ChangeSetPending,
44    #[error("A3S Code Agent run has no Git-compatible change set")]
45    ChangeSetUnavailable,
46    #[error(transparent)]
47    Code(#[from] CodeError),
48}
49
50impl AgentProtocolHostError {
51    pub const fn code(&self) -> &'static str {
52        match self {
53            Self::Protocol(error) => error.code(),
54            Self::ReleaseMismatch => "a3s.code.agent_protocol.release_mismatch",
55            Self::ReleaseProtocolMismatch => "a3s.code.agent_protocol.release_protocol_mismatch",
56            Self::SessionMismatch => "a3s.code.agent_protocol.session_mismatch",
57            Self::RunNotFound => "a3s.code.agent_protocol.run_not_found",
58            Self::RunUnavailable => "a3s.code.agent_protocol.run_unavailable",
59            Self::SequenceOverflow => "a3s.code.agent_protocol.sequence_overflow",
60            Self::ChangeSetPending => "a3s.code.agent_protocol.change_set_pending",
61            Self::ChangeSetUnavailable => "a3s.code.agent_protocol.change_set_unavailable",
62            Self::Code(error) => error.code(),
63        }
64    }
65}
66
67/// One release- and session-bound A3S Code headless protocol host.
68///
69/// Cloud, Fleet, and other callers may transport commands and receipts, but
70/// this adapter is the sole mapping into Code's run lifecycle and event store.
71#[derive(Clone)]
72pub struct AgentProtocolHost {
73    agent_release_identity: String,
74    session: Arc<AgentSession>,
75    change_set_admission: Arc<Mutex<()>>,
76    change_set_states: Arc<RwLock<HashMap<String, ChangeSetCaptureState>>>,
77}
78
79#[derive(Debug, Clone, Copy, PartialEq, Eq)]
80enum ChangeSetCaptureState {
81    Capturing,
82    Unavailable,
83}
84
85impl AgentProtocolHost {
86    pub fn new(
87        agent_release_identity: impl Into<String>,
88        session: Arc<AgentSession>,
89    ) -> Result<Self, AgentProtocolHostError> {
90        let agent_release_identity = agent_release_identity.into();
91        validate_lower_sha256("agent_release_identity", &agent_release_identity)?;
92        Ok(Self {
93            agent_release_identity,
94            session,
95            change_set_admission: Arc::new(Mutex::new(())),
96            change_set_states: Arc::new(RwLock::new(HashMap::new())),
97        })
98    }
99
100    /// Bind an admitted v1 release manifest to its Code session.
101    ///
102    /// Capability compatibility remains an activation concern for the process
103    /// host, but a manifest for another protocol can never enter this v1 host.
104    pub fn from_manifest(
105        manifest: &AgentReleaseManifest,
106        session: Arc<AgentSession>,
107    ) -> Result<Self, AgentProtocolHostError> {
108        if manifest.protocol() != AGENT_PROTOCOL_V1 {
109            return Err(AgentProtocolHostError::ReleaseProtocolMismatch);
110        }
111        Ok(Self {
112            agent_release_identity: manifest.artifact().digest().to_string(),
113            session,
114            change_set_admission: Arc::new(Mutex::new(())),
115            change_set_states: Arc::new(RwLock::new(HashMap::new())),
116        })
117    }
118
119    pub fn agent_release_identity(&self) -> &str {
120        &self.agent_release_identity
121    }
122
123    pub fn session(&self) -> &Arc<AgentSession> {
124        &self.session
125    }
126
127    /// Execute, cancel, or recover one exact run and return a digest-bound
128    /// receipt. Start and recovery return after Code has admitted the detached
129    /// worker; progress is observed through [`Self::event_page`].
130    pub async fn execute(
131        &self,
132        command: &AgentProtocolCommandV1,
133    ) -> Result<AgentProtocolCommandReceiptV1, AgentProtocolHostError> {
134        command.validate()?;
135        self.validate_identity(command.identity())?;
136
137        let _change_set_admission = self.change_set_admission.lock().await;
138        let replayed = match command {
139            AgentProtocolCommandV1::Start { request } => {
140                let baseline = self.prepare_change_set(&request.identity).await;
141                let spawned = match self
142                    .session
143                    .spawn_run_with_id(&request.identity.run_id, &request.prompt)
144                    .await
145                {
146                    Ok(spawned) => spawned,
147                    Err(error) => {
148                        self.mark_change_set_unavailable(&request.identity).await;
149                        return Err(error.into());
150                    }
151                };
152                self.detach_with_change_set_capture(&request.identity, spawned, baseline)
153                    .await
154            }
155            AgentProtocolCommandV1::Recover { request } => {
156                let baseline = self.prepare_change_set(&request.identity).await;
157                let spawned = match self
158                    .session
159                    .spawn_recovery_with_run_id(
160                        &request.checkpoint_run_id,
161                        &request.identity.run_id,
162                    )
163                    .await
164                {
165                    Ok(spawned) => spawned,
166                    Err(error) => {
167                        self.mark_change_set_unavailable(&request.identity).await;
168                        return Err(error.into());
169                    }
170                };
171                self.detach_with_change_set_capture(&request.identity, spawned, baseline)
172                    .await
173            }
174            AgentProtocolCommandV1::Cancel { request } => {
175                let snapshot = self.snapshot(&request.identity).await?;
176                if snapshot.status.is_terminal() {
177                    true
178                } else if self.session.cancel_run(&request.identity.run_id).await {
179                    false
180                } else if self.snapshot(&request.identity).await?.status.is_terminal() {
181                    true
182                } else {
183                    return Err(AgentProtocolHostError::RunUnavailable);
184                }
185            }
186        };
187
188        let snapshot = self.snapshot(command.identity()).await?;
189        let receipt = AgentProtocolCommandReceiptV1 {
190            schema: AgentProtocolCommandReceiptV1::SCHEMA.into(),
191            action: command.action(),
192            request_id: command.request_id().into(),
193            identity: command.identity().clone(),
194            command_digest: command.digest()?,
195            state: snapshot.status.into(),
196            latest_event_sequence_exclusive: u64::try_from(snapshot.event_count)
197                .map_err(|_| AgentProtocolHostError::SequenceOverflow)?,
198            observed_at_ms: now_ms().max(snapshot.updated_at_ms),
199            replayed,
200        };
201        receipt.validate_for(command)?;
202        Ok(receipt)
203    }
204
205    /// Project a bounded cursor page directly from Code's authoritative run
206    /// store without introducing a second provider event model.
207    pub async fn event_page(
208        &self,
209        identity: &AgentProtocolRunIdentityV1,
210        after_event_sequence: Option<u64>,
211        limit: usize,
212    ) -> Result<AgentProtocolEventPageV1, AgentProtocolHostError> {
213        identity.validate()?;
214        self.validate_identity(identity)?;
215        if limit == 0 || limit > AGENT_PROTOCOL_MAX_EVENTS_PER_PAGE {
216            return Err(AgentProtocolError::InvalidField("limit").into());
217        }
218        let after_sequence = after_event_sequence
219            .map(|sequence| {
220                usize::try_from(sequence).map_err(|_| AgentProtocolHostError::SequenceOverflow)
221            })
222            .transpose()?;
223        let snapshot = self.snapshot(identity).await?;
224        let page = self
225            .session
226            .run_event_page(&identity.run_id, after_sequence, limit)
227            .await
228            .ok_or(AgentProtocolHostError::RunNotFound)?;
229        AgentProtocolEventPageV1::from_run_page(
230            identity.clone(),
231            snapshot.status,
232            now_ms().max(snapshot.updated_at_ms),
233            after_sequence,
234            &page,
235        )
236        .map_err(Into::into)
237    }
238
239    /// Execute the canonical transport-facing event page query.
240    pub async fn event_page_for(
241        &self,
242        request: &AgentProtocolEventPageRequestV1,
243    ) -> Result<AgentProtocolEventPageV1, AgentProtocolHostError> {
244        request.validate()?;
245        self.event_page(
246            &request.identity,
247            request.after_event_sequence,
248            usize::from(request.limit),
249        )
250        .await
251    }
252
253    /// Read the immutable Git-compatible patch captured for one terminal run.
254    pub async fn change_set_for(
255        &self,
256        request: &AgentProtocolChangeSetRequestV1,
257    ) -> Result<AgentProtocolChangeSetV1, AgentProtocolHostError> {
258        request.validate()?;
259        self.validate_identity(&request.identity)?;
260        let snapshot = self.snapshot(&request.identity).await?;
261        if !snapshot.status.is_terminal() {
262            return Err(AgentProtocolHostError::ChangeSetPending);
263        }
264        if let Some(change_set) = snapshot.workspace_change_set {
265            let response = AgentProtocolChangeSetV1 {
266                schema: AgentProtocolChangeSetV1::SCHEMA.into(),
267                identity: request.identity.clone(),
268                state: snapshot.status.into(),
269                format: AGENT_PROTOCOL_CHANGE_SET_FORMAT_V1.into(),
270                encoding: AGENT_PROTOCOL_CHANGE_SET_ENCODING_V1.into(),
271                base_tree: change_set.base_tree,
272                result_tree: change_set.result_tree,
273                patch_digest: change_set.patch_digest,
274                patch_bytes: change_set.patch_bytes,
275                patch_base64: change_set.patch_base64,
276                observed_at_ms: change_set.observed_at_ms,
277            };
278            response.validate()?;
279            return Ok(response);
280        }
281        match self
282            .change_set_states
283            .read()
284            .await
285            .get(&request.identity.run_id)
286            .copied()
287        {
288            Some(ChangeSetCaptureState::Capturing) => Err(AgentProtocolHostError::ChangeSetPending),
289            Some(ChangeSetCaptureState::Unavailable) | None => {
290                Err(AgentProtocolHostError::ChangeSetUnavailable)
291            }
292        }
293    }
294
295    async fn prepare_change_set(
296        &self,
297        identity: &AgentProtocolRunIdentityV1,
298    ) -> Option<crate::git::WorkspaceTreeSnapshot> {
299        if self.session.run_snapshot(&identity.run_id).await.is_some() {
300            return None;
301        }
302        let workspace = self.session.workspace().to_path_buf();
303        let baseline =
304            tokio::task::spawn_blocking(move || crate::git::snapshot_workspace_tree(&workspace))
305                .await
306                .ok()
307                .and_then(Result::ok);
308        self.change_set_states.write().await.insert(
309            identity.run_id.clone(),
310            if baseline.is_some() {
311                ChangeSetCaptureState::Capturing
312            } else {
313                ChangeSetCaptureState::Unavailable
314            },
315        );
316        baseline
317    }
318
319    async fn mark_change_set_unavailable(&self, identity: &AgentProtocolRunIdentityV1) {
320        self.change_set_states
321            .write()
322            .await
323            .insert(identity.run_id.clone(), ChangeSetCaptureState::Unavailable);
324    }
325
326    async fn detach_with_change_set_capture(
327        &self,
328        identity: &AgentProtocolRunIdentityV1,
329        spawned: AgentRunSpawn,
330        baseline: Option<crate::git::WorkspaceTreeSnapshot>,
331    ) -> bool {
332        match spawned {
333            AgentRunSpawn::Started { worker, .. } => {
334                let Some(baseline) = baseline else {
335                    drop(worker);
336                    return false;
337                };
338                let workspace = self.session.workspace().to_path_buf();
339                let session = Arc::clone(&self.session);
340                let run_id = identity.run_id.clone();
341                let pin_identity = format!(
342                    "{}:{}:{}",
343                    self.agent_release_identity, identity.session_id, identity.run_id
344                );
345                let states = Arc::clone(&self.change_set_states);
346                tokio::spawn(async move {
347                    let _ = worker.await;
348                    let evidence = tokio::task::spawn_blocking(move || {
349                        let result = crate::git::snapshot_workspace_tree(&workspace)?;
350                        let patch = crate::git::diff_workspace_trees(
351                            &workspace,
352                            &baseline,
353                            &result,
354                            AGENT_PROTOCOL_MAX_CHANGE_SET_BYTES,
355                        )?;
356                        crate::git::pin_workspace_tree(&workspace, &pin_identity, &result)?;
357                        let patch_bytes = u64::try_from(patch.len())
358                            .map_err(|_| anyhow::anyhow!("change-set byte count overflowed"))?;
359                        Ok::<_, anyhow::Error>(RunWorkspaceChangeSet {
360                            base_tree: format!("git-tree:{}", baseline.tree),
361                            result_tree: format!("git-tree:{}", result.tree),
362                            patch_digest: format!("sha256:{:x}", Sha256::digest(&patch)),
363                            patch_bytes,
364                            patch_base64: base64::engine::general_purpose::STANDARD.encode(patch),
365                            observed_at_ms: now_ms(),
366                        })
367                    })
368                    .await
369                    .ok()
370                    .and_then(Result::ok);
371                    let available = match evidence {
372                        Some(evidence) => session
373                            .record_workspace_change_set(&run_id, evidence)
374                            .await
375                            .is_ok(),
376                        None => false,
377                    };
378                    let mut states = states.write().await;
379                    if available {
380                        states.remove(&run_id);
381                    } else {
382                        states.insert(run_id, ChangeSetCaptureState::Unavailable);
383                    }
384                });
385                false
386            }
387            AgentRunSpawn::Replayed { .. } => true,
388        }
389    }
390
391    fn validate_identity(
392        &self,
393        identity: &AgentProtocolRunIdentityV1,
394    ) -> Result<(), AgentProtocolHostError> {
395        if identity.agent_release_identity != self.agent_release_identity {
396            return Err(AgentProtocolHostError::ReleaseMismatch);
397        }
398        if identity.session_id != self.session.session_id() {
399            return Err(AgentProtocolHostError::SessionMismatch);
400        }
401        Ok(())
402    }
403
404    async fn snapshot(
405        &self,
406        identity: &AgentProtocolRunIdentityV1,
407    ) -> Result<RunSnapshot, AgentProtocolHostError> {
408        let snapshot = self
409            .session
410            .run_snapshot(&identity.run_id)
411            .await
412            .ok_or(AgentProtocolHostError::RunNotFound)?;
413        if snapshot.session_id != identity.session_id {
414            return Err(AgentProtocolHostError::SessionMismatch);
415        }
416        Ok(snapshot)
417    }
418}
419
420fn now_ms() -> u64 {
421    std::time::SystemTime::now()
422        .duration_since(std::time::UNIX_EPOCH)
423        .map(|duration| duration.as_millis() as u64)
424        .unwrap_or_default()
425}
a3s_code_core/agent_protocol_host.rs

a3s_code_core/
agent_protocol_host.rs
