zentinel_agent_protocol/protocol.rs
1//! Agent protocol types and constants.
2//!
3//! This module defines the wire protocol types for communication between
4//! the proxy dataplane and external processing agents.
5
6use bytes::Bytes;
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9
10/// Agent protocol version
11pub const PROTOCOL_VERSION: u32 = 2;
12
13/// Maximum message size for gRPC transport (10MB)
14pub const MAX_MESSAGE_SIZE: usize = 10 * 1024 * 1024;
15
16/// Agent event type
17#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
18#[serde(rename_all = "snake_case")]
19pub enum EventType {
20 /// Agent configuration (sent once when agent connects)
21 Configure,
22 /// Request headers received
23 RequestHeaders,
24 /// Request body chunk received
25 RequestBodyChunk,
26 /// Response headers received
27 ResponseHeaders,
28 /// Response body chunk received
29 ResponseBodyChunk,
30 /// Request/response complete (for logging)
31 RequestComplete,
32 /// WebSocket frame received (after upgrade)
33 WebSocketFrame,
34 /// Guardrail content inspection (prompt injection, PII detection)
35 GuardrailInspect,
36}
37
38/// Agent response decision indicating how to handle a request or response.
39///
40/// This enum represents the decision an agent makes when processing a request or response.
41/// It allows agents to allow, block, redirect, or challenge requests based on their processing logic.
42///
43/// # Variants
44///
45/// - **Allow**: Continue normal processing without modification
46/// - **Block**: Reject the request/response with a custom error response
47/// - **Redirect**: Send the client to a different URL
48/// - **Challenge**: Request additional verification from the client
49///
50/// # Examples
51///
52/// ```rust
53/// use std::collections::HashMap;
54/// use zentinel_agent_protocol::Decision;
55///
56/// // Allow request to proceed normally
57/// let allow = Decision::Allow;
58///
59/// // Block with 403 and custom body
60/// let block = Decision::Block {
61/// status: 403,
62/// body: Some("Access denied".to_string()),
63/// headers: None,
64/// };
65///
66/// // Redirect to login page
67/// let redirect = Decision::Redirect {
68/// url: "https://example.com/login".to_string(),
69/// status: 302,
70/// };
71///
72/// // Challenge with CAPTCHA
73/// let mut params = HashMap::new();
74/// params.insert("type".to_string(), "recaptcha".to_string());
75/// let challenge = Decision::Challenge {
76/// challenge_type: "captcha".to_string(),
77/// params,
78/// };
79/// ```
80#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
81#[serde(rename_all = "snake_case")]
82pub enum Decision {
83 /// Allow the request/response to continue without modification.
84 ///
85 /// This is the default decision that indicates normal processing should proceed.
86 #[default]
87 Allow,
88 /// Block the request/response with a custom error response.
89 ///
90 /// The proxy will return the specified status code and optional body/headers
91 /// instead of forwarding the request to the upstream or returning the original response.
92 Block {
93 /// HTTP status code to return (typically 4xx or 5xx)
94 status: u16,
95 /// Optional response body content
96 body: Option<String>,
97 /// Optional response headers to include
98 headers: Option<HashMap<String, String>>,
99 },
100 /// Redirect the client to a different URL.
101 ///
102 /// The proxy will return a redirect response with the specified URL and status code.
103 Redirect {
104 /// Target URL for the redirect
105 url: String,
106 /// HTTP redirect status code (301, 302, 303, 307, or 308)
107 status: u16,
108 },
109 /// Request additional verification from the client.
110 ///
111 /// This can be used to implement CAPTCHA, multi-factor authentication,
112 /// or other challenge-response mechanisms.
113 Challenge {
114 /// Type of challenge (e.g., "captcha", "otp", "totp")
115 challenge_type: String,
116 /// Challenge-specific parameters and configuration
117 params: HashMap<String, String>,
118 },
119}
120
121/// Header modification operation
122#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
123#[serde(rename_all = "snake_case")]
124pub enum HeaderOp {
125 /// Set a header (replace if exists)
126 Set { name: String, value: String },
127 /// Add a header (append if exists)
128 Add { name: String, value: String },
129 /// Remove a header
130 Remove { name: String },
131}
132
133// ============================================================================
134// Body Mutation
135// ============================================================================
136
137/// Body mutation from agent
138///
139/// Allows agents to modify body content during streaming:
140/// - `None` data: pass through original chunk unchanged
141/// - `Some(empty)`: drop the chunk entirely
142/// - `Some(data)`: replace chunk with modified content
143#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
144pub struct BodyMutation {
145 /// Modified body data (base64 encoded for JSON transport)
146 ///
147 /// - `None`: use original chunk unchanged
148 /// - `Some("")`: drop this chunk
149 /// - `Some(data)`: replace chunk with this data
150 pub data: Option<String>,
151
152 /// Chunk index this mutation applies to
153 ///
154 /// Must match the `chunk_index` from the body chunk event.
155 #[serde(default)]
156 pub chunk_index: u32,
157}
158
159impl BodyMutation {
160 /// Create a pass-through mutation (no change)
161 pub fn pass_through(chunk_index: u32) -> Self {
162 Self {
163 data: None,
164 chunk_index,
165 }
166 }
167
168 /// Create a mutation that drops the chunk
169 pub fn drop_chunk(chunk_index: u32) -> Self {
170 Self {
171 data: Some(String::new()),
172 chunk_index,
173 }
174 }
175
176 /// Create a mutation that replaces the chunk
177 pub fn replace(chunk_index: u32, data: String) -> Self {
178 Self {
179 data: Some(data),
180 chunk_index,
181 }
182 }
183
184 /// Check if this mutation passes through unchanged
185 pub fn is_pass_through(&self) -> bool {
186 self.data.is_none()
187 }
188
189 /// Check if this mutation drops the chunk
190 pub fn is_drop(&self) -> bool {
191 matches!(&self.data, Some(d) if d.is_empty())
192 }
193}
194
195/// Request metadata sent to agents
196#[derive(Debug, Clone, Serialize, Deserialize)]
197pub struct RequestMetadata {
198 /// Correlation ID for request tracing
199 pub correlation_id: String,
200 /// Request ID (internal)
201 pub request_id: String,
202 /// Client IP address
203 pub client_ip: String,
204 /// Client port
205 pub client_port: u16,
206 /// Server name (SNI or Host header)
207 pub server_name: Option<String>,
208 /// Protocol (HTTP/1.1, HTTP/2, etc.)
209 pub protocol: String,
210 /// TLS version if applicable
211 pub tls_version: Option<String>,
212 /// TLS cipher suite if applicable
213 pub tls_cipher: Option<String>,
214 /// Route ID that matched
215 pub route_id: Option<String>,
216 /// Upstream ID
217 pub upstream_id: Option<String>,
218 /// Request start timestamp (RFC3339)
219 pub timestamp: String,
220 /// W3C Trace Context traceparent header (for distributed tracing)
221 ///
222 /// Format: `{version}-{trace-id}-{parent-id}-{trace-flags}`
223 /// Example: `00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01`
224 ///
225 /// Agents can use this to create child spans that link to the proxy's span.
226 #[serde(skip_serializing_if = "Option::is_none")]
227 pub traceparent: Option<String>,
228}
229
230/// Request headers event
231#[derive(Debug, Clone, Serialize, Deserialize)]
232pub struct RequestHeadersEvent {
233 /// Event metadata
234 pub metadata: RequestMetadata,
235 /// HTTP method
236 pub method: String,
237 /// Request URI
238 pub uri: String,
239 /// HTTP headers
240 pub headers: HashMap<String, Vec<String>>,
241}
242
243/// Request body chunk event
244#[derive(Debug, Clone, Serialize, Deserialize)]
245pub struct RequestBodyChunkEvent {
246 /// Correlation ID
247 pub correlation_id: String,
248 /// Body chunk data (base64 encoded for JSON transport)
249 pub data: String,
250 /// Is this the last chunk?
251 pub is_last: bool,
252 /// Total body size if known
253 pub total_size: Option<usize>,
254 /// Chunk index for ordering (0-based)
255 ///
256 /// Used to match mutations to chunks and ensure ordering.
257 #[serde(default)]
258 pub chunk_index: u32,
259 /// Bytes received so far (cumulative)
260 #[serde(default)]
261 pub bytes_received: usize,
262}
263
264/// Response headers event
265#[derive(Debug, Clone, Serialize, Deserialize)]
266pub struct ResponseHeadersEvent {
267 /// Correlation ID
268 pub correlation_id: String,
269 /// HTTP status code
270 pub status: u16,
271 /// HTTP headers
272 pub headers: HashMap<String, Vec<String>>,
273}
274
275/// Response body chunk event
276#[derive(Debug, Clone, Serialize, Deserialize)]
277pub struct ResponseBodyChunkEvent {
278 /// Correlation ID
279 pub correlation_id: String,
280 /// Body chunk data (base64 encoded for JSON transport)
281 pub data: String,
282 /// Is this the last chunk?
283 pub is_last: bool,
284 /// Total body size if known
285 pub total_size: Option<usize>,
286 /// Chunk index for ordering (0-based)
287 #[serde(default)]
288 pub chunk_index: u32,
289 /// Bytes sent so far (cumulative)
290 #[serde(default)]
291 pub bytes_sent: usize,
292}
293
294// ============================================================================
295// Binary Body Chunk Events (Zero-Copy)
296// ============================================================================
297
298/// Binary request body chunk event.
299///
300/// This type uses `Bytes` for zero-copy body streaming, avoiding the base64
301/// encode/decode overhead of `RequestBodyChunkEvent`. Use this type for:
302/// - Binary UDS transport (with `binary-uds` feature)
303/// - gRPC transport (protobuf already uses bytes)
304/// - Any transport that supports raw binary data
305///
306/// For JSON transport, use `RequestBodyChunkEvent` with base64-encoded data.
307#[derive(Debug, Clone)]
308pub struct BinaryRequestBodyChunkEvent {
309 /// Correlation ID
310 pub correlation_id: String,
311 /// Body chunk data (raw bytes, no encoding)
312 pub data: Bytes,
313 /// Is this the last chunk?
314 pub is_last: bool,
315 /// Total body size if known
316 pub total_size: Option<usize>,
317 /// Chunk index for ordering (0-based)
318 pub chunk_index: u32,
319 /// Bytes received so far (cumulative)
320 pub bytes_received: usize,
321}
322
323/// Binary response body chunk event.
324///
325/// This type uses `Bytes` for zero-copy body streaming, avoiding the base64
326/// encode/decode overhead of `ResponseBodyChunkEvent`. Use this type for:
327/// - Binary UDS transport (with `binary-uds` feature)
328/// - gRPC transport (protobuf already uses bytes)
329/// - Any transport that supports raw binary data
330///
331/// For JSON transport, use `ResponseBodyChunkEvent` with base64-encoded data.
332#[derive(Debug, Clone)]
333pub struct BinaryResponseBodyChunkEvent {
334 /// Correlation ID
335 pub correlation_id: String,
336 /// Body chunk data (raw bytes, no encoding)
337 pub data: Bytes,
338 /// Is this the last chunk?
339 pub is_last: bool,
340 /// Total body size if known
341 pub total_size: Option<usize>,
342 /// Chunk index for ordering (0-based)
343 pub chunk_index: u32,
344 /// Bytes sent so far (cumulative)
345 pub bytes_sent: usize,
346}
347
348impl BinaryRequestBodyChunkEvent {
349 /// Create a new binary request body chunk event.
350 pub fn new(
351 correlation_id: impl Into<String>,
352 data: impl Into<Bytes>,
353 chunk_index: u32,
354 is_last: bool,
355 ) -> Self {
356 let data = data.into();
357 Self {
358 correlation_id: correlation_id.into(),
359 bytes_received: data.len(),
360 data,
361 is_last,
362 total_size: None,
363 chunk_index,
364 }
365 }
366
367 /// Set the total body size.
368 pub fn with_total_size(mut self, size: usize) -> Self {
369 self.total_size = Some(size);
370 self
371 }
372
373 /// Set cumulative bytes received.
374 pub fn with_bytes_received(mut self, bytes: usize) -> Self {
375 self.bytes_received = bytes;
376 self
377 }
378}
379
380impl BinaryResponseBodyChunkEvent {
381 /// Create a new binary response body chunk event.
382 pub fn new(
383 correlation_id: impl Into<String>,
384 data: impl Into<Bytes>,
385 chunk_index: u32,
386 is_last: bool,
387 ) -> Self {
388 let data = data.into();
389 Self {
390 correlation_id: correlation_id.into(),
391 bytes_sent: data.len(),
392 data,
393 is_last,
394 total_size: None,
395 chunk_index,
396 }
397 }
398
399 /// Set the total body size.
400 pub fn with_total_size(mut self, size: usize) -> Self {
401 self.total_size = Some(size);
402 self
403 }
404
405 /// Set cumulative bytes sent.
406 pub fn with_bytes_sent(mut self, bytes: usize) -> Self {
407 self.bytes_sent = bytes;
408 self
409 }
410}
411
412// ============================================================================
413// Conversions between String (base64) and Binary body chunk types
414// ============================================================================
415
416impl From<BinaryRequestBodyChunkEvent> for RequestBodyChunkEvent {
417 /// Convert binary body chunk to base64-encoded JSON-compatible type.
418 fn from(event: BinaryRequestBodyChunkEvent) -> Self {
419 use base64::{engine::general_purpose::STANDARD, Engine as _};
420 Self {
421 correlation_id: event.correlation_id,
422 data: STANDARD.encode(&event.data),
423 is_last: event.is_last,
424 total_size: event.total_size,
425 chunk_index: event.chunk_index,
426 bytes_received: event.bytes_received,
427 }
428 }
429}
430
431impl From<&RequestBodyChunkEvent> for BinaryRequestBodyChunkEvent {
432 /// Convert base64-encoded body chunk to binary type.
433 ///
434 /// If base64 decoding fails, falls back to treating data as raw UTF-8 bytes.
435 fn from(event: &RequestBodyChunkEvent) -> Self {
436 use base64::{engine::general_purpose::STANDARD, Engine as _};
437 let data = STANDARD
438 .decode(&event.data)
439 .map(Bytes::from)
440 .unwrap_or_else(|_| Bytes::copy_from_slice(event.data.as_bytes()));
441 Self {
442 correlation_id: event.correlation_id.clone(),
443 data,
444 is_last: event.is_last,
445 total_size: event.total_size,
446 chunk_index: event.chunk_index,
447 bytes_received: event.bytes_received,
448 }
449 }
450}
451
452impl From<BinaryResponseBodyChunkEvent> for ResponseBodyChunkEvent {
453 /// Convert binary body chunk to base64-encoded JSON-compatible type.
454 fn from(event: BinaryResponseBodyChunkEvent) -> Self {
455 use base64::{engine::general_purpose::STANDARD, Engine as _};
456 Self {
457 correlation_id: event.correlation_id,
458 data: STANDARD.encode(&event.data),
459 is_last: event.is_last,
460 total_size: event.total_size,
461 chunk_index: event.chunk_index,
462 bytes_sent: event.bytes_sent,
463 }
464 }
465}
466
467impl From<&ResponseBodyChunkEvent> for BinaryResponseBodyChunkEvent {
468 /// Convert base64-encoded body chunk to binary type.
469 ///
470 /// If base64 decoding fails, falls back to treating data as raw UTF-8 bytes.
471 fn from(event: &ResponseBodyChunkEvent) -> Self {
472 use base64::{engine::general_purpose::STANDARD, Engine as _};
473 let data = STANDARD
474 .decode(&event.data)
475 .map(Bytes::from)
476 .unwrap_or_else(|_| Bytes::copy_from_slice(event.data.as_bytes()));
477 Self {
478 correlation_id: event.correlation_id.clone(),
479 data,
480 is_last: event.is_last,
481 total_size: event.total_size,
482 chunk_index: event.chunk_index,
483 bytes_sent: event.bytes_sent,
484 }
485 }
486}
487
488/// Request complete event (for logging/audit)
489#[derive(Debug, Clone, Serialize, Deserialize)]
490pub struct RequestCompleteEvent {
491 /// Correlation ID
492 pub correlation_id: String,
493 /// Final HTTP status code
494 pub status: u16,
495 /// Request duration in milliseconds
496 pub duration_ms: u64,
497 /// Request body size
498 pub request_body_size: usize,
499 /// Response body size
500 pub response_body_size: usize,
501 /// Upstream attempts
502 pub upstream_attempts: u32,
503 /// Error if any
504 pub error: Option<String>,
505}
506
507// ============================================================================
508// WebSocket Frame Events
509// ============================================================================
510
511/// WebSocket frame event
512///
513/// Sent to agents after a WebSocket upgrade when frame inspection is enabled.
514/// Each frame is sent individually for inspection.
515#[derive(Debug, Clone, Serialize, Deserialize)]
516pub struct WebSocketFrameEvent {
517 /// Correlation ID (same as the original HTTP upgrade request)
518 pub correlation_id: String,
519 /// Frame opcode: "text", "binary", "ping", "pong", "close", "continuation"
520 pub opcode: String,
521 /// Frame payload (base64 encoded for JSON transport)
522 pub data: String,
523 /// Direction: true = client->server, false = server->client
524 pub client_to_server: bool,
525 /// Frame index for this connection (0-based, per direction)
526 pub frame_index: u64,
527 /// FIN bit - true if final frame of message (for fragmented messages)
528 pub fin: bool,
529 /// Route ID
530 pub route_id: Option<String>,
531 /// Client IP
532 pub client_ip: String,
533}
534
535/// WebSocket opcode
536#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
537#[serde(rename_all = "snake_case")]
538pub enum WebSocketOpcode {
539 /// Continuation frame (0x0)
540 Continuation,
541 /// Text frame (0x1)
542 Text,
543 /// Binary frame (0x2)
544 Binary,
545 /// Connection close (0x8)
546 Close,
547 /// Ping (0x9)
548 Ping,
549 /// Pong (0xA)
550 Pong,
551}
552
553impl WebSocketOpcode {
554 /// Convert opcode to string representation
555 pub fn as_str(&self) -> &'static str {
556 match self {
557 Self::Continuation => "continuation",
558 Self::Text => "text",
559 Self::Binary => "binary",
560 Self::Close => "close",
561 Self::Ping => "ping",
562 Self::Pong => "pong",
563 }
564 }
565
566 /// Parse from byte value
567 pub fn from_u8(value: u8) -> Option<Self> {
568 match value {
569 0x0 => Some(Self::Continuation),
570 0x1 => Some(Self::Text),
571 0x2 => Some(Self::Binary),
572 0x8 => Some(Self::Close),
573 0x9 => Some(Self::Ping),
574 0xA => Some(Self::Pong),
575 _ => None,
576 }
577 }
578
579 /// Convert to byte value
580 pub fn as_u8(&self) -> u8 {
581 match self {
582 Self::Continuation => 0x0,
583 Self::Text => 0x1,
584 Self::Binary => 0x2,
585 Self::Close => 0x8,
586 Self::Ping => 0x9,
587 Self::Pong => 0xA,
588 }
589 }
590}
591
592/// WebSocket frame decision
593///
594/// Agents return this decision for WebSocket frame events.
595#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
596#[serde(rename_all = "snake_case")]
597pub enum WebSocketDecision {
598 /// Allow frame to pass through
599 #[default]
600 Allow,
601 /// Drop this frame silently (don't forward)
602 Drop,
603 /// Close the WebSocket connection
604 Close {
605 /// Close code (RFC 6455 section 7.4.1)
606 code: u16,
607 /// Close reason
608 reason: String,
609 },
610}
611
612/// Agent response message
613#[derive(Debug, Clone, Serialize, Deserialize)]
614pub struct AgentResponse {
615 /// Protocol version
616 pub version: u32,
617 /// Decision
618 pub decision: Decision,
619 /// Header modifications for request
620 #[serde(default)]
621 pub request_headers: Vec<HeaderOp>,
622 /// Header modifications for response
623 #[serde(default)]
624 pub response_headers: Vec<HeaderOp>,
625 /// Routing metadata modifications
626 #[serde(default)]
627 pub routing_metadata: HashMap<String, String>,
628 /// Audit metadata
629 #[serde(default)]
630 pub audit: AuditMetadata,
631
632 // ========================================================================
633 // Streaming-specific fields
634 // ========================================================================
635 /// Agent needs more data to make a final decision
636 ///
637 /// When `true`, the current `decision` is provisional and may change
638 /// after processing more body chunks. The proxy should continue
639 /// streaming body data to this agent.
640 ///
641 /// When `false` (default), the decision is final.
642 #[serde(default)]
643 pub needs_more: bool,
644
645 /// Request body mutation (for streaming mode)
646 ///
647 /// If present, applies the mutation to the current request body chunk.
648 /// Only valid for `RequestBodyChunk` events.
649 #[serde(default)]
650 pub request_body_mutation: Option<BodyMutation>,
651
652 /// Response body mutation (for streaming mode)
653 ///
654 /// If present, applies the mutation to the current response body chunk.
655 /// Only valid for `ResponseBodyChunk` events.
656 #[serde(default)]
657 pub response_body_mutation: Option<BodyMutation>,
658
659 /// WebSocket frame decision
660 ///
661 /// Only valid for `WebSocketFrame` events. If not set, defaults to Allow.
662 #[serde(default)]
663 pub websocket_decision: Option<WebSocketDecision>,
664}
665
666impl AgentResponse {
667 /// Create a default allow response
668 pub fn default_allow() -> Self {
669 Self {
670 version: PROTOCOL_VERSION,
671 decision: Decision::Allow,
672 request_headers: vec![],
673 response_headers: vec![],
674 routing_metadata: HashMap::new(),
675 audit: AuditMetadata::default(),
676 needs_more: false,
677 request_body_mutation: None,
678 response_body_mutation: None,
679 websocket_decision: None,
680 }
681 }
682
683 /// Create a block response
684 pub fn block(status: u16, body: Option<String>) -> Self {
685 Self {
686 version: PROTOCOL_VERSION,
687 decision: Decision::Block {
688 status,
689 body,
690 headers: None,
691 },
692 request_headers: vec![],
693 response_headers: vec![],
694 routing_metadata: HashMap::new(),
695 audit: AuditMetadata::default(),
696 needs_more: false,
697 request_body_mutation: None,
698 response_body_mutation: None,
699 websocket_decision: None,
700 }
701 }
702
703 /// Create a redirect response
704 pub fn redirect(url: String, status: u16) -> Self {
705 Self {
706 version: PROTOCOL_VERSION,
707 decision: Decision::Redirect { url, status },
708 request_headers: vec![],
709 response_headers: vec![],
710 routing_metadata: HashMap::new(),
711 audit: AuditMetadata::default(),
712 needs_more: false,
713 request_body_mutation: None,
714 response_body_mutation: None,
715 websocket_decision: None,
716 }
717 }
718
719 /// Create a streaming response indicating more data is needed
720 pub fn needs_more_data() -> Self {
721 Self {
722 version: PROTOCOL_VERSION,
723 decision: Decision::Allow,
724 request_headers: vec![],
725 response_headers: vec![],
726 routing_metadata: HashMap::new(),
727 audit: AuditMetadata::default(),
728 needs_more: true,
729 request_body_mutation: None,
730 response_body_mutation: None,
731 websocket_decision: None,
732 }
733 }
734
735 /// Create a WebSocket allow response
736 pub fn websocket_allow() -> Self {
737 Self {
738 websocket_decision: Some(WebSocketDecision::Allow),
739 ..Self::default_allow()
740 }
741 }
742
743 /// Create a WebSocket drop response (drop the frame, don't forward)
744 pub fn websocket_drop() -> Self {
745 Self {
746 websocket_decision: Some(WebSocketDecision::Drop),
747 ..Self::default_allow()
748 }
749 }
750
751 /// Create a WebSocket close response (close the connection)
752 pub fn websocket_close(code: u16, reason: String) -> Self {
753 Self {
754 websocket_decision: Some(WebSocketDecision::Close { code, reason }),
755 ..Self::default_allow()
756 }
757 }
758
759 /// Set WebSocket decision
760 pub fn with_websocket_decision(mut self, decision: WebSocketDecision) -> Self {
761 self.websocket_decision = Some(decision);
762 self
763 }
764
765 /// Create a streaming response with body mutation
766 pub fn with_request_body_mutation(mut self, mutation: BodyMutation) -> Self {
767 self.request_body_mutation = Some(mutation);
768 self
769 }
770
771 /// Create a streaming response with response body mutation
772 pub fn with_response_body_mutation(mut self, mutation: BodyMutation) -> Self {
773 self.response_body_mutation = Some(mutation);
774 self
775 }
776
777 /// Set needs_more flag
778 pub fn set_needs_more(mut self, needs_more: bool) -> Self {
779 self.needs_more = needs_more;
780 self
781 }
782
783 /// Add a request header modification
784 pub fn add_request_header(mut self, op: HeaderOp) -> Self {
785 self.request_headers.push(op);
786 self
787 }
788
789 /// Add a response header modification
790 pub fn add_response_header(mut self, op: HeaderOp) -> Self {
791 self.response_headers.push(op);
792 self
793 }
794
795 /// Add audit metadata
796 pub fn with_audit(mut self, audit: AuditMetadata) -> Self {
797 self.audit = audit;
798 self
799 }
800}
801
802/// Audit metadata from agent
803#[derive(Debug, Clone, Default, Serialize, Deserialize)]
804pub struct AuditMetadata {
805 /// Tags for logging/metrics
806 #[serde(default)]
807 pub tags: Vec<String>,
808 /// Rule IDs that matched
809 #[serde(default)]
810 pub rule_ids: Vec<String>,
811 /// Confidence score (0.0 - 1.0)
812 pub confidence: Option<f32>,
813 /// Reason codes
814 #[serde(default)]
815 pub reason_codes: Vec<String>,
816 /// Custom metadata
817 #[serde(default)]
818 pub custom: HashMap<String, serde_json::Value>,
819}
820
821// ============================================================================
822// Guardrail Inspection Types
823// ============================================================================
824
825/// Type of guardrail inspection to perform
826#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
827#[serde(rename_all = "snake_case")]
828pub enum GuardrailInspectionType {
829 /// Prompt injection detection (analyze request content)
830 PromptInjection,
831 /// PII detection (analyze response content)
832 PiiDetection,
833}
834
835/// Guardrail inspection event
836///
837/// Sent to guardrail agents for semantic content analysis.
838/// Used for prompt injection detection on requests and PII detection on responses.
839#[derive(Debug, Clone, Serialize, Deserialize)]
840pub struct GuardrailInspectEvent {
841 /// Correlation ID for request tracing
842 pub correlation_id: String,
843 /// Type of inspection to perform
844 pub inspection_type: GuardrailInspectionType,
845 /// Content to inspect (request body or response content)
846 pub content: String,
847 /// Model name if available (for context)
848 #[serde(skip_serializing_if = "Option::is_none")]
849 pub model: Option<String>,
850 /// PII categories to check (for PII detection)
851 /// e.g., ["ssn", "credit_card", "email", "phone"]
852 #[serde(default)]
853 pub categories: Vec<String>,
854 /// Route ID
855 #[serde(skip_serializing_if = "Option::is_none")]
856 pub route_id: Option<String>,
857 /// Additional metadata for context
858 #[serde(default)]
859 pub metadata: HashMap<String, String>,
860}
861
862/// Guardrail inspection response from agent
863#[derive(Debug, Clone, Serialize, Deserialize)]
864pub struct GuardrailResponse {
865 /// Whether any issues were detected
866 pub detected: bool,
867 /// Confidence score (0.0 - 1.0)
868 #[serde(default)]
869 pub confidence: f64,
870 /// List of detections found
871 #[serde(default)]
872 pub detections: Vec<GuardrailDetection>,
873 /// Redacted content (for PII, if requested)
874 #[serde(skip_serializing_if = "Option::is_none")]
875 pub redacted_content: Option<String>,
876}
877
878impl Default for GuardrailResponse {
879 fn default() -> Self {
880 Self {
881 detected: false,
882 confidence: 0.0,
883 detections: Vec::new(),
884 redacted_content: None,
885 }
886 }
887}
888
889impl GuardrailResponse {
890 /// Create a response indicating nothing detected
891 pub fn clean() -> Self {
892 Self::default()
893 }
894
895 /// Create a response with a detection
896 pub fn with_detection(detection: GuardrailDetection) -> Self {
897 Self {
898 detected: true,
899 confidence: detection.confidence.unwrap_or(1.0),
900 detections: vec![detection],
901 redacted_content: None,
902 }
903 }
904
905 /// Add a detection to the response
906 pub fn add_detection(&mut self, detection: GuardrailDetection) {
907 self.detected = true;
908 if let Some(conf) = detection.confidence {
909 self.confidence = self.confidence.max(conf);
910 }
911 self.detections.push(detection);
912 }
913}
914
915/// A single guardrail detection (prompt injection attempt, PII instance, etc.)
916#[derive(Debug, Clone, Serialize, Deserialize)]
917pub struct GuardrailDetection {
918 /// Category of detection (e.g., "prompt_injection", "ssn", "credit_card")
919 pub category: String,
920 /// Human-readable description of what was detected
921 pub description: String,
922 /// Severity level
923 #[serde(default)]
924 pub severity: DetectionSeverity,
925 /// Confidence score for this detection (0.0 - 1.0)
926 #[serde(skip_serializing_if = "Option::is_none")]
927 pub confidence: Option<f64>,
928 /// Location in content where detection occurred
929 #[serde(skip_serializing_if = "Option::is_none")]
930 pub span: Option<TextSpan>,
931}
932
933impl GuardrailDetection {
934 /// Create a new detection
935 pub fn new(category: impl Into<String>, description: impl Into<String>) -> Self {
936 Self {
937 category: category.into(),
938 description: description.into(),
939 severity: DetectionSeverity::Medium,
940 confidence: None,
941 span: None,
942 }
943 }
944
945 /// Set severity
946 pub fn with_severity(mut self, severity: DetectionSeverity) -> Self {
947 self.severity = severity;
948 self
949 }
950
951 /// Set confidence
952 pub fn with_confidence(mut self, confidence: f64) -> Self {
953 self.confidence = Some(confidence);
954 self
955 }
956
957 /// Set span
958 pub fn with_span(mut self, start: usize, end: usize) -> Self {
959 self.span = Some(TextSpan { start, end });
960 self
961 }
962}
963
964/// Text span indicating location in content
965#[derive(Debug, Clone, Serialize, Deserialize)]
966pub struct TextSpan {
967 /// Start position (byte offset)
968 pub start: usize,
969 /// End position (byte offset, exclusive)
970 pub end: usize,
971}
972
973/// Severity level for guardrail detections
974#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
975#[serde(rename_all = "lowercase")]
976pub enum DetectionSeverity {
977 /// Low severity (informational)
978 Low,
979 /// Medium severity (default)
980 #[default]
981 Medium,
982 /// High severity (should likely block)
983 High,
984 /// Critical severity (must block)
985 Critical,
986}