Files
next-tjwater-server/app/mcp/contexts.py
T
jiang be6317eb48 feat(server): integrate MCP queries and investigation messaging
Add project-bound OAuth MCP queries and controlled analysis with scoped evidence,
limits, timeouts and a nearby-model spatial query using EPSG:4326 inputs.
Default nearby searches to 200 meters and five results.

Integrate investigation messaging/presentation, health/readiness and typed history;
remove duplicate web search and fix online simulation and compressed baseline writes.
Include project migrations 023-025; keep local credentials and demo seeding excluded.

Validation: 319 focused pytest checks; real PostGIS temporary-table regression;
read-only Jinqian bridge radius comparison; cross-language MCP discovery/calls.
2026-10-19 18:46:58 +08:00

113 lines
4.5 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),
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