Files
KVM-privacy/docs/plans/2026-03-03-agent-gateway-implementation.md
T

49 KiB
Raw Blame History

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:

  1. Add _load_mode() function that reads state.json
  2. In request(): check mode first, return if "off"
  3. In _process_upload(): branch on mode for action and redaction
  4. 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 314370):

// ─── 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 3552).

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.