agentos_native_sidecar_core/
router.rs

1use crate::frames::{reject, DispatchResult};
2use agentos_sidecar_protocol::protocol::{
3    AuthenticateRequest, BootstrapRootFilesystemRequest, CloseStdinRequest, ConfigureVmRequest,
4    CreateLayerRequest, CreateOverlayRequest, CreateVmRequest, DisposeVmRequest, ExecuteRequest,
5    ExportSnapshotRequest, ExtEnvelope, FindBoundUdpRequest, FindListenerRequest,
6    GetProcessSnapshotRequest, GetResourceSnapshotRequest, GetSignalStateRequest,
7    GetZombieTimerCountRequest, GuestFilesystemCallRequest, GuestKernelCallRequest,
8    ImportSnapshotRequest, KillProcessRequest, LinkPackageRequest, OpenSessionRequest,
9    OwnershipScope, ProvidedCommandsRequest, RegisterHostCallbacksRequest, RequestFrame,
10    RequestPayload, ResizePtyRequest, SealLayerRequest, SnapshotRootFilesystemRequest,
11    VmFetchRequest, WriteStdinRequest,
12};
13use agentos_sidecar_protocol::wire as generated_wire;
14
15pub const UNSUPPORTED_HOST_CALLBACK_DIRECTION_CODE: &str = "unsupported_direction";
16pub const UNSUPPORTED_HOST_CALLBACK_DIRECTION_MESSAGE: &str =
17    "host callback request categories are sidecar-to-host only in this scaffold";
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub enum RequestDispatchMode {
21    Immediate,
22    Async,
23}
24
25// Request payload variants intentionally vary widely in size (small acks next
26// to bulky create/exec payloads); boxing is a wire-adjacent refactor.
27#[allow(clippy::large_enum_variant)]
28#[derive(Debug, Clone)]
29pub enum RequestRoute {
30    Authenticate(AuthenticateRequest),
31    OpenSession(OpenSessionRequest),
32    CreateVm(CreateVmRequest),
33    DisposeVm(DisposeVmRequest),
34    BootstrapRootFilesystem(BootstrapRootFilesystemRequest),
35    ConfigureVm(ConfigureVmRequest),
36    RegisterHostCallbacks(RegisterHostCallbacksRequest),
37    CreateLayer(CreateLayerRequest),
38    SealLayer(SealLayerRequest),
39    ImportSnapshot(ImportSnapshotRequest),
40    ExportSnapshot(ExportSnapshotRequest),
41    CreateOverlay(CreateOverlayRequest),
42    GuestFilesystemCall(GuestFilesystemCallRequest),
43    GuestKernelCall(GuestKernelCallRequest),
44    SnapshotRootFilesystem(SnapshotRootFilesystemRequest),
45    Execute(ExecuteRequest),
46    WriteStdin(WriteStdinRequest),
47    ResizePty(ResizePtyRequest),
48    CloseStdin(CloseStdinRequest),
49    KillProcess(KillProcessRequest),
50    GetProcessSnapshot(GetProcessSnapshotRequest),
51    GetResourceSnapshot(GetResourceSnapshotRequest),
52    FindListener(FindListenerRequest),
53    FindBoundUdp(FindBoundUdpRequest),
54    VmFetch(VmFetchRequest),
55    GetSignalState(GetSignalStateRequest),
56    GetZombieTimerCount(GetZombieTimerCountRequest),
57    LinkPackage(LinkPackageRequest),
58    ProvidedCommands(ProvidedCommandsRequest),
59    Ext(ExtEnvelope),
60    UnsupportedHostCallbackDirection,
61}
62
63#[derive(Debug, Clone, Copy, PartialEq, Eq)]
64pub enum BlockingExtensionInterrupt<'a> {
65    ExtensionPayload(&'a [u8]),
66    KillProcess,
67}
68
69pub fn route_request_payload(request: &RequestFrame) -> RequestRoute {
70    match request.payload.clone() {
71        RequestPayload::Authenticate(payload) => RequestRoute::Authenticate(payload),
72        RequestPayload::OpenSession(payload) => RequestRoute::OpenSession(payload),
73        RequestPayload::CreateVm(payload) => RequestRoute::CreateVm(payload),
74        RequestPayload::DisposeVm(payload) => RequestRoute::DisposeVm(payload),
75        RequestPayload::BootstrapRootFilesystem(payload) => {
76            RequestRoute::BootstrapRootFilesystem(payload)
77        }
78        RequestPayload::ConfigureVm(payload) => RequestRoute::ConfigureVm(payload),
79        RequestPayload::RegisterHostCallbacks(payload) => {
80            RequestRoute::RegisterHostCallbacks(payload)
81        }
82        RequestPayload::CreateLayer(payload) => RequestRoute::CreateLayer(payload),
83        RequestPayload::SealLayer(payload) => RequestRoute::SealLayer(payload),
84        RequestPayload::ImportSnapshot(payload) => RequestRoute::ImportSnapshot(payload),
85        RequestPayload::ExportSnapshot(payload) => RequestRoute::ExportSnapshot(payload),
86        RequestPayload::CreateOverlay(payload) => RequestRoute::CreateOverlay(payload),
87        RequestPayload::GuestFilesystemCall(payload) => RequestRoute::GuestFilesystemCall(payload),
88        RequestPayload::GuestKernelCall(payload) => RequestRoute::GuestKernelCall(payload),
89        RequestPayload::SnapshotRootFilesystem(payload) => {
90            RequestRoute::SnapshotRootFilesystem(payload)
91        }
92        RequestPayload::Execute(payload) => RequestRoute::Execute(payload),
93        RequestPayload::WriteStdin(payload) => RequestRoute::WriteStdin(payload),
94        RequestPayload::ResizePty(payload) => RequestRoute::ResizePty(payload),
95        RequestPayload::CloseStdin(payload) => RequestRoute::CloseStdin(payload),
96        RequestPayload::KillProcess(payload) => RequestRoute::KillProcess(payload),
97        RequestPayload::GetProcessSnapshot(payload) => RequestRoute::GetProcessSnapshot(payload),
98        RequestPayload::GetResourceSnapshot(payload) => RequestRoute::GetResourceSnapshot(payload),
99        RequestPayload::FindListener(payload) => RequestRoute::FindListener(payload),
100        RequestPayload::FindBoundUdp(payload) => RequestRoute::FindBoundUdp(payload),
101        RequestPayload::VmFetch(payload) => RequestRoute::VmFetch(payload),
102        RequestPayload::GetSignalState(payload) => RequestRoute::GetSignalState(payload),
103        RequestPayload::GetZombieTimerCount(payload) => RequestRoute::GetZombieTimerCount(payload),
104        RequestPayload::LinkPackage(payload) => RequestRoute::LinkPackage(payload),
105        RequestPayload::ProvidedCommands(payload) => RequestRoute::ProvidedCommands(payload),
106        RequestPayload::HostFilesystemCall(_)
107        | RequestPayload::PersistenceLoad(_)
108        | RequestPayload::PersistenceFlush(_) => RequestRoute::UnsupportedHostCallbackDirection,
109        RequestPayload::Ext(payload) => RequestRoute::Ext(payload),
110    }
111}
112
113pub fn generated_wire_blocking_extension_interrupt<'a>(
114    active_request: &generated_wire::RequestFrame,
115    blocking_namespace: &str,
116    interrupting_request: &'a generated_wire::RequestFrame,
117) -> Option<BlockingExtensionInterrupt<'a>> {
118    if interrupting_request.ownership != active_request.ownership {
119        return None;
120    }
121
122    match &interrupting_request.payload {
123        generated_wire::RequestPayload::ExtEnvelope(envelope)
124            if envelope.namespace == blocking_namespace =>
125        {
126            Some(BlockingExtensionInterrupt::ExtensionPayload(
127                &envelope.payload,
128            ))
129        }
130        generated_wire::RequestPayload::ExtEnvelope(_) => None,
131        generated_wire::RequestPayload::KillProcessRequest(_) => {
132            Some(BlockingExtensionInterrupt::KillProcess)
133        }
134        _ => None,
135    }
136}
137
138pub fn request_dispatch_mode(request: &RequestFrame) -> RequestDispatchMode {
139    match request.payload {
140        RequestPayload::DisposeVm(_) | RequestPayload::Ext(_) => RequestDispatchMode::Async,
141        RequestPayload::Authenticate(_)
142        | RequestPayload::OpenSession(_)
143        | RequestPayload::CreateVm(_)
144        | RequestPayload::BootstrapRootFilesystem(_)
145        | RequestPayload::ConfigureVm(_)
146        | RequestPayload::RegisterHostCallbacks(_)
147        | RequestPayload::CreateLayer(_)
148        | RequestPayload::SealLayer(_)
149        | RequestPayload::ImportSnapshot(_)
150        | RequestPayload::ExportSnapshot(_)
151        | RequestPayload::CreateOverlay(_)
152        | RequestPayload::GuestFilesystemCall(_)
153        | RequestPayload::GuestKernelCall(_)
154        | RequestPayload::SnapshotRootFilesystem(_)
155        | RequestPayload::Execute(_)
156        | RequestPayload::WriteStdin(_)
157        | RequestPayload::ResizePty(_)
158        | RequestPayload::CloseStdin(_)
159        | RequestPayload::KillProcess(_)
160        | RequestPayload::GetProcessSnapshot(_)
161        | RequestPayload::GetResourceSnapshot(_)
162        | RequestPayload::FindListener(_)
163        | RequestPayload::FindBoundUdp(_)
164        | RequestPayload::VmFetch(_)
165        | RequestPayload::GetSignalState(_)
166        | RequestPayload::GetZombieTimerCount(_)
167        | RequestPayload::LinkPackage(_)
168        | RequestPayload::ProvidedCommands(_)
169        | RequestPayload::HostFilesystemCall(_)
170        | RequestPayload::PersistenceLoad(_)
171        | RequestPayload::PersistenceFlush(_) => RequestDispatchMode::Immediate,
172    }
173}
174
175pub fn request_is_unsupported_host_callback_direction(request: &RequestFrame) -> bool {
176    matches!(
177        request.payload,
178        RequestPayload::HostFilesystemCall(_)
179            | RequestPayload::PersistenceLoad(_)
180            | RequestPayload::PersistenceFlush(_)
181    )
182}
183
184pub fn unsupported_host_callback_direction_dispatch(request: &RequestFrame) -> DispatchResult {
185    debug_assert!(request_is_unsupported_host_callback_direction(request));
186    DispatchResult {
187        response: reject(
188            request,
189            UNSUPPORTED_HOST_CALLBACK_DIRECTION_CODE,
190            UNSUPPORTED_HOST_CALLBACK_DIRECTION_MESSAGE,
191        ),
192        events: Vec::new(),
193    }
194}
195
196pub fn connection_id_of(ownership: &OwnershipScope) -> Option<String> {
197    match ownership {
198        OwnershipScope::ConnectionOwnership(ownership) => Some(ownership.connection_id.clone()),
199        OwnershipScope::SessionOwnership(ownership) => Some(ownership.connection_id.clone()),
200        OwnershipScope::VmOwnership(ownership) => Some(ownership.connection_id.clone()),
201    }
202}
203
204pub fn session_scope_of(ownership: &OwnershipScope) -> Option<(String, String)> {
205    match ownership {
206        OwnershipScope::SessionOwnership(ownership) => Some((
207            ownership.connection_id.clone(),
208            ownership.session_id.clone(),
209        )),
210        OwnershipScope::VmOwnership(ownership) => Some((
211            ownership.connection_id.clone(),
212            ownership.session_id.clone(),
213        )),
214        OwnershipScope::ConnectionOwnership(_) => None,
215    }
216}
217
218pub fn vm_id_of(ownership: &OwnershipScope) -> Option<String> {
219    match ownership {
220        OwnershipScope::VmOwnership(ownership) => Some(ownership.vm_id.clone()),
221        OwnershipScope::ConnectionOwnership(_) | OwnershipScope::SessionOwnership(_) => None,
222    }
223}
224
225#[cfg(test)]
226mod tests {
227    use super::*;
228    use agentos_sidecar_protocol::protocol::{
229        AuthenticateRequest, ExtEnvelope, FilesystemOperation, HostFilesystemCallRequest,
230        OwnershipScope, PersistenceFlushRequest, PersistenceLoadRequest, ResponsePayload,
231        PROTOCOL_VERSION,
232    };
233    use agentos_sidecar_protocol::wire as generated_wire;
234
235    fn request(payload: RequestPayload) -> RequestFrame {
236        RequestFrame::new(7, OwnershipScope::connection("conn"), payload)
237    }
238
239    fn generated_request(
240        request_id: i64,
241        ownership: generated_wire::OwnershipScope,
242        payload: generated_wire::RequestPayload,
243    ) -> generated_wire::RequestFrame {
244        generated_wire::RequestFrame {
245            schema: generated_wire::protocol_schema(),
246            request_id,
247            ownership,
248            payload,
249        }
250    }
251
252    fn reverse_host_callback_payloads() -> Vec<RequestPayload> {
253        vec![
254            RequestPayload::HostFilesystemCall(HostFilesystemCallRequest {
255                operation: FilesystemOperation::Read,
256                path: String::from("/state"),
257                payload_size_bytes: 0,
258            }),
259            RequestPayload::PersistenceLoad(PersistenceLoadRequest {
260                key: String::from("state"),
261            }),
262            RequestPayload::PersistenceFlush(PersistenceFlushRequest {
263                key: String::from("state"),
264                payload_size_bytes: 0,
265            }),
266        ]
267    }
268
269    #[test]
270    fn dispose_and_ext_requests_are_async() {
271        let ext = request(RequestPayload::Ext(ExtEnvelope {
272            namespace: String::from("test"),
273            payload: Vec::new(),
274        }));
275        assert_eq!(request_dispatch_mode(&ext), RequestDispatchMode::Async);
276    }
277
278    #[test]
279    fn normal_requests_are_immediate() {
280        let authenticate = request(RequestPayload::Authenticate(AuthenticateRequest {
281            client_name: String::from("test"),
282            auth_token: String::from("token"),
283            protocol_version: PROTOCOL_VERSION,
284            bridge_version: 1,
285        }));
286        assert_eq!(
287            request_dispatch_mode(&authenticate),
288            RequestDispatchMode::Immediate
289        );
290    }
291
292    #[test]
293    fn host_callback_requests_are_identified_as_reverse_direction_only() {
294        for payload in reverse_host_callback_payloads() {
295            let host_call = request(payload);
296
297            assert!(request_is_unsupported_host_callback_direction(&host_call));
298            assert_eq!(
299                request_dispatch_mode(&host_call),
300                RequestDispatchMode::Immediate
301            );
302        }
303    }
304
305    #[test]
306    fn routes_protocol_payloads_through_shared_enum() {
307        let authenticate = request(RequestPayload::Authenticate(AuthenticateRequest {
308            client_name: String::from("test"),
309            auth_token: String::from("token"),
310            protocol_version: PROTOCOL_VERSION,
311            bridge_version: 1,
312        }));
313        assert!(matches!(
314            route_request_payload(&authenticate),
315            RequestRoute::Authenticate(_)
316        ));
317
318        let extension = request(RequestPayload::Ext(ExtEnvelope {
319            namespace: String::from("test"),
320            payload: vec![1, 2, 3],
321        }));
322        assert!(matches!(
323            route_request_payload(&extension),
324            RequestRoute::Ext(_)
325        ));
326
327        for payload in reverse_host_callback_payloads() {
328            let host_call = request(payload);
329            assert!(matches!(
330                route_request_payload(&host_call),
331                RequestRoute::UnsupportedHostCallbackDirection
332            ));
333        }
334    }
335
336    #[test]
337    fn unsupported_host_callback_dispatch_rejects_with_shared_code() {
338        for payload in reverse_host_callback_payloads() {
339            let host_call = request(payload);
340
341            let dispatch = unsupported_host_callback_direction_dispatch(&host_call);
342
343            assert!(dispatch.events.is_empty());
344            assert_eq!(dispatch.response.request_id, host_call.request_id);
345            assert_eq!(dispatch.response.ownership, host_call.ownership);
346            match dispatch.response.payload {
347                ResponsePayload::Rejected(rejected) => {
348                    assert_eq!(rejected.code, UNSUPPORTED_HOST_CALLBACK_DIRECTION_CODE);
349                    assert_eq!(
350                        rejected.message,
351                        UNSUPPORTED_HOST_CALLBACK_DIRECTION_MESSAGE
352                    );
353                }
354                other => panic!("unexpected response payload: {other:?}"),
355            }
356        }
357    }
358
359    #[test]
360    fn generated_wire_prompt_interrupt_classifier_matches_only_same_scope_interrupts() {
361        let ownership = generated_wire::OwnershipScope::VmOwnership(generated_wire::VmOwnership {
362            connection_id: String::from("conn"),
363            session_id: String::from("session"),
364            vm_id: String::from("vm"),
365        });
366        let active = generated_request(
367            1,
368            ownership.clone(),
369            generated_wire::RequestPayload::ExtEnvelope(generated_wire::ExtEnvelope {
370                namespace: String::from("prompt"),
371                payload: b"active".to_vec(),
372            }),
373        );
374
375        let same_namespace = generated_request(
376            2,
377            ownership.clone(),
378            generated_wire::RequestPayload::ExtEnvelope(generated_wire::ExtEnvelope {
379                namespace: String::from("prompt"),
380                payload: b"cancel".to_vec(),
381            }),
382        );
383        assert_eq!(
384            generated_wire_blocking_extension_interrupt(&active, "prompt", &same_namespace),
385            Some(BlockingExtensionInterrupt::ExtensionPayload(b"cancel"))
386        );
387
388        let kill = generated_request(
389            3,
390            ownership.clone(),
391            generated_wire::RequestPayload::KillProcessRequest(
392                generated_wire::KillProcessRequest {
393                    process_id: String::from("proc"),
394                    signal: String::from("SIGTERM"),
395                },
396            ),
397        );
398        assert_eq!(
399            generated_wire_blocking_extension_interrupt(&active, "prompt", &kill),
400            Some(BlockingExtensionInterrupt::KillProcess)
401        );
402
403        let other_namespace = generated_request(
404            4,
405            ownership.clone(),
406            generated_wire::RequestPayload::ExtEnvelope(generated_wire::ExtEnvelope {
407                namespace: String::from("other"),
408                payload: b"cancel".to_vec(),
409            }),
410        );
411        assert_eq!(
412            generated_wire_blocking_extension_interrupt(&active, "prompt", &other_namespace),
413            None
414        );
415
416        let other_scope = generated_request(
417            5,
418            generated_wire::OwnershipScope::VmOwnership(generated_wire::VmOwnership {
419                connection_id: String::from("conn"),
420                session_id: String::from("session"),
421                vm_id: String::from("other-vm"),
422            }),
423            generated_wire::RequestPayload::KillProcessRequest(
424                generated_wire::KillProcessRequest {
425                    process_id: String::from("proc"),
426                    signal: String::from("SIGTERM"),
427                },
428            ),
429        );
430        assert_eq!(
431            generated_wire_blocking_extension_interrupt(&active, "prompt", &other_scope),
432            None
433        );
434    }
435
436    #[test]
437    fn ownership_scope_helpers_extract_shared_ids() {
438        let connection = OwnershipScope::connection("conn-1");
439        let session = OwnershipScope::session("conn-1", "session-1");
440        let vm = OwnershipScope::vm("conn-1", "session-1", "vm-1");
441
442        assert_eq!(connection_id_of(&connection).as_deref(), Some("conn-1"));
443        assert_eq!(connection_id_of(&session).as_deref(), Some("conn-1"));
444        assert_eq!(connection_id_of(&vm).as_deref(), Some("conn-1"));
445        assert_eq!(
446            session_scope_of(&session),
447            Some((String::from("conn-1"), String::from("session-1")))
448        );
449        assert_eq!(
450            session_scope_of(&vm),
451            Some((String::from("conn-1"), String::from("session-1")))
452        );
453        assert_eq!(session_scope_of(&connection), None);
454        assert_eq!(vm_id_of(&vm).as_deref(), Some("vm-1"));
455        assert_eq!(vm_id_of(&session), None);
456    }
457}
agentos_native_sidecar_core/router.rs

agentos_native_sidecar_core/
router.rs
