From fe0225f33d67db43fd7e13644cf01071cd5d43bf Mon Sep 17 00:00:00 2001
From: ghzhost <patricknasci@ghzhost.com>
Date: Thu, 13 Aug 2026 21:53:57 +0000
Subject: [PATCH] Add session budget None regression coverage

---
 agentshield/eval_gym.py | 22 ++++++++++++++++------
 tests/eval_gym.py       | 22 ++++++++++++++++------
 2 files changed, 32 insertions(+), 12 deletions(-)

diff --git a/agentshield/eval_gym.py b/agentshield/eval_gym.py
index 31473cf..0a36624 100644
--- a/agentshield/eval_gym.py
+++ b/agentshield/eval_gym.py
@@ -405,25 +405,35 @@ def _prior(agent_id, amount, timestamp):
      "expected": "APPROVED",
      "description": "Prior transaction in different session, not counted"},
 
+    {"id": 54, "category": "session_budget",
+     "transaction": {**_txn("t054", amount=100.00), "session_id": None},
+     "rules": [{"id": "sb4", "type": "session_budget", "priority": 1,
+                "params": {"max_session": 500}, "action": "BLOCK"}],
+     "prior_transactions": [
+         {**_txn("t054a", amount=450.00), "session_id": None},
+     ],
+     "expected": "BLOCKED",
+     "description": "Missing session_id cannot bypass session budget evaluation"},
+
     # ─── Cascade Cost (inspired by HeartFlow / @yun520-1) ───
-    {"id": 54, "category": "cascade_cost",
-     "transaction": {**_txn("t054", amount=50.00), "fail_probability": 0.3, "reversal_cost": 200},
+    {"id": 55, "category": "cascade_cost",
+     "transaction": {**_txn("t055", amount=50.00), "fail_probability": 0.3, "reversal_cost": 200},
      "rules": [{"id": "cc1", "type": "cascade_cost", "priority": 1,
                 "params": {"max_cascade_cost": 100}, "action": "BLOCK"}],
      "prior_transactions": [],
      "expected": "BLOCKED",
      "description": "Cascade cost $110 ($50 + 30% × $200) exceeds $100 limit"},
 
-    {"id": 55, "category": "cascade_cost",
-     "transaction": {**_txn("t055", amount=10.00), "fail_probability": 0.1, "reversal_cost": 50},
+    {"id": 56, "category": "cascade_cost",
+     "transaction": {**_txn("t056", amount=10.00), "fail_probability": 0.1, "reversal_cost": 50},
      "rules": [{"id": "cc2", "type": "cascade_cost", "priority": 1,
                 "params": {"max_cascade_cost": 100}, "action": "BLOCK"}],
      "prior_transactions": [],
      "expected": "APPROVED",
      "description": "Cascade cost $15 ($10 + 10% × $50) under $100 limit"},
 
-    {"id": 56, "category": "cascade_cost",
-     "transaction": {**_txn("t056", amount=10.00), "estimated_cascade_cost": 150},
+    {"id": 57, "category": "cascade_cost",
+     "transaction": {**_txn("t057", amount=10.00), "estimated_cascade_cost": 150},
      "rules": [{"id": "cc3", "type": "cascade_cost", "priority": 1,
                 "params": {"max_cascade_cost": 100}, "action": "BLOCK"}],
      "prior_transactions": [],
diff --git a/tests/eval_gym.py b/tests/eval_gym.py
index 7854e7d..dfc179c 100644
--- a/tests/eval_gym.py
+++ b/tests/eval_gym.py
@@ -449,24 +449,34 @@ def _prior(agent_id, amount, timestamp):
      "expected": "BLOCKED",
      "description": "Negative prior session amount ($-95) should not reduce session total, $190 > $100 → BLOCKED"},
 
-    {"id": 59, "category": "transaction_limit_block",
-     "transaction": _txn("t059", amount=-1000000.00),
+    {"id": 59, "category": "session_budget",
+     "transaction": {**_txn("t059", amount=100.00), "session_id": None},
+     "rules": [{"id": "sb4", "type": "session_budget", "priority": 1,
+                "params": {"max_session": 500}, "action": "BLOCK"}],
+     "prior_transactions": [
+         {**_txn("t059a", amount=450.00), "session_id": None},
+     ],
+     "expected": "BLOCKED",
+     "description": "Missing session_id cannot bypass session budget evaluation"},
+
+    {"id": 60, "category": "transaction_limit_block",
+     "transaction": _txn("t060", amount=-1000000.00),
      "rules": [{"id": "r1", "type": "transaction_limit", "priority": 1,
                 "params": {"max_amount": 500}, "action": "BLOCK"}],
      "prior_transactions": [],
      "expected": "BLOCKED",
      "description": "Negative amount should be rejected by transaction_limit"},
 
-    {"id": 60, "category": "transaction_limit_block",
-     "transaction": _txn("t060", amount=0.00),
+    {"id": 61, "category": "transaction_limit_block",
+     "transaction": _txn("t061", amount=0.00),
      "rules": [{"id": "r1", "type": "transaction_limit", "priority": 1,
                 "params": {"max_amount": 500}, "action": "BLOCK"}],
      "prior_transactions": [],
      "expected": "BLOCKED",
      "description": "Zero amount transaction should be rejected"},
 
-    {"id": 61, "category": "cascade_cost",
-     "transaction": {**_txn("t061", amount=50.00), "estimated_cascade_cost": -10},
+    {"id": 62, "category": "cascade_cost",
+     "transaction": {**_txn("t062", amount=50.00), "estimated_cascade_cost": -10},
      "rules": [{"id": "cc4", "type": "cascade_cost", "priority": 1,
                 "params": {"max_cascade_cost": 100}, "action": "BLOCK"}],
      "prior_transactions": [],