tollgate 1.0.3__tar.gz → 1.0.4__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (60) hide show
  1. {tollgate-1.0.3 → tollgate-1.0.4}/CHANGELOG.md +8 -0
  2. {tollgate-1.0.3 → tollgate-1.0.4}/PKG-INFO +25 -1
  3. {tollgate-1.0.3 → tollgate-1.0.4}/README.md +24 -0
  4. {tollgate-1.0.3 → tollgate-1.0.4}/pyproject.toml +1 -1
  5. {tollgate-1.0.3 → tollgate-1.0.4}/src/tollgate/__init__.py +6 -2
  6. tollgate-1.0.4/src/tollgate/grants.py +103 -0
  7. {tollgate-1.0.3 → tollgate-1.0.4}/src/tollgate/tower.py +58 -37
  8. {tollgate-1.0.3 → tollgate-1.0.4}/src/tollgate/types.py +25 -0
  9. tollgate-1.0.4/tests/test_grants.py +286 -0
  10. {tollgate-1.0.3 → tollgate-1.0.4}/.claude/settings.local.json +0 -0
  11. {tollgate-1.0.3 → tollgate-1.0.4}/.gitignore +0 -0
  12. {tollgate-1.0.3 → tollgate-1.0.4}/COMPARISON.md +0 -0
  13. {tollgate-1.0.3 → tollgate-1.0.4}/CONTRIBUTING.md +0 -0
  14. {tollgate-1.0.3 → tollgate-1.0.4}/LICENSE +0 -0
  15. {tollgate-1.0.3 → tollgate-1.0.4}/Makefile +0 -0
  16. {tollgate-1.0.3 → tollgate-1.0.4}/QUICKSTART.md +0 -0
  17. {tollgate-1.0.3 → tollgate-1.0.4}/SECURITY.md +0 -0
  18. {tollgate-1.0.3 → tollgate-1.0.4}/examples/mcp_minimal/audit.jsonl +0 -0
  19. {tollgate-1.0.3 → tollgate-1.0.4}/examples/mcp_minimal/demo.py +0 -0
  20. {tollgate-1.0.3 → tollgate-1.0.4}/examples/mcp_minimal/manifest.yaml +0 -0
  21. {tollgate-1.0.3 → tollgate-1.0.4}/examples/mcp_minimal/policy.yaml +0 -0
  22. {tollgate-1.0.3 → tollgate-1.0.4}/examples/mock_tickets/README.md +0 -0
  23. {tollgate-1.0.3 → tollgate-1.0.4}/examples/mock_tickets/agent.py +0 -0
  24. {tollgate-1.0.3 → tollgate-1.0.4}/examples/mock_tickets/demo.py +0 -0
  25. {tollgate-1.0.3 → tollgate-1.0.4}/examples/mock_tickets/manifest.yaml +0 -0
  26. {tollgate-1.0.3 → tollgate-1.0.4}/examples/mock_tickets/tickets.json +0 -0
  27. {tollgate-1.0.3 → tollgate-1.0.4}/examples/mock_tickets/tools.py +0 -0
  28. {tollgate-1.0.3 → tollgate-1.0.4}/examples/strands_minimal/audit.jsonl +0 -0
  29. {tollgate-1.0.3 → tollgate-1.0.4}/examples/strands_minimal/demo.py +0 -0
  30. {tollgate-1.0.3 → tollgate-1.0.4}/examples/strands_minimal/manifest.yaml +0 -0
  31. {tollgate-1.0.3 → tollgate-1.0.4}/examples/strands_minimal/policy.yaml +0 -0
  32. {tollgate-1.0.3 → tollgate-1.0.4}/policies/default.yaml +0 -0
  33. {tollgate-1.0.3 → tollgate-1.0.4}/specs/audit_event.schema.json +0 -0
  34. {tollgate-1.0.3 → tollgate-1.0.4}/specs/decision.schema.json +0 -0
  35. {tollgate-1.0.3 → tollgate-1.0.4}/specs/identity.schema.json +0 -0
  36. {tollgate-1.0.3 → tollgate-1.0.4}/specs/intent.schema.json +0 -0
  37. {tollgate-1.0.3 → tollgate-1.0.4}/specs/tool_request.schema.json +0 -0
  38. {tollgate-1.0.3 → tollgate-1.0.4}/src/tollgate/approvals.py +0 -0
  39. {tollgate-1.0.3 → tollgate-1.0.4}/src/tollgate/audit.py +0 -0
  40. {tollgate-1.0.3 → tollgate-1.0.4}/src/tollgate/exceptions.py +0 -0
  41. {tollgate-1.0.3 → tollgate-1.0.4}/src/tollgate/helpers.py +0 -0
  42. {tollgate-1.0.3 → tollgate-1.0.4}/src/tollgate/integrations/__init__.py +0 -0
  43. {tollgate-1.0.3 → tollgate-1.0.4}/src/tollgate/integrations/mcp.py +0 -0
  44. {tollgate-1.0.3 → tollgate-1.0.4}/src/tollgate/integrations/strands.py +0 -0
  45. {tollgate-1.0.3 → tollgate-1.0.4}/src/tollgate/interceptors/__init__.py +0 -0
  46. {tollgate-1.0.3 → tollgate-1.0.4}/src/tollgate/interceptors/base.py +0 -0
  47. {tollgate-1.0.3 → tollgate-1.0.4}/src/tollgate/interceptors/langchain.py +0 -0
  48. {tollgate-1.0.3 → tollgate-1.0.4}/src/tollgate/interceptors/openai.py +0 -0
  49. {tollgate-1.0.3 → tollgate-1.0.4}/src/tollgate/policy.py +0 -0
  50. {tollgate-1.0.3 → tollgate-1.0.4}/src/tollgate/registry.py +0 -0
  51. {tollgate-1.0.3 → tollgate-1.0.4}/tests/test_adapters_v1.py +0 -0
  52. {tollgate-1.0.3 → tollgate-1.0.4}/tests/test_audit_integrity_v1.py +0 -0
  53. {tollgate-1.0.3 → tollgate-1.0.4}/tests/test_deferred_v1.py +0 -0
  54. {tollgate-1.0.3 → tollgate-1.0.4}/tests/test_helpers_v1.py +0 -0
  55. {tollgate-1.0.3 → tollgate-1.0.4}/tests/test_integrations_v1.py +0 -0
  56. {tollgate-1.0.3 → tollgate-1.0.4}/tests/test_policy_v1.py +0 -0
  57. {tollgate-1.0.3 → tollgate-1.0.4}/tests/test_registry_v1.py +0 -0
  58. {tollgate-1.0.3 → tollgate-1.0.4}/tests/test_security_v1.py +0 -0
  59. {tollgate-1.0.3 → tollgate-1.0.4}/tests/test_tower_v1.py +0 -0
  60. {tollgate-1.0.3 → tollgate-1.0.4}/tests/test_v1_integrations.py +0 -0
