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, ExactRecoveryError, ExactRecoveryPreparation};
8use crate::agent_protocol::{
9    validate_lower_sha256, AgentProtocolChangeSetRequestV1, AgentProtocolChangeSetV1,
10    AgentProtocolCommandReceiptV1, AgentProtocolCommandV1, AgentProtocolError,
11    AgentProtocolEventPageRequestV1, AgentProtocolEventPageV1, AgentProtocolRunIdentityV1,
12    AgentProtocolRunRecoverExactV1, AGENT_PROTOCOL_CHANGE_SET_ENCODING_V1,
13    AGENT_PROTOCOL_CHANGE_SET_FORMAT_V1, AGENT_PROTOCOL_MAX_CHANGE_SET_BYTES,
14    AGENT_PROTOCOL_MAX_EVENTS_PER_PAGE,
15};
16use crate::error::CodeError;
17use crate::release::{AgentReleaseManifest, AGENT_PROTOCOL_V1};
18use crate::run::{RunSnapshot, RunWorkspaceChangeSet};
19use crate::session_checkpoint::SessionCheckpointError;
20use base64::Engine as _;
21use sha2::{Digest, Sha256};
22use std::collections::HashMap;
23use std::sync::Arc;
24use thiserror::Error;
25use tokio::sync::{Mutex, RwLock};
26
27/// Stable failures returned by the Code-owned headless protocol adapter.
28#[derive(Debug, Error)]
29pub enum AgentProtocolHostError {
30    #[error(transparent)]
31    Protocol(#[from] AgentProtocolError),
32    #[error("A3S Code Agent command targets another release")]
33    ReleaseMismatch,
34    #[error("A3S Code Agent release declares another protocol")]
35    ReleaseProtocolMismatch,
36    #[error("A3S Code Agent command targets another session")]
37    SessionMismatch,
38    #[error("A3S Code Agent run was not found")]
39    RunNotFound,
40    #[error("A3S Code Agent run is not active; recover it from a durable checkpoint")]
41    RunUnavailable,
42    #[error("A3S Code Agent sequence cannot be represented on this host")]
43    SequenceOverflow,
44    #[error("A3S Code Agent change set is still being captured")]
45    ChangeSetPending,
46    #[error("A3S Code Agent run has no Git-compatible change set")]
47    ChangeSetUnavailable,
48    #[error(transparent)]
49    Code(#[from] CodeError),
50}
51
52impl AgentProtocolHostError {
53    pub const fn code(&self) -> &'static str {
54        match self {
55            Self::Protocol(error) => error.code(),
56            Self::ReleaseMismatch => "a3s.code.agent_protocol.release_mismatch",
57            Self::ReleaseProtocolMismatch => "a3s.code.agent_protocol.release_protocol_mismatch",
58            Self::SessionMismatch => "a3s.code.agent_protocol.session_mismatch",
59            Self::RunNotFound => "a3s.code.agent_protocol.run_not_found",
60            Self::RunUnavailable => "a3s.code.agent_protocol.run_unavailable",
61            Self::SequenceOverflow => "a3s.code.agent_protocol.sequence_overflow",
62            Self::ChangeSetPending => "a3s.code.agent_protocol.change_set_pending",
63            Self::ChangeSetUnavailable => "a3s.code.agent_protocol.change_set_unavailable",
64            Self::Code(error) => error.code(),
65        }
66    }
67}
68
69/// Failures specific to the additive evidence-bound recovery entry point.
70///
71/// Keeping checkpoint drift in this adjacent error preserves the variant set
72/// of [`AgentProtocolHostError`] for existing v1 command callers.
73#[derive(Debug, Error)]
74pub enum AgentProtocolExactRecoveryError {
75    #[error(transparent)]
76    Host(#[from] AgentProtocolHostError),
77    #[error(transparent)]
78    Checkpoint(#[from] SessionCheckpointError),
79}
80
81impl AgentProtocolExactRecoveryError {
82    pub const fn code(&self) -> &'static str {
83        match self {
84            Self::Host(error) => error.code(),
85            Self::Checkpoint(error) => error.code(),
86        }
87    }
88}
89
90impl From<AgentProtocolError> for AgentProtocolExactRecoveryError {
91    fn from(error: AgentProtocolError) -> Self {
92        Self::Host(AgentProtocolHostError::Protocol(error))
93    }
94}
95
96impl From<ExactRecoveryError> for AgentProtocolExactRecoveryError {
97    fn from(error: ExactRecoveryError) -> Self {
98        match error {
99            ExactRecoveryError::Checkpoint(error) => Self::Checkpoint(error),
100            ExactRecoveryError::Code(error) => Self::Host(AgentProtocolHostError::Code(error)),
101        }
102    }
103}
104
105/// One release- and session-bound A3S Code headless protocol host.
106///
107/// Cloud, Fleet, and other callers may transport commands and receipts, but
108/// this adapter is the sole mapping into Code's run lifecycle and event store.
109#[derive(Clone)]
110pub struct AgentProtocolHost {
111    agent_release_identity: String,
112    session: Arc<AgentSession>,
113    change_set_admission: Arc<Mutex<()>>,
114    change_set_states: Arc<RwLock<HashMap<String, ChangeSetCaptureState>>>,
115}
116
117#[derive(Debug, Clone, Copy, PartialEq, Eq)]
118enum ChangeSetCaptureState {
119    Capturing,
120}
121
122impl AgentProtocolHost {
123    pub fn new(
124        agent_release_identity: impl Into<String>,
125        session: Arc<AgentSession>,
126    ) -> Result<Self, AgentProtocolHostError> {
127        let agent_release_identity = agent_release_identity.into();
128        validate_lower_sha256("agent_release_identity", &agent_release_identity)?;
129        Ok(Self {
130            agent_release_identity,
131            session,
132            change_set_admission: Arc::new(Mutex::new(())),
133            change_set_states: Arc::new(RwLock::new(HashMap::new())),
134        })
135    }
136
137    /// Bind an admitted v1 release manifest to its Code session.
138    ///
139    /// Capability compatibility remains an activation concern for the process
140    /// host, but a manifest for another protocol can never enter this v1 host.
141    pub fn from_manifest(
142        manifest: &AgentReleaseManifest,
143        session: Arc<AgentSession>,
144    ) -> Result<Self, AgentProtocolHostError> {
145        if manifest.protocol() != AGENT_PROTOCOL_V1 {
146            return Err(AgentProtocolHostError::ReleaseProtocolMismatch);
147        }
148        Ok(Self {
149            agent_release_identity: manifest.artifact().digest().to_string(),
150            session,
151            change_set_admission: Arc::new(Mutex::new(())),
152            change_set_states: Arc::new(RwLock::new(HashMap::new())),
153        })
154    }
155
156    pub fn agent_release_identity(&self) -> &str {
157        &self.agent_release_identity
158    }
159
160    pub fn session(&self) -> &Arc<AgentSession> {
161        &self.session
162    }
163
164    /// Execute, cancel, or recover one exact run and return a digest-bound
165    /// receipt. Start and recovery return after Code has admitted the detached
166    /// worker; progress is observed through [`Self::event_page`].
167    pub async fn execute(
168        &self,
169        command: &AgentProtocolCommandV1,
170    ) -> Result<AgentProtocolCommandReceiptV1, AgentProtocolHostError> {
171        command.validate()?;
172        self.validate_identity(command.identity())?;
173
174        let _change_set_admission = self.change_set_admission.lock().await;
175        let replayed = match command {
176            AgentProtocolCommandV1::Start { request } => {
177                let baseline = self.prepare_change_set(&request.identity).await;
178                let spawned = match self
179                    .session
180                    .spawn_run_with_id(&request.identity.run_id, &request.prompt)
181                    .await
182                {
183                    Ok(spawned) => spawned,
184                    Err(error) => {
185                        self.mark_change_set_unavailable(&request.identity).await;
186                        return Err(error.into());
187                    }
188                };
189                self.detach_with_change_set_capture(&request.identity, spawned, baseline)
190                    .await
191            }
192            AgentProtocolCommandV1::Recover { request } => {
193                let baseline = self.prepare_change_set(&request.identity).await;
194                let spawned = match self
195                    .session
196                    .spawn_recovery_with_run_id(
197                        &request.checkpoint_run_id,
198                        &request.identity.run_id,
199                    )
200                    .await
201                {
202                    Ok(spawned) => spawned,
203                    Err(error) => {
204                        self.mark_change_set_unavailable(&request.identity).await;
205                        return Err(error.into());
206                    }
207                };
208                self.detach_with_change_set_capture(&request.identity, spawned, baseline)
209                    .await
210            }
211            AgentProtocolCommandV1::Cancel { request } => {
212                let snapshot = self.snapshot(&request.identity).await?;
213                if snapshot.status.is_terminal() {
214                    true
215                } else if self.session.cancel_run(&request.identity.run_id).await {
216                    false
217                } else if self.snapshot(&request.identity).await?.status.is_terminal() {
218                    true
219                } else {
220                    return Err(AgentProtocolHostError::RunUnavailable);
221                }
222            }
223        };
224
225        let snapshot = self.snapshot(command.identity()).await?;
226        let receipt = AgentProtocolCommandReceiptV1 {
227            schema: AgentProtocolCommandReceiptV1::SCHEMA.into(),
228            action: command.action(),
229            request_id: command.request_id().into(),
230            identity: command.identity().clone(),
231            command_digest: command.digest()?,
232            state: snapshot.status.into(),
233            latest_event_sequence_exclusive: u64::try_from(snapshot.event_count)
234                .map_err(|_| AgentProtocolHostError::SequenceOverflow)?,
235            observed_at_ms: now_ms().max(snapshot.updated_at_ms),
236            replayed,
237        };
238        receipt.validate_for(command)?;
239        Ok(receipt)
240    }
241
242    /// Recover only when the locally loadable loop boundary matches the
243    /// logical component of the complete portable-checkpoint descriptor.
244    ///
245    /// This additive entry point leaves [`AgentProtocolCommandV1`] and its v1
246    /// transport shape unchanged. Validation pins the immutable checkpoint in
247    /// memory before workspace baseline capture and before the target Run is
248    /// admitted, so a concurrently overwritten store entry cannot change what
249    /// the worker resumes.
250    pub async fn execute_exact_recovery(
251        &self,
252        request: &AgentProtocolRunRecoverExactV1,
253    ) -> Result<AgentProtocolCommandReceiptV1, AgentProtocolExactRecoveryError> {
254        request.validate()?;
255        self.validate_identity(&request.identity)?;
256        let logical_resume = request.logical_resume()?;
257
258        let _change_set_admission = self.change_set_admission.lock().await;
259        let prepared = self
260            .session
261            .prepare_recovery_with_evidence(
262                logical_resume,
263                &request.checkpoint.descriptor_digest,
264                &request.identity.run_id,
265            )
266            .await?;
267        self.finish_exact_recovery(request, prepared).await
268    }
269
270    /// Recover from the logical value decoded from the same validated portable
271    /// checkpoint as `request`, without first publishing split store writes.
272    pub async fn execute_exact_recovery_from_checkpoint(
273        &self,
274        request: &AgentProtocolRunRecoverExactV1,
275        checkpoint: crate::loop_checkpoint::LoopCheckpoint,
276    ) -> Result<AgentProtocolCommandReceiptV1, AgentProtocolExactRecoveryError> {
277        request.validate()?;
278        self.validate_identity(&request.identity)?;
279        let logical_resume = request.logical_resume()?;
280
281        let _change_set_admission = self.change_set_admission.lock().await;
282        let prepared = self
283            .session
284            .prepare_recovery_from_checkpoint(
285                logical_resume,
286                &request.checkpoint.descriptor_digest,
287                &request.identity.run_id,
288                checkpoint,
289            )
290            .await?;
291        self.finish_exact_recovery(request, prepared).await
292    }
293
294    async fn finish_exact_recovery(
295        &self,
296        request: &AgentProtocolRunRecoverExactV1,
297        prepared: ExactRecoveryPreparation,
298    ) -> Result<AgentProtocolCommandReceiptV1, AgentProtocolExactRecoveryError> {
299        let replayed = match prepared {
300            ExactRecoveryPreparation::Replayed(spawned) => spawned.replayed(),
301            ExactRecoveryPreparation::Ready(prepared) => {
302                let baseline = self.prepare_change_set(&request.identity).await;
303                let spawned = match self.session.spawn_prepared_recovery(prepared).await {
304                    Ok(spawned) => spawned,
305                    Err(error) => {
306                        self.mark_change_set_unavailable(&request.identity).await;
307                        return Err(error.into());
308                    }
309                };
310                self.detach_with_change_set_capture(&request.identity, spawned, baseline)
311                    .await
312            }
313        };
314
315        let snapshot = self.snapshot(&request.identity).await?;
316        let receipt = AgentProtocolCommandReceiptV1 {
317            schema: AgentProtocolCommandReceiptV1::SCHEMA.into(),
318            action: crate::agent_protocol::AgentProtocolCommandActionV1::Recover,
319            request_id: request.request_id.clone(),
320            identity: request.identity.clone(),
321            command_digest: request.digest()?,
322            state: snapshot.status.into(),
323            latest_event_sequence_exclusive: u64::try_from(snapshot.event_count)
324                .map_err(|_| AgentProtocolHostError::SequenceOverflow)?,
325            observed_at_ms: now_ms().max(snapshot.updated_at_ms),
326            replayed,
327        };
328        receipt.validate_for_exact_recovery(request)?;
329        Ok(receipt)
330    }
331
332    /// Project a bounded cursor page directly from Code's authoritative run
333    /// store without introducing a second provider event model.
334    pub async fn event_page(
335        &self,
336        identity: &AgentProtocolRunIdentityV1,
337        after_event_sequence: Option<u64>,
338        limit: usize,
339    ) -> Result<AgentProtocolEventPageV1, AgentProtocolHostError> {
340        identity.validate()?;
341        self.validate_identity(identity)?;
342        if limit == 0 || limit > AGENT_PROTOCOL_MAX_EVENTS_PER_PAGE {
343            return Err(AgentProtocolError::InvalidField("limit").into());
344        }
345        let after_sequence = after_event_sequence
346            .map(|sequence| {
347                usize::try_from(sequence).map_err(|_| AgentProtocolHostError::SequenceOverflow)
348            })
349            .transpose()?;
350        let observation = self
351            .session
352            .run_event_observation(&identity.run_id, after_sequence, limit)
353            .await
354            .ok_or(AgentProtocolHostError::RunNotFound)?;
355        if observation.snapshot.session_id != identity.session_id {
356            return Err(AgentProtocolHostError::SessionMismatch);
357        }
358        AgentProtocolEventPageV1::from_run_page(
359            identity.clone(),
360            observation.snapshot.status,
361            now_ms().max(observation.snapshot.updated_at_ms),
362            after_sequence,
363            &observation.page,
364        )
365        .map_err(Into::into)
366    }
367
368    /// Execute the canonical transport-facing event page query.
369    pub async fn event_page_for(
370        &self,
371        request: &AgentProtocolEventPageRequestV1,
372    ) -> Result<AgentProtocolEventPageV1, AgentProtocolHostError> {
373        request.validate()?;
374        self.event_page(
375            &request.identity,
376            request.after_event_sequence,
377            usize::from(request.limit),
378        )
379        .await
380    }
381
382    /// Read the immutable Git-compatible patch captured for one terminal run.
383    pub async fn change_set_for(
384        &self,
385        request: &AgentProtocolChangeSetRequestV1,
386    ) -> Result<AgentProtocolChangeSetV1, AgentProtocolHostError> {
387        request.validate()?;
388        self.validate_identity(&request.identity)?;
389        let snapshot = self.snapshot(&request.identity).await?;
390        if !snapshot.status.is_terminal() {
391            return Err(AgentProtocolHostError::ChangeSetPending);
392        }
393        if let Some(change_set) = snapshot.workspace_change_set {
394            let response = AgentProtocolChangeSetV1 {
395                schema: AgentProtocolChangeSetV1::SCHEMA.into(),
396                identity: request.identity.clone(),
397                state: snapshot.status.into(),
398                format: AGENT_PROTOCOL_CHANGE_SET_FORMAT_V1.into(),
399                encoding: AGENT_PROTOCOL_CHANGE_SET_ENCODING_V1.into(),
400                base_tree: change_set.base_tree,
401                result_tree: change_set.result_tree,
402                patch_digest: change_set.patch_digest,
403                patch_bytes: change_set.patch_bytes,
404                patch_base64: change_set.patch_base64,
405                observed_at_ms: change_set.observed_at_ms,
406            };
407            response.validate()?;
408            return Ok(response);
409        }
410        match self
411            .change_set_states
412            .read()
413            .await
414            .get(&request.identity.run_id)
415            .copied()
416        {
417            Some(ChangeSetCaptureState::Capturing) => Err(AgentProtocolHostError::ChangeSetPending),
418            None => Err(AgentProtocolHostError::ChangeSetUnavailable),
419        }
420    }
421
422    async fn prepare_change_set(
423        &self,
424        identity: &AgentProtocolRunIdentityV1,
425    ) -> Option<crate::git::WorkspaceTreeSnapshot> {
426        if self.session.run_snapshot(&identity.run_id).await.is_some() {
427            return None;
428        }
429        let workspace = self.session.workspace().to_path_buf();
430        let baseline =
431            tokio::task::spawn_blocking(move || crate::git::snapshot_workspace_tree(&workspace))
432                .await
433                .ok()
434                .and_then(Result::ok);
435        if baseline.is_some() {
436            self.change_set_states
437                .write()
438                .await
439                .insert(identity.run_id.clone(), ChangeSetCaptureState::Capturing);
440        }
441        baseline
442    }
443
444    async fn mark_change_set_unavailable(&self, identity: &AgentProtocolRunIdentityV1) {
445        self.change_set_states
446            .write()
447            .await
448            .remove(&identity.run_id);
449    }
450
451    async fn detach_with_change_set_capture(
452        &self,
453        identity: &AgentProtocolRunIdentityV1,
454        spawned: AgentRunSpawn,
455        baseline: Option<crate::git::WorkspaceTreeSnapshot>,
456    ) -> bool {
457        match spawned {
458            AgentRunSpawn::Started { worker, .. } => {
459                let Some(baseline) = baseline else {
460                    drop(worker);
461                    return false;
462                };
463                let workspace = self.session.workspace().to_path_buf();
464                let session = Arc::clone(&self.session);
465                let run_id = identity.run_id.clone();
466                let pin_identity = format!(
467                    "{}:{}:{}",
468                    self.agent_release_identity, identity.session_id, identity.run_id
469                );
470                let states = Arc::clone(&self.change_set_states);
471                tokio::spawn(async move {
472                    let _ = worker.await;
473                    let run_id_for_state = run_id.clone();
474                    let evidence = tokio::task::spawn_blocking(move || {
475                        let result = crate::git::snapshot_workspace_tree(&workspace)?;
476                        let patch = crate::git::diff_workspace_trees(
477                            &workspace,
478                            &baseline,
479                            &result,
480                            AGENT_PROTOCOL_MAX_CHANGE_SET_BYTES,
481                        )?;
482                        crate::git::pin_workspace_tree(&workspace, &pin_identity, &result)?;
483                        let patch_bytes = u64::try_from(patch.len())
484                            .map_err(|_| anyhow::anyhow!("change-set byte count overflowed"))?;
485                        Ok::<_, anyhow::Error>(RunWorkspaceChangeSet {
486                            base_tree: format!("git-tree:{}", baseline.tree),
487                            result_tree: format!("git-tree:{}", result.tree),
488                            patch_digest: format!("sha256:{:x}", Sha256::digest(&patch)),
489                            patch_bytes,
490                            patch_base64: base64::engine::general_purpose::STANDARD.encode(patch),
491                            observed_at_ms: now_ms(),
492                        })
493                    })
494                    .await
495                    .ok()
496                    .and_then(Result::ok);
497                    if let Some(evidence) = evidence {
498                        let _ = session
499                            .record_workspace_change_set(&run_id_for_state, evidence)
500                            .await;
501                    }
502                    // A completed or failed capture is represented by the
503                    // authoritative Run snapshot. Retain only in-flight
504                    // entries so this map cannot grow with session age.
505                    states.write().await.remove(&run_id_for_state);
506                });
507                false
508            }
509            AgentRunSpawn::Replayed { .. } => true,
510        }
511    }
512
513    fn validate_identity(
514        &self,
515        identity: &AgentProtocolRunIdentityV1,
516    ) -> Result<(), AgentProtocolHostError> {
517        if identity.agent_release_identity != self.agent_release_identity {
518            return Err(AgentProtocolHostError::ReleaseMismatch);
519        }
520        if identity.session_id != self.session.session_id() {
521            return Err(AgentProtocolHostError::SessionMismatch);
522        }
523        Ok(())
524    }
525
526    async fn snapshot(
527        &self,
528        identity: &AgentProtocolRunIdentityV1,
529    ) -> Result<RunSnapshot, AgentProtocolHostError> {
530        let snapshot = self
531            .session
532            .run_snapshot(&identity.run_id)
533            .await
534            .ok_or(AgentProtocolHostError::RunNotFound)?;
535        if snapshot.session_id != identity.session_id {
536            return Err(AgentProtocolHostError::SessionMismatch);
537        }
538        Ok(snapshot)
539    }
540}
541
542fn now_ms() -> u64 {
543    std::time::SystemTime::now()
544        .duration_since(std::time::UNIX_EPOCH)
545        .map(|duration| duration.as_millis() as u64)
546        .unwrap_or_default()
547}
a3s_code_core/agent_protocol_host.rs

a3s_code_core/
agent_protocol_host.rs
