115 lines
4.6 KiB
Python
115 lines
4.6 KiB
Python
"""An OAuth identity and a Server-owned binding are both required for data queries."""
|
|
|
|
import secrets
|
|
from datetime import datetime, timedelta, timezone
|
|
from uuid import UUID
|
|
from fastapi import HTTPException
|
|
from app.auth.permissions import (
|
|
permissions_for_context,
|
|
OPERATIONAL_TASK_EXECUTE,
|
|
OPERATIONAL_TASK_VIEW,
|
|
)
|
|
from app.infra.db.postgresql.operational_tasks import OperationalTaskRepository
|
|
from app.infra.db.postgresql.mcp_queries import McpQueryRepository
|
|
from app.mcp.models import BindInput, BoundContext
|
|
from app.mcp.scope import QueryScope
|
|
from app.mcp.config import McpSettings
|
|
from mcp.server.auth.provider import AccessToken
|
|
|
|
|
|
def task_repository(scope: QueryScope) -> OperationalTaskRepository:
|
|
return OperationalTaskRepository(scope.connection, scope.project.project_id)
|
|
|
|
|
|
async def verify_task(scope: QueryScope, binding: BoundContext) -> None:
|
|
permissions = permissions_for_context(scope.project)
|
|
if OPERATIONAL_TASK_VIEW not in permissions:
|
|
raise HTTPException(403, "Task read permission is required")
|
|
repo = task_repository(scope)
|
|
async with scope.connection.transaction():
|
|
task = await repo.get_task(str(binding.task_id))
|
|
if not task:
|
|
raise HTTPException(404, "Task is not in this project")
|
|
if binding.purpose == "task_collaboration":
|
|
return
|
|
if OPERATIONAL_TASK_EXECUTE not in permissions:
|
|
raise HTTPException(403, "Task execution permission is required")
|
|
run = await repo.get_active_run(str(binding.task_id))
|
|
now = datetime.now(timezone.utc)
|
|
owner = binding.parent_session_id or binding.native_session_id
|
|
if (
|
|
not run
|
|
or run["status"] not in {"claimed", "running"}
|
|
or run.get("agent_session_id") != owner
|
|
or not run.get("lease_expires_at")
|
|
or run["lease_expires_at"] <= now
|
|
or binding.run_id is not None
|
|
and str(run["run_id"]) != str(binding.run_id)
|
|
or binding.execution_request_id is not None
|
|
and str(run["execution_request_id"]) != str(binding.execution_request_id)
|
|
):
|
|
raise HTTPException(
|
|
409, "Investigation run is no longer active for this session"
|
|
)
|
|
binding.run_id = UUID(str(run["run_id"]))
|
|
binding.execution_request_id = UUID(str(run["execution_request_id"]))
|
|
|
|
|
|
async def bind_context(
|
|
scope: QueryScope, principal: AccessToken, command: BindInput, settings: McpSettings
|
|
) -> BoundContext:
|
|
if principal.client_id != settings.host_client_id:
|
|
raise HTTPException(
|
|
403, "Only the trusted product host may bind query contexts"
|
|
)
|
|
from app.mcp.catalog import TOOLS
|
|
|
|
binding = BoundContext(
|
|
context_handle=secrets.token_urlsafe(32),
|
|
project_id=scope.project.project_id,
|
|
subject=principal.subject,
|
|
client_id=principal.client_id,
|
|
tool_cap=sorted(name for name, tool in TOOLS.items()
|
|
if (command.purpose != "task_reader" or not (tool.computation or tool.mutating))
|
|
and (not tool.mutating or command.purpose == "ordinary_chat")),
|
|
permission_cap=sorted(permissions_for_context(scope.project)),
|
|
native_session_id=command.native_session_id,
|
|
purpose=command.purpose,
|
|
parent_session_id=command.parent_session_id,
|
|
task_id=command.task_id,
|
|
expires_at=datetime.now(timezone.utc)
|
|
+ timedelta(seconds=settings.context_ttl_seconds),
|
|
)
|
|
if binding.task_id:
|
|
await verify_task(scope, binding)
|
|
from app.infra.db.postgresql.read_queries import task_data_scopes
|
|
|
|
async with scope.connection.transaction():
|
|
binding.task_scope_cap = await task_data_scopes(scope.connection,
|
|
project_id=binding.project_id, task_id=binding.task_id,
|
|
permissions=permissions_for_context(scope.project))
|
|
await McpQueryRepository(scope.connection, scope.project.project_id).create(binding)
|
|
return binding
|
|
|
|
|
|
async def resolve_context(
|
|
scope: QueryScope, principal: AccessToken, handle: str
|
|
) -> BoundContext:
|
|
binding = await McpQueryRepository(scope.connection, scope.project.project_id).get(
|
|
handle
|
|
)
|
|
if (
|
|
not binding
|
|
or binding.revoked_at
|
|
or binding.expires_at <= datetime.now(timezone.utc)
|
|
or binding.subject != principal.subject
|
|
or binding.client_id != principal.client_id
|
|
or binding.project_id != scope.project.project_id
|
|
):
|
|
raise HTTPException(
|
|
403, "Query context is unavailable for this identity and project"
|
|
)
|
|
if binding.task_id:
|
|
await verify_task(scope, binding)
|
|
return binding
|