Skip to main content

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, 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/// Stable failures returned by the Code-owned headless protocol adapter.
20#[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/// One release- and session-bound A3S Code headless protocol host.
56///
57/// Cloud, Fleet, and other callers may transport commands and receipts, but
58/// this adapter is the sole mapping into Code's run lifecycle and event store.
59#[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    /// Bind an admitted v1 release manifest to its Code session.
79    ///
80    /// Capability compatibility remains an activation concern for the process
81    /// host, but a manifest for another protocol can never enter this v1 host.
82    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    /// Execute, cancel, or recover one exact run and return a digest-bound
104    /// receipt. Start and recovery return after Code has admitted the detached
105    /// worker; progress is observed through [`Self::event_page`].
106    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    /// Project a bounded cursor page directly from Code's authoritative run
163    /// store without introducing a second provider event model.
164    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    /// Execute the canonical transport-facing event page query.
197    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            // Tokio tasks continue after their JoinHandle is dropped. The
243            // session remains the cancellation and graceful-shutdown owner.
244            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}