From 8e0f7ced5f260fb7a307acf7d0332ce6bd068f34 Mon Sep 17 00:00:00 2001 From: Alex Date: Thu, 28 May 2026 18:23:25 +0100 Subject: [PATCH] feat: remote device connection as tool (#2506) * feat: remote device connection as tool * fix: test * fix: fix issues * fix: ruff * fix: mini issues * fix: bugs * fix: tests and mini issues --- .gitignore | 2 + application/agents/tool_executor.py | 65 ++- application/agents/tools/remote_device.py | 282 ++++++++++ application/agents/tools/tool_manager.py | 4 +- .../alembic/versions/0012_remote_devices.py | 117 ++++ .../versions/0013_devices_approval_modes.py | 52 ++ .../versions/0014_device_token_hash_index.py | 29 + application/api/devices/__init__.py | 11 + application/api/devices/auth.py | 141 +++++ application/api/devices/pairing.py | 518 ++++++++++++++++++ application/api/devices/routes.py | 327 +++++++++++ application/api/devices/session.py | 243 ++++++++ application/app.py | 15 + application/core/settings.py | 5 + application/devices/__init__.py | 1 + application/devices/broker.py | 344 ++++++++++++ application/devices/denylist.py | 172 ++++++ application/devices/normalizer.py | 84 +++ application/devices/splitter.py | 110 ++++ application/llm/handlers/base.py | 20 +- application/storage/db/models.py | 66 +++ .../db/repositories/device_audit_log.py | 126 +++++ .../device_auto_approve_patterns.py | 81 +++ .../storage/db/repositories/devices.py | 142 +++++ .../storage/db/repositories/message_events.py | 42 +- docs/content/Tools/_meta.js | 4 + docs/content/Tools/remote-device.mdx | 111 ++++ .../public/toolIcons/tool_remote_device.svg | 7 + frontend/src/agents/NewAgent.tsx | 56 +- frontend/src/api/client.ts | 2 + frontend/src/api/endpoints.ts | 7 + frontend/src/api/services/devicesService.ts | 172 ++++++ .../src/components/MultiSelectPopover.tsx | 8 +- .../src/conversation/ConversationBubble.tsx | 38 ++ frontend/src/conversation/types/index.ts | 3 + frontend/src/locale/en.json | 54 ++ frontend/src/modals/AddToolModal.tsx | 21 + frontend/src/settings/PairDeviceModal.tsx | 274 +++++++++ frontend/src/settings/RemoteDeviceConfig.tsx | 466 ++++++++++++++++ frontend/src/settings/Tools.tsx | 68 ++- frontend/src/settings/index.tsx | 4 + frontend/vite.config.ts | 38 +- ...86e12c2b0def48bc79fd4468533864c268441.json | 17 + ...c5f96cea35175091612545e24ff91ed65c8a3.json | 17 + ...a9c878800e030223be5189ebe13274c65d21f.json | 17 + ...8526da3813c82b0bb57ceaff298e959193d1f.json | 17 + tests/agents/test_tool_executor.py | 340 ++++++++++++ tests/core/test_openai_compatible.py | 9 +- tests/devices/__init__.py | 0 tests/devices/test_auto_approve_validation.py | 44 ++ tests/devices/test_broker_cleanup.py | 77 +++ tests/devices/test_broker_drain.py | 51 ++ tests/devices/test_denylist.py | 212 +++++++ tests/devices/test_me_endpoint.py | 123 +++++ tests/devices/test_normalizer.py | 81 +++ tests/devices/test_pairing_create.py | 283 ++++++++++ tests/devices/test_pairing_redeem_atomic.py | 269 +++++++++ tests/devices/test_remote_device_decision.py | 58 ++ tests/devices/test_session_ticket.py | 217 ++++++++ tests/devices/test_signature.py | 42 ++ tests/devices/test_splitter.py | 75 +++ tests/devices/test_update_device.py | 109 ++++ .../test_message_events_repository.py | 213 +++++++ tests/test_continuation.py | 186 +++++++ 64 files changed, 6733 insertions(+), 56 deletions(-) create mode 100644 application/agents/tools/remote_device.py create mode 100644 application/alembic/versions/0012_remote_devices.py create mode 100644 application/alembic/versions/0013_devices_approval_modes.py create mode 100644 application/alembic/versions/0014_device_token_hash_index.py create mode 100644 application/api/devices/__init__.py create mode 100644 application/api/devices/auth.py create mode 100644 application/api/devices/pairing.py create mode 100644 application/api/devices/routes.py create mode 100644 application/api/devices/session.py create mode 100644 application/devices/__init__.py create mode 100644 application/devices/broker.py create mode 100644 application/devices/denylist.py create mode 100644 application/devices/normalizer.py create mode 100644 application/devices/splitter.py create mode 100644 application/storage/db/repositories/device_audit_log.py create mode 100644 application/storage/db/repositories/device_auto_approve_patterns.py create mode 100644 application/storage/db/repositories/devices.py create mode 100644 docs/content/Tools/remote-device.mdx create mode 100644 frontend/public/toolIcons/tool_remote_device.svg create mode 100644 frontend/src/api/services/devicesService.ts create mode 100644 frontend/src/settings/PairDeviceModal.tsx create mode 100644 frontend/src/settings/RemoteDeviceConfig.tsx create mode 100644 scripts/e2e/mock_llm_fixtures/01c733b9660eed626247ffb527486e12c2b0def48bc79fd4468533864c268441.json create mode 100644 scripts/e2e/mock_llm_fixtures/b196e9532b6a979d959a570dcdec5f96cea35175091612545e24ff91ed65c8a3.json create mode 100644 scripts/e2e/mock_llm_fixtures/d42a7dccbd1f218bb095a28e2b4a9c878800e030223be5189ebe13274c65d21f.json create mode 100644 scripts/e2e/mock_llm_fixtures/ed4ff698efb5a8c5f4d6f8ad1a68526da3813c82b0bb57ceaff298e959193d1f.json create mode 100644 tests/devices/__init__.py create mode 100644 tests/devices/test_auto_approve_validation.py create mode 100644 tests/devices/test_broker_cleanup.py create mode 100644 tests/devices/test_broker_drain.py create mode 100644 tests/devices/test_denylist.py create mode 100644 tests/devices/test_me_endpoint.py create mode 100644 tests/devices/test_normalizer.py create mode 100644 tests/devices/test_pairing_create.py create mode 100644 tests/devices/test_pairing_redeem_atomic.py create mode 100644 tests/devices/test_remote_device_decision.py create mode 100644 tests/devices/test_session_ticket.py create mode 100644 tests/devices/test_signature.py create mode 100644 tests/devices/test_splitter.py create mode 100644 tests/devices/test_update_device.py diff --git a/.gitignore b/.gitignore index c3fc5b9f..3eef457a 100644 --- a/.gitignore +++ b/.gitignore @@ -107,6 +107,8 @@ celerybeat.pid # Environments .env +.env.local +.env.*.local .venv # Machine-specific Claude Code guidance (see CLAUDE.md preamble) CLAUDE.md diff --git a/application/agents/tool_executor.py b/application/agents/tool_executor.py index a7cac646..ae23f6fd 100644 --- a/application/agents/tool_executor.py +++ b/application/agents/tool_executor.py @@ -403,10 +403,30 @@ class ToolExecutor: {}, ) - if action_data.get("require_approval"): + require_approval = bool(action_data.get("require_approval")) + # ``denylist_forced`` marks a prompt the hard denylist mandates; a + # headless allowlist must never bypass it (see below). + denylist_forced = False + # ``remote_device`` decides per-invocation based on the live device + # state (``approval_mode``, sticky patterns, allow/denylist). The + # cached ``user_tools.actions[].require_approval`` snapshot does + # not reflect later approval-mode changes nor command-level + # heuristics, so consult the tool directly. + if tool_data.get("name") == "remote_device": + require_approval, denylist_forced = ( + self._remote_device_requires_approval( + tool_data, action_name, arguments, + ) + ) + + if require_approval: if self.headless: tool_row_id = str(tool_data.get("id") or tool_id) - if tool_row_id in self.tool_allowlist: + # A denylist-forced prompt is never pre-authorizable: a + # scheduled/headless run with the device allowlisted must + # still be denied a denylisted command. Only non-forced + # approvals honor the allowlist bypass. + if tool_row_id in self.tool_allowlist and not denylist_forced: # Pre-authorized for headless execution — fall through. return None return { @@ -425,7 +445,7 @@ class ToolExecutor: "error_type": "tool_not_allowed", "thought_signature": getattr(call, "thought_signature", None), } - return { + payload = { "call_id": call_id, "name": llm_name, "tool_name": tool_data.get("name", "unknown"), @@ -436,9 +456,42 @@ class ToolExecutor: "pause_type": "awaiting_approval", "thought_signature": getattr(call, "thought_signature", None), } + # Surface the device id so the approval UI can offer a + # "don't ask again" sticky-pattern action for remote devices. + if tool_data.get("name") == "remote_device": + config = tool_data.get("config") or {} + if config.get("device_id"): + payload["device_id"] = config["device_id"] + return payload return None + def _remote_device_requires_approval( + self, tool_data: Dict, action_name: str, arguments: Dict, + ) -> tuple[bool, bool]: + """Live approval decision for a ``remote_device`` invocation. + + Instantiates ``RemoteDeviceTool`` with the cached config and the + executor's user context, then asks it to evaluate the command. + Returns ``(requires_approval, denylist_forced)``. Falls back to a + denylist-forced prompt on any error so a misconfigured device never + silently bypasses the prompt — not even via the headless allowlist. + """ + try: + from application.agents.tools.remote_device import RemoteDeviceTool + + tool = RemoteDeviceTool( + config=tool_data.get("config") or {}, + user_id=self.user, + ) + return tool.preview_decision(action_name, arguments) + except Exception: + logger.exception( + "remote_device preview_decision failed; defaulting to a " + "forced prompt", + ) + return True, True + def execute(self, tools_dict: Dict, call, llm_class_name: str): """Execute a tool call. Yields status events, returns (result, call_id).""" parser = ToolActionParser(llm_class_name, name_mapping=self._name_to_tool) @@ -499,6 +552,12 @@ class ToolExecutor: "arguments": call_args, } tool_data = tools_dict[tool_id] + # Surface the device id on remote_device tool-call events so the + # approval UI can wire up the sticky "don't ask again" button. + if tool_data.get("name") == "remote_device": + config = tool_data.get("config") or {} + if config.get("device_id"): + tool_call_data["device_id"] = config["device_id"] # Journal first so the reconciler sees malformed calls and any # subsequent ``_mark_failed`` actually updates a real row. proposed_ok = _record_proposed( diff --git a/application/agents/tools/remote_device.py b/application/agents/tools/remote_device.py new file mode 100644 index 00000000..f728a0b7 --- /dev/null +++ b/application/agents/tools/remote_device.py @@ -0,0 +1,282 @@ +"""Remote Device tool. + +Run shell commands on a paired remote device via the DeviceBroker. +""" + +from __future__ import annotations + +import logging +import time +import uuid +from datetime import datetime, timezone +from typing import Any, Dict, Optional + +from application.agents.tools.base import Tool +from application.devices.broker import get_broker +from application.devices.denylist import check_denylist +from application.devices.normalizer import normalize_command +from application.storage.db.repositories.device_audit_log import ( + DeviceAuditLogRepository, +) +from application.storage.db.repositories.device_auto_approve_patterns import ( + DeviceAutoApprovePatternsRepository, +) +from application.storage.db.repositories.devices import DevicesRepository +from application.storage.db.session import db_readonly, db_session + + +logger = logging.getLogger(__name__) + + +_DEFAULT_TIMEOUT_MS = 30_000 +_MAX_TIMEOUT_MS = 600_000 + + +class RemoteDeviceTool(Tool): + """Remote Device + Run shell commands on a paired remote machine via docsgpt-cli host. + """ + + def __init__(self, config: Optional[dict] = None, user_id: Optional[str] = None): + self.config = config or {} + self.user_id = user_id + self.device_id = self.config.get("device_id") or "" + self._device: Optional[dict] = None + if self.device_id and self.user_id: + self._device = self._load_device() + + def _load_device(self) -> Optional[dict]: + try: + with db_readonly() as conn: + return DevicesRepository(conn).get(self.device_id, user_id=self.user_id) + except Exception: + logger.exception("failed to load device %s", self.device_id) + return None + + # ------------------------------------------------------------------ + # Tool ABC + # ------------------------------------------------------------------ + def get_actions_metadata(self): + device = self._device or {} + device_name = device.get("name") or "remote device" + description = device.get("description") or "" + approval_mode = device.get("approval_mode") or "ask" + return [ + { + "name": "run_command", + "description": ( + f"Execute a shell command on the remote device " + f"'{device_name}'. {description}".strip() + ), + "active": True, + "require_approval": approval_mode != "full", + "parameters": { + "type": "object", + "properties": { + "command": { + "type": "string", + "description": "Shell command to run.", + "filled_by_llm": True, + "value": "", + }, + "working_directory": { + "type": "string", + "description": "Working directory on the remote.", + "filled_by_llm": True, + "value": "", + }, + "timeout_ms": { + "type": "integer", + "description": "Timeout in milliseconds (max 600000).", + "filled_by_llm": True, + "value": "", + }, + }, + "required": ["command"], + }, + } + ] + + def get_config_requirements(self): + return { + "device_id": { + "type": "string", + "label": "Device", + "description": "Paired remote device id.", + "required": True, + "source": "devices", + } + } + + def preview_requires_approval(self, action_name: str, params: dict) -> bool: + """Live approval decision for a specific invocation. + + The tool_executor gate calls this for ``remote_device`` so the + decision considers the device's current ``approval_mode``, sticky + patterns, and the denylist — rather than trusting the static + ``user_tools.actions[].require_approval`` snapshot stored at pair + time. Returns ``True`` when a prompt is required. + """ + requires_approval, _denylist_forced = self.preview_decision( + action_name, params, + ) + return requires_approval + + def preview_decision( + self, action_name: str, params: dict, + ) -> tuple[bool, bool]: + """Live approval decision plus whether it's a denylist-forced prompt. + + Returns ``(requires_approval, denylist_forced)``. ``denylist_forced`` + is True only when the prompt is mandated by the hard denylist, which + a headless allowlist must never bypass. Unknown / inactive devices + and missing commands require approval but are NOT denylist-forced. + """ + if action_name != "run_command": + return True, False + if not self.device_id or not self.user_id: + return True, False + if self._device is None: + self._device = self._load_device() + device = self._device + if device is None or device.get("status") != "active": + # Don't bypass the prompt for an unknown / inactive device; + # execute_action will surface the error. + return True, False + command = ((params or {}).get("command") or "").strip() + if not command: + return True, False + reason, effective_mode = self._decide_approval(device, command) + denylist_forced = reason == "denylist_forced_prompt" + return effective_mode != "full", denylist_forced + + def execute_action(self, action_name: str, **kwargs): + if action_name != "run_command": + return {"error": f"unknown action: {action_name}"} + if not self.device_id or not self.user_id: + return {"error": "device_id and user_id required"} + if self._device is None: + self._device = self._load_device() + device = self._device + if device is None: + return {"error": "device not found"} + if device.get("status") != "active": + return {"error": f"device status: {device.get('status')}"} + + command = (kwargs.get("command") or "").strip() + if not command: + return {"error": "command is required"} + working_directory = kwargs.get("working_directory") or "" + timeout_ms = kwargs.get("timeout_ms") + try: + timeout_ms = int(timeout_ms) if timeout_ms else _DEFAULT_TIMEOUT_MS + except (TypeError, ValueError): + timeout_ms = _DEFAULT_TIMEOUT_MS + timeout_ms = min(max(timeout_ms, 1), _MAX_TIMEOUT_MS) + + decision_reason, effective_mode = self._decide_approval(device, command) + denied = self._denylist_label(command) + + envelope = { + "invocation_id": "inv_" + uuid.uuid4().hex, + "action": "run_command", + "params": { + "command": command, + "working_directory": working_directory, + "timeout_ms": timeout_ms, + }, + "approval_mode": effective_mode, + "issued_at": datetime.now(timezone.utc).isoformat(), + } + broker = get_broker() + inv = broker.dispatch_invocation(self.device_id, self.user_id, envelope) + + try: + with db_session() as conn: + DeviceAuditLogRepository(conn).record_dispatch( + device_id=self.device_id, + user_id=self.user_id, + invocation_id=inv.invocation_id, + command=command, + working_dir=working_directory, + approval_mode=effective_mode, + decision="dispatched", + decision_reason=decision_reason or ("denylist:" + denied if denied else None), + issued_at=datetime.now(timezone.utc), + ) + except Exception: + logger.exception("audit record_dispatch failed for %s", inv.invocation_id) + + return self._collect_result(broker, inv, device, timeout_ms) + + # ------------------------------------------------------------------ + # Internals + # ------------------------------------------------------------------ + def _decide_approval(self, device: dict, command: str) -> tuple[Optional[str], str]: + """Resolve the effective approval mode + a short audit reason. + + Effective mode is ``full`` (auto-run, no prompt) or ``ask`` (prompt). + """ + mode = device.get("approval_mode") or "ask" + # Denylist forces a prompt on every path — full access and the + # ask-mode sticky auto-approve alike. + if check_denylist(command): + return ("denylist_forced_prompt", "ask") + if mode == "full": + return ("full_access_passthrough", "full") + # mode == "ask" + if self._matches_sticky(command): + return ("sticky_auto_approve", "full") + return ("user_approval_required", "ask") + + def _denylist_label(self, command: str) -> Optional[str]: + return check_denylist(command) + + def _matches_sticky(self, command: str) -> bool: + pattern = normalize_command(command) + if not pattern: + return False + try: + with db_readonly() as conn: + return DeviceAutoApprovePatternsRepository(conn).has_pattern( + self.device_id, self.user_id, pattern, + ) + except Exception: + logger.exception("sticky lookup failed") + return False + + def _collect_result(self, broker, inv, device: dict, timeout_ms: int) -> Dict[str, Any]: + """Drain output from the broker until the control chunk arrives.""" + deadline = time.time() + (timeout_ms / 1000.0) + 5.0 + stdout = [] + stderr = [] + try: + for chunk in broker.drain_output( + inv.invocation_id, timeout=1.0, deadline=deadline + ): + if time.time() > deadline: + break + stream = chunk.get("stream") + if stream == "stdout": + stdout.append(chunk.get("chunk", "")) + elif stream == "stderr": + stderr.append(chunk.get("chunk", "")) + elif stream == "control": + # control chunks include exit_code; drain loop will stop next iter + pass + finally: + broker.cleanup_invocation(inv.invocation_id) + + # Deadline hit with no control chunk: the device never connected or + # never finished. Surface a clear timeout instead of empty success. + if not inv.completed.is_set() and inv.exit_code is None and not inv.error: + inv.error = "device did not respond (timed out)" + + return { + "exit_code": inv.exit_code, + "stdout": "".join(stdout) if stdout else "".join(inv.stdout_parts), + "stderr": "".join(stderr) if stderr else "".join(inv.stderr_parts), + "duration_ms": inv.duration_ms, + "device_name": device.get("name"), + "error": inv.error, + } diff --git a/application/agents/tools/tool_manager.py b/application/agents/tools/tool_manager.py index 34924511..f969e87f 100644 --- a/application/agents/tools/tool_manager.py +++ b/application/agents/tools/tool_manager.py @@ -29,7 +29,7 @@ class ToolManager: for member_name, obj in inspect.getmembers(module, inspect.isclass): if issubclass(obj, Tool) and obj is not Tool: if ( - tool_name in {"mcp_tool", "notes", "memory", "todo_list", "scheduler"} + tool_name in {"mcp_tool", "notes", "memory", "todo_list", "scheduler", "remote_device"} and user_id ): return obj(tool_config, user_id) @@ -40,7 +40,7 @@ class ToolManager: if tool_name not in self.tools: raise ValueError(f"Tool '{tool_name}' not loaded") if ( - tool_name in {"mcp_tool", "memory", "todo_list", "notes", "scheduler"} + tool_name in {"mcp_tool", "memory", "todo_list", "notes", "scheduler", "remote_device"} and user_id ): tool_config = self.config.get(tool_name, {}) diff --git a/application/alembic/versions/0012_remote_devices.py b/application/alembic/versions/0012_remote_devices.py new file mode 100644 index 00000000..bf4eb007 --- /dev/null +++ b/application/alembic/versions/0012_remote_devices.py @@ -0,0 +1,117 @@ +"""0012 remote devices — devices, device_audit_log, device_auto_approve_patterns. + +Adds tables for the Remote Device feature: paired remote hosts that DocsGPT +agents can drive via shell tool calls. + +Revision ID: 0012_remote_devices +Revises: 0011_schedules_nullable_agent +""" + +from typing import Sequence, Union + +from alembic import op + + +revision: str = "0012_remote_devices" +down_revision: Union[str, None] = "0011_schedules_nullable_agent" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.execute( + """ + CREATE TABLE devices ( + id TEXT PRIMARY KEY, + user_id TEXT NOT NULL, + name TEXT NOT NULL, + hostname TEXT, + os TEXT, + arch TEXT, + cli_version TEXT, + machine_pubkey_fingerprint TEXT NOT NULL, + token_hash TEXT NOT NULL, + approval_mode TEXT NOT NULL DEFAULT 'ask', + description TEXT, + status TEXT NOT NULL DEFAULT 'active', + paired_at TIMESTAMPTZ NOT NULL DEFAULT now(), + last_seen_at TIMESTAMPTZ, + revoked_at TIMESTAMPTZ, + revoke_reason TEXT, + CONSTRAINT devices_approval_mode_chk + CHECK (approval_mode IN ('ask', 'writes-only', 'never')), + CONSTRAINT devices_status_chk + CHECK (status IN ('active', 'revoked')), + CONSTRAINT devices_user_name_uidx UNIQUE (user_id, name) + ); + """ + ) + op.execute( + "CREATE INDEX devices_user_active_idx ON devices(user_id) " + "WHERE status = 'active';" + ) + + op.execute( + """ + CREATE TABLE device_audit_log ( + id BIGSERIAL PRIMARY KEY, + device_id TEXT NOT NULL REFERENCES devices(id) ON DELETE CASCADE, + user_id TEXT NOT NULL, + agent_id TEXT, + conversation_id TEXT, + invocation_id TEXT NOT NULL, + action TEXT NOT NULL, + command TEXT NOT NULL, + working_dir TEXT, + approval_mode TEXT NOT NULL, + decision TEXT NOT NULL, + decision_reason TEXT, + issued_at TIMESTAMPTZ NOT NULL, + started_at TIMESTAMPTZ, + finished_at TIMESTAMPTZ, + exit_code INTEGER, + duration_ms INTEGER, + stdout_sha256 CHAR(64), + stderr_sha256 CHAR(64), + stdout_bytes INTEGER, + stderr_bytes INTEGER, + error TEXT, + created_at TIMESTAMPTZ NOT NULL DEFAULT now() + ); + """ + ) + op.execute( + "CREATE INDEX device_audit_device_idx " + "ON device_audit_log(device_id, created_at DESC);" + ) + op.execute( + "CREATE INDEX device_audit_user_idx " + "ON device_audit_log(user_id, created_at DESC);" + ) + + # Per-device, per-user sticky "don't ask again" patterns. Normalized + # form: command head + first sub-token, wildcard rest (see + # application/devices/normalizer.py). + op.execute( + """ + CREATE TABLE device_auto_approve_patterns ( + id BIGSERIAL PRIMARY KEY, + device_id TEXT NOT NULL REFERENCES devices(id) ON DELETE CASCADE, + user_id TEXT NOT NULL, + pattern TEXT NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + CONSTRAINT device_auto_approve_uidx + UNIQUE (device_id, user_id, pattern) + ); + """ + ) + op.execute( + "CREATE INDEX device_auto_approve_lookup_idx " + "ON device_auto_approve_patterns(device_id, user_id);" + ) + + +def downgrade() -> None: + op.execute("DROP TABLE IF EXISTS device_auto_approve_patterns CASCADE;") + op.execute("DROP TABLE IF EXISTS device_audit_log CASCADE;") + op.execute("DROP TABLE IF EXISTS devices CASCADE;") diff --git a/application/alembic/versions/0013_devices_approval_modes.py b/application/alembic/versions/0013_devices_approval_modes.py new file mode 100644 index 00000000..eeb7e0e2 --- /dev/null +++ b/application/alembic/versions/0013_devices_approval_modes.py @@ -0,0 +1,52 @@ +"""0013 devices approval modes — collapse to ask/full. + +Drops the ``writes-only`` mode and renames ``never`` to ``full``. Existing +rows are converted (``writes-only`` -> ``ask``, ``never`` -> ``full``) and +the CHECK constraint is rewritten to allow only ``ask`` / ``full``. + +Revision ID: 0013_devices_approval_modes +Revises: 0012_remote_devices +""" + +from typing import Sequence, Union + +from alembic import op + + +revision: str = "0013_devices_approval_modes" +down_revision: Union[str, None] = "0012_remote_devices" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # Drop the old constraint before rewriting values so the conversion + # never trips the stale CHECK. + op.execute( + "ALTER TABLE devices DROP CONSTRAINT IF EXISTS devices_approval_mode_chk;" + ) + op.execute( + "UPDATE devices SET approval_mode = 'ask' WHERE approval_mode = 'writes-only';" + ) + op.execute( + "UPDATE devices SET approval_mode = 'full' WHERE approval_mode = 'never';" + ) + op.execute( + "ALTER TABLE devices ADD CONSTRAINT devices_approval_mode_chk " + "CHECK (approval_mode IN ('ask', 'full'));" + ) + + +def downgrade() -> None: + # ``never`` is restored from ``full``. ``writes-only`` is gone — rows + # that were converted to ``ask`` stay ``ask`` (no reverse mapping). + op.execute( + "ALTER TABLE devices DROP CONSTRAINT IF EXISTS devices_approval_mode_chk;" + ) + op.execute( + "UPDATE devices SET approval_mode = 'never' WHERE approval_mode = 'full';" + ) + op.execute( + "ALTER TABLE devices ADD CONSTRAINT devices_approval_mode_chk " + "CHECK (approval_mode IN ('ask', 'writes-only', 'never'));" + ) diff --git a/application/alembic/versions/0014_device_token_hash_index.py b/application/alembic/versions/0014_device_token_hash_index.py new file mode 100644 index 00000000..2df9df77 --- /dev/null +++ b/application/alembic/versions/0014_device_token_hash_index.py @@ -0,0 +1,29 @@ +"""0014 device token_hash index — index the per-request token lookup. + +``find_by_token_hash`` runs on every CLI request but ``token_hash`` was +unindexed. Token uniqueness is guaranteed at generation, so a UNIQUE +index is safe (and doubles as a guard against accidental collisions). + +Revision ID: 0014_device_token_hash_index +Revises: 0013_devices_approval_modes +""" + +from typing import Sequence, Union + +from alembic import op + + +revision: str = "0014_device_token_hash_index" +down_revision: Union[str, None] = "0013_devices_approval_modes" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.execute( + "CREATE UNIQUE INDEX devices_token_hash_uidx ON devices(token_hash);" + ) + + +def downgrade() -> None: + op.execute("DROP INDEX IF EXISTS devices_token_hash_uidx;") diff --git a/application/api/devices/__init__.py b/application/api/devices/__init__.py new file mode 100644 index 00000000..c0f108a1 --- /dev/null +++ b/application/api/devices/__init__.py @@ -0,0 +1,11 @@ +"""Flask blueprint for the Remote Device feature.""" + +from __future__ import annotations + +from flask import Blueprint + +from .routes import register as register_routes + + +devices_bp = Blueprint("devices", __name__) +register_routes(devices_bp) diff --git a/application/api/devices/auth.py b/application/api/devices/auth.py new file mode 100644 index 00000000..d4338fe3 --- /dev/null +++ b/application/api/devices/auth.py @@ -0,0 +1,141 @@ +"""Device session-token verification + machine-key signature check.""" + +from __future__ import annotations + +import base64 +import hashlib +import logging +import time +from typing import Optional, Tuple + +from flask import jsonify, make_response, request + +from application.core.settings import settings +from application.storage.db.repositories.devices import DevicesRepository +from application.storage.db.session import db_readonly, db_session + + +logger = logging.getLogger(__name__) + + +def hash_session_token(token: str) -> str: + """SHA-256 over the opaque session token. Stored in ``devices.token_hash``.""" + return hashlib.sha256(token.encode("utf-8")).hexdigest() + + +def fingerprint_pubkey(pubkey_b64: str) -> str: + """SHA-256 over the raw public-key bytes; stored as the fingerprint.""" + raw = base64.b64decode(pubkey_b64) + return hashlib.sha256(raw).hexdigest() + + +def _canonical_payload(method: str, path: str, ts: str, body: bytes) -> str: + """Canonical string the device signs / the server verifies. + + Format: ``"{method} {path} {ts} {sha256_hex(body)}"``. The body hash + binds the request body into the signature so a captured signature can't + be replayed with a tampered body inside the timestamp window. For a GET + (empty body) this is the SHA-256 of the empty string. + + KEEP IN SYNC with the DocsGPT-cli signer + (``internal/host/identity.go`` ``CanonicalPayload`` / ``SignRequest``). + The hex encoding and single-space separators must match byte-for-byte. + """ + body_hash = hashlib.sha256(body or b"").hexdigest() + return f"{method} {path} {ts} {body_hash}" + + +def verify_device_session(*, touch: bool = True) -> Tuple[Optional[dict], Optional[tuple]]: + """Validate the device session token on the current request. + + Returns ``(device_row, None)`` on success or ``(None, response)`` on + failure, where ``response`` is a ``(json, status)`` Flask tuple. + + Args: + touch: When True, bump ``last_seen_at`` on the device row. + """ + auth = request.headers.get("Authorization", "") + if not auth.lower().startswith("bearer "): + return None, _error("missing_token", 401) + token = auth.split(" ", 1)[1].strip() + if not token: + return None, _error("missing_token", 401) + + token_hash = hash_session_token(token) + with db_readonly() as conn: + device = DevicesRepository(conn).find_by_token_hash(token_hash) + if device is None: + return None, _error("invalid_token", 401) + + if settings.REMOTE_DEVICE_REQUIRE_SIGNATURE: + sig_error = _verify_signature(device) + if sig_error is not None: + return None, sig_error + + if touch: + try: + with db_session() as conn: + DevicesRepository(conn).touch_last_seen(device["id"]) + except Exception: + logger.exception("touch_last_seen failed for device %s", device["id"]) + + return device, None + + +def _verify_signature(device: dict) -> Optional[tuple]: + """Verify ``X-Device-Signature`` against the stored machine pubkey.""" + sig_b64 = request.headers.get("X-Device-Signature") + ts = request.headers.get("X-Device-Timestamp") + fp = request.headers.get("X-Device-Machine-Key") + if not sig_b64 or not ts or not fp: + return _error("missing_signature", 401) + + try: + ts_int = int(ts) + except (TypeError, ValueError): + return _error("invalid_timestamp", 401) + if abs(time.time() - ts_int) > 300: + return _error("timestamp_skew", 401) + + if fp != device.get("machine_pubkey_fingerprint"): + return _error("fingerprint_mismatch", 401) + + # Defer cryptography import to keep cold-start light when signatures + # are disabled (the dev default). + try: + from cryptography.hazmat.primitives.asymmetric.ed25519 import ( + Ed25519PublicKey, + ) + from cryptography.exceptions import InvalidSignature + except ImportError: + logger.error( + "Signature verification requested but ``cryptography`` is not installed." + ) + return _error("signature_unsupported", 500) + + # The full public key isn't stored — only the fingerprint. For signature + # verification we accept the pubkey in a header too; we trust it iff + # its fingerprint matches the stored one. This avoids a separate pubkey + # column for MVP. + pubkey_b64 = request.headers.get("X-Device-Machine-Pubkey") + if not pubkey_b64: + return _error("missing_pubkey", 401) + if fingerprint_pubkey(pubkey_b64) != device["machine_pubkey_fingerprint"]: + return _error("pubkey_fingerprint_mismatch", 401) + + payload = _canonical_payload( + request.method, request.path, ts, request.get_data() + ).encode("utf-8") + try: + pubkey = Ed25519PublicKey.from_public_bytes(base64.b64decode(pubkey_b64)) + pubkey.verify(base64.b64decode(sig_b64), payload) + except (InvalidSignature, ValueError): + return _error("invalid_signature", 401) + + return None + + +def _error(code: str, status: int) -> tuple: + return make_response( + jsonify({"success": False, "error": code}), status + ) diff --git a/application/api/devices/pairing.py b/application/api/devices/pairing.py new file mode 100644 index 00000000..5eafb016 --- /dev/null +++ b/application/api/devices/pairing.py @@ -0,0 +1,518 @@ +"""Pairing flow (RFC 8628 device authorization grant).""" + +from __future__ import annotations + +import json +import logging +import os +import secrets +import time +import uuid +from typing import Optional + +from flask import jsonify, make_response, request +from sqlalchemy.exc import IntegrityError + +from application.api.devices.auth import fingerprint_pubkey, hash_session_token +from application.cache import get_redis_instance +from application.core.settings import settings +from application.storage.db.repositories.devices import DevicesRepository +from application.storage.db.repositories.user_tools import UserToolsRepository +from application.storage.db.session import db_session + + +logger = logging.getLogger(__name__) + + +_USER_CODE_ALPHABET = "ABCDEFGHJKLMNPQRSTUVWXYZ23456789" +_USER_CODE_LEN = 8 # ABCD-WXYZ (8 chars, displayed with a single dash) +_PAIRING_REDIS_PREFIX = "docsgpt:device_pairing:" +_ALLOWED_APPROVAL_MODES = {"ask", "full"} +# Attempts to find a non-colliding device name before releasing the claim. +_DEVICE_NAME_RETRIES = 5 + +# Atomic redeem claim: read the pairing JSON, and only if its ``status`` is +# ``pending`` flip it to ``redeemed`` and write it back (TTL preserved). +# Returns 1 if this caller won the claim, 0 otherwise (missing/not pending). +# Run server-side so two concurrent redeems of one code can't both win. +_CLAIM_PAIRING_LUA = """ +local raw = redis.call('GET', KEYS[1]) +if not raw then + return 0 +end +local ok, state = pcall(cjson.decode, raw) +if not ok then + return 0 +end +if state['status'] ~= 'pending' then + return 0 +end +state['status'] = 'redeemed' +local ttl = redis.call('TTL', KEYS[1]) +local encoded = cjson.encode(state) +if ttl and ttl > 0 then + redis.call('SET', KEYS[1], encoded, 'EX', ttl) +else + redis.call('SET', KEYS[1], encoded) +end +return 1 +""" + + +def _redis(): + return get_redis_instance() + + +def _generate_user_code() -> str: + return "".join(secrets.choice(_USER_CODE_ALPHABET) for _ in range(_USER_CODE_LEN)) + + +def _format_user_code(code: str) -> str: + return f"{code[:4]}-{code[4:]}" + + +def _normalize_user_code(code: str) -> str: + return code.replace("-", "").replace(" ", "").upper() + + +def _user_code_index_key(user_code: str) -> str: + return f"docsgpt:device_pairing_code:{user_code}" + + +def create_pairing() -> tuple: + """Create a new pairing code for the calling user. + + Optional body fields ``name``, ``description``, and ``approval_mode`` + are stashed in Redis alongside the pairing state and consumed by + ``redeem_pairing`` when the device row is created. Unknown values + fall back to: name -> CLI hostname; description -> empty; approval_mode + -> ``ask``. + """ + decoded = getattr(request, "decoded_token", None) + user_id = decoded.get("sub") if isinstance(decoded, dict) else None + if not user_id: + return _error("authentication_required", 401) + + body = request.get_json(silent=True) or {} + raw_name = body.get("name") + if raw_name is not None and not isinstance(raw_name, str): + return _error("invalid_name", 400) + requested_name = (raw_name or "").strip() or None + requested_description = body.get("description") + if isinstance(requested_description, str): + requested_description = requested_description.strip() or None + else: + requested_description = None + requested_mode = body.get("approval_mode") + if requested_mode is not None and requested_mode not in _ALLOWED_APPROVAL_MODES: + return _error("invalid_approval_mode", 400) + + redis_client = _redis() + if redis_client is None: + return _error("redis_unavailable", 503) + + device_code = f"dc_{uuid.uuid4().hex}" + user_code = _generate_user_code() + ttl = int(settings.REMOTE_DEVICE_PAIRING_TTL_SECONDS) + + state = { + "device_code": device_code, + "user_code": _format_user_code(user_code), + "user_code_raw": user_code, + "user_id": user_id, + "status": "pending", + "created_at": int(time.time()), + "requested_name": requested_name, + "requested_description": requested_description, + "requested_approval_mode": requested_mode, + } + try: + redis_client.setex( + _PAIRING_REDIS_PREFIX + device_code, ttl, json.dumps(state) + ) + redis_client.setex(_user_code_index_key(user_code), ttl, device_code) + except Exception: + logger.exception("redis setex failed during pairing create") + return _error("redis_unavailable", 503) + + base_url = settings.API_URL or "" + return make_response( + jsonify( + { + "device_code": device_code, + "user_code": _format_user_code(user_code), + "verification_uri": base_url, + "expires_in": ttl, + "interval": 3, + } + ), + 200, + ) + + +def get_pairing(device_code: str) -> tuple: + """Poll pairing status (UI side).""" + decoded = getattr(request, "decoded_token", None) + user_id = decoded.get("sub") if isinstance(decoded, dict) else None + if not user_id: + return _error("authentication_required", 401) + + state = _load_pairing(device_code) + if state is None: + return _error("pairing_not_found", 404) + if state.get("user_id") != user_id: + return _error("pairing_not_found", 404) + + return make_response( + jsonify( + { + "device_code": state["device_code"], + "user_code": state.get("user_code"), + "status": state.get("status", "pending"), + "device_id": state.get("device_id"), + "device_name": state.get("device_name"), + } + ), + 200, + ) + + +def delete_pairing(device_code: str) -> tuple: + decoded = getattr(request, "decoded_token", None) + user_id = decoded.get("sub") if isinstance(decoded, dict) else None + if not user_id: + return _error("authentication_required", 401) + state = _load_pairing(device_code) + if state is None: + return _error("pairing_not_found", 404) + if state.get("user_id") != user_id: + return _error("pairing_not_found", 404) + redis_client = _redis() + if redis_client is None: + return _error("redis_unavailable", 503) + try: + redis_client.delete(_PAIRING_REDIS_PREFIX + device_code) + if state.get("user_code_raw"): + redis_client.delete(_user_code_index_key(state["user_code_raw"])) + except Exception: + logger.exception("redis delete failed during pairing cancel") + return make_response(jsonify({"success": True}), 200) + + +def redeem_pairing() -> tuple: + """CLI submits the user_code and machine info to claim a device row. + + Returns ``device_id`` + opaque ``session_token`` (shown once). The + token is stored only as ``token_hash`` so a Redis snapshot leak + + DB snapshot leak still can't reconstruct it. + """ + body = request.get_json(silent=True) or {} + user_code_raw = _normalize_user_code(str(body.get("user_code", ""))) + if not user_code_raw or len(user_code_raw) != _USER_CODE_LEN: + return _error("invalid_user_code", 400) + hostname = body.get("hostname") or "unknown" + os_name = body.get("os") or "" + arch = body.get("arch") or "" + cli_version = body.get("cli_version") or "" + pubkey_b64 = body.get("machine_pubkey") + if not pubkey_b64: + return _error("missing_machine_pubkey", 400) + + redis_client = _redis() + if redis_client is None: + return _error("redis_unavailable", 503) + try: + device_code = redis_client.get(_user_code_index_key(user_code_raw)) + if device_code is None: + return _error("pairing_not_found", 404) + if isinstance(device_code, (bytes, bytearray)): + device_code = device_code.decode("utf-8") + except Exception: + logger.exception("redis get failed during pairing redeem") + return _error("redis_unavailable", 503) + + state = _load_pairing(device_code) + if state is None: + return _error("pairing_not_found", 404) + if state.get("status") != "pending": + return _error("pairing_already_redeemed", 409) + + # Atomically claim the pairing before doing any work. Two concurrent + # redeems of the same code can both pass the check above; only the one + # that flips ``pending`` -> ``redeemed`` here proceeds to create a device. + claimed = _claim_pairing(redis_client, device_code) + if claimed is None: + return _error("redis_unavailable", 503) + if not claimed: + return _error("pairing_already_redeemed", 409) + + user_id = state["user_id"] + + try: + fingerprint = fingerprint_pubkey(pubkey_b64) + except Exception: + return _error("invalid_machine_pubkey", 400) + session_token = "tok_" + secrets.token_urlsafe(32) + token_hash = hash_session_token(session_token) + device_id = "dev_" + uuid.uuid4().hex + + # Apply UI-side overrides stashed at pairing create time, with + # sensible fallbacks. + requested_name = state.get("requested_name") or None + description = state.get("requested_description") or None + approval_mode = state.get("requested_approval_mode") or "ask" + if approval_mode not in _ALLOWED_APPROVAL_MODES: + approval_mode = "ask" + + name_source = requested_name or hostname + # The claim has already been consumed, so a failed device insert would + # strand a one-time code. ``_next_device_name`` only adds a 16-bit + # suffix, so a ``UNIQUE (user_id, name)`` collision is possible; retry + # with a fresh name a few times, and release the claim if every attempt + # collides so the user can redeem the code again. + name = None + created = False + for _attempt in range(_DEVICE_NAME_RETRIES): + candidate = _next_device_name(user_id, name_source) + try: + with db_session() as conn: + DevicesRepository(conn).create( + device_id=device_id, + user_id=user_id, + name=candidate, + machine_pubkey_fingerprint=fingerprint, + token_hash=token_hash, + hostname=hostname, + os=os_name, + arch=arch, + cli_version=cli_version, + approval_mode=approval_mode, + description=description, + ) + _upsert_remote_device_user_tool( + conn, user_id=user_id, device_id=device_id, + device_name=candidate, description=description, + ) + name = candidate + created = True + break + except IntegrityError: + # Name collision on this attempt; the transaction rolled back. + # Loop to try a freshly-generated name. + logger.warning( + "device name collision on redeem (attempt %d); retrying", + _attempt + 1, + ) + continue + except Exception: + logger.exception("failed to create device row during redeem") + _release_claim(redis_client, device_code, state, user_code_raw) + return _error("internal_error", 500) + + if not created: + logger.error( + "device create exhausted name retries during redeem; " + "releasing claim", + ) + _release_claim(redis_client, device_code, state, user_code_raw) + return _error("internal_error", 500) + + state["status"] = "redeemed" + state["device_id"] = device_id + state["device_name"] = name + try: + redis_client.setex( + _PAIRING_REDIS_PREFIX + device_code, 300, json.dumps(state) + ) + if state.get("user_code_raw"): + redis_client.delete(_user_code_index_key(state["user_code_raw"])) + except Exception: + logger.exception("redis update failed after redeem (non-fatal)") + + return make_response( + jsonify( + { + "device_id": device_id, + "session_token": session_token, + "name": name, + } + ), + 200, + ) + + +def _claim_pairing(redis_client, device_code: str) -> Optional[bool]: + """Atomically transition a pairing from ``pending`` to ``redeemed``. + + Returns ``True`` if this caller won the claim, ``False`` if the pairing + is already non-pending or gone, and ``None`` if the atomic op itself + failed (caller should treat that as Redis-unavailable rather than + proceeding, to avoid a non-atomic double-redeem). + """ + key = _PAIRING_REDIS_PREFIX + device_code + try: + result = redis_client.eval(_CLAIM_PAIRING_LUA, 1, key) + except Exception: + logger.exception("redis eval failed during pairing claim") + return None + return bool(result) + + +def _release_claim( + redis_client, device_code: str, state: dict, user_code_raw: str, +) -> None: + """Revert a consumed claim so the one-time code is redeemable again. + + Called when the device row can't be created after the atomic claim + already flipped the pairing to ``redeemed``. Restores ``pending`` and + re-points the user-code index at the device code, preserving the + remaining TTL where possible. Best-effort: failures are logged, not + raised, since the caller is already returning an error. + """ + key = _PAIRING_REDIS_PREFIX + device_code + try: + ttl = redis_client.ttl(key) + revert = dict(state) + revert["status"] = "pending" + revert.pop("device_id", None) + revert.pop("device_name", None) + encoded = json.dumps(revert) + if isinstance(ttl, int) and ttl > 0: + redis_client.setex(key, ttl, encoded) + if user_code_raw: + redis_client.setex( + _user_code_index_key(user_code_raw), ttl, device_code + ) + else: + redis_client.set(key, encoded) + if user_code_raw: + redis_client.set(_user_code_index_key(user_code_raw), device_code) + except Exception: + logger.exception("failed to release pairing claim for %s", device_code) + + +def _load_pairing(device_code: str) -> Optional[dict]: + redis_client = _redis() + if redis_client is None: + return None + try: + raw = redis_client.get(_PAIRING_REDIS_PREFIX + device_code) + except Exception: + return None + if raw is None: + return None + if isinstance(raw, (bytes, bytearray)): + raw = raw.decode("utf-8") + try: + return json.loads(raw) + except Exception: + return None + + +def _next_device_name(user_id: str, hostname: str) -> str: + """Return a candidate name for the new device row. + + Adds a random 16-bit suffix to reduce collisions on the DB + ``UNIQUE (user_id, name)`` constraint. Collisions are still possible, + so ``redeem_pairing`` retries with a fresh candidate on conflict. + """ + base = hostname or "device" + base = base.strip() or "device" + return f"{base}-{os.urandom(2).hex()}" + + +def _upsert_remote_device_user_tool( + conn, *, user_id: str, device_id: str, + device_name: str, description: Optional[str], +) -> None: + """Create or update the ``user_tools`` row that fronts this device.""" + repo = UserToolsRepository(conn) + existing = _find_remote_device_tool_row(conn, user_id, device_id) + actions = [ + { + "name": "run_command", + "description": ( + f"Execute a shell command on the remote device '{device_name}'. " + f"{description or ''}".strip() + ), + "active": True, + # ``tool_executor`` consults ``RemoteDeviceTool.preview_requires_approval`` + # live for ``remote_device`` calls, so the static snapshot must + # stay neutral. The tool's own ``execute_action`` still yields + # the awaiting-approval pause when the live decision says so. + "require_approval": False, + "parameters": { + "type": "object", + "properties": { + "command": { + "type": "string", + "description": "Shell command to run.", + "filled_by_llm": True, + "value": "", + }, + "working_directory": { + "type": "string", + "description": "Working directory on the remote.", + "filled_by_llm": True, + "value": "", + }, + "timeout_ms": { + "type": "integer", + "description": "Timeout in milliseconds (max 600000).", + "filled_by_llm": True, + "value": "", + }, + }, + "required": ["command"], + }, + } + ] + config = {"device_id": device_id} + display_name = f"Remote device — {device_name}" + if existing: + repo.update( + str(existing["id"]), + user_id, + { + "display_name": display_name, + "description": description or "", + "config": config, + "actions": actions, + "status": True, + }, + ) + else: + repo.create( + user_id=user_id, + name="remote_device", + display_name=display_name, + description=description or "", + config=config, + actions=actions, + status=True, + ) + + +def _find_remote_device_tool_row(conn, user_id: str, device_id: str): + """Locate the user_tools row whose config.device_id matches.""" + from sqlalchemy import text + row = conn.execute( + text( + """ + SELECT * + FROM user_tools + WHERE user_id = :user_id + AND name = 'remote_device' + AND config ->> 'device_id' = :device_id + LIMIT 1 + """ + ), + {"user_id": user_id, "device_id": device_id}, + ).fetchone() + if row is None: + return None + return dict(row._mapping) + + +def _error(code: str, status: int) -> tuple: + return make_response(jsonify({"success": False, "error": code}), status) diff --git a/application/api/devices/routes.py b/application/api/devices/routes.py new file mode 100644 index 00000000..27ac2d61 --- /dev/null +++ b/application/api/devices/routes.py @@ -0,0 +1,327 @@ +"""Device CRUD + auto-approve pattern routes.""" + +from __future__ import annotations + +import logging + +from flask import Blueprint, jsonify, make_response, request + +from application.api.devices.pairing import ( + create_pairing, + delete_pairing, + get_pairing, + redeem_pairing, +) +from application.api.devices.session import ( + ack_invocation, + me, + poll, + session_events, + submit_output, +) +from application.devices.normalizer import normalize_command +from application.storage.db.repositories.device_audit_log import ( + DeviceAuditLogRepository, +) +from application.storage.db.repositories.device_auto_approve_patterns import ( + DeviceAutoApprovePatternsRepository, +) +from application.storage.db.repositories.devices import DevicesRepository +from application.storage.db.session import db_readonly, db_session + + +logger = logging.getLogger(__name__) + + +_ALLOWED_APPROVAL_MODES = {"ask", "full"} + + +def _authed_user_id(): + decoded = getattr(request, "decoded_token", None) + if isinstance(decoded, dict): + return decoded.get("sub") + return None + + +def _serialize_device(row: dict) -> dict: + """Strip server-internal fields before returning to the UI.""" + out = { + "id": row.get("id"), + "name": row.get("name"), + "hostname": row.get("hostname"), + "os": row.get("os"), + "arch": row.get("arch"), + "cli_version": row.get("cli_version"), + "approval_mode": row.get("approval_mode"), + "description": row.get("description"), + "status": row.get("status"), + "paired_at": row.get("paired_at"), + "last_seen_at": row.get("last_seen_at"), + "revoked_at": row.get("revoked_at"), + } + # JSON-ify datetimes if SQLAlchemy returned them raw. + for key in ("paired_at", "last_seen_at", "revoked_at"): + value = out.get(key) + if value is not None and not isinstance(value, str): + out[key] = value.isoformat() + return out + + +def list_devices(): + user_id = _authed_user_id() + if not user_id: + return make_response(jsonify({"success": False, "error": "auth"}), 401) + with db_readonly() as conn: + rows = DevicesRepository(conn).list_for_user(user_id) + return make_response( + jsonify({"devices": [_serialize_device(r) for r in rows]}), + 200, + ) + + +def get_device(device_id: str): + user_id = _authed_user_id() + if not user_id: + return make_response(jsonify({"success": False, "error": "auth"}), 401) + with db_readonly() as conn: + row = DevicesRepository(conn).get(device_id, user_id=user_id) + if row is None: + return make_response(jsonify({"success": False, "error": "not_found"}), 404) + return make_response(jsonify(_serialize_device(row)), 200) + + +def update_device(device_id: str): + user_id = _authed_user_id() + if not user_id: + return make_response(jsonify({"success": False, "error": "auth"}), 401) + body = request.get_json(silent=True) or {} + update_fields: dict = {} + if "name" in body: + name = body["name"] + if not isinstance(name, str) or not name.strip(): + return make_response( + jsonify({"success": False, "error": "invalid_name"}), 400 + ) + update_fields["name"] = name.strip() + if "description" in body: + description = body["description"] + if not isinstance(description, str): + return make_response( + jsonify({"success": False, "error": "invalid_description"}), 400 + ) + update_fields["description"] = description.strip() + if "approval_mode" in body: + mode = body["approval_mode"] + if mode not in _ALLOWED_APPROVAL_MODES: + return make_response( + jsonify({"success": False, "error": "invalid_approval_mode"}), 400 + ) + update_fields["approval_mode"] = mode + if not update_fields: + return make_response(jsonify({"success": False, "error": "no_fields"}), 400) + with db_session() as conn: + ok = DevicesRepository(conn).update(device_id, user_id, update_fields) + if not ok: + return make_response( + jsonify({"success": False, "error": "not_found"}), 404 + ) + # Reflect name/description into the user_tools row that fronts this device. + if "name" in update_fields or "description" in update_fields: + from application.api.devices.pairing import ( + _upsert_remote_device_user_tool, + ) + row = DevicesRepository(conn).get(device_id, user_id=user_id) + if row is not None: + _upsert_remote_device_user_tool( + conn, + user_id=user_id, + device_id=device_id, + device_name=row.get("name") or "device", + description=row.get("description"), + ) + row = DevicesRepository(conn).get(device_id, user_id=user_id) + return make_response(jsonify(_serialize_device(row)), 200) + + +def delete_device(device_id: str): + user_id = _authed_user_id() + if not user_id: + return make_response(jsonify({"success": False, "error": "auth"}), 401) + with db_session() as conn: + ok = DevicesRepository(conn).revoke(device_id, user_id) + if not ok: + return make_response( + jsonify({"success": False, "error": "not_found"}), 404 + ) + # Drop the corresponding user_tools row so the tool picker stops showing it. + from sqlalchemy import text + conn.execute( + text( + """ + DELETE FROM user_tools + WHERE user_id = :user_id + AND name = 'remote_device' + AND config ->> 'device_id' = :device_id + """ + ), + {"user_id": user_id, "device_id": device_id}, + ) + DeviceAutoApprovePatternsRepository(conn).clear_for_device(device_id) + return make_response(jsonify({"success": True}), 200) + + +def add_auto_approve_pattern(device_id: str): + """Server-side normalization: client sends the raw command, we store the pattern.""" + user_id = _authed_user_id() + if not user_id: + return make_response(jsonify({"success": False, "error": "auth"}), 401) + body = request.get_json(silent=True) or {} + command = body.get("command") or body.get("pattern") + if not command: + return make_response( + jsonify({"success": False, "error": "missing_command"}), 400 + ) + if not isinstance(command, str): + return make_response( + jsonify({"success": False, "error": "invalid_command"}), 400 + ) + pattern = normalize_command(command) + if not pattern: + return make_response( + jsonify({"success": False, "error": "invalid_command"}), 400 + ) + with db_session() as conn: + if DevicesRepository(conn).get(device_id, user_id=user_id) is None: + return make_response( + jsonify({"success": False, "error": "not_found"}), 404 + ) + DeviceAutoApprovePatternsRepository(conn).add(device_id, user_id, pattern) + return make_response(jsonify({"pattern": pattern}), 200) + + +def list_auto_approve_patterns(device_id: str): + user_id = _authed_user_id() + if not user_id: + return make_response(jsonify({"success": False, "error": "auth"}), 401) + with db_readonly() as conn: + if DevicesRepository(conn).get(device_id, user_id=user_id) is None: + return make_response( + jsonify({"success": False, "error": "not_found"}), 404 + ) + patterns = DeviceAutoApprovePatternsRepository(conn).list_for_device( + device_id, user_id + ) + return make_response(jsonify({"patterns": patterns}), 200) + + +def delete_auto_approve_pattern(device_id: str): + user_id = _authed_user_id() + if not user_id: + return make_response(jsonify({"success": False, "error": "auth"}), 401) + body = request.get_json(silent=True) or {} + pattern = body.get("pattern") + if not pattern: + return make_response( + jsonify({"success": False, "error": "missing_pattern"}), 400 + ) + with db_session() as conn: + ok = DeviceAutoApprovePatternsRepository(conn).remove( + device_id, user_id, pattern + ) + return make_response(jsonify({"success": ok}), 200) + + +def list_audit(device_id: str): + user_id = _authed_user_id() + if not user_id: + return make_response(jsonify({"success": False, "error": "auth"}), 401) + with db_readonly() as conn: + if DevicesRepository(conn).get(device_id, user_id=user_id) is None: + return make_response( + jsonify({"success": False, "error": "not_found"}), 404 + ) + rows = DeviceAuditLogRepository(conn).list_for_device(device_id, user_id) + # Normalize datetimes to ISO strings. + serialized = [] + for row in rows: + out = dict(row) + for key in ("issued_at", "started_at", "finished_at", "created_at"): + v = out.get(key) + if v is not None and not isinstance(v, str): + out[key] = v.isoformat() + serialized.append(out) + return make_response(jsonify({"entries": serialized}), 200) + + +def register(bp: Blueprint) -> None: + """Attach all device routes to ``bp``.""" + bp.add_url_rule( + "/api/devices", view_func=list_devices, methods=["GET"], endpoint="list_devices", + ) + bp.add_url_rule( + "/api/devices/", view_func=get_device, methods=["GET"], + endpoint="get_device", + ) + bp.add_url_rule( + "/api/devices/", view_func=update_device, methods=["PATCH"], + endpoint="update_device", + ) + bp.add_url_rule( + "/api/devices/", view_func=delete_device, methods=["DELETE"], + endpoint="delete_device", + ) + bp.add_url_rule( + "/api/devices//auto-approve", + view_func=add_auto_approve_pattern, methods=["POST"], + endpoint="add_auto_approve_pattern", + ) + bp.add_url_rule( + "/api/devices//auto-approve", + view_func=list_auto_approve_patterns, methods=["GET"], + endpoint="list_auto_approve_patterns", + ) + bp.add_url_rule( + "/api/devices//auto-approve", + view_func=delete_auto_approve_pattern, methods=["DELETE"], + endpoint="delete_auto_approve_pattern", + ) + bp.add_url_rule( + "/api/devices//audit", view_func=list_audit, methods=["GET"], + endpoint="list_audit", + ) + # Pairing endpoints + bp.add_url_rule( + "/api/devices/pairings", view_func=create_pairing, methods=["POST"], + endpoint="create_pairing", + ) + bp.add_url_rule( + "/api/devices/pairings/redeem", view_func=redeem_pairing, methods=["POST"], + endpoint="redeem_pairing", + ) + bp.add_url_rule( + "/api/devices/pairings/", view_func=get_pairing, methods=["GET"], + endpoint="get_pairing", + ) + bp.add_url_rule( + "/api/devices/pairings/", view_func=delete_pairing, + methods=["DELETE"], endpoint="delete_pairing", + ) + # Session endpoints (device-token auth) + bp.add_url_rule( + "/api/devices/poll", view_func=poll, methods=["GET"], endpoint="device_poll", + ) + bp.add_url_rule( + "/api/devices/me", view_func=me, methods=["GET"], endpoint="device_me", + ) + bp.add_url_rule( + "/api/devices/sessions//events", view_func=session_events, + methods=["GET"], endpoint="device_session_events", + ) + bp.add_url_rule( + "/api/devices/sessions//invocations//ack", + view_func=ack_invocation, methods=["POST"], endpoint="device_ack", + ) + bp.add_url_rule( + "/api/devices/sessions//invocations//output", + view_func=submit_output, methods=["POST"], endpoint="device_output", + ) diff --git a/application/api/devices/session.py b/application/api/devices/session.py new file mode 100644 index 00000000..02530d15 --- /dev/null +++ b/application/api/devices/session.py @@ -0,0 +1,243 @@ +"""Device session endpoints: poll, SSE, ack, output.""" + +from __future__ import annotations + +import gzip +import io +import json +import logging +import time +from queue import Empty +from typing import Iterator + +from flask import Response, jsonify, make_response, request, stream_with_context + +from application.api.devices.auth import verify_device_session +from application.core.settings import settings +from application.devices.broker import get_broker +from application.storage.db.repositories.device_audit_log import ( + DeviceAuditLogRepository, +) +from application.storage.db.session import db_session + + +logger = logging.getLogger(__name__) + + +# Window (seconds) the CLI has to upgrade a poll-issued ticket to an SSE +# stream. Advertised to the CLI as ``expires_in`` and used as the broker +# ticket TTL so the two never drift. +_SESSION_TICKET_TTL_SECONDS = 30 + + +def poll() -> Response: + """``GET /api/devices/poll`` — long-poll for queued invocations. + + Returns ``202`` with empty body when nothing is queued and ``200`` with + a session ticket payload when work is waiting. + """ + device, err = verify_device_session() + if err is not None: + return err + broker = get_broker() + ticket = broker.claim_ticket(device["id"], _SESSION_TICKET_TTL_SECONDS) + if ticket is None: + return make_response("", 202) + return make_response( + jsonify( + { + "session_ticket": ticket, + "session_url": f"/api/devices/sessions/{ticket}/events", + "expires_in": _SESSION_TICKET_TTL_SECONDS, + } + ), + 200, + ) + + +def me() -> Response: + """``GET /api/devices/me`` — return the calling device's own record. + + Auth: device session token (same as ``/poll``). Used by ``docsgpt-cli + host status`` to show live server state. + """ + device, err = verify_device_session() + if err is not None: + return err + out = { + "id": device.get("id"), + "name": device.get("name"), + "hostname": device.get("hostname"), + "os": device.get("os"), + "status": device.get("status"), + "approval_mode": device.get("approval_mode"), + "description": device.get("description"), + "paired_at": device.get("paired_at"), + "last_seen_at": device.get("last_seen_at"), + } + for key in ("paired_at", "last_seen_at"): + value = out.get(key) + if value is not None and not isinstance(value, str): + out[key] = value.isoformat() + return make_response(jsonify(out), 200) + + +def session_events(session_id: str) -> Response: + """``GET /api/devices/sessions/{id}/events`` — SSE invocation stream. + + The ``session_id`` must be the ``session_ticket`` the device's own + ``/poll`` just issued (the path it was handed as ``session_url``). A + stale, mismatched, or fabricated ticket is rejected with ``410 Gone`` + before any stream is opened. + """ + device, err = verify_device_session() + if err is not None: + return err + broker = get_broker() + if not broker.validate_ticket(device["id"], session_id): + return make_response( + jsonify({"success": False, "error": "session_ticket_invalid"}), 410 + ) + sess = broker.register_session(device["id"], device["user_id"]) + + keepalive_interval = float(settings.SSE_KEEPALIVE_SECONDS) + idle_seconds = float(settings.REMOTE_DEVICE_SESSION_IDLE_SECONDS) + + @stream_with_context + def generate() -> Iterator[str]: + try: + last_keepalive = time.time() + while not sess.closed.is_set(): + now = time.time() + if now - sess.last_activity_at > idle_seconds: + yield _sse_event( + "session_end", + {"reason": "inactivity_timeout"}, + sess.last_event_id + 1, + ) + sess.last_event_id += 1 + broker.close_session(sess.session_id, reason="idle") + return + try: + inv = sess.invocation_queue.get(timeout=1.0) + except Empty: + if time.time() - last_keepalive >= keepalive_interval: + last_keepalive = time.time() + yield ": heartbeat\n\n" + continue + sess.last_event_id += 1 + sess.last_activity_at = time.time() + yield _sse_event("invocation", inv.envelope, sess.last_event_id) + last_keepalive = time.time() + except GeneratorExit: + logger.debug("device SSE generator exiting for session %s", sess.session_id) + raise + finally: + broker.close_session(sess.session_id, reason="generator_exit") + + return Response( + generate(), + mimetype="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "X-Accel-Buffering": "no", + }, + ) + + +def ack_invocation(session_id: str, invocation_id: str) -> Response: + """CLI acks (accepted / denied / auto_approved) the invocation.""" + device, err = verify_device_session() + if err is not None: + return err + body = request.get_json(silent=True) or {} + decision = body.get("decision") + reason = body.get("reason") + if decision not in {"accepted", "denied", "auto_approved"}: + return make_response( + jsonify({"success": False, "error": "invalid_decision"}), 400 + ) + broker = get_broker() + inv = broker.get_invocation(invocation_id) + if inv is None or inv.device_id != device["id"]: + return make_response( + jsonify({"success": False, "error": "invocation_not_found"}), 404 + ) + broker.submit_ack(invocation_id, decision, reason) + return make_response(jsonify({"success": True}), 200) + + +def submit_output(session_id: str, invocation_id: str) -> Response: + """CLI streams stdout/stderr/control chunks (NDJSON, gzip-aware).""" + device, err = verify_device_session() + if err is not None: + return err + broker = get_broker() + inv = broker.get_invocation(invocation_id) + if inv is None or inv.device_id != device["id"]: + return make_response( + jsonify({"success": False, "error": "invocation_not_found"}), 404 + ) + + body = request.get_data() or b"" + if request.headers.get("Content-Encoding", "").lower() == "gzip": + try: + body = gzip.decompress(body) + except OSError: + return make_response( + jsonify({"success": False, "error": "invalid_gzip"}), 400 + ) + + received = 0 + for line in io.BytesIO(body): + line = line.strip() + if not line: + continue + try: + chunk = json.loads(line.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError): + continue + if not isinstance(chunk, dict): + continue + broker.submit_output_chunk(invocation_id, chunk) + received += 1 + + # Persist outcome on the audit row when a control chunk closed the stream. + if inv.completed.is_set(): + try: + from datetime import datetime, timezone + with db_session() as conn: + DeviceAuditLogRepository(conn).record_result( + invocation_id, + started_at=( + datetime.fromtimestamp(inv.started_at, tz=timezone.utc) + if inv.started_at else None + ), + finished_at=( + datetime.fromtimestamp(inv.finished_at, tz=timezone.utc) + if inv.finished_at else None + ), + exit_code=inv.exit_code, + duration_ms=inv.duration_ms, + stdout_bytes=sum( + len(c) for c in inv.stdout_parts + ), + stderr_bytes=sum( + len(c) for c in inv.stderr_parts + ), + error=inv.error, + ) + except Exception: + logger.exception("audit record_result failed for %s", invocation_id) + + return make_response( + jsonify({"success": True, "received": received}), 200 + ) + + +def _sse_event(name: str, payload: dict, event_id: int) -> str: + return ( + f"event: {name}\n" + f"id: {event_id}\n" + f"data: {json.dumps(payload)}\n\n" + ) diff --git a/application/app.py b/application/app.py index 6956d06d..0d3f37bc 100644 --- a/application/app.py +++ b/application/app.py @@ -17,6 +17,7 @@ setup_logging() from application.api import api # noqa: E402 from application.api.answer import answer # noqa: E402 from application.api.answer.routes.messages import messages_bp # noqa: E402 +from application.api.devices import devices_bp # noqa: E402 from application.api.events.routes import events # noqa: E402 from application.api.internal.routes import internal # noqa: E402 from application.api.user.routes import user # noqa: E402 @@ -61,6 +62,7 @@ app.register_blueprint(events) app.register_blueprint(messages_bp) app.register_blueprint(internal) app.register_blueprint(connector) +app.register_blueprint(devices_bp) app.register_blueprint(v1_bp) app.config.update( UPLOAD_FOLDER="inputs", @@ -182,6 +184,19 @@ def authenticate_request(): if request.path.startswith("/v1/"): request.decoded_token = None return None + # Remote-device CLI endpoints carry opaque ``tok_…`` session tokens + # (not JWTs); ``verify_device_session`` runs inside the route handler. + # The redeem endpoint is tokenless — it authenticates via the one-time + # ``user_code`` inside ``redeem_pairing`` — so it's exempt too. Pairing + # create + status stay JWT-protected (UI calls). + if ( + request.path.startswith("/api/devices/poll") + or request.path.startswith("/api/devices/sessions/") + or request.path == "/api/devices/me" + or request.path == "/api/devices/pairings/redeem" + ): + request.decoded_token = None + return None decoded_token = handle_auth(request) if not decoded_token: request.decoded_token = None diff --git a/application/core/settings.py b/application/core/settings.py index b4373f2b..b93dbbf0 100644 --- a/application/core/settings.py +++ b/application/core/settings.py @@ -237,6 +237,11 @@ class Settings(BaseSettings): # flows without unbounded table growth. MESSAGE_EVENTS_RETENTION_DAYS: int = 14 + # Remote Device feature. + REMOTE_DEVICE_SESSION_IDLE_SECONDS: int = 60 + REMOTE_DEVICE_REQUIRE_SIGNATURE: bool = False + REMOTE_DEVICE_PAIRING_TTL_SECONDS: int = 600 + # Scheduler (see scheduler.md). SCHEDULE_DISPATCHER_INTERVAL: int = 30 SCHEDULE_MIN_INTERVAL: int = 900 diff --git a/application/devices/__init__.py b/application/devices/__init__.py new file mode 100644 index 00000000..57f4066c --- /dev/null +++ b/application/devices/__init__.py @@ -0,0 +1 @@ +"""Remote device backend: broker, safety checks, and pairing state.""" diff --git a/application/devices/broker.py b/application/devices/broker.py new file mode 100644 index 00000000..52305ee8 --- /dev/null +++ b/application/devices/broker.py @@ -0,0 +1,344 @@ +"""In-process broker that routes invocations to active device sessions. + +Single-worker assumption (per spec 13.1). Each device has at most one +active SSE session; concurrent invocations queue and drain in FIFO order. + +Lifecycle: +1. Agent calls ``RemoteDeviceTool.execute_action`` → ``broker.dispatch_invocation``. +2. If a session is active, the envelope goes on the session's queue; the + SSE handler reads it and emits to the wire. Otherwise it sits on the + pending-invocation queue until ``register_session`` is called by the + next ``GET /api/devices/poll`` -> SSE upgrade. +3. CLI POSTs ack and chunked output back. Output chunks land in the + invocation's stdout/stderr queue, which the tool's generator drains + to yield streaming events. +4. Final ``control`` chunk closes the invocation. The tool resolves with + the aggregated result. +""" + +from __future__ import annotations + +import logging +import threading +import time +import uuid +from dataclasses import dataclass, field +from queue import Empty, Queue +from typing import Any, Dict, List, Optional + + +logger = logging.getLogger(__name__) + + +# Sentinel object signalling the end of an invocation's output stream. +INVOCATION_DONE = object() + + +@dataclass +class InvocationState: + """In-flight invocation: envelope + back-streaming queues.""" + + invocation_id: str + device_id: str + session_id: str + envelope: Dict[str, Any] + output_queue: Queue = field(default_factory=Queue) + stdout_parts: List[str] = field(default_factory=list) + stderr_parts: List[str] = field(default_factory=list) + exit_code: Optional[int] = None + duration_ms: Optional[int] = None + error: Optional[str] = None + decision: Optional[str] = None + decision_reason: Optional[str] = None + started_at: Optional[float] = None + finished_at: Optional[float] = None + completed: threading.Event = field(default_factory=threading.Event) + + +@dataclass +class SessionState: + """Active SSE session for one device.""" + + session_id: str + device_id: str + user_id: str + invocation_queue: Queue = field(default_factory=Queue) + last_event_id: int = 0 + invocations: Dict[str, InvocationState] = field(default_factory=dict) + last_activity_at: float = field(default_factory=time.time) + closed: threading.Event = field(default_factory=threading.Event) + + +class DeviceBroker: + """In-process device registry; thread-safe.""" + + def __init__(self) -> None: + self._lock = threading.Lock() + # device_id -> SessionState (active SSE). + self._sessions: Dict[str, SessionState] = {} + # session_id -> device_id (reverse map for fast session-by-id lookup). + self._session_by_id: Dict[str, str] = {} + # device_id -> pending invocations (no active session yet). + self._pending: Dict[str, List[InvocationState]] = {} + # invocation_id -> InvocationState for ack/output lookup. + self._invocations: Dict[str, InvocationState] = {} + # device_id -> (session ticket, monotonic expiry) for the next poll. + self._tickets: Dict[str, tuple[str, float]] = {} + + # ------------------------------------------------------------------ + # Session lifecycle (CLI side) + # ------------------------------------------------------------------ + def register_session(self, device_id: str, user_id: str) -> SessionState: + """Open or refresh the active SSE session for ``device_id``. + + Drains any pending invocations into the session's queue so the + next SSE read sees them in order. If a previous session is still + registered, it is closed and replaced (the CLI reconnected). + """ + with self._lock: + existing = self._sessions.get(device_id) + if existing is not None: + existing.closed.set() + self._session_by_id.pop(existing.session_id, None) + issued = self._tickets.pop(device_id, None) + session_id = issued[0] if issued else f"st_{uuid.uuid4().hex}" + sess = SessionState( + session_id=session_id, + device_id=device_id, + user_id=user_id, + ) + for inv in self._pending.pop(device_id, []): + inv.session_id = session_id + sess.invocations[inv.invocation_id] = inv + sess.invocation_queue.put(inv) + self._sessions[device_id] = sess + self._session_by_id[session_id] = device_id + return sess + + def close_session(self, session_id: str, *, reason: str = "closed") -> None: + """Close the session by id; outstanding invocations get an error.""" + with self._lock: + device_id = self._session_by_id.pop(session_id, None) + if device_id is None: + return + sess = self._sessions.pop(device_id, None) + if sess is None: + return + sess.closed.set() + for inv in sess.invocations.values(): + if not inv.completed.is_set(): + inv.error = inv.error or f"session_{reason}" + inv.output_queue.put(INVOCATION_DONE) + inv.completed.set() + logger.debug("device session closed: %s (%s)", session_id, reason) + + def get_session(self, session_id: str) -> Optional[SessionState]: + with self._lock: + device_id = self._session_by_id.get(session_id) + if device_id is None: + return None + return self._sessions.get(device_id) + + # ------------------------------------------------------------------ + # Polling + # ------------------------------------------------------------------ + def claim_ticket(self, device_id: str, ttl_seconds: float) -> Optional[str]: + """Return a session ticket if the device has work waiting, else None. + + A pending invocation queued before the CLI polls allocates a ticket + eagerly; the CLI then upgrades to SSE with that ticket within + ``ttl_seconds``. The ticket is remembered (with its expiry) so the SSE + ``session_id`` can be validated against the one this poll issued. + """ + now = time.monotonic() + with self._lock: + if not self._pending.get(device_id): + return None + existing = self._tickets.get(device_id) + if existing and existing[1] > now: + return existing[0] + ticket = f"st_{uuid.uuid4().hex}" + self._tickets[device_id] = (ticket, now + ttl_seconds) + return ticket + + def validate_ticket(self, device_id: str, session_id: str) -> bool: + """True iff ``session_id`` is the unexpired ticket issued to ``device_id``. + + Gates the SSE upgrade: the CLI must open the events stream with the + exact ``session_ticket`` returned by its ``/poll`` (which became the + ``session_url``). An absent, mismatched, or expired ticket is rejected. + Expired entries are evicted so a stale ticket can't be reused. + """ + if not session_id: + return False + now = time.monotonic() + with self._lock: + issued = self._tickets.get(device_id) + if issued is None: + return False + ticket, expires_at = issued + if expires_at <= now: + self._tickets.pop(device_id, None) + return False + return session_id == ticket + + # ------------------------------------------------------------------ + # Dispatch (server-issued) + # ------------------------------------------------------------------ + def dispatch_invocation( + self, + device_id: str, + user_id: str, + envelope: Dict[str, Any], + ) -> InvocationState: + """Enqueue an invocation for ``device_id``. + + If there's an active session, drops it onto the session's + invocation queue. Otherwise parks it in the pending list and + the next poll will see a ticket. + """ + invocation_id = envelope["invocation_id"] + with self._lock: + sess = self._sessions.get(device_id) + session_id = sess.session_id if sess else "" + inv = InvocationState( + invocation_id=invocation_id, + device_id=device_id, + session_id=session_id, + envelope=envelope, + ) + self._invocations[invocation_id] = inv + if sess is not None: + sess.invocations[invocation_id] = inv + sess.invocation_queue.put(inv) + sess.last_activity_at = time.time() + else: + self._pending.setdefault(device_id, []).append(inv) + return inv + + def get_invocation(self, invocation_id: str) -> Optional[InvocationState]: + with self._lock: + return self._invocations.get(invocation_id) + + # ------------------------------------------------------------------ + # Output streaming (CLI side) + # ------------------------------------------------------------------ + def submit_output_chunk( + self, + invocation_id: str, + chunk: Dict[str, Any], + ) -> bool: + """Forward a chunk from the CLI to the agent stream's generator.""" + inv = self.get_invocation(invocation_id) + if inv is None: + return False + with self._lock: + sess = self._sessions.get(inv.device_id) + if sess is not None: + sess.last_activity_at = time.time() + stream = chunk.get("stream") + if stream == "stdout": + inv.stdout_parts.append(chunk.get("chunk", "")) + elif stream == "stderr": + inv.stderr_parts.append(chunk.get("chunk", "")) + elif stream == "control": + inv.exit_code = chunk.get("exit_code") + inv.duration_ms = chunk.get("duration_ms") + inv.error = chunk.get("error") or inv.error + if inv.started_at is None: + inv.started_at = time.time() + inv.finished_at = time.time() + inv.output_queue.put(chunk) + if chunk.get("stream") == "control": + inv.output_queue.put(INVOCATION_DONE) + inv.completed.set() + return True + + def submit_ack( + self, + invocation_id: str, + decision: str, + reason: Optional[str] = None, + ) -> bool: + inv = self.get_invocation(invocation_id) + if inv is None: + return False + with self._lock: + inv.decision = decision + inv.decision_reason = reason + inv.started_at = inv.started_at or time.time() + if decision == "denied": + inv.error = inv.error or "denied" + inv.output_queue.put(INVOCATION_DONE) + inv.completed.set() + return True + + def drain_output( + self, + invocation_id: str, + timeout: float = 0.5, + deadline: Optional[float] = None, + ): + """Yield queued chunks for an invocation; stops on INVOCATION_DONE. + + Used by the tool's generator. Blocks up to ``timeout`` per chunk + so the caller can interleave with other work. ``deadline`` is an + absolute ``time.time()`` value; once it passes with no output and + the invocation never completed, the generator returns so a device + that never connects can't loop forever. + """ + inv = self.get_invocation(invocation_id) + if inv is None: + return + while True: + try: + item = inv.output_queue.get(timeout=timeout) + except Empty: + if inv.completed.is_set(): + return + if deadline is not None and time.time() >= deadline: + return + continue + if item is INVOCATION_DONE: + return + yield item + + def cleanup_invocation(self, invocation_id: str) -> None: + """Drop an invocation from all registries. + + Removes it from ``_invocations``, the active session's + ``invocations`` map, and — if it was still parked for an offline + device — ``_pending`` too. Dropping it from the active session keeps + a long-lived SSE session from accumulating completed invocations + (unbounded memory). Without the pending removal a timed-out + invocation would be drained and executed by a later + ``register_session`` — running a command the user already saw fail. + """ + with self._lock: + inv = self._invocations.pop(invocation_id, None) + if inv is None: + return + sess = self._sessions.get(inv.device_id) + if sess is not None: + sess.invocations.pop(invocation_id, None) + pending = self._pending.get(inv.device_id) + if pending: + remaining = [p for p in pending if p.invocation_id != invocation_id] + if remaining: + self._pending[inv.device_id] = remaining + else: + self._pending.pop(inv.device_id, None) + + +_broker_instance: Optional[DeviceBroker] = None +_broker_lock = threading.Lock() + + +def get_broker() -> DeviceBroker: + """Return the process-wide ``DeviceBroker`` instance.""" + global _broker_instance + if _broker_instance is None: + with _broker_lock: + if _broker_instance is None: + _broker_instance = DeviceBroker() + return _broker_instance diff --git a/application/devices/denylist.py b/application/devices/denylist.py new file mode 100644 index 00000000..32192bb5 --- /dev/null +++ b/application/devices/denylist.py @@ -0,0 +1,172 @@ +"""Hard denylist — always triggers a forced prompt, even under ``never``.""" + +from __future__ import annotations + +import re +import shlex +from typing import List, Optional, Tuple + +from application.devices.splitter import split_command, strip_wrappers + + +# Label/reason for the rm-root wipe, shared by the regex rule and the +# token-based helper so both report identically. +_RM_ROOT_LABEL = "rm -rf /" + +# Filesystem-root targets that, with recursive+force, mean a full wipe. +_RM_ROOT_TARGETS = frozenset({"/", "/*"}) + + +# Each entry: (label, compiled regex). Regex is matched against each +# segment's whole text (case-insensitive, after whitespace normalization). +_PATTERNS: List[Tuple[str, re.Pattern]] = [ + ( + _RM_ROOT_LABEL, + # ``rm`` with recursive+force targeting the filesystem root. Flags + # (incl. ``--no-preserve-root``) may appear in any order before or + # after the recursive/force flags; the target must be exactly ``/`` + # or ``/*`` so safe paths (``/tmp/foo``, ``./build``, ``/home/x``) + # don't match. ``(?:-\S+\s+)*`` absorbs any extra leading options. + re.compile( + r"\brm\s+(?:-\S+\s+)*" + r"(?:-[a-zA-Z]*r[a-zA-Z]*f|-[a-zA-Z]*f[a-zA-Z]*r" + r"|--recursive\s+(?:-\S+\s+)*--force|--force\s+(?:-\S+\s+)*--recursive)" + r"\s+(?:-\S+\s+)*/\*?\s*(?:$|\s)" + ), + ), + ( + "rm -rf ~", + re.compile( + r"\brm\s+(-[a-zA-Z]*r[a-zA-Z]*f|-[a-zA-Z]*f[a-zA-Z]*r|-rf|-fr)\s+~" + ), + ), + ( + "rm -rf $HOME", + re.compile( + r"\brm\s+(-[a-zA-Z]*r[a-zA-Z]*f|-[a-zA-Z]*f[a-zA-Z]*r|-rf|-fr)\s+\$HOME\b" + ), + ), + ( + "fork bomb", + re.compile(r":\(\)\s*\{\s*:\s*\|\s*:\s*&\s*\}\s*;\s*:"), + ), + ( + "dd to block device", + re.compile( + r"\bdd\s+.*\bif=/dev/(zero|random|urandom).*\bof=/dev/(sd|nvme|hd|disk|mmcblk)", + re.DOTALL, + ), + ), + ( + "mkfs", + re.compile(r"\bmkfs(\.[a-zA-Z0-9]+)?\b"), + ), + ( + "shutdown", + re.compile(r"\bshutdown\b"), + ), + ( + "halt", + re.compile(r"\bhalt\b"), + ), + ( + "poweroff", + re.compile(r"\bpoweroff\b"), + ), + ( + "init 0/6", + re.compile(r"\binit\s+(0|6)\b"), + ), + ( + "git push --force", + # ``--force-with-lease`` is OK; treat ``-f`` short form as denied. + re.compile( + r"\bgit\s+push\b.*(--force(?!-with-lease)|--mirror|\s-f(?:\s|$))" + ), + ), +] + + +def _is_rm_root_wipe(segment: str) -> bool: + """True if ``segment`` is ``rm`` with recursive+force targeting root. + + Token-based so separated flags (``rm -r -f /``, ``rm -f -r /``) and + bundles (``-rf``, ``-fr``, ``-rfv``) are all caught regardless of order. + Stays conservative: requires BOTH recursive and force AND a root target + (``/`` or ``/*``), so safe paths (``/tmp/foo``, ``./build``) and + single-flag forms (``rm -r /tmp``, ``rm -f /etc/x``) are not denied. + + Args: + segment: A single shell command segment (no compound connectors). + + Returns: + ``True`` if the segment is a recursive+force wipe of filesystem root. + """ + try: + tokens = shlex.split(segment, posix=True) + except ValueError: + # Unbalanced quotes — let the regex rules handle it; don't crash. + return False + if not tokens or tokens[0] != "rm": + return False + + recursive = False + force = False + targets: List[str] = [] + for tok in tokens[1:]: + if tok in ("--recursive", "-R"): + recursive = True + elif tok == "--force": + force = True + elif tok == "--no-preserve-root": + # A flag, not a target; ignore for target detection. + continue + elif tok.startswith("--"): + # Unknown long option — neither recursive/force nor a target. + continue + elif tok.startswith("-") and len(tok) > 1: + # Short flag bundle, e.g. ``-rf``, ``-fr``, ``-rfv``, ``-r``, ``-f``. + letters = tok[1:] + if "r" in letters: + recursive = True + if "f" in letters: + force = True + else: + targets.append(tok) + + if not (recursive and force): + return False + return any(t in _RM_ROOT_TARGETS for t in targets) + + +def check_denylist(command: str) -> Optional[str]: + """Return the matched pattern label if ``command`` hits the hard denylist. + + Splits the command into segments first (see ``splitter.py``) so + ``echo safe && rm -rf /`` still trips. Returns ``None`` if no segment + matches. + + Args: + command: Raw shell command string. + + Returns: + The human-readable pattern label, or ``None`` if no match. + """ + if not command: + return None + for segment in split_command(command): + # ``split_command`` keeps wrappers (``timeout 5 rm -rf /``, + # ``nohup rm -rf /``); strip them so the token-based rm check sees + # the real inner command and can't be evaded by a wrapper. + if _is_rm_root_wipe(strip_wrappers(segment)) or _is_rm_root_wipe(segment): + return _RM_ROOT_LABEL + for label, pattern in _PATTERNS: + if pattern.search(segment): + return label + # Also try matching the whole command as a single string, so a + # multi-line / unusual-whitespace fork bomb that splits weirdly still + # gets caught. + for label, pattern in _PATTERNS: + if pattern.search(command): + return label + return None diff --git a/application/devices/normalizer.py b/application/devices/normalizer.py new file mode 100644 index 00000000..393ab068 --- /dev/null +++ b/application/devices/normalizer.py @@ -0,0 +1,84 @@ +"""Normalize a command segment to its sticky-approval pattern. + +Rule: head + first sub-token (when present), wildcard the rest. + ``git checkout main -- file.txt`` -> ``git checkout *`` + ``npm install foo bar`` -> ``npm install *`` + ``ls -la /tmp`` -> ``ls *`` + ``cat /etc/passwd`` -> ``cat *`` + +Compound commands are normalized segment-by-segment and rejoined with +`` && ``, so a sticky only matches a command with the same normalized +segment sequence: + ``ls /tmp && whoami`` -> ``ls * && whoami`` +""" + +from __future__ import annotations + +import shlex +from typing import Optional + +from application.devices.splitter import split_command + + +# Commands whose first arg is a sub-command (so the pattern is "cmd subcmd *"). +_TWO_WORD_HEADS = frozenset( + { + "git", "npm", "pnpm", "yarn", "pip", "uv", "poetry", "cargo", + "docker", "kubectl", "brew", "apt", "apt-get", "dnf", "yum", + "go", "rustup", "rustc", "make", "gradle", "mvn", + "systemctl", "service", + } +) + + +def normalize_segment(segment: str) -> str: + """Reduce a segment to its sticky-pattern form. + + Args: + segment: A single shell command segment (no compound connectors). + + Returns: + Pattern like ``git checkout *``, ``ls *``, or just ``cat`` if the + original had no args. Empty string for empty input. + """ + if not segment or not segment.strip(): + return "" + try: + tokens = shlex.split(segment.strip(), posix=True) + except ValueError: + tokens = segment.strip().split() + if not tokens: + return "" + + head = tokens[0] + + if head in _TWO_WORD_HEADS and len(tokens) >= 2: + subcmd = tokens[1] + if len(tokens) > 2: + return f"{head} {subcmd} *" + return f"{head} {subcmd}" + + if len(tokens) > 1: + return f"{head} *" + return head + + +def normalize_command(command: str) -> Optional[str]: + """Normalize an entire command into its joined sticky pattern. + + Every segment is normalized and the results are joined with `` && ``, + so a compound command's sticky only matches a command with the same + normalized segment sequence (approving ``ls /tmp && whoami`` won't + auto-approve ``ls /tmp && rm /tmp/x``). Empty segments are skipped; + returns ``None`` if there are no non-empty segments. + """ + if not command: + return None + patterns = [ + normalize_segment(segment) + for segment in split_command(command) + ] + patterns = [p for p in patterns if p] + if not patterns: + return None + return " && ".join(patterns) diff --git a/application/devices/splitter.py b/application/devices/splitter.py new file mode 100644 index 00000000..0dcf02a8 --- /dev/null +++ b/application/devices/splitter.py @@ -0,0 +1,110 @@ +"""Compound-command splitter for safety evaluation.""" + +from __future__ import annotations + +import shlex +import re +from typing import List + + +# Splitters from spec 9.4: && || ; | & |& \n +# Order matters: longest first so e.g. ``|&`` doesn't get caught by ``|``. +_SPLITTER_RE = re.compile(r"\|\&|\&\&|\|\||;|\||\&|\n") + +# Wrappers stripped before extracting the head token. +_WRAPPERS = frozenset({"timeout", "nice", "nohup", "stdbuf", "xargs", "env"}) + + +def split_command(command: str) -> List[str]: + """Split a command string on shell connectors. + + Each returned segment is a single "logical" command; the segments are + independently evaluated against the hard denylist. ``env`` is + intentionally in ``_WRAPPERS`` ONLY at the head position. + + Args: + command: Raw shell command string. + + Returns: + List of trimmed, non-empty segments. + """ + if not command: + return [] + parts = _SPLITTER_RE.split(command) + return [p.strip() for p in parts if p and p.strip()] + + +def strip_wrappers(segment: str) -> str: + """Return ``segment`` with leading command wrappers removed. + + ``timeout 30s rm -rf /`` → ``rm -rf /``; ``nice -n 5 rm -rf /`` → + ``rm -rf /``. Wrappers and their own arguments (e.g. ``timeout 30s``, + ``nice -n 5``) are dropped so safety checks see the real inner command. + ``env`` is intentionally left in place so the denylist sees the head + token rather than env assignments. Returns the original segment when + there is nothing to strip or on a parse failure. + """ + if not segment: + return "" + try: + tokens = shlex.split(segment, posix=True) + except ValueError: + # Unbalanced quotes — let callers fall back to the raw segment. + return segment + if not tokens: + return "" + i = _wrapper_offset(tokens) + if i == 0: + return segment + return " ".join(tokens[i:]) + + +def _wrapper_offset(tokens: List[str]) -> int: + """Index of the first non-wrapper token (skipping wrapper args). + + ``env`` is treated as a non-wrapper so it is never stripped. + """ + i = 0 + while i < len(tokens) and tokens[i] in _WRAPPERS and tokens[i] != "env": + current = tokens[i] + i += 1 + if i >= len(tokens): + break + # Skip wrapper-specific arguments before the wrapped command. + if current == "timeout": + # ``timeout 30`` or ``timeout 30s`` -- skip the duration token. + if re.match(r"^\d+[smhd]?$", tokens[i]): + i += 1 + elif current == "nice": + # ``nice -n 10 cmd`` -- skip ``-n`` and its value. + if tokens[i].startswith("-n"): + i += 1 + if i < len(tokens) and re.match(r"^-?\d+$", tokens[i]): + i += 1 + return i + + +def head_token(segment: str) -> str: + """Extract the head command token from a segment, stripping wrappers. + + ``timeout 30s ls`` → ``ls``. ``env X=Y cmd`` is NOT stripped here. + """ + if not segment: + return "" + try: + tokens = shlex.split(segment, posix=True) + except ValueError: + # Unbalanced quotes — fall back to a naive split. Don't crash. + tokens = segment.split() + if not tokens: + return "" + i = _wrapper_offset(tokens) + if i >= len(tokens): + # Segment was only wrappers/args; nothing meaningful to return. + return tokens[0] + return tokens[i] + + +def head_tokens(command: str) -> List[str]: + """Return the head token of every segment.""" + return [head_token(seg) for seg in split_command(command)] diff --git a/application/llm/handlers/base.py b/application/llm/handlers/base.py index 45959e21..8466b8dd 100644 --- a/application/llm/handlers/base.py +++ b/application/llm/handlers/base.py @@ -938,16 +938,18 @@ class LLMHandler(ABC): } continue # Yield pause event so the client knows this tool is waiting - yield { - "type": "tool_call", - "data": { - "tool_name": pause_info["tool_name"], - "call_id": pause_info["call_id"], - "action_name": pause_info.get("llm_name", pause_info["name"]), - "arguments": pause_info["arguments"], - "status": pause_info["pause_type"], - }, + pause_data = { + "tool_name": pause_info["tool_name"], + "call_id": pause_info["call_id"], + "action_name": pause_info.get("llm_name", pause_info["name"]), + "arguments": pause_info["arguments"], + "status": pause_info["pause_type"], } + # Surface device_id for remote_device pauses so the approval UI + # can wire the sticky "don't ask again" button. + if pause_info.get("device_id"): + pause_data["device_id"] = pause_info["device_id"] + yield {"type": "tool_call", "data": pause_data} pending_actions.append(pause_info) # Do NOT add messages for pending tools here. # They will be added on resume to keep call/result pairs together. diff --git a/application/storage/db/models.py b/application/storage/db/models.py index 7b04b599..81f2c7e3 100644 --- a/application/storage/db/models.py +++ b/application/storage/db/models.py @@ -18,6 +18,7 @@ declarations to keep this file readable; the DB is the authority. from sqlalchemy import ( BigInteger, Boolean, + CHAR, Column, DateTime, ForeignKey, @@ -676,3 +677,68 @@ schedule_runs_table = Table( Column("updated_at", DateTime(timezone=True), nullable=False, server_default=func.now()), UniqueConstraint("schedule_id", "scheduled_for", name="schedule_runs_dedup_uidx"), ) + + +# --- Remote devices (migration 0012) ---------------------------------------- + +devices_table = Table( + "devices", + metadata, + Column("id", Text, primary_key=True), + Column("user_id", Text, nullable=False), + Column("name", Text, nullable=False), + Column("hostname", Text), + Column("os", Text), + Column("arch", Text), + Column("cli_version", Text), + Column("machine_pubkey_fingerprint", Text, nullable=False), + Column("token_hash", Text, nullable=False), + # CHECK (approval_mode IN ('ask', 'full')) lives in 0013. + Column("approval_mode", Text, nullable=False, server_default="ask"), + Column("description", Text), + Column("status", Text, nullable=False, server_default="active"), + Column("paired_at", DateTime(timezone=True), nullable=False, server_default=func.now()), + Column("last_seen_at", DateTime(timezone=True)), + Column("revoked_at", DateTime(timezone=True)), + Column("revoke_reason", Text), + UniqueConstraint("user_id", "name", name="devices_user_name_uidx"), +) + +device_audit_log_table = Table( + "device_audit_log", + metadata, + Column("id", BigInteger, primary_key=True, autoincrement=True), + Column("device_id", Text, ForeignKey("devices.id", ondelete="CASCADE"), nullable=False), + Column("user_id", Text, nullable=False), + Column("agent_id", Text), + Column("conversation_id", Text), + Column("invocation_id", Text, nullable=False), + Column("action", Text, nullable=False), + Column("command", Text, nullable=False), + Column("working_dir", Text), + Column("approval_mode", Text, nullable=False), + Column("decision", Text, nullable=False), + Column("decision_reason", Text), + Column("issued_at", DateTime(timezone=True), nullable=False), + Column("started_at", DateTime(timezone=True)), + Column("finished_at", DateTime(timezone=True)), + Column("exit_code", Integer), + Column("duration_ms", Integer), + Column("stdout_sha256", CHAR(64)), + Column("stderr_sha256", CHAR(64)), + Column("stdout_bytes", Integer), + Column("stderr_bytes", Integer), + Column("error", Text), + Column("created_at", DateTime(timezone=True), nullable=False, server_default=func.now()), +) + +device_auto_approve_patterns_table = Table( + "device_auto_approve_patterns", + metadata, + Column("id", BigInteger, primary_key=True, autoincrement=True), + Column("device_id", Text, ForeignKey("devices.id", ondelete="CASCADE"), nullable=False), + Column("user_id", Text, nullable=False), + Column("pattern", Text, nullable=False), + Column("created_at", DateTime(timezone=True), nullable=False, server_default=func.now()), + UniqueConstraint("device_id", "user_id", "pattern", name="device_auto_approve_uidx"), +) diff --git a/application/storage/db/repositories/device_audit_log.py b/application/storage/db/repositories/device_audit_log.py new file mode 100644 index 00000000..93b6dd8d --- /dev/null +++ b/application/storage/db/repositories/device_audit_log.py @@ -0,0 +1,126 @@ +"""Append-only audit log for remote-device tool invocations.""" + +from __future__ import annotations + +from datetime import datetime +from typing import Optional + +from sqlalchemy import Connection, text + +from application.storage.db.base_repository import row_to_dict + + +class DeviceAuditLogRepository: + """Server-side canonical record of every remote-device invocation.""" + + def __init__(self, conn: Connection) -> None: + self._conn = conn + + def record_dispatch( + self, + *, + device_id: str, + user_id: str, + invocation_id: str, + command: str, + approval_mode: str, + decision: str, + decision_reason: Optional[str], + issued_at: datetime, + action: str = "run_command", + working_dir: Optional[str] = None, + agent_id: Optional[str] = None, + conversation_id: Optional[str] = None, + ) -> dict: + row = self._conn.execute( + text( + """ + INSERT INTO device_audit_log ( + device_id, user_id, agent_id, conversation_id, + invocation_id, action, command, working_dir, + approval_mode, decision, decision_reason, issued_at + ) VALUES ( + :device_id, :user_id, :agent_id, :conversation_id, + :invocation_id, :action, :command, :working_dir, + :approval_mode, :decision, :decision_reason, :issued_at + ) RETURNING * + """ + ), + { + "device_id": device_id, + "user_id": user_id, + "agent_id": agent_id, + "conversation_id": conversation_id, + "invocation_id": invocation_id, + "action": action, + "command": command, + "working_dir": working_dir, + "approval_mode": approval_mode, + "decision": decision, + "decision_reason": decision_reason, + "issued_at": issued_at, + }, + ).fetchone() + return row_to_dict(row) + + def record_result( + self, + invocation_id: str, + *, + started_at: Optional[datetime] = None, + finished_at: Optional[datetime] = None, + exit_code: Optional[int] = None, + duration_ms: Optional[int] = None, + stdout_sha256: Optional[str] = None, + stderr_sha256: Optional[str] = None, + stdout_bytes: Optional[int] = None, + stderr_bytes: Optional[int] = None, + error: Optional[str] = None, + ) -> bool: + """Update the dispatched row with execution outcome.""" + result = self._conn.execute( + text( + """ + UPDATE device_audit_log + SET started_at = COALESCE(:started_at, started_at), + finished_at = COALESCE(:finished_at, finished_at), + exit_code = COALESCE(:exit_code, exit_code), + duration_ms = COALESCE(:duration_ms, duration_ms), + stdout_sha256 = COALESCE(:stdout_sha256, stdout_sha256), + stderr_sha256 = COALESCE(:stderr_sha256, stderr_sha256), + stdout_bytes = COALESCE(:stdout_bytes, stdout_bytes), + stderr_bytes = COALESCE(:stderr_bytes, stderr_bytes), + error = COALESCE(:error, error) + WHERE invocation_id = :invocation_id + """ + ), + { + "invocation_id": invocation_id, + "started_at": started_at, + "finished_at": finished_at, + "exit_code": exit_code, + "duration_ms": duration_ms, + "stdout_sha256": stdout_sha256, + "stderr_sha256": stderr_sha256, + "stdout_bytes": stdout_bytes, + "stderr_bytes": stderr_bytes, + "error": error, + }, + ) + return result.rowcount > 0 + + def list_for_device( + self, device_id: str, user_id: str, *, limit: int = 100 + ) -> list[dict]: + result = self._conn.execute( + text( + """ + SELECT * FROM device_audit_log + WHERE device_id = :device_id AND user_id = :user_id + ORDER BY created_at DESC + LIMIT :limit + """ + ), + {"device_id": device_id, "user_id": user_id, "limit": int(limit)}, + ) + return [row_to_dict(r) for r in result.fetchall()] diff --git a/application/storage/db/repositories/device_auto_approve_patterns.py b/application/storage/db/repositories/device_auto_approve_patterns.py new file mode 100644 index 00000000..31d15014 --- /dev/null +++ b/application/storage/db/repositories/device_auto_approve_patterns.py @@ -0,0 +1,81 @@ +"""Per-device, per-user sticky "don't ask again" patterns.""" + +from __future__ import annotations + + +from sqlalchemy import Connection, text + + +class DeviceAutoApprovePatternsRepository: + """Normalized sticky-approval patterns scoped to (device, user).""" + + def __init__(self, conn: Connection) -> None: + self._conn = conn + + def add(self, device_id: str, user_id: str, pattern: str) -> bool: + """Idempotently add a pattern; returns True if newly inserted.""" + result = self._conn.execute( + text( + """ + INSERT INTO device_auto_approve_patterns + (device_id, user_id, pattern) + VALUES (:device_id, :user_id, :pattern) + ON CONFLICT (device_id, user_id, pattern) DO NOTHING + """ + ), + {"device_id": device_id, "user_id": user_id, "pattern": pattern}, + ) + return result.rowcount > 0 + + def remove(self, device_id: str, user_id: str, pattern: str) -> bool: + result = self._conn.execute( + text( + """ + DELETE FROM device_auto_approve_patterns + WHERE device_id = :device_id + AND user_id = :user_id + AND pattern = :pattern + """ + ), + {"device_id": device_id, "user_id": user_id, "pattern": pattern}, + ) + return result.rowcount > 0 + + def list_for_device(self, device_id: str, user_id: str) -> list[str]: + result = self._conn.execute( + text( + """ + SELECT pattern FROM device_auto_approve_patterns + WHERE device_id = :device_id AND user_id = :user_id + ORDER BY created_at + """ + ), + {"device_id": device_id, "user_id": user_id}, + ) + return [row[0] for row in result.fetchall()] + + def has_pattern( + self, device_id: str, user_id: str, pattern: str + ) -> bool: + row = self._conn.execute( + text( + """ + SELECT 1 FROM device_auto_approve_patterns + WHERE device_id = :device_id + AND user_id = :user_id + AND pattern = :pattern + LIMIT 1 + """ + ), + {"device_id": device_id, "user_id": user_id, "pattern": pattern}, + ).fetchone() + return row is not None + + def clear_for_device(self, device_id: str) -> int: + result = self._conn.execute( + text( + "DELETE FROM device_auto_approve_patterns WHERE device_id = :device_id" + ), + {"device_id": device_id}, + ) + return result.rowcount or 0 diff --git a/application/storage/db/repositories/devices.py b/application/storage/db/repositories/devices.py new file mode 100644 index 00000000..b31b447e --- /dev/null +++ b/application/storage/db/repositories/devices.py @@ -0,0 +1,142 @@ +"""Repository for the ``devices`` table.""" + +from __future__ import annotations + +from typing import Optional + +from sqlalchemy import Connection, text + +from application.storage.db.base_repository import row_to_dict + + +_ALLOWED_UPDATES = frozenset( + { + "name", "description", "approval_mode", "cli_version", "hostname", + "os", "arch", + } +) + + +class DevicesRepository: + """CRUD for paired remote devices.""" + + def __init__(self, conn: Connection) -> None: + self._conn = conn + + def create( + self, + device_id: str, + user_id: str, + name: str, + *, + machine_pubkey_fingerprint: str, + token_hash: str, + hostname: Optional[str] = None, + os: Optional[str] = None, + arch: Optional[str] = None, + cli_version: Optional[str] = None, + approval_mode: str = "ask", + description: Optional[str] = None, + ) -> dict: + row = self._conn.execute( + text( + """ + INSERT INTO devices ( + id, user_id, name, hostname, os, arch, cli_version, + machine_pubkey_fingerprint, token_hash, approval_mode, + description + ) VALUES ( + :id, :user_id, :name, :hostname, :os, :arch, :cli_version, + :fp, :token_hash, :approval_mode, :description + ) RETURNING * + """ + ), + { + "id": device_id, + "user_id": user_id, + "name": name, + "hostname": hostname, + "os": os, + "arch": arch, + "cli_version": cli_version, + "fp": machine_pubkey_fingerprint, + "token_hash": token_hash, + "approval_mode": approval_mode, + "description": description, + }, + ).fetchone() + return row_to_dict(row) + + def get(self, device_id: str, user_id: Optional[str] = None) -> Optional[dict]: + sql = "SELECT * FROM devices WHERE id = :id" + params: dict = {"id": device_id} + if user_id is not None: + sql += " AND user_id = :user_id" + params["user_id"] = user_id + row = self._conn.execute(text(sql), params).fetchone() + return row_to_dict(row) if row is not None else None + + def list_for_user( + self, user_id: str, *, include_revoked: bool = False + ) -> list[dict]: + sql = "SELECT * FROM devices WHERE user_id = :user_id" + if not include_revoked: + sql += " AND status = 'active'" + sql += " ORDER BY paired_at DESC" + result = self._conn.execute(text(sql), {"user_id": user_id}) + return [row_to_dict(r) for r in result.fetchall()] + + def update(self, device_id: str, user_id: str, fields: dict) -> bool: + """Update name/description/approval_mode (and cli/host metadata).""" + filtered = {k: v for k, v in fields.items() if k in _ALLOWED_UPDATES} + if not filtered: + return False + set_clauses = [f"{col} = :{col}" for col in filtered] + params = {**filtered, "id": device_id, "user_id": user_id} + result = self._conn.execute( + text( + f""" + UPDATE devices + SET {', '.join(set_clauses)} + WHERE id = :id AND user_id = :user_id + """ + ), + params, + ) + return result.rowcount > 0 + + def touch_last_seen(self, device_id: str) -> None: + """Bump ``last_seen_at`` to now(). Called on every poll/SSE open.""" + self._conn.execute( + text("UPDATE devices SET last_seen_at = now() WHERE id = :id"), + {"id": device_id}, + ) + + def revoke( + self, device_id: str, user_id: str, *, reason: str = "user_revoked" + ) -> bool: + result = self._conn.execute( + text( + """ + UPDATE devices + SET status = 'revoked', + revoked_at = now(), + revoke_reason = :reason + WHERE id = :id AND user_id = :user_id AND status = 'active' + """ + ), + {"id": device_id, "user_id": user_id, "reason": reason}, + ) + return result.rowcount > 0 + + def find_by_token_hash(self, token_hash: str) -> Optional[dict]: + """Used by the session-token verifier on each device request.""" + row = self._conn.execute( + text( + "SELECT * FROM devices " + "WHERE token_hash = :token_hash AND status = 'active' " + "LIMIT 1" + ), + {"token_hash": token_hash}, + ).fetchone() + return row_to_dict(row) if row is not None else None diff --git a/application/storage/db/repositories/message_events.py b/application/storage/db/repositories/message_events.py index 553efd01..c35e47c6 100644 --- a/application/storage/db/repositories/message_events.py +++ b/application/storage/db/repositories/message_events.py @@ -176,8 +176,14 @@ class MessageEventsRepository: def reconstruct_partial(self, message_id: str) -> dict: """Rebuild partial response/thought/sources/tool_calls from journal events. - ``answer``/``thought`` chunks concat in seq order; ``source``/ - ``tool_calls`` carry the full list at emit time (last-wins). + ``answer``/``thought`` chunks concat in seq order; ``source`` + carries the full list at emit time (last-wins). ``tool_calls`` + and per-call ``tool_call`` events are merged by ``call_id`` — + the most recent event for each call wins, and first-seen order + is preserved. An empty bulk ``tool_calls: []`` is a no-op (the + classic agent yields one at end-of-turn even when paused, and + wiping the overlay there would erase the live awaiting-approval + entry). """ rows = self._conn.execute( text( @@ -195,6 +201,23 @@ class MessageEventsRepository: thought_parts: list[str] = [] sources: list = [] tool_calls: list = [] + # Per-call overlay: maps call_id -> index into ``tool_calls`` so a + # later event for the same call replaces the earlier one in place + # and preserves first-seen ordering. Bulk ``tool_calls`` emits + # merge into the same overlay rather than reseeding so they can't + # erase a per-call entry that arrived earlier in the stream. + tool_call_index: dict[str, int] = {} + + def _overlay(entry: dict) -> None: + call_id = entry.get("call_id") + if not call_id: + return + existing = tool_call_index.get(call_id) + if existing is None: + tool_call_index[call_id] = len(tool_calls) + tool_calls.append(entry) + else: + tool_calls[existing] = entry for row in rows: payload = row.payload @@ -215,8 +238,19 @@ class MessageEventsRepository: sources = src elif etype == "tool_calls": tcs = payload.get("tool_calls") - if isinstance(tcs, list): - tool_calls = tcs + if not isinstance(tcs, list) or not tcs: + # Empty bulk is a no-op: the classic-agent end-of-turn + # yield emits one even on a paused turn where + # ``self.tool_calls`` is empty, and the per-call + # overlay carries the awaiting-approval entry. + continue + for tc in tcs: + if isinstance(tc, dict): + _overlay(tc) + elif etype == "tool_call": + data = payload.get("data") + if isinstance(data, dict): + _overlay(data) return { "response": "".join(response_parts), diff --git a/docs/content/Tools/_meta.js b/docs/content/Tools/_meta.js index c98154a2..66b923b4 100644 --- a/docs/content/Tools/_meta.js +++ b/docs/content/Tools/_meta.js @@ -7,6 +7,10 @@ export default { "title": "🗝️ API Tool", "href": "/Tools/api-tool" }, + "remote-device": { + "title": "🖥️ Remote Device", + "href": "/Tools/remote-device" + }, "creating-a-tool": { "title": "🛠️ Creating a Custom Tool", "href": "/Tools/creating-a-tool" diff --git a/docs/content/Tools/remote-device.mdx b/docs/content/Tools/remote-device.mdx new file mode 100644 index 00000000..5cc8f40e --- /dev/null +++ b/docs/content/Tools/remote-device.mdx @@ -0,0 +1,111 @@ +--- +title: 🖥️ Remote Device +description: Run shell commands on a paired remote machine from a DocsGPT agent using docsgpt-cli host. +--- + +import { Callout } from 'nextra/components'; + +# Remote Device Tool + +The Remote Device tool lets a DocsGPT agent run shell commands on a machine you control, such as a server, a Raspberry Pi, or your own laptop. You install `docsgpt-cli` on that machine, run it in `host` mode, and pair it with your DocsGPT account. The paired device then shows up as a tool you can attach to any agent. + +The machine connects outward to DocsGPT, so it works behind NAT or a firewall without opening any inbound ports. + +## How it works + +1. You run `docsgpt-cli host` on the target machine. It pairs to your account and keeps a lightweight connection open: it polls while idle and streams while a command is running. +2. When an agent calls the device's `run_command` action, DocsGPT sends the command down to the daemon. +3. The daemon runs the command locally, streams stdout and stderr back, and the agent uses the output to continue. +4. Every invocation is recorded in the device's activity log. + +## Prerequisites + +- `docsgpt-cli` installed on the machine you want to control. See the [installation instructions](https://github.com/arc53/DocsGPT-cli#installation): download a binary from the [releases page](https://github.com/arc53/DocsGPT-cli/releases), use Homebrew (`brew tap arc53/docsgpt-cli && brew install docsgpt-cli`), or build from source. +- Outbound internet access from that machine to your DocsGPT instance. + + + The `host` commands require a docsgpt-cli build with remote-device support. If `docsgpt-cli host` is not recognized, update to the latest release or build from source. + + +## Pair a device + +Pairing uses a short one-time code (a device-code style flow). + +1. In DocsGPT, go to **Settings -> Tools**, click **Add tool**, and choose **Remote Device**. +2. Give it a name, an optional description (the agent sees this, so describe what the machine is for), and pick an approval mode. DocsGPT then shows a pairing code such as `ABCD-WXYZ` and the command to run. +3. On the target machine, run: + ```bash + docsgpt-cli host pair --url https://your-docsgpt-instance + ``` + Enter the code when prompted. (Omit `--url` to use the default cloud instance.) +4. The UI switches to "Paired" once the code is redeemed. If you ran the command at a terminal, the CLI then offers to start the daemon or install it as a service. + +## Run the daemon + +Start it in the foreground: + +```bash +docsgpt-cli host +``` + +You will see a startup banner, then periodic "idle" heartbeats. Press Ctrl-C to stop. + +To keep it running across reboots, install it as a service: + +```bash +# Linux (systemd). As root this installs a system service; otherwise a user service. +docsgpt-cli host install-service + +# macOS (launchd). The default is a per-user LaunchAgent that starts on login. +docsgpt-cli host install-service + +# Always-on machine that starts at boot (Linux example): +sudo docsgpt-cli host install-service --system --user $USER +``` + +Remove the service with `docsgpt-cli host uninstall-service`. + +## Approval modes + +The approval mode is set per device, either in the pairing form or later on the device's page under **Settings -> Tools**. There are two modes: + +- **Ask** (default): every command pauses for your approval before it runs. You approve it from the chat. You can also choose "approve and don't ask again" to auto-approve that command pattern in the future. +- **Full access**: commands run without asking. + + + A built-in safety denylist always applies, even in Full access mode. Catastrophic commands (for example `rm -rf /`, fork bombs, writing directly to a disk device, or `git push --force`) still pause for explicit approval and cannot be bypassed. + + +Compound commands are split on operators such as `&&`, `||`, `;`, and `|`, and each part is checked on its own, so a dangerous part cannot be hidden behind a safe one. + +## Attach to an agent + +A paired device behaves like any other tool. Open or create an agent, add the device from the tool picker, and the agent can call its `run_command` action. The picker shows whether the device is currently online. + +If you attach more than one device to an agent, give each a clear description so the agent can pick the right one. + +## Manage a device + +Open the device from **Settings -> Tools** to: + +- See its status (online or offline, with last-seen time), host, OS, and CLI version. +- Change its name, description, or approval mode. +- View recent activity (the command audit log). +- Revoke it. + +From the CLI: + +```bash +docsgpt-cli host status # live status from the server +docsgpt-cli host revoke # revoke on the server and clear local state +docsgpt-cli host reset # clear local pairing only (leaves the server-side device) +``` + +Revoking a device stops its daemon: the next time it checks in it sees the revocation, prints a message, and exits. Under a service manager it will not be restarted. + +## Security notes + +- The machine connects outward only. No inbound ports are opened. +- Each device has its own token, stored hashed on the server and revocable at any time. +- Prefer **Ask** mode for any machine with sensitive data on it. Use **Full access** only on machines you are comfortable letting an agent drive unattended, and remember that the denylist is the only automatic guard in that mode. +- Every command is logged on the server and visible in the device's activity log. diff --git a/frontend/public/toolIcons/tool_remote_device.svg b/frontend/public/toolIcons/tool_remote_device.svg new file mode 100644 index 00000000..7f55d903 --- /dev/null +++ b/frontend/public/toolIcons/tool_remote_device.svg @@ -0,0 +1,7 @@ + + + + + + + diff --git a/frontend/src/agents/NewAgent.tsx b/frontend/src/agents/NewAgent.tsx index 21d692c6..e2a49a62 100644 --- a/frontend/src/agents/NewAgent.tsx +++ b/frontend/src/agents/NewAgent.tsx @@ -23,6 +23,7 @@ import { } from '@/components/ui/select'; import { Switch } from '@/components/ui/switch'; +import devicesService from '../api/services/devicesService'; import modelService from '../api/services/modelService'; import userService from '../api/services/userService'; import SourceIcon from '../assets/source.svg'; @@ -455,9 +456,25 @@ export default function NewAgent({ mode }: { mode: 'new' | 'edit' | 'draft' }) { useEffect(() => { const getTools = async () => { - const response = await userService.getUserTools(token); - if (!response.ok) throw new Error('Failed to fetch tools'); - const data = await response.json(); + const [toolsResponse, devicesResult] = await Promise.all([ + userService.getUserTools(token), + // Tolerate failures here: the picker should still render the + // tool list even if /api/devices returns an error or 401. + devicesService.list(token).catch(() => ({ devices: [] })), + ]); + if (!toolsResponse.ok) throw new Error('Failed to fetch tools'); + const data = await toolsResponse.json(); + const devicesById = new Map< + string, + { online: boolean; last_seen_at: string | null | undefined } + >(); + const onlineWindowMs = 30_000; + (devicesResult.devices || []).forEach((d) => { + const seen = d.last_seen_at ? Date.parse(d.last_seen_at) : NaN; + const online = + !Number.isNaN(seen) && Date.now() - seen < onlineWindowMs; + devicesById.set(d.id, { online, last_seen_at: d.last_seen_at }); + }); // Group ordering: builtins -> defaults -> user tools (sorted via the // MultiSelectPopover first-appearance grouping). const groupFor = (tool: UserToolType): string => { @@ -466,12 +483,33 @@ export default function NewAgent({ mode }: { mode: 'new' | 'edit' | 'draft' }) { return t('agents.form.toolsPopup.groupCustom'); }; const tools: MultiSelectPopoverItem[] = data.tools.map( - (tool: UserToolType) => ({ - id: tool.id, - label: getToolDisplayName(tool), - icon: `/toolIcons/tool_${tool.name}.svg`, - group: groupFor(tool), - }), + (tool: UserToolType) => { + const base: MultiSelectPopoverItem = { + id: tool.id, + label: getToolDisplayName(tool), + icon: `/toolIcons/tool_${tool.name}.svg`, + group: groupFor(tool), + }; + if (tool.name === 'remote_device') { + const deviceId = (tool.config?.device_id as string) || ''; + const meta = devicesById.get(deviceId); + const online = meta?.online ?? false; + base.descriptionNode = ( + + {online + ? t('settings.devices.online') + : t('settings.devices.offline')} + + ); + } + return base; + }, ); const groupOrder = [ t('agents.form.toolsPopup.groupBuiltin'), diff --git a/frontend/src/api/client.ts b/frontend/src/api/client.ts index 37a6d08e..c7e599ac 100644 --- a/frontend/src/api/client.ts +++ b/frontend/src/api/client.ts @@ -110,12 +110,14 @@ const createClient = (transport: FetchLike) => { delete: ( url: string, token: string | null, + data?: any, headers = {}, signal?: AbortSignal, ): Promise => request(url, { method: 'DELETE', headers: getHeaders(token, headers), + body: data !== undefined ? JSON.stringify(data) : undefined, signal, }), }; diff --git a/frontend/src/api/endpoints.ts b/frontend/src/api/endpoints.ts index 2ad8f220..178ef7f7 100644 --- a/frontend/src/api/endpoints.ts +++ b/frontend/src/api/endpoints.ts @@ -91,6 +91,13 @@ const endpoints = { `/api/schedules/${id}/runs?limit=${limit ?? 50}&offset=${offset ?? 0}`, SCHEDULE_RUN: (id: string, runId: string) => `/api/schedules/${id}/runs/${runId}`, + DEVICES: '/api/devices', + DEVICE: (id: string) => `/api/devices/${id}`, + DEVICE_AUTO_APPROVE: (id: string) => `/api/devices/${id}/auto-approve`, + DEVICE_AUDIT: (id: string) => `/api/devices/${id}/audit`, + DEVICE_PAIRINGS: '/api/devices/pairings', + DEVICE_PAIRING: (deviceCode: string) => + `/api/devices/pairings/${deviceCode}`, }, V1: { CHAT_COMPLETIONS: '/v1/chat/completions', diff --git a/frontend/src/api/services/devicesService.ts b/frontend/src/api/services/devicesService.ts new file mode 100644 index 00000000..4ea8ce23 --- /dev/null +++ b/frontend/src/api/services/devicesService.ts @@ -0,0 +1,172 @@ +import apiClient from '../client'; +import endpoints from '../endpoints'; + +export type ApprovalMode = 'ask' | 'full'; + +export interface Device { + id: string; + name: string; + hostname?: string | null; + os?: string | null; + arch?: string | null; + cli_version?: string | null; + approval_mode: ApprovalMode; + description?: string | null; + status: 'active' | 'revoked'; + paired_at?: string | null; + last_seen_at?: string | null; + revoked_at?: string | null; +} + +export interface PairingResponse { + device_code: string; + user_code: string; + verification_uri?: string; + expires_in: number; + interval: number; +} + +export interface PairingStatus { + device_code: string; + user_code?: string; + status: 'pending' | 'redeemed' | 'cancelled' | 'expired'; + device_id?: string | null; + device_name?: string | null; +} + +export interface AuditEntry { + id: number; + device_id: string; + invocation_id: string; + command: string; + working_dir?: string | null; + approval_mode: string; + decision: string; + decision_reason?: string | null; + exit_code?: number | null; + duration_ms?: number | null; + issued_at: string; + started_at?: string | null; + finished_at?: string | null; + created_at: string; +} + +const json = async (response: Response | unknown) => { + const r = response as Response; + if (!('json' in r) || typeof r.json !== 'function') return r as unknown; + return r.json(); +}; + +const devicesService = { + list: async (token: string | null): Promise<{ devices: Device[] }> => { + const r = await apiClient.get(endpoints.USER.DEVICES, token); + return (await json(r)) as { devices: Device[] }; + }, + + get: async (id: string, token: string | null): Promise => { + const r = await apiClient.get(endpoints.USER.DEVICE(id), token); + return (await json(r)) as Device; + }, + + update: async ( + id: string, + payload: Partial>, + token: string | null, + ): Promise => { + const r = await apiClient.patch(endpoints.USER.DEVICE(id), payload, token); + return (await json(r)) as Device; + }, + + revoke: async ( + id: string, + token: string | null, + ): Promise<{ success: boolean }> => { + const r = await apiClient.delete(endpoints.USER.DEVICE(id), token); + return (await json(r)) as { success: boolean }; + }, + + startPairing: async ( + token: string | null, + body?: { + name?: string; + description?: string; + approval_mode?: ApprovalMode; + }, + ): Promise => { + const r = await apiClient.post( + endpoints.USER.DEVICE_PAIRINGS, + body ?? {}, + token, + ); + return (await json(r)) as PairingResponse; + }, + + pollPairing: async ( + deviceCode: string, + token: string | null, + ): Promise => { + const r = await apiClient.get( + endpoints.USER.DEVICE_PAIRING(deviceCode), + token, + ); + return (await json(r)) as PairingStatus; + }, + + cancelPairing: async ( + deviceCode: string, + token: string | null, + ): Promise<{ success: boolean }> => { + const r = await apiClient.delete( + endpoints.USER.DEVICE_PAIRING(deviceCode), + token, + ); + return (await json(r)) as { success: boolean }; + }, + + addAutoApprovePattern: async ( + id: string, + command: string, + token: string | null, + ): Promise<{ pattern: string }> => { + const r = await apiClient.post( + endpoints.USER.DEVICE_AUTO_APPROVE(id), + { command }, + token, + ); + return (await json(r)) as { pattern: string }; + }, + + listAutoApprovePatterns: async ( + id: string, + token: string | null, + ): Promise<{ patterns: string[] }> => { + const r = await apiClient.get( + endpoints.USER.DEVICE_AUTO_APPROVE(id), + token, + ); + return (await json(r)) as { patterns: string[] }; + }, + + removeAutoApprovePattern: async ( + id: string, + pattern: string, + token: string | null, + ): Promise<{ success: boolean }> => { + const r = await apiClient.delete( + endpoints.USER.DEVICE_AUTO_APPROVE(id), + token, + { pattern }, + ); + return (await json(r)) as { success: boolean }; + }, + + listAudit: async ( + id: string, + token: string | null, + ): Promise<{ entries: AuditEntry[] }> => { + const r = await apiClient.get(endpoints.USER.DEVICE_AUDIT(id), token); + return (await json(r)) as { entries: AuditEntry[] }; + }, +}; + +export default devicesService; diff --git a/frontend/src/components/MultiSelectPopover.tsx b/frontend/src/components/MultiSelectPopover.tsx index 9fbc87bb..17262cc6 100644 --- a/frontend/src/components/MultiSelectPopover.tsx +++ b/frontend/src/components/MultiSelectPopover.tsx @@ -22,6 +22,8 @@ export type MultiSelectPopoverItem = { id: string; label: string; description?: string; + /** Rich variant of ``description`` — wins when both are set. */ + descriptionNode?: React.ReactNode; icon?: React.ReactNode | string; group?: string; disabled?: boolean; @@ -138,11 +140,13 @@ export function MultiSelectPopover({ > {item.label}

- {item.description && ( + {item.descriptionNode ? ( +
{item.descriptionNode}
+ ) : item.description ? (

{item.description}

- )} + ) : null}
60 ? argPreview.slice(0, 57) + '...' : argPreview; + const isRemoteDevice = + toolCall.tool_name === 'remote_device' && toolCall.device_id; + const handleApproveSticky = async () => { + if (!isRemoteDevice || !toolCall.device_id) return; + const command = + (toolCall.arguments && (toolCall.arguments.command as string)) || ''; + if (command) { + try { + const { default: devicesService } = await import( + '../api/services/devicesService' + ); + await devicesService.addAutoApprovePattern( + toolCall.device_id, + command, + token, + ); + } catch (err) { + console.error('auto-approve register failed', err); + } + } + onToolAction?.(toolCall.call_id, 'approved'); + }; + return (
@@ -964,6 +989,19 @@ function ToolCallApprovalBar({ > Approve + {isRemoteDevice && ( + + )} + + + ); + } + return ( + + ); + })(); + + return ( + { + if (!open) close(); + }} + title={t('settings.devices.pairing.title')} + footer={footer} + contentClassName="px-1" + > + {stage === 'form' && renderForm()} + {stage === 'waiting' && renderWaiting()} + + ); +} diff --git a/frontend/src/settings/RemoteDeviceConfig.tsx b/frontend/src/settings/RemoteDeviceConfig.tsx new file mode 100644 index 00000000..b830a433 --- /dev/null +++ b/frontend/src/settings/RemoteDeviceConfig.tsx @@ -0,0 +1,466 @@ +import React from 'react'; +import { useTranslation } from 'react-i18next'; +import { useSelector } from 'react-redux'; + +import devicesService, { + ApprovalMode, + AuditEntry, + Device, +} from '../api/services/devicesService'; +import ArrowLeft from '../assets/arrow-left.svg'; +import CopyButton from '../components/CopyButton'; +import Spinner from '../components/Spinner'; +import { + Accordion, + AccordionContent, + AccordionItem, + AccordionTrigger, +} from '../components/ui/accordion'; +import { Alert, AlertDescription } from '../components/ui/alert'; +import { Button } from '../components/ui/button'; +import { Input } from '../components/ui/input'; +import { + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, +} from '../components/ui/select'; +import ConfirmationModal from '../modals/ConfirmationModal'; +import { ActiveState } from '../models/misc'; +import { selectToken } from '../preferences/preferenceSlice'; +import { UserToolType } from './types'; + +/** ms-since-last-seen threshold for the online pill. */ +const ONLINE_WINDOW_MS = 30_000; + +function isOnline(device: Device | null): boolean { + if (!device?.last_seen_at) return false; + const t = Date.parse(device.last_seen_at); + if (Number.isNaN(t)) return false; + return Date.now() - t < ONLINE_WINDOW_MS; +} + +function formatTimestamp(value: string | null | undefined): string { + if (!value) return '-'; + try { + return new Date(value).toLocaleString(); + } catch { + return value; + } +} + +/** Compact relative span (e.g. "12s", "5m", "3h", "2d") since `value`. */ +function formatRelative(value: string | null | undefined): string | null { + if (!value) return null; + const t = Date.parse(value); + if (Number.isNaN(t)) return null; + const diff = Math.max(0, Date.now() - t); + const s = Math.floor(diff / 1000); + if (s < 60) return `${s}s`; + const m = Math.floor(s / 60); + if (m < 60) return `${m}m`; + const h = Math.floor(m / 60); + if (h < 24) return `${h}h`; + const d = Math.floor(h / 24); + return `${d}d`; +} + +interface Props { + tool: UserToolType; + handleGoBack: () => void; +} + +export default function RemoteDeviceConfig({ tool, handleGoBack }: Props) { + const { t } = useTranslation(); + const token = useSelector(selectToken); + const deviceId = React.useMemo( + () => (tool.config?.device_id as string | undefined) || '', + [tool.config?.device_id], + ); + + const [device, setDevice] = React.useState(null); + const [loading, setLoading] = React.useState(false); + const [name, setName] = React.useState(''); + const [description, setDescription] = React.useState(''); + const [approvalMode, setApprovalMode] = React.useState('ask'); + const [saving, setSaving] = React.useState(false); + const [error, setError] = React.useState(null); + + const [audit, setAudit] = React.useState(null); + const [auditLoading, setAuditLoading] = React.useState(false); + + const [revokeState, setRevokeState] = React.useState('INACTIVE'); + + const applyDevice = React.useCallback((d: Device) => { + setDevice(d); + setName(d.name || ''); + setDescription(d.description || ''); + setApprovalMode(d.approval_mode); + }, []); + + const loadDevice = React.useCallback(() => { + if (!deviceId) return; + setLoading(true); + devicesService + .get(deviceId, token) + .then((d) => applyDevice(d)) + .catch((err) => { + console.error('load device failed', err); + setError(err instanceof Error ? err.message : 'load failed'); + }) + .finally(() => setLoading(false)); + }, [deviceId, token, applyDevice]); + + React.useEffect(() => { + loadDevice(); + }, [loadDevice]); + + const loadAudit = React.useCallback(() => { + if (!deviceId) return; + setAuditLoading(true); + devicesService + .listAudit(deviceId, token) + .then((res) => setAudit(res.entries || [])) + .catch((err) => { + console.error('load audit failed', err); + setAudit([]); + }) + .finally(() => setAuditLoading(false)); + }, [deviceId, token]); + + const handleAuditToggle = (value: string) => { + // Radix Accordion (type="single") passes the open item id, or empty + // when closed. Lazy-fetch on first open. + if (value === 'audit' && audit === null) { + loadAudit(); + } + }; + + const hasUnsavedChanges = + !!device && + !!name.trim() && + ((device.name || '') !== name || + (device.description || '') !== description || + device.approval_mode !== approvalMode); + + const handleSaveChanges = async () => { + if (!deviceId || !device || !hasUnsavedChanges) return; + const payload: Partial< + Pick + > = {}; + if ((device.name || '') !== name) payload.name = name; + if ((device.description || '') !== description) + payload.description = description; + if (device.approval_mode !== approvalMode) + payload.approval_mode = approvalMode; + + setSaving(true); + setError(null); + try { + const next = await devicesService.update(deviceId, payload, token); + applyDevice(next); + } catch (err) { + setError(err instanceof Error ? err.message : 'update failed'); + } finally { + setSaving(false); + } + }; + + const handleRevoke = async () => { + if (!deviceId) return; + try { + await devicesService.revoke(deviceId, token); + setRevokeState('INACTIVE'); + handleGoBack(); + } catch (err) { + console.error('revoke failed', err); + setRevokeState('INACTIVE'); + } + }; + + const online = isOnline(device); + const lastSeenAgo = formatRelative(device?.last_seen_at); + const pillText = + (online ? t('settings.devices.online') : t('settings.devices.offline')) + + (lastSeenAgo + ? ` · ${t('settings.devices.seenAgo', { time: lastSeenAgo })}` + : ''); + + if (loading && !device) { + return ( +
+ +
+ ); + } + + return ( +
+
+
+ +

{t('settings.tools.backToAllTools')}

+
+ +
+ + {error && ( + + {error} + + )} + +
+ {`${tool.displayName} +

+ {device?.name || tool.customName || tool.displayName || 'device'} +

+ + {pillText} + + {approvalMode === 'full' && ( + + {t('settings.devices.approvalFull')} + + )} +
+ + {/* Identity */} +
+

+ {t('settings.devices.identity')} +

+
+ +
+ setName(e.target.value)} + placeholder={t('settings.devices.nameLabel')} + className="rounded-xl" + /> +
+
+
+ +
+ setDescription(e.target.value)} + placeholder={t('settings.devices.descriptionPlaceholder')} + className="rounded-xl" + /> +
+
+
+ + {/* Access */} +
+

+ {t('settings.devices.access')} +

+
+ +
+ +

+ {approvalMode === 'full' + ? t('settings.devices.approvalFullDescription') + : t('settings.devices.approvalAskDescription')} +

+ {approvalMode === 'full' && ( +
+ {t('settings.devices.fullAccessWarning')} +
+ )} +
+
+
+ + {/* Connection */} +
+

+ {t('settings.devices.connection')} +

+
+ + {t('settings.devices.hostLabel')} + +
+ + {device?.hostname || 'unknown host'} + + {device?.os ? ` · ${device.os}` : ''} + {device?.arch ? ` · ${device.arch}` : ''} + {device?.cli_version ? ` · cli ${device.cli_version}` : ''} +
+
+
+ + {t('settings.devices.deviceIdLabel')} + +
+ + {deviceId || '-'} + + {deviceId && } +
+
+
+ + {t('settings.devices.lastSeenLabel')} + +
+ {formatTimestamp(device?.last_seen_at)} +
+
+ {!online && ( +

+ {t('settings.devices.offlineHint')} +

+ )} +
+ + {/* Recent activity */} +
+

+ {t('settings.devices.auditTitle')} +

+ + + + {t('settings.devices.auditTitle')} + {audit !== null ? ` (${audit.length})` : ''} + + + {auditLoading ? ( +

+ {t('settings.devices.auditLoading')} +

+ ) : !audit || audit.length === 0 ? ( +

+ {t('settings.devices.auditEmpty')} +

+ ) : ( +
    + {audit.map((entry) => ( +
  • + + {entry.command} + +
    + + {t('settings.devices.auditDecision')}:{' '} + {entry.decision} + + + {t('settings.devices.auditExit')}:{' '} + {entry.exit_code ?? '-'} + + + {t('settings.devices.auditDuration')}:{' '} + {entry.duration_ms ?? '-'} ms + + {formatTimestamp(entry.created_at)} +
    +
    • + ))} +
+ )} + + + +
+ + {/* Danger zone */} +
+

+ {t('settings.devices.dangerZone')} +

+
+

+ {t('settings.devices.dangerZoneDescription')} +

+ +
+
+ + +
+ ); +} diff --git a/frontend/src/settings/Tools.tsx b/frontend/src/settings/Tools.tsx index 5515869a..12db8ff3 100644 --- a/frontend/src/settings/Tools.tsx +++ b/frontend/src/settings/Tools.tsx @@ -4,6 +4,7 @@ import React from 'react'; import { useTranslation } from 'react-i18next'; import { useSelector } from 'react-redux'; +import devicesService from '../api/services/devicesService'; import userService from '../api/services/userService'; import Edit from '../assets/edit.svg'; import NoFilesDarkIcon from '../assets/no-files-dark.svg'; @@ -26,6 +27,7 @@ import ConfirmationModal from '../modals/ConfirmationModal'; import MCPServerModal from '../modals/MCPServerModal'; import { ActiveState } from '../models/misc'; import { selectToken } from '../preferences/preferenceSlice'; +import RemoteDeviceConfig from './RemoteDeviceConfig'; import ToolConfig from './ToolConfig'; import { APIToolType, UserToolType } from './types'; @@ -70,14 +72,29 @@ export default function Tools() { }; const confirmDeleteTool = () => { - if (toolToDelete) { - userService.deleteTool({ id: toolToDelete.id }, token).then(() => { - getUserTools(); - fetchMcpStatuses(); - setDeleteModalState('INACTIVE'); - setToolToDelete(null); - }); + if (!toolToDelete) return; + const afterDelete = () => { + getUserTools(); + fetchMcpStatuses(); + setDeleteModalState('INACTIVE'); + setToolToDelete(null); + }; + // Remote-device tools front a paired device + live session token. Revoke + // the device server-side (marks revoked, closes any session, invalidates + // the token, and drops the user_tools row) instead of deleting only the + // tool row, which would leave the daemon polling with a live token. + const deviceId = + toolToDelete.name === 'remote_device' + ? toolToDelete.config?.device_id + : undefined; + if (deviceId) { + devicesService + .revoke(deviceId, token) + .then(afterDelete) + .catch((error) => console.error('Failed to revoke device:', error)); + return; } + userService.deleteTool({ id: toolToDelete.id }, token).then(afterDelete); }; const handleReconnect = (tool: UserToolType) => { @@ -152,7 +169,7 @@ export default function Tools() { .then((data) => { // Pure builtins (agent-only, e.g. a future builtin without an // agentless path) carry no per-user state and only apply when - // added to an agent — hide them from the management page. Dual- + // added to an agent, so hide them from the management page. Dual- // registered tools (``scheduler``: builtin + default) stay visible // here so the user can toggle the default off in agentless chats. const filtered = (data.tools || []).filter( @@ -209,6 +226,23 @@ export default function Tools() { .catch((error) => console.error('Error fetching tools:', error)); }; + const handleDevicePaired = (deviceId: string) => { + setAddToolModalState('INACTIVE'); + userService + .getUserTools(token) + .then((res) => res.json()) + .then((data) => { + const newTool = data.tools.find( + (toolItem: UserToolType) => + toolItem.name === 'remote_device' && + toolItem.config?.device_id === deviceId, + ); + if (newTool) setSelectedTool(newTool); + else console.error('Paired device tool not found'); + }) + .catch((error) => console.error('Error fetching tools:', error)); + }; + React.useEffect(() => { getUserTools(); fetchMcpStatuses(); @@ -216,11 +250,18 @@ export default function Tools() { return (
{selectedTool ? ( - + selectedTool.name === 'remote_device' ? ( + + ) : ( + + ) ) : (
@@ -437,6 +478,7 @@ export default function Tools() { setModalState={setAddToolModalState} getUserTools={getUserTools} onToolAdded={handleToolAdded} + onDevicePaired={handleDevicePaired} /> } /> } /> } /> + } + /> } /> } /> diff --git a/frontend/vite.config.ts b/frontend/vite.config.ts index 31a2720c..ae507f82 100644 --- a/frontend/vite.config.ts +++ b/frontend/vite.config.ts @@ -1,21 +1,33 @@ /// -import { defineConfig } from 'vite'; +import { defineConfig, loadEnv } from 'vite'; import react from '@vitejs/plugin-react'; import svgr from 'vite-plugin-svgr'; import path from 'path'; // https://vitejs.dev/config/ -export default defineConfig({ - plugins: [react(), svgr()], - resolve: { - alias: { - '@': path.resolve(__dirname, './src'), +export default defineConfig(({ mode }) => { + const env = loadEnv(mode, process.cwd(), ''); + return { + plugins: [react(), svgr()], + resolve: { + alias: { + '@': path.resolve(__dirname, './src'), + }, }, - }, - test: { - environment: 'happy-dom', - globals: true, - include: ['src/**/*.test.{ts,tsx}'], - setupFiles: ['./vitest.setup.ts'], - }, + server: { + // Extra dev hosts (e.g. a tailscale name) come from VITE_ALLOWED_HOSTS in + // an untracked .env.local; machine-specific names stay out of the repo. + allowedHosts: env.VITE_ALLOWED_HOSTS + ? env.VITE_ALLOWED_HOSTS.split(',') + .map((h) => h.trim()) + .filter(Boolean) + : [], + }, + test: { + environment: 'happy-dom', + globals: true, + include: ['src/**/*.test.{ts,tsx}'], + setupFiles: ['./vitest.setup.ts'], + }, + }; }); diff --git a/scripts/e2e/mock_llm_fixtures/01c733b9660eed626247ffb527486e12c2b0def48bc79fd4468533864c268441.json b/scripts/e2e/mock_llm_fixtures/01c733b9660eed626247ffb527486e12c2b0def48bc79fd4468533864c268441.json new file mode 100644 index 00000000..19b90b15 --- /dev/null +++ b/scripts/e2e/mock_llm_fixtures/01c733b9660eed626247ffb527486e12c2b0def48bc79fd4468533864c268441.json @@ -0,0 +1,17 @@ +{ + "_comment": "Tester fixture: 'please run whoami on my remote device' -> run_command tool call (whoami). Used to verify full mode auto-approves a safe non-denylisted command.", + "response": { + "content": "", + "tool_calls": [ + { + "id": "call_remote_tester_whoami_1", + "type": "function", + "function": { + "name": "run_command", + "arguments": "{\"command\":\"whoami\",\"working_directory\":\"/tmp\",\"timeout_ms\":5000}" + } + } + ], + "finish_reason": "tool_calls" + } +} diff --git a/scripts/e2e/mock_llm_fixtures/b196e9532b6a979d959a570dcdec5f96cea35175091612545e24ff91ed65c8a3.json b/scripts/e2e/mock_llm_fixtures/b196e9532b6a979d959a570dcdec5f96cea35175091612545e24ff91ed65c8a3.json new file mode 100644 index 00000000..20facd7b --- /dev/null +++ b/scripts/e2e/mock_llm_fixtures/b196e9532b6a979d959a570dcdec5f96cea35175091612545e24ff91ed65c8a3.json @@ -0,0 +1,17 @@ +{ + "_comment": "Tester fixture: 'please run mkdir /tmp/rd_v3_test_dir on my remote device' -> run_command tool call (mkdir /tmp/rd_v3_test_dir). Used to verify a write command in ask mode still pauses on the approval bar.", + "response": { + "content": "", + "tool_calls": [ + { + "id": "call_remote_tester_mkdir_1", + "type": "function", + "function": { + "name": "run_command", + "arguments": "{\"command\":\"mkdir /tmp/rd_v3_test_dir\",\"working_directory\":\"/tmp\",\"timeout_ms\":5000}" + } + } + ], + "finish_reason": "tool_calls" + } +} diff --git a/scripts/e2e/mock_llm_fixtures/d42a7dccbd1f218bb095a28e2b4a9c878800e030223be5189ebe13274c65d21f.json b/scripts/e2e/mock_llm_fixtures/d42a7dccbd1f218bb095a28e2b4a9c878800e030223be5189ebe13274c65d21f.json new file mode 100644 index 00000000..e52cc2f2 --- /dev/null +++ b/scripts/e2e/mock_llm_fixtures/d42a7dccbd1f218bb095a28e2b4a9c878800e030223be5189ebe13274c65d21f.json @@ -0,0 +1,17 @@ +{ + "_comment": "Tester fixture: 'please run rm -rf / on my remote device' -> run_command tool call (rm -rf /). Used to verify full mode still pauses on the denylist-forced-prompt circuit-breaker.", + "response": { + "content": "", + "tool_calls": [ + { + "id": "call_remote_tester_rm_1", + "type": "function", + "function": { + "name": "run_command", + "arguments": "{\"command\":\"rm -rf /\",\"working_directory\":\"/tmp\",\"timeout_ms\":5000}" + } + } + ], + "finish_reason": "tool_calls" + } +} diff --git a/scripts/e2e/mock_llm_fixtures/ed4ff698efb5a8c5f4d6f8ad1a68526da3813c82b0bb57ceaff298e959193d1f.json b/scripts/e2e/mock_llm_fixtures/ed4ff698efb5a8c5f4d6f8ad1a68526da3813c82b0bb57ceaff298e959193d1f.json new file mode 100644 index 00000000..512f0eb1 --- /dev/null +++ b/scripts/e2e/mock_llm_fixtures/ed4ff698efb5a8c5f4d6f8ad1a68526da3813c82b0bb57ceaff298e959193d1f.json @@ -0,0 +1,17 @@ +{ + "_comment": "Tester fixture: 'please run ls -la /tmp on my remote device' -> run_command tool call (ls -la /tmp).", + "response": { + "content": "", + "tool_calls": [ + { + "id": "call_remote_tester_1", + "type": "function", + "function": { + "name": "run_command", + "arguments": "{\"command\":\"ls -la /tmp\",\"working_directory\":\"/tmp\",\"timeout_ms\":5000}" + } + } + ], + "finish_reason": "tool_calls" + } +} diff --git a/tests/agents/test_tool_executor.py b/tests/agents/test_tool_executor.py index d37c5eae..80e5777e 100644 --- a/tests/agents/test_tool_executor.py +++ b/tests/agents/test_tool_executor.py @@ -457,6 +457,346 @@ class TestCheckPause: assert result["action_name"] == "delete_all" +@pytest.mark.unit +class TestCheckPauseRemoteDevice: + """The gate must consult the live ``RemoteDeviceTool`` decision for + ``remote_device`` so ``approval_mode`` changes after pair time and + per-command heuristics (denylist, sticky) are respected. + """ + + def _make_call(self, name="run_command", call_id="c1", command="ls -la /tmp"): + import json + + call = Mock() + call.name = name + call.id = call_id + call.arguments = json.dumps({"command": command}) + call.thought_signature = None + return call + + def _tools_dict(self, *, device_id="dev_abc"): + return { + "rd0": { + "id": "rd0", + "name": "remote_device", + "config": {"device_id": device_id}, + "actions": [ + { + "name": "run_command", + "description": "Execute on device", + "active": True, + "require_approval": False, + "parameters": { + "type": "object", + "properties": { + "command": { + "type": "string", + "filled_by_llm": True, + "value": "", + }, + }, + "required": ["command"], + }, + } + ], + } + } + + def _patch_device(self, monkeypatch, device): + """Stub ``RemoteDeviceTool._load_device`` to return ``device``.""" + from application.agents.tools import remote_device + + monkeypatch.setattr( + remote_device.RemoteDeviceTool, + "_load_device", + lambda self: device, + ) + + def _patch_sticky(self, monkeypatch, patterns): + """Stub the sticky lookup to match any normalized pattern in ``patterns``.""" + from application.agents.tools import remote_device + + monkeypatch.setattr( + remote_device, + "normalize_command", + lambda cmd: cmd.split()[0] + " *" if cmd else "", + ) + captured = {"patterns": set(patterns)} + + class _StubRepo: + def __init__(self, conn): + pass + + def has_pattern(self, device_id, user_id, pattern): + return pattern in captured["patterns"] + + monkeypatch.setattr( + remote_device, + "DeviceAutoApprovePatternsRepository", + _StubRepo, + ) + + class _StubConn: + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + monkeypatch.setattr(remote_device, "db_readonly", lambda: _StubConn()) + + def test_full_non_deny_no_pause(self, monkeypatch): + """``full`` + a non-denylisted command auto-approves.""" + executor = ToolExecutor(user="alice") + self._patch_device( + monkeypatch, + { + "id": "dev_abc", + "approval_mode": "full", + "status": "active", + }, + ) + self._patch_sticky(monkeypatch, set()) + + tools_dict = self._tools_dict() + executor.prepare_tools_for_llm(tools_dict) + call = self._make_call(command="whoami") + assert executor.check_pause(tools_dict, call, "OpenAILLM") is None + + def test_full_writing_command_no_pause(self, monkeypatch): + """``full`` auto-approves writes too (only the denylist stops it).""" + executor = ToolExecutor(user="alice") + self._patch_device( + monkeypatch, + { + "id": "dev_abc", + "approval_mode": "full", + "status": "active", + }, + ) + self._patch_sticky(monkeypatch, set()) + + tools_dict = self._tools_dict() + executor.prepare_tools_for_llm(tools_dict) + call = self._make_call(command="rm /tmp/foo") + assert executor.check_pause(tools_dict, call, "OpenAILLM") is None + + def test_full_denylist_forces_pause(self, monkeypatch): + """``full`` + a denylisted command still pauses (forced prompt).""" + executor = ToolExecutor(user="alice") + self._patch_device( + monkeypatch, + { + "id": "dev_abc", + "approval_mode": "full", + "status": "active", + }, + ) + self._patch_sticky(monkeypatch, set()) + + tools_dict = self._tools_dict() + executor.prepare_tools_for_llm(tools_dict) + call = self._make_call(command="rm -rf /") + result = executor.check_pause(tools_dict, call, "OpenAILLM") + assert result is not None + assert result["pause_type"] == "awaiting_approval" + + def test_ask_mode_pauses(self, monkeypatch): + """``ask`` + any command pauses by default.""" + executor = ToolExecutor(user="alice") + self._patch_device( + monkeypatch, + { + "id": "dev_abc", + "approval_mode": "ask", + "status": "active", + }, + ) + self._patch_sticky(monkeypatch, set()) + + tools_dict = self._tools_dict() + executor.prepare_tools_for_llm(tools_dict) + call = self._make_call(command="ls -la /tmp") + result = executor.check_pause(tools_dict, call, "OpenAILLM") + assert result is not None + assert result["pause_type"] == "awaiting_approval" + + def test_ask_mode_sticky_match_no_pause(self, monkeypatch): + """``ask`` + a command matching a stored sticky pattern auto-approves.""" + executor = ToolExecutor(user="alice") + self._patch_device( + monkeypatch, + { + "id": "dev_abc", + "approval_mode": "ask", + "status": "active", + }, + ) + # The stub normalizer turns ``ls -la /tmp`` into ``ls *``. + self._patch_sticky(monkeypatch, {"ls *"}) + + tools_dict = self._tools_dict() + executor.prepare_tools_for_llm(tools_dict) + call = self._make_call(command="ls -la /tmp") + assert executor.check_pause(tools_dict, call, "OpenAILLM") is None + + +@pytest.mark.unit +class TestCheckPauseRemoteDeviceHeadless: + """Headless gating for ``remote_device``. + + A denylist-forced prompt must NOT be bypassed by the run's + ``tool_allowlist`` — otherwise a scheduled/headless run with the device + allowlisted would auto-execute even a denylisted command. Normal + (non-forced) approvals keep the allowlist bypass. + """ + + def _make_call(self, name="run_command", call_id="c1", command="whoami"): + import json + + call = Mock() + call.name = name + call.id = call_id + call.arguments = json.dumps({"command": command}) + call.thought_signature = None + return call + + def _tools_dict(self, *, device_id="dev_abc"): + return { + "rd0": { + "id": "rd0", + "name": "remote_device", + "config": {"device_id": device_id}, + "actions": [ + { + "name": "run_command", + "description": "Execute on device", + "active": True, + "require_approval": False, + "parameters": { + "type": "object", + "properties": { + "command": { + "type": "string", + "filled_by_llm": True, + "value": "", + }, + }, + "required": ["command"], + }, + } + ], + } + } + + def _patch_device(self, monkeypatch, device): + from application.agents.tools import remote_device + + monkeypatch.setattr( + remote_device.RemoteDeviceTool, + "_load_device", + lambda self: device, + ) + + def _patch_sticky(self, monkeypatch, patterns): + from application.agents.tools import remote_device + + monkeypatch.setattr( + remote_device, + "normalize_command", + lambda cmd: cmd.split()[0] + " *" if cmd else "", + ) + captured = {"patterns": set(patterns)} + + class _StubRepo: + def __init__(self, conn): + pass + + def has_pattern(self, device_id, user_id, pattern): + return pattern in captured["patterns"] + + monkeypatch.setattr( + remote_device, "DeviceAutoApprovePatternsRepository", _StubRepo + ) + + class _StubConn: + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + monkeypatch.setattr(remote_device, "db_readonly", lambda: _StubConn()) + + def test_headless_allowlisted_denylisted_command_denied(self, monkeypatch): + """Device allowlisted, but ``rm -rf /`` is a forced prompt -> denied.""" + executor = ToolExecutor( + user="alice", headless=True, tool_allowlist={"rd0"} + ) + self._patch_device( + monkeypatch, + {"id": "dev_abc", "approval_mode": "full", "status": "active"}, + ) + self._patch_sticky(monkeypatch, set()) + + tools_dict = self._tools_dict() + executor.prepare_tools_for_llm(tools_dict) + call = self._make_call(command="rm -rf /") + result = executor.check_pause(tools_dict, call, "OpenAILLM") + assert result is not None + assert result["pause_type"] == "headless_denied" + assert result["error_type"] == "tool_not_allowed" + + def test_headless_allowlisted_normal_command_executes(self, monkeypatch): + """Device allowlisted + a normal command in ``full`` -> executes.""" + executor = ToolExecutor( + user="alice", headless=True, tool_allowlist={"rd0"} + ) + self._patch_device( + monkeypatch, + {"id": "dev_abc", "approval_mode": "full", "status": "active"}, + ) + self._patch_sticky(monkeypatch, set()) + + tools_dict = self._tools_dict() + executor.prepare_tools_for_llm(tools_dict) + call = self._make_call(command="whoami") + assert executor.check_pause(tools_dict, call, "OpenAILLM") is None + + def test_headless_not_allowlisted_denied(self, monkeypatch): + """Without the allowlist, an ask-mode command is denied headless.""" + executor = ToolExecutor(user="alice", headless=True, tool_allowlist=set()) + self._patch_device( + monkeypatch, + {"id": "dev_abc", "approval_mode": "ask", "status": "active"}, + ) + self._patch_sticky(monkeypatch, set()) + + tools_dict = self._tools_dict() + executor.prepare_tools_for_llm(tools_dict) + call = self._make_call(command="ls -la /tmp") + result = executor.check_pause(tools_dict, call, "OpenAILLM") + assert result is not None + assert result["pause_type"] == "headless_denied" + + def test_headless_allowlisted_ask_mode_normal_executes(self, monkeypatch): + """ask-mode + allowlisted + non-denylisted -> bypass allowed.""" + executor = ToolExecutor( + user="alice", headless=True, tool_allowlist={"rd0"} + ) + self._patch_device( + monkeypatch, + {"id": "dev_abc", "approval_mode": "ask", "status": "active"}, + ) + self._patch_sticky(monkeypatch, set()) + + tools_dict = self._tools_dict() + executor.prepare_tools_for_llm(tools_dict) + call = self._make_call(command="ls -la /tmp") + # Not denylist-forced, so the allowlist bypass applies. + assert executor.check_pause(tools_dict, call, "OpenAILLM") is None + + @pytest.mark.unit class TestToolExecutorExecute: diff --git a/tests/core/test_openai_compatible.py b/tests/core/test_openai_compatible.py index f8f52987..4001197a 100644 --- a/tests/core/test_openai_compatible.py +++ b/tests/core/test_openai_compatible.py @@ -74,7 +74,14 @@ def _write_together_yaml(directory: Path) -> Path: @pytest.fixture(autouse=True) -def _reset_registry(): +def _reset_registry(monkeypatch): + # Builtin openai_compatible catalogs (e.g. deepseek.yaml) activate when + # their api_key_env is present in os.environ, which the openai_compatible + # provider reads directly, independent of the patched settings. Clear it + # so a key leaked by another test (or present in the dev .env) doesn't add + # extra models to the exact-match assertions below. Mirrors the delenv + # guard in test_model_registry_yaml.py. + monkeypatch.delenv("DEEPSEEK_API_KEY", raising=False) ModelRegistry.reset() yield ModelRegistry.reset() diff --git a/tests/devices/__init__.py b/tests/devices/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/devices/test_auto_approve_validation.py b/tests/devices/test_auto_approve_validation.py new file mode 100644 index 00000000..bab1d240 --- /dev/null +++ b/tests/devices/test_auto_approve_validation.py @@ -0,0 +1,44 @@ +"""Input validation for the add-auto-approve-pattern route. + +These exercise ``add_auto_approve_pattern`` directly. The non-string and +missing-command cases short-circuit with a 400 before any DB access, so no +Postgres fixture is needed. +""" + +from __future__ import annotations + +import pytest +from flask import Flask + +from application.api.devices import routes as routes_module + + +@pytest.fixture +def app(): + return Flask(__name__) + + +def _call(app: Flask, body): + with app.test_request_context( + "/api/devices/dev_x/auto-approve", + method="POST", + json=body, + ): + from flask import request as flask_request + + flask_request.decoded_token = {"sub": "user_abc"} + return routes_module.add_auto_approve_pattern("dev_x") + + +@pytest.mark.parametrize("bad_command", [123, 12.5, True, ["ls"], {"c": 1}]) +def test_add_auto_approve_rejects_non_string_command(app, bad_command): + # A truthy-but-non-string command must 400 rather than crash in shlex. + response = _call(app, {"command": bad_command}) + assert response.status_code == 400 + assert response.get_json()["error"] == "invalid_command" + + +def test_add_auto_approve_missing_command(app): + response = _call(app, {}) + assert response.status_code == 400 + assert response.get_json()["error"] == "missing_command" diff --git a/tests/devices/test_broker_cleanup.py b/tests/devices/test_broker_cleanup.py new file mode 100644 index 00000000..561cc744 --- /dev/null +++ b/tests/devices/test_broker_cleanup.py @@ -0,0 +1,77 @@ +"""Tests for ``DeviceBroker.cleanup_invocation`` pending removal.""" + +from __future__ import annotations + +from queue import Empty + +from application.devices.broker import DeviceBroker + + +def _dispatch_offline(broker: DeviceBroker): + # No session registered -> the invocation lands in _pending. + envelope = {"invocation_id": "inv_stale", "action": "run_command"} + return broker.dispatch_invocation("dev_x", "user_x", envelope) + + +def test_cleanup_removes_pending_so_register_session_doesnt_replay(): + # A timed-out invocation cleaned up while still pending must NOT be + # drained/executed by a later session registration. + broker = DeviceBroker() + inv = _dispatch_offline(broker) + assert broker._pending.get("dev_x") == [inv] + + broker.cleanup_invocation(inv.invocation_id) + # Pending key gone entirely once the list empties. + assert "dev_x" not in broker._pending + assert broker.get_invocation(inv.invocation_id) is None + + sess = broker.register_session("dev_x", "user_x") + # The stale invocation must not have been re-queued. + assert inv.invocation_id not in sess.invocations + try: + sess.invocation_queue.get_nowait() + raise AssertionError("stale invocation was dispatched after cleanup") + except Empty: + pass + + +def test_cleanup_removes_from_active_session(): + # A completed invocation must not linger in the active session's map, + # or a long-lived SSE session accumulates them unboundedly. + broker = DeviceBroker() + sess = broker.register_session("dev_live", "user_live") + inv = broker.dispatch_invocation( + "dev_live", "user_live", + {"invocation_id": "inv_live", "action": "run_command"}, + ) + # Dispatched into the active session. + assert "inv_live" in sess.invocations + assert broker.get_invocation("inv_live") is inv + + # Drive it to completion via a control chunk, then clean up. + broker.submit_output_chunk( + "inv_live", {"stream": "control", "exit_code": 0, "duration_ms": 1} + ) + broker.cleanup_invocation("inv_live") + + assert "inv_live" not in sess.invocations + assert broker.get_invocation("inv_live") is None + + +def test_cleanup_keeps_other_pending_invocations(): + # Cleaning one invocation must leave a sibling pending invocation intact. + broker = DeviceBroker() + inv1 = broker.dispatch_invocation( + "dev_y", "user_y", {"invocation_id": "inv_1", "action": "run_command"} + ) + inv2 = broker.dispatch_invocation( + "dev_y", "user_y", {"invocation_id": "inv_2", "action": "run_command"} + ) + assert broker._pending.get("dev_y") == [inv1, inv2] + + broker.cleanup_invocation(inv1.invocation_id) + assert broker._pending.get("dev_y") == [inv2] + + sess = broker.register_session("dev_y", "user_y") + assert "inv_2" in sess.invocations + assert "inv_1" not in sess.invocations diff --git a/tests/devices/test_broker_drain.py b/tests/devices/test_broker_drain.py new file mode 100644 index 00000000..16dc784f --- /dev/null +++ b/tests/devices/test_broker_drain.py @@ -0,0 +1,51 @@ +"""Tests for ``DeviceBroker.drain_output`` deadline handling.""" + +from __future__ import annotations + +import time + +from application.devices.broker import DeviceBroker, INVOCATION_DONE + + +def _dispatch(broker: DeviceBroker): + envelope = {"invocation_id": "inv_test", "action": "run_command"} + return broker.dispatch_invocation("dev_x", "user_x", envelope) + + +def test_drain_output_returns_after_deadline_with_no_output(): + # Device never connects / posts output and ``completed`` is never set: + # the generator must stop once the deadline passes instead of looping. + broker = DeviceBroker() + inv = _dispatch(broker) + deadline = time.time() + 0.2 + start = time.time() + chunks = list( + broker.drain_output(inv.invocation_id, timeout=0.05, deadline=deadline) + ) + elapsed = time.time() - start + assert chunks == [] + assert not inv.completed.is_set() + assert elapsed < 2.0 + + +def test_drain_output_yields_then_stops_on_done(): + # Sanity: a normal stream still drains and stops on INVOCATION_DONE, + # with a generous deadline that should not fire. + broker = DeviceBroker() + inv = _dispatch(broker) + inv.output_queue.put({"stream": "stdout", "chunk": "hello"}) + inv.output_queue.put(INVOCATION_DONE) + deadline = time.time() + 5.0 + chunks = list( + broker.drain_output(inv.invocation_id, timeout=0.05, deadline=deadline) + ) + assert chunks == [{"stream": "stdout", "chunk": "hello"}] + + +def test_drain_output_no_deadline_stops_when_completed(): + # Without a deadline, the completed event still terminates the loop. + broker = DeviceBroker() + inv = _dispatch(broker) + inv.completed.set() + chunks = list(broker.drain_output(inv.invocation_id, timeout=0.05)) + assert chunks == [] diff --git a/tests/devices/test_denylist.py b/tests/devices/test_denylist.py new file mode 100644 index 00000000..17e9d390 --- /dev/null +++ b/tests/devices/test_denylist.py @@ -0,0 +1,212 @@ +"""Tests for application.devices.denylist.""" + +from application.devices.denylist import check_denylist + + +def test_rm_rf_slash(): + assert check_denylist("rm -rf /") == "rm -rf /" + + +def test_rm_rf_short_flag_order(): + assert check_denylist("rm -fr /") == "rm -rf /" + + +def test_rm_rf_slash_star(): + assert check_denylist("rm -rf /*") == "rm -rf /" + + +def test_rm_rf_long_flags(): + assert check_denylist("rm --recursive --force /") == "rm -rf /" + + +def test_rm_rf_no_preserve_root(): + assert check_denylist("rm -rf --no-preserve-root /") == "rm -rf /" + + +def test_rm_rf_no_preserve_root_before_flags(): + assert check_denylist("rm --no-preserve-root -rf /") == "rm -rf /" + + +def test_rm_rf_no_preserve_root_slash_star(): + assert check_denylist("rm -rf --no-preserve-root /*") == "rm -rf /" + + +def test_rm_rf_long_flags_with_no_preserve_root(): + assert ( + check_denylist("rm --recursive --no-preserve-root --force /") + == "rm -rf /" + ) + + +def test_rm_separated_flags_r_then_f(): + # ``rm -r -f /`` (flags as separate tokens) must still trip. + assert check_denylist("rm -r -f /") == "rm -rf /" + + +def test_rm_separated_flags_f_then_r(): + assert check_denylist("rm -f -r /") == "rm -rf /" + + +def test_rm_rfv_bundle_root(): + assert check_denylist("rm -rfv /") == "rm -rf /" + + +def test_rm_short_recursive_long_force(): + assert check_denylist("rm -r --force /") == "rm -rf /" + + +def test_rm_long_recursive_short_force_slash_star(): + assert check_denylist("rm --recursive -f /*") == "rm -rf /" + + +def test_rm_rf_tmp_subpath_safe(): + assert check_denylist("rm -rf /tmp/foo") is None + + +def test_rm_rf_relative_build_safe(): + assert check_denylist("rm -rf ./build") is None + + +def test_rm_rf_home_subpath_safe(): + assert check_denylist("rm -rf /home/x") is None + + +def test_rm_recursive_only_safe(): + # Recursive without force isn't the catastrophic root-wipe form. + assert check_denylist("rm -r /tmp") is None + + +def test_rm_force_only_slash_safe(): + # Force without recursive (``rm -f /``) fails on a non-empty dir; not denied. + assert check_denylist("rm -f /") is None + + +def test_rm_force_only_etc_safe(): + # Force without recursive on a real path is not the root-wipe form. + assert check_denylist("rm -f /etc/hosts") is None + + +def test_rm_separated_flags_subpath_safe(): + # Separated recursive+force on a subpath stays safe. + assert check_denylist("rm -r -f /tmp/x") is None + + +def test_rm_wrapped_timeout_separated_flags(): + # ``timeout 5 rm -r -f /`` — wrapper + separated flags must still trip. + assert check_denylist("timeout 5 rm -r -f /") == "rm -rf /" + + +def test_rm_wrapped_nice_separated_flags(): + assert check_denylist("nice rm -f -r /") == "rm -rf /" + + +def test_rm_wrapped_nohup_bundle(): + assert check_denylist("nohup rm -rf /") == "rm -rf /" + + +def test_rm_wrapped_timeout_duration_slash_star(): + # ``timeout 30s rm -rf /*`` — duration token skipped, root-star target. + assert check_denylist("timeout 30s rm -rf /*") == "rm -rf /" + + +def test_wrapped_safe_ls_not_denied(): + # ``timeout 5 ls`` is harmless even though it's wrapped. + assert check_denylist("timeout 5 ls") is None + + +def test_wrapped_rm_subpath_safe(): + # ``nice rm -rf /tmp/foo`` is a subpath wipe; stays safe. + assert check_denylist("nice rm -rf /tmp/foo") is None + + +def test_rm_rf_home_tilde(): + assert check_denylist("rm -rf ~") == "rm -rf ~" + + +def test_rm_rf_home_var(): + assert check_denylist("rm -rf $HOME") == "rm -rf $HOME" + + +def test_fork_bomb(): + assert check_denylist(":(){:|:&};:") == "fork bomb" + + +def test_fork_bomb_with_spaces(): + assert check_denylist(":(){ :|:& };:") == "fork bomb" + + +def test_dd_to_block_device(): + assert ( + check_denylist("dd if=/dev/zero of=/dev/sda") + == "dd to block device" + ) + + +def test_dd_with_random_to_nvme(): + assert ( + check_denylist("dd if=/dev/random of=/dev/nvme0n1") + == "dd to block device" + ) + + +def test_mkfs(): + assert check_denylist("mkfs.ext4 /dev/sda1") == "mkfs" + + +def test_mkfs_bare(): + assert check_denylist("mkfs /dev/sda1") == "mkfs" + + +def test_shutdown(): + assert check_denylist("shutdown -h now") == "shutdown" + + +def test_halt(): + assert check_denylist("halt") == "halt" + + +def test_poweroff(): + assert check_denylist("poweroff") == "poweroff" + + +def test_init_0(): + assert check_denylist("init 0") == "init 0/6" + + +def test_init_6(): + assert check_denylist("init 6") == "init 0/6" + + +def test_init_other_levels_safe(): + # ``init 3`` shouldn't trigger. + assert check_denylist("init 3") is None + + +def test_git_push_force_long(): + assert check_denylist("git push --force origin main") == "git push --force" + + +def test_git_push_force_with_lease_safe(): + # --force-with-lease is allowed. + assert check_denylist("git push --force-with-lease origin main") is None + + +def test_git_push_short_flag(): + assert check_denylist("git push -f origin main") == "git push --force" + + +def test_git_push_mirror(): + assert check_denylist("git push --mirror origin") == "git push --force" + + +def test_safe_command_no_match(): + assert check_denylist("ls -la /tmp") is None + + +def test_compound_with_dangerous_segment(): + # ``echo safe && rm -rf /`` must still trip. + assert check_denylist("echo safe && rm -rf /") == "rm -rf /" + + +def test_empty_command(): + assert check_denylist("") is None diff --git a/tests/devices/test_me_endpoint.py b/tests/devices/test_me_endpoint.py new file mode 100644 index 00000000..bb42de86 --- /dev/null +++ b/tests/devices/test_me_endpoint.py @@ -0,0 +1,123 @@ +"""Tests for ``GET /api/devices/me``. + +Exercises the route handler directly through a Flask test-request context +so we can mock ``DevicesRepository.find_by_token_hash`` without spinning +up Postgres. +""" + +from __future__ import annotations + +from datetime import datetime, timezone +from unittest.mock import patch + +import pytest +from flask import Flask + +from application.api.devices import auth as auth_module +from application.api.devices import session as session_module + + +@pytest.fixture +def app(): + return Flask(__name__) + + +def _device_row() -> dict: + return { + "id": "dev_abc", + "user_id": "user_xyz", + "name": "laptop", + "hostname": "laptop.local", + "os": "linux", + "arch": "amd64", + "status": "active", + "approval_mode": "full", + "description": "primary desktop", + "paired_at": datetime(2026, 5, 26, 14, 23, tzinfo=timezone.utc), + "last_seen_at": datetime(2026, 5, 26, 14, 30, tzinfo=timezone.utc), + "machine_pubkey_fingerprint": "fpfpfp", + "token_hash": "tokhash", + } + + +def _call_me(app: Flask, token: str | None): + headers = {} + if token is not None: + headers["Authorization"] = f"Bearer {token}" + with app.test_request_context( + "/api/devices/me", method="GET", headers=headers, + ): + return session_module.me() + + +def test_me_returns_device_record_on_valid_token(app): + """Valid token returns the device's own public fields, datetimes ISO.""" + row = _device_row() + + class _Repo: + def __init__(self, _conn): + self._row = row + + def find_by_token_hash(self, _token_hash): + return self._row + + def touch_last_seen(self, _device_id): + pass + + class _Ctx: + def __enter__(self): + return None + + def __exit__(self, *a): + return False + + with patch.object(auth_module, "DevicesRepository", _Repo), \ + patch.object(auth_module, "db_readonly", _Ctx), \ + patch.object(auth_module, "db_session", _Ctx): + response = _call_me(app, "tok_good") + assert response.status_code == 200 + payload = response.get_json() + assert payload["id"] == "dev_abc" + assert payload["name"] == "laptop" + assert payload["hostname"] == "laptop.local" + assert payload["os"] == "linux" + assert payload["status"] == "active" + assert payload["approval_mode"] == "full" + assert payload["description"] == "primary desktop" + assert payload["paired_at"] == "2026-05-26T14:23:00+00:00" + assert payload["last_seen_at"] == "2026-05-26T14:30:00+00:00" + # Internal fields never leak. + assert "token_hash" not in payload + assert "machine_pubkey_fingerprint" not in payload + assert "user_id" not in payload + + +def test_me_rejects_missing_authorization(app): + """No Authorization header -> 401 missing_token, route handler not entered.""" + response = _call_me(app, token=None) + assert response.status_code == 401 + assert response.get_json()["error"] == "missing_token" + + +def test_me_rejects_invalid_token(app): + """Unknown token -> 401 invalid_token.""" + + class _Repo: + def __init__(self, _conn): + pass + + def find_by_token_hash(self, _token_hash): + return None + + class _Ctx: + def __enter__(self): + return None + + def __exit__(self, *a): + return False + + with patch.object(auth_module, "DevicesRepository", _Repo), \ + patch.object(auth_module, "db_readonly", _Ctx): + response = _call_me(app, "tok_bogus") + assert response.status_code == 401 + assert response.get_json()["error"] == "invalid_token" diff --git a/tests/devices/test_normalizer.py b/tests/devices/test_normalizer.py new file mode 100644 index 00000000..f77a9b8f --- /dev/null +++ b/tests/devices/test_normalizer.py @@ -0,0 +1,81 @@ +"""Tests for application.devices.normalizer.""" + +from application.devices.normalizer import normalize_command, normalize_segment + + +def test_normalize_git_checkout(): + assert normalize_segment("git checkout main -- file.txt") == "git checkout *" + + +def test_normalize_npm_install(): + assert normalize_segment("npm install foo bar") == "npm install *" + + +def test_normalize_ls_with_args(): + assert normalize_segment("ls -la /tmp") == "ls *" + + +def test_normalize_cat_path(): + assert normalize_segment("cat /etc/passwd") == "cat *" + + +def test_normalize_no_args(): + assert normalize_segment("pwd") == "pwd" + + +def test_normalize_two_word_no_args(): + # ``git status`` has no additional args; pattern is just ``git status``. + assert normalize_segment("git status") == "git status" + + +def test_normalize_empty(): + assert normalize_segment("") == "" + + +def test_normalize_command_compound_joins_all_segments(): + # All segments normalized and joined; not just the first. + assert normalize_command("ls -la && rm -rf /tmp") == "ls * && rm *" + + +def test_normalize_command_compound_distinguishes_tail(): + # Sticky for ``ls /tmp && whoami`` must not match ``ls /tmp && rm /tmp/x``. + assert normalize_command("ls /tmp && whoami") == "ls * && whoami" + assert ( + normalize_command("ls /tmp && rm /tmp/x") + != normalize_command("ls /tmp && whoami") + ) + + +def test_normalize_command_single_segment_unchanged(): + # A non-compound command yields just its own pattern (no `` && ``). + assert normalize_command("ls -la /tmp") == "ls *" + + +def test_normalize_command_skips_empty_segments(): + # Trailing/empty connector segments are dropped, not joined as blanks. + assert normalize_command("ls -la && ") == "ls *" + + +def test_normalize_command_none_when_empty(): + assert normalize_command("") is None + + +def test_normalize_command_none_when_only_connectors(): + assert normalize_command("&& ;") is None + + +def test_normalize_quoted_args(): + assert normalize_segment('echo "hello world"') == "echo *" + + +def test_normalize_docker_run(): + assert normalize_segment("docker run -it nginx") == "docker run *" + + +def test_normalize_kubectl_apply(): + assert normalize_segment("kubectl apply -f deploy.yaml") == "kubectl apply *" + + +def test_normalize_known_two_word_minimal(): + # Just ``npm`` with no subcmd. + assert normalize_segment("npm") == "npm" diff --git a/tests/devices/test_pairing_create.py b/tests/devices/test_pairing_create.py new file mode 100644 index 00000000..97a81c81 --- /dev/null +++ b/tests/devices/test_pairing_create.py @@ -0,0 +1,283 @@ +"""Tests for the create-pairing endpoint's optional preset fields. + +Covers the friction-reducer A change: ``POST /api/devices/pairings`` now +accepts ``name``, ``description``, and ``approval_mode`` in the body and +stashes them in Redis so ``redeem_pairing`` can apply them to the new +``devices`` row. These tests exercise the function directly (no Flask +test client) because the broader pairing flow needs Redis + Postgres +fixtures that aren't wired into this module yet. +""" + +from __future__ import annotations + +import json +from typing import Optional +from unittest.mock import MagicMock, patch + +import pytest + +from flask import Flask + +from application.api.devices import pairing as pairing_module + + +class _StubRedis: + """Minimal Redis stand-in capturing setex payloads. + + ``eval`` re-implements the redeem-claim Lua: read the pairing JSON and, + only if ``status == 'pending'``, flip it to ``redeemed`` (returning 1), + else return 0. This lets the redeem tests exercise the atomic path. + """ + + def __init__(self) -> None: + self.setex_calls: list[tuple[str, int, str]] = [] + self.store: dict[str, str] = {} + + def setex(self, key: str, ttl: int, value: str) -> None: + self.setex_calls.append((key, ttl, value)) + self.store[key] = value + + def get(self, key: str) -> Optional[str]: + return self.store.get(key) + + def delete(self, key: str) -> None: + self.store.pop(key, None) + + def eval(self, _script: str, _numkeys: int, key: str, *_args) -> int: + raw = self.store.get(key) + if raw is None: + return 0 + try: + state = json.loads(raw) + except Exception: + return 0 + if state.get("status") != "pending": + return 0 + state["status"] = "redeemed" + self.store[key] = json.dumps(state) + return 1 + + +@pytest.fixture +def app(): + return Flask(__name__) + + +def _call_create(app: Flask, body: dict, redis: _StubRedis): + with app.test_request_context( + "/api/devices/pairings", + method="POST", + json=body, + ): + from flask import request as flask_request + + flask_request.decoded_token = {"sub": "user_abc"} + with patch.object(pairing_module, "_redis", return_value=redis): + return pairing_module.create_pairing() + + +def _stashed_state(redis: _StubRedis) -> dict: + assert redis.setex_calls, "no redis writes recorded" + # The first setex is the device_code -> state hash; the second is the + # user_code -> device_code index. + raw = redis.setex_calls[0][2] + return json.loads(raw) + + +def test_create_pairing_stashes_name_description_and_mode(app): + redis = _StubRedis() + response = _call_create( + app, + { + "name": " laptop ", + "description": " desk machine ", + "approval_mode": "full", + }, + redis, + ) + payload = response.get_json() + assert response.status_code == 200 + assert payload["user_code"], "expected a user_code in the response" + state = _stashed_state(redis) + assert state["requested_name"] == "laptop" + assert state["requested_description"] == "desk machine" + assert state["requested_approval_mode"] == "full" + + +def test_create_pairing_defaults_when_body_omitted(app): + redis = _StubRedis() + response = _call_create(app, {}, redis) + assert response.status_code == 200 + state = _stashed_state(redis) + assert state["requested_name"] is None + assert state["requested_description"] is None + assert state["requested_approval_mode"] is None + + +@pytest.mark.parametrize("mode", ["bogus", "writes-only", "never"]) +def test_create_pairing_rejects_invalid_approval_mode(app, mode): + redis = _StubRedis() + response = _call_create(app, {"approval_mode": mode}, redis) + assert response.status_code == 400 + body = response.get_json() + assert body["error"] == "invalid_approval_mode" + assert not redis.setex_calls + + +@pytest.mark.parametrize("bad_name", [123, 12.5, True, ["x"], {"a": 1}]) +def test_create_pairing_rejects_non_string_name(app, bad_name): + # A non-string ``name`` must 400 instead of crashing on ``.strip()``. + redis = _StubRedis() + response = _call_create(app, {"name": bad_name}, redis) + assert response.status_code == 400 + assert response.get_json()["error"] == "invalid_name" + assert not redis.setex_calls + + +def test_create_pairing_blank_strings_normalize_to_none(app): + redis = _StubRedis() + response = _call_create( + app, + {"name": " ", "description": ""}, + redis, + ) + assert response.status_code == 200 + state = _stashed_state(redis) + assert state["requested_name"] is None + assert state["requested_description"] is None + + +class _StubDevicesRepo: + def __init__(self) -> None: + self.create_args: dict = {} + + def create(self, **kwargs) -> dict: + self.create_args = kwargs + return {"id": kwargs.get("device_id"), **kwargs} + + +class _StubConn: + def __init__(self) -> None: + self.executed: list = [] + + def execute(self, *args, **kwargs): + self.executed.append((args, kwargs)) + # Return an object with .fetchone()/.rowcount where needed; for + # the upsert helper we only need .fetchone() returning None. + return MagicMock(fetchone=lambda: None, rowcount=0) + + +def test_redeem_applies_stashed_preferences(app, monkeypatch): + """End-to-end: stash via create, then redeem applies the values.""" + redis = _StubRedis() + # Step 1 — create with preset name/description/mode + create_response = _call_create( + app, + { + "name": "edgebox", + "description": "raspberry pi", + "approval_mode": "full", + }, + redis, + ) + create_payload = create_response.get_json() + user_code = create_payload["user_code"].replace("-", "") + # Step 2 — simulate the CLI hitting redeem. + devices_repo = _StubDevicesRepo() + monkeypatch.setattr( + pairing_module, "DevicesRepository", lambda conn: devices_repo + ) + monkeypatch.setattr( + pairing_module, + "_upsert_remote_device_user_tool", + lambda conn, **kwargs: None, + ) + + class _Sess: + def __enter__(self): + return _StubConn() + + def __exit__(self, *args): + return False + + monkeypatch.setattr(pairing_module, "db_session", _Sess) + monkeypatch.setattr( + pairing_module, + "fingerprint_pubkey", + lambda _pubkey: "fpfpfp", + ) + monkeypatch.setattr( + pairing_module, + "hash_session_token", + lambda _tok: "tokhash", + ) + + with app.test_request_context( + "/api/devices/pairings/redeem", + method="POST", + json={ + "user_code": user_code, + "hostname": "myhost", + "os": "linux", + "arch": "amd64", + "cli_version": "0.5.0", + "machine_pubkey": "Zm9v", + }, + ): + with patch.object(pairing_module, "_redis", return_value=redis): + response = pairing_module.redeem_pairing() + assert response.status_code == 200, response.get_json() + args = devices_repo.create_args + assert args["approval_mode"] == "full" + assert args["description"] == "raspberry pi" + # Name prefix is the user-supplied value (collision suffix appended). + assert args["name"].startswith("edgebox") + + +def test_redeem_falls_back_to_hostname_and_ask(app, monkeypatch): + redis = _StubRedis() + create_response = _call_create(app, {}, redis) + user_code = create_response.get_json()["user_code"].replace("-", "") + devices_repo = _StubDevicesRepo() + monkeypatch.setattr( + pairing_module, "DevicesRepository", lambda conn: devices_repo + ) + monkeypatch.setattr( + pairing_module, + "_upsert_remote_device_user_tool", + lambda conn, **kwargs: None, + ) + + class _Sess: + def __enter__(self): + return _StubConn() + + def __exit__(self, *args): + return False + + monkeypatch.setattr(pairing_module, "db_session", _Sess) + monkeypatch.setattr( + pairing_module, "fingerprint_pubkey", lambda _pubkey: "fpfpfp", + ) + monkeypatch.setattr( + pairing_module, "hash_session_token", lambda _tok: "tokhash", + ) + + with app.test_request_context( + "/api/devices/pairings/redeem", + method="POST", + json={ + "user_code": user_code, + "hostname": "fallback-host", + "os": "linux", + "arch": "amd64", + "machine_pubkey": "Zm9v", + }, + ): + with patch.object(pairing_module, "_redis", return_value=redis): + response = pairing_module.redeem_pairing() + assert response.status_code == 200 + args = devices_repo.create_args + assert args["approval_mode"] == "ask" + assert args["description"] is None + assert args["name"].startswith("fallback-host") diff --git a/tests/devices/test_pairing_redeem_atomic.py b/tests/devices/test_pairing_redeem_atomic.py new file mode 100644 index 00000000..ed3d09c3 --- /dev/null +++ b/tests/devices/test_pairing_redeem_atomic.py @@ -0,0 +1,269 @@ +"""Concurrency test: two redeems of one user_code yield exactly one device. + +Exercises the real Lua claim against the local dev Redis (127.0.0.1:6379). +Skips cleanly if Redis isn't reachable so CI without Redis still passes. +The DB layer is stubbed; we only count how many device rows get created. +""" + +from __future__ import annotations + +import json +import time +import uuid +from unittest.mock import MagicMock, patch + +import pytest +from flask import Flask +from sqlalchemy.exc import IntegrityError + +from application.api.devices import pairing as pairing_module + + +def _redis_or_skip(): + redis = pytest.importorskip("redis") + try: + client = redis.Redis(host="127.0.0.1", port=6379, socket_connect_timeout=1) + client.ping() + except Exception: # noqa: BLE001 - any connection failure -> skip + pytest.skip("local Redis (127.0.0.1:6379) not available") + return client + + +class _CountingDevicesRepo: + creates = 0 + + def __init__(self, conn): + pass + + def create(self, **kwargs): + type(self).creates += 1 + return {"id": kwargs.get("device_id"), **kwargs} + + +class _StubConn: + def execute(self, *args, **kwargs): + return MagicMock(fetchone=lambda: None, rowcount=0) + + +class _Sess: + def __enter__(self): + return _StubConn() + + def __exit__(self, *args): + return False + + +@pytest.fixture +def app(): + return Flask(__name__) + + +def _seed_pairing(redis_client, user_id: str) -> str: + """Write a pending pairing + user-code index; return the raw user_code.""" + device_code = f"dc_{uuid.uuid4().hex}" + user_code = "".join( + pairing_module.secrets.choice(pairing_module._USER_CODE_ALPHABET) + for _ in range(pairing_module._USER_CODE_LEN) + ) + state = { + "device_code": device_code, + "user_code": pairing_module._format_user_code(user_code), + "user_code_raw": user_code, + "user_id": user_id, + "status": "pending", + "created_at": int(time.time()), + "requested_name": None, + "requested_description": None, + "requested_approval_mode": None, + } + redis_client.setex( + pairing_module._PAIRING_REDIS_PREFIX + device_code, 300, json.dumps(state) + ) + redis_client.setex( + pairing_module._user_code_index_key(user_code), 300, device_code + ) + return user_code + + +def _redeem_once(app: Flask, redis_client, user_code: str): + with app.test_request_context( + "/api/devices/pairings/redeem", + method="POST", + json={ + "user_code": user_code, + "hostname": "concurrent-host", + "os": "linux", + "arch": "amd64", + "machine_pubkey": "Zm9v", + }, + ): + with patch.object(pairing_module, "_redis", return_value=redis_client): + return pairing_module.redeem_pairing() + + +def test_two_redeems_of_one_code_create_exactly_one_device(app, monkeypatch): + redis_client = _redis_or_skip() + _CountingDevicesRepo.creates = 0 + monkeypatch.setattr(pairing_module, "DevicesRepository", _CountingDevicesRepo) + monkeypatch.setattr( + pairing_module, "_upsert_remote_device_user_tool", lambda conn, **k: None + ) + monkeypatch.setattr(pairing_module, "db_session", _Sess) + monkeypatch.setattr(pairing_module, "fingerprint_pubkey", lambda _p: "fpfpfp") + monkeypatch.setattr(pairing_module, "hash_session_token", lambda _t: "tokhash") + + user_code = _seed_pairing(redis_client, f"user_{uuid.uuid4().hex}") + try: + first = _redeem_once(app, redis_client, user_code) + second = _redeem_once(app, redis_client, user_code) + finally: + # Best-effort cleanup of the seeded keys. + device_code = redis_client.get( + pairing_module._user_code_index_key(user_code) + ) + if device_code: + if isinstance(device_code, (bytes, bytearray)): + device_code = device_code.decode() + redis_client.delete(pairing_module._PAIRING_REDIS_PREFIX + device_code) + redis_client.delete(pairing_module._user_code_index_key(user_code)) + + statuses = [first.status_code, second.status_code] + # Exactly one redeem wins (200); the other is rejected. Sequentially the + # winner deletes the user-code index so the loser sees 404; in a true + # race the loser loses the atomic claim and sees 409. Either way, no + # second device. + assert first.status_code == 200, first.get_json() + assert second.status_code in (404, 409), second.get_json() + assert statuses.count(200) == 1 + assert _CountingDevicesRepo.creates == 1 + + +def _fake_integrity_error() -> IntegrityError: + return IntegrityError("INSERT ...", {}, Exception("duplicate key")) + + +class _CollideOnceRepo: + """``create`` raises IntegrityError on its first call, then succeeds.""" + + calls = 0 + names: list = [] + + def __init__(self, conn): + pass + + def create(self, **kwargs): + type(self).calls += 1 + type(self).names.append(kwargs.get("name")) + if type(self).calls == 1: + raise _fake_integrity_error() + return {"id": kwargs.get("device_id"), **kwargs} + + +class _AlwaysCollideRepo: + """``create`` always raises IntegrityError (every name collides).""" + + calls = 0 + + def __init__(self, conn): + pass + + def create(self, **kwargs): + type(self).calls += 1 + raise _fake_integrity_error() + + +def test_redeem_retries_name_collision_then_succeeds(app, monkeypatch): + redis_client = _redis_or_skip() + _CollideOnceRepo.calls = 0 + _CollideOnceRepo.names = [] + monkeypatch.setattr(pairing_module, "DevicesRepository", _CollideOnceRepo) + monkeypatch.setattr( + pairing_module, "_upsert_remote_device_user_tool", lambda conn, **k: None + ) + monkeypatch.setattr(pairing_module, "db_session", _Sess) + monkeypatch.setattr(pairing_module, "fingerprint_pubkey", lambda _p: "fpfpfp") + monkeypatch.setattr(pairing_module, "hash_session_token", lambda _t: "tokhash") + + user_code = _seed_pairing(redis_client, f"user_{uuid.uuid4().hex}") + try: + resp = _redeem_once(app, redis_client, user_code) + finally: + device_code = redis_client.get( + pairing_module._user_code_index_key(user_code) + ) + if device_code: + if isinstance(device_code, (bytes, bytearray)): + device_code = device_code.decode() + redis_client.delete(pairing_module._PAIRING_REDIS_PREFIX + device_code) + redis_client.delete(pairing_module._user_code_index_key(user_code)) + + # Redeem still succeeds; the second name attempt won. + assert resp.status_code == 200, resp.get_json() + assert _CollideOnceRepo.calls == 2 + # The winning name differs from the first (colliding) candidate. + assert _CollideOnceRepo.names[0] != _CollideOnceRepo.names[1] + assert resp.get_json()["name"] == _CollideOnceRepo.names[1] + + +def test_redeem_releases_claim_when_all_names_collide(app, monkeypatch): + redis_client = _redis_or_skip() + _AlwaysCollideRepo.calls = 0 + monkeypatch.setattr(pairing_module, "DevicesRepository", _AlwaysCollideRepo) + monkeypatch.setattr( + pairing_module, "_upsert_remote_device_user_tool", lambda conn, **k: None + ) + monkeypatch.setattr(pairing_module, "db_session", _Sess) + monkeypatch.setattr(pairing_module, "fingerprint_pubkey", lambda _p: "fpfpfp") + monkeypatch.setattr(pairing_module, "hash_session_token", lambda _t: "tokhash") + + user_code = _seed_pairing(redis_client, f"user_{uuid.uuid4().hex}") + device_code = redis_client.get( + pairing_module._user_code_index_key(user_code) + ) + if isinstance(device_code, (bytes, bytearray)): + device_code = device_code.decode() + try: + resp = _redeem_once(app, redis_client, user_code) + + # Every attempt collided -> error returned. + assert resp.status_code == 500, resp.get_json() + assert _AlwaysCollideRepo.calls == pairing_module._DEVICE_NAME_RETRIES + + # The claim was released: status reverted to pending and the + # user-code index restored, so the code can be redeemed again. + state = json.loads( + redis_client.get(pairing_module._PAIRING_REDIS_PREFIX + device_code) + ) + assert state["status"] == "pending" + restored = redis_client.get( + pairing_module._user_code_index_key(user_code) + ) + if isinstance(restored, (bytes, bytearray)): + restored = restored.decode() + assert restored == device_code + finally: + redis_client.delete(pairing_module._PAIRING_REDIS_PREFIX + device_code) + redis_client.delete(pairing_module._user_code_index_key(user_code)) + + +def test_claim_pairing_wins_exactly_once(): + # Direct test of the atomic claim against real Redis: two claims on the + # same pending pairing key -> first True, second False. This is the 409 + # race path (no index deletion involved). + redis_client = _redis_or_skip() + device_code = f"dc_{uuid.uuid4().hex}" + key = pairing_module._PAIRING_REDIS_PREFIX + device_code + redis_client.setex(key, 60, json.dumps({"status": "pending", "user_id": "u"})) + try: + first = pairing_module._claim_pairing(redis_client, device_code) + second = pairing_module._claim_pairing(redis_client, device_code) + assert first is True + assert second is False + # TTL preserved (claim must not drop expiry). + assert redis_client.ttl(key) > 0 + # State flipped to redeemed, other fields intact. + stored = json.loads(redis_client.get(key)) + assert stored["status"] == "redeemed" + assert stored["user_id"] == "u" + finally: + redis_client.delete(key) diff --git a/tests/devices/test_remote_device_decision.py b/tests/devices/test_remote_device_decision.py new file mode 100644 index 00000000..a5452a90 --- /dev/null +++ b/tests/devices/test_remote_device_decision.py @@ -0,0 +1,58 @@ +"""Tests for ``RemoteDeviceTool._decide_approval`` (two-mode model).""" + +from __future__ import annotations + +import pytest + +from application.agents.tools.remote_device import RemoteDeviceTool + + +def _tool(monkeypatch, *, sticky=False): + tool = RemoteDeviceTool.__new__(RemoteDeviceTool) + tool.device_id = "dev_abc" + tool.user_id = "alice" + monkeypatch.setattr(tool, "_matches_sticky", lambda command: sticky) + return tool + + +@pytest.mark.unit +class TestDecideApproval: + def test_full_passthrough(self, monkeypatch): + tool = _tool(monkeypatch) + reason, mode = tool._decide_approval({"approval_mode": "full"}, "rm /tmp/x") + assert (reason, mode) == ("full_access_passthrough", "full") + + def test_full_denylist_forces_prompt(self, monkeypatch): + tool = _tool(monkeypatch) + reason, mode = tool._decide_approval({"approval_mode": "full"}, "rm -rf /") + assert (reason, mode) == ("denylist_forced_prompt", "ask") + + def test_ask_requires_approval(self, monkeypatch): + tool = _tool(monkeypatch, sticky=False) + reason, mode = tool._decide_approval({"approval_mode": "ask"}, "ls -la") + assert (reason, mode) == ("user_approval_required", "ask") + + def test_ask_sticky_auto_approves(self, monkeypatch): + tool = _tool(monkeypatch, sticky=True) + reason, mode = tool._decide_approval({"approval_mode": "ask"}, "ls -la") + assert (reason, mode) == ("sticky_auto_approve", "full") + + def test_ask_sticky_denylist_still_forces_prompt(self, monkeypatch): + # Regression: a "don't ask again" sticky pattern must not let a + # denylisted command (e.g. sticky ``rm *`` -> ``rm -rf /``, or sticky + # ``git push *`` -> ``git push --force --mirror``) auto-run. The + # denylist forces a prompt on every path, including the sticky one. + from application.devices.denylist import check_denylist + + assert check_denylist("rm -rf /") is not None + assert check_denylist("git push --force --mirror") is not None + + tool = _tool(monkeypatch, sticky=True) + for command in ("rm -rf /", "git push --force --mirror"): + reason, mode = tool._decide_approval({"approval_mode": "ask"}, command) + assert (reason, mode) == ("denylist_forced_prompt", "ask") + + def test_missing_mode_defaults_to_ask(self, monkeypatch): + tool = _tool(monkeypatch, sticky=False) + reason, mode = tool._decide_approval({}, "ls -la") + assert (reason, mode) == ("user_approval_required", "ask") diff --git a/tests/devices/test_session_ticket.py b/tests/devices/test_session_ticket.py new file mode 100644 index 00000000..9a235af0 --- /dev/null +++ b/tests/devices/test_session_ticket.py @@ -0,0 +1,217 @@ +"""Tests for poll-issued session-ticket enforcement on the SSE upgrade. + +Two layers: + +* Broker unit tests for ``claim_ticket`` / ``validate_ticket`` (issue, match, + mismatch, expiry, absence). +* Route tests proving the real CLI loop still works: ``/poll`` issues a + ticket and ``session_events`` accepts *that* ticket, while a mismatched + ticket is rejected with ``410`` before any stream opens. +""" + +from __future__ import annotations + +from unittest.mock import patch + +import pytest +from flask import Flask + +from application.api.devices import auth as auth_module +from application.api.devices import session as session_module +from application.devices.broker import DeviceBroker + + +# --------------------------------------------------------------------------- +# Broker unit tests +# --------------------------------------------------------------------------- +def _queue_work(broker: DeviceBroker, device_id: str = "dev_x") -> None: + # No active session -> the invocation parks in _pending, which is what + # makes claim_ticket hand out a ticket. + broker.dispatch_invocation( + device_id, "user_x", {"invocation_id": "inv_1", "action": "run_command"} + ) + + +def test_claim_ticket_none_when_no_work(): + broker = DeviceBroker() + assert broker.claim_ticket("dev_x", 30) is None + + +def test_validate_accepts_issued_ticket(): + broker = DeviceBroker() + _queue_work(broker) + ticket = broker.claim_ticket("dev_x", 30) + assert ticket is not None + assert broker.validate_ticket("dev_x", ticket) is True + + +def test_validate_rejects_wrong_ticket(): + broker = DeviceBroker() + _queue_work(broker) + broker.claim_ticket("dev_x", 30) + assert broker.validate_ticket("dev_x", "st_not_the_one") is False + + +def test_validate_rejects_when_no_ticket_issued(): + broker = DeviceBroker() + # Never polled -> nothing issued. + assert broker.validate_ticket("dev_x", "st_anything") is False + + +def test_validate_rejects_empty_session_id(): + broker = DeviceBroker() + _queue_work(broker) + broker.claim_ticket("dev_x", 30) + assert broker.validate_ticket("dev_x", "") is False + + +def test_validate_rejects_expired_ticket_and_evicts(): + broker = DeviceBroker() + _queue_work(broker) + with patch("application.devices.broker.time.monotonic", return_value=1000.0): + ticket = broker.claim_ticket("dev_x", 30) + # 31s later the ticket is past its 30s TTL. + with patch("application.devices.broker.time.monotonic", return_value=1031.0): + assert broker.validate_ticket("dev_x", ticket) is False + # Evicted: a second check can't resurrect it. + assert broker.validate_ticket("dev_x", ticket) is False + + +def test_register_session_consumes_issued_ticket_as_session_id(): + # The issued ticket becomes the session_id, so the URL the CLI opens + # (= the ticket) matches the live session. This is the invariant the + # route enforcement relies on. + broker = DeviceBroker() + _queue_work(broker) + ticket = broker.claim_ticket("dev_x", 30) + sess = broker.register_session("dev_x", "user_x") + assert sess.session_id == ticket + + +# --------------------------------------------------------------------------- +# Route tests (poll -> SSE upgrade) +# --------------------------------------------------------------------------- +@pytest.fixture +def app(): + return Flask(__name__) + + +def _device_row() -> dict: + return { + "id": "dev_route", + "user_id": "user_route", + "name": "laptop", + "status": "active", + "approval_mode": "ask", + "machine_pubkey_fingerprint": "fp", + "token_hash": "tokhash", + } + + +class _Repo: + def __init__(self, _conn): + pass + + def find_by_token_hash(self, _token_hash): + return _device_row() + + def touch_last_seen(self, _device_id): + pass + + +class _Ctx: + def __enter__(self): + return None + + def __exit__(self, *a): + return False + + +def _patched_auth(): + return [ + patch.object(auth_module, "DevicesRepository", _Repo), + patch.object(auth_module, "db_readonly", _Ctx), + patch.object(auth_module, "db_session", _Ctx), + ] + + +def _call(app: Flask, view, path: str, *args): + with app.test_request_context( + path, method="GET", headers={"Authorization": "Bearer tok_good"} + ): + ctxs = _patched_auth() + for c in ctxs: + c.start() + try: + return view(*args) + finally: + for c in ctxs: + c.stop() + + +def test_poll_to_sse_upgrade_with_issued_ticket_works(app): + """The legitimate loop: /poll issues a ticket, session_events accepts it.""" + broker = DeviceBroker() + # Queue work so /poll returns a ticket rather than 202. + broker.dispatch_invocation( + "dev_route", "user_route", + {"invocation_id": "inv_a", "action": "run_command"}, + ) + with patch.object(session_module, "get_broker", return_value=broker): + poll_resp = _call(app, session_module.poll, "/api/devices/poll") + assert poll_resp.status_code == 200 + payload = poll_resp.get_json() + ticket = payload["session_ticket"] + assert payload["session_url"] == f"/api/devices/sessions/{ticket}/events" + assert payload["expires_in"] == 30 + + # CLI opens the exact session_url it was handed. + sse_resp = _call( + app, + session_module.session_events, + f"/api/devices/sessions/{ticket}/events", + ticket, + ) + try: + assert sse_resp.status_code == 200 + assert sse_resp.mimetype == "text/event-stream" + finally: + # Close the streamed (stream_with_context) generator so its + # deferred app-context teardown runs now instead of leaking into + # the next test as a "popped wrong app context" error. + sse_resp.close() + + +def test_session_events_rejects_mismatched_ticket(app): + """A fabricated/mismatched session_id is 410 Gone, no stream opened.""" + broker = DeviceBroker() + broker.dispatch_invocation( + "dev_route", "user_route", + {"invocation_id": "inv_b", "action": "run_command"}, + ) + with patch.object(session_module, "get_broker", return_value=broker): + # Poll issues the real ticket... + _call(app, session_module.poll, "/api/devices/poll") + # ...but the client opens a different one. + resp = _call( + app, + session_module.session_events, + "/api/devices/sessions/st_bogus/events", + "st_bogus", + ) + assert resp.status_code == 410 + assert resp.get_json()["error"] == "session_ticket_invalid" + + +def test_session_events_rejects_when_never_polled(app): + """Opening the SSE stream without a prior poll is rejected (410).""" + broker = DeviceBroker() + with patch.object(session_module, "get_broker", return_value=broker): + resp = _call( + app, + session_module.session_events, + "/api/devices/sessions/st_anything/events", + "st_anything", + ) + assert resp.status_code == 410 + assert resp.get_json()["error"] == "session_ticket_invalid" diff --git a/tests/devices/test_signature.py b/tests/devices/test_signature.py new file mode 100644 index 00000000..9d41a49d --- /dev/null +++ b/tests/devices/test_signature.py @@ -0,0 +1,42 @@ +"""Canonical signed-payload tests. + +The string here MUST match the DocsGPT-cli signer byte-for-byte (see +``internal/host/identity.go`` ``CanonicalPayload`` and the Go test +``TestCanonicalPayloadExactString``). If you change the format, change both +repos together. +""" + +from __future__ import annotations + +from application.api.devices.auth import _canonical_payload + + +def test_canonical_payload_empty_body(): + # sha256("") = e3b0c442...b855 + got = _canonical_payload("GET", "/api/devices/poll", "1700000000", b"") + assert got == ( + "GET /api/devices/poll 1700000000 " + "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855" + ) + + +def test_canonical_payload_none_body_equals_empty(): + # The verifier passes request.get_data(); guard against a None slipping in. + assert _canonical_payload("GET", "/x", "1", None) == _canonical_payload( + "GET", "/x", "1", b"" + ) + + +def test_canonical_payload_with_body_matches_cli(): + # sha256("hello") = 2cf24dba...9824 — identical to the Go test constant. + got = _canonical_payload("POST", "/api/devices/x", "1700000000", b"hello") + assert got == ( + "POST /api/devices/x 1700000000 " + "2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824" + ) + + +def test_canonical_payload_body_changes_hash(): + a = _canonical_payload("POST", "/x", "1", b'{"decision":"accept"}') + b = _canonical_payload("POST", "/x", "1", b'{"decision":"deny"}') + assert a != b diff --git a/tests/devices/test_splitter.py b/tests/devices/test_splitter.py new file mode 100644 index 00000000..74f8209f --- /dev/null +++ b/tests/devices/test_splitter.py @@ -0,0 +1,75 @@ +"""Tests for application.devices.splitter.""" + +from application.devices.splitter import head_token, head_tokens, split_command + + +def test_single_segment(): + assert split_command("ls -la") == ["ls -la"] + + +def test_and_split(): + assert split_command("ls && rm -rf /") == ["ls", "rm -rf /"] + + +def test_or_split(): + assert split_command("foo || bar") == ["foo", "bar"] + + +def test_semicolon_split(): + assert split_command("a; b; c") == ["a", "b", "c"] + + +def test_pipe_split(): + assert split_command("cat /etc/passwd | grep root") == [ + "cat /etc/passwd", + "grep root", + ] + + +def test_background_amp_split(): + assert split_command("a & b") == ["a", "b"] + + +def test_pipe_amp_split(): + assert split_command("foo |& bar") == ["foo", "bar"] + + +def test_newline_split(): + assert split_command("a\nb\nc") == ["a", "b", "c"] + + +def test_empty_command(): + assert split_command("") == [] + + +def test_whitespace_only(): + assert split_command(" ") == [] + + +def test_head_token_simple(): + assert head_token("ls -la /tmp") == "ls" + + +def test_head_token_wrapper_timeout(): + assert head_token("timeout 30s ls -la") == "ls" + + +def test_head_token_wrapper_nice(): + # ``nice -n 10 ls`` should strip ``nice -n 10``. + assert head_token("nice -n 10 ls") == "ls" + + +def test_head_token_wrapper_nohup(): + assert head_token("nohup ls") == "ls" + + +def test_head_token_empty(): + assert head_token("") == "" + + +def test_head_tokens_compound(): + assert head_tokens("ls && rm -rf / | grep foo") == ["ls", "rm", "grep"] + + +def test_compound_with_wrappers(): + assert head_tokens("timeout 5 ls && nohup git push") == ["ls", "git"] diff --git a/tests/devices/test_update_device.py b/tests/devices/test_update_device.py new file mode 100644 index 00000000..95af78bb --- /dev/null +++ b/tests/devices/test_update_device.py @@ -0,0 +1,109 @@ +"""Tests for ``PATCH /api/devices/`` body validation. + +The validation rejects bad ``name``/``description`` values before any DB +access, so the 400 cases exercise the handler directly without Postgres. +A valid PATCH is covered with the repository + session mocked. +""" + +from __future__ import annotations + + +import pytest +from flask import Flask + +from application.api.devices import routes as routes_module + + +@pytest.fixture +def app(): + return Flask(__name__) + + +def _call_update(app: Flask, body, device_id: str = "dev_abc"): + with app.test_request_context( + f"/api/devices/{device_id}", method="PATCH", json=body, + ): + from flask import request as flask_request + + flask_request.decoded_token = {"sub": "user_abc"} + return routes_module.update_device(device_id) + + +def test_patch_rejects_null_name(app): + response = _call_update(app, {"name": None}) + assert response.status_code == 400 + assert response.get_json()["error"] == "invalid_name" + + +def test_patch_rejects_blank_name(app): + response = _call_update(app, {"name": " "}) + assert response.status_code == 400 + assert response.get_json()["error"] == "invalid_name" + + +def test_patch_rejects_non_string_name(app): + response = _call_update(app, {"name": 123}) + assert response.status_code == 400 + assert response.get_json()["error"] == "invalid_name" + + +def test_patch_rejects_null_description(app): + response = _call_update(app, {"description": None}) + assert response.status_code == 400 + assert response.get_json()["error"] == "invalid_description" + + +def test_patch_rejects_non_string_description(app): + response = _call_update(app, {"description": ["x"]}) + assert response.status_code == 400 + assert response.get_json()["error"] == "invalid_description" + + +def test_patch_still_validates_approval_mode(app): + response = _call_update(app, {"approval_mode": "bogus"}) + assert response.status_code == 400 + assert response.get_json()["error"] == "invalid_approval_mode" + + +def test_patch_empty_body_is_no_fields(app): + response = _call_update(app, {}) + assert response.status_code == 400 + assert response.get_json()["error"] == "no_fields" + + +def test_patch_valid_name_trimmed_and_persisted(app, monkeypatch): + """A valid PATCH trims the name and writes the stripped value.""" + captured: dict = {} + + class _Repo: + def __init__(self, _conn): + pass + + def update(self, _device_id, _user_id, fields): + captured.update(fields) + return True + + def get(self, device_id, user_id=None): + return {"id": device_id, "name": fields_name(), "description": ""} + + def fields_name(): + return captured.get("name") + + class _Sess: + def __enter__(self): + return None + + def __exit__(self, *a): + return False + + monkeypatch.setattr(routes_module, "DevicesRepository", _Repo) + monkeypatch.setattr(routes_module, "db_session", _Sess) + monkeypatch.setattr( + "application.api.devices.pairing._upsert_remote_device_user_tool", + lambda conn, **kwargs: None, + ) + + response = _call_update(app, {"name": " laptop ", "description": " box "}) + assert response.status_code == 200 + assert captured["name"] == "laptop" + assert captured["description"] == "box" diff --git a/tests/storage/db/repositories/test_message_events_repository.py b/tests/storage/db/repositories/test_message_events_repository.py index 0dcbef1f..aa04f123 100644 --- a/tests/storage/db/repositories/test_message_events_repository.py +++ b/tests/storage/db/repositories/test_message_events_repository.py @@ -134,6 +134,219 @@ class TestMessageEventsRepository: "tool_calls": [], } + def test_reconstruct_partial_includes_paused_tool_call(self, pg_conn): + """A paused (awaiting_approval) ``tool_call`` event must surface + in ``tool_calls`` so a conversation refresh re-renders the + approval bar.""" + message_id = _seed_message(pg_conn) + repo = MessageEventsRepository(pg_conn) + repo.record( + message_id, 0, "tool_call", + { + "type": "tool_call", + "data": { + "tool_name": "remote_device", + "call_id": "c1", + "action_name": "run_command", + "arguments": {"command": "ls"}, + "status": "awaiting_approval", + "device_id": "dev_abc", + }, + }, + ) + partial = repo.reconstruct_partial(message_id) + assert len(partial["tool_calls"]) == 1 + tc = partial["tool_calls"][0] + assert tc["call_id"] == "c1" + assert tc["status"] == "awaiting_approval" + assert tc["device_id"] == "dev_abc" + + def test_reconstruct_partial_tool_call_completed_replaces_pause(self, pg_conn): + """Later events for the same ``call_id`` overlay earlier ones, + so ``completed`` displaces ``awaiting_approval`` on refresh.""" + message_id = _seed_message(pg_conn) + repo = MessageEventsRepository(pg_conn) + repo.record( + message_id, 0, "tool_call", + { + "type": "tool_call", + "data": { + "tool_name": "remote_device", + "call_id": "c1", + "status": "awaiting_approval", + }, + }, + ) + repo.record( + message_id, 1, "tool_call", + { + "type": "tool_call", + "data": { + "tool_name": "remote_device", + "call_id": "c1", + "status": "completed", + "result": "ok", + }, + }, + ) + partial = repo.reconstruct_partial(message_id) + assert len(partial["tool_calls"]) == 1 + assert partial["tool_calls"][0]["status"] == "completed" + assert partial["tool_calls"][0]["result"] == "ok" + + def test_reconstruct_partial_multiple_tool_calls_preserve_order(self, pg_conn): + """First-seen order is preserved when multiple calls interleave.""" + message_id = _seed_message(pg_conn) + repo = MessageEventsRepository(pg_conn) + repo.record( + message_id, 0, "tool_call", + {"type": "tool_call", "data": {"call_id": "c1", "status": "pending"}}, + ) + repo.record( + message_id, 1, "tool_call", + {"type": "tool_call", "data": {"call_id": "c2", "status": "pending"}}, + ) + repo.record( + message_id, 2, "tool_call", + {"type": "tool_call", "data": {"call_id": "c1", "status": "completed"}}, + ) + partial = repo.reconstruct_partial(message_id) + assert [tc["call_id"] for tc in partial["tool_calls"]] == ["c1", "c2"] + assert partial["tool_calls"][0]["status"] == "completed" + assert partial["tool_calls"][1]["status"] == "pending" + + def test_reconstruct_partial_tool_call_overlays_bulk_tool_calls(self, pg_conn): + """Per-call ``tool_call`` events overlay an earlier bulk + ``tool_calls`` emit by ``call_id``.""" + message_id = _seed_message(pg_conn) + repo = MessageEventsRepository(pg_conn) + repo.record( + message_id, 0, "tool_calls", + { + "type": "tool_calls", + "tool_calls": [ + {"call_id": "c1", "status": "pending"}, + {"call_id": "c2", "status": "pending"}, + ], + }, + ) + repo.record( + message_id, 1, "tool_call", + {"type": "tool_call", "data": {"call_id": "c1", "status": "completed"}}, + ) + partial = repo.reconstruct_partial(message_id) + assert len(partial["tool_calls"]) == 2 + by_id = {tc["call_id"]: tc for tc in partial["tool_calls"]} + assert by_id["c1"]["status"] == "completed" + assert by_id["c2"]["status"] == "pending" + + def test_reconstruct_partial_tool_call_skips_missing_call_id(self, pg_conn): + """Malformed per-call events without ``call_id`` are dropped, not crashing.""" + message_id = _seed_message(pg_conn) + repo = MessageEventsRepository(pg_conn) + repo.record( + message_id, 0, "tool_call", + {"type": "tool_call", "data": {"status": "pending"}}, + ) + repo.record( + message_id, 1, "tool_call", + {"type": "tool_call", "data": {"call_id": "c1", "status": "completed"}}, + ) + partial = repo.reconstruct_partial(message_id) + assert len(partial["tool_calls"]) == 1 + assert partial["tool_calls"][0]["call_id"] == "c1" + + def test_reconstruct_partial_paused_then_empty_bulk_preserves_overlay( + self, pg_conn + ): + """Realistic pause sequence: per-call ``awaiting_approval`` followed + by an empty bulk ``tool_calls: []`` from the classic-agent + end-of-turn yield. The empty bulk must not wipe the per-call entry. + """ + message_id = _seed_message(pg_conn) + repo = MessageEventsRepository(pg_conn) + repo.record( + message_id, 0, "tool_call", + { + "type": "tool_call", + "data": { + "tool_name": "remote_device", + "call_id": "c1", + "status": "awaiting_approval", + "device_id": "dev_abc", + }, + }, + ) + repo.record( + message_id, 1, "tool_calls", + {"type": "tool_calls", "tool_calls": []}, + ) + partial = repo.reconstruct_partial(message_id) + assert len(partial["tool_calls"]) == 1 + assert partial["tool_calls"][0]["status"] == "awaiting_approval" + assert partial["tool_calls"][0]["device_id"] == "dev_abc" + + def test_reconstruct_partial_paused_then_bulk_other_call_preserves_overlay( + self, pg_conn + ): + """A bulk emit for a different ``call_id`` must merge alongside the + per-call awaiting_approval entry rather than wipe it. + """ + message_id = _seed_message(pg_conn) + repo = MessageEventsRepository(pg_conn) + repo.record( + message_id, 0, "tool_call", + { + "type": "tool_call", + "data": { + "tool_name": "remote_device", + "call_id": "c1", + "status": "awaiting_approval", + }, + }, + ) + repo.record( + message_id, 1, "tool_calls", + { + "type": "tool_calls", + "tool_calls": [ + {"call_id": "c2", "status": "completed"}, + ], + }, + ) + partial = repo.reconstruct_partial(message_id) + by_id = {tc["call_id"]: tc for tc in partial["tool_calls"]} + assert set(by_id) == {"c1", "c2"} + assert by_id["c1"]["status"] == "awaiting_approval" + assert by_id["c2"]["status"] == "completed" + + def test_reconstruct_partial_completed_then_empty_bulk_preserves_overlay( + self, pg_conn + ): + """A completed per-call entry survives a subsequent empty bulk emit.""" + message_id = _seed_message(pg_conn) + repo = MessageEventsRepository(pg_conn) + repo.record( + message_id, 0, "tool_call", + { + "type": "tool_call", + "data": { + "tool_name": "remote_device", + "call_id": "c1", + "status": "completed", + "result": "ok", + }, + }, + ) + repo.record( + message_id, 1, "tool_calls", + {"type": "tool_calls", "tool_calls": []}, + ) + partial = repo.reconstruct_partial(message_id) + assert len(partial["tool_calls"]) == 1 + assert partial["tool_calls"][0]["status"] == "completed" + assert partial["tool_calls"][0]["result"] == "ok" + def test_reconstruct_partial_skips_malformed_payload(self, pg_conn): message_id = _seed_message(pg_conn) repo = MessageEventsRepository(pg_conn) diff --git a/tests/test_continuation.py b/tests/test_continuation.py index e48a1cc0..d7304fc3 100644 --- a/tests/test_continuation.py +++ b/tests/test_continuation.py @@ -409,6 +409,80 @@ class TestHandlerPauseSignaling: ] assert len(pause_events) == 1 + def test_pause_propagates_device_id_for_remote_device(self): + """``pause_info['device_id']`` (set in tool_executor for the + remote_device tool) must be copied into the emitted ``tool_call`` + event so the approval UI can render the "don't ask again" button.""" + handler = ConcreteHandler() + agent = self._make_agent() + agent.tool_executor.check_pause = Mock(return_value={ + "call_id": "c1", + "name": "run_command_0", + "tool_name": "remote_device", + "tool_id": "0", + "action_name": "run_command", + "arguments": {"command": "ls"}, + "pause_type": "awaiting_approval", + "device_id": "dev_abc", + "thought_signature": None, + }) + + call = ToolCall(id="c1", name="run_command_0", arguments='{"command": "ls"}') + gen = handler.handle_tool_calls( + agent, [call], {"0": {"name": "remote_device"}}, [] + ) + + events = [] + try: + while True: + events.append(next(gen)) + except StopIteration: + pass + + pause_events = [ + e for e in events + if e.get("type") == "tool_call" + and e.get("data", {}).get("status") == "awaiting_approval" + ] + assert len(pause_events) == 1 + assert pause_events[0]["data"].get("device_id") == "dev_abc" + + def test_pause_omits_device_id_for_non_remote_tools(self): + """``device_id`` must NOT leak into pause events for tools that + don't ship one in ``pause_info``.""" + handler = ConcreteHandler() + agent = self._make_agent() + agent.tool_executor.check_pause = Mock(return_value={ + "call_id": "c1", + "name": "send_msg_0", + "tool_name": "telegram", + "tool_id": "0", + "action_name": "send_msg", + "arguments": {"text": "hello"}, + "pause_type": "awaiting_approval", + "thought_signature": None, + }) + + call = ToolCall(id="c1", name="send_msg_0", arguments='{"text": "hello"}') + gen = handler.handle_tool_calls( + agent, [call], {"0": {"name": "telegram"}}, [] + ) + + events = [] + try: + while True: + events.append(next(gen)) + except StopIteration: + pass + + pause_events = [ + e for e in events + if e.get("type") == "tool_call" + and e.get("data", {}).get("status") == "awaiting_approval" + ] + assert len(pause_events) == 1 + assert "device_id" not in pause_events[0]["data"] + def test_mixed_execute_and_pause(self): """One tool executes, another needs approval.""" handler = ConcreteHandler() @@ -989,3 +1063,115 @@ class TestContinuationServiceMarkResuming: svc = cont_mod.ContinuationService() # Not a UUID and no legacy row exists. assert svc.mark_resuming("not-a-uuid", "alice") is False + + +# --------------------------------------------------------------------------- +# Refresh during pause: per-call ``tool_call`` events must reconstruct +# into ``tool_calls`` so the approval bar re-renders. +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +class TestReconstructPartialToolCallReplay: + """The pause path in ``llm/handlers/base.py`` emits a per-call + ``tool_call`` event (status ``awaiting_approval``) — not a bulk + ``tool_calls`` snapshot. ``reconstruct_partial`` must overlay these + so a conversation refresh while paused still shows the approval bar. + """ + + @staticmethod + def _seed_message(conn): + from sqlalchemy import text + + user_id = f"user-{uuid.uuid4().hex[:8]}" + conv_id = uuid.uuid4() + msg_id = uuid.uuid4() + conn.execute( + text("INSERT INTO users (user_id) VALUES (:u)"), + {"u": user_id}, + ) + conn.execute( + text( + "INSERT INTO conversations (id, user_id, name) " + "VALUES (:id, :u, 'test')" + ), + {"id": conv_id, "u": user_id}, + ) + conn.execute( + text( + "INSERT INTO conversation_messages (id, conversation_id, user_id, position) " + "VALUES (:id, :c, :u, 0)" + ), + {"id": msg_id, "c": conv_id, "u": user_id}, + ) + return str(msg_id) + + def test_paused_tool_call_event_lands_in_tool_calls(self, pg_conn): + from application.storage.db.repositories.message_events import ( + MessageEventsRepository, + ) + + message_id = self._seed_message(pg_conn) + repo = MessageEventsRepository(pg_conn) + # Mirror the exact shape emitted at base.py:941-950 on pause. + repo.record( + message_id, + 0, + "tool_call", + { + "type": "tool_call", + "data": { + "tool_name": "remote_device", + "call_id": "call_remote_test_1", + "action_name": "run_command", + "arguments": {"command": "ls -la /tmp"}, + "status": "awaiting_approval", + "device_id": "dev_abc", + }, + }, + ) + partial = repo.reconstruct_partial(message_id) + assert len(partial["tool_calls"]) == 1 + tc = partial["tool_calls"][0] + assert tc["status"] == "awaiting_approval" + assert tc["device_id"] == "dev_abc" + assert tc["call_id"] == "call_remote_test_1" + + def test_completed_event_replaces_paused_event(self, pg_conn): + from application.storage.db.repositories.message_events import ( + MessageEventsRepository, + ) + + message_id = self._seed_message(pg_conn) + repo = MessageEventsRepository(pg_conn) + repo.record( + message_id, + 0, + "tool_call", + { + "type": "tool_call", + "data": { + "tool_name": "remote_device", + "call_id": "c1", + "status": "awaiting_approval", + }, + }, + ) + repo.record( + message_id, + 1, + "tool_call", + { + "type": "tool_call", + "data": { + "tool_name": "remote_device", + "call_id": "c1", + "status": "completed", + "result": "/tmp listing here", + }, + }, + ) + partial = repo.reconstruct_partial(message_id) + assert len(partial["tool_calls"]) == 1 + assert partial["tool_calls"][0]["status"] == "completed" + assert partial["tool_calls"][0]["result"] == "/tmp listing here"