1use agentos_sidecar_protocol::protocol::{
2 AgentosProjectedAgent, AuthenticateRequest, AuthenticatedResponse, BoundUdpSnapshotResponse,
3 EventFrame, EventPayload, LayerCreatedResponse, LayerSealedResponse, ListenerSnapshotResponse,
4 OverlayCreatedResponse, OwnershipScope, PackageCommands, PackageLinkedResponse,
5 ProcessExitedEvent, ProcessKilledResponse, ProcessOutputEvent, ProcessSnapshotEntry,
6 ProcessSnapshotResponse, ProcessStartedResponse, ProjectedCommand, ProtocolSchema,
7 ProvidedCommandsResponse, RejectedResponse, RequestFrame, RequestId, ResponseFrame,
8 ResponsePayload, RootFilesystemBootstrappedResponse, RootFilesystemEntry,
9 RootFilesystemSnapshotResponse, SessionOpenedResponse, SignalHandlerRegistration,
10 SignalStateResponse, SnapshotExportedResponse, SnapshotImportedResponse, SocketStateEntry,
11 StdinClosedResponse, StdinWrittenResponse, StreamChannel, StructuredEvent,
12 VmConfiguredResponse, VmCreatedResponse, VmDisposedResponse, VmLifecycleEvent,
13 VmLifecycleState, ZombieTimerCountResponse, PROTOCOL_VERSION,
14};
15use std::collections::HashMap;
16
17pub const UNSUPPORTED_GUEST_KERNEL_CALL_EVENT: &str = "guest.kernel_call.unsupported";
18
19#[derive(Debug, Clone)]
20pub struct DispatchResult {
21 pub response: ResponseFrame,
22 pub events: Vec<EventFrame>,
23}
24
25pub fn response_with_ownership(
26 request_id: RequestId,
27 ownership: OwnershipScope,
28 payload: ResponsePayload,
29) -> ResponseFrame {
30 ResponseFrame {
31 schema: ProtocolSchema::current(),
32 request_id,
33 ownership,
34 payload,
35 }
36}
37
38#[derive(Debug, Clone, PartialEq, Eq)]
39pub enum AuthenticateVersionError {
40 ProtocolVersionMismatch(String),
41 BridgeVersionMismatch(String),
42}
43
44impl AuthenticateVersionError {
45 pub fn code(&self) -> &'static str {
46 match self {
47 Self::ProtocolVersionMismatch(_) => "protocol_version_mismatch",
48 Self::BridgeVersionMismatch(_) => "bridge_version_mismatch",
49 }
50 }
51
52 pub fn message(&self) -> &str {
53 match self {
54 Self::ProtocolVersionMismatch(message) | Self::BridgeVersionMismatch(message) => {
55 message
56 }
57 }
58 }
59}
60
61pub fn validate_authenticate_versions(
62 payload: &AuthenticateRequest,
63) -> Result<(), AuthenticateVersionError> {
64 if payload.protocol_version != PROTOCOL_VERSION {
65 return Err(AuthenticateVersionError::ProtocolVersionMismatch(format!(
66 "sidecar protocol version mismatch: expected {}, got {}",
67 PROTOCOL_VERSION, payload.protocol_version
68 )));
69 }
70
71 let expected_bridge_version = agentos_bridge::bridge_contract().version;
72 if payload.bridge_version != expected_bridge_version {
73 return Err(AuthenticateVersionError::BridgeVersionMismatch(format!(
74 "bridge contract version mismatch: expected {expected_bridge_version}, got {}",
75 payload.bridge_version
76 )));
77 }
78
79 Ok(())
80}
81
82pub fn authenticated_response(
83 request_id: RequestId,
84 sidecar_id: impl Into<String>,
85 connection_id: String,
86 max_frame_bytes: u32,
87) -> ResponseFrame {
88 response_with_ownership(
89 request_id,
90 OwnershipScope::connection(&connection_id),
91 ResponsePayload::Authenticated(AuthenticatedResponse {
92 sidecar_id: sidecar_id.into(),
93 connection_id,
94 max_frame_bytes,
95 }),
96 )
97}
98
99pub fn session_opened_response(
100 request_id: RequestId,
101 owner_connection_id: String,
102 session_id: String,
103) -> ResponseFrame {
104 response_with_ownership(
105 request_id,
106 OwnershipScope::session(&owner_connection_id, &session_id),
107 ResponsePayload::SessionOpened(SessionOpenedResponse {
108 session_id,
109 owner_connection_id,
110 }),
111 )
112}
113
114pub fn respond(request: &RequestFrame, payload: ResponsePayload) -> ResponseFrame {
115 response_with_ownership(request.request_id, request.ownership.clone(), payload)
116}
117
118pub fn reject(request: &RequestFrame, code: &str, message: &str) -> ResponseFrame {
119 respond(
120 request,
121 ResponsePayload::Rejected(RejectedResponse {
122 code: code.to_owned(),
123 message: message.to_owned(),
124 }),
125 )
126}
127
128pub fn vm_created_response(request: &RequestFrame, vm_id: String) -> ResponseFrame {
129 respond(
130 request,
131 ResponsePayload::VmCreated(VmCreatedResponse { vm_id }),
132 )
133}
134
135pub fn vm_disposed_response(request: &RequestFrame, vm_id: String) -> ResponseFrame {
136 respond(
137 request,
138 ResponsePayload::VmDisposed(VmDisposedResponse { vm_id }),
139 )
140}
141
142pub fn root_filesystem_bootstrapped_response(
143 request: &RequestFrame,
144 entry_count: u32,
145) -> ResponseFrame {
146 respond(
147 request,
148 ResponsePayload::RootFilesystemBootstrapped(RootFilesystemBootstrappedResponse {
149 entry_count,
150 }),
151 )
152}
153
154pub fn vm_configured_response(
155 request: &RequestFrame,
156 applied_mounts: u32,
157 applied_software: u32,
158 projected_commands: Vec<ProjectedCommand>,
159 agents: Vec<AgentosProjectedAgent>,
160) -> ResponseFrame {
161 respond(
162 request,
163 ResponsePayload::VmConfigured(VmConfiguredResponse {
164 applied_mounts,
165 applied_software,
166 projected_commands,
167 agents,
168 }),
169 )
170}
171
172pub fn package_linked_response(
173 request: &RequestFrame,
174 projected_commands: Vec<ProjectedCommand>,
175 agents: Vec<AgentosProjectedAgent>,
176) -> ResponseFrame {
177 respond(
178 request,
179 ResponsePayload::PackageLinked(PackageLinkedResponse {
180 projected_commands,
181 agents,
182 }),
183 )
184}
185
186pub fn provided_commands_response(
187 request: &RequestFrame,
188 packages: Vec<PackageCommands>,
189) -> ResponseFrame {
190 respond(
191 request,
192 ResponsePayload::ProvidedCommands(ProvidedCommandsResponse { packages }),
193 )
194}
195
196pub fn layer_created_response(request: &RequestFrame, layer_id: String) -> ResponseFrame {
197 respond(
198 request,
199 ResponsePayload::LayerCreated(LayerCreatedResponse { layer_id }),
200 )
201}
202
203pub fn layer_sealed_response(request: &RequestFrame, layer_id: String) -> ResponseFrame {
204 respond(
205 request,
206 ResponsePayload::LayerSealed(LayerSealedResponse { layer_id }),
207 )
208}
209
210pub fn snapshot_imported_response(request: &RequestFrame, layer_id: String) -> ResponseFrame {
211 respond(
212 request,
213 ResponsePayload::SnapshotImported(SnapshotImportedResponse { layer_id }),
214 )
215}
216
217pub fn snapshot_exported_response(
218 request: &RequestFrame,
219 layer_id: String,
220 entries: Vec<RootFilesystemEntry>,
221) -> ResponseFrame {
222 respond(
223 request,
224 ResponsePayload::SnapshotExported(SnapshotExportedResponse { layer_id, entries }),
225 )
226}
227
228pub fn overlay_created_response(request: &RequestFrame, layer_id: String) -> ResponseFrame {
229 respond(
230 request,
231 ResponsePayload::OverlayCreated(OverlayCreatedResponse { layer_id }),
232 )
233}
234
235pub fn root_filesystem_snapshot_response(
236 request: &RequestFrame,
237 entries: Vec<RootFilesystemEntry>,
238) -> ResponseFrame {
239 respond(
240 request,
241 ResponsePayload::RootFilesystemSnapshot(RootFilesystemSnapshotResponse { entries }),
242 )
243}
244
245pub fn process_started_response(
246 request: &RequestFrame,
247 process_id: String,
248 pid: Option<u32>,
249) -> ResponseFrame {
250 respond(
251 request,
252 ResponsePayload::ProcessStarted(ProcessStartedResponse { process_id, pid }),
253 )
254}
255
256pub fn stdin_written_response(
257 request: &RequestFrame,
258 process_id: String,
259 accepted_bytes: u64,
260) -> ResponseFrame {
261 respond(
262 request,
263 ResponsePayload::StdinWritten(StdinWrittenResponse {
264 process_id,
265 accepted_bytes,
266 }),
267 )
268}
269
270pub fn stdin_closed_response(request: &RequestFrame, process_id: String) -> ResponseFrame {
271 respond(
272 request,
273 ResponsePayload::StdinClosed(StdinClosedResponse { process_id }),
274 )
275}
276
277pub fn process_killed_response(request: &RequestFrame, process_id: String) -> ResponseFrame {
278 respond(
279 request,
280 ResponsePayload::ProcessKilled(ProcessKilledResponse { process_id }),
281 )
282}
283
284pub fn process_snapshot_response(
285 request: &RequestFrame,
286 processes: Vec<ProcessSnapshotEntry>,
287) -> ResponseFrame {
288 respond(
289 request,
290 ResponsePayload::ProcessSnapshot(ProcessSnapshotResponse { processes }),
291 )
292}
293
294pub fn listener_snapshot_response(
295 request: &RequestFrame,
296 listener: Option<SocketStateEntry>,
297) -> ResponseFrame {
298 respond(
299 request,
300 ResponsePayload::ListenerSnapshot(ListenerSnapshotResponse { listener }),
301 )
302}
303
304pub fn bound_udp_snapshot_response(
305 request: &RequestFrame,
306 socket: Option<SocketStateEntry>,
307) -> ResponseFrame {
308 respond(
309 request,
310 ResponsePayload::BoundUdpSnapshot(BoundUdpSnapshotResponse { socket }),
311 )
312}
313
314pub fn signal_state_response(
315 request: &RequestFrame,
316 process_id: String,
317 handlers: impl IntoIterator<Item = (u32, SignalHandlerRegistration)>,
318) -> ResponseFrame {
319 respond(
320 request,
321 ResponsePayload::SignalState(SignalStateResponse {
322 process_id,
323 handlers: handlers.into_iter().collect(),
324 }),
325 )
326}
327
328pub fn zombie_timer_count_response(request: &RequestFrame, count: u64) -> ResponseFrame {
329 respond(
330 request,
331 ResponsePayload::ZombieTimerCount(ZombieTimerCountResponse { count }),
332 )
333}
334
335pub fn event(ownership: OwnershipScope, payload: EventPayload) -> EventFrame {
336 EventFrame::new(ownership, payload)
337}
338
339pub fn vm_lifecycle_event(
340 connection_id: &str,
341 session_id: &str,
342 vm_id: &str,
343 state: VmLifecycleState,
344) -> EventFrame {
345 event(
346 OwnershipScope::vm(connection_id, session_id, vm_id),
347 EventPayload::VmLifecycle(VmLifecycleEvent { state }),
348 )
349}
350
351pub fn process_output_event(
352 ownership: OwnershipScope,
353 process_id: &str,
354 channel: StreamChannel,
355 chunk: Vec<u8>,
356) -> EventFrame {
357 event(
358 ownership,
359 EventPayload::ProcessOutput(ProcessOutputEvent {
360 process_id: process_id.to_owned(),
361 channel,
362 chunk,
363 }),
364 )
365}
366
367pub fn process_exited_event(
368 ownership: OwnershipScope,
369 process_id: &str,
370 exit_code: i32,
371) -> EventFrame {
372 event(
373 ownership,
374 EventPayload::ProcessExited(ProcessExitedEvent {
375 process_id: process_id.to_owned(),
376 exit_code,
377 }),
378 )
379}
380
381pub fn unsupported_guest_kernel_call_event(
382 ownership: OwnershipScope,
383 process_id: &str,
384 execution_id: &str,
385 operation: &str,
386 payload_size_bytes: usize,
387) -> EventFrame {
388 event(
389 ownership,
390 EventPayload::Structured(StructuredEvent {
391 name: String::from(UNSUPPORTED_GUEST_KERNEL_CALL_EVENT),
392 detail: unsupported_guest_kernel_call_detail(
393 Some(process_id),
394 execution_id,
395 operation,
396 payload_size_bytes,
397 ),
398 }),
399 )
400}
401
402pub fn unsupported_guest_kernel_call_detail(
403 process_id: Option<&str>,
404 execution_id: &str,
405 operation: &str,
406 payload_size_bytes: usize,
407) -> HashMap<String, String> {
408 let mut detail = HashMap::from([
409 (String::from("execution_id"), execution_id.to_owned()),
410 (String::from("operation"), operation.to_owned()),
411 (
412 String::from("payload_size_bytes"),
413 payload_size_bytes.to_string(),
414 ),
415 ]);
416 if let Some(process_id) = process_id {
417 detail.insert(String::from("process_id"), process_id.to_owned());
418 }
419 detail
420}
421
422#[cfg(test)]
423mod tests {
424 use super::*;
425 use agentos_sidecar_protocol::protocol::RequestPayload;
426
427 fn authenticate_request() -> AuthenticateRequest {
428 AuthenticateRequest {
429 client_name: String::from("test"),
430 auth_token: String::from("token"),
431 protocol_version: agentos_sidecar_protocol::protocol::PROTOCOL_VERSION,
432 bridge_version: agentos_bridge::bridge_contract().version,
433 }
434 }
435
436 #[test]
437 fn reject_preserves_request_identity_and_ownership() {
438 let request = RequestFrame::new(
439 42,
440 OwnershipScope::connection("conn-1"),
441 RequestPayload::Authenticate(authenticate_request()),
442 );
443
444 let response = reject(&request, "bad_request", "nope");
445
446 assert_eq!(response.request_id, request.request_id);
447 assert_eq!(response.ownership, request.ownership);
448 match response.payload {
449 ResponsePayload::Rejected(rejected) => {
450 assert_eq!(rejected.code, "bad_request");
451 assert_eq!(rejected.message, "nope");
452 }
453 other => panic!("unexpected response payload: {other:?}"),
454 }
455 }
456
457 #[test]
458 fn validates_authenticate_versions() {
459 validate_authenticate_versions(&authenticate_request()).expect("current versions");
460
461 let mut stale_protocol = authenticate_request();
462 stale_protocol.protocol_version = stale_protocol.protocol_version.saturating_sub(1);
463 let error = validate_authenticate_versions(&stale_protocol).expect_err("protocol mismatch");
464 assert_eq!(error.code(), "protocol_version_mismatch");
465 assert!(error
466 .message()
467 .contains("sidecar protocol version mismatch"));
468
469 let mut stale_bridge = authenticate_request();
470 stale_bridge.bridge_version = stale_bridge.bridge_version.saturating_sub(1);
471 let error = validate_authenticate_versions(&stale_bridge).expect_err("bridge mismatch");
472 assert_eq!(error.code(), "bridge_version_mismatch");
473 assert!(error.message().contains("bridge contract version mismatch"));
474 }
475
476 #[test]
477 fn authenticated_response_sets_connection_ownership() {
478 let response =
479 authenticated_response(7, "secure-exec-test", String::from("conn-test"), 1024);
480
481 assert_eq!(response.request_id, 7);
482 assert_eq!(response.ownership, OwnershipScope::connection("conn-test"));
483 match response.payload {
484 ResponsePayload::Authenticated(authenticated) => {
485 assert_eq!(authenticated.sidecar_id, "secure-exec-test");
486 assert_eq!(authenticated.connection_id, "conn-test");
487 assert_eq!(authenticated.max_frame_bytes, 1024);
488 }
489 other => panic!("unexpected response payload: {other:?}"),
490 }
491 }
492
493 #[test]
494 fn session_opened_response_sets_session_ownership() {
495 let response =
496 session_opened_response(8, String::from("conn-1"), String::from("session-1"));
497
498 assert_eq!(response.request_id, 8);
499 assert_eq!(
500 response.ownership,
501 OwnershipScope::session("conn-1", "session-1")
502 );
503 match response.payload {
504 ResponsePayload::SessionOpened(opened) => {
505 assert_eq!(opened.owner_connection_id, "conn-1");
506 assert_eq!(opened.session_id, "session-1");
507 }
508 other => panic!("unexpected response payload: {other:?}"),
509 }
510 }
511
512 #[test]
513 fn lifecycle_response_helpers_preserve_request_ownership() {
514 let request = RequestFrame::new(
515 43,
516 OwnershipScope::vm("conn-1", "session-1", "vm-1"),
517 RequestPayload::Authenticate(authenticate_request()),
518 );
519
520 let created = vm_created_response(&request, String::from("vm-1"));
521 assert_eq!(created.request_id, request.request_id);
522 assert_eq!(created.ownership, request.ownership);
523 match created.payload {
524 ResponsePayload::VmCreated(created) => assert_eq!(created.vm_id, "vm-1"),
525 other => panic!("unexpected response payload: {other:?}"),
526 }
527
528 let bootstrapped = root_filesystem_bootstrapped_response(&request, 3);
529 assert_eq!(bootstrapped.request_id, request.request_id);
530 assert_eq!(bootstrapped.ownership, request.ownership);
531 match bootstrapped.payload {
532 ResponsePayload::RootFilesystemBootstrapped(bootstrapped) => {
533 assert_eq!(bootstrapped.entry_count, 3);
534 }
535 other => panic!("unexpected response payload: {other:?}"),
536 }
537
538 let disposed = vm_disposed_response(&request, String::from("vm-1"));
539 assert_eq!(disposed.request_id, request.request_id);
540 assert_eq!(disposed.ownership, request.ownership);
541 match disposed.payload {
542 ResponsePayload::VmDisposed(disposed) => assert_eq!(disposed.vm_id, "vm-1"),
543 other => panic!("unexpected response payload: {other:?}"),
544 }
545 }
546
547 #[test]
548 fn process_started_response_preserves_pid() {
549 let request = RequestFrame::new(
550 44,
551 OwnershipScope::vm("conn-1", "session-1", "vm-1"),
552 RequestPayload::Authenticate(authenticate_request()),
553 );
554
555 let started = process_started_response(&request, String::from("proc-1"), Some(123));
556 assert_eq!(started.request_id, request.request_id);
557 assert_eq!(started.ownership, request.ownership);
558 match started.payload {
559 ResponsePayload::ProcessStarted(started) => {
560 assert_eq!(started.process_id, "proc-1");
561 assert_eq!(started.pid, Some(123));
562 }
563 other => panic!("unexpected response payload: {other:?}"),
564 }
565
566 let started_without_pid = process_started_response(&request, String::from("proc-2"), None);
567 match started_without_pid.payload {
568 ResponsePayload::ProcessStarted(started) => {
569 assert_eq!(started.process_id, "proc-2");
570 assert_eq!(started.pid, None);
571 }
572 other => panic!("unexpected response payload: {other:?}"),
573 }
574 }
575
576 #[test]
577 fn shared_response_helpers_preserve_payloads() {
578 let request = RequestFrame::new(
579 45,
580 OwnershipScope::vm("conn-1", "session-1", "vm-1"),
581 RequestPayload::Authenticate(authenticate_request()),
582 );
583
584 match vm_configured_response(&request, 2, 3, Vec::new(), Vec::new()).payload {
585 ResponsePayload::VmConfigured(configured) => {
586 assert_eq!(configured.applied_mounts, 2);
587 assert_eq!(configured.applied_software, 3);
588 assert!(configured.projected_commands.is_empty());
589 }
590 other => panic!("unexpected response payload: {other:?}"),
591 }
592
593 match snapshot_exported_response(&request, String::from("layer-1"), Vec::new()).payload {
594 ResponsePayload::SnapshotExported(exported) => {
595 assert_eq!(exported.layer_id, "layer-1");
596 assert!(exported.entries.is_empty());
597 }
598 other => panic!("unexpected response payload: {other:?}"),
599 }
600
601 match stdin_written_response(&request, String::from("proc-1"), 9).payload {
602 ResponsePayload::StdinWritten(written) => {
603 assert_eq!(written.process_id, "proc-1");
604 assert_eq!(written.accepted_bytes, 9);
605 }
606 other => panic!("unexpected response payload: {other:?}"),
607 }
608
609 match process_killed_response(&request, String::from("proc-1")).payload {
610 ResponsePayload::ProcessKilled(killed) => {
611 assert_eq!(killed.process_id, "proc-1");
612 }
613 other => panic!("unexpected response payload: {other:?}"),
614 }
615
616 match signal_state_response(&request, String::from("proc-1"), []).payload {
617 ResponsePayload::SignalState(state) => {
618 assert_eq!(state.process_id, "proc-1");
619 assert!(state.handlers.is_empty());
620 }
621 other => panic!("unexpected response payload: {other:?}"),
622 }
623
624 match zombie_timer_count_response(&request, 4).payload {
625 ResponsePayload::ZombieTimerCount(count) => assert_eq!(count.count, 4),
626 other => panic!("unexpected response payload: {other:?}"),
627 }
628 }
629
630 #[test]
631 fn process_event_helpers_build_vm_owned_events() {
632 let ownership = OwnershipScope::vm("conn-1", "session-1", "vm-1");
633
634 let output = process_output_event(
635 ownership.clone(),
636 "proc-1",
637 StreamChannel::Stdout,
638 b"hello".to_vec(),
639 );
640 assert_eq!(output.ownership, ownership);
641 match output.payload {
642 EventPayload::ProcessOutput(event) => {
643 assert_eq!(event.process_id, "proc-1");
644 assert_eq!(event.channel, StreamChannel::Stdout);
645 assert_eq!(event.chunk, b"hello");
646 }
647 other => panic!("unexpected event payload: {other:?}"),
648 }
649
650 let exited = process_exited_event(output.ownership, "proc-1", 7);
651 match exited.payload {
652 EventPayload::ProcessExited(event) => {
653 assert_eq!(event.process_id, "proc-1");
654 assert_eq!(event.exit_code, 7);
655 }
656 other => panic!("unexpected event payload: {other:?}"),
657 }
658 }
659
660 #[test]
661 fn unsupported_guest_kernel_call_event_preserves_execution_identity() {
662 let ownership = OwnershipScope::vm("conn-1", "session-1", "vm-1");
663 let event = unsupported_guest_kernel_call_event(
664 ownership.clone(),
665 "proc-1",
666 "exec-1",
667 "fs.read",
668 17,
669 );
670
671 assert_eq!(event.ownership, ownership);
672 match event.payload {
673 EventPayload::Structured(event) => {
674 assert_eq!(event.name, "guest.kernel_call.unsupported");
675 assert_eq!(event.detail["process_id"], "proc-1");
676 assert_eq!(event.detail["execution_id"], "exec-1");
677 assert_eq!(event.detail["operation"], "fs.read");
678 assert_eq!(event.detail["payload_size_bytes"], "17");
679 }
680 other => panic!("unexpected event payload: {other:?}"),
681 }
682 }
683
684 #[test]
685 fn unsupported_guest_kernel_call_detail_can_omit_process_identity() {
686 let detail = unsupported_guest_kernel_call_detail(None, "exec-1", "fs.read", 17);
687
688 assert_eq!(detail["execution_id"], "exec-1");
689 assert_eq!(detail["operation"], "fs.read");
690 assert_eq!(detail["payload_size_bytes"], "17");
691 assert!(!detail.contains_key("process_id"));
692 }
693}