49 KiB
AI Agent Command Bar + Privacy Gateway Implementation Plan
For Claude: REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.
Goal: Add AI agent command bar (bottom of ConsolePage) and three-mode privacy gateway selector (PrivacyPage) with complete frontend/backend wiring.
Architecture: Go reverse proxy routes /api/v1/agent/* → port 8890 and /api/v1/privacy/* → port 8889. Privacy API redesigned with mode enum ("off"/"audit"/"redact"). Runner supports task cancellation via asyncio.Task.cancel().
Tech Stack: Go 1.22 (net/http/httputil), Python 3.12 asyncio/aiohttp, React 18 + TypeScript + Zustand, TailwindCSS
Working directory: /data/project/KVM-privacy
Task 1: Go Reverse Proxy
Files:
- Create:
deps/KVM/go/internal/api/proxy_handler.go - Modify:
deps/KVM/go/internal/api/router.go
Step 1: Write failing test
Create deps/KVM/go/internal/api/proxy_handler_test.go:
package api
import (
"net/http"
"net/http/httptest"
"testing"
)
func TestAgentProxy_ForwardsRequest(t *testing.T) {
// Fake upstream (Python agent)
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/v1/agent/status" {
t.Errorf("unexpected path: %s", r.URL.Path)
}
w.Header().Set("Content-Type", "application/json")
w.Write([]byte(`{"running":true}`))
}))
defer upstream.Close()
proxy := NewReverseProxy(upstream.URL)
req := httptest.NewRequest("GET", "/api/v1/agent/status", nil)
rec := httptest.NewRecorder()
proxy.ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
t.Errorf("expected 200, got %d", rec.Code)
}
}
func TestPrivacyProxy_ForwardsRequest(t *testing.T) {
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/v1/privacy/mode" {
t.Errorf("unexpected path: %s", r.URL.Path)
}
w.Write([]byte(`{"mode":"off"}`))
}))
defer upstream.Close()
proxy := NewReverseProxy(upstream.URL)
req := httptest.NewRequest("GET", "/api/v1/privacy/mode", nil)
rec := httptest.NewRecorder()
proxy.ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
t.Errorf("expected 200, got %d", rec.Code)
}
}
func TestProxy_UpstreamDown_Returns502(t *testing.T) {
proxy := NewReverseProxy("http://127.0.0.1:19999") // nothing listening
req := httptest.NewRequest("GET", "/api/v1/agent/status", nil)
rec := httptest.NewRecorder()
proxy.ServeHTTP(rec, req)
if rec.Code != http.StatusBadGateway {
t.Errorf("expected 502, got %d", rec.Code)
}
}
Step 2: Run to verify it fails
cd deps/KVM && go test ./go/internal/api/... -run TestAgentProxy -v 2>&1 | head -20
Expected: compile error "NewReverseProxy undefined"
Step 3: Implement proxy_handler.go
Create deps/KVM/go/internal/api/proxy_handler.go:
package api
import (
"log"
"net/http"
"net/http/httputil"
"net/url"
)
// NewReverseProxy creates a reverse proxy to the given target base URL.
// Path is forwarded as-is (no stripping). Upstream CORS headers are removed
// since the Go CORS middleware handles them.
func NewReverseProxy(targetBase string) http.HandlerFunc {
target, err := url.Parse(targetBase)
if err != nil {
log.Fatalf("invalid proxy target %q: %v", targetBase, err)
}
proxy := httputil.NewSingleHostReverseProxy(target)
// Strip upstream CORS headers — Go middleware owns them
proxy.ModifyResponse = func(resp *http.Response) error {
resp.Header.Del("Access-Control-Allow-Origin")
resp.Header.Del("Access-Control-Allow-Methods")
resp.Header.Del("Access-Control-Allow-Headers")
return nil
}
proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
log.Printf("proxy error for %s: %v", r.URL.Path, err)
http.Error(w, "Bad Gateway", http.StatusBadGateway)
}
return proxy.ServeHTTP
}
Step 4: Run tests to verify they pass
cd deps/KVM && go test ./go/internal/api/... -run "TestAgentProxy|TestPrivacyProxy|TestProxy_Upstream" -v
Expected: all 3 PASS
Step 5: Register routes in router.go
In deps/KVM/go/internal/api/router.go, add before the // Terminal routes comment:
// Privacy Gateway proxy (→ port 8889)
mux.Handle("/api/v1/privacy/", withJWT(deps, NewReverseProxy("http://127.0.0.1:8889")))
// KVM Agent proxy (→ port 8890)
mux.Handle("/api/v1/agent/", withJWT(deps, NewReverseProxy("http://127.0.0.1:8890")))
Step 6: Build to verify no compile errors
cd deps/KVM && go build ./...
Expected: exits 0, no errors
Step 7: Commit
cd deps/KVM && git add go/internal/api/proxy_handler.go go/internal/api/proxy_handler_test.go go/internal/api/router.go
git commit -m "feat: add reverse proxy routes for agent (8890) and privacy gateway (8889)"
Task 2: Privacy API Redesign (mode enum + stats + audit format)
Files:
- Modify:
services/privacy_gateway/privacy_api.py - Modify:
services/privacy_gateway/audit_logger.py
Context: Current API uses {enabled: bool}. Redesign to {mode: "off"|"audit"|"redact"}. stats_today() already returns rich data — just need correct response mapping. Audit endpoint returns entries but frontend expects logs.
Step 1: Write failing tests
Create services/privacy_gateway/tests/test_privacy_api_redesign.py:
"""Tests for redesigned Privacy API with mode enum."""
import json
import pytest
import pytest_asyncio
from unittest.mock import MagicMock, patch
from pathlib import Path
# Import after patching state path
@pytest.fixture
def mock_state_path(tmp_path):
p = tmp_path / "state.json"
return p
@pytest.fixture
def api_server(mock_state_path):
with patch("services.privacy_gateway.privacy_api._STATE_PATH", mock_state_path):
from services.privacy_gateway.privacy_api import PrivacyAPIServer
server = PrivacyAPIServer.__new__(PrivacyAPIServer)
server._state = {"mode": "off"}
server._audit = MagicMock()
server._state_path = mock_state_path
return server
@pytest.mark.asyncio
async def test_get_mode_returns_enum(api_server):
req = MagicMock()
resp = await api_server._get_mode(req)
data = json.loads(resp.body)
assert data["mode"] in ("off", "audit", "redact")
assert "enabled" not in data # old field removed
@pytest.mark.asyncio
async def test_set_mode_accepts_enum(api_server, mock_state_path):
req = MagicMock()
req.json = pytest.AsyncMock(return_value={"mode": "redact"})
resp = await api_server._set_mode(req)
data = json.loads(resp.body)
assert data["mode"] == "redact"
assert api_server._state["mode"] == "redact"
@pytest.mark.asyncio
async def test_set_mode_rejects_invalid(api_server):
req = MagicMock()
req.json = pytest.AsyncMock(return_value={"mode": "invalid"})
resp = await api_server._set_mode(req)
assert resp.status == 400
@pytest.mark.asyncio
async def test_get_stats_returns_rich_format(api_server):
api_server._audit.stats_today.return_value = {
"requests": 10, "files": 5, "pii_total": 3,
"by_action": {"auto_redact": 2, "bypass": 8},
"by_domain": {"api.openai.com": 5, "claude.ai": 3},
"pii_by_type": {"id_card": 2, "phone": 1},
}
req = MagicMock()
resp = await api_server._get_stats(req)
data = json.loads(resp.body)
assert data["requests"] == 10
assert data["pii_by_type"]["id_card"] == 2
assert data["actions"]["redact"] == 2
assert data["actions"]["allow"] == 8
assert data["top_domains"][0]["domain"] == "api.openai.com"
@pytest.mark.asyncio
async def test_get_audit_returns_logs_field(api_server):
api_server._audit.query.return_value = [
{"id": 1, "ts": "2026-03-03T10:00:00", "domain": "openai.com",
"pii_types": '{"id_card":1}', "action": "redact",
"filename": "doc.pdf", "file_size": 1024, "client_ip": "127.0.0.1",
"request_url": "/v1/files"}
]
api_server._audit.count.return_value = 1
req = MagicMock()
req.query = {}
resp = await api_server._get_audit(req)
data = json.loads(resp.body)
assert "logs" in data # not "entries"
assert data["total"] == 1
assert data["logs"][0]["pii_types"] == {"id_card": 1} # dict, not list
Step 2: Run to verify it fails
cd services && python -m pytest privacy_gateway/tests/test_privacy_api_redesign.py -v 2>&1 | head -30
Expected: FAIL — test_get_mode_returns_enum fails (returns enabled not mode)
Step 3: Rewrite privacy_api.py
Replace services/privacy_gateway/privacy_api.py with:
"""Privacy Gateway REST API — aiohttp server on port 8889.
Provides /api/v1/privacy/* endpoints consumed by the KVM WebUI frontend.
Mode enum: "off" | "audit" | "redact"
"""
from __future__ import annotations
import json
import logging
from pathlib import Path
from aiohttp import web
from .audit_logger import AuditLogger
from .cert_manager import get_ca_pem
log = logging.getLogger("privacy-api")
_STATE_PATH = Path("/var/lib/kvm-privacy/state.json")
_VALID_MODES = {"off", "audit", "redact"}
def _load_state() -> dict:
try:
return json.loads(_STATE_PATH.read_text())
except (FileNotFoundError, json.JSONDecodeError):
return {"mode": "off"}
def _save_state(state: dict) -> None:
_STATE_PATH.parent.mkdir(parents=True, exist_ok=True)
_STATE_PATH.write_text(json.dumps(state))
class PrivacyAPIServer:
def __init__(self, host: str = "0.0.0.0", port: int = 8889) -> None:
self._host = host
self._port = port
self._audit = AuditLogger()
self._state = _load_state()
self._runner: web.AppRunner | None = None
async def start(self) -> None:
app = web.Application(middlewares=[self._cors_middleware])
app.router.add_get("/api/v1/privacy/mode", self._get_mode)
app.router.add_post("/api/v1/privacy/mode", self._set_mode)
app.router.add_put("/api/v1/privacy/mode", self._set_mode)
app.router.add_get("/api/v1/privacy/stats", self._get_stats)
app.router.add_get("/api/v1/privacy/audit", self._get_audit)
app.router.add_get("/api/v1/privacy/cert", self._get_cert)
self._runner = web.AppRunner(app)
await self._runner.setup()
site = web.TCPSite(self._runner, self._host, self._port)
await site.start()
log.info("Privacy API listening on %s:%s", self._host, self._port)
async def stop(self) -> None:
if self._runner:
await self._runner.cleanup()
def get_mode(self) -> str:
"""Return current mode string — called by addon.py to check state."""
return self._state.get("mode", "off")
# ── middleware ──────────────────────────────────────────
@web.middleware
async def _cors_middleware(self, request: web.Request, handler):
if request.method == "OPTIONS":
resp = web.Response(status=204)
else:
try:
resp = await handler(request)
except web.HTTPException as exc:
resp = exc
resp.headers["Access-Control-Allow-Origin"] = "*"
resp.headers["Access-Control-Allow-Methods"] = "GET, POST, PUT, OPTIONS"
resp.headers["Access-Control-Allow-Headers"] = "Content-Type, Authorization"
return resp
# ── handlers ───────────────────────────────────────────
async def _get_mode(self, _request: web.Request) -> web.Response:
return web.json_response({"mode": self._state.get("mode", "off")})
async def _set_mode(self, request: web.Request) -> web.Response:
body = await request.json()
mode = body.get("mode", "")
if mode not in _VALID_MODES:
return web.json_response(
{"error": f"mode must be one of {sorted(_VALID_MODES)}"},
status=400,
)
self._state["mode"] = mode
_save_state(self._state)
log.info("Privacy mode set to %s", mode)
return web.json_response({"mode": mode})
async def _get_stats(self, _request: web.Request) -> web.Response:
raw = self._audit.stats_today()
by_action = raw.get("by_action", {})
by_domain = raw.get("by_domain", {})
return web.json_response({
"requests": raw.get("requests", 0),
"files": raw.get("files", 0),
"pii_by_type": raw.get("pii_by_type", {}),
"actions": {
"allow": by_action.get("allow", 0),
"block": by_action.get("block", 0),
"redact": by_action.get("redact", 0),
},
"top_domains": [
{"domain": d, "count": c}
for d, c in sorted(by_domain.items(), key=lambda x: -x[1])
],
})
async def _get_audit(self, request: web.Request) -> web.Response:
page = int(request.query.get("page", "1"))
page_size = int(request.query.get("page_size", "20"))
domain = request.query.get("domain") or None
offset = (page - 1) * page_size
rows = self._audit.query(limit=page_size, offset=offset, domain=domain)
total = self._audit.count(domain=domain)
logs = []
for r in rows:
pii_types: dict = {}
if r.get("pii_types"):
try:
pii_types = json.loads(r["pii_types"])
except (json.JSONDecodeError, TypeError):
pass
logs.append({
"id": r["id"],
"ts": r.get("ts", ""),
"domain": r.get("domain", ""),
"action": r.get("action", ""),
"pii_types": pii_types,
"filename": r.get("filename", ""),
"file_size": r.get("file_size", 0),
"client_ip": r.get("client_ip", ""),
"request_url": r.get("request_url", ""),
})
return web.json_response({"logs": logs, "total": total, "page": page})
async def _get_cert(self, _request: web.Request) -> web.Response:
try:
pem = get_ca_pem()
except Exception as exc:
log.error("Failed to read CA cert: %s", exc)
return web.Response(status=500, text="CA certificate not available")
return web.Response(
body=pem,
content_type="application/x-pem-file",
headers={"Content-Disposition": 'attachment; filename="kvm-privacy-ca.crt"'},
)
async def main() -> None:
logging.basicConfig(level=logging.INFO, format="%(name)s %(message)s")
server = PrivacyAPIServer()
await server.start()
import asyncio
try:
await asyncio.Event().wait()
except (KeyboardInterrupt, asyncio.CancelledError):
pass
finally:
await server.stop()
if __name__ == "__main__":
import asyncio
asyncio.run(main())
Step 4: Run tests to verify they pass
cd services && python -m pytest privacy_gateway/tests/test_privacy_api_redesign.py -v
Expected: all 5 PASS
Step 5: Commit
git add services/privacy_gateway/privacy_api.py
git commit -m "feat: redesign privacy API - mode enum off/audit/redact, fix stats and audit response format"
Task 3: addon.py Three-Mode Logic
Files:
- Modify:
services/privacy_gateway/addon.py
Context: addon.py currently always redacts. Need three branches: off=skip, audit=log-only (no redact), redact=log+redact. Action values standardized: "allow", "redact" (matches frontend PrivacyAuditEntry.action).
Step 1: Write failing test
Create services/privacy_gateway/tests/test_addon_modes.py:
"""Tests for addon.py three-mode logic."""
import pytest
from unittest.mock import MagicMock, patch, AsyncMock
import json
@pytest.fixture
def mock_state_path(tmp_path):
p = tmp_path / "state.json"
return p
def make_addon(mode: str, tmp_path):
state_path = tmp_path / "state.json"
state_path.write_text(json.dumps({"mode": mode}))
with patch("addon._STATE_PATH", state_path), \
patch("addon._audit") as mock_audit, \
patch("addon.ctx") as mock_ctx:
import importlib
import addon
importlib.reload(addon)
return addon.PrivacyGatewayAddon(), mock_audit, mock_ctx
def make_flow(host="api.openai.com", has_file=True):
flow = MagicMock()
flow.request.pretty_host = host
flow.client_conn.peername = ("127.0.0.1", 12345)
flow.request.url = "https://api.openai.com/v1/files"
return flow
def test_off_mode_skips_all(tmp_path):
"""In off mode, even AI domain file uploads are skipped."""
with patch("addon._STATE_PATH", tmp_path / "state.json"), \
patch("addon._AI_DOMAINS", {"api.openai.com"}), \
patch("addon.is_file_upload", return_value=True), \
patch("addon.extract_files", return_value=[MagicMock()]), \
patch("addon._load_mode", return_value="off") as mock_mode, \
patch("addon.asyncio.ensure_future") as mock_future:
import importlib, addon
importlib.reload(addon)
a = addon.PrivacyGatewayAddon()
flow = make_flow()
a.request(flow)
mock_future.assert_not_called()
@pytest.mark.asyncio
async def test_audit_mode_logs_no_redact(tmp_path):
"""In audit mode, files are scanned but NOT redacted."""
mock_audit = MagicMock()
mock_scan_result = MagicMock(
pii_found=True,
pii_types={"id_card": 1},
redacted_bytes=None,
)
with patch("addon._load_mode", return_value="audit"), \
patch("addon._audit", mock_audit), \
patch("addon.scan_and_redact", AsyncMock(return_value=mock_scan_result)), \
patch("addon.ctx", MagicMock()):
import importlib, addon
importlib.reload(addon)
a = addon.PrivacyGatewayAddon()
mock_file = MagicMock(filename="test.pdf", data=b"data", field_name="file")
flow = make_flow()
flow.request.content = b"original"
await a._process_upload(flow, "api.openai.com", [mock_file])
# logged
mock_audit.log.assert_called_once()
call_kwargs = mock_audit.log.call_args
assert call_kwargs[0][2] == "allow" # action = allow in audit mode
# NOT redacted — flow content unchanged
assert flow.request.content == b"original"
@pytest.mark.asyncio
async def test_redact_mode_logs_and_redacts(tmp_path):
"""In redact mode, PII files are redacted."""
mock_audit = MagicMock()
redacted_bytes = b"[REDACTED]"
mock_scan_result = MagicMock(
pii_found=True,
pii_types={"id_card": 1},
redacted_bytes=redacted_bytes,
)
with patch("addon._load_mode", return_value="redact"), \
patch("addon._audit", mock_audit), \
patch("addon.scan_and_redact", AsyncMock(return_value=mock_scan_result)), \
patch("addon.rebuild_multipart", return_value=b"rebuilt"), \
patch("addon.ctx", MagicMock()):
import importlib, addon
importlib.reload(addon)
a = addon.PrivacyGatewayAddon()
mock_file = MagicMock(filename="test.pdf", data=b"data", field_name="file")
flow = make_flow()
await a._process_upload(flow, "api.openai.com", [mock_file])
# logged with action=redact
call_kwargs = mock_audit.log.call_args
assert call_kwargs[0][2] == "redact"
# flow content was rebuilt
assert flow.request.content == b"rebuilt"
Step 2: Run to verify it fails
cd services/privacy_gateway && python -m pytest tests/test_addon_modes.py::test_off_mode_skips_all -v 2>&1 | head -20
Expected: FAIL or import error
Step 3: Update addon.py
Replace the addon.py body with updated version. Key changes:
- Add
_load_mode()function that reads state.json - In
request(): check mode first, return if "off" - In
_process_upload(): branch on mode for action and redaction - Standardize action values: "allow" (audit mode or no PII), "redact" (redact mode + PII found)
"""KVM Privacy Gateway —— mitmproxy addon 核心。
Mode logic:
off → pass through, no logging
audit → log PII detections, do NOT redact
redact → log PII detections AND redact files
"""
from __future__ import annotations
import asyncio
import hashlib
import json
import logging
from pathlib import Path
import httpx
from mitmproxy import ctx, http
_MAX_RETRIES = 3
_RETRY_BACKOFF = (0.5, 1.0, 2.0)
from audit_logger import AuditLogger
from interceptor import extract_files, is_file_upload, rebuild_multipart
from upload_scanner import scan_and_redact
logger = logging.getLogger(__name__)
_DOMAINS_FILE = Path(__file__).parent / "ai_domains.txt"
_STATE_FILE = Path("/var/lib/kvm-privacy/state.json")
_audit = AuditLogger()
_KVM_AUDIT_URL = "http://127.0.0.1:8080/api/internal/privacy-event"
_http_client: httpx.AsyncClient | None = None
def _get_http_client() -> httpx.AsyncClient:
global _http_client
if _http_client is None:
_http_client = httpx.AsyncClient(timeout=5.0)
return _http_client
def _load_mode() -> str:
"""Read current privacy mode from state.json. Defaults to 'off'."""
try:
state = json.loads(_STATE_FILE.read_text())
mode = state.get("mode", "off")
return mode if mode in ("off", "audit", "redact") else "off"
except (FileNotFoundError, json.JSONDecodeError):
return "off"
def _load_domains() -> set[str]:
if not _DOMAINS_FILE.exists():
return set()
return {
line.strip()
for line in _DOMAINS_FILE.read_text().splitlines()
if line.strip() and not line.startswith("#")
}
_AI_DOMAINS: set[str] = _load_domains()
class PrivacyGatewayAddon:
def request(self, flow: http.HTTPFlow) -> None:
mode = _load_mode()
if mode == "off":
return # pass through, no logging
host = flow.request.pretty_host
if host not in _AI_DOMAINS:
return
if not is_file_upload(flow):
return
files = extract_files(flow)
if not files:
return
ctx.log.info(f"[privacy-gw] mode={mode} intercepted {host} {len(files)} file(s)")
asyncio.ensure_future(self._process_upload(flow, host, files, mode))
async def _process_upload(self, flow, host: str, files, mode: str) -> None:
replacements: dict[str, bytes] = {}
total_files = len(files)
client_ip = flow.client_conn.peername[0] if flow.client_conn.peername else ""
request_url = flow.request.url
for f in files:
try:
result = await scan_and_redact(f.filename, f.data)
except Exception as e:
ctx.log.warn(f"[privacy-gw] scan failed {f.filename}: {e}")
_audit.log(
host, {}, "allow", f.data, total_files,
client_ip=client_ip, request_url=request_url,
filename=f.filename, file_size=len(f.data),
)
continue
for pii_type, count in result.pii_types.items():
pass # just for logging totals
if mode == "audit":
# Log only — do not redact
action = "allow"
_audit.log(
host, result.pii_types, action, f.data, total_files,
client_ip=client_ip, request_url=request_url,
filename=f.filename, file_size=len(f.data),
)
if result.pii_found:
ctx.log.info(f"[privacy-gw] audit: {f.filename} has PII {result.pii_types}")
elif mode == "redact":
if result.pii_found and result.redacted_bytes:
action = "redact"
replacements[f.field_name] = result.redacted_bytes
ctx.log.info(f"[privacy-gw] redacted {f.filename}: {result.pii_types}")
else:
action = "allow"
_audit.log(
host, result.pii_types, action, f.data, total_files,
client_ip=client_ip, request_url=request_url,
filename=f.filename, file_size=len(f.data),
)
doc_hash = hashlib.sha256(f.data).hexdigest()
await self._post_to_kvm_audit(
domain=host, pii_types=result.pii_types, action=action,
doc_hash=doc_hash, file_count=total_files,
client_ip=client_ip, request_url=request_url,
filename=f.filename, file_size=len(f.data),
)
if replacements:
flow.request.content = rebuild_multipart(flow, replacements)
@staticmethod
async def _post_to_kvm_audit(**kwargs) -> None:
for attempt in range(_MAX_RETRIES):
try:
resp = await _get_http_client().post(_KVM_AUDIT_URL, json=kwargs)
if resp.status_code < 400:
return
logger.warning("[privacy-gw] KVM audit POST failed: %d (attempt %d)", resp.status_code, attempt + 1)
except Exception as e:
logger.debug("[privacy-gw] KVM audit POST unreachable: %s (attempt %d)", e, attempt + 1)
if attempt < _MAX_RETRIES - 1:
await asyncio.sleep(_RETRY_BACKOFF[attempt])
addons = [PrivacyGatewayAddon()]
Step 4: Run tests
cd services/privacy_gateway && python -m pytest tests/test_addon_modes.py -v
Expected: all 3 tests PASS
Step 5: Commit
git add services/privacy_gateway/addon.py services/privacy_gateway/tests/test_addon_modes.py
git commit -m "feat: addon.py three-mode logic (off/audit/redact), standardize action values"
Task 4: Runner Task Cancellation
Files:
- Modify:
services/kvm_agent/runner.py - Modify:
services/kvm_agent/api_server.py
Context: _handle_cancel_task only cancels "pending" tasks. Need to cancel "running" tasks by cancelling their asyncio.Task. Runner wraps _execute_task in asyncio.create_task() and stores it.
Step 1: Write failing test
Create services/kvm_agent/tests/test_runner_cancel.py:
"""Tests for runner task cancellation."""
import asyncio
import pytest
from unittest.mock import AsyncMock, MagicMock
from services.kvm_agent.runner import TaskQueue, AutonomousRunner, TaskStatus
@pytest.fixture
def queue(tmp_path):
q = TaskQueue(str(tmp_path / "test.db"))
q.open()
yield q
q.close()
def test_queue_cancel_pending(queue):
task_id = queue.enqueue("task", "test task")
queue.cancel(task_id, "User cancelled")
task = queue.get_task(task_id)
assert task.status == "cancelled"
def test_queue_cancel_running(queue):
task_id = queue.enqueue("task", "test task")
# simulate dequeue sets status to running
queue.dequeue()
queue.cancel(task_id, "User cancelled")
task = queue.get_task(task_id)
assert task.status == "cancelled"
@pytest.mark.asyncio
async def test_runner_cancel_current_task(queue):
"""Cancelling current task sets status to cancelled."""
agent = MagicMock()
cancel_event = asyncio.Event()
async def slow_run(description):
await asyncio.sleep(10) # will be cancelled
return MagicMock(success=True, final_reason="done")
agent.run_task = slow_run
runner = AutonomousRunner(agent, queue, poll_interval=0.1)
task_id = queue.enqueue("task", "slow task")
# dequeue to mark as running
queued = queue.dequeue()
# Start execution in background
exec_coro = asyncio.create_task(runner._execute_task(queued))
await asyncio.sleep(0.05) # let it start
# Cancel
success = runner.cancel_current(task_id)
assert success is True
await asyncio.gather(exec_coro, return_exceptions=True)
task = queue.get_task(task_id)
assert task.status == "cancelled"
Step 2: Run to verify it fails
cd services && python -m pytest kvm_agent/tests/test_runner_cancel.py -v 2>&1 | head -30
Expected: FAIL — queue.cancel not found
Step 3: Update runner.py
Add to TaskQueue class after complete():
def cancel(self, task_id: int, result: str = "Cancelled by user") -> None:
"""Mark a task as cancelled regardless of current status."""
assert self._conn is not None
self._conn.execute(
"UPDATE tasks SET status = 'cancelled', finished_at = ?, result = ? WHERE id = ?",
(time.time(), result, task_id),
)
self._conn.commit()
Update AutonomousRunner.__init__ to track current execution:
def __init__(self, agent, queue: TaskQueue, poll_interval: float = 5.0):
self._agent = agent
self._queue = queue
self._poll_interval = poll_interval
self._shutdown = False
self._current_task_id: int | None = None
self._running_coro: asyncio.Task | None = None
Update start() to use create_task:
async def start(self):
logger.info("Autonomous runner started (poll=%.1fs)", self._poll_interval)
loop = asyncio.get_event_loop()
for sig in (signal.SIGINT, signal.SIGTERM):
loop.add_signal_handler(sig, self._handle_signal)
while not self._shutdown:
task = self._queue.dequeue()
if task:
logger.info("Dequeued task #%d: [%s] %s", task.task_id, task.task_type, task.description)
self._running_coro = asyncio.create_task(self._execute_task(task))
try:
await self._running_coro
except asyncio.CancelledError:
pass
finally:
self._running_coro = None
self._current_task_id = None
else:
await asyncio.sleep(self._poll_interval)
logger.info("Autonomous runner stopped")
Update _execute_task():
async def _execute_task(self, task: QueuedTask):
self._current_task_id = task.task_id
try:
result = await self._agent.run_task(task.description)
self._queue.complete(task.task_id, result=result.final_reason, success=result.success)
logger.info("Task #%d %s: %s", task.task_id, "succeeded" if result.success else "failed", result.final_reason)
except asyncio.CancelledError:
self._queue.cancel(task.task_id, result="Cancelled by user")
logger.info("Task #%d cancelled", task.task_id)
raise # re-raise so asyncio knows it was cancelled
except Exception as e:
logger.error("Task #%d crashed: %s", task.task_id, e)
self._queue.complete(task.task_id, result=str(e), success=False)
Add cancel_current() method:
def cancel_current(self, task_id: int) -> bool:
"""Cancel the currently running task by asyncio cancellation."""
if self._current_task_id == task_id and self._running_coro is not None:
self._running_coro.cancel()
return True
return False
Step 4: Update api_server.py
Add runner parameter to AgentAPIServer.__init__:
def __init__(
self,
queue: TaskQueue,
config: AgentConfig,
runner: "AutonomousRunner | None" = None,
host: str = "0.0.0.0",
port: int = 8890,
):
self._queue = queue
self._config = config
self._runner = runner
self._host = host
self._port = port
self._runner_instance: Optional[web.AppRunner] = None
Update _handle_cancel_task:
async def _handle_cancel_task(self, request: web.Request) -> web.Response:
try:
task_id = int(request.match_info["id"])
except ValueError:
return _error("Invalid task ID")
task = self._queue.get_task(task_id)
if not task:
return _error("Task not found", status=404)
if task.status == "pending":
self._queue.cancel(task_id, result="Cancelled by user")
return _json({"cancelled": True})
if task.status == "running":
if self._runner and self._runner.cancel_current(task_id):
return _json({"cancelled": True})
return _error("Task is running but runner not available", status=503)
return _error(f"Cannot cancel task in '{task.status}' status")
Step 5: Run tests
cd services && python -m pytest kvm_agent/tests/test_runner_cancel.py -v
Expected: all 3 PASS
Step 6: Commit
git add services/kvm_agent/runner.py services/kvm_agent/api_server.py services/kvm_agent/tests/test_runner_cancel.py
git commit -m "feat: runner task cancellation via asyncio.Task.cancel(), api_server cancel_current support"
Task 5: agentStore.ts + i18n keys
Files:
- Create:
deps/KVM/web/src/stores/agentStore.ts - Modify:
deps/KVM/web/public/locales/zh/translation.json - Modify:
deps/KVM/web/public/locales/en/translation.json
Context: Zustand store for agent state. Polls task status every 2s after submission. Stops polling on terminal state (completed/failed/cancelled).
Step 1: Create agentStore.ts
// deps/KVM/web/src/stores/agentStore.ts
import { create } from 'zustand'
import apiClient from '@/services/api/client'
type AgentStatus = 'idle' | 'running' | 'error'
export interface AgentTask {
id: number
description: string
status: string
success: boolean
result: string
started_at: number | null
finished_at: number | null
}
interface AgentStore {
status: AgentStatus
currentTask: AgentTask | null
taskHistory: AgentTask[]
queueDepth: number
submitTask: (description: string) => Promise<void>
cancelTask: (id: number) => Promise<void>
_pollTask: (id: number) => void
_stopPolling: () => void
}
const TERMINAL_STATES = new Set(['completed', 'failed', 'cancelled'])
let _pollTimer: ReturnType<typeof setInterval> | null = null
export const useAgentStore = create<AgentStore>((set, get) => ({
status: 'idle',
currentTask: null,
taskHistory: [],
queueDepth: 0,
submitTask: async (description) => {
try {
const res = await apiClient.post<{ task_id: number }>('/api/v1/agent/tasks', {
description,
task_type: 'task',
})
const id = res.data.task_id
set({ status: 'running', currentTask: { id, description, status: 'pending', success: false, result: '', started_at: null, finished_at: null } })
get()._pollTask(id)
} catch (err: any) {
set({ status: 'error' })
throw err
}
},
cancelTask: async (id) => {
get()._stopPolling()
try {
await apiClient.delete(`/api/v1/agent/tasks/${id}`)
} catch {
// ignore — task may already be done
}
set({ status: 'idle', currentTask: null })
},
_pollTask: (id) => {
get()._stopPolling()
_pollTimer = setInterval(async () => {
try {
const res = await apiClient.get<AgentTask>(`/api/v1/agent/tasks/${id}`)
const task = res.data
set({ currentTask: task })
if (TERMINAL_STATES.has(task.status)) {
get()._stopPolling()
set(s => ({
status: 'idle',
currentTask: null,
taskHistory: [task, ...s.taskHistory].slice(0, 10),
}))
}
} catch {
// service temporarily unavailable — keep polling
}
}, 2000)
},
_stopPolling: () => {
if (_pollTimer !== null) {
clearInterval(_pollTimer)
_pollTimer = null
}
},
}))
Step 2: Add i18n keys to zh translation
In deps/KVM/web/public/locales/zh/translation.json, add inside the top-level object:
"agent": {
"command_placeholder": "描述任务,例如:打开记事本",
"submit": "执行",
"cancel": "终止",
"status_idle": "待命",
"status_running": "执行中",
"task_completed": "已完成",
"task_failed": "失败",
"task_cancelled": "已终止"
},
And add privacy gateway mode keys inside "privacy" object:
"mode_label": "隐私网关模式",
"mode_off": "关闭",
"mode_audit": "审计",
"mode_redact": "脱敏",
"mode_off_desc": "不拦截任何请求",
"mode_audit_desc": "记录 PII 但不拦截",
"mode_redact_desc": "记录并自动脱敏上传文件",
Step 3: Add i18n keys to en translation
In deps/KVM/web/public/locales/en/translation.json, add:
"agent": {
"command_placeholder": "Describe a task, e.g. open notepad",
"submit": "Run",
"cancel": "Cancel",
"status_idle": "Idle",
"status_running": "Running",
"task_completed": "Completed",
"task_failed": "Failed",
"task_cancelled": "Cancelled"
},
And in "privacy":
"mode_label": "Privacy Gateway Mode",
"mode_off": "Off",
"mode_audit": "Audit",
"mode_redact": "Redact",
"mode_off_desc": "Pass through all requests",
"mode_audit_desc": "Log PII but do not intercept",
"mode_redact_desc": "Log and auto-redact uploaded files",
Step 4: Also update kvmStore.ts fetchPrivacyMode
In deps/KVM/web/src/stores/kvmStore.ts, update fetchPrivacyMode:
fetchPrivacyMode: async () => {
try {
const res = await api.get('/api/v1/privacy/mode')
// mode is "off"|"audit"|"redact"; any non-off mode enables video overlay
const mode = res.data.mode as string
set({ privacyMode: mode !== 'off' })
} catch {
// service unavailable — default false
}
},
Step 5: Verify TypeScript compiles
cd deps/KVM/web && npx tsc --noEmit 2>&1 | head -20
Expected: no errors (or only pre-existing errors)
Step 6: Commit
cd deps/KVM && git add web/src/stores/agentStore.ts web/src/stores/kvmStore.ts web/public/locales/zh/translation.json web/public/locales/en/translation.json
git commit -m "feat: agentStore Zustand store with submit/poll/cancel, i18n keys for agent + privacy modes"
Task 6: AgentCommandBar Component + ConsolePage Wiring
Files:
- Create:
deps/KVM/web/src/components/console/AgentCommandBar.tsx - Modify:
deps/KVM/web/src/pages/ConsolePage.tsx - Modify:
deps/KVM/web/src/pages/ConsolePage.css
Step 1: Create AgentCommandBar.tsx
// deps/KVM/web/src/components/console/AgentCommandBar.tsx
import { useState, KeyboardEvent } from 'react'
import { useTranslation } from 'react-i18next'
import { useAgentStore } from '../../stores/agentStore'
export default function AgentCommandBar() {
const { t } = useTranslation()
const { status, currentTask, taskHistory, submitTask, cancelTask } = useAgentStore()
const [input, setInput] = useState('')
const [error, setError] = useState('')
const handleSubmit = async () => {
const desc = input.trim()
if (!desc) return
setError('')
try {
await submitTask(desc)
setInput('')
} catch {
setError(t('agent.task_failed'))
}
}
const handleKeyDown = (e: KeyboardEvent<HTMLInputElement>) => {
if (e.key === 'Enter') handleSubmit()
}
const isRunning = status === 'running'
return (
<div className="agent-command-bar">
<div className="agent-input-row">
<input
className="agent-input"
value={input}
onChange={e => setInput(e.target.value)}
onKeyDown={handleKeyDown}
placeholder={t('agent.command_placeholder')}
disabled={isRunning}
/>
{isRunning ? (
<button
className="agent-btn agent-btn-cancel"
onClick={() => currentTask && cancelTask(currentTask.id)}
>
{t('agent.cancel')}
</button>
) : (
<button
className="agent-btn agent-btn-submit"
onClick={handleSubmit}
disabled={!input.trim()}
>
{t('agent.submit')}
</button>
)}
</div>
<div className="agent-status-row">
<span className={`agent-status-dot ${isRunning ? 'running' : 'idle'}`} />
<span className="agent-status-text">
{isRunning
? `${t('agent.status_running')} #${currentTask?.id}`
: t('agent.status_idle')}
</span>
{error && <span className="agent-error">{error}</span>}
{taskHistory.slice(0, 3).map(task => (
<span key={task.id} className={`agent-history-item ${task.status === 'completed' && task.success ? 'success' : 'fail'}`}>
{task.status === 'completed' && task.success ? '✓' : '✗'}
{' '}
{task.description.length > 18 ? task.description.slice(0, 18) + '…' : task.description}
</span>
))}
</div>
</div>
)
}
Step 2: Add CSS to ConsolePage.css
Append to deps/KVM/web/src/pages/ConsolePage.css:
/* Agent Command Bar */
.agent-command-bar {
flex-shrink: 0;
background: var(--bg-dark, #111);
border-top: 1px solid var(--border, #333);
padding: 0.5rem 1rem;
display: flex;
flex-direction: column;
gap: 0.25rem;
}
.agent-input-row {
display: flex;
gap: 0.5rem;
align-items: center;
}
.agent-input {
flex: 1;
background: var(--bg-gray, #1e1e2e);
border: 1px solid var(--border, #333);
border-radius: 0.375rem;
color: var(--text, #eee);
padding: 0.375rem 0.75rem;
font-size: 0.875rem;
height: 2rem;
}
.agent-input:disabled {
opacity: 0.6;
}
.agent-input:focus {
outline: none;
border-color: var(--primary, #3b82f6);
}
.agent-btn {
height: 2rem;
padding: 0 0.75rem;
border-radius: 0.375rem;
font-size: 0.8rem;
cursor: pointer;
white-space: nowrap;
width: auto;
}
.agent-btn-submit {
background: var(--primary, #3b82f6);
color: #fff;
}
.agent-btn-submit:disabled {
opacity: 0.4;
cursor: not-allowed;
}
.agent-btn-cancel {
background: #b91c1c;
color: #fff;
}
.agent-status-row {
display: flex;
align-items: center;
gap: 0.5rem;
font-size: 0.75rem;
color: var(--text-secondary, #888);
}
.agent-status-dot {
width: 0.5rem;
height: 0.5rem;
border-radius: 50%;
flex-shrink: 0;
}
.agent-status-dot.idle { background: #4ade80; }
.agent-status-dot.running { background: #facc15; animation: pulse 1s ease-in-out infinite; }
@keyframes pulse {
0%, 100% { opacity: 1; }
50% { opacity: 0.4; }
}
.agent-status-text { color: var(--text-secondary, #888); }
.agent-error { color: #f87171; }
.agent-history-item {
padding: 0.125rem 0.375rem;
border-radius: 0.25rem;
font-size: 0.7rem;
background: rgba(255,255,255,0.05);
}
.agent-history-item.success { color: #4ade80; }
.agent-history-item.fail { color: #f87171; }
Step 3: Wire AgentCommandBar into ConsolePage
In deps/KVM/web/src/pages/ConsolePage.tsx:
Add import at top:
import AgentCommandBar from '../components/console/AgentCommandBar'
Add <AgentCommandBar /> after the closing </div> of .console-content and before {keyboardVisible && ...}:
return (
<div className="console-page">
<StatusBar />
<OCRAlertBanner />
<div className="console-content">
{/* ... existing content ... */}
</div>
<AgentCommandBar />
{keyboardVisible && (
<div className="keyboard-panel">
<VirtualKeyboard />
</div>
)}
{/* ... rest ... */}
</div>
)
Step 4: Verify build
cd deps/KVM/web && npm run build 2>&1 | tail -20
Expected: "built in Xs" — no TypeScript errors
Step 5: Commit
cd deps/KVM && git add web/src/components/console/AgentCommandBar.tsx web/src/pages/ConsolePage.tsx web/src/pages/ConsolePage.css
git commit -m "feat: AgentCommandBar at bottom of ConsolePage with submit/cancel/history"
Task 7: PrivacyPage CertificateTab — Three-Mode Selector
Files:
- Modify:
deps/KVM/web/src/pages/PrivacyPage.tsx
Context: Replace binary ToggleCard (privacy_enabled bool) with three-segment mode selector (off/audit/redact). Remove masking toggle — masking is now implicit in mode. Keep video overlay toggle (privacyMode in kvmStore) as separate concern.
Step 1: Update CertificateTab in PrivacyPage.tsx
Replace the CertificateTab function (lines 314–370):
// ─── CertificateTab ───────────────────────────────────────────────────────
type GatewayMode = 'off' | 'audit' | 'redact'
const MODES: { value: GatewayMode; labelKey: string; descKey: string }[] = [
{ value: 'off', labelKey: 'privacy.mode_off', descKey: 'privacy.mode_off_desc' },
{ value: 'audit', labelKey: 'privacy.mode_audit', descKey: 'privacy.mode_audit_desc' },
{ value: 'redact', labelKey: 'privacy.mode_redact', descKey: 'privacy.mode_redact_desc' },
]
function CertificateTab() {
const { t } = useTranslation()
const { privacyMode, setPrivacyMode } = useKvmStore()
const [gatewayMode, setGatewayMode] = useState<GatewayMode>('off')
const [modeLoading, setModeLoading] = useState(false)
useEffect(() => {
apiClient.get<{ mode: GatewayMode }>('/api/v1/privacy/mode')
.then(r => {
const mode = r.data.mode ?? 'off'
setGatewayMode(mode)
setPrivacyMode(mode !== 'off')
})
.catch(() => {})
}, [setPrivacyMode])
async function handleModeChange(mode: GatewayMode) {
setModeLoading(true)
try {
await apiClient.post('/api/v1/privacy/mode', { mode })
setGatewayMode(mode)
setPrivacyMode(mode !== 'off')
} finally {
setModeLoading(false)
}
}
const selectedDesc = MODES.find(m => m.value === gatewayMode)?.descKey ?? ''
return (
<div className="space-y-4 max-w-lg">
{/* Gateway mode selector */}
<div className="bg-gray-800 rounded-xl p-5 border border-gray-700">
<div className="font-medium text-white mb-3">{t('privacy.mode_label')}</div>
<div className="flex rounded-lg overflow-hidden border border-gray-600">
{MODES.map(({ value, labelKey }) => (
<button
key={value}
onClick={() => handleModeChange(value)}
disabled={modeLoading}
className={`flex-1 py-2 text-sm font-medium transition-colors disabled:opacity-50 ${
gatewayMode === value
? 'bg-blue-600 text-white'
: 'text-gray-400 hover:text-white hover:bg-gray-700'
}`}
>
{t(labelKey)}
</button>
))}
</div>
<p className="text-xs text-gray-400 mt-2">{t(selectedDesc)}</p>
</div>
{/* CA Certificate download */}
<div className="bg-gray-800 rounded-xl p-5 border border-gray-700">
<div className="font-medium text-white mb-1">{t('privacy.ca_cert')}</div>
<div className="text-sm text-gray-400 mb-3">{t('privacy.ca_cert_desc')}</div>
<a href="/api/v1/privacy/cert" download="kvm-privacy-ca.pem"
className="inline-flex items-center gap-2 px-4 py-2 bg-blue-600 hover:bg-blue-500 text-white text-sm rounded-lg transition-colors">
↓ {t('privacy.ca_cert')}
</a>
</div>
{/* Proxy setup instructions */}
<div className="bg-gray-800/50 rounded-xl p-5 border border-gray-700 text-sm text-gray-400">
<div className="font-medium text-gray-300 mb-2">{t('privacy.proxy_setup')}</div>
<code className="block bg-gray-900 rounded p-2 text-green-400 text-xs">
{t('privacy.proxy_host')}
</code>
</div>
</div>
)
}
Step 2: Remove unused ToggleCard if no other tab uses it
Check if ToggleCard is used elsewhere in the file. If only used in the old CertificateTab, remove the function definition (lines 35–52).
Step 3: Verify TypeScript
cd deps/KVM/web && npx tsc --noEmit 2>&1 | head -20
Expected: no new errors
Step 4: Verify build
cd deps/KVM/web && npm run build 2>&1 | tail -10
Expected: builds successfully
Step 5: Commit
cd deps/KVM && git add web/src/pages/PrivacyPage.tsx
git commit -m "feat: replace PrivacyPage mode toggle with three-segment selector (off/audit/redact)"
Task 8: Parent Repo Sync + Final Verification
Files:
- Modify:
deps/KVM(submodule pointer update)
Step 1: Run backend tests
cd services && python -m pytest privacy_gateway/tests/ kvm_agent/tests/test_runner_cancel.py -v
Expected: all pass
Step 2: Run Go tests
cd deps/KVM && go test ./go/internal/api/... -v 2>&1 | grep -E "PASS|FAIL|---"
Expected: all PASS
Step 3: Verify frontend build
cd deps/KVM/web && npm run build 2>&1 | tail -5
Expected: "built in Xs"
Step 4: Update parent repo submodule pointer
git add deps/KVM
git commit -m "chore: update KVM submodule (agent command bar + privacy mode selector)"
Step 5: Final git status check
git status
Expected: clean working tree or only untracked non-essential files
Implementation Notes
Go proxy auth
Routes use withJWT(deps, ...) — requires valid JWT. Python services have CORS disabled for proxied requests (handled by Go middleware).
Privacy mode race condition
addon.py calls _load_mode() on every request (reads state.json from disk). This is safe because writes are atomic JSON.write(), but adds ~1ms per intercepted request. Acceptable given intercept frequency.
Agent cancellation timing
Between runner.cancel_current() and actual cancellation, the task continues for one more step. Frontend polls every 2s — worst case visible lag is ~4s.
ToggleCard removal
If ToggleCard is imported/used nowhere else after CertificateTab update, remove it from PrivacyPage.tsx to avoid dead code.
kvmStore.privacyMode meaning
privacyMode: boolean in kvmStore controls the video PII overlay (bounding boxes on screen). It is now derived from gateway mode: true when mode is "audit" or "redact". These are separate concerns — video overlay can work even without gateway interception.