mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-12 22:14:05 +00:00
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
This commit is contained in:
64 files changed
+6733
-56
No files matched your search
@@ -107,6 +107,8 @@ celerybeat.pid
|
||||
|
||||
# Environments
|
||||
.env
|
||||
.env.local
|
||||
.env.*.local
|
||||
.venv
|
||||
# Machine-specific Claude Code guidance (see CLAUDE.md preamble)
|
||||
CLAUDE.md
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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, {})
|
||||
|
||||
@@ -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;")
|
||||
@@ -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'));"
|
||||
)
|
||||
@@ -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;")
|
||||
@@ -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)
|
||||
@@ -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
|
||||
)
|
||||
@@ -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)
|
||||
@@ -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/<device_id>", view_func=get_device, methods=["GET"],
|
||||
endpoint="get_device",
|
||||
)
|
||||
bp.add_url_rule(
|
||||
"/api/devices/<device_id>", view_func=update_device, methods=["PATCH"],
|
||||
endpoint="update_device",
|
||||
)
|
||||
bp.add_url_rule(
|
||||
"/api/devices/<device_id>", view_func=delete_device, methods=["DELETE"],
|
||||
endpoint="delete_device",
|
||||
)
|
||||
bp.add_url_rule(
|
||||
"/api/devices/<device_id>/auto-approve",
|
||||
view_func=add_auto_approve_pattern, methods=["POST"],
|
||||
endpoint="add_auto_approve_pattern",
|
||||
)
|
||||
bp.add_url_rule(
|
||||
"/api/devices/<device_id>/auto-approve",
|
||||
view_func=list_auto_approve_patterns, methods=["GET"],
|
||||
endpoint="list_auto_approve_patterns",
|
||||
)
|
||||
bp.add_url_rule(
|
||||
"/api/devices/<device_id>/auto-approve",
|
||||
view_func=delete_auto_approve_pattern, methods=["DELETE"],
|
||||
endpoint="delete_auto_approve_pattern",
|
||||
)
|
||||
bp.add_url_rule(
|
||||
"/api/devices/<device_id>/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/<device_code>", view_func=get_pairing, methods=["GET"],
|
||||
endpoint="get_pairing",
|
||||
)
|
||||
bp.add_url_rule(
|
||||
"/api/devices/pairings/<device_code>", 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/<session_id>/events", view_func=session_events,
|
||||
methods=["GET"], endpoint="device_session_events",
|
||||
)
|
||||
bp.add_url_rule(
|
||||
"/api/devices/sessions/<session_id>/invocations/<invocation_id>/ack",
|
||||
view_func=ack_invocation, methods=["POST"], endpoint="device_ack",
|
||||
)
|
||||
bp.add_url_rule(
|
||||
"/api/devices/sessions/<session_id>/invocations/<invocation_id>/output",
|
||||
view_func=submit_output, methods=["POST"], endpoint="device_output",
|
||||
)
|
||||
@@ -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"
|
||||
)
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""Remote device backend: broker, safety checks, and pairing state."""
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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)]
|
||||
@@ -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.
|
||||
|
||||
@@ -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"),
|
||||
)
|
||||
@@ -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()]
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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),
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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.
|
||||
|
||||
<Callout type="info">
|
||||
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.
|
||||
</Callout>
|
||||
|
||||
## 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.
|
||||
|
||||
<Callout type="warning">
|
||||
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.
|
||||
</Callout>
|
||||
|
||||
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.
|
||||
@@ -0,0 +1,7 @@
|
||||
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" width="64" height="64" color="#e3e3e3" fill="none">
|
||||
<rect x="2.5" y="4" width="19" height="13" rx="2" stroke="currentColor" stroke-width="1.5"/>
|
||||
<path d="M8 21h8" stroke="currentColor" stroke-width="1.5" stroke-linecap="round"/>
|
||||
<path d="M12 17v4" stroke="currentColor" stroke-width="1.5" stroke-linecap="round"/>
|
||||
<path d="M6.5 9l2.5 2-2.5 2" stroke="currentColor" stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round"/>
|
||||
<path d="M11 13h5" stroke="currentColor" stroke-width="1.5" stroke-linecap="round"/>
|
||||
</svg>
|
||||
|
After Width: | Height: | Size: 605 B |
@@ -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 = (
|
||||
<span
|
||||
className={`mt-0.5 inline-flex items-center rounded-full px-2 py-0.5 text-[10px] font-medium ${
|
||||
online
|
||||
? 'bg-green-100 text-green-900 dark:bg-green-900/30 dark:text-green-300'
|
||||
: 'bg-gray-200 text-gray-700 dark:bg-gray-700/40 dark:text-gray-300'
|
||||
}`}
|
||||
>
|
||||
{online
|
||||
? t('settings.devices.online')
|
||||
: t('settings.devices.offline')}
|
||||
</span>
|
||||
);
|
||||
}
|
||||
return base;
|
||||
},
|
||||
);
|
||||
const groupOrder = [
|
||||
t('agents.form.toolsPopup.groupBuiltin'),
|
||||
|
||||
@@ -110,12 +110,14 @@ const createClient = (transport: FetchLike) => {
|
||||
delete: (
|
||||
url: string,
|
||||
token: string | null,
|
||||
data?: any,
|
||||
headers = {},
|
||||
signal?: AbortSignal,
|
||||
): Promise<any> =>
|
||||
request(url, {
|
||||
method: 'DELETE',
|
||||
headers: getHeaders(token, headers),
|
||||
body: data !== undefined ? JSON.stringify(data) : undefined,
|
||||
signal,
|
||||
}),
|
||||
};
|
||||
|
||||
@@ -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',
|
||||
|
||||
@@ -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<Device> => {
|
||||
const r = await apiClient.get(endpoints.USER.DEVICE(id), token);
|
||||
return (await json(r)) as Device;
|
||||
},
|
||||
|
||||
update: async (
|
||||
id: string,
|
||||
payload: Partial<Pick<Device, 'name' | 'description' | 'approval_mode'>>,
|
||||
token: string | null,
|
||||
): Promise<Device> => {
|
||||
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<PairingResponse> => {
|
||||
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<PairingStatus> => {
|
||||
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;
|
||||
@@ -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}
|
||||
</p>
|
||||
{item.description && (
|
||||
{item.descriptionNode ? (
|
||||
<div className="overflow-hidden">{item.descriptionNode}</div>
|
||||
) : item.description ? (
|
||||
<p className="text-muted-foreground overflow-hidden text-xs text-ellipsis whitespace-nowrap">
|
||||
{item.description}
|
||||
</p>
|
||||
)}
|
||||
) : null}
|
||||
</div>
|
||||
</div>
|
||||
<div
|
||||
|
||||
@@ -43,6 +43,7 @@ import { useDarkTheme, useOutsideAlerter } from '../hooks';
|
||||
import {
|
||||
selectChunks,
|
||||
selectSelectedDocs,
|
||||
selectToken,
|
||||
} from '../preferences/preferenceSlice';
|
||||
import classes from './ConversationBubble.module.css';
|
||||
import { FEEDBACK, MESSAGE_TYPE, ResearchState } from './conversationModels';
|
||||
@@ -927,6 +928,7 @@ function ToolCallApprovalBar({
|
||||
}) {
|
||||
const [expanded, setExpanded] = useState(false);
|
||||
const [comment, setComment] = useState('');
|
||||
const token = useSelector(selectToken);
|
||||
const actionLabel = toolCall.action_name.substring(
|
||||
0,
|
||||
toolCall.action_name.lastIndexOf('_'),
|
||||
@@ -935,6 +937,29 @@ function ToolCallApprovalBar({
|
||||
const truncated =
|
||||
argPreview.length > 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 (
|
||||
<div className="border-border bg-muted dark:bg-card mb-2 w-full overflow-hidden rounded-2xl border">
|
||||
<div className="flex items-center gap-3 px-4 py-2.5">
|
||||
@@ -964,6 +989,19 @@ function ToolCallApprovalBar({
|
||||
>
|
||||
Approve
|
||||
</Button>
|
||||
{isRemoteDevice && (
|
||||
<Button
|
||||
type="button"
|
||||
variant="outline"
|
||||
className="h-auto rounded-full border bg-transparent px-4 py-1 text-xs font-medium shadow-none"
|
||||
onClick={() => {
|
||||
void handleApproveSticky();
|
||||
}}
|
||||
title="Approve and don't ask again for this command shape"
|
||||
>
|
||||
Approve, don't ask again
|
||||
</Button>
|
||||
)}
|
||||
<Button
|
||||
type="button"
|
||||
variant="outline"
|
||||
|
||||
@@ -13,4 +13,7 @@ export type ToolCallsType = {
|
||||
| 'denied'
|
||||
| 'requires_client_execution';
|
||||
artifact_id?: string;
|
||||
// Remote-device tool calls carry the device id so the approval UI can
|
||||
// offer a "don't ask again" sticky-pattern action without a lookup.
|
||||
device_id?: string;
|
||||
};
|
||||
@@ -260,6 +260,60 @@
|
||||
}
|
||||
}
|
||||
},
|
||||
"devices": {
|
||||
"label": "Devices",
|
||||
"subtitle": "Pair remote machines that agents can drive via shell commands.",
|
||||
"pair": "Pair a device",
|
||||
"empty": "No paired devices.",
|
||||
"revoke": "Revoke",
|
||||
"edit": "Edit",
|
||||
"revokeWarning": "Revoke {{name}}? The CLI will be disconnected.",
|
||||
"approvalMode": "Approval mode",
|
||||
"approvalAsk": "Ask",
|
||||
"approvalFull": "Full access",
|
||||
"online": "Online",
|
||||
"offline": "Offline",
|
||||
"lastSeen": "Last seen: {{time}}",
|
||||
"lastSeenLabel": "Last seen",
|
||||
"seenAgo": "seen {{time}} ago",
|
||||
"nameLabel": "Device name",
|
||||
"descriptionLabel": "Description (visible to the agent)",
|
||||
"descriptionPlaceholder": "What is this machine for?",
|
||||
"identity": "Identity",
|
||||
"access": "Access",
|
||||
"connection": "Connection",
|
||||
"hostLabel": "Host",
|
||||
"dangerZone": "Danger zone",
|
||||
"approvalAskDescription": "Approve every command before it runs.",
|
||||
"approvalFullDescription": "Runs every command without asking, except the safety denylist.",
|
||||
"fullAccessWarning": "Full access: agents can run any command on this machine without asking. The safety denylist still applies.",
|
||||
"deviceIdLabel": "Device ID",
|
||||
"offlineHint": "Not connected. Run `docsgpt-cli host` on the machine to bring it online.",
|
||||
"dangerZoneDescription": "Revoking stops the daemon on the device and removes the pairing.",
|
||||
"auditTitle": "Recent activity",
|
||||
"auditLoading": "Loading…",
|
||||
"auditEmpty": "No activity yet.",
|
||||
"auditDecision": "Decision",
|
||||
"auditExit": "Exit",
|
||||
"auditDuration": "Duration",
|
||||
"pairing": {
|
||||
"title": "Pair a remote machine",
|
||||
"subtitle": "Run the command below on the remote machine and enter the code when prompted.",
|
||||
"formSubtitle": "Set a name, description, and approval mode for the new device.",
|
||||
"nameLabel": "Device name (optional)",
|
||||
"namePlaceholder": "e.g. edge-server",
|
||||
"descriptionLabel": "Description (optional)",
|
||||
"descriptionPlaceholder": "Visible to the agent so it can pick the right device.",
|
||||
"fullAccessWarning": "Full access: agents can run any command on this machine without asking. The safety denylist still applies.",
|
||||
"stepOne": "1. Run this on your remote machine:",
|
||||
"installLink": "Don't have docsgpt-cli? Install it",
|
||||
"stepTwo": "2. Enter this code when prompted:",
|
||||
"start": "Generate code",
|
||||
"starting": "Working…",
|
||||
"waitingForCli": "Waiting for the CLI to redeem the code…",
|
||||
"cancel": "Cancel"
|
||||
}
|
||||
},
|
||||
"customModels": {
|
||||
"label": "Custom Models",
|
||||
"subtitle": "Bring your own OpenAI-compatible model.",
|
||||
|
||||
@@ -6,6 +6,7 @@ import userService from '../api/services/userService';
|
||||
import SkeletonLoader from '../components/SkeletonLoader';
|
||||
import { Modal } from '../components/ui/modal';
|
||||
import { useLoaderState } from '../hooks';
|
||||
import PairDeviceModal from '../settings/PairDeviceModal';
|
||||
import { ActiveState } from '../models/misc';
|
||||
import { selectToken } from '../preferences/preferenceSlice';
|
||||
import ConfigToolModal from './ConfigToolModal';
|
||||
@@ -18,12 +19,14 @@ export default function AddToolModal({
|
||||
setModalState,
|
||||
getUserTools,
|
||||
onToolAdded,
|
||||
onDevicePaired,
|
||||
}: {
|
||||
message: string;
|
||||
modalState: ActiveState;
|
||||
setModalState: (state: ActiveState) => void;
|
||||
getUserTools: () => void;
|
||||
onToolAdded: (toolId: string) => void;
|
||||
onDevicePaired?: (deviceId: string) => void;
|
||||
}) {
|
||||
const { t } = useTranslation();
|
||||
const token = useSelector(selectToken);
|
||||
@@ -36,6 +39,8 @@ export default function AddToolModal({
|
||||
React.useState<ActiveState>('INACTIVE');
|
||||
const [mcpModalState, setMcpModalState] =
|
||||
React.useState<ActiveState>('INACTIVE');
|
||||
const [pairModalState, setPairModalState] =
|
||||
React.useState<ActiveState>('INACTIVE');
|
||||
const [loading, setLoading] = useLoaderState(false);
|
||||
|
||||
const getAvailableTools = () => {
|
||||
@@ -52,6 +57,13 @@ export default function AddToolModal({
|
||||
};
|
||||
|
||||
const handleAddTool = (tool: AvailableToolType) => {
|
||||
// ``remote_device`` is created server-side via the pairing redeem
|
||||
// endpoint, not the standard create_tool path.
|
||||
if (tool.name === 'remote_device') {
|
||||
setModalState('INACTIVE');
|
||||
setPairModalState('ACTIVE');
|
||||
return;
|
||||
}
|
||||
if (Object.keys(tool.configRequirements).length === 0) {
|
||||
userService
|
||||
.createTool(
|
||||
@@ -174,6 +186,15 @@ export default function AddToolModal({
|
||||
setModalState={setMcpModalState}
|
||||
onServerSaved={handleMcpServerAdded}
|
||||
/>
|
||||
<PairDeviceModal
|
||||
modalState={pairModalState}
|
||||
setModalState={setPairModalState}
|
||||
onPaired={(deviceId) => {
|
||||
setPairModalState('INACTIVE');
|
||||
setModalState('INACTIVE');
|
||||
onDevicePaired?.(deviceId);
|
||||
}}
|
||||
/>
|
||||
</>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,274 @@
|
||||
import React from 'react';
|
||||
import { useTranslation } from 'react-i18next';
|
||||
import { useSelector } from 'react-redux';
|
||||
|
||||
import { baseURL } from '../api/client';
|
||||
import devicesService, {
|
||||
ApprovalMode,
|
||||
PairingResponse,
|
||||
} from '../api/services/devicesService';
|
||||
import CopyButton from '../components/CopyButton';
|
||||
import { Alert, AlertDescription } from '../components/ui/alert';
|
||||
import { Button } from '../components/ui/button';
|
||||
import { Input } from '../components/ui/input';
|
||||
import { Modal } from '../components/ui/modal';
|
||||
import {
|
||||
Select,
|
||||
SelectContent,
|
||||
SelectItem,
|
||||
SelectTrigger,
|
||||
SelectValue,
|
||||
} from '../components/ui/select';
|
||||
import { ActiveState } from '../models/misc';
|
||||
import { selectToken } from '../preferences/preferenceSlice';
|
||||
|
||||
interface Props {
|
||||
modalState: ActiveState;
|
||||
setModalState: (state: ActiveState) => void;
|
||||
onPaired?: (deviceId: string, deviceName: string) => void;
|
||||
}
|
||||
|
||||
type Stage = 'form' | 'waiting';
|
||||
|
||||
export default function PairDeviceModal({
|
||||
modalState,
|
||||
setModalState,
|
||||
onPaired,
|
||||
}: Props) {
|
||||
const { t } = useTranslation();
|
||||
const token = useSelector(selectToken);
|
||||
|
||||
const [stage, setStage] = React.useState<Stage>('form');
|
||||
const [name, setName] = React.useState('');
|
||||
const [description, setDescription] = React.useState('');
|
||||
const [approvalMode, setApprovalMode] = React.useState<ApprovalMode>('ask');
|
||||
const [submitting, setSubmitting] = React.useState(false);
|
||||
|
||||
const [pairing, setPairing] = React.useState<PairingResponse | null>(null);
|
||||
const [error, setError] = React.useState<string | null>(null);
|
||||
const pollTimerRef = React.useRef<ReturnType<typeof setInterval> | null>(
|
||||
null,
|
||||
);
|
||||
|
||||
const reset = React.useCallback(() => {
|
||||
setStage('form');
|
||||
setName('');
|
||||
setDescription('');
|
||||
setApprovalMode('ask');
|
||||
setSubmitting(false);
|
||||
setPairing(null);
|
||||
setError(null);
|
||||
if (pollTimerRef.current) {
|
||||
clearInterval(pollTimerRef.current);
|
||||
pollTimerRef.current = null;
|
||||
}
|
||||
}, []);
|
||||
|
||||
React.useEffect(() => {
|
||||
if (modalState === 'INACTIVE') {
|
||||
reset();
|
||||
}
|
||||
}, [modalState, reset]);
|
||||
|
||||
const installCommand = React.useMemo(() => {
|
||||
const url = baseURL || '';
|
||||
return `docsgpt-cli host pair --url ${url}`;
|
||||
}, []);
|
||||
|
||||
const handleStart = async () => {
|
||||
setError(null);
|
||||
setSubmitting(true);
|
||||
try {
|
||||
const trimmedName = name.trim();
|
||||
const trimmedDescription = description.trim();
|
||||
const res = await devicesService.startPairing(token, {
|
||||
name: trimmedName || undefined,
|
||||
description: trimmedDescription || undefined,
|
||||
approval_mode: approvalMode,
|
||||
});
|
||||
setPairing(res);
|
||||
setStage('waiting');
|
||||
} catch (err: unknown) {
|
||||
setError(err instanceof Error ? err.message : 'pairing failed');
|
||||
} finally {
|
||||
setSubmitting(false);
|
||||
}
|
||||
};
|
||||
|
||||
// Poll while waiting for the CLI to redeem.
|
||||
React.useEffect(() => {
|
||||
if (stage !== 'waiting' || !pairing) return;
|
||||
pollTimerRef.current = setInterval(async () => {
|
||||
try {
|
||||
const s = await devicesService.pollPairing(pairing.device_code, token);
|
||||
if (s.status === 'redeemed') {
|
||||
if (pollTimerRef.current) {
|
||||
clearInterval(pollTimerRef.current);
|
||||
pollTimerRef.current = null;
|
||||
}
|
||||
if (s.device_id && s.device_name) {
|
||||
onPaired?.(s.device_id, s.device_name);
|
||||
}
|
||||
setModalState('INACTIVE');
|
||||
}
|
||||
} catch (err) {
|
||||
console.error('pairing poll failed', err);
|
||||
}
|
||||
}, 3000);
|
||||
return () => {
|
||||
if (pollTimerRef.current) {
|
||||
clearInterval(pollTimerRef.current);
|
||||
pollTimerRef.current = null;
|
||||
}
|
||||
};
|
||||
}, [stage, pairing, token, onPaired, setModalState]);
|
||||
|
||||
const close = () => {
|
||||
if (pairing && stage === 'waiting') {
|
||||
void devicesService
|
||||
.cancelPairing(pairing.device_code, token)
|
||||
.catch(() => {});
|
||||
}
|
||||
setModalState('INACTIVE');
|
||||
};
|
||||
|
||||
const userCode = pairing?.user_code || '';
|
||||
|
||||
const renderForm = () => (
|
||||
<div className="flex flex-col gap-4">
|
||||
<p className="text-muted-foreground text-sm">
|
||||
{t('settings.devices.pairing.formSubtitle')}
|
||||
</p>
|
||||
<Input
|
||||
type="text"
|
||||
value={name}
|
||||
onChange={(e) => setName(e.target.value)}
|
||||
label={t('settings.devices.pairing.nameLabel')}
|
||||
placeholder={t('settings.devices.pairing.namePlaceholder')}
|
||||
labelBgClassName="bg-card"
|
||||
/>
|
||||
<Input
|
||||
type="text"
|
||||
value={description}
|
||||
onChange={(e) => setDescription(e.target.value)}
|
||||
label={t('settings.devices.pairing.descriptionLabel')}
|
||||
placeholder={t('settings.devices.pairing.descriptionPlaceholder')}
|
||||
labelBgClassName="bg-card"
|
||||
/>
|
||||
<div className="flex flex-col gap-2">
|
||||
<span className="text-muted-foreground text-xs">
|
||||
{t('settings.devices.approvalMode')}
|
||||
</span>
|
||||
<Select
|
||||
value={approvalMode}
|
||||
onValueChange={(value) => setApprovalMode(value as ApprovalMode)}
|
||||
>
|
||||
<SelectTrigger className="w-full">
|
||||
<SelectValue />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
<SelectItem value="ask">
|
||||
{t('settings.devices.approvalAsk')}
|
||||
</SelectItem>
|
||||
<SelectItem value="full">
|
||||
{t('settings.devices.approvalFull')}
|
||||
</SelectItem>
|
||||
</SelectContent>
|
||||
</Select>
|
||||
{approvalMode === 'full' && (
|
||||
<p className="text-xs text-red-600 dark:text-red-400">
|
||||
{t('settings.devices.pairing.fullAccessWarning')}
|
||||
</p>
|
||||
)}
|
||||
</div>
|
||||
{error && (
|
||||
<Alert variant="destructive">
|
||||
<AlertDescription>{error}</AlertDescription>
|
||||
</Alert>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
|
||||
const renderWaiting = () => (
|
||||
<div className="flex flex-col gap-4">
|
||||
<p className="text-muted-foreground text-sm">
|
||||
{t('settings.devices.pairing.subtitle')}
|
||||
</p>
|
||||
<div className="flex flex-col gap-1">
|
||||
<span className="text-muted-foreground text-xs">
|
||||
{t('settings.devices.pairing.stepOne')}
|
||||
</span>
|
||||
<div className="bg-muted flex items-center gap-2 rounded-md p-3">
|
||||
<pre className="grow font-mono text-sm break-all whitespace-pre-wrap select-all">
|
||||
{installCommand}
|
||||
</pre>
|
||||
<CopyButton textToCopy={installCommand} />
|
||||
</div>
|
||||
<a
|
||||
href="https://github.com/arc53/DocsGPT-cli#installation"
|
||||
target="_blank"
|
||||
rel="noreferrer"
|
||||
className="text-muted-foreground hover:text-foreground text-xs underline"
|
||||
>
|
||||
{t('settings.devices.pairing.installLink')}
|
||||
</a>
|
||||
</div>
|
||||
<div className="flex flex-col gap-1">
|
||||
<span className="text-muted-foreground text-xs">
|
||||
{t('settings.devices.pairing.stepTwo')}
|
||||
</span>
|
||||
<div className="bg-muted flex items-center gap-2 rounded-md p-3">
|
||||
<div className="grow text-center font-mono text-3xl tracking-widest select-all">
|
||||
{userCode || '...'}
|
||||
</div>
|
||||
{userCode && <CopyButton textToCopy={userCode} />}
|
||||
</div>
|
||||
</div>
|
||||
<p className="text-muted-foreground text-xs">
|
||||
{t('settings.devices.pairing.waitingForCli')}
|
||||
</p>
|
||||
</div>
|
||||
);
|
||||
|
||||
const footer = (() => {
|
||||
if (stage === 'form') {
|
||||
return (
|
||||
<>
|
||||
<Button variant="ghost" onClick={close} className="rounded-3xl px-5">
|
||||
{t('settings.devices.pairing.cancel')}
|
||||
</Button>
|
||||
<Button
|
||||
type="button"
|
||||
onClick={handleStart}
|
||||
disabled={submitting}
|
||||
className="rounded-3xl px-5 text-white"
|
||||
>
|
||||
{submitting
|
||||
? t('settings.devices.pairing.starting')
|
||||
: t('settings.devices.pairing.start')}
|
||||
</Button>
|
||||
</>
|
||||
);
|
||||
}
|
||||
return (
|
||||
<Button variant="ghost" onClick={close} className="rounded-3xl px-5">
|
||||
{t('settings.devices.pairing.cancel')}
|
||||
</Button>
|
||||
);
|
||||
})();
|
||||
|
||||
return (
|
||||
<Modal
|
||||
open={modalState === 'ACTIVE'}
|
||||
onOpenChange={(open) => {
|
||||
if (!open) close();
|
||||
}}
|
||||
title={t('settings.devices.pairing.title')}
|
||||
footer={footer}
|
||||
contentClassName="px-1"
|
||||
>
|
||||
{stage === 'form' && renderForm()}
|
||||
{stage === 'waiting' && renderWaiting()}
|
||||
</Modal>
|
||||
);
|
||||
}
|
||||
@@ -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<Device | null>(null);
|
||||
const [loading, setLoading] = React.useState(false);
|
||||
const [name, setName] = React.useState<string>('');
|
||||
const [description, setDescription] = React.useState<string>('');
|
||||
const [approvalMode, setApprovalMode] = React.useState<ApprovalMode>('ask');
|
||||
const [saving, setSaving] = React.useState(false);
|
||||
const [error, setError] = React.useState<string | null>(null);
|
||||
|
||||
const [audit, setAudit] = React.useState<AuditEntry[] | null>(null);
|
||||
const [auditLoading, setAuditLoading] = React.useState(false);
|
||||
|
||||
const [revokeState, setRevokeState] = React.useState<ActiveState>('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<Device, 'name' | 'description' | 'approval_mode'>
|
||||
> = {};
|
||||
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 (
|
||||
<div className="mt-8 flex items-center justify-center py-16">
|
||||
<Spinner size="large" />
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="scrollbar-overlay mt-8 flex flex-col gap-6">
|
||||
<div className="mb-4 flex items-center justify-between">
|
||||
<div className="text-foreground dark:text-foreground flex items-center gap-3 text-sm">
|
||||
<Button
|
||||
type="button"
|
||||
variant="outline"
|
||||
size="icon"
|
||||
className="text-muted-foreground rounded-full p-3"
|
||||
onClick={handleGoBack}
|
||||
>
|
||||
<img src={ArrowLeft} alt="left-arrow" className="h-3 w-3" />
|
||||
</Button>
|
||||
<p className="mt-px">{t('settings.tools.backToAllTools')}</p>
|
||||
</div>
|
||||
<Button
|
||||
type="button"
|
||||
className="rounded-full px-3 py-2 text-xs text-nowrap text-white sm:px-4 sm:py-2"
|
||||
onClick={handleSaveChanges}
|
||||
disabled={!hasUnsavedChanges || saving || loading}
|
||||
>
|
||||
{saving ? t('settings.tools.saving') : t('settings.tools.save')}
|
||||
</Button>
|
||||
</div>
|
||||
|
||||
{error && (
|
||||
<Alert variant="destructive">
|
||||
<AlertDescription>{error}</AlertDescription>
|
||||
</Alert>
|
||||
)}
|
||||
|
||||
<div className="flex flex-wrap items-center gap-3">
|
||||
<img
|
||||
src={`/toolIcons/tool_${tool.name}.svg`}
|
||||
alt={`${tool.displayName} icon`}
|
||||
className="h-7 w-7"
|
||||
/>
|
||||
<h2 className="text-foreground dark:text-foreground text-xl font-semibold">
|
||||
{device?.name || tool.customName || tool.displayName || 'device'}
|
||||
</h2>
|
||||
<span
|
||||
className={`rounded-full px-3 py-0.5 text-xs font-medium ${
|
||||
online
|
||||
? 'bg-green-100 text-green-900 dark:bg-green-900/30 dark:text-green-300'
|
||||
: 'bg-gray-200 text-gray-700 dark:bg-gray-700/40 dark:text-gray-300'
|
||||
}`}
|
||||
>
|
||||
{pillText}
|
||||
</span>
|
||||
{approvalMode === 'full' && (
|
||||
<span className="rounded-full bg-red-200 px-3 py-0.5 text-xs font-medium text-red-900 dark:bg-red-900/40 dark:text-red-300">
|
||||
{t('settings.devices.approvalFull')}
|
||||
</span>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{/* Identity */}
|
||||
<section className="flex flex-col gap-4">
|
||||
<h3 className="text-foreground dark:text-foreground text-sm font-semibold">
|
||||
{t('settings.devices.identity')}
|
||||
</h3>
|
||||
<div className="flex flex-col gap-1 sm:flex-row sm:items-start sm:gap-4">
|
||||
<label className="text-muted-foreground w-32 shrink-0 pt-2 text-sm">
|
||||
{t('settings.devices.nameLabel')}
|
||||
</label>
|
||||
<div className="w-full max-w-[340px]">
|
||||
<Input
|
||||
type="text"
|
||||
value={name}
|
||||
onChange={(e) => setName(e.target.value)}
|
||||
placeholder={t('settings.devices.nameLabel')}
|
||||
className="rounded-xl"
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
<div className="flex flex-col gap-1 sm:flex-row sm:items-start sm:gap-4">
|
||||
<label className="text-muted-foreground w-32 shrink-0 pt-2 text-sm">
|
||||
{t('settings.devices.descriptionLabel')}
|
||||
</label>
|
||||
<div className="w-full max-w-[340px]">
|
||||
<Input
|
||||
type="text"
|
||||
value={description}
|
||||
onChange={(e) => setDescription(e.target.value)}
|
||||
placeholder={t('settings.devices.descriptionPlaceholder')}
|
||||
className="rounded-xl"
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
</section>
|
||||
|
||||
{/* Access */}
|
||||
<section className="flex flex-col gap-4">
|
||||
<h3 className="text-foreground dark:text-foreground text-sm font-semibold">
|
||||
{t('settings.devices.access')}
|
||||
</h3>
|
||||
<div className="flex flex-col gap-1 sm:flex-row sm:items-start sm:gap-4">
|
||||
<label className="text-muted-foreground w-32 shrink-0 pt-2 text-sm">
|
||||
{t('settings.devices.approvalMode')}
|
||||
</label>
|
||||
<div className="flex w-full max-w-[340px] flex-col gap-2">
|
||||
<Select
|
||||
value={approvalMode}
|
||||
onValueChange={(value) => setApprovalMode(value as ApprovalMode)}
|
||||
>
|
||||
<SelectTrigger className="w-full rounded-xl px-4 py-2" size="lg">
|
||||
<SelectValue />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
<SelectItem value="ask">
|
||||
{t('settings.devices.approvalAsk')}
|
||||
</SelectItem>
|
||||
<SelectItem value="full">
|
||||
{t('settings.devices.approvalFull')}
|
||||
</SelectItem>
|
||||
</SelectContent>
|
||||
</Select>
|
||||
<p className="text-muted-foreground text-xs">
|
||||
{approvalMode === 'full'
|
||||
? t('settings.devices.approvalFullDescription')
|
||||
: t('settings.devices.approvalAskDescription')}
|
||||
</p>
|
||||
{approvalMode === 'full' && (
|
||||
<div className="rounded-lg bg-red-50 px-4 py-2 text-sm text-red-600 dark:bg-red-900/20 dark:text-red-400">
|
||||
{t('settings.devices.fullAccessWarning')}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
</section>
|
||||
|
||||
{/* Connection */}
|
||||
<section className="flex flex-col gap-4">
|
||||
<h3 className="text-foreground dark:text-foreground text-sm font-semibold">
|
||||
{t('settings.devices.connection')}
|
||||
</h3>
|
||||
<div className="flex flex-col gap-1 sm:flex-row sm:items-start sm:gap-4">
|
||||
<span className="text-muted-foreground w-32 shrink-0 pt-px text-sm">
|
||||
{t('settings.devices.hostLabel')}
|
||||
</span>
|
||||
<div className="text-muted-foreground w-full max-w-md text-sm">
|
||||
<span className="font-mono">
|
||||
{device?.hostname || 'unknown host'}
|
||||
</span>
|
||||
{device?.os ? ` · ${device.os}` : ''}
|
||||
{device?.arch ? ` · ${device.arch}` : ''}
|
||||
{device?.cli_version ? ` · cli ${device.cli_version}` : ''}
|
||||
</div>
|
||||
</div>
|
||||
<div className="flex flex-col gap-1 sm:flex-row sm:items-start sm:gap-4">
|
||||
<span className="text-muted-foreground w-32 shrink-0 pt-1 text-sm">
|
||||
{t('settings.devices.deviceIdLabel')}
|
||||
</span>
|
||||
<div className="flex w-full max-w-md flex-wrap items-center gap-2">
|
||||
<code className="text-foreground dark:text-foreground bg-muted max-w-full truncate rounded px-2 py-0.5 font-mono text-xs">
|
||||
{deviceId || '-'}
|
||||
</code>
|
||||
{deviceId && <CopyButton textToCopy={deviceId} />}
|
||||
</div>
|
||||
</div>
|
||||
<div className="flex flex-col gap-1 sm:flex-row sm:items-start sm:gap-4">
|
||||
<span className="text-muted-foreground w-32 shrink-0 pt-px text-sm">
|
||||
{t('settings.devices.lastSeenLabel')}
|
||||
</span>
|
||||
<div className="text-muted-foreground w-full max-w-md text-sm">
|
||||
{formatTimestamp(device?.last_seen_at)}
|
||||
</div>
|
||||
</div>
|
||||
{!online && (
|
||||
<p className="text-muted-foreground text-xs">
|
||||
{t('settings.devices.offlineHint')}
|
||||
</p>
|
||||
)}
|
||||
</section>
|
||||
|
||||
{/* Recent activity */}
|
||||
<section className="flex flex-col gap-3">
|
||||
<h3 className="text-foreground dark:text-foreground text-sm font-semibold">
|
||||
{t('settings.devices.auditTitle')}
|
||||
</h3>
|
||||
<Accordion
|
||||
type="single"
|
||||
collapsible
|
||||
onValueChange={handleAuditToggle}
|
||||
className="border-border dark:border-border w-full rounded-xl border"
|
||||
>
|
||||
<AccordionItem value="audit" className="border-b-0">
|
||||
<AccordionTrigger className="px-4 py-3 text-sm font-semibold">
|
||||
{t('settings.devices.auditTitle')}
|
||||
{audit !== null ? ` (${audit.length})` : ''}
|
||||
</AccordionTrigger>
|
||||
<AccordionContent className="px-4 pb-4">
|
||||
{auditLoading ? (
|
||||
<p className="text-muted-foreground py-2 text-sm">
|
||||
{t('settings.devices.auditLoading')}
|
||||
</p>
|
||||
) : !audit || audit.length === 0 ? (
|
||||
<p className="text-muted-foreground py-2 text-sm">
|
||||
{t('settings.devices.auditEmpty')}
|
||||
</p>
|
||||
) : (
|
||||
<ul className="flex flex-col gap-2">
|
||||
{audit.map((entry) => (
|
||||
<li
|
||||
key={entry.id}
|
||||
className="bg-muted/40 flex flex-col gap-1 rounded-md px-3 py-2 text-xs"
|
||||
>
|
||||
<code className="text-foreground dark:text-foreground block font-mono break-all whitespace-pre-wrap">
|
||||
{entry.command}
|
||||
</code>
|
||||
<div className="text-muted-foreground flex flex-wrap gap-3">
|
||||
<span>
|
||||
{t('settings.devices.auditDecision')}:{' '}
|
||||
{entry.decision}
|
||||
</span>
|
||||
<span>
|
||||
{t('settings.devices.auditExit')}:{' '}
|
||||
{entry.exit_code ?? '-'}
|
||||
</span>
|
||||
<span>
|
||||
{t('settings.devices.auditDuration')}:{' '}
|
||||
{entry.duration_ms ?? '-'} ms
|
||||
</span>
|
||||
<span>{formatTimestamp(entry.created_at)}</span>
|
||||
</div>
|
||||
</li>
|
||||
))}
|
||||
</ul>
|
||||
)}
|
||||
</AccordionContent>
|
||||
</AccordionItem>
|
||||
</Accordion>
|
||||
</section>
|
||||
|
||||
{/* Danger zone */}
|
||||
<section className="flex flex-col gap-3">
|
||||
<h3 className="text-destructive text-sm font-semibold">
|
||||
{t('settings.devices.dangerZone')}
|
||||
</h3>
|
||||
<div className="border-destructive/40 flex flex-col gap-3 rounded-xl border p-4 sm:flex-row sm:items-center sm:justify-between">
|
||||
<p className="text-muted-foreground text-sm">
|
||||
{t('settings.devices.dangerZoneDescription')}
|
||||
</p>
|
||||
<Button
|
||||
type="button"
|
||||
variant="destructive"
|
||||
className="shrink-0 rounded-full px-5 text-white"
|
||||
onClick={() => setRevokeState('ACTIVE')}
|
||||
>
|
||||
{t('settings.devices.revoke')}
|
||||
</Button>
|
||||
</div>
|
||||
</section>
|
||||
|
||||
<ConfirmationModal
|
||||
message={
|
||||
device
|
||||
? t('settings.devices.revokeWarning', { name: device.name })
|
||||
: ''
|
||||
}
|
||||
modalState={revokeState}
|
||||
setModalState={setRevokeState}
|
||||
handleSubmit={handleRevoke}
|
||||
submitLabel={t('settings.devices.revoke')}
|
||||
variant="danger"
|
||||
/>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -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 (
|
||||
<div>
|
||||
{selectedTool ? (
|
||||
<ToolConfig
|
||||
tool={selectedTool}
|
||||
setTool={setSelectedTool}
|
||||
handleGoBack={handleGoBack}
|
||||
/>
|
||||
selectedTool.name === 'remote_device' ? (
|
||||
<RemoteDeviceConfig
|
||||
tool={selectedTool as UserToolType}
|
||||
handleGoBack={handleGoBack}
|
||||
/>
|
||||
) : (
|
||||
<ToolConfig
|
||||
tool={selectedTool}
|
||||
setTool={setSelectedTool}
|
||||
handleGoBack={handleGoBack}
|
||||
/>
|
||||
)
|
||||
) : (
|
||||
<div className="mt-8">
|
||||
<div className="relative flex flex-col">
|
||||
@@ -437,6 +478,7 @@ export default function Tools() {
|
||||
setModalState={setAddToolModalState}
|
||||
getUserTools={getUserTools}
|
||||
onToolAdded={handleToolAdded}
|
||||
onDevicePaired={handleDevicePaired}
|
||||
/>
|
||||
<ConfirmationModal
|
||||
message={t('settings.tools.deleteWarning', {
|
||||
|
||||
@@ -193,6 +193,10 @@ export default function Settings() {
|
||||
<Route path="analytics" element={<Analytics />} />
|
||||
<Route path="logs" element={<Logs />} />
|
||||
<Route path="tools" element={<Tools />} />
|
||||
<Route
|
||||
path="devices"
|
||||
element={<Navigate to="/settings/tools" replace />}
|
||||
/>
|
||||
<Route path="custom-models" element={<CustomModels />} />
|
||||
<Route path="*" element={<Navigate to="/settings" replace />} />
|
||||
</Routes>
|
||||
|
||||
+25
-13
@@ -1,21 +1,33 @@
|
||||
/// <reference types="vitest" />
|
||||
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'],
|
||||
},
|
||||
};
|
||||
});
|
||||
+17
@@ -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"
|
||||
}
|
||||
}
|
||||
+17
@@ -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"
|
||||
}
|
||||
}
|
||||
+17
@@ -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"
|
||||
}
|
||||
}
|
||||
+17
@@ -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"
|
||||
}
|
||||
}
|
||||
@@ -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:
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
Whitespace-only changes.
@@ -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"
|
||||
@@ -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
|
||||
@@ -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 == []
|
||||
@@ -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
|
||||
@@ -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"
|
||||
@@ -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"
|
||||
@@ -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")
|
||||
@@ -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)
|
||||
@@ -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")
|
||||
@@ -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"
|
||||
@@ -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
|
||||
@@ -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"]
|
||||
@@ -0,0 +1,109 @@
|
||||
"""Tests for ``PATCH /api/devices/<id>`` 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"
|
||||
@@ -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)
|
||||
|
||||
@@ -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"
|
||||
Reference in new issue
Block a user