agent-loop-guard-runtime 0.6.0a2__py3-none-any.whl
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.
- agent_loop_guard_runtime-0.6.0a2.dist-info/METADATA +407 -0
- agent_loop_guard_runtime-0.6.0a2.dist-info/RECORD +76 -0
- agent_loop_guard_runtime-0.6.0a2.dist-info/WHEEL +5 -0
- agent_loop_guard_runtime-0.6.0a2.dist-info/entry_points.txt +2 -0
- agent_loop_guard_runtime-0.6.0a2.dist-info/licenses/LICENSE +202 -0
- agent_loop_guard_runtime-0.6.0a2.dist-info/top_level.txt +1 -0
- app/__init__.py +6 -0
- app/api/__init__.py +1 -0
- app/api/admin_routes.py +179 -0
- app/api/anthropic_routes.py +21 -0
- app/api/common.py +457 -0
- app/api/mcp_routes.py +218 -0
- app/api/openai_routes.py +26 -0
- app/api/replay_routes.py +202 -0
- app/api/ui_routes.py +295 -0
- app/benchmark/__init__.py +2 -0
- app/benchmark/adapters.py +97 -0
- app/benchmark/data/starter-v1.json +38 -0
- app/benchmark/dataset.py +59 -0
- app/benchmark/models.py +46 -0
- app/benchmark/runner.py +63 -0
- app/benchmark/scorers.py +34 -0
- app/benchmark/statistics.py +48 -0
- app/benchmark/storage.py +78 -0
- app/cli.py +624 -0
- app/core/config.py +196 -0
- app/core/demo.py +109 -0
- app/core/loop_detector.py +120 -0
- app/core/policy_engine.py +167 -0
- app/core/redaction.py +84 -0
- app/core/security.py +54 -0
- app/core/token_meter.py +67 -0
- app/db/models.py +296 -0
- app/db/repository.py +1488 -0
- app/db/session.py +57 -0
- app/main.py +59 -0
- app/mcp/__init__.py +1 -0
- app/mcp/gateway.py +230 -0
- app/mcp/policy.py +281 -0
- app/mcp/presets/development.yml +25 -0
- app/mcp/presets/filesystem.yml +18 -0
- app/mcp/stdio.py +142 -0
- app/platform/__init__.py +1 -0
- app/platform/alembic/__init__.py +2 -0
- app/platform/alembic/env.py +38 -0
- app/platform/alembic/script.py.mako +24 -0
- app/platform/alembic/versions/0001_initial_schema.py +18 -0
- app/platform/alembic/versions/__init__.py +2 -0
- app/platform/events.py +44 -0
- app/platform/maintenance.py +138 -0
- app/platform/migrations.py +24 -0
- app/platform/setup.py +92 -0
- app/providers/__init__.py +5 -0
- app/providers/base.py +23 -0
- app/providers/mock.py +169 -0
- app/providers/upstream.py +101 -0
- app/replay/__init__.py +1 -0
- app/replay/costs.py +37 -0
- app/replay/formats.py +87 -0
- app/replay/sdk.py +102 -0
- app/sandbox/__init__.py +2 -0
- app/sandbox/policy.py +22 -0
- app/sandbox/workspace.py +281 -0
- app/static/styles.css +327 -0
- app/templates/agents.html +48 -0
- app/templates/base.html +27 -0
- app/templates/dashboard.html +55 -0
- app/templates/demo.html +36 -0
- app/templates/mcp.html +69 -0
- app/templates/policies.html +30 -0
- app/templates/replay.html +63 -0
- app/templates/replay_compare.html +74 -0
- app/templates/replay_detail.html +130 -0
- app/templates/session_detail.html +71 -0
- app/templates/sessions.html +24 -0
- app/templates/settings.html +20 -0
app/api/common.py
ADDED
|
@@ -0,0 +1,457 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import json
|
|
4
|
+
import time
|
|
5
|
+
from collections.abc import AsyncIterator
|
|
6
|
+
from typing import Any
|
|
7
|
+
|
|
8
|
+
from fastapi import HTTPException, Request
|
|
9
|
+
from fastapi.responses import JSONResponse, Response, StreamingResponse
|
|
10
|
+
from sqlalchemy.orm import Session
|
|
11
|
+
|
|
12
|
+
from app.core.config import AppConfig
|
|
13
|
+
from app.core.loop_detector import (
|
|
14
|
+
error_fingerprint,
|
|
15
|
+
request_fingerprint,
|
|
16
|
+
tool_call_fingerprints,
|
|
17
|
+
)
|
|
18
|
+
from app.core.policy_engine import MESSAGE_BY_REASON, PolicyDecision, evaluate
|
|
19
|
+
from app.core.redaction import safe_json_bytes, safe_preview
|
|
20
|
+
from app.core.security import bearer_token, external_session_id
|
|
21
|
+
from app.core.token_meter import estimate_tokens, usage_or_estimate
|
|
22
|
+
from app.db.models import Agent, GuardSession, Policy, Project
|
|
23
|
+
from app.db.repository import Repository, new_id
|
|
24
|
+
from app.providers import MockProvider, ProviderResult, ProviderStream, UpstreamProvider
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
def _incoming_headers(request: Request) -> dict[str, str]:
|
|
28
|
+
return {key: value for key, value in request.headers.items()}
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
async def read_json_body(request: Request, config: AppConfig) -> tuple[bytes, dict[str, Any]]:
|
|
32
|
+
raw = await request.body()
|
|
33
|
+
if len(raw) > config.body_limit_bytes:
|
|
34
|
+
raise HTTPException(status_code=413, detail={"error": "Request body too large."})
|
|
35
|
+
if not raw:
|
|
36
|
+
return raw, {}
|
|
37
|
+
try:
|
|
38
|
+
parsed = json.loads(raw)
|
|
39
|
+
except json.JSONDecodeError as exc:
|
|
40
|
+
raise HTTPException(status_code=400, detail={"error": "Malformed JSON body."}) from exc
|
|
41
|
+
if not isinstance(parsed, dict):
|
|
42
|
+
raise HTTPException(status_code=400, detail={"error": "JSON body must be an object."})
|
|
43
|
+
return raw, parsed
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
def mode_for_request(request: Request, config: AppConfig, project: Project) -> str:
|
|
47
|
+
mode = project.mode or config.default_mode
|
|
48
|
+
if config.allow_mode_header:
|
|
49
|
+
requested = request.headers.get("x-alg-mode")
|
|
50
|
+
if requested in {"shadow", "enforce"}:
|
|
51
|
+
mode = requested
|
|
52
|
+
return "enforce" if mode == "enforce" else "shadow"
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
def provider_for(project: Project, config: AppConfig, protocol: str):
|
|
56
|
+
if (project.provider or config.default_provider) == "mock":
|
|
57
|
+
return MockProvider()
|
|
58
|
+
if protocol == "anthropic":
|
|
59
|
+
return UpstreamProvider(
|
|
60
|
+
base_url=config.anthropic_base_url,
|
|
61
|
+
api_key=config.anthropic_api_key,
|
|
62
|
+
protocol="anthropic",
|
|
63
|
+
)
|
|
64
|
+
return UpstreamProvider(
|
|
65
|
+
base_url=config.openai_base_url,
|
|
66
|
+
api_key=config.openai_api_key,
|
|
67
|
+
protocol="openai",
|
|
68
|
+
)
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
def response_headers(base: dict[str, str], decision: PolicyDecision, request_id: str) -> dict[str, str]:
|
|
72
|
+
headers = dict(base)
|
|
73
|
+
headers["x-alg-request-id"] = request_id
|
|
74
|
+
headers["x-alg-decision"] = decision.decision
|
|
75
|
+
if decision.reason_code:
|
|
76
|
+
headers["x-alg-reason"] = decision.reason_code
|
|
77
|
+
return headers
|
|
78
|
+
|
|
79
|
+
|
|
80
|
+
def _block_payload(decision: PolicyDecision, session: GuardSession) -> dict[str, Any]:
|
|
81
|
+
return {
|
|
82
|
+
"error": {
|
|
83
|
+
"type": "agent_loop_guard_block",
|
|
84
|
+
"code": decision.reason_code,
|
|
85
|
+
"message": decision.message,
|
|
86
|
+
"session_id": session.id,
|
|
87
|
+
"rule_id": decision.rule_id,
|
|
88
|
+
}
|
|
89
|
+
}
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
def block_response(decision: PolicyDecision, session: GuardSession, request_id: str) -> JSONResponse:
|
|
93
|
+
payload = _block_payload(decision, session)
|
|
94
|
+
return JSONResponse(
|
|
95
|
+
payload,
|
|
96
|
+
status_code=decision.http_status,
|
|
97
|
+
headers=response_headers({}, decision, request_id),
|
|
98
|
+
)
|
|
99
|
+
|
|
100
|
+
|
|
101
|
+
def _record_policy_event(repo: Repository, record_id: str | None, session_id: str, decision: PolicyDecision) -> None:
|
|
102
|
+
if not decision.reason_code:
|
|
103
|
+
return
|
|
104
|
+
repo.event(
|
|
105
|
+
session_id,
|
|
106
|
+
record_id,
|
|
107
|
+
"policy_triggered",
|
|
108
|
+
"error" if decision.blocked else "warning",
|
|
109
|
+
decision.reason_code,
|
|
110
|
+
decision.mode,
|
|
111
|
+
decision.details or {},
|
|
112
|
+
decision.rule_id,
|
|
113
|
+
)
|
|
114
|
+
|
|
115
|
+
|
|
116
|
+
def _policies_by_type(policies: list[Policy]) -> dict[str, Policy]:
|
|
117
|
+
return {policy.rule_type: policy for policy in policies if policy.enabled}
|
|
118
|
+
|
|
119
|
+
|
|
120
|
+
def _post_error_decision(
|
|
121
|
+
*,
|
|
122
|
+
mode: str,
|
|
123
|
+
policies: list[Policy],
|
|
124
|
+
error_fp: str | None,
|
|
125
|
+
previous_error: tuple[str | None, int],
|
|
126
|
+
) -> PolicyDecision | None:
|
|
127
|
+
if not error_fp:
|
|
128
|
+
return None
|
|
129
|
+
policy = _policies_by_type(policies).get("error_retry")
|
|
130
|
+
if policy is None:
|
|
131
|
+
return None
|
|
132
|
+
previous_fp, previous_count = previous_error
|
|
133
|
+
count = previous_count + 1 if previous_fp == error_fp else 1
|
|
134
|
+
if count < policy.threshold:
|
|
135
|
+
return None
|
|
136
|
+
return PolicyDecision(
|
|
137
|
+
decision="shadow_flag",
|
|
138
|
+
mode=mode,
|
|
139
|
+
reason_code="LOOP_ERROR_RETRY",
|
|
140
|
+
rule_id=policy.id,
|
|
141
|
+
message=MESSAGE_BY_REASON["LOOP_ERROR_RETRY"],
|
|
142
|
+
details={"count": count, "threshold": policy.threshold, "fingerprint": error_fp},
|
|
143
|
+
http_status=200,
|
|
144
|
+
)
|
|
145
|
+
|
|
146
|
+
|
|
147
|
+
def _preview_response(result: ProviderResult, config: AppConfig) -> str:
|
|
148
|
+
if result.json_body is not None:
|
|
149
|
+
return safe_preview(result.json_body, config.full_content_logging)
|
|
150
|
+
if config.full_content_logging:
|
|
151
|
+
return safe_json_bytes(result.content)
|
|
152
|
+
return safe_preview(
|
|
153
|
+
{"bytes": len(result.content), "content_type": result.media_type or "unknown"},
|
|
154
|
+
full_content_logging=False,
|
|
155
|
+
)
|
|
156
|
+
|
|
157
|
+
|
|
158
|
+
def _model_from_body(body: dict[str, Any]) -> str | None:
|
|
159
|
+
model = body.get("model")
|
|
160
|
+
return str(model) if model is not None else None
|
|
161
|
+
|
|
162
|
+
|
|
163
|
+
def _tool_policy_fps(tool_calls: list[tuple[str, str]]) -> list[str]:
|
|
164
|
+
return [f"{name}:{args_hash}" for name, args_hash in tool_calls]
|
|
165
|
+
|
|
166
|
+
|
|
167
|
+
def _authenticate(repo: Repository, request: Request) -> Agent:
|
|
168
|
+
raw_key = bearer_token(request)
|
|
169
|
+
agent = repo.agent_for_key(raw_key)
|
|
170
|
+
if agent is None:
|
|
171
|
+
raise HTTPException(
|
|
172
|
+
status_code=403,
|
|
173
|
+
detail={"error": {"type": "agent_loop_guard_auth", "message": "Invalid gateway key."}},
|
|
174
|
+
)
|
|
175
|
+
return agent
|
|
176
|
+
|
|
177
|
+
|
|
178
|
+
def _project(repo: Repository, request: Request, agent: Agent) -> Project:
|
|
179
|
+
requested = request.headers.get("x-alg-project-id")
|
|
180
|
+
project_id = requested.strip() if requested else agent.project_id
|
|
181
|
+
project = repo.project(project_id)
|
|
182
|
+
if project is None or project.id != agent.project_id:
|
|
183
|
+
raise HTTPException(
|
|
184
|
+
status_code=403,
|
|
185
|
+
detail={"error": {"type": "agent_loop_guard_auth", "message": "Project is not available."}},
|
|
186
|
+
)
|
|
187
|
+
return project
|
|
188
|
+
|
|
189
|
+
|
|
190
|
+
def _initial_decision(
|
|
191
|
+
*,
|
|
192
|
+
mode: str,
|
|
193
|
+
repo: Repository,
|
|
194
|
+
session: GuardSession,
|
|
195
|
+
agent: Agent,
|
|
196
|
+
policies: list[Policy],
|
|
197
|
+
body: dict[str, Any],
|
|
198
|
+
req_fp: str,
|
|
199
|
+
tool_fps: list[str],
|
|
200
|
+
) -> PolicyDecision:
|
|
201
|
+
return evaluate(
|
|
202
|
+
mode=mode,
|
|
203
|
+
session=session,
|
|
204
|
+
agent=agent,
|
|
205
|
+
policies=policies,
|
|
206
|
+
request_fp=req_fp,
|
|
207
|
+
tool_fps=tool_fps,
|
|
208
|
+
recent_request_fps=repo.recent_request_fingerprints(session.id, limit=30),
|
|
209
|
+
recent_tool_fps=repo.recent_tool_fingerprints(session.id, limit=30),
|
|
210
|
+
last_error_count=repo.consecutive_error_count(session.id)[1],
|
|
211
|
+
)
|
|
212
|
+
|
|
213
|
+
|
|
214
|
+
def _final_decision(
|
|
215
|
+
*,
|
|
216
|
+
current: PolicyDecision,
|
|
217
|
+
mode: str,
|
|
218
|
+
policies: list[Policy],
|
|
219
|
+
error_fp: str | None,
|
|
220
|
+
previous_error: tuple[str | None, int],
|
|
221
|
+
) -> PolicyDecision:
|
|
222
|
+
if current.reason_code:
|
|
223
|
+
return current
|
|
224
|
+
error_decision = _post_error_decision(
|
|
225
|
+
mode=mode, policies=policies, error_fp=error_fp, previous_error=previous_error
|
|
226
|
+
)
|
|
227
|
+
return error_decision or current
|
|
228
|
+
|
|
229
|
+
|
|
230
|
+
def _record_request(
|
|
231
|
+
*,
|
|
232
|
+
repo: Repository,
|
|
233
|
+
session: GuardSession,
|
|
234
|
+
protocol: str,
|
|
235
|
+
endpoint: str,
|
|
236
|
+
model: str | None,
|
|
237
|
+
req_fp: str,
|
|
238
|
+
status: int,
|
|
239
|
+
latency_ms: int,
|
|
240
|
+
tokens: dict[str, int | bool],
|
|
241
|
+
decision: PolicyDecision,
|
|
242
|
+
request_preview: str | None,
|
|
243
|
+
response_preview: str | None,
|
|
244
|
+
tool_calls: list[tuple[str, str]],
|
|
245
|
+
error_fp: str | None,
|
|
246
|
+
) -> None:
|
|
247
|
+
record = repo.record_request(
|
|
248
|
+
session=session,
|
|
249
|
+
protocol=protocol,
|
|
250
|
+
endpoint=endpoint,
|
|
251
|
+
model=model,
|
|
252
|
+
fingerprint=req_fp,
|
|
253
|
+
status=status,
|
|
254
|
+
latency_ms=latency_ms,
|
|
255
|
+
tokens=tokens,
|
|
256
|
+
decision=decision.decision,
|
|
257
|
+
reason_code=decision.reason_code,
|
|
258
|
+
request_preview=request_preview,
|
|
259
|
+
response_preview=response_preview,
|
|
260
|
+
tool_calls=tool_calls,
|
|
261
|
+
error_fingerprint=error_fp,
|
|
262
|
+
)
|
|
263
|
+
_record_policy_event(repo, record.id, session.id, decision)
|
|
264
|
+
|
|
265
|
+
|
|
266
|
+
async def guarded_proxy(request: Request, db: Session, protocol: str, endpoint: str) -> Response:
|
|
267
|
+
config: AppConfig = request.app.state.config
|
|
268
|
+
request_id = new_id("algreq")
|
|
269
|
+
raw_body, body = await read_json_body(request, config)
|
|
270
|
+
repo = Repository(db)
|
|
271
|
+
agent = _authenticate(repo, request)
|
|
272
|
+
project = _project(repo, request, agent)
|
|
273
|
+
session = repo.get_or_create_session(
|
|
274
|
+
project.id,
|
|
275
|
+
agent.id,
|
|
276
|
+
external_session_id(request, protocol),
|
|
277
|
+
config.inactive_timeout_seconds,
|
|
278
|
+
)
|
|
279
|
+
mode = mode_for_request(request, config, project)
|
|
280
|
+
policies = repo.policies(project.id)
|
|
281
|
+
req_fp = request_fingerprint(body)
|
|
282
|
+
tool_calls = tool_call_fingerprints(body)
|
|
283
|
+
tool_fps = _tool_policy_fps(tool_calls)
|
|
284
|
+
decision = _initial_decision(
|
|
285
|
+
mode=mode,
|
|
286
|
+
repo=repo,
|
|
287
|
+
session=session,
|
|
288
|
+
agent=agent,
|
|
289
|
+
policies=policies,
|
|
290
|
+
body=body,
|
|
291
|
+
req_fp=req_fp,
|
|
292
|
+
tool_fps=tool_fps,
|
|
293
|
+
)
|
|
294
|
+
request_preview = safe_preview(body, config.full_content_logging)
|
|
295
|
+
|
|
296
|
+
if decision.blocked:
|
|
297
|
+
payload = _block_payload(decision, session)
|
|
298
|
+
_record_request(
|
|
299
|
+
repo=repo,
|
|
300
|
+
session=session,
|
|
301
|
+
protocol=protocol,
|
|
302
|
+
endpoint=endpoint,
|
|
303
|
+
model=_model_from_body(body),
|
|
304
|
+
req_fp=req_fp,
|
|
305
|
+
status=decision.http_status,
|
|
306
|
+
latency_ms=0,
|
|
307
|
+
tokens=usage_or_estimate(protocol, body, payload),
|
|
308
|
+
decision=decision,
|
|
309
|
+
request_preview=request_preview,
|
|
310
|
+
response_preview=safe_preview(payload, config.full_content_logging),
|
|
311
|
+
tool_calls=tool_calls,
|
|
312
|
+
error_fp=None,
|
|
313
|
+
)
|
|
314
|
+
return block_response(decision, session, request_id)
|
|
315
|
+
|
|
316
|
+
provider = provider_for(project, config, protocol)
|
|
317
|
+
wants_stream = bool(body.get("stream"))
|
|
318
|
+
if wants_stream:
|
|
319
|
+
if isinstance(provider, MockProvider):
|
|
320
|
+
stream = await provider.stream(protocol, endpoint, body)
|
|
321
|
+
else:
|
|
322
|
+
stream = await provider.stream(endpoint, raw_body, _incoming_headers(request))
|
|
323
|
+
return streaming_response(
|
|
324
|
+
request=request,
|
|
325
|
+
stream=stream,
|
|
326
|
+
protocol=protocol,
|
|
327
|
+
endpoint=endpoint,
|
|
328
|
+
body=body,
|
|
329
|
+
session_id=session.id,
|
|
330
|
+
model=_model_from_body(body),
|
|
331
|
+
req_fp=req_fp,
|
|
332
|
+
decision=decision,
|
|
333
|
+
policies=policies,
|
|
334
|
+
previous_error=repo.consecutive_error_count(session.id),
|
|
335
|
+
request_preview=request_preview,
|
|
336
|
+
tool_calls=tool_calls,
|
|
337
|
+
started_at=time.perf_counter(),
|
|
338
|
+
request_id=request_id,
|
|
339
|
+
)
|
|
340
|
+
|
|
341
|
+
started_at = time.perf_counter()
|
|
342
|
+
if isinstance(provider, MockProvider):
|
|
343
|
+
result = await provider.request(protocol, endpoint, body)
|
|
344
|
+
else:
|
|
345
|
+
result = await provider.request(endpoint, raw_body, _incoming_headers(request))
|
|
346
|
+
latency_ms = int((time.perf_counter() - started_at) * 1000)
|
|
347
|
+
err_fp = error_fingerprint(result.status_code, result.json_body or result.content.decode("utf-8", "replace")) if result.status_code >= 400 else None
|
|
348
|
+
final = _final_decision(
|
|
349
|
+
current=decision,
|
|
350
|
+
mode=mode,
|
|
351
|
+
policies=policies,
|
|
352
|
+
error_fp=err_fp,
|
|
353
|
+
previous_error=repo.consecutive_error_count(session.id),
|
|
354
|
+
)
|
|
355
|
+
_record_request(
|
|
356
|
+
repo=repo,
|
|
357
|
+
session=session,
|
|
358
|
+
protocol=protocol,
|
|
359
|
+
endpoint=endpoint,
|
|
360
|
+
model=_model_from_body(body),
|
|
361
|
+
req_fp=req_fp,
|
|
362
|
+
status=result.status_code,
|
|
363
|
+
latency_ms=latency_ms,
|
|
364
|
+
tokens=usage_or_estimate(protocol, body, result.json_body),
|
|
365
|
+
decision=final,
|
|
366
|
+
request_preview=request_preview,
|
|
367
|
+
response_preview=_preview_response(result, config),
|
|
368
|
+
tool_calls=tool_calls,
|
|
369
|
+
error_fp=err_fp,
|
|
370
|
+
)
|
|
371
|
+
return Response(
|
|
372
|
+
content=result.content,
|
|
373
|
+
status_code=result.status_code,
|
|
374
|
+
headers=response_headers(result.headers, final, request_id),
|
|
375
|
+
media_type=result.media_type,
|
|
376
|
+
)
|
|
377
|
+
|
|
378
|
+
|
|
379
|
+
def streaming_response(
|
|
380
|
+
*,
|
|
381
|
+
request: Request,
|
|
382
|
+
stream: ProviderStream,
|
|
383
|
+
protocol: str,
|
|
384
|
+
endpoint: str,
|
|
385
|
+
body: dict[str, Any],
|
|
386
|
+
session_id: str,
|
|
387
|
+
model: str | None,
|
|
388
|
+
req_fp: str,
|
|
389
|
+
decision: PolicyDecision,
|
|
390
|
+
policies: list[Policy],
|
|
391
|
+
previous_error: tuple[str | None, int],
|
|
392
|
+
request_preview: str | None,
|
|
393
|
+
tool_calls: list[tuple[str, str]],
|
|
394
|
+
started_at: float,
|
|
395
|
+
request_id: str,
|
|
396
|
+
) -> StreamingResponse:
|
|
397
|
+
config: AppConfig = request.app.state.config
|
|
398
|
+
session_factory = request.app.state.SessionLocal
|
|
399
|
+
|
|
400
|
+
async def iterator() -> AsyncIterator[bytes]:
|
|
401
|
+
total_bytes = 0
|
|
402
|
+
preview = bytearray()
|
|
403
|
+
try:
|
|
404
|
+
async for chunk in stream.chunks:
|
|
405
|
+
total_bytes += len(chunk)
|
|
406
|
+
if len(preview) < 1000:
|
|
407
|
+
preview.extend(chunk[: 1000 - len(preview)])
|
|
408
|
+
yield chunk
|
|
409
|
+
finally:
|
|
410
|
+
latency_ms = int((time.perf_counter() - started_at) * 1000)
|
|
411
|
+
preview_text = bytes(preview).decode("utf-8", "replace")
|
|
412
|
+
err_fp = error_fingerprint(stream.status_code, preview_text) if stream.status_code >= 400 else None
|
|
413
|
+
final = _final_decision(
|
|
414
|
+
current=decision,
|
|
415
|
+
mode=decision.mode,
|
|
416
|
+
policies=policies,
|
|
417
|
+
error_fp=err_fp,
|
|
418
|
+
previous_error=previous_error,
|
|
419
|
+
)
|
|
420
|
+
tokens = {
|
|
421
|
+
"input_tokens": estimate_tokens(body),
|
|
422
|
+
"output_tokens": max(1, (total_bytes + 3) // 4),
|
|
423
|
+
"total_tokens": estimate_tokens(body) + max(1, (total_bytes + 3) // 4),
|
|
424
|
+
"estimated": True,
|
|
425
|
+
}
|
|
426
|
+
response_preview = (
|
|
427
|
+
safe_json_bytes(bytes(preview))
|
|
428
|
+
if config.full_content_logging
|
|
429
|
+
else safe_preview({"stream_bytes": total_bytes, "status": stream.status_code}, False)
|
|
430
|
+
)
|
|
431
|
+
with session_factory() as db:
|
|
432
|
+
repo = Repository(db)
|
|
433
|
+
session = repo.get_session(session_id)
|
|
434
|
+
if session is not None:
|
|
435
|
+
_record_request(
|
|
436
|
+
repo=repo,
|
|
437
|
+
session=session,
|
|
438
|
+
protocol=protocol,
|
|
439
|
+
endpoint=endpoint,
|
|
440
|
+
model=model,
|
|
441
|
+
req_fp=req_fp,
|
|
442
|
+
status=stream.status_code,
|
|
443
|
+
latency_ms=latency_ms,
|
|
444
|
+
tokens=tokens,
|
|
445
|
+
decision=final,
|
|
446
|
+
request_preview=request_preview,
|
|
447
|
+
response_preview=response_preview,
|
|
448
|
+
tool_calls=tool_calls,
|
|
449
|
+
error_fp=err_fp,
|
|
450
|
+
)
|
|
451
|
+
|
|
452
|
+
return StreamingResponse(
|
|
453
|
+
iterator(),
|
|
454
|
+
status_code=stream.status_code,
|
|
455
|
+
headers=response_headers(stream.headers, decision, request_id),
|
|
456
|
+
media_type=stream.media_type,
|
|
457
|
+
)
|
app/api/mcp_routes.py
ADDED
|
@@ -0,0 +1,218 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import json
|
|
4
|
+
from typing import Any
|
|
5
|
+
|
|
6
|
+
import httpx
|
|
7
|
+
from fastapi import APIRouter, Depends, HTTPException, Request
|
|
8
|
+
from fastapi.responses import JSONResponse, Response, StreamingResponse
|
|
9
|
+
from pydantic import BaseModel
|
|
10
|
+
from sqlalchemy.orm import Session
|
|
11
|
+
|
|
12
|
+
from app.core.security import bearer_token, filtered_upstream_headers
|
|
13
|
+
from app.db.repository import (
|
|
14
|
+
Repository,
|
|
15
|
+
mcp_approval_dict,
|
|
16
|
+
mcp_event_dict,
|
|
17
|
+
mcp_server_dict,
|
|
18
|
+
)
|
|
19
|
+
from app.db.session import get_db
|
|
20
|
+
from app.mcp.gateway import MCPGateway, jsonrpc_error, mock_response
|
|
21
|
+
from app.mcp.policy import MCPPolicyEngine, validate_policy
|
|
22
|
+
|
|
23
|
+
router = APIRouter()
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
class ApprovalDecision(BaseModel):
|
|
27
|
+
action: str
|
|
28
|
+
scope: str = "once"
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
class PolicyValidation(BaseModel):
|
|
32
|
+
policy: dict[str, Any]
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
def _authenticate(repo: Repository, request: Request) -> None:
|
|
36
|
+
if repo.agent_for_key(bearer_token(request)) is None:
|
|
37
|
+
raise HTTPException(403, "Invalid gateway key.")
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
def _valid_origin(request: Request) -> bool:
|
|
41
|
+
origin = request.headers.get("origin")
|
|
42
|
+
if not origin:
|
|
43
|
+
return True
|
|
44
|
+
for allowed in request.app.state.config.mcp_allowed_origins:
|
|
45
|
+
if origin == allowed or origin.startswith(allowed.rstrip("/") + ":"):
|
|
46
|
+
return True
|
|
47
|
+
return False
|
|
48
|
+
|
|
49
|
+
|
|
50
|
+
def _gateway(request: Request, db: Session, server_id: str) -> tuple[Repository, MCPGateway]:
|
|
51
|
+
repo = Repository(db)
|
|
52
|
+
server = repo.mcp_server(server_id)
|
|
53
|
+
if server is None or not server.enabled:
|
|
54
|
+
raise HTTPException(404, "MCP server not found.")
|
|
55
|
+
policy = MCPPolicyEngine(request.app.state.config.mcp_policy_path)
|
|
56
|
+
session_id = request.headers.get("mcp-session-id")
|
|
57
|
+
gateway = MCPGateway(
|
|
58
|
+
repo,
|
|
59
|
+
policy,
|
|
60
|
+
server_id,
|
|
61
|
+
session_id=session_id,
|
|
62
|
+
approval_timeout_seconds=request.app.state.config.mcp_approval_timeout_seconds,
|
|
63
|
+
)
|
|
64
|
+
return repo, gateway
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
@router.api_route("/mcp/{server_id}", methods=["GET", "DELETE"])
|
|
68
|
+
async def mcp_stream(server_id: str, request: Request, db: Session = Depends(get_db)):
|
|
69
|
+
repo, gateway = _gateway(request, db, server_id)
|
|
70
|
+
_authenticate(repo, request)
|
|
71
|
+
if not _valid_origin(request):
|
|
72
|
+
return Response(status_code=403)
|
|
73
|
+
if request.method == "DELETE":
|
|
74
|
+
session_id = request.headers.get("mcp-session-id")
|
|
75
|
+
if not session_id:
|
|
76
|
+
return Response(status_code=404)
|
|
77
|
+
server = repo.mcp_server(server_id)
|
|
78
|
+
assert server is not None
|
|
79
|
+
upstream_session = request.app.state.mcp_upstream_sessions.pop(session_id, None)
|
|
80
|
+
if server.transport != "mock" and not server.target.startswith("mock://"):
|
|
81
|
+
headers = filtered_upstream_headers(dict(request.headers))
|
|
82
|
+
if upstream_session:
|
|
83
|
+
headers["mcp-session-id"] = upstream_session
|
|
84
|
+
async with httpx.AsyncClient(timeout=15) as client:
|
|
85
|
+
await client.delete(server.target, headers=headers)
|
|
86
|
+
if repo.end_mcp_session(session_id) is None:
|
|
87
|
+
return Response(status_code=404)
|
|
88
|
+
return Response(status_code=204)
|
|
89
|
+
|
|
90
|
+
server = repo.mcp_server(server_id)
|
|
91
|
+
assert server is not None
|
|
92
|
+
if server.transport != "mock" and not server.target.startswith("mock://"):
|
|
93
|
+
headers = filtered_upstream_headers(dict(request.headers))
|
|
94
|
+
if gateway.session_id:
|
|
95
|
+
upstream_session = request.app.state.mcp_upstream_sessions.get(gateway.session_id)
|
|
96
|
+
if upstream_session:
|
|
97
|
+
headers["mcp-session-id"] = upstream_session
|
|
98
|
+
|
|
99
|
+
async def upstream_events():
|
|
100
|
+
async with httpx.AsyncClient(timeout=None) as client:
|
|
101
|
+
async with client.stream("GET", server.target, headers=headers) as upstream:
|
|
102
|
+
async for chunk in upstream.aiter_bytes():
|
|
103
|
+
yield chunk
|
|
104
|
+
|
|
105
|
+
return StreamingResponse(upstream_events(), media_type="text/event-stream")
|
|
106
|
+
|
|
107
|
+
async def ping():
|
|
108
|
+
yield "event: ping\ndata:\n\n"
|
|
109
|
+
|
|
110
|
+
return StreamingResponse(ping(), media_type="text/event-stream")
|
|
111
|
+
|
|
112
|
+
|
|
113
|
+
@router.post("/mcp/{server_id}")
|
|
114
|
+
async def mcp_post(server_id: str, request: Request, db: Session = Depends(get_db)):
|
|
115
|
+
repo, gateway = _gateway(request, db, server_id)
|
|
116
|
+
_authenticate(repo, request)
|
|
117
|
+
if not _valid_origin(request):
|
|
118
|
+
return Response(status_code=403)
|
|
119
|
+
try:
|
|
120
|
+
message = await request.json()
|
|
121
|
+
except (json.JSONDecodeError, UnicodeDecodeError):
|
|
122
|
+
return JSONResponse(jsonrpc_error(None, -32700, "Parse error"), status_code=400)
|
|
123
|
+
if not isinstance(message, dict) or message.get("jsonrpc") != "2.0":
|
|
124
|
+
return JSONResponse(jsonrpc_error(message.get("id") if isinstance(message, dict) else None, -32600, "Invalid Request"), status_code=400)
|
|
125
|
+
|
|
126
|
+
method = message.get("method")
|
|
127
|
+
if method != "initialize" and not gateway.session_id:
|
|
128
|
+
return JSONResponse(jsonrpc_error(message.get("id"), -32001, "MCP-Session-Id is required"), status_code=400)
|
|
129
|
+
if method == "initialize":
|
|
130
|
+
params = message.get("params") or {}
|
|
131
|
+
client = params.get("clientInfo") or {}
|
|
132
|
+
gateway.ensure_session(
|
|
133
|
+
client_name=str(client.get("name") or "unknown"),
|
|
134
|
+
protocol_version=str(params.get("protocolVersion") or "2025-11-25"),
|
|
135
|
+
)
|
|
136
|
+
|
|
137
|
+
interception = gateway.intercept(message)
|
|
138
|
+
if not interception.forward:
|
|
139
|
+
return JSONResponse(interception.response or {}, headers={"MCP-Session-Id": gateway.session_id or ""})
|
|
140
|
+
|
|
141
|
+
server = repo.mcp_server(server_id)
|
|
142
|
+
assert server is not None
|
|
143
|
+
if server.transport == "mock" or server.target.startswith("mock://"):
|
|
144
|
+
response_message = mock_response(interception.message)
|
|
145
|
+
if response_message is None:
|
|
146
|
+
return Response(status_code=202)
|
|
147
|
+
if method == "tools/list":
|
|
148
|
+
response_message = gateway.filter_tools(response_message)
|
|
149
|
+
return JSONResponse(response_message, headers={"MCP-Session-Id": gateway.session_id or ""})
|
|
150
|
+
|
|
151
|
+
headers = filtered_upstream_headers(dict(request.headers))
|
|
152
|
+
headers["accept"] = "application/json, text/event-stream"
|
|
153
|
+
if gateway.session_id and method != "initialize":
|
|
154
|
+
upstream_session = request.app.state.mcp_upstream_sessions.get(gateway.session_id)
|
|
155
|
+
if upstream_session:
|
|
156
|
+
headers["mcp-session-id"] = upstream_session
|
|
157
|
+
async with httpx.AsyncClient(timeout=60) as client:
|
|
158
|
+
upstream = await client.post(server.target, json=interception.message, headers=headers)
|
|
159
|
+
content_type = upstream.headers.get("content-type", "application/json")
|
|
160
|
+
if "application/json" not in content_type:
|
|
161
|
+
return Response(upstream.content, status_code=upstream.status_code, media_type=content_type)
|
|
162
|
+
response_message = upstream.json()
|
|
163
|
+
if method == "initialize" and gateway.session_id:
|
|
164
|
+
upstream_session = upstream.headers.get("mcp-session-id")
|
|
165
|
+
if upstream_session:
|
|
166
|
+
request.app.state.mcp_upstream_sessions[gateway.session_id] = upstream_session
|
|
167
|
+
if method == "tools/list" and isinstance(response_message, dict):
|
|
168
|
+
response_message = gateway.filter_tools(response_message)
|
|
169
|
+
response_headers = {}
|
|
170
|
+
if gateway.session_id:
|
|
171
|
+
response_headers["MCP-Session-Id"] = gateway.session_id
|
|
172
|
+
return JSONResponse(response_message, status_code=upstream.status_code, headers=response_headers)
|
|
173
|
+
|
|
174
|
+
|
|
175
|
+
@router.get("/api/v1/mcp/servers")
|
|
176
|
+
def mcp_servers(db: Session = Depends(get_db)) -> list[dict[str, Any]]:
|
|
177
|
+
return [mcp_server_dict(item) for item in Repository(db).list_mcp_servers()]
|
|
178
|
+
|
|
179
|
+
|
|
180
|
+
@router.get("/api/v1/mcp/events")
|
|
181
|
+
def mcp_events(limit: int = 200, db: Session = Depends(get_db)) -> list[dict[str, Any]]:
|
|
182
|
+
return [mcp_event_dict(item) for item in Repository(db).list_mcp_events(min(limit, 1000))]
|
|
183
|
+
|
|
184
|
+
|
|
185
|
+
@router.get("/api/v1/mcp/approvals")
|
|
186
|
+
def mcp_approvals(pending: bool = False, db: Session = Depends(get_db)) -> list[dict[str, Any]]:
|
|
187
|
+
return [mcp_approval_dict(item) for item in Repository(db).list_mcp_approvals(pending)]
|
|
188
|
+
|
|
189
|
+
|
|
190
|
+
@router.post("/api/v1/mcp/approvals/{approval_id}")
|
|
191
|
+
def decide_mcp_approval(
|
|
192
|
+
approval_id: str, payload: ApprovalDecision, db: Session = Depends(get_db)
|
|
193
|
+
) -> dict[str, Any]:
|
|
194
|
+
if payload.action not in {"allow", "deny"}:
|
|
195
|
+
raise HTTPException(422, "action must be allow or deny")
|
|
196
|
+
approval = Repository(db).decide_mcp_approval(approval_id, payload.action, payload.scope)
|
|
197
|
+
if approval is None:
|
|
198
|
+
raise HTTPException(404, "Approval is missing, expired, or already decided.")
|
|
199
|
+
return mcp_approval_dict(approval)
|
|
200
|
+
|
|
201
|
+
|
|
202
|
+
@router.post("/api/v1/mcp/policies/validate")
|
|
203
|
+
def validate_mcp_policy(payload: PolicyValidation) -> dict[str, Any]:
|
|
204
|
+
errors = validate_policy(payload.policy)
|
|
205
|
+
return {"valid": not errors, "errors": errors}
|
|
206
|
+
|
|
207
|
+
|
|
208
|
+
@router.get("/api/v1/mcp/export")
|
|
209
|
+
def export_mcp_events(db: Session = Depends(get_db)) -> JSONResponse:
|
|
210
|
+
repo = Repository(db)
|
|
211
|
+
return JSONResponse(
|
|
212
|
+
{
|
|
213
|
+
"schema_version": "mcp.audit.v1",
|
|
214
|
+
"servers": [mcp_server_dict(item) for item in repo.list_mcp_servers()],
|
|
215
|
+
"events": [mcp_event_dict(item) for item in repo.list_mcp_events(10000)],
|
|
216
|
+
"approvals": [mcp_approval_dict(item) for item in repo.list_mcp_approvals()],
|
|
217
|
+
}
|
|
218
|
+
)
|
app/api/openai_routes.py
ADDED
|
@@ -0,0 +1,26 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from fastapi import APIRouter, Depends, Request
|
|
4
|
+
from fastapi.responses import Response
|
|
5
|
+
from sqlalchemy.orm import Session
|
|
6
|
+
|
|
7
|
+
from app.api.common import guarded_proxy
|
|
8
|
+
from app.db.session import get_db
|
|
9
|
+
|
|
10
|
+
router = APIRouter()
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
@router.post("/v1/responses")
|
|
14
|
+
async def responses(request: Request, db: Session = Depends(get_db)) -> Response:
|
|
15
|
+
return await guarded_proxy(request, db, "openai", "/v1/responses")
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
@router.post("/v1/chat/completions")
|
|
19
|
+
async def chat_completions(request: Request, db: Session = Depends(get_db)) -> Response:
|
|
20
|
+
return await guarded_proxy(request, db, "openai", "/v1/chat/completions")
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
@router.get("/v1/models")
|
|
24
|
+
async def models(request: Request, db: Session = Depends(get_db)) -> Response:
|
|
25
|
+
return await guarded_proxy(request, db, "openai", "/v1/models")
|
|
26
|
+
|