diff --git a/backend/druks/accounts/constants.py b/backend/druks/accounts/constants.py index 4e0431bb..de2a95ed 100644 --- a/backend/druks/accounts/constants.py +++ b/backend/druks/accounts/constants.py @@ -5,6 +5,14 @@ # is the row's lookup key and the only part that may appear in errors, logs, or # lists; the secret exists only inside the copy-once plaintext. PAT_TOKEN_TAG = "druks_pat" +# A call-scoped operator credential serializes as druks_call_. It is +# Redis-only: the plaintext never sits in Settings, and revoke drops it when +# the agent call ends. +OPERATOR_TOKEN_TAG = "druks_call" +OPERATOR_TOKEN_PREFIX = "operator_token:" +OPERATOR_TOKEN_CALL_PREFIX = "operator_token_call:" +OPERATOR_DEFERRED_PREFIX = "operator_deferred:" +OPERATOR_WRITES = frozenset({"deny", "defer", "allow"}) PAT_NAME_LENGTH = 80 PAT_PREFIX_LENGTH = 12 # No separator characters, so the serialized token splits unambiguously on "_". diff --git a/backend/druks/accounts/dependencies.py b/backend/druks/accounts/dependencies.py index 1cd84872..65005d0b 100644 --- a/backend/druks/accounts/dependencies.py +++ b/backend/druks/accounts/dependencies.py @@ -11,7 +11,7 @@ InvalidPatError, ) from druks.accounts.jwt import verify_assertion -from druks.accounts.models import Account, PersonalAccessToken +from druks.accounts.models import Account, OperatorToken, PersonalAccessToken _BEARER_CHALLENGE = 'Bearer realm="druks"' # auto_error=False: absence and malformed both come back None — presence is @@ -24,7 +24,7 @@ async def resolve_pat_account(credentials: HTTPAuthorizationCredentials | None) """A present Authorization must authenticate — never a fall-through.""" if credentials: try: - return (await PersonalAccessToken.authenticate(credentials.credentials)).account + return await resolve_bearer_account(credentials.credentials) except InvalidPatError as error: raise HTTPException( status_code=401, @@ -38,6 +38,17 @@ async def resolve_pat_account(credentials: HTTPAuthorizationCredentials | None) ) +async def resolve_bearer_account(credential: str) -> Account: + """The one bearer door: a live call-scoped operator token, else a PAT.""" + operator = await OperatorToken.lookup(credential) + if operator: + account = await Account.get(operator.account_id) + if account: + return account + raise InvalidPatError("Not a recognized operator token.") + return (await PersonalAccessToken.authenticate(credential)).account + + async def resolve_single_operator() -> Account | None: """None while zero accounts exist (setup); more than one refuses rather than guesses.""" diff --git a/backend/druks/accounts/models.py b/backend/druks/accounts/models.py index b7c5536b..f8ec618b 100644 --- a/backend/druks/accounts/models.py +++ b/backend/druks/accounts/models.py @@ -1,14 +1,23 @@ import base64 import hashlib import hmac +import json import secrets +from dataclasses import dataclass from datetime import datetime +from typing import Any +import httpx2 from sqlalchemy import ForeignKey, Index, LargeBinary, String, select, text from sqlalchemy.dialects.postgresql import CITEXT, insert from sqlalchemy.orm import Mapped, mapped_column, relationship from druks.accounts.constants import ( + OPERATOR_DEFERRED_PREFIX, + OPERATOR_TOKEN_CALL_PREFIX, + OPERATOR_TOKEN_PREFIX, + OPERATOR_TOKEN_TAG, + OPERATOR_WRITES, PAT_LAST_USED_RESOLUTION, PAT_LIFETIME, PAT_NAME_LENGTH, @@ -18,9 +27,11 @@ PAT_TOKEN_TAG, ) from druks.accounts.exceptions import AuthConfigurationError, InvalidPatError -from druks.core.models import Uuid7Pk +from druks.core.models import Uuid7Pk, uuid7_str from druks.database import db_session from druks.models import Base +from druks.redis import get_client +from druks.sandbox.constants import MAX_AGENT_TIMEOUT_SECONDS from druks.settings import load_settings from druks.user_settings.models import InstallationSettings @@ -196,3 +207,127 @@ async def revoke(self) -> None: # Keep the first revocation instant — a repeat revoke changes nothing. self.revoked_at = self.revoked_at or Base.utc_now() await db_session().flush() + + +def _hash_operator_token(token: str) -> str: + return hashlib.sha256(token.encode()).hexdigest() + + +_operator_api = None + + +@dataclass(frozen=True) +class OperatorToken: + """A call-scoped bearer with PAT authority. Redis holds it for the agent + call; Settings never does. ``writes`` is deny (read tools only), defer + (stash mutating calls), or allow (execute as ``account_id``).""" + + account_id: str + agent_call_id: str + run_id: str + writes: str + + @classmethod + def bind_api(cls, api: Any) -> None: + global _operator_api + _operator_api = api + + @classmethod + async def mint( + cls, + *, + account_id: str, + agent_call_id: str, + run_id: str, + writes: str, + ) -> str: + if writes not in OPERATOR_WRITES: + raise ValueError( + f"operator token writes must be one of {sorted(OPERATOR_WRITES)}, not {writes!r}" + ) + token = f"{OPERATOR_TOKEN_TAG}_{secrets.token_urlsafe(32)}" + payload = json.dumps( + { + "account_id": account_id, + "agent_call_id": agent_call_id, + "run_id": run_id, + "writes": writes, + } + ) + digest = _hash_operator_token(token) + redis = get_client() + await redis.set(f"{OPERATOR_TOKEN_PREFIX}{digest}", payload, ex=MAX_AGENT_TIMEOUT_SECONDS) + await redis.set( + f"{OPERATOR_TOKEN_CALL_PREFIX}{agent_call_id}", digest, ex=MAX_AGENT_TIMEOUT_SECONDS + ) + return token + + @classmethod + async def lookup(cls, credential: str) -> "OperatorToken | None": + if not credential.startswith(f"{OPERATOR_TOKEN_TAG}_"): + return + raw = await get_client().get(f"{OPERATOR_TOKEN_PREFIX}{_hash_operator_token(credential)}") + if raw: + return cls(**json.loads(raw)) + return + + @classmethod + async def authenticate(cls, credential: str) -> "OperatorToken": + found = await cls.lookup(credential) + if found: + return found + raise InvalidPatError("Not a recognized operator token.") + + @classmethod + async def revoke(cls, agent_call_id: str) -> None: + redis = get_client() + call_key = f"{OPERATOR_TOKEN_CALL_PREFIX}{agent_call_id}" + digest = await redis.get(call_key) + if digest: + await redis.delete(f"{OPERATOR_TOKEN_PREFIX}{digest.decode()}", call_key) + + @classmethod + async def defer_write(cls, run_id: str, write: dict[str, str]) -> None: + redis = get_client() + key = f"{OPERATOR_DEFERRED_PREFIX}{run_id}" + await redis.rpush(key, json.dumps(write)) + await redis.expire(key, MAX_AGENT_TIMEOUT_SECONDS) + + @classmethod + async def take_deferred(cls, run_id: str) -> list[dict[str, str]]: + redis = get_client() + key = f"{OPERATOR_DEFERRED_PREFIX}{run_id}" + items = await redis.lrange(key, 0, -1) + await redis.delete(key) + return [json.loads(item) for item in items] + + @classmethod + async def play_deferred(cls, account_id: str, writes: list[dict[str, str]]) -> None: + if not _operator_api: + raise RuntimeError("operator token replay needs the API bound at MCP boot") + call_id = uuid7_str() + token = await cls.mint( + account_id=account_id, agent_call_id=call_id, run_id=call_id, writes="allow" + ) + try: + async with httpx2.AsyncClient( + transport=httpx2.ASGITransport(app=_operator_api, raise_app_exceptions=False), + base_url="http://druks", + ) as client: + for write in writes: + response = await client.request( + write["method"], + write["path"], + content=write["body"] or None, + headers={ + "Authorization": f"Bearer {token}", + "Content-Type": write["content_type"], + }, + ) + if response.status_code >= 400: + raise RuntimeError( + f"deferred {write['method']} {write['path']} failed: " + f"{response.status_code} {response.text}" + ) + finally: + await cls.revoke(call_id) diff --git a/backend/druks/api/server.py b/backend/druks/api/server.py index e7e5c019..197d7b2c 100644 --- a/backend/druks/api/server.py +++ b/backend/druks/api/server.py @@ -13,6 +13,7 @@ from druks.accounts.dependencies import current_account, resolve_single_operator from druks.accounts.exceptions import AuthConfigurationError +from druks.accounts.models import OperatorToken from druks.accounts.routes import router as auth_router from druks.api.artifacts import router as artifacts_router from druks.api.dashboard import router as dashboard_router @@ -317,6 +318,7 @@ async def _unhandled_exception_handler( # A bare Route at exactly /mcp (a Mount would 307 the no-slash path); PATs # authenticate it, so it sits outside the identity gate. mcp_app = create_mcp_app(app) +OperatorToken.bind_api(app) app.router.routes.append( Route("/mcp", mcp_app, methods=["POST", "DELETE"], include_in_schema=False) ) diff --git a/backend/druks/contrib/chat/templates/talk.md b/backend/druks/contrib/chat/templates/talk.md index 4f7fd17b..4c6c240a 100644 --- a/backend/druks/contrib/chat/templates/talk.md +++ b/backend/druks/contrib/chat/templates/talk.md @@ -1,4 +1,6 @@ -Reply to the operator. +You are this operator. Act through the live Druks MCP catalog. + +Discover work with list_open_subjects. Echo parkedAt from get_gate unchanged when you answer a gate. Do not invent run ids. Autonomy: {{ autonomy }} diff --git a/backend/druks/contrib/chat/workflows.py b/backend/druks/contrib/chat/workflows.py index b0088e89..734afe90 100644 --- a/backend/druks/contrib/chat/workflows.py +++ b/backend/druks/contrib/chat/workflows.py @@ -1,9 +1,25 @@ +from dataclasses import dataclass from datetime import timedelta +from typing import TYPE_CHECKING, Any, Literal +from druks.accounts.models import OperatorToken from druks.contrib.chat.app import Chat from druks.contrib.chat.enums import Role from druks.contrib.chat.models import Conversation -from druks.workflows import Gate, Workflow, step +from druks.mcp.constants import THIS_APPLIANCE +from druks.mcp.helpers import get_bearer_token_env_var +from druks.sandbox.datastructures import McpServer +from druks.workflows import FatalError, Gate, Workflow, step +from druks.workspaces import Workspace, this_appliance_mcp_url + +if TYPE_CHECKING: + from druks.sandbox.host import Host + +_WRITES = { + "propose": "deny", + "confirm": "defer", + "full": "allow", +} class ChatTurn(Gate): @@ -15,12 +31,62 @@ class ChatTurn(Gate): stop: bool = False +class ConfirmTool(Gate): + """The operator approves or skips the mutating MCP call the agent proposed. + Parks after the agent step and before the next ChatTurn — a run holds one + gate at a time.""" + + name = "confirm_tool" + action: Literal["approve", "reject"] + + +@dataclass(frozen=True, kw_only=True) +class TalkWorkspace(Workspace): + account_id: str + run_id: str + writes: str + + async def with_mcp_servers(self, account_id: str | None, **kwargs: Any) -> dict[str, Any]: + # Call-scoped: minted here, not a vault row. The box cannot hold it — + # it dies with the agent call. + kwargs = await super().with_mcp_servers(account_id, **kwargs) + token = await OperatorToken.mint( + account_id=self.account_id, + agent_call_id=kwargs["call_id"], + run_id=self.run_id, + writes=self.writes, + ) + variable = get_bearer_token_env_var(THIS_APPLIANCE) + servers = [ + server for server in kwargs.get("mcp_servers") or () if server.name != THIS_APPLIANCE + ] + servers.append( + McpServer( + name=THIS_APPLIANCE, + url=this_appliance_mcp_url(self.host), + bearer_token_env_var=variable, + ) + ) + kwargs["mcp_servers"] = tuple(servers) + env = dict(kwargs.get("extra_env") or {}) + env[variable] = token + kwargs["extra_env"] = env + return kwargs + + async def run_agent(self, *, account_id: str | None, **kwargs: Any): + try: + return await super().run_agent(account_id=account_id, **kwargs) + finally: + await OperatorToken.revoke(kwargs["call_id"]) + + class Talk(Workflow): """One conversation: agent reply, park, operator line, repeat until stop.""" subject = Conversation steps_reuse_sandbox = True sandbox_hold = timedelta(minutes=15) + workspace_class = TalkWorkspace async def run_multistep(self) -> None: while True: @@ -31,6 +97,18 @@ async def run_multistep(self) -> None: messages=[{"role": message.role, "body": message.body} for message in messages], ) await self.record_message(Role.ASSISTANT, result.text) + deferred = await self.deferred_writes() + if deferred: + decision = await ConfirmTool.wait( + input_request={ + "presentation": "in_app", + "label": "Confirm the proposed action", + "controls": ["approve", "reject"], + }, + hold_sandbox=self.sandbox_hold, + ) + if decision.action == "approve": + await self.apply_deferred_writes(deferred) reply = await ChatTurn.wait( input_request={"presentation": "in_app", "label": "Chat turn"}, hold_sandbox=self.sandbox_hold, @@ -39,11 +117,32 @@ async def run_multistep(self) -> None: return await self.record_message(Role.USER, reply.text) + async def get_workspace_kwargs(self, host: "Host") -> dict[str, Any]: + conversation = await self.subject + if not self.account_id: + raise FatalError("Talk runs as the operator who started the conversation.") + return { + **await super().get_workspace_kwargs(host), + "account_id": self.account_id, + "run_id": self.workflow_id, + "writes": _WRITES[conversation.autonomy], + } + @step async def record_message(self, role: Role, body: str) -> None: conversation = await self.subject await conversation.add_message(role=role, body=body) + @step + async def deferred_writes(self) -> list[dict[str, str]]: + return await OperatorToken.take_deferred(self.workflow_id) + + @step + async def apply_deferred_writes(self, writes: list[dict[str, str]]) -> None: + if not self.account_id: + raise FatalError("Talk runs as the operator who started the conversation.") + await OperatorToken.play_deferred(self.account_id, writes) + @classmethod async def dispatch(cls, *, conversation: Conversation) -> str: # One Talk per conversation: start() already dedups live runs on the diff --git a/backend/druks/mcp/constants.py b/backend/druks/mcp/constants.py index 30195666..f2d05be5 100644 --- a/backend/druks/mcp/constants.py +++ b/backend/druks/mcp/constants.py @@ -4,6 +4,9 @@ # the var name ever lands in emitted config, never the value. TOKEN_ENV_PREFIX = "MCP_" TOKEN_ENV_SUFFIX = "_TOKEN" +# The name a workspace uses when it injects this appliance's /mcp into a +# sandbox — one config key, one bearer env var. +THIS_APPLIANCE = "druks" # A server name is one identifier reused as the MCP config key (a bare TOML path # segment for codex, a JSON object key for claude) and the stem of the bearer env diff --git a/backend/druks/mcp/server.py b/backend/druks/mcp/server.py index 74dfddf3..8486fe9b 100644 --- a/backend/druks/mcp/server.py +++ b/backend/druks/mcp/server.py @@ -3,21 +3,24 @@ # operation's single declaration — schema, docstring, operation_id — and a # tagged app route joins the surface the same way. import inspect -from collections.abc import Generator +from collections.abc import Generator, Sequence import httpx2 from fastapi import FastAPI from fastapi.routing import APIRoute, iter_route_contexts from fastmcp import FastMCP from fastmcp.server.auth import AccessToken, TokenVerifier -from fastmcp.server.dependencies import get_http_request +from fastmcp.server.dependencies import get_access_token, get_http_request from fastmcp.server.http import StarletteWithLifespan from fastmcp.server.providers.openapi import MCPType, OpenAPIProvider, OpenAPITool, RouteMap +from fastmcp.server.transforms import GetToolNext, Transform +from fastmcp.tools.base import Tool from fastmcp.utilities.openapi import HTTPRoute +from fastmcp.utilities.versions import VersionSpec from mcp.types import ToolAnnotations from druks.accounts.exceptions import InvalidPatError -from druks.accounts.models import PersonalAccessToken +from druks.accounts.models import OperatorToken, PersonalAccessToken from druks.apps.loader import iter_apps from druks.database import db_session from druks.mcp.exceptions import InvalidAgentToolError @@ -41,14 +44,28 @@ class PatTokenVerifier(TokenVerifier): async def verify_token(self, token: str) -> AccessToken | None: # Auth middleware runs outside the request session boundary, so this - # owns one — authenticate stamps last_used_at. + # owns one — authenticate stamps last_used_at on a PAT; a call token + # is Redis-only and needs no commit. try: + operator = await OperatorToken.lookup(token) + if operator: + return AccessToken( + token=token, + client_id=operator.agent_call_id, + scopes=[], + claims={ + "account_id": operator.account_id, + "agent_call_id": operator.agent_call_id, + "run_id": operator.run_id, + "writes": operator.writes, + }, + ) pat = await PersonalAccessToken.authenticate(token) access = AccessToken( token=token, client_id=pat.token_prefix, scopes=[], - claims={"account_id": pat.account_id, "pat_id": pat.id}, + claims={"account_id": pat.account_id, "pat_id": pat.id, "writes": "allow"}, ) await db_session().commit() return access @@ -75,6 +92,87 @@ def auth_flow( yield request +def _writes_claim() -> str: + token = get_access_token() + if token: + return token.claims.get("writes", "allow") + return "allow" + + +def _is_read_tool(tool: Tool) -> bool: + return bool(tool.annotations and tool.annotations.read_only_hint) + + +class OperatorWritesFilter(Transform): + """Propose (writes=deny) lists only GET-derived tools. Confirm and full + keep the live catalog; mutating calls are intercepted on the HTTP hop.""" + + async def list_tools(self, tools: Sequence[Tool]) -> Sequence[Tool]: + if _writes_claim() == "deny": + return [tool for tool in tools if _is_read_tool(tool)] + return tools + + async def get_tool( + self, name: str, call_next: GetToolNext, *, version: VersionSpec | None = None + ) -> Tool | None: + tool = await call_next(name, version=version) + if tool and _writes_claim() == "deny" and not _is_read_tool(tool): + return + return tool + + +def _bearer_credential(header: str | None) -> str: + if header: + scheme, _, credential = header.partition(" ") + if scheme.lower() == "bearer" and credential: + return credential + return "" + + +class OperatorWritesTransport(httpx2.AsyncBaseTransport): + """Mutating OpenAPI hops: deny 403s, defer stashes, allow passes through.""" + + def __init__(self, inner: httpx2.AsyncBaseTransport): + self._inner = inner + + async def handle_async_request(self, request: httpx2.Request) -> httpx2.Response: + if request.method.upper() != "GET": + operator = await OperatorToken.lookup( + _bearer_credential(request.headers.get("authorization")) + ) + if operator and operator.writes == "deny": + return httpx2.Response( + 403, + json={ + "code": "AUTONOMY_READ_ONLY", + "message": ( + "This conversation's autonomy is propose; mutating tools do not run." + ), + "retryable": False, + }, + request=request, + ) + if operator and operator.writes == "defer": + await OperatorToken.defer_write( + operator.run_id, + { + "method": request.method, + "path": request.url.path, + "body": request.content.decode() if request.content else "", + "content_type": request.headers.get("content-type") or "application/json", + }, + ) + return httpx2.Response( + 200, + json={ + "result": "deferred", + "message": "Proposed. The operator must confirm before this runs.", + }, + request=request, + ) + return await self._inner.handle_async_request(request) + + def _validate_agent_tools(api: FastAPI) -> None: # The provider logs component-fn errors instead of raising, so derived tools # cannot refuse boot; validate the routes before derivation. Inclusion is @@ -160,6 +258,11 @@ def _annotate(route: HTTPRoute, component: object) -> None: destructive_hint=not is_read and route.extensions.get("x-destructive", True), idempotent_hint=route.extensions.get("x-idempotent", False), ) + # Confirm (writes=defer) intercepts mutating hops with a deferred stub, + # not the route's 201 model. Advertising that model as outputSchema + # makes MCP reject the stub (`identifier` required on create_ticket). + if not is_read: + component.output_schema = None def create_mcp_app(api: FastAPI) -> StarletteWithLifespan: @@ -169,7 +272,9 @@ def create_mcp_app(api: FastAPI) -> StarletteWithLifespan: # raise_app_exceptions=False makes an app crash reach the tool as the # app's sanitized 500, so no masking is needed and the taxonomy travels. client = httpx2.AsyncClient( - transport=httpx2.ASGITransport(app=api, raise_app_exceptions=False), + transport=OperatorWritesTransport( + httpx2.ASGITransport(app=api, raise_app_exceptions=False) + ), base_url="http://druks", auth=CallerPat(), ) @@ -185,6 +290,7 @@ def create_mcp_app(api: FastAPI) -> StarletteWithLifespan: server = FastMCP( name="druks", providers=[provider], + transforms=[OperatorWritesFilter()], instructions=_INSTRUCTIONS, auth=PatTokenVerifier(), ) diff --git a/backend/druks/workspaces.py b/backend/druks/workspaces.py index c660b14b..dd394f14 100644 --- a/backend/druks/workspaces.py +++ b/backend/druks/workspaces.py @@ -5,7 +5,7 @@ from dataclasses import dataclass from pathlib import Path, PurePosixPath from typing import TYPE_CHECKING, Any, ClassVar -from urllib.parse import urlsplit +from urllib.parse import urlparse, urlsplit, urlunparse from druks.accounts.models import Account from druks.core.apis.github import get_github_client @@ -28,11 +28,30 @@ from druks.sandbox.exceptions import ExecFailed from druks.sandbox.layout import get_repo_root, get_work_root from druks.sandbox.models import SecretRef +from druks.settings import load_settings if TYPE_CHECKING: from druks.sandbox.host import Host +def this_appliance_mcp_url(host: "Host") -> str: + """The /mcp hop a sandbox uses to reach this process. + + Docker sibling containers cannot use the host loopback; the engine + publishes that address as host.docker.internal:8001. An exe VM is a + different machine and uses the dashboard URL (urls.endpoint), never + webhook_host. + """ + base = (load_settings().urls.endpoint or "http://127.0.0.1:8001").rstrip("/") + parsed = urlparse(base) + if host.record.provider == "docker" and parsed.hostname in {"127.0.0.1", "localhost", "::1"}: + port = parsed.port + if not port: + port = 8001 if parsed.scheme == "http" else 443 + base = urlunparse(parsed._replace(netloc=f"host.docker.internal:{port}")).rstrip("/") + return f"{base}/mcp" + + @dataclass(frozen=True) class Workspace: # What an agent runs in: the VM it abstracts. diff --git a/backend/tests/chat/test_operator_mcp.py b/backend/tests/chat/test_operator_mcp.py new file mode 100644 index 00000000..06d25c8e --- /dev/null +++ b/backend/tests/chat/test_operator_mcp.py @@ -0,0 +1,92 @@ +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +from druks.accounts.models import Account, OperatorToken +from druks.contrib.chat.enums import Autonomy +from druks.contrib.chat.models import Conversation +from druks.contrib.chat.workflows import Talk, TalkWorkspace +from druks.mcp.constants import THIS_APPLIANCE +from druks.mcp.helpers import get_bearer_token_env_var +from druks.workspaces import this_appliance_mcp_url + + +class _FakeHost: + ssh_username = "exedev" + + def __init__(self, provider: str = "docker"): + self.record = SimpleNamespace(provider=provider) + self.run_agent = AsyncMock(return_value="ok") + + +def test_docker_sandbox_uses_host_docker_internal_for_loopback(monkeypatch): + settings = MagicMock() + settings.urls.endpoint = "http://127.0.0.1:8001" + monkeypatch.setattr("druks.workspaces.load_settings", lambda: settings) + + assert this_appliance_mcp_url(_FakeHost("docker")) == "http://host.docker.internal:8001/mcp" + assert this_appliance_mcp_url(_FakeHost("exe.dev")) == "http://127.0.0.1:8001/mcp" + + +def test_appliance_mcp_url_never_uses_webhook_host(monkeypatch): + settings = MagicMock() + settings.urls.endpoint = "https://druks.example.com" + settings.urls.webhook_host = "hooks.example.com" + monkeypatch.setattr("druks.workspaces.load_settings", lambda: settings) + + assert this_appliance_mcp_url(_FakeHost("docker")) == "https://druks.example.com/mcp" + + +async def test_talk_workspace_injects_this_appliance_mcp(druks_db, monkeypatch): + account = await Account.get_or_create("op@example.com") + settings = MagicMock() + settings.urls.endpoint = "http://127.0.0.1:8001" + monkeypatch.setattr("druks.workspaces.load_settings", lambda: settings) + workspace = TalkWorkspace( + host=_FakeHost("docker"), # type: ignore[arg-type] + account_id=account.id, + run_id="run-1", + writes="deny", + ) + kwargs = await workspace.with_mcp_servers(account.id, call_id="call-1") + server = next(s for s in kwargs["mcp_servers"] if s.name == THIS_APPLIANCE) + token = kwargs["extra_env"][get_bearer_token_env_var(THIS_APPLIANCE)] + + assert server.url == "http://host.docker.internal:8001/mcp" + found = await OperatorToken.lookup(token) + assert found.account_id == account.id + assert found.writes == "deny" + assert found.agent_call_id == "call-1" + + +async def test_talk_workspace_revokes_the_token_after_the_call(druks_db, monkeypatch): + account = await Account.get_or_create("op@example.com") + settings = MagicMock() + settings.urls.endpoint = "http://127.0.0.1:8001" + monkeypatch.setattr("druks.workspaces.load_settings", lambda: settings) + host = _FakeHost("exe.dev") + workspace = TalkWorkspace( + host=host, # type: ignore[arg-type] + account_id=account.id, + run_id="run-1", + writes="allow", + ) + + await workspace.run_agent(account_id=account.id, call_id="call-done") + token = host.run_agent.await_args.kwargs["extra_env"][get_bearer_token_env_var(THIS_APPLIANCE)] + assert await OperatorToken.lookup(token) is None + + +async def test_talk_reads_live_autonomy_for_the_next_call(druks_db): + account = await Account.get_or_create("op@example.com") + conversation = await Conversation.create(account_id=account.id, title="") + flow = Talk() + flow.subject = conversation + flow.account_id = account.id + flow._workflow_id = "run-1" + + kwargs = await flow.get_workspace_kwargs(_FakeHost()) # type: ignore[arg-type] + assert kwargs["writes"] == "deny" + + conversation.autonomy = Autonomy.FULL + kwargs = await flow.get_workspace_kwargs(_FakeHost()) # type: ignore[arg-type] + assert kwargs["writes"] == "allow" diff --git a/backend/tests/chat/test_workflows.py b/backend/tests/chat/test_workflows.py index 4c25db4f..98a866e0 100644 --- a/backend/tests/chat/test_workflows.py +++ b/backend/tests/chat/test_workflows.py @@ -6,13 +6,14 @@ from druks.contrib.chat.contracts import TurnOutput from druks.contrib.chat.enums import Role from druks.contrib.chat.models import Conversation -from druks.contrib.chat.workflows import ChatTurn, Talk +from druks.contrib.chat.workflows import ChatTurn, ConfirmTool, Talk from druks.workflows import current_workflow async def _run_talk(conversation: Conversation) -> None: flow = Talk() flow.subject = conversation + flow.account_id = conversation.account_id token = current_workflow.set(flow) try: await flow.run_multistep() @@ -46,6 +47,7 @@ async def wait(cls, **kwargs): monkeypatch.setattr(ChatTurn, "wait", classmethod(wait)) monkeypatch.setattr(Talk, "record_message", Talk.record_message.__wrapped__) + monkeypatch.setattr(Talk, "deferred_writes", mock.AsyncMock(return_value=[])) await _run_talk(conversation) @@ -77,6 +79,7 @@ async def wait(cls, **kwargs): monkeypatch.setattr(ChatTurn, "wait", classmethod(wait)) monkeypatch.setattr(Talk, "record_message", Talk.record_message.__wrapped__) + monkeypatch.setattr(Talk, "deferred_writes", mock.AsyncMock(return_value=[])) await _run_talk(conversation) @@ -90,3 +93,74 @@ async def wait(cls, **kwargs): ] assert turns[0]["autonomy"] == conversation.autonomy assert turns[1]["messages"][-1] == {"role": Role.USER, "body": "and then?"} + + +async def test_talk_confirms_deferred_writes_before_the_next_line(druks_db, monkeypatch): + account = await Account.get_or_create("op@example.com") + conversation = await Conversation.create(account_id=account.id, title="") + await conversation.add_message(role=Role.USER, body="hello") + + monkeypatch.setattr(Chat, "reply", mock.AsyncMock(return_value=TurnOutput(text="hi"))) + monkeypatch.setattr(Talk, "record_message", Talk.record_message.__wrapped__) + proposed = [ + { + "method": "POST", + "path": "/api/gates/run-1/answer", + "body": "{}", + "content_type": "application/json", + } + ] + monkeypatch.setattr(Talk, "deferred_writes", mock.AsyncMock(return_value=proposed)) + applied = [] + + async def apply(self, writes): + applied.extend(writes) + + monkeypatch.setattr(Talk, "apply_deferred_writes", apply) + parked = [] + + async def confirm_wait(cls, **kwargs): + parked.append(kwargs) + return ConfirmTool(action="approve") + + async def turn_wait(cls, **kwargs): + return ChatTurn(text="", stop=True) + + monkeypatch.setattr(ConfirmTool, "wait", classmethod(confirm_wait)) + monkeypatch.setattr(ChatTurn, "wait", classmethod(turn_wait)) + + await _run_talk(conversation) + + assert parked[0]["input_request"]["controls"] == ["approve", "reject"] + assert applied == proposed + + +async def test_talk_skips_deferred_writes_when_the_operator_rejects(druks_db, monkeypatch): + account = await Account.get_or_create("op@example.com") + conversation = await Conversation.create(account_id=account.id, title="") + await conversation.add_message(role=Role.USER, body="hello") + + monkeypatch.setattr(Chat, "reply", mock.AsyncMock(return_value=TurnOutput(text="hi"))) + monkeypatch.setattr(Talk, "record_message", Talk.record_message.__wrapped__) + monkeypatch.setattr( + Talk, + "deferred_writes", + mock.AsyncMock( + return_value=[{"method": "POST", "path": "/x", "body": "", "content_type": ""}] + ), + ) + apply = mock.AsyncMock() + monkeypatch.setattr(Talk, "apply_deferred_writes", apply) + + async def confirm_wait(cls, **kwargs): + return ConfirmTool(action="reject") + + async def turn_wait(cls, **kwargs): + return ChatTurn(text="", stop=True) + + monkeypatch.setattr(ConfirmTool, "wait", classmethod(confirm_wait)) + monkeypatch.setattr(ChatTurn, "wait", classmethod(turn_wait)) + + await _run_talk(conversation) + + apply.assert_not_awaited() diff --git a/backend/tests/test_mcp_endpoint.py b/backend/tests/test_mcp_endpoint.py index 4f99a03f..b183b175 100644 --- a/backend/tests/test_mcp_endpoint.py +++ b/backend/tests/test_mcp_endpoint.py @@ -8,7 +8,7 @@ import httpx2 import pytest from conftest import finish_agent_run, make_test_note, seed_note_agent_run, seed_note_run -from druks.accounts.models import Account, PersonalAccessToken +from druks.accounts.models import Account, OperatorToken, PersonalAccessToken from druks.api.server import mcp_app from druks.contrib.software_factory.app import SoftwareFactory from druks.core.apis.exceptions import UnknownTicketError @@ -20,6 +20,7 @@ from fastapi import APIRouter, FastAPI from fastmcp.client import Client from fastmcp.client.transports import StreamableHttpTransport +from pydantic import BaseModel from starlette.routing import Route _IN_APP_ASK = { @@ -272,6 +273,43 @@ async def software_factory_scan(): create_mcp_app(api) +class _Issued(BaseModel): + identifier: str + + +def _create_shaped_app() -> FastAPI: + # A 201 whose model requires identifier — the confirm stub does not have one. + held: dict[str, object] = {} + + @asynccontextmanager + async def lifespan(scope_app): + async with held["mcp"].lifespan(scope_app): + yield + + api = FastAPI(lifespan=lifespan) + router = APIRouter() + + async def mint_ticket() -> _Issued: + """Write a ticket down.""" + return _Issued(identifier="BOX-1") + + router.add_api_route( + "/tickets", + mint_ticket, + methods=["POST"], + status_code=201, + operation_id="mint_ticket", + tags=["agent"], + ) + api.include_router(router, prefix="/api/software_factory", tags=["software_factory"]) + mcp = create_mcp_app(api) + held["mcp"] = mcp + api.router.routes.append( + Route("/mcp", mcp, methods=["POST", "DELETE"], include_in_schema=False) + ) + return api + + def _agent_route_app(operation_id: str) -> FastAPI: # A synthetic agent route owned by the installed software_factory app — the # loader-stamped app tag names the owner, exactly as a real router does. @@ -555,3 +593,158 @@ async def test_mcp_server_registry_routes_stay_untouched(tmp_path, druks_db, mon assert listed.status_code == 200 # The inbound endpoint never joins the outbound server registry. assert "druks" not in {server["name"] for server in listed.json()} + + +async def _operator_token(account, *, writes: str, run_id: str = "run-op"): + return await OperatorToken.mint( + account_id=account.id, + agent_call_id=f"call-{writes}", + run_id=run_id, + writes=writes, + ) + + +async def test_propose_cannot_hit_answer_gate(app, account, druks_db, resume_spy): + item = await make_test_note() + run = await seed_note_run( + druks_db, + note=item, + state="parked", + input_gate="review", + input_request=dict(_IN_APP_ASK), + ) + run.input_requested_at = datetime.now(UTC) + await druks_db.flush() + token = await _operator_token(account, writes="deny") + + async with live(app), _client(app, token) as client: + tools = {tool.name for tool in await client.list_tools()} + assert "list_open_subjects" in tools + assert "get_gate" in tools + assert "get_agent_call" in tools + assert "get_usage" in tools + assert "answer_gate" not in tools + assert "cancel_run" not in tools + result = await client.call_tool( + "answer_gate", + {"run": run.id, "parkedAt": run.input_requested_at.isoformat(), "control": "approve"}, + raise_on_error=False, + ) + + assert result.is_error + assert resume_spy == [] + + +async def test_confirm_defers_answer_gate_until_play(app, account, druks_db, resume_spy): + item = await make_test_note() + run = await seed_note_run( + druks_db, + note=item, + state="parked", + input_gate="review", + input_request=dict(_IN_APP_ASK), + ) + run.input_requested_at = datetime.now(UTC) + await druks_db.flush() + token = await _operator_token(account, writes="defer", run_id="run-confirm") + + async with live(app), _client(app, token) as client: + tools = {tool.name for tool in await client.list_tools()} + assert "answer_gate" in tools + parked_at = run.input_requested_at.isoformat() + await client.call_tool( + "answer_gate", + {"run": run.id, "parkedAt": parked_at, "control": "approve"}, + raise_on_error=False, + ) + + assert resume_spy == [] + writes = await OperatorToken.take_deferred("run-confirm") + assert writes[0]["method"] == "POST" + await OperatorToken.play_deferred(account.id, writes) + assert resume_spy == [{"id": run.id, "action": "approve", "answers": {}, "note": ""}] + + +async def test_mutating_agent_tools_omit_output_schema(app, pat_token): + async with live(app), _client(app, pat_token) as client: + tools = {tool.name: tool for tool in await client.list_tools()} + assert tools["get_gate"].output_schema + assert tools["answer_gate"].output_schema is None + assert tools["cancel_run"].output_schema is None + + +async def test_confirm_defers_a_create_shaped_write(account): + token = await _operator_token(account, writes="defer", run_id="run-confirm-create") + api = _create_shaped_app() + + async with live(api), _client(api, token) as client: + result = await client.call_tool("software_factory_mint_ticket", {}, raise_on_error=False) + + if result.is_error: + raise AssertionError(result.content[0].text) + assert result.structured_content["result"] == "deferred" + writes = await OperatorToken.take_deferred("run-confirm-create") + assert writes[0]["method"] == "POST" + assert writes[0]["path"] == "/api/software_factory/tickets" + + +async def test_full_answer_gate_runs_as_the_token_account(app, account, druks_db, resume_spy): + item = await make_test_note() + run = await seed_note_run( + druks_db, + note=item, + state="parked", + input_gate="review", + input_request=dict(_IN_APP_ASK), + ) + run.input_requested_at = datetime.now(UTC) + await druks_db.flush() + token = await _operator_token(account, writes="allow") + + async with live(app), _client(app, token) as client: + gate = (await client.call_tool("get_gate", {"run": run.id})).structured_content + answered = ( + await client.call_tool( + "answer_gate", + {"run": run.id, "parkedAt": gate["parkedAt"], "control": "approve"}, + ) + ).structured_content + + assert answered["result"] == "answered" + assert resume_spy == [{"id": run.id, "action": "approve", "answers": {}, "note": ""}] + + +async def test_operator_token_cannot_act_as_another_account(app, druks_db): + mine = await Account.get_or_create("op@example.com") + theirs = await Account.get_or_create("peer@example.com") + druks_db.add( + UsageScrape( + provider="openai", + account_id=theirs.id, + scraped_at=datetime.now(UTC), + five_hour_percent_left=42, + ) + ) + await druks_db.flush() + token = await _operator_token(mine, writes="allow") + + async with live(app), _client(app, token) as client: + usage = (await client.call_tool("get_usage", {})).structured_content + codex = next(h for h in usage["providers"] if h["id"] == "openai") + assert codex["fiveHourPercentLeft"] is None + + +async def test_operator_token_is_gone_after_the_call(app, account): + token = await _operator_token(account, writes="allow", run_id="run-gone") + assert (await OperatorToken.lookup(token)).account_id == account.id + await OperatorToken.revoke("call-allow") + assert await OperatorToken.lookup(token) is None + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://druks.test" + ) as wire: + gone = await wire.post( + "/mcp", + json=_INIT, + headers={**_WIRE_HEADERS, "Authorization": f"Bearer {token}"}, + ) + assert gone.status_code == 401