99 lines
4.3 KiB
Python
99 lines
4.3 KiB
Python
"""Persistent MCP context and receipt storage in the selected project's BizDB."""
|
|
|
|
from datetime import datetime, timezone
|
|
from psycopg import AsyncConnection
|
|
from psycopg.types.json import Jsonb
|
|
from app.mcp.models import BoundContext
|
|
|
|
|
|
class McpQueryRepository:
|
|
def __init__(self, conn: AsyncConnection, project_id):
|
|
self.conn, self.project_id = conn, project_id
|
|
|
|
async def create(self, context: BoundContext):
|
|
async with self.conn.transaction(), self.conn.cursor() as cur:
|
|
# Bounded maintenance; expired/revoked handles remain invalid even if cleanup fails.
|
|
await cur.execute(
|
|
"""DELETE FROM workbench.mcp_query_contexts WHERE context_handle IN
|
|
(SELECT context_handle FROM workbench.mcp_query_contexts
|
|
WHERE project_id=%s AND (expires_at < now() OR revoked_at IS NOT NULL)
|
|
ORDER BY expires_at LIMIT 100 FOR UPDATE SKIP LOCKED)""",
|
|
(self.project_id,),
|
|
)
|
|
await cur.execute(
|
|
"""INSERT INTO workbench.mcp_query_contexts
|
|
(context_handle,project_id,subject,client_id,document,expires_at)
|
|
VALUES (%s,%s,%s,%s,%s,%s)""",
|
|
(
|
|
context.context_handle,
|
|
self.project_id,
|
|
context.subject,
|
|
context.client_id,
|
|
Jsonb(context.model_dump(mode="json")),
|
|
context.expires_at,
|
|
),
|
|
)
|
|
|
|
async def get(self, handle: str) -> BoundContext | None:
|
|
async with self.conn.transaction(), self.conn.cursor() as cur:
|
|
await cur.execute(
|
|
"""SELECT document,revoked_at FROM workbench.mcp_query_contexts
|
|
WHERE context_handle=%s AND project_id=%s""",
|
|
(handle, self.project_id),
|
|
)
|
|
row = await cur.fetchone()
|
|
if not row:
|
|
return None
|
|
document = dict(row["document"])
|
|
document["revoked_at"] = row["revoked_at"]
|
|
return BoundContext.model_validate(document)
|
|
|
|
async def revoke(self, handle: str, subject: str, client_id: str):
|
|
async with self.conn.transaction(), self.conn.cursor() as cur:
|
|
await cur.execute(
|
|
"""UPDATE workbench.mcp_query_contexts SET revoked_at=COALESCE(revoked_at,%s)
|
|
WHERE context_handle=%s AND project_id=%s AND subject=%s AND client_id=%s""",
|
|
(
|
|
datetime.now(timezone.utc),
|
|
handle,
|
|
self.project_id,
|
|
subject,
|
|
client_id,
|
|
),
|
|
)
|
|
|
|
async def receipt(self, scope_key: str, call_id: str) -> dict | None:
|
|
async with self.conn.transaction(), self.conn.cursor() as cur:
|
|
await cur.execute(
|
|
"""SELECT parameters,result FROM workbench.mcp_query_receipts
|
|
WHERE project_id=%s AND scope_key=%s AND call_id=%s""",
|
|
(self.project_id, scope_key, call_id),
|
|
)
|
|
return await cur.fetchone()
|
|
|
|
async def save_receipt(
|
|
self, scope_key: str, call_id: str, parameters: dict, result: dict
|
|
):
|
|
# Caller owns the short evidence + receipt transaction.
|
|
async with self.conn.cursor() as cur:
|
|
await cur.execute(
|
|
"""INSERT INTO workbench.mcp_query_receipts
|
|
(project_id,scope_key,call_id,parameters,result) VALUES (%s,%s,%s,%s,%s)
|
|
ON CONFLICT (project_id,scope_key,call_id) DO NOTHING""",
|
|
(self.project_id, scope_key, call_id, Jsonb(parameters), Jsonb(result)),
|
|
)
|
|
|
|
async def lock_receipt(self, scope_key: str, call_id: str):
|
|
async with self.conn.cursor() as cur:
|
|
await cur.execute(
|
|
"SELECT pg_advisory_xact_lock(hashtextextended(%s,0))",
|
|
(f"mcp:{self.project_id}:{scope_key}:{call_id}",),
|
|
)
|
|
|
|
async def finish_receipt(self, scope_key: str, call_id: str, result: dict):
|
|
async with self.conn.cursor() as cur:
|
|
await cur.execute(
|
|
"UPDATE workbench.mcp_query_receipts SET result=%s WHERE project_id=%s AND scope_key=%s AND call_id=%s",
|
|
(Jsonb(result), self.project_id, scope_key, call_id),
|
|
)
|