arbiter-session 0.0.43

Task session management with budget and tool whitelisting for Arbiter
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
//! In-memory session store with TTL-based cleanup.

use chrono::Utc;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use uuid::Uuid;

use crate::error::SessionError;
use crate::model::{DataSensitivity, SessionId, SessionStatus, TaskSession};

/// Request to create a new task session.
pub struct CreateSessionRequest {
    /// The agent ID for this session.
    pub agent_id: Uuid,
    /// Delegation chain snapshot (serialized).
    pub delegation_chain_snapshot: Vec<String>,
    /// Declared intent for the session.
    pub declared_intent: String,
    /// Tools authorized by policy evaluation.
    pub authorized_tools: Vec<String>,
    /// Credential references this session may resolve.
    /// Empty means no credentials (deny-by-default for credential injection).
    #[allow(dead_code)]
    pub authorized_credentials: Vec<String>,
    /// Session time limit.
    pub time_limit: chrono::Duration,
    /// Maximum number of tool calls.
    pub call_budget: u64,
    /// Per-minute rate limit. `None` means no rate limit.
    pub rate_limit_per_minute: Option<u64>,
    /// Duration of the rate-limit window in seconds. Defaults to 60.
    pub rate_limit_window_secs: u64,
    /// Data sensitivity ceiling.
    pub data_sensitivity_ceiling: DataSensitivity,
}

/// In-memory session store with TTL-based cleanup.
#[derive(Clone)]
pub struct SessionStore {
    sessions: Arc<RwLock<HashMap<SessionId, TaskSession>>>,
}

