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    Unavailable,
121}
122
123impl AgentProtocolHost {
124    pub fn new(
125        agent_release_identity: impl Into<String>,
126        session: Arc<AgentSession>,
127    ) -> Result<Self, AgentProtocolHostError> {
128        let agent_release_identity = agent_release_identity.into();
129        validate_lower_sha256("agent_release_identity", &agent_release_identity)?;
130        Ok(Self {
131            agent_release_identity,
132            session,
133            change_set_admission: Arc::new(Mutex::new(())),
134            change_set_states: Arc::new(RwLock::new(HashMap::new())),
135        })
136    }
137
138    /// Bind an admitted v1 release manifest to its Code session.
139    ///
140    /// Capability compatibility remains an activation concern for the process
141    /// host, but a manifest for another protocol can never enter this v1 host.
142    pub fn from_manifest(
143        manifest: &AgentReleaseManifest,
144        session: Arc<AgentSession>,
145    ) -> Result<Self, AgentProtocolHostError> {
146        if manifest.protocol() != AGENT_PROTOCOL_V1 {
147            return Err(AgentProtocolHostError::ReleaseProtocolMismatch);
148        }
149        Ok(Self {
150            agent_release_identity: manifest.artifact().digest().to_string(),
151            session,
152            change_set_admission: Arc::new(Mutex::new(())),
153            change_set_states: Arc::new(RwLock::new(HashMap::new())),
154        })
155    }
156
157    pub fn agent_release_identity(&self) -> &str {
158        &self.agent_release_identity
159    }
160
161    pub fn session(&self) -> &Arc<AgentSession> {
162        &self.session
163    }
164
165    /// Execute, cancel, or recover one exact run and return a digest-bound
166    /// receipt. Start and recovery return after Code has admitted the detached
167    /// worker; progress is observed through [`Self::event_page`].
168    pub async fn execute(
169        &self,
170        command: &AgentProtocolCommandV1,
171    ) -> Result<AgentProtocolCommandReceiptV1, AgentProtocolHostError> {
172        command.validate()?;
173        self.validate_identity(command.identity())?;
174
175        let _change_set_admission = self.change_set_admission.lock().await;
176        let replayed = match command {
177            AgentProtocolCommandV1::Start { request } => {
178                let baseline = self.prepare_change_set(&request.identity).await;
179                let spawned = match self
180                    .session
181                    .spawn_run_with_id(&request.identity.run_id, &request.prompt)
182                    .await
183                {
184                    Ok(spawned) => spawned,
185                    Err(error) => {
186                        self.mark_change_set_unavailable(&request.identity).await;
187                        return Err(error.into());
188                    }
189                };
190                self.detach_with_change_set_capture(&request.identity, spawned, baseline)
191                    .await
192            }
193            AgentProtocolCommandV1::Recover { request } => {
194                let baseline = self.prepare_change_set(&request.identity).await;
195                let spawned = match self
196                    .session
197                    .spawn_recovery_with_run_id(
198                        &request.checkpoint_run_id,
199                        &request.identity.run_id,
200                    )
201                    .await
202                {
203                    Ok(spawned) => spawned,
204                    Err(error) => {
205                        self.mark_change_set_unavailable(&request.identity).await;
206                        return Err(error.into());
207                    }
208                };
209                self.detach_with_change_set_capture(&request.identity, spawned, baseline)
210                    .await
211            }
212            AgentProtocolCommandV1::Cancel { request } => {
213                let snapshot = self.snapshot(&request.identity).await?;
214                if snapshot.status.is_terminal() {
215                    true
216                } else if self.session.cancel_run(&request.identity.run_id).await {
217                    false
218                } else if self.snapshot(&request.identity).await?.status.is_terminal() {
219                    true
220                } else {
221                    return Err(AgentProtocolHostError::RunUnavailable);
222                }
223            }
224        };
225
226        let snapshot = self.snapshot(command.identity()).await?;
227        let receipt = AgentProtocolCommandReceiptV1 {
228            schema: AgentProtocolCommandReceiptV1::SCHEMA.into(),
229            action: command.action(),
230            request_id: command.request_id().into(),
231            identity: command.identity().clone(),
232            command_digest: command.digest()?,
233            state: snapshot.status.into(),
234            latest_event_sequence_exclusive: u64::try_from(snapshot.event_count)
235                .map_err(|_| AgentProtocolHostError::SequenceOverflow)?,
236            observed_at_ms: now_ms().max(snapshot.updated_at_ms),
237            replayed,
238        };
239        receipt.validate_for(command)?;
240        Ok(receipt)
241    }
242
243    /// Recover only when the locally loadable loop boundary matches the
244    /// logical component of the complete portable-checkpoint descriptor.
245    ///
246    /// This additive entry point leaves [`AgentProtocolCommandV1`] and its v1
247    /// transport shape unchanged. Validation pins the immutable checkpoint in
248    /// memory before workspace baseline capture and before the target Run is
249    /// admitted, so a concurrently overwritten store entry cannot change what
250    /// the worker resumes.
251    pub async fn execute_exact_recovery(
252        &self,
253        request: &AgentProtocolRunRecoverExactV1,
254    ) -> Result<AgentProtocolCommandReceiptV1, AgentProtocolExactRecoveryError> {
255        request.validate()?;
256        self.validate_identity(&request.identity)?;
257        let logical_resume = request.logical_resume()?;
258
259        let _change_set_admission = self.change_set_admission.lock().await;
260        let prepared = self
261            .session
262            .prepare_recovery_with_evidence(
263                logical_resume,
264                &request.checkpoint.descriptor_digest,
265                &request.identity.run_id,
266            )
267            .await?;
268        self.finish_exact_recovery(request, prepared).await
269    }
270
271    /// Recover from the logical value decoded from the same validated portable
272    /// checkpoint as `request`, without first publishing split store writes.
273    pub async fn execute_exact_recovery_from_checkpoint(
274        &self,
275        request: &AgentProtocolRunRecoverExactV1,
276        checkpoint: crate::loop_checkpoint::LoopCheckpoint,
277    ) -> Result<AgentProtocolCommandReceiptV1, AgentProtocolExactRecoveryError> {
278        request.validate()?;
279        self.validate_identity(&request.identity)?;
280        let logical_resume = request.logical_resume()?;
281
282        let _change_set_admission = self.change_set_admission.lock().await;
283        let prepared = self
284            .session
285            .prepare_recovery_from_checkpoint(
286                logical_resume,
287                &request.checkpoint.descriptor_digest,
288                &request.identity.run_id,
289                checkpoint,
290            )
291            .await?;
292        self.finish_exact_recovery(request, prepared).await
293    }
294
295    async fn finish_exact_recovery(
296        &self,
297        request: &AgentProtocolRunRecoverExactV1,
298        prepared: ExactRecoveryPreparation,
299    ) -> Result<AgentProtocolCommandReceiptV1, AgentProtocolExactRecoveryError> {
300        let replayed = match prepared {
301            ExactRecoveryPreparation::Replayed(spawned) => spawned.replayed(),
302            ExactRecoveryPreparation::Ready(prepared) => {
303                let baseline = self.prepare_change_set(&request.identity).await;
304                let spawned = match self.session.spawn_prepared_recovery(prepared).await {
305                    Ok(spawned) => spawned,
306                    Err(error) => {
307                        self.mark_change_set_unavailable(&request.identity).await;
308                        return Err(error.into());
309                    }
310                };
311                self.detach_with_change_set_capture(&request.identity, spawned, baseline)
312                    .await
313            }
314        };
315
316        let snapshot = self.snapshot(&request.identity).await?;
317        let receipt = AgentProtocolCommandReceiptV1 {
318            schema: AgentProtocolCommandReceiptV1::SCHEMA.into(),
319            action: crate::agent_protocol::AgentProtocolCommandActionV1::Recover,
320            request_id: request.request_id.clone(),
321            identity: request.identity.clone(),
322            command_digest: request.digest()?,
323            state: snapshot.status.into(),
324            latest_event_sequence_exclusive: u64::try_from(snapshot.event_count)
325                .map_err(|_| AgentProtocolHostError::SequenceOverflow)?,
326            observed_at_ms: now_ms().max(snapshot.updated_at_ms),
327            replayed,
328        };
329        receipt.validate_for_exact_recovery(request)?;
330        Ok(receipt)
331    }
332
333    /// Project a bounded cursor page directly from Code's authoritative run
334    /// store without introducing a second provider event model.
335    pub async fn event_page(
336        &self,
337        identity: &AgentProtocolRunIdentityV1,
338        after_event_sequence: Option<u64>,
339        limit: usize,
340    ) -> Result<AgentProtocolEventPageV1, AgentProtocolHostError> {
341        identity.validate()?;
342        self.validate_identity(identity)?;
343        if limit == 0 || limit > AGENT_PROTOCOL_MAX_EVENTS_PER_PAGE {
344            return Err(AgentProtocolError::InvalidField("limit").into());
345        }
346        let after_sequence = after_event_sequence
347            .map(|sequence| {
348                usize::try_from(sequence).map_err(|_| AgentProtocolHostError::SequenceOverflow)
349            })
350            .transpose()?;
351        let observation = self
352            .session
353            .run_event_observation(&identity.run_id, after_sequence, limit)
354            .await
355            .ok_or(AgentProtocolHostError::RunNotFound)?;
356        if observation.snapshot.session_id != identity.session_id {
357            return Err(AgentProtocolHostError::SessionMismatch);
358        }
359        AgentProtocolEventPageV1::from_run_page(
360            identity.clone(),
361            observation.snapshot.status,
362            now_ms().max(observation.snapshot.updated_at_ms),
363            after_sequence,
364            &observation.page,
365        )
366        .map_err(Into::into)
367    }
368
369    /// Execute the canonical transport-facing event page query.
370    pub async fn event_page_for(
371        &self,
372        request: &AgentProtocolEventPageRequestV1,
373    ) -> Result<AgentProtocolEventPageV1, AgentProtocolHostError> {
374        request.validate()?;
375        self.event_page(
376            &request.identity,
377            request.after_event_sequence,
378            usize::from(request.limit),
379        )
380        .await
381    }
382
383    /// Read the immutable Git-compatible patch captured for one terminal run.
384    pub async fn change_set_for(
385        &self,
386        request: &AgentProtocolChangeSetRequestV1,
387    ) -> Result<AgentProtocolChangeSetV1, AgentProtocolHostError> {
388        request.validate()?;
389        self.validate_identity(&request.identity)?;
390        let snapshot = self.snapshot(&request.identity).await?;
391        if !snapshot.status.is_terminal() {
392            return Err(AgentProtocolHostError::ChangeSetPending);
393        }
394        if let Some(change_set) = snapshot.workspace_change_set {
395            let response = AgentProtocolChangeSetV1 {
396                schema: AgentProtocolChangeSetV1::SCHEMA.into(),
397                identity: request.identity.clone(),
398                state: snapshot.status.into(),
399                format: AGENT_PROTOCOL_CHANGE_SET_FORMAT_V1.into(),
400                encoding: AGENT_PROTOCOL_CHANGE_SET_ENCODING_V1.into(),
401                base_tree: change_set.base_tree,
402                result_tree: change_set.result_tree,
403                patch_digest: change_set.patch_digest,
404                patch_bytes: change_set.patch_bytes,
405                patch_base64: change_set.patch_base64,
406                observed_at_ms: change_set.observed_at_ms,
407            };
408            response.validate()?;
409            return Ok(response);
410        }
411        match self
412            .change_set_states
413            .read()
414            .await
415            .get(&request.identity.run_id)
416            .copied()
417        {
418            Some(ChangeSetCaptureState::Capturing) => Err(AgentProtocolHostError::ChangeSetPending),
419            Some(ChangeSetCaptureState::Unavailable) | None => {
420                Err(AgentProtocolHostError::ChangeSetUnavailable)
421            }
422        }
423    }
424
425    async fn prepare_change_set(
426        &self,
427        identity: &AgentProtocolRunIdentityV1,
428    ) -> Option<crate::git::WorkspaceTreeSnapshot> {
429        if self.session.run_snapshot(&identity.run_id).await.is_some() {
430            return None;
431        }
432        let workspace = self.session.workspace().to_path_buf();
433        let baseline =
434            tokio::task::spawn_blocking(move || crate::git::snapshot_workspace_tree(&workspace))
435                .await
436                .ok()
437                .and_then(Result::ok);
438        self.change_set_states.write().await.insert(
439            identity.run_id.clone(),
440            if baseline.is_some() {
441                ChangeSetCaptureState::Capturing
442            } else {
443                ChangeSetCaptureState::Unavailable
444            },
445        );
446        baseline
447    }
448
449    async fn mark_change_set_unavailable(&self, identity: &AgentProtocolRunIdentityV1) {
450        self.change_set_states
451            .write()
452            .await
453            .insert(identity.run_id.clone(), ChangeSetCaptureState::Unavailable);
454    }
455
456    async fn detach_with_change_set_capture(
457        &self,
458        identity: &AgentProtocolRunIdentityV1,
459        spawned: AgentRunSpawn,
460        baseline: Option<crate::git::WorkspaceTreeSnapshot>,
461    ) -> bool {
462        match spawned {
463            AgentRunSpawn::Started { worker, .. } => {
464                let Some(baseline) = baseline else {
465                    drop(worker);
466                    return false;
467                };
468                let workspace = self.session.workspace().to_path_buf();
469                let session = Arc::clone(&self.session);
470                let run_id = identity.run_id.clone();
471                let pin_identity = format!(
472                    "{}:{}:{}",
473                    self.agent_release_identity, identity.session_id, identity.run_id
474                );
475                let states = Arc::clone(&self.change_set_states);
476                tokio::spawn(async move {
477                    let _ = worker.await;
478                    let evidence = tokio::task::spawn_blocking(move || {
479                        let result = crate::git::snapshot_workspace_tree(&workspace)?;
480                        let patch = crate::git::diff_workspace_trees(
481                            &workspace,
482                            &baseline,
483                            &result,
484                            AGENT_PROTOCOL_MAX_CHANGE_SET_BYTES,
485                        )?;
486                        crate::git::pin_workspace_tree(&workspace, &pin_identity, &result)?;
487                        let patch_bytes = u64::try_from(patch.len())
488                            .map_err(|_| anyhow::anyhow!("change-set byte count overflowed"))?;
489                        Ok::<_, anyhow::Error>(RunWorkspaceChangeSet {
490                            base_tree: format!("git-tree:{}", baseline.tree),
491                            result_tree: format!("git-tree:{}", result.tree),
492                            patch_digest: format!("sha256:{:x}", Sha256::digest(&patch)),
493                            patch_bytes,
494                            patch_base64: base64::engine::general_purpose::STANDARD.encode(patch),
495                            observed_at_ms: now_ms(),
496                        })
497                    })
498                    .await
499                    .ok()
500                    .and_then(Result::ok);
501                    let available = match evidence {
502                        Some(evidence) => session
503                            .record_workspace_change_set(&run_id, evidence)
504                            .await
505                            .is_ok(),
506                        None => false,
507                    };
508                    let mut states = states.write().await;
509                    if available {
510                        states.remove(&run_id);
511                    } else {
512                        states.insert(run_id, ChangeSetCaptureState::Unavailable);
513                    }
514                });
515                false
516            }
517            AgentRunSpawn::Replayed { .. } => true,
518        }
519    }
520
521    fn validate_identity(
522        &self,
523        identity: &AgentProtocolRunIdentityV1,
524    ) -> Result<(), AgentProtocolHostError> {
525        if identity.agent_release_identity != self.agent_release_identity {
526            return Err(AgentProtocolHostError::ReleaseMismatch);
527        }
528        if identity.session_id != self.session.session_id() {
529            return Err(AgentProtocolHostError::SessionMismatch);
530        }
531        Ok(())
532    }
533
534    async fn snapshot(
535        &self,
536        identity: &AgentProtocolRunIdentityV1,
537    ) -> Result<RunSnapshot, AgentProtocolHostError> {
538        let snapshot = self
539            .session
540            .run_snapshot(&identity.run_id)
541            .await
542            .ok_or(AgentProtocolHostError::RunNotFound)?;
543        if snapshot.session_id != identity.session_id {
544            return Err(AgentProtocolHostError::SessionMismatch);
545        }
546        Ok(snapshot)
547    }
548}
549
550fn now_ms() -> u64 {
551    std::time::SystemTime::now()
552        .duration_since(std::time::UNIX_EPOCH)
553        .map(|duration| duration.as_millis() as u64)
554        .unwrap_or_default()
555}
a3s_code_core/agent_protocol_host.rs

a3s_code_core/
agent_protocol_host.rs
