agentos_native_sidecar/
extension.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::time::Duration;
4
5use crate::protocol::{
6    CloseStdinRequest, EventFrame, EventPayload, ExecuteRequest, ExtEnvelope,
7    GuestFilesystemCallRequest, GuestFilesystemResultResponse, KillProcessRequest, OwnershipScope,
8    ProcessKilledResponse, ProcessStartedResponse, SidecarRequestPayload, SidecarResponsePayload,
9    StdinClosedResponse, StdinWrittenResponse, WriteStdinRequest,
10};
11use crate::state::{SharedEventSink, SharedSidecarRequestClient, SidecarError};
12
13pub type ExtensionFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, SidecarError>> + 'a>>;
14
15/// One projected agent package's launch surface, served from sidecar-owned VM
16/// state (sourced from packed vbare manifests; packed packages ship no
17/// `agentos-package.json` for extensions to read from the guest filesystem).
18#[derive(Debug, Clone)]
19pub struct ProjectedAgentLaunchEntry {
20    pub id: String,
21    pub acp_entrypoint: String,
22    pub env: std::collections::BTreeMap<String, String>,
23    pub launch_args: Vec<String>,
24}
25
26pub trait ExtensionHost {
27    fn projected_agents<'a>(
28        &'a mut self,
29        ownership: OwnershipScope,
30    ) -> ExtensionFuture<'a, Vec<ProjectedAgentLaunchEntry>>;
31
32    fn spawn_process<'a>(
33        &'a mut self,
34        ownership: OwnershipScope,
35        request: ExecuteRequest,
36    ) -> ExtensionFuture<'a, ProcessStartedResponse>;
37
38    fn write_stdin<'a>(
39        &'a mut self,
40        ownership: OwnershipScope,
41        request: WriteStdinRequest,
42    ) -> ExtensionFuture<'a, StdinWrittenResponse>;
43
44    fn close_stdin<'a>(
45        &'a mut self,
46        ownership: OwnershipScope,
47        request: CloseStdinRequest,
48    ) -> ExtensionFuture<'a, StdinClosedResponse>;
49
50    fn kill_process<'a>(
51        &'a mut self,
52        ownership: OwnershipScope,
53        request: KillProcessRequest,
54    ) -> ExtensionFuture<'a, ProcessKilledResponse>;
55
56    fn poll_event<'a>(
57        &'a mut self,
58        ownership: OwnershipScope,
59        timeout: Duration,
60    ) -> ExtensionFuture<'a, Option<EventFrame>>;
61
62    fn guest_filesystem_call<'a>(
63        &'a mut self,
64        ownership: OwnershipScope,
65        request: GuestFilesystemCallRequest,
66    ) -> ExtensionFuture<'a, GuestFilesystemResultResponse>;
67
68    fn bind_process_to_session<'a>(
69        &'a mut self,
70        ownership: OwnershipScope,
71        namespace: String,
72        ext_session_id: String,
73        process_id: String,
74    ) -> ExtensionFuture<'a, ()>;
75
76    fn bind_vm_to_session<'a>(
77        &'a mut self,
78        ownership: OwnershipScope,
79        namespace: String,
80        ext_session_id: String,
81    ) -> ExtensionFuture<'a, ()>;
82
83    fn dispose_session_resources<'a>(
84        &'a mut self,
85        ownership: OwnershipScope,
86        namespace: String,
87        ext_session_id: String,
88    ) -> ExtensionFuture<'a, Vec<EventFrame>>;
89
90    fn start_buffering_process_output<'a>(
91        &'a mut self,
92        ownership: OwnershipScope,
93        process_id: String,
94    ) -> ExtensionFuture<'a, ()>;
95
96    fn handoff_buffered_process_output<'a>(
97        &'a mut self,
98        ownership: OwnershipScope,
99        namespace: String,
100        ext_session_id: String,
101        process_id: String,
102        timeout: Duration,
103    ) -> ExtensionFuture<'a, ExtensionBufferedProcessOutput>;
104}
105
106#[derive(Debug, Clone, Default, PartialEq, Eq)]
107pub struct ExtensionBufferedProcessOutput {
108    pub stdout: Vec<u8>,
109    pub stderr: Vec<u8>,
110    pub stdout_truncated: bool,
111    pub stderr_truncated: bool,
112}
113
114impl ExtensionBufferedProcessOutput {
115    pub(crate) fn append_stdout(&mut self, chunk: &[u8], cap: usize) {
116        self.stdout_truncated |= append_bounded_bytes(&mut self.stdout, chunk, cap);
117    }
118
119    pub(crate) fn append_stderr(&mut self, chunk: &[u8], cap: usize) {
120        self.stderr_truncated |= append_bounded_bytes(&mut self.stderr, chunk, cap);
121    }
122}
123
124fn append_bounded_bytes(buffer: &mut Vec<u8>, chunk: &[u8], cap: usize) -> bool {
125    buffer.extend_from_slice(chunk);
126    if buffer.len() <= cap {
127        return false;
128    }
129    let remove_len = buffer.len() - cap;
130    buffer.drain(..remove_len);
131    true
132}
133
134#[derive(Debug, Clone)]
135pub struct ExtensionResponse {
136    pub payload: Vec<u8>,
137    pub events: Vec<EventFrame>,
138}
139
140impl ExtensionResponse {
141    pub fn new(payload: Vec<u8>) -> Self {
142        Self {
143            payload,
144            events: Vec::new(),
145        }
146    }
147
148    pub fn with_events(payload: Vec<u8>, events: Vec<EventFrame>) -> Self {
149        Self { payload, events }
150    }
151
152    pub fn with_wire_events(
153        payload: Vec<u8>,
154        events: Vec<crate::wire::EventFrame>,
155    ) -> Result<Self, SidecarError> {
156        let events = events
157            .into_iter()
158            .map(crate::wire::event_frame_to_compat)
159            .collect::<Result<Vec<_>, _>>()
160            .map_err(wire_protocol_error)?;
161        Ok(Self { payload, events })
162    }
163}
164
165#[derive(Clone)]
166pub struct ExtensionSnapshot {
167    namespace: String,
168    ownership: OwnershipScope,
169    sidecar_requests: SharedSidecarRequestClient,
170    event_sink: SharedEventSink,
171}
172
173pub struct ExtensionContext<'a> {
174    snapshot: ExtensionSnapshot,
175    host: &'a mut dyn ExtensionHost,
176}
177
178impl ExtensionSnapshot {
179    pub(crate) fn new(
180        namespace: String,
181        ownership: OwnershipScope,
182        sidecar_requests: SharedSidecarRequestClient,
183        event_sink: SharedEventSink,
184    ) -> Self {
185        Self {
186            namespace,
187            ownership,
188            sidecar_requests,
189            event_sink,
190        }
191    }
192
193    pub fn namespace(&self) -> &str {
194        &self.namespace
195    }
196
197    pub fn ownership(&self) -> &OwnershipScope {
198        &self.ownership
199    }
200
201    pub fn ext_event(&self, payload: Vec<u8>) -> EventFrame {
202        EventFrame::new(
203            self.ownership.clone(),
204            EventPayload::Ext(ExtEnvelope {
205                namespace: self.namespace.clone(),
206                payload,
207            }),
208        )
209    }
210
211    pub fn ext_event_wire(
212        &self,
213        payload: Vec<u8>,
214    ) -> Result<crate::wire::EventFrame, SidecarError> {
215        crate::wire::event_frame_from_compat(self.ext_event(payload)).map_err(wire_protocol_error)
216    }
217
218    /// Emit a wire event frame to the host the instant it is produced, rather
219    /// than collecting it into the dispatch result and waiting for the whole
220    /// request to resolve. Returns `Ok(None)` when delivered live, or
221    /// `Ok(Some(event))` when no live sink is configured (in-process sidecar) so
222    /// the caller can fall back to returning it in the dispatch's batch.
223    pub fn emit_event_wire(
224        &self,
225        event: crate::wire::EventFrame,
226    ) -> Result<Option<crate::wire::EventFrame>, SidecarError> {
227        self.event_sink.try_emit(event)
228    }
229
230    /// Build a namespaced ext-event from `payload` and emit it live (see
231    /// [`Self::emit_event_wire`]).
232    pub fn emit_ext_event(
233        &self,
234        payload: Vec<u8>,
235    ) -> Result<Option<crate::wire::EventFrame>, SidecarError> {
236        let event = self.ext_event_wire(payload)?;
237        self.emit_event_wire(event)
238    }
239
240    pub fn invoke_callback(
241        &self,
242        payload: Vec<u8>,
243        timeout: Duration,
244    ) -> Result<Vec<u8>, SidecarError> {
245        let response = self.sidecar_requests.invoke(
246            self.ownership.clone(),
247            SidecarRequestPayload::Ext(ExtEnvelope {
248                namespace: self.namespace.clone(),
249                payload,
250            }),
251            timeout,
252        )?;
253        extension_callback_response_payload(&self.namespace, response)
254    }
255}
256
257impl<'a> ExtensionContext<'a> {
258    pub(crate) fn new(snapshot: ExtensionSnapshot, host: &'a mut dyn ExtensionHost) -> Self {
259        Self { snapshot, host }
260    }
261
262    pub fn snapshot(&self) -> ExtensionSnapshot {
263        self.snapshot.clone()
264    }
265
266    pub fn namespace(&self) -> &str {
267        self.snapshot.namespace()
268    }
269
270    pub fn ownership(&self) -> &OwnershipScope {
271        self.snapshot.ownership()
272    }
273
274    pub fn ext_event(&self, payload: Vec<u8>) -> EventFrame {
275        self.snapshot.ext_event(payload)
276    }
277
278    pub fn ext_event_wire(
279        &self,
280        payload: Vec<u8>,
281    ) -> Result<crate::wire::EventFrame, SidecarError> {
282        self.snapshot.ext_event_wire(payload)
283    }
284
285    /// Emit a wire event frame to the host live (see
286    /// [`ExtensionSnapshot::emit_event_wire`]).
287    pub fn emit_event_wire(
288        &self,
289        event: crate::wire::EventFrame,
290    ) -> Result<Option<crate::wire::EventFrame>, SidecarError> {
291        self.snapshot.emit_event_wire(event)
292    }
293
294    /// Build a namespaced ext-event from `payload` and emit it live (see
295    /// [`ExtensionSnapshot::emit_ext_event`]).
296    pub fn emit_ext_event(
297        &self,
298        payload: Vec<u8>,
299    ) -> Result<Option<crate::wire::EventFrame>, SidecarError> {
300        self.snapshot.emit_ext_event(payload)
301    }
302
303    pub fn invoke_callback(
304        &self,
305        payload: Vec<u8>,
306        timeout: Duration,
307    ) -> Result<Vec<u8>, SidecarError> {
308        self.snapshot.invoke_callback(payload, timeout)
309    }
310
311    pub async fn spawn_process(
312        &mut self,
313        request: ExecuteRequest,
314    ) -> Result<ProcessStartedResponse, SidecarError> {
315        self.host
316            .spawn_process(self.snapshot.ownership.clone(), request)
317            .await
318    }
319
320    pub async fn spawn_process_wire(
321        &mut self,
322        request: crate::wire::ExecuteRequest,
323    ) -> Result<crate::wire::ProcessStartedResponse, SidecarError> {
324        let payload = crate::wire::request_payload_to_compat(
325            self.snapshot.ownership(),
326            crate::wire::RequestPayload::ExecuteRequest(request),
327        )
328        .map_err(wire_protocol_error)?;
329        let crate::protocol::RequestPayload::Execute(request) = payload else {
330            return Err(unexpected_wire_request_payload("execute"));
331        };
332        let response = self.spawn_process(request).await?;
333        let payload = crate::wire::response_payload_from_compat(
334            self.snapshot.ownership(),
335            crate::protocol::ResponsePayload::ProcessStarted(response),
336        )
337        .map_err(wire_protocol_error)?;
338        let crate::wire::ResponsePayload::ProcessStartedResponse(response) = payload else {
339            return Err(unexpected_wire_response_payload("process started"));
340        };
341        Ok(response)
342    }
343
344    pub async fn write_stdin(
345        &mut self,
346        request: WriteStdinRequest,
347    ) -> Result<StdinWrittenResponse, SidecarError> {
348        self.host
349            .write_stdin(self.snapshot.ownership.clone(), request)
350            .await
351    }
352
353    pub async fn write_stdin_wire(
354        &mut self,
355        request: crate::wire::WriteStdinRequest,
356    ) -> Result<crate::wire::StdinWrittenResponse, SidecarError> {
357        let payload = crate::wire::request_payload_to_compat(
358            self.snapshot.ownership(),
359            crate::wire::RequestPayload::WriteStdinRequest(request),
360        )
361        .map_err(wire_protocol_error)?;
362        let crate::protocol::RequestPayload::WriteStdin(request) = payload else {
363            return Err(unexpected_wire_request_payload("write stdin"));
364        };
365        let response = self.write_stdin(request).await?;
366        let payload = crate::wire::response_payload_from_compat(
367            self.snapshot.ownership(),
368            crate::protocol::ResponsePayload::StdinWritten(response),
369        )
370        .map_err(wire_protocol_error)?;
371        let crate::wire::ResponsePayload::StdinWrittenResponse(response) = payload else {
372            return Err(unexpected_wire_response_payload("stdin written"));
373        };
374        Ok(response)
375    }
376
377    pub async fn close_stdin(
378        &mut self,
379        request: CloseStdinRequest,
380    ) -> Result<StdinClosedResponse, SidecarError> {
381        self.host
382            .close_stdin(self.snapshot.ownership.clone(), request)
383            .await
384    }
385
386    pub async fn close_stdin_wire(
387        &mut self,
388        request: crate::wire::CloseStdinRequest,
389    ) -> Result<crate::wire::StdinClosedResponse, SidecarError> {
390        let payload = crate::wire::request_payload_to_compat(
391            self.snapshot.ownership(),
392            crate::wire::RequestPayload::CloseStdinRequest(request),
393        )
394        .map_err(wire_protocol_error)?;
395        let crate::protocol::RequestPayload::CloseStdin(request) = payload else {
396            return Err(unexpected_wire_request_payload("close stdin"));
397        };
398        let response = self.close_stdin(request).await?;
399        let payload = crate::wire::response_payload_from_compat(
400            self.snapshot.ownership(),
401            crate::protocol::ResponsePayload::StdinClosed(response),
402        )
403        .map_err(wire_protocol_error)?;
404        let crate::wire::ResponsePayload::StdinClosedResponse(response) = payload else {
405            return Err(unexpected_wire_response_payload("stdin closed"));
406        };
407        Ok(response)
408    }
409
410    pub async fn kill_process(
411        &mut self,
412        request: KillProcessRequest,
413    ) -> Result<ProcessKilledResponse, SidecarError> {
414        self.host
415            .kill_process(self.snapshot.ownership.clone(), request)
416            .await
417    }
418
419    pub async fn kill_process_wire(
420        &mut self,
421        request: crate::wire::KillProcessRequest,
422    ) -> Result<crate::wire::ProcessKilledResponse, SidecarError> {
423        let payload = crate::wire::request_payload_to_compat(
424            self.snapshot.ownership(),
425            crate::wire::RequestPayload::KillProcessRequest(request),
426        )
427        .map_err(wire_protocol_error)?;
428        let crate::protocol::RequestPayload::KillProcess(request) = payload else {
429            return Err(unexpected_wire_request_payload("kill process"));
430        };
431        let response = self.kill_process(request).await?;
432        let payload = crate::wire::response_payload_from_compat(
433            self.snapshot.ownership(),
434            crate::protocol::ResponsePayload::ProcessKilled(response),
435        )
436        .map_err(wire_protocol_error)?;
437        let crate::wire::ResponsePayload::ProcessKilledResponse(response) = payload else {
438            return Err(unexpected_wire_response_payload("process killed"));
439        };
440        Ok(response)
441    }
442
443    pub async fn poll_event(
444        &mut self,
445        timeout: Duration,
446    ) -> Result<Option<EventFrame>, SidecarError> {
447        self.host
448            .poll_event(self.snapshot.ownership.clone(), timeout)
449            .await
450    }
451
452    pub async fn poll_event_wire(
453        &mut self,
454        timeout: Duration,
455    ) -> Result<Option<crate::wire::EventFrame>, SidecarError> {
456        self.poll_event(timeout)
457            .await?
458            .map(crate::wire::event_frame_from_compat)
459            .transpose()
460            .map_err(wire_protocol_error)
461    }
462
463    pub async fn guest_filesystem_call(
464        &mut self,
465        request: GuestFilesystemCallRequest,
466    ) -> Result<GuestFilesystemResultResponse, SidecarError> {
467        self.host
468            .guest_filesystem_call(self.snapshot.ownership.clone(), request)
469            .await
470    }
471
472    /// Enumerate the VM's projected agent packages (id + launch surface) from
473    /// sidecar-owned state. This is the agent source of truth for extensions;
474    /// it reflects `ConfigureVm` and live `LinkPackage` updates.
475    pub async fn projected_agents(
476        &mut self,
477    ) -> Result<Vec<ProjectedAgentLaunchEntry>, SidecarError> {
478        let ownership = self.snapshot.ownership().clone();
479        self.host.projected_agents(ownership).await
480    }
481
482    pub async fn guest_filesystem_call_wire(
483        &mut self,
484        request: crate::wire::GuestFilesystemCallRequest,
485    ) -> Result<crate::wire::GuestFilesystemResultResponse, SidecarError> {
486        let payload = crate::wire::request_payload_to_compat(
487            self.snapshot.ownership(),
488            crate::wire::RequestPayload::GuestFilesystemCallRequest(request),
489        )
490        .map_err(wire_protocol_error)?;
491        let crate::protocol::RequestPayload::GuestFilesystemCall(request) = payload else {
492            return Err(unexpected_wire_request_payload("guest filesystem call"));
493        };
494        let response = self.guest_filesystem_call(request).await?;
495        let payload = crate::wire::response_payload_from_compat(
496            self.snapshot.ownership(),
497            crate::protocol::ResponsePayload::GuestFilesystemResult(response),
498        )
499        .map_err(wire_protocol_error)?;
500        let crate::wire::ResponsePayload::GuestFilesystemResultResponse(response) = payload else {
501            return Err(unexpected_wire_response_payload("guest filesystem result"));
502        };
503        Ok(response)
504    }
505
506    pub async fn bind_process_to_session(
507        &mut self,
508        ext_session_id: impl Into<String>,
509        process_id: impl Into<String>,
510    ) -> Result<(), SidecarError> {
511        self.host
512            .bind_process_to_session(
513                self.snapshot.ownership.clone(),
514                self.snapshot.namespace.clone(),
515                ext_session_id.into(),
516                process_id.into(),
517            )
518            .await
519    }
520
521    pub async fn bind_vm_to_session(
522        &mut self,
523        ext_session_id: impl Into<String>,
524    ) -> Result<(), SidecarError> {
525        self.host
526            .bind_vm_to_session(
527                self.snapshot.ownership.clone(),
528                self.snapshot.namespace.clone(),
529                ext_session_id.into(),
530            )
531            .await
532    }
533
534    pub async fn dispose_session_resources(
535        &mut self,
536        ext_session_id: impl Into<String>,
537    ) -> Result<Vec<EventFrame>, SidecarError> {
538        self.host
539            .dispose_session_resources(
540                self.snapshot.ownership.clone(),
541                self.snapshot.namespace.clone(),
542                ext_session_id.into(),
543            )
544            .await
545    }
546
547    pub async fn dispose_session_resources_wire(
548        &mut self,
549        ext_session_id: impl Into<String>,
550    ) -> Result<Vec<crate::wire::EventFrame>, SidecarError> {
551        self.dispose_session_resources(ext_session_id)
552            .await?
553            .into_iter()
554            .map(crate::wire::event_frame_from_compat)
555            .collect::<Result<Vec<_>, _>>()
556            .map_err(wire_protocol_error)
557    }
558
559    pub async fn start_buffering_process_output(
560        &mut self,
561        process_id: impl Into<String>,
562    ) -> Result<(), SidecarError> {
563        self.host
564            .start_buffering_process_output(self.snapshot.ownership.clone(), process_id.into())
565            .await
566    }
567
568    pub async fn handoff_buffered_process_output(
569        &mut self,
570        ext_session_id: impl Into<String>,
571        process_id: impl Into<String>,
572        timeout: Duration,
573    ) -> Result<ExtensionBufferedProcessOutput, SidecarError> {
574        self.host
575            .handoff_buffered_process_output(
576                self.snapshot.ownership.clone(),
577                self.snapshot.namespace.clone(),
578                ext_session_id.into(),
579                process_id.into(),
580                timeout,
581            )
582            .await
583    }
584}
585
586fn wire_protocol_error(error: crate::wire::ProtocolCodecError) -> SidecarError {
587    SidecarError::InvalidState(format!("invalid generated wire protocol frame: {error}"))
588}
589
590fn unexpected_wire_request_payload(operation: &str) -> SidecarError {
591    SidecarError::InvalidState(format!(
592        "generated wire {operation} request converted to the wrong compatibility payload"
593    ))
594}
595
596fn unexpected_wire_response_payload(operation: &str) -> SidecarError {
597    SidecarError::InvalidState(format!(
598        "compatibility {operation} response converted to the wrong generated wire payload"
599    ))
600}
601
602fn extension_callback_response_payload(
603    namespace: &str,
604    response: SidecarResponsePayload,
605) -> Result<Vec<u8>, SidecarError> {
606    match response {
607        SidecarResponsePayload::ExtResult(envelope) if envelope.namespace == namespace => {
608            Ok(envelope.payload)
609        }
610        SidecarResponsePayload::ExtResult(envelope) => Err(SidecarError::InvalidState(format!(
611            "extension callback response namespace {} did not match {}",
612            envelope.namespace, namespace
613        ))),
614        SidecarResponsePayload::HostCallbackResult(_)
615        | SidecarResponsePayload::JsBridgeResult(_) => Err(SidecarError::InvalidState(
616            String::from("extension callback received a non-extension response"),
617        )),
618    }
619}
620
621pub enum ExtensionInterruptRequest<'a> {
622    ExtensionPayload(&'a [u8]),
623    KillProcess,
624}
625
626#[derive(Debug, Clone)]
627pub struct ExtensionInterruptResponse {
628    pub interrupted_response_payload: Vec<u8>,
629    pub interrupting_response_payload: Option<Vec<u8>>,
630}
631
632pub trait Extension: Send + Sync {
633    fn namespace(&self) -> &str;
634
635    fn handle_request<'a>(
636        &'a self,
637        ctx: ExtensionContext<'a>,
638        payload: Vec<u8>,
639    ) -> ExtensionFuture<'a, ExtensionResponse>;
640
641    fn on_vm_created<'a>(&'a self, _ctx: ExtensionSnapshot) -> ExtensionFuture<'a, ()> {
642        Box::pin(async { Ok(()) })
643    }
644
645    /// Per-session teardown hook. The host invokes this for every registered
646    /// extension when a session is disposed because its connection closed
647    /// (`DisposeReason::ConnectionClosed`), giving the extension the disposed
648    /// session's ownership scope so it can release the per-session state it
649    /// keyed on that session. Default is a no-op. This is the only signal an
650    /// extension receives that a client has disconnected, so it is what lets an
651    /// ACP-style extension free per-session state instead of leaking it for the
652    /// process lifetime.
653    fn on_session_disposed<'a>(&'a self, _ctx: ExtensionSnapshot) -> ExtensionFuture<'a, ()> {
654        Box::pin(async { Ok(()) })
655    }
656
657    fn is_blocking_request(&self, _payload: &[u8]) -> bool {
658        false
659    }
660
661    fn interrupt_blocking_request(
662        &self,
663        _blocking_payload: &[u8],
664        _interrupt: ExtensionInterruptRequest<'_>,
665    ) -> Option<ExtensionInterruptResponse> {
666        None
667    }
668
669    fn on_dispose<'a>(&'a self) -> ExtensionFuture<'a, ()> {
670        Box::pin(async { Ok(()) })
671    }
672}
673
674#[cfg(test)]
675mod live_event_tests {
676    use super::*;
677    use crate::state::EventSinkTransport;
678    use std::sync::{Arc, Mutex};
679
680    /// Records every event handed to the live sink, standing in for the stdio
681    /// `FrameEventTransport` that writes `ProtocolFrame::EventFrame`s to stdout.
682    #[derive(Default)]
683    struct RecordingEventSink {
684        events: Arc<Mutex<Vec<crate::wire::EventFrame>>>,
685    }
686
687    impl EventSinkTransport for RecordingEventSink {
688        fn emit_event(&self, event: crate::wire::EventFrame) -> Result<(), SidecarError> {
689            self.events.lock().unwrap().push(event);
690            Ok(())
691        }
692    }
693
694    fn snapshot_with_sink(event_sink: SharedEventSink) -> ExtensionSnapshot {
695        ExtensionSnapshot::new(
696            String::from("dev.rivet.test.live-event"),
697            OwnershipScope::session("conn-live", "sess-live"),
698            SharedSidecarRequestClient::default(),
699            event_sink,
700        )
701    }
702
703    // With a transport configured (the stdio path), an ext event is emitted live
704    // and `emit_ext_event` reports nothing left to batch.
705    #[test]
706    fn emit_ext_event_streams_live_when_sink_configured() {
707        let recorded = Arc::new(Mutex::new(Vec::new()));
708        let mut sink = SharedEventSink::default();
709        sink.set_transport(Arc::new(RecordingEventSink {
710            events: recorded.clone(),
711        }));
712
713        let snapshot = snapshot_with_sink(sink);
714        let leftover = snapshot
715            .emit_ext_event(b"live-update".to_vec())
716            .expect("emit must succeed");
717
718        assert!(
719            leftover.is_none(),
720            "a configured sink consumes the event live, leaving nothing to batch"
721        );
722        let recorded = recorded.lock().unwrap();
723        assert_eq!(recorded.len(), 1, "the event must reach the live transport");
724        // The live frame round-trips back to the namespaced ext payload.
725        let compat = crate::wire::event_frame_to_compat(recorded[0].clone())
726            .expect("recorded frame converts back to compat");
727        match compat.payload {
728            EventPayload::Ext(envelope) => {
729                assert_eq!(envelope.namespace, "dev.rivet.test.live-event");
730                assert_eq!(envelope.payload, b"live-update");
731            }
732            other => panic!("unexpected live event payload: {other:?}"),
733        }
734    }
735
736    // With no transport (an in-process sidecar), the event is handed back so the
737    // caller can fall back to the dispatch-result batch — preserving delivery.
738    #[test]
739    fn emit_ext_event_falls_back_to_batch_without_sink() {
740        let snapshot = snapshot_with_sink(SharedEventSink::default());
741        let leftover = snapshot
742            .emit_ext_event(b"batched-update".to_vec())
743            .expect("emit must succeed");
744
745        let frame = leftover.expect("without a sink the event is returned for batching");
746        let compat = crate::wire::event_frame_to_compat(frame)
747            .expect("returned frame converts back to compat");
748        match compat.payload {
749            EventPayload::Ext(envelope) => assert_eq!(envelope.payload, b"batched-update"),
750            other => panic!("unexpected batched event payload: {other:?}"),
751        }
752    }
753}
agentos_native_sidecar/extension.rs

agentos_native_sidecar/
extension.rs
