Skip to main content

arbiter_session/
middleware.rs

1//! Session middleware utilities for proxy integration.
2//!
3//! Provides session ID parsing from the `X-Arbiter-Session` header and
4//! error-to-status-code mapping. Session constraint validation is performed
5//! directly by the handler stages in `crates/arbiter/src/stages/`.
6
7use uuid::Uuid;
8
9use crate::error::SessionError;
10
11/// HTTP status code that should be returned for each session error type.
12pub fn status_code_for_error(err: &SessionError) -> u16 {
13    match err {
14        SessionError::NotFound(_) => 404,
15        SessionError::Expired(_) => 408,
16        SessionError::BudgetExceeded { .. } => 429,
17        SessionError::ToolNotAuthorized { .. } => 403,
18        SessionError::AlreadyClosed(_) => 410,
19        SessionError::RateLimited { .. } => 429,
20        SessionError::TooManySessions { .. } => 429,
21        SessionError::AgentMismatch { .. } => 403,
22        SessionError::StorageWriteThrough { .. } => 503,
23    }
24}
25
26/// Extract a session ID from the `X-Arbiter-Session` header value.
27///
28/// Returns `None` if the header is missing or not a valid UUID.
29pub fn parse_session_header(header_value: &str) -> Option<Uuid> {
30    Uuid::parse_str(header_value.trim()).ok()
31}
32
33#[cfg(test)]
34mod tests {
35    use super::*;
36
37    #[test]
38    fn parse_session_header_valid() {
39        let id = Uuid::new_v4();
40        assert_eq!(parse_session_header(&id.to_string()), Some(id));
41    }
42
43    #[test]
44    fn parse_session_header_invalid() {
45        assert_eq!(parse_session_header("not-a-uuid"), None);
46        assert_eq!(parse_session_header(""), None);
47    }
48
49    /// Test that each SessionError variant maps to the correct HTTP status code.
50    #[test]
51    fn status_code_for_all_error_variants() {
52        let id = Uuid::new_v4();
53
54        assert_eq!(status_code_for_error(&SessionError::NotFound(id)), 404);
55        assert_eq!(status_code_for_error(&SessionError::Expired(id)), 408);
56        assert_eq!(
57            status_code_for_error(&SessionError::BudgetExceeded {
58                session_id: id,
59                limit: 10,
60                used: 10,
61            }),
62            429
63        );
64        assert_eq!(
65            status_code_for_error(&SessionError::ToolNotAuthorized {
66                session_id: id,
67                tool: "delete".into(),
68            }),
69            403
70        );
71        assert_eq!(status_code_for_error(&SessionError::AlreadyClosed(id)), 410);
72        assert_eq!(
73            status_code_for_error(&SessionError::RateLimited {
74                session_id: id,
75                limit_per_minute: 5,
76            }),
77            429
78        );
79        assert_eq!(
80            status_code_for_error(&SessionError::TooManySessions {
81                agent_id: id.to_string(),
82                max: 10,
83                current: 10,
84            }),
85            429
86        );
87    }
88
89    /// Whitespace-padded UUIDs should be accepted (the trim() in parse_session_header).
90    #[test]
91    fn parse_session_header_with_whitespace() {
92        let id = Uuid::new_v4();
93        let padded = format!("  {}  ", id);
94        assert_eq!(
95            parse_session_header(&padded),
96            Some(id),
97            "whitespace-padded UUID should be parsed correctly"
98        );
99    }
100}