impl SessionStore {
    /// Create a new empty session store.
    pub fn new() -> Self {
        Self {
            sessions: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Create a new task session and return it.
    pub async fn create(&self, req: CreateSessionRequest) -> TaskSession {
        // Enforce minimum time limit to prevent zero-duration sessions.
        let time_limit = if req.time_limit < chrono::Duration::seconds(1) {
            tracing::warn!(
                requested = ?req.time_limit,
                "session time_limit below minimum, clamping to 1 second"
            );
            chrono::Duration::seconds(1)
        } else {
            req.time_limit
        };
        let session = TaskSession {
            session_id: Uuid::new_v4(),
            agent_id: req.agent_id,
            delegation_chain_snapshot: req.delegation_chain_snapshot,
            declared_intent: req.declared_intent,
            authorized_tools: req.authorized_tools,
            authorized_credentials: req.authorized_credentials,
            time_limit,
            call_budget: req.call_budget,
            calls_made: 0,
            rate_limit_per_minute: req.rate_limit_per_minute,
            rate_window_start: Utc::now(),
            rate_window_calls: 0,
            rate_limit_window_secs: req.rate_limit_window_secs,
            data_sensitivity_ceiling: req.data_sensitivity_ceiling,
            created_at: Utc::now(),
            status: SessionStatus::Active,
        };

        tracing::info!(
            session_id = %session.session_id,
            agent_id = %session.agent_id,
            intent = %session.declared_intent,
            budget = session.call_budget,
            "created task session"
        );

        let mut sessions = self.sessions.write().await;
        sessions.insert(session.session_id, session.clone());
        session
    }

    /// Atomically check per-agent session cap and create if under the limit.
    /// Prevents the TOCTOU race where two concurrent requests both pass the
    /// count check before either creates a session.
    pub async fn create_if_under_cap(
        &self,
        req: CreateSessionRequest,
        max_sessions: u64,
    ) -> Result<TaskSession, SessionError> {
        let mut sessions = self.sessions.write().await;

        let active_count = sessions
            .values()
            .filter(|s| s.agent_id == req.agent_id && s.status == SessionStatus::Active)
            .count() as u64;

        if active_count >= max_sessions {
            return Err(SessionError::TooManySessions {
                agent_id: req.agent_id.to_string(),
                max: max_sessions,
                current: active_count,
            });
        }

        let session = TaskSession {
            session_id: Uuid::new_v4(),
            agent_id: req.agent_id,
            delegation_chain_snapshot: req.delegation_chain_snapshot,
            declared_intent: req.declared_intent,
            authorized_tools: req.authorized_tools,
            authorized_credentials: req.authorized_credentials,
            time_limit: req.time_limit,
            call_budget: req.call_budget,
            calls_made: 0,
            rate_limit_per_minute: req.rate_limit_per_minute,
            rate_window_start: Utc::now(),
            rate_window_calls: 0,
            rate_limit_window_secs: req.rate_limit_window_secs,
            data_sensitivity_ceiling: req.data_sensitivity_ceiling,
            created_at: Utc::now(),
            status: SessionStatus::Active,
        };

        sessions.insert(session.session_id, session.clone());
        Ok(session)
    }

    /// Record a tool call against the session, checking all constraints.
    ///
    /// Returns the updated session on success, or an error if:
    /// - Session not found
    /// - Session expired (would return 408)
    /// - Budget exceeded (would return 429)
    /// - Tool not authorized (would return 403)
    pub async fn use_session(
        &self,
        session_id: SessionId,
        tool_name: &str,
        requesting_agent_id: Option<Uuid>,
    ) -> Result<TaskSession, SessionError> {
        let mut sessions = self.sessions.write().await;
        let session = sessions
            .get_mut(&session_id)
            .ok_or(SessionError::NotFound(session_id))?;

        // Verify agent binding to prevent session fixation.
        if let Some(agent_id) = requesting_agent_id
            && agent_id != session.agent_id
        {
            return Err(SessionError::AgentMismatch {
                session_id,
                expected: session.agent_id,
                actual: agent_id,
            });
        }

        if session.status == SessionStatus::Closed {
            return Err(SessionError::AlreadyClosed(session_id));
        }

        // Check expiry.
        if session.is_expired() {
            session.status = SessionStatus::Expired;
            return Err(SessionError::Expired(session_id));
        }

        // Check budget.
        if session.is_budget_exceeded() {
            return Err(SessionError::BudgetExceeded {
                session_id,
                limit: session.call_budget,
                used: session.calls_made,
            });
        }

        // Check tool authorization.
        if !session.is_tool_authorized(tool_name) {
            return Err(SessionError::ToolNotAuthorized {
                session_id,
                tool: tool_name.into(),
            });
        }

        // Check rate limit.
        if session.check_rate_limit() {
            return Err(SessionError::RateLimited {
                session_id,
                limit_per_minute: session.rate_limit_per_minute.unwrap_or(0),
            });
        }

        // All checks passed. Increment counter.
        session.calls_made += 1;

        tracing::debug!(
            session_id = %session_id,
            tool = tool_name,
            calls = session.calls_made,
            budget = session.call_budget,
            "session tool call recorded"
        );

        Ok(session.clone())
    }

    /// Atomically validate and record a batch of tool calls against the session.
    ///
    /// Acquires the write lock once, validates ALL tools against
    /// the whitelist and budget, and only increments `calls_made` by the full
    /// batch count if every tool passes. If any tool fails validation, no
    /// budget is consumed for any of them.
    pub async fn use_session_batch(
        &self,
        session_id: SessionId,
        tool_names: &[&str],
        requesting_agent_id: Option<Uuid>,
    ) -> Result<TaskSession, SessionError> {
        let mut sessions = self.sessions.write().await;
        let session = sessions
            .get_mut(&session_id)
            .ok_or(SessionError::NotFound(session_id))?;

        // Verify agent binding to prevent session fixation.
        if let Some(agent_id) = requesting_agent_id
            && agent_id != session.agent_id
        {
            return Err(SessionError::AgentMismatch {
                session_id,
                expected: session.agent_id,
                actual: agent_id,
            });
        }

        if session.status == SessionStatus::Closed {
            return Err(SessionError::AlreadyClosed(session_id));
        }

        // Check expiry.
        if session.is_expired() {
            session.status = SessionStatus::Expired;
            return Err(SessionError::Expired(session_id));
        }

        let batch_size = tool_names.len() as u64;

        // Check budget for the entire batch.
        if session.calls_made + batch_size > session.call_budget {
            return Err(SessionError::BudgetExceeded {
                session_id,
                limit: session.call_budget,
                used: session.calls_made,
            });
        }

        // Check tool authorization for every tool before consuming any budget.
        for tool_name in tool_names {
            if !session.is_tool_authorized(tool_name) {
                return Err(SessionError::ToolNotAuthorized {
                    session_id,
                    tool: (*tool_name).into(),
                });
            }
        }

        // Check rate limit for the entire batch.
        // We check whether adding batch_size calls would exceed the limit,
        // without mutating state until we know it's safe.
        if let Some(limit) = session.rate_limit_per_minute {
            let now = chrono::Utc::now();
            let elapsed = now - session.rate_window_start;
            if elapsed >= chrono::Duration::seconds(session.rate_limit_window_secs as i64) {
                // New window; will be reset below after all checks pass.
            } else if session.rate_window_calls + batch_size > limit {
                return Err(SessionError::RateLimited {
                    session_id,
                    limit_per_minute: limit,
                });
            }
        }

        // All checks passed. Atomically increment counters.
        // Update rate limit window.
        if let Some(_limit) = session.rate_limit_per_minute {
            let now = chrono::Utc::now();
            let elapsed = now - session.rate_window_start;
            if elapsed >= chrono::Duration::seconds(session.rate_limit_window_secs as i64) {
                session.rate_window_start = now;
                session.rate_window_calls = batch_size;
            } else {
                session.rate_window_calls += batch_size;
            }
        }

        session.calls_made += batch_size;

        tracing::debug!(
            session_id = %session_id,
            batch_size = batch_size,
            calls = session.calls_made,
            budget = session.call_budget,
            "session batch tool calls recorded"
        );

        Ok(session.clone())
    }

    /// Close a session, preventing further use.
    pub async fn close(&self, session_id: SessionId) -> Result<TaskSession, SessionError> {
        let mut sessions = self.sessions.write().await;
        let session = sessions
            .get_mut(&session_id)
            .ok_or(SessionError::NotFound(session_id))?;

        if session.status == SessionStatus::Closed {
            return Err(SessionError::AlreadyClosed(session_id));
        }

        session.status = SessionStatus::Closed;
        tracing::info!(session_id = %session_id, "session closed");
        Ok(session.clone())
    }

    /// Get a session by ID without modifying it.
    pub async fn get(&self, session_id: SessionId) -> Result<TaskSession, SessionError> {
        let sessions = self.sessions.read().await;
        sessions
            .get(&session_id)
            .cloned()
            .ok_or(SessionError::NotFound(session_id))
    }

    /// List all sessions currently in the store (active, expired, and closed).
    pub async fn list_all(&self) -> Vec<TaskSession> {
        let sessions = self.sessions.read().await;
        sessions.values().cloned().collect()
    }

    /// List only sessions belonging to a specific agent.
    /// Use this instead of list_all() when agent-scoped access is needed
    /// to prevent cross-agent session data exposure.
    pub async fn list_for_agent(&self, agent_id: Uuid) -> Vec<TaskSession> {
        let sessions = self.sessions.read().await;
        sessions
            .values()
            .filter(|s| s.agent_id == agent_id)
            .cloned()
            .collect()
    }

    /// Count the number of active sessions for a given agent.
    ///
    /// P0: Used to enforce per-agent concurrent session caps.
    pub async fn count_active_for_agent(&self, agent_id: uuid::Uuid) -> u64 {
        let sessions = self.sessions.read().await;
        sessions
            .values()
            .filter(|s| s.agent_id == agent_id && s.status == SessionStatus::Active)
            .count() as u64
    }

    /// Close all active sessions belonging to a specific agent.
    ///
    /// When an agent is deactivated via cascade_deactivate,
    /// all its sessions must be immediately closed.
    pub async fn close_sessions_for_agent(&self, agent_id: uuid::Uuid) -> usize {
        let mut sessions = self.sessions.write().await;
        let mut closed = 0usize;
        for session in sessions.values_mut() {
            if session.agent_id == agent_id && session.status == SessionStatus::Active {
                session.status = SessionStatus::Closed;
                closed += 1;
                tracing::info!(
                    session_id = %session.session_id,
                    agent_id = %agent_id,
                    "closed session due to agent deactivation"
                );
            }
        }
        closed
    }

    /// Remove expired sessions from the store. Returns the number removed.
    pub async fn cleanup_expired(&self) -> usize {
        let mut sessions = self.sessions.write().await;
        let before = sessions.len();
        // Also clean up closed sessions, not just expired ones.
        // Previously, closed sessions accumulated indefinitely, growing the store without bound.
        sessions.retain(|_, s| {
            if s.is_expired() {
                tracing::debug!(session_id = %s.session_id, "cleaning up expired session");
                false
            } else if s.status == SessionStatus::Closed {
                tracing::debug!(session_id = %s.session_id, "cleaning up closed session");
                false
            } else {
                true
            }
        });
        let removed = before - sessions.len();
        if removed > 0 {
            tracing::info!(removed, "cleaned up expired/closed sessions");
        }
        removed
    }
}

impl Default for SessionStore {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn test_create_request() -> CreateSessionRequest {
        CreateSessionRequest {
            agent_id: Uuid::new_v4(),
            delegation_chain_snapshot: vec![],
            declared_intent: "read and analyze files".into(),
            authorized_tools: vec!["read_file".into(), "list_dir".into()],
            authorized_credentials: vec![],
            time_limit: chrono::Duration::hours(1),
            call_budget: 5,
            rate_limit_per_minute: None,
            rate_limit_window_secs: 60,
            data_sensitivity_ceiling: DataSensitivity::Internal,
        }
    }

    #[tokio::test]
    async fn create_and_use_session() {
        let store = SessionStore::new();
        let session = store.create(test_create_request()).await;

        assert_eq!(session.calls_made, 0);
        assert!(session.is_active());

        let updated = store
            .use_session(session.session_id, "read_file", None)
            .await
            .unwrap();
        assert_eq!(updated.calls_made, 1);
    }

    #[tokio::test]
    async fn budget_enforcement() {
        let store = SessionStore::new();
        let mut req = test_create_request();
        req.call_budget = 2;
        let session = store.create(req).await;

        // Use up the budget.
        store
            .use_session(session.session_id, "read_file", None)
            .await
            .unwrap();
        store
            .use_session(session.session_id, "read_file", None)
            .await
            .unwrap();

        // Third call should fail.
        let result = store
            .use_session(session.session_id, "read_file", None)
            .await;
        assert!(matches!(result, Err(SessionError::BudgetExceeded { .. })));
    }

    #[tokio::test]
    async fn tool_whitelist_enforcement() {
        let store = SessionStore::new();
        let session = store.create(test_create_request()).await;

        // Authorized tool works.
        store
            .use_session(session.session_id, "read_file", None)
            .await
            .unwrap();

        // Unauthorized tool is rejected.
        let result = store
            .use_session(session.session_id, "delete_file", None)
            .await;
        assert!(matches!(
            result,
            Err(SessionError::ToolNotAuthorized { .. })
        ));
    }

    #[tokio::test]
    async fn session_expiry() {
        let store = SessionStore::new();
        let mut req = test_create_request();
        // Set a 1-second time limit (minimum enforced by create()).
        // Previously used zero, but minimum is now clamped to 1s.
        req.time_limit = chrono::Duration::seconds(1);
        let session = store.create(req).await;

        // Wait for the session to expire.
        tokio::time::sleep(std::time::Duration::from_millis(1100)).await;

        let result = store
            .use_session(session.session_id, "read_file", None)
            .await;
        assert!(matches!(result, Err(SessionError::Expired(_))));
    }

    #[tokio::test]
    async fn close_and_reuse() {
        let store = SessionStore::new();
        let session = store.create(test_create_request()).await;

        store.close(session.session_id).await.unwrap();

        let result = store
            .use_session(session.session_id, "read_file", None)
            .await;
        assert!(matches!(result, Err(SessionError::AlreadyClosed(_))));
    }

    #[tokio::test]
    async fn cleanup_expired_sessions() {
        let store = SessionStore::new();

        // Create a short-lived session (1s minimum).
        let mut req = test_create_request();
        req.time_limit = chrono::Duration::seconds(1);
        store.create(req).await;

        // Create a valid session with longer limit.
        let valid_req = test_create_request();
        store.create(valid_req).await;

        // Wait for the short session to expire.
        tokio::time::sleep(std::time::Duration::from_millis(1100)).await;

        let removed = store.cleanup_expired().await;
        assert_eq!(removed, 1);
    }

    #[tokio::test]
    async fn session_not_found() {
        let store = SessionStore::new();
        let fake_id = Uuid::new_v4();
        let result = store.use_session(fake_id, "anything", None).await;
        assert!(matches!(result, Err(SessionError::NotFound(_))));
    }

    #[tokio::test]
    async fn rate_limit_enforcement() {
        let store = SessionStore::new();
        let mut req = test_create_request();
        req.rate_limit_per_minute = Some(3);
        req.call_budget = 100; // high budget, rate limit should trigger first
        let session = store.create(req).await;

        // First 3 calls succeed (within rate limit).
        store
            .use_session(session.session_id, "read_file", None)
            .await
            .unwrap();
        store
            .use_session(session.session_id, "read_file", None)
            .await
            .unwrap();
        store
            .use_session(session.session_id, "read_file", None)
            .await
            .unwrap();

        // 4th call hits rate limit.
        let result = store
            .use_session(session.session_id, "read_file", None)
            .await;
        assert!(
            matches!(result, Err(SessionError::RateLimited { .. })),
            "expected RateLimited, got {result:?}"
        );
    }

    #[tokio::test]
    async fn no_rate_limit_when_unset() {
        let store = SessionStore::new();
        let mut req = test_create_request();
        req.rate_limit_per_minute = None;
        req.call_budget = 100;
        let session = store.create(req).await;

        // All calls succeed without rate limiting.
        for _ in 0..10 {
            store
                .use_session(session.session_id, "read_file", None)
                .await
                .unwrap();
        }
    }

    /// batch with one unauthorized tool must consume zero budget.
    #[tokio::test]
    async fn batch_validation_atomicity() {
        let store = SessionStore::new();
        let mut req = test_create_request();
        req.call_budget = 10;
        req.authorized_tools = vec!["read_file".into(), "list_dir".into()];
        let session = store.create(req).await;

        // Batch contains one unauthorized tool ("delete_file").
        let result = store
            .use_session_batch(session.session_id, &["read_file", "delete_file"], None)
            .await;
        assert!(
            matches!(result, Err(SessionError::ToolNotAuthorized { .. })),
            "expected ToolNotAuthorized, got {result:?}"
        );

        // Budget must remain untouched.
        let s = store.get(session.session_id).await.unwrap();
        assert_eq!(
            s.calls_made, 0,
            "no budget should be consumed on batch failure"
        );
    }

    #[tokio::test]
    async fn batch_budget_enforcement() {
        let store = SessionStore::new();
        let mut req = test_create_request();
        req.call_budget = 3;
        req.authorized_tools = vec!["read_file".into()];
        let session = store.create(req).await;

        // Batch of 4 exceeds budget of 3.
        let result = store
            .use_session_batch(
                session.session_id,
                &["read_file", "read_file", "read_file", "read_file"],
                None,
            )
            .await;
        assert!(
            matches!(result, Err(SessionError::BudgetExceeded { .. })),
            "expected BudgetExceeded, got {result:?}"
        );

        // Budget must remain at 0.
        let s = store.get(session.session_id).await.unwrap();
        assert_eq!(
            s.calls_made, 0,
            "no budget should be consumed on batch failure"
        );
    }

    #[tokio::test]
    async fn batch_rate_limit_enforcement() {
        let store = SessionStore::new();
        let mut req = test_create_request();
        req.call_budget = 100;
        req.rate_limit_per_minute = Some(3);
        req.authorized_tools = vec!["read_file".into()];
        let session = store.create(req).await;

        // Batch of 4 exceeds rate limit of 3.
        let result = store
            .use_session_batch(
                session.session_id,
                &["read_file", "read_file", "read_file", "read_file"],
                None,
            )
            .await;
        assert!(
            matches!(result, Err(SessionError::RateLimited { .. })),
            "expected RateLimited, got {result:?}"
        );
    }

    #[tokio::test]
    async fn empty_batch_succeeds() {
        let store = SessionStore::new();
        let session = store.create(test_create_request()).await;

        // Empty batch should succeed without consuming budget.
        let result = store
            .use_session_batch(session.session_id, &[], None)
            .await
            .unwrap();
        assert_eq!(result.calls_made, 0, "empty batch must not consume budget");
    }

    /// cleanup should also remove closed sessions.
    #[tokio::test]
    async fn cleanup_also_removes_closed() {
        let store = SessionStore::new();
        let session = store.create(test_create_request()).await;

        // Close it.
        store.close(session.session_id).await.unwrap();

        // Cleanup should remove the closed session.
        let removed = store.cleanup_expired().await;
        assert_eq!(removed, 1, "closed session should be cleaned up");

        // It should be gone.
        let result = store.get(session.session_id).await;
        assert!(
            matches!(result, Err(SessionError::NotFound(_))),
            "closed session should be removed after cleanup"
        );
    }

    /// A session created with call_budget=0 should immediately fail on use.
    #[tokio::test]
    async fn zero_budget_session() {
        let store = SessionStore::new();
        let mut req = test_create_request();
        req.call_budget = 0;
        let session = store.create(req).await;

        let result = store
            .use_session(session.session_id, "read_file", None)
            .await;
        assert!(
            matches!(result, Err(SessionError::BudgetExceeded { .. })),
            "zero-budget session must reject the first call, got {result:?}"
        );
    }

    /// Agent deactivation must close all agent sessions.
    #[tokio::test]
    async fn deactivation_closes_agent_sessions() {
        let store = SessionStore::new();
        let agent_id = Uuid::new_v4();
        let other_agent = Uuid::new_v4();

        for _ in 0..3 {
            let mut req = test_create_request();
            req.agent_id = agent_id;
            store.create(req).await;
        }
        let mut other_req = test_create_request();
        other_req.agent_id = other_agent;
        let other_session = store.create(other_req).await;

        let closed = store.close_sessions_for_agent(agent_id).await;
        assert_eq!(closed, 3);

        let all = store.list_all().await;
        for s in &all {
            if s.agent_id == agent_id {
                assert_eq!(s.status, SessionStatus::Closed);
            }
        }
        let other = store.get(other_session.session_id).await.unwrap();
        assert_eq!(other.status, SessionStatus::Active);
    }

    /// Concurrent budget enforcement.
    /// Spawn 10 tasks each calling use_session once on a session with budget=5.
    /// Exactly 5 must succeed and 5 must fail with BudgetExceeded.
    #[tokio::test]
    async fn concurrent_budget_enforcement() {
        let store = SessionStore::new();
        let mut req = test_create_request();
        req.call_budget = 5;
        req.authorized_tools = vec!["read_file".into()];
        let session = store.create(req).await;

        let successes = Arc::new(std::sync::atomic::AtomicU64::new(0));
        let failures = Arc::new(std::sync::atomic::AtomicU64::new(0));

        let mut handles = Vec::new();
        for _ in 0..10 {
            let store = store.clone();
            let sid = session.session_id;
            let s = successes.clone();
            let f = failures.clone();
            handles.push(tokio::spawn(async move {
                match store.use_session(sid, "read_file", None).await {
                    Ok(_) => {
                        s.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                    }
                    Err(SessionError::BudgetExceeded { .. }) => {
                        f.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                    }
                    Err(e) => panic!("unexpected error: {e:?}"),
                }
            }));
        }

        for h in handles {
            h.await.unwrap();
        }

        assert_eq!(
            successes.load(std::sync::atomic::Ordering::Relaxed),
            5,
            "exactly 5 calls should succeed"
        );
        assert_eq!(
            failures.load(std::sync::atomic::Ordering::Relaxed),
            5,
            "exactly 5 calls should fail with BudgetExceeded"
        );
    }

    /// Session fixation prevention: a different agent must not be able to use
    /// another agent's session by presenting its session ID.
    #[tokio::test]
    async fn agent_mismatch_rejected() {
        let store = SessionStore::new();
        let session = store.create(test_create_request()).await;
        let attacker_id = Uuid::new_v4();

        // Attacker presents a different agent_id than the session owner.
        let result = store
            .use_session(session.session_id, "read_file", Some(attacker_id))
            .await;
        assert!(
            matches!(result, Err(SessionError::AgentMismatch { .. })),
            "different agent must be rejected, got {result:?}"
        );

        // Legitimate agent succeeds.
        let result = store
            .use_session(session.session_id, "read_file", Some(session.agent_id))
            .await;
        assert!(result.is_ok(), "session owner should succeed");
    }

    /// Batch variant of agent mismatch check.
    #[tokio::test]
    async fn batch_agent_mismatch_rejected() {
        let store = SessionStore::new();
        let session = store.create(test_create_request()).await;
        let attacker_id = Uuid::new_v4();

        let result = store
            .use_session_batch(session.session_id, &["read_file"], Some(attacker_id))
            .await;
        assert!(
            matches!(result, Err(SessionError::AgentMismatch { .. })),
            "batch with wrong agent must be rejected, got {result:?}"
        );
    }
}