@@ -5,6 +5,14 @@ All notable changes to this project will be documented in this file.
5
5
  The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
6
6
  and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
7
7
 
8
+ ## [1.1.0] - 2026-01-28
9
+ ### Added
10
+ - **Session Grants**: Introduced `Grant` and `InMemoryGrantStore` to allow bypassing human approval for specific, pre-authorized actions.
11
+ - **Grant ID Auto-generation**: Grant IDs are now automatically generated if not provided.
12
+ - **Grant Usage Tracking**: Added usage counters to `InMemoryGrantStore`.
13
+ - **Audit Integration**: Audit events now include `grant_id` when a grant is used.
14
+ - **Security**: Re-added and enforced exception sanitization in audit logs.
15
+
8
16
  ## [1.0.0] - 2026-01-27
9
17
  ### Added
10
18
  - **Interception-First Architecture**: Added `TollgateInterceptor` and framework adapters for LangChain and OpenAI.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: tollgate
3
- Version: 1.0.3
3
+ Version: 1.0.4
4
4
  Summary: Runtime enforcement layer for AI agent tool calls using Identity + Intent + Policy
5
5
  Author: Tollgate Maintainers
6
6
  License-Expression: Apache-2.0
@@ -77,6 +77,30 @@ Runtime enforcement layer for AI agent tool calls using **Identity + Intent + Po
77
77
 
78
78
  ## 🚀 v1 Integrations
79
79
 
80
+ ### 🎟️ Session Grants
81
+ Grants allow you to pre-authorize specific actions for an agent session, bypassing human-in-the-loop approvals for repetitive or low-risk tasks.
82
+
83
+ ```python
84
+ from tollgate import Grant, InMemoryGrantStore, Effect
85
+
86
+ # 1. Setup a grant store
87
+ grant_store = InMemoryGrantStore()
88
+ tower = ControlTower(..., grant_store=grant_store)
89
+
90
+ # 2. Issue a grant (e.g., after initial human approval)
91
+ grant = Grant(
92
+ agent_id="my-agent",
93
+ effect=Effect.WRITE,
94
+ tool="mcp:*", # Wildcard prefix: matches any MCP tool (e.g., "mcp:server.write")
95
+ action=None, # Wildcard: matches any action
96
+ resource_type=None,
97
+ expires_at=time.time() + 3600, # Valid for 1 hour
98
+ granted_by="admin-user",
99
+ created_at=time.time()
100
+ )
101
+ await grant_store.create_grant(grant)
102
+ ```
103
+
80
104
  ### MCP (Model Context Protocol)
81
105
  Wrap an MCP client to gate all tool calls:
82
106
  ```python
@@ -52,6 +52,30 @@ Runtime enforcement layer for AI agent tool calls using **Identity + Intent + Po
52
52
 
53
53
  ## 🚀 v1 Integrations
54
54
 
55
+ ### 🎟️ Session Grants
56
+ Grants allow you to pre-authorize specific actions for an agent session, bypassing human-in-the-loop approvals for repetitive or low-risk tasks.
57
+
58
+ ```python
59
+ from tollgate import Grant, InMemoryGrantStore, Effect
60
+
61
+ # 1. Setup a grant store
62
+ grant_store = InMemoryGrantStore()
63
+ tower = ControlTower(..., grant_store=grant_store)
64
+
65
+ # 2. Issue a grant (e.g., after initial human approval)
66
+ grant = Grant(
67
+ agent_id="my-agent",
68
+ effect=Effect.WRITE,
69
+ tool="mcp:*", # Wildcard prefix: matches any MCP tool (e.g., "mcp:server.write")
70
+ action=None, # Wildcard: matches any action
71
+ resource_type=None,
72
+ expires_at=time.time() + 3600, # Valid for 1 hour
73
+ granted_by="admin-user",
74
+ created_at=time.time()
75
+ )
76
+ await grant_store.create_grant(grant)
77
+ ```
78
+
55
79
  ### MCP (Model Context Protocol)
56
80
  Wrap an MCP client to gate all tool calls:
57
81
  ```python
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "tollgate"
7
- version = "1.0.3"
7
+ version = "1.0.4"
8
8
  description = "Runtime enforcement layer for AI agent tool calls using Identity + Intent + Policy"
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.10"
@@ -1,5 +1,4 @@
1
1
  from .approvals import (
2
- ApprovalOutcome,
3
2
  ApprovalStore,
4
3
  Approver,
5
4
  AsyncQueueApprover,
@@ -15,23 +14,26 @@ from .exceptions import (
15
14
  TollgateDenied,
16
15
  TollgateError,
17
16
  )
17
+ from .grants import InMemoryGrantStore
18
18
  from .helpers import guard, wrap_tool
19
19
  from .policy import PolicyEvaluator, YamlPolicyEvaluator
20
20
  from .registry import ToolRegistry
21
21
  from .tower import ControlTower
22
22
  from .types import (
23
23
  AgentContext,
24
+ ApprovalOutcome,
24
25
  AuditEvent,
25
26
  Decision,
26
27
  DecisionType,
27
28
  Effect,
29
+ Grant,
28
30
  Intent,
29
31
  NormalizedToolCall,
30
32
  Outcome,
31
33
  ToolRequest,
32
34
  )
33
35
 
34
- __version__ = "1.0.3"
36
+ __version__ = "1.0.4"
35
37
 
36
38
  __all__ = [
37
39
  "ControlTower",
@@ -42,6 +44,7 @@ __all__ = [
42
44
  "Decision",
43
45
  "DecisionType",
44
46
  "Effect",
47
+ "Grant",
45
48
  "AuditEvent",
46
49
  "Outcome",
47
50
  "ApprovalOutcome",
@@ -57,6 +60,7 @@ __all__ = [
57
60
  "ToolRegistry",
58
61
  "PolicyEvaluator",
59
62
  "YamlPolicyEvaluator",
63
+ "InMemoryGrantStore",
60
64
  "TollgateError",
61
65
  "TollgateDenied",
62
66
  "TollgateApprovalDenied",
@@ -0,0 +1,103 @@
1
+ import asyncio
2
+ import time
3
+
4
+ from .types import AgentContext, Grant, ToolRequest
5
+
6
+
7
+ class InMemoryGrantStore:
8
+ """In-memory store for action grants with thread-safe matching logic."""
9
+
10
+ def __init__(self):
11
+ self._grants: dict[str, Grant] = {}
12
+ self._usage_counts: dict[str, int] = {}
13
+ self._lock = asyncio.Lock()
14
+
15
+ async def create_grant(self, grant: Grant) -> str:
16
+ """Store a new grant."""
17
+ async with self._lock:
18
+ self._grants[grant.id] = grant
19
+ self._usage_counts[grant.id] = 0
20
+ return grant.id
21
+
22
+ async def find_matching_grant(
23
+ self, agent_ctx: AgentContext, tool_request: ToolRequest
24
+ ) -> Grant | None:
25
+ """Find a non-expired grant that matches the request."""
26
+ now = time.time()
27
+ async with self._lock:
28
+ for grant in self._grants.values():
29
+ # 1. Skip expired
30
+ if grant.expires_at <= now:
31
+ continue
32
+
33
+ # 2. Match agent_id
34
+ if grant.agent_id is not None and grant.agent_id != agent_ctx.agent_id:
35
+ continue
36
+
37
+ # 3. Match effect
38
+ if grant.effect is not None and grant.effect != tool_request.effect:
39
+ continue
40
+
41
+ # 4. Match tool (exact or prefix with *)
42
+ if grant.tool is not None:
43
+ if grant.tool.endswith("*"):
44
+ prefix = grant.tool[:-1]
45
+ if not tool_request.tool.startswith(prefix):
46
+ continue
47
+ elif grant.tool != tool_request.tool:
48
+ continue
49
+
50
+ # 5. Match action
51
+ if grant.action is not None and grant.action != tool_request.action:
52
+ continue
53
+
54
+ # 6. Match resource_type
55
+ if (
56
+ grant.resource_type is not None
57
+ and grant.resource_type != tool_request.resource_type
58
+ ):
59
+ continue
60
+
61
+ # Match found! Increment usage count
62
+ self._usage_counts[grant.id] += 1
63
+ return grant
64
+ return None
65
+
66
+ async def get_usage_count(self, grant_id: str) -> int:
67
+ """Get the number of times a grant has been used."""
68
+ async with self._lock:
69
+ return self._usage_counts.get(grant_id, 0)
70
+
71
+ async def revoke_grant(self, grant_id: str) -> bool:
72
+ """Remove a grant by ID."""
73
+ async with self._lock:
74
+ if grant_id in self._grants:
75
+ del self._grants[grant_id]
76
+ if grant_id in self._usage_counts:
77
+ del self._usage_counts[grant_id]
78
+ return True
79
+ return False
80
+
81
+ async def list_active_grants(self, agent_id: str | None = None) -> list[Grant]:
82
+ """List all non-expired grants, optionally filtered by agent."""
83
+ now = time.time()
84
+ async with self._lock:
85
+ active = []
86
+ for grant in self._grants.values():
87
+ if grant.expires_at > now:
88
+ if agent_id is None or grant.agent_id == agent_id:
89
+ active.append(grant)
90
+ return active
91
+
92
+ async def cleanup_expired(self) -> int:
93
+ """Remove all expired grants from the store."""
94
+ now = time.time()
95
+ async with self._lock:
96
+ to_remove = [
97
+ gid for gid, g in self._grants.items() if g.expires_at <= now
98
+ ]
99
+ for gid in to_remove:
100
+ del self._grants[gid]
101
+ if gid in self._usage_counts:
102
+ del self._usage_counts[gid]
103
+ return len(to_remove)
@@ -32,11 +32,13 @@ class ControlTower:
32
32
  policy: PolicyEvaluator,
33
33
  approver: Approver,
34
34
  audit: AuditSink,
35
+ grant_store: Any | None = None,
35
36
  redact_fn: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
36
37
  ):
37
38
  self.policy = policy
38
39
  self.approver = approver
39
40
  self.audit = audit
41
+ self.grant_store = grant_store
40
42
  self.redact_fn = redact_fn or self._default_redact
41
43
 
42
44
  @staticmethod
@@ -86,6 +88,26 @@ class ControlTower:
86
88
 
87
89
  # 3. Handle ASK
88
90
  if decision.decision == DecisionType.ASK:
91
+ # 3.1 Check Grants
92
+ if self.grant_store:
93
+ matching_grant = await self.grant_store.find_matching_grant(
94
+ agent_ctx, tool_request
95
+ )
96
+ if matching_grant:
97
+ # Grant found! Proceed to execution without asking approver
98
+ result = await self._execute_and_log(
99
+ correlation_id,
100
+ request_hash,
101
+ agent_ctx,
102
+ intent,
103
+ tool_request,
104
+ decision,
105
+ exec_async,
106
+ grant_id=matching_grant.id,
107
+ )
108
+ return result
109
+
110
+ # 3.2 Request Approval if no grant found
89
111
  outcome = await self.approver.request_approval_async(
90
112
  agent_ctx, intent, tool_request, request_hash, decision.reason
91
113
  )
@@ -120,14 +142,35 @@ class ControlTower:
120
142
  )
121
143
  raise TollgateApprovalDenied(f"Approval failed: {outcome.value}")
122
144
 
123
- # 4. Execute tool
145
+ # 4. Execute tool (Policy ALLOW or Approval APPROVED)
146
+ return await self._execute_and_log(
147
+ correlation_id,
148
+ request_hash,
149
+ agent_ctx,
150
+ intent,
151
+ tool_request,
152
+ decision,
153
+ exec_async,
154
+ )
155
+
156
+ async def _execute_and_log(
157
+ self,
158
+ correlation_id: str,
159
+ request_hash: str,
160
+ agent_ctx: AgentContext,
161
+ intent: Intent,
162
+ tool_request: ToolRequest,
163
+ decision: Decision,
164
+ exec_async: Callable[[], Awaitable[Any]],
165
+ grant_id: str | None = None,
166
+ ) -> Any:
167
+ """Internal helper to execute tool and log result."""
124
168
  result = None
125
169
  outcome = Outcome.EXECUTED
126
170
  try:
127
171
  result = await exec_async()
128
172
  except Exception as e:
129
173
  outcome = Outcome.FAILED
130
- # Security: Sanitize exception message to avoid info disclosure
131
174
  result_summary = self._sanitize_exception(e)
132
175
  self._log(
133
176
  correlation_id,
@@ -137,11 +180,12 @@ class ControlTower:
137
180
  tool_request,
138
181
  decision,
139
182
  outcome,
183
+ grant_id=grant_id,
140
184
  result_summary=result_summary,
141
185
  )
142
186
  raise
143
187
 
144
- # 5. Final Audit
188
+ # Final Audit
145
189
  result_summary = self._truncate_result(result)
146
190
  self._log(
147
191
  correlation_id,
@@ -151,6 +195,7 @@ class ControlTower:
151
195
  tool_request,
152
196
  decision,
153
197
  outcome,
198
+ grant_id=grant_id,
154
199
  result_summary=result_summary,
155
200
  )
156
201
 
@@ -177,7 +222,9 @@ class ControlTower:
177
222
  async def _exec():
178
223
  return exec_sync()
179
224
 
180
- return asyncio.run(self.execute_async(agent_ctx, intent, tool_request, _exec))
225
+ return asyncio.run(
226
+ self.execute_async(agent_ctx, intent, tool_request, _exec)
227
+ )
181
228
 
182
229
  def _log(
183
230
  self,
@@ -189,6 +236,7 @@ class ControlTower:
189
236
  decision: Decision,
190
237
  outcome: Outcome,
191
238
  approval_id: str | None = None,
239
+ grant_id: str | None = None,
192
240
  result_summary: str | None = None,
193
241
  ):
194
242
  # Redact params before logging
@@ -212,47 +260,20 @@ class ControlTower:
212
260
  decision=decision,
213
261
  outcome=outcome,
214
262
  approval_id=approval_id,
263
+ grant_id=grant_id,
215
264
  result_summary=result_summary,
216
265
  policy_version=decision.policy_version,
217
266
  manifest_version=req.manifest_version,
218
267
  )
219
268
  self.audit.emit(event)
220
269
 
270
+ def _sanitize_exception(self, e: Exception) -> str:
271
+ """Sanitize exception message to avoid leaking sensitive data."""
272
+ # Only include the exception type and a generic message for security
273
+ return f"{type(e).__name__}: Execution failed"
274
+
221
275
  def _truncate_result(self, result: Any, max_chars: int = 200) -> str | None:
222
276
  if result is None:
223
277
  return None
224
278
  s = str(result)
225
279
  return s[:max_chars] + "..." if len(s) > max_chars else s
226
-
227
- def _sanitize_exception(self, e: Exception, max_chars: int = 100) -> str:
228
- """
229
- Sanitize exception message for audit logging.
230
-
231
- Removes potentially sensitive information like file paths, stack traces,
232
- and internal details while preserving the exception type.
233
- """
234
- exc_type = type(e).__name__
235
-
236
- # For common safe exceptions, include a truncated message
237
- safe_exceptions = {
238
- "ValueError",
239
- "TypeError",
240
- "KeyError",
241
- "IndexError",
242
- "AttributeError",
243
- "RuntimeError",
244
- }
245
-
246
- if exc_type in safe_exceptions:
247
- msg = str(e)
248
- # Remove file paths (Unix and Windows style)
249
- import re
250
-
251
- msg = re.sub(r"['\"]?(/[^\s'\"]+|[A-Za-z]:\\[^\s'\"]+)['\"]?", "[PATH]", msg)
252
- # Truncate to avoid leaking too much info
253
- if len(msg) > max_chars:
254
- msg = msg[:max_chars] + "..."
255
- return f"{exc_type}: {msg}"
256
-
257
- # For unknown exceptions, only log the type
258
- return f"{exc_type}: [details redacted]"
@@ -1,3 +1,4 @@
1
+ import uuid
1
2
  from collections.abc import Awaitable, Callable
2
3
  from dataclasses import asdict, dataclass, field
3
4
  from enum import Enum
@@ -92,6 +93,28 @@ class Decision:
92
93
  return d
93
94
 
94
95
 
96
+ @dataclass(frozen=True)
97
+ class Grant:
98
+ """A grant that allows bypassing human approval for specific actions."""
99
+
100
+ agent_id: str | None # None = any agent
101
+ effect: Effect | None # None = any effect
102
+ tool: str | None # None = any tool, supports prefix like "mcp:*"
103
+ action: str | None # None = any action
104
+ resource_type: str | None # None = any resource
105
+ expires_at: float # Unix timestamp
106
+ granted_by: str
107
+ created_at: float
108
+ id: str = field(default_factory=lambda: str(uuid.uuid4()))
109
+ reason: str | None = None
110
+
111
+ def to_dict(self) -> dict[str, Any]:
112
+ d = asdict(self)
113
+ if self.effect:
114
+ d["effect"] = self.effect.value
115
+ return d
116
+
117
+
95
118
  @dataclass(frozen=True)
96
119
  class AuditEvent:
97
120
  timestamp: str
@@ -103,6 +126,7 @@ class AuditEvent:
103
126
  decision: Decision
104
127
  outcome: Outcome
105
128
  approval_id: str | None = None
129
+ grant_id: str | None = None
106
130
  result_summary: str | None = None
107
131
  policy_version: str | None = None
108
132
  manifest_version: str | None = None
@@ -118,6 +142,7 @@ class AuditEvent:
118
142
  "decision": self.decision.to_dict(),
119
143
  "outcome": self.outcome.value,
120
144
  "approval_id": self.approval_id,
145
+ "grant_id": self.grant_id,
121
146
  "result_summary": self.result_summary,
122
147
  "policy_version": self.policy_version,
123
148
  "manifest_version": self.manifest_version,
@@ -0,0 +1,286 @@
1
+ import asyncio
2
+ import time
3
+
4
+ import pytest
5
+
6
+ from tollgate import (
7
+ AgentContext,
8
+ ControlTower,
9
+ Decision,
10
+ DecisionType,
11
+ Effect,
12
+ Grant,
13
+ InMemoryGrantStore,
14
+ Intent,
15
+ Outcome,
16
+ ToolRequest,
17
+ )
18
+
19
+
20
+ class MockPolicy:
21
+ def evaluate(self, _ctx, _intent, _req):
22
+ return Decision(
23
+ decision=DecisionType.ASK, reason="Needs approval", policy_version="1"
24
+ )
25
+
26
+
27
+ class MockApprover:
28
+ async def request_approval_async(self, *_args):
29
+ # This should NOT be called if a grant matches
30
+ pytest.fail("Approver should not be called when a grant matches")
31
+
32
+
33
+ class MockAudit:
34
+ def __init__(self):
35
+ self.events = []
36
+
37
+ def emit(self, event):
38
+ self.events.append(event)
39
+
40
+
41
+ @pytest.fixture
42
+ def agent_ctx():
43
+ return AgentContext(agent_id="agent-1", version="1.0.0", owner="user-1")
44
+
45
+
46
+ @pytest.fixture
47
+ def tool_req():
48
+ return ToolRequest(
49
+ tool="mcp:server.read",
50
+ action="read",
51
+ resource_type="data",
52
+ effect=Effect.READ,
53
+ params={"id": 1},
54
+ )
55
+
56
+
57
+ @pytest.mark.asyncio
58
+ async def test_create_and_find_grant(agent_ctx, tool_req):
59
+ store = InMemoryGrantStore()
60
+ # Test auto-generating ID
61
+ grant = Grant(
62
+ agent_id="agent-1",
63
+ effect=Effect.READ,
64
+ tool="mcp:server.read",
65
+ action="read",
66
+ resource_type="data",
67
+ expires_at=time.time() + 3600,
68
+ granted_by="admin",
69
+ created_at=time.time(),
70
+ )
71
+ assert grant.id is not None
72
+ await store.create_grant(grant)
73
+
74
+ found = await store.find_matching_grant(agent_ctx, tool_req)
75
+ assert found is not None
76
+ assert found.id == grant.id
77
+
78
+
79
+ @pytest.mark.asyncio
80
+ async def test_grant_usage_counter(agent_ctx, tool_req):
81
+ store = InMemoryGrantStore()
82
+ grant = Grant(
83
+ agent_id="agent-1",
84
+ effect=Effect.READ,
85
+ tool="mcp:server.read",
86
+ action="read",
87
+ resource_type="data",
88
+ expires_at=time.time() + 3600,
89
+ granted_by="admin",
90
+ created_at=time.time(),
91
+ )
92
+ await store.create_grant(grant)
93
+
94
+ assert await store.get_usage_count(grant.id) == 0
95
+ await store.find_matching_grant(agent_ctx, tool_req)
96
+ assert await store.get_usage_count(grant.id) == 1
97
+ await store.find_matching_grant(agent_ctx, tool_req)
98
+ assert await store.get_usage_count(grant.id) == 2
99
+
100
+
101
+ @pytest.mark.asyncio
102
+ async def test_grant_expiry(agent_ctx, tool_req):
103
+ store = InMemoryGrantStore()
104
+ grant = Grant(
105
+ agent_id="agent-1",
106
+ effect=Effect.READ,
107
+ tool="mcp:server.read",
108
+ action="read",
109
+ resource_type="data",
110
+ expires_at=time.time() - 1, # Expired
111
+ granted_by="admin",
112
+ created_at=time.time() - 3600,
113
+ )
114
+ await store.create_grant(grant)
115
+
116
+ found = await store.find_matching_grant(agent_ctx, tool_req)
117
+ assert found is None
118
+
119
+
120
+ @pytest.mark.asyncio
121
+ async def test_grant_wildcard_agent(tool_req):
122
+ store = InMemoryGrantStore()
123
+ grant = Grant(
124
+ agent_id=None, # Wildcard
125
+ effect=Effect.READ,
126
+ tool="mcp:server.read",
127
+ action="read",
128
+ resource_type="data",
129
+ expires_at=time.time() + 3600,
130
+ granted_by="admin",
131
+ created_at=time.time(),
132
+ )
133
+ await store.create_grant(grant)
134
+
135
+ ctx2 = AgentContext(agent_id="any-agent", version="1", owner="o")
136
+ found = await store.find_matching_grant(ctx2, tool_req)
137
+ assert found is not None
138
+
139
+
140
+ @pytest.mark.asyncio
141
+ async def test_grant_wildcard_effect(agent_ctx, tool_req):
142
+ store = InMemoryGrantStore()
143
+ grant = Grant(
144
+ agent_id="agent-1",
145
+ effect=None, # Wildcard
146
+ tool="mcp:server.read",
147
+ action="read",
148
+ resource_type="data",
149
+ expires_at=time.time() + 3600,
150
+ granted_by="admin",
151
+ created_at=time.time(),
152
+ )
153
+ await store.create_grant(grant)
154
+
155
+ found = await store.find_matching_grant(agent_ctx, tool_req)
156
+ assert found is not None
157
+
158
+
159
+ @pytest.mark.asyncio
160
+ async def test_grant_tool_prefix_match(agent_ctx, tool_req):
161
+ store = InMemoryGrantStore()
162
+ grant = Grant(
163
+ agent_id="agent-1",
164
+ effect=Effect.READ,
165
+ tool="mcp:*", # Prefix match
166
+ action="read",
167
+ resource_type="data",
168
+ expires_at=time.time() + 3600,
169
+ granted_by="admin",
170
+ created_at=time.time(),
171
+ )
172
+ await store.create_grant(grant)
173
+
174
+ found = await store.find_matching_grant(agent_ctx, tool_req)
175
+ assert found is not None
176
+
177
+
178
+ @pytest.mark.asyncio
179
+ async def test_grant_no_match(agent_ctx):
180
+ store = InMemoryGrantStore()
181
+ grant = Grant(
182
+ agent_id="agent-1",
183
+ effect=Effect.WRITE, # Mismatch (READ vs WRITE)
184
+ tool="mcp:server.read",
185
+ action="read",
186
+ resource_type="data",
187
+ expires_at=time.time() + 3600,
188
+ granted_by="admin",
189
+ created_at=time.time(),
190
+ )
191
+ await store.create_grant(grant)
192
+
193
+ req = ToolRequest(
194
+ tool="mcp:server.read",
195
+ action="read",
196
+ resource_type="data",
197
+ effect=Effect.READ,
198
+ params={},
199
+ )
200
+ found = await store.find_matching_grant(agent_ctx, req)
201
+ assert found is None
202
+
203
+
204
+ @pytest.mark.asyncio
205
+ async def test_revoke_grant(agent_ctx, tool_req):
206
+ store = InMemoryGrantStore()
207
+ grant = Grant(
208
+ agent_id="agent-1",
209
+ effect=Effect.READ,
210
+ tool="mcp:server.read",
211
+ action="read",
212
+ resource_type="data",
213
+ expires_at=time.time() + 3600,
214
+ granted_by="admin",
215
+ created_at=time.time(),
216
+ )
217
+ await store.create_grant(grant)
218
+ await store.revoke_grant(grant.id)
219
+
220
+ found = await store.find_matching_grant(agent_ctx, tool_req)
221
+ assert found is None
222
+
223
+
224
+ @pytest.mark.asyncio
225
+ async def test_cleanup_expired():
226
+ store = InMemoryGrantStore()
227
+ g1 = Grant(
228
+ agent_id=None,
229
+ effect=None,
230
+ tool=None,
231
+ action=None,
232
+ resource_type=None,
233
+ expires_at=time.time() - 1,
234
+ granted_by="a",
235
+ created_at=time.time(),
236
+ )
237
+ g2 = Grant(
238
+ agent_id=None,
239
+ effect=None,
240
+ tool=None,
241
+ action=None,
242
+ resource_type=None,
243
+ expires_at=time.time() + 10,
244
+ granted_by="a",
245
+ created_at=time.time(),
246
+ )
247
+ await store.create_grant(g1)
248
+ await store.create_grant(g2)
249
+
250
+ cleaned = await store.cleanup_expired()
251
+ assert cleaned == 1
252
+ active = await store.list_active_grants()
253
+ assert len(active) == 1
254
+ assert active[0].id == g2.id
255
+
256
+
257
+ @pytest.mark.asyncio
258
+ async def test_tower_uses_grant(agent_ctx, tool_req):
259
+ store = InMemoryGrantStore()
260
+ audit = MockAudit()
261
+ approver = MockApprover()
262
+ tower = ControlTower(MockPolicy(), approver, audit, grant_store=store)
263
+
264
+ grant = Grant(
265
+ agent_id="agent-1",
266
+ effect=Effect.READ,
267
+ tool="mcp:server.read",
268
+ action="read",
269
+ resource_type="data",
270
+ expires_at=time.time() + 3600,
271
+ granted_by="admin",
272
+ created_at=time.time(),
273
+ )
274
+ await store.create_grant(grant)
275
+
276
+ async def tool_fn():
277
+ return "ok"
278
+
279
+ intent = Intent(action="test", reason="test")
280
+ result = await tower.execute_async(agent_ctx, intent, tool_req, tool_fn)
281
+
282
+ assert result == "ok"
283
+ assert len(audit.events) > 0
284
+ event = audit.events[-1]
285
+ assert event.grant_id == grant.id
286
+ assert event.outcome == Outcome.EXECUTED
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes