codex-app-server/cloud/tests/test_audit_fixes4.py

"""Audit fix 4: third-audit gate findings.

- Explicit gateway capability split: known-missing blocks (active),
  unknown audits only; shadow never blocks.
- Explicit failures surfaced in the shadow report as an independent
  field, distinct from the two explicit_n counts.
- N2: coordinator requeue CAS rowcount=0 branch, via a proxy
  connection (the auditor's low-cost approach, no refactor).
- N3: shadow endpoint test on the create_server + real HTTP skeleton.
- N4: creation-layer rejections for tools / instructions /
  tool_approvals on goal/coordinator tasks, in one pass.
- N5: model writeback failure logs a warning (best-effort kept).

Run: python3 -m unittest discover -s cloud/tests -p "test_audit_fixes4.py" -v
No real models, no network, no real gateway: gateway steps are
stubbed exactly like test_audit_fixes.py / test_audit_fixes3.py.
"""

import http.client
import json
import os
import sqlite3
import sys
import tempfile
import threading
import time
import unittest
from types import SimpleNamespace
from unittest import mock

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))

import cloud.worker as worker_module
from cloud.auth import AuthStore
from cloud.models import split_model_capability_missing
from cloud.server import create_server, init_db
from cloud.worker import Worker


def _completed_runner(cmd):
    if "dispatch" in cmd:
        name = cmd[cmd.index("--name") + 1]
    else:
        name = cmd[cmd.index("--task") + 1]
    return SimpleNamespace(
        returncode=0,
        stdout='{"status": "completed", "result": "done",'
               ' "task_id": "%s", "thread_id": "th1"}' % name,
        stderr="")


def _insert_task(db_path, task_id, model=None, status="queued",
                 project_id="p1", coordinator_parts=None,
                 parent_task_id=None, coordinator_part=None,
                 result_text=None, error_text=None, paused=0,
                 interrupted_reason=None):
    now = time.time()
    conn = sqlite3.connect(db_path)
    try:
        conn.execute(
            "INSERT INTO tasks(id, project_id, prompt, status,"
            " idempotency_key, depends_on, priority, created_by,"
            " created_at, updated_at, last_heartbeat, coordinator_parts,"
            " parent_task_id, coordinator_part, result_text, error_text,"
            " model, paused_by_user, interrupted_reason)"
            " VALUES(?, ?, 'do work', ?, NULL, '[]', 0, 'tester',"
            " ?, ?, NULL, ?, ?, ?, ?, ?, ?, ?, ?)",
            (task_id, project_id, status, now, now,
             json.dumps(coordinator_parts) if coordinator_parts else None,
             parent_task_id, coordinator_part, result_text, error_text,
             model, paused, interrupted_reason),
        )
        conn.commit()
    finally:
        conn.close()


def _row(db_path, task_id):
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    try:
        row = conn.execute("SELECT * FROM tasks WHERE id = ?",
                           (task_id,)).fetchone()
        return dict(row) if row else None
    finally:
        conn.close()


def _audit_actions(db_path, project_id="p1"):
    store = AuthStore(db_path)
    try:
        return [(e["action"], e["detail"]) for e in store.list_audit(
            project_id)]
    finally:
        store.conn.close()


def _worker(db_path, policy_doc=None, env=None, gateway=None):
    calls = []

    def recording(cmd):
        calls.append(list(cmd))
        return _completed_runner(cmd)

    worker = Worker(db_path, "/fake/host.py", runner=recording,
                    env=env or {"XCM_API_KEY": "test-placeholder"},
                    routing_policy=policy_doc)
    if gateway is not None:
        worker._ensure_gateway = lambda task: gateway
        worker._sub2api_default_model = lambda e: "gateway-model"
    worker._calls = calls
    return worker


class SplitMissingTestCase(unittest.TestCase):
    def test_absent_model_is_all_unknown(self):
        known, unknown = split_model_capability_missing(
            {}, "no-such-model", {"tool_calling": True,
                                   "min_context_tokens": 1000})
        self.assertEqual(known, [])
        self.assertEqual(unknown, ["min_context_tokens", "tool_calling"])

    def test_explicit_false_is_known(self):
        known, unknown = split_model_capability_missing(
            {"m": {"tool_calling": False}},
            "m", {"tool_calling": True})
        self.assertEqual(known, ["tool_calling"])
        self.assertEqual(unknown, [])

    def test_unfilled_field_is_unknown(self):
        known, unknown = split_model_capability_missing(
            {"m": {}}, "m", {"tool_calling": True})
        self.assertEqual(known, [])
        self.assertEqual(unknown, ["tool_calling"])

    def test_short_context_is_known_and_missing_context_is_unknown(self):
        known, unknown = split_model_capability_missing(
            {"m": {"context_tokens": 100}},
            "m", {"min_context_tokens": 1000})
        self.assertEqual(known, ["min_context_tokens"])
        known2, unknown2 = split_model_capability_missing(
            {"m": {}}, "m", {"min_context_tokens": 1000})
        self.assertEqual(known2, [])
        self.assertEqual(unknown2, ["min_context_tokens"])

    def test_mixed_known_and_unknown_keeps_known(self):
        known, unknown = split_model_capability_missing(
            {"m": {"tool_calling": False}},
            "m", {"tool_calling": True, "vision": True})
        self.assertEqual(known, ["tool_calling"])
        self.assertEqual(unknown, ["vision"])


class ExplicitGatewaySplitTestCase(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.db_path = os.path.join(self.tmp.name, "control.db")
        init_db(self.db_path)

    def tearDown(self):
        self.tmp.cleanup()

    def _policy(self, mode, capabilities):
        return {
            "mode": mode,
            "tiers": {"default": {"model": "policy-model",
                                  "requires": {"tool_calling": True}}},
            "model_capabilities": capabilities,
        }

    def test_known_missing_explicit_is_blocked(self):
        _insert_task(self.db_path, "E1", model="explicit-model")
        worker = _worker(self.db_path, self._policy("active", {
            "explicit-model": {"tool_calling": False},
        }), gateway=("sub2api", {}))
        result = worker.run_one()
        self.assertEqual(result["status"], "failed")
        self.assertFalse(worker._calls)
        self.assertIn("tool_calling", result["error"])
        actions = _audit_actions(self.db_path)
        explicit = [d for a, d in actions
                    if a == "worker_capability_gateway_explicit"]
        self.assertTrue(explicit)
        self.assertIn("blocked=True", explicit[0])
        self.assertTrue([a for a, _ in actions
                         if a == "worker_capability_gate_rejected"])

    def test_unknown_explicit_is_audited_not_blocked(self):
        _insert_task(self.db_path, "E1", model="explicit-model")
        worker = _worker(self.db_path, self._policy("active", {}),
                         gateway=("sub2api", {}))
        result = worker.run_one()
        self.assertEqual(result["status"], "completed")
        self.assertTrue(worker._calls)
        actions = _audit_actions(self.db_path)
        explicit = [d for a, d in actions
                    if a == "worker_capability_gateway_explicit"]
        self.assertTrue(explicit)
        self.assertIn("blocked=False", explicit[0])
        self.assertFalse([a for a, _ in actions
                          if a == "worker_capability_gate_rejected"])

    def test_mixed_known_and_unknown_explicit_is_blocked(self):
        # One known-missing capability is enough to block: the task
        # cannot succeed regardless of the unknown field.
        _insert_task(self.db_path, "E1", model="explicit-model")
        policy = self._policy("active", {
            "explicit-model": {"tool_calling": False},
        })
        policy["tiers"]["default"]["requires"] = {
            "tool_calling": True, "vision": True}
        worker = _worker(self.db_path, policy, gateway=("sub2api", {}))
        result = worker.run_one()
        self.assertEqual(result["status"], "failed")
        self.assertFalse(worker._calls)

    def test_shadow_known_missing_explicit_is_not_blocked(self):
        _insert_task(self.db_path, "E1", model="explicit-model")
        worker = _worker(self.db_path, self._policy("shadow", {
            "explicit-model": {"tool_calling": False},
        }), gateway=("sub2api", {}))
        result = worker.run_one()
        self.assertEqual(result["status"], "completed")
        self.assertTrue(worker._calls)
        actions = _audit_actions(self.db_path)
        explicit = [d for a, d in actions
                    if a == "worker_capability_gateway_explicit"]
        self.assertTrue(explicit)
        self.assertIn("blocked=False", explicit[0])
        self.assertFalse([a for a, _ in actions
                          if a == "worker_capability_gate_rejected"])


class I7RowcountZeroTestCase(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.db_path = os.path.join(self.tmp.name, "control.db")
        init_db(self.db_path)

    def tearDown(self):
        self.tmp.cleanup()

    def test_rowcount_zero_skips_requeue_and_audit(self):
        parts = [{"name": "后端", "prompt": "做后端"}]
        _insert_task(self.db_path, "coord", coordinator_parts=parts)
        _insert_task(self.db_path, "c-old", status="failed",
                     parent_task_id="coord", coordinator_part="后端",
                     result_text="旧输出", error_text="旧错")
        worker = _worker(self.db_path, gateway=(None, None))
        task = worker.get_task("coord")
        orig_connect = worker_module._db_connect
        db_path = self.db_path

        class _FlippingConn:
            """Proxy: flip the child to completed just before the
            requeue UPDATE executes, so its rowcount is 0. The
            snapshot connection is wrapped too, but only the
            requeue UPDATE triggers the flip."""

            def __init__(self, real):
                self._real = real
                self._flipped = False

            def execute(self, sql, *args):
                if not self._flipped and \
                        "UPDATE tasks SET status = 'queued'" in sql:
                    flip = sqlite3.connect(db_path)
                    try:
                        flip.execute(
                            "UPDATE tasks SET status = 'completed'"
                            " WHERE id = 'c-old'")
                        flip.commit()
                    finally:
                        flip.close()
                    self._flipped = True
                return self._real.execute(sql, *args)

            def commit(self):
                return self._real.commit()

            def close(self):
                return self._real.close()

            def __getattr__(self, name):
                return getattr(self._real, name)

        with mock.patch.object(
                worker_module, "_db_connect",
                side_effect=lambda p: _FlippingConn(orig_connect(p))):
            children = worker._coord_sync_children(
                task, task["coordinator_parts"])
        self.assertEqual(children, [{"name": "后端", "id": "c-old"}])
        old = _row(self.db_path, "c-old")
        self.assertEqual(old["status"], "completed")
        self.assertEqual(old["result_text"], "旧输出")
        actions = _audit_actions(self.db_path)
        self.assertFalse([a for a, _ in actions
                          if a == "coordinator_child_requeued"])


class _HttpServerCase(unittest.TestCase):
    def setUp(self):
        self._old_policy_json = os.environ.get(
            "CLOUD_ROUTING_POLICY_JSON")
        self._old_policy_file = os.environ.get(
            "CLOUD_ROUTING_POLICY_FILE")
        os.environ.pop("CLOUD_ROUTING_POLICY_FILE", None)
        self.tmp = tempfile.TemporaryDirectory()
        self.db_path = os.path.join(self.tmp.name, "control.db")
        self.auth = AuthStore(self.db_path)
        self.owner_id, self.owner_key = self.auth.create_user("owner")
        self.server = create_server(self.db_path, "127.0.0.1", 0)
        self.port = self.server.server_address[1]
        self.thread = threading.Thread(
            target=self.server.serve_forever, daemon=True)
        self.thread.start()

    def tearDown(self):
        self.server.shutdown()
        self.server.server_close()
        self.thread.join(timeout=5)
        if self._old_policy_json is None:
            os.environ.pop("CLOUD_ROUTING_POLICY_JSON", None)
        else:
            os.environ["CLOUD_ROUTING_POLICY_JSON"] = \
                self._old_policy_json
        if self._old_policy_file is None:
            os.environ.pop("CLOUD_ROUTING_POLICY_FILE", None)
        else:
            os.environ["CLOUD_ROUTING_POLICY_FILE"] = \
                self._old_policy_file
        try:
            self.auth.conn.close()
        except Exception:
            pass
        self.tmp.cleanup()

    def _req(self, method, path, body=None, key=None):
        conn = http.client.HTTPConnection(
            "127.0.0.1", self.port, timeout=10)
        hdrs = {}
        payload = None
        if body is not None:
            payload = json.dumps(body).encode("utf-8")
            hdrs["Content-Type"] = "application/json"
        if key is not None:
            hdrs["X-API-Key"] = key
        try:
            conn.request(method, path, body=payload, headers=hdrs)
            resp = conn.getresponse()
            data = resp.read()
        finally:
            conn.close()
        try:
            parsed = json.loads(data.decode("utf-8")) if data else None
        except ValueError:
            parsed = None
        return resp.status, parsed

    def _project(self, name="p"):
        status, project = self._req(
            "POST", "/projects", {"name": name, "auto_pr": 1},
            self.owner_key)
        self.assertEqual(status, 201, project)
        return project["id"]

    def _insert_audit(self, project_id, action, detail):
        conn = sqlite3.connect(self.db_path)
        try:
            conn.execute(
                "INSERT INTO audit_log(user_id, project_id, action,"
                " detail, created_at) VALUES(?, ?, ?, ?, ?)",
                ("tester", project_id, action, detail, time.time()))
            conn.commit()
        finally:
            conn.close()


class ShadowEndpointTestCase(_HttpServerCase):
    def test_report_carries_capability_and_explicit_failures(self):
        os.environ["CLOUD_ROUTING_POLICY_JSON"] = json.dumps({
            "mode": "active",
            "tiers": {"default": {"model": "policy-model",
                                  "requires": {"tool_calling": True}}},
            "model_capabilities": {
                "policy-model": {"tool_calling": False},
            },
        })
        pid = self._project()
        _insert_task(self.db_path, "T1", model="policy-model",
                     status="completed", project_id=pid)
        self._insert_audit(
            pid, "worker_routing",
            "T1 tier=default source=policy applied=true"
            " suggested=policy-model chosen=policy-model")
        self._insert_audit(
            pid, "worker_capability_gateway_explicit",
            "E1 model=explicit-model tier=default mode=active"
            " missing=tool_calling unknown=False blocked=True")
        status, report = self._req(
            "GET", "/routing-shadow?project_id=" + pid,
            key=self.owner_key)
        self.assertEqual(status, 200, report)
        self.assertEqual(
            report["tiers"]["default"]["capability"]["status"], "fail")
        self.assertEqual(
            report["tiers"]["default"]["capability"]["missing"],
            ["tool_calling"])
        self.assertFalse(report["truncated"])
        failures = report["explicit_capability_failures"]
        self.assertEqual(len(failures), 1)
        self.assertEqual(failures[0], {
            "task_id": "E1", "model": "explicit-model",
            "tier": "default", "mode": "active",
            "missing": ["tool_calling"], "unknown": False,
            "blocked": True})
        self.assertFalse(report["explicit_capability_failures_truncated"])
        self.assertIn("top level", report["explicit_counts_note"])
        self.assertIn("per-tier", report["explicit_counts_note"])
        self.assertIn("worker_capability_gateway_explicit",
                      report["explicit_counts_note"])

    def test_report_truncated_flag(self):
        pid = self._project()
        _insert_task(self.db_path, "T1", model="policy-model",
                     status="completed", project_id=pid)
        conn = sqlite3.connect(self.db_path)
        try:
            conn.executemany(
                "INSERT INTO audit_log(user_id, project_id, action,"
                " detail, created_at) VALUES(?, ?, ?, ?, ?)",
                [("tester", pid, "worker_routing",
                  "T1 tier=default source=policy applied=true"
                  " suggested=policy-model chosen=policy-model",
                  time.time())] * 10001)
            conn.commit()
        finally:
            conn.close()
        status, report = self._req(
            "GET", "/routing-shadow?project_id=" + pid,
            key=self.owner_key)
        self.assertEqual(status, 200, report)
        self.assertTrue(report["truncated"])
        self.assertIn("truncated", report["recommendation"])


class CreateRejectionTestCase(_HttpServerCase):
    def _create(self, extra):
        pid = self._project()
        body = {"project_id": pid, "prompt": "协调干活"}
        body.update(extra)
        return self._req("POST", "/tasks", body, self.owner_key)

    def test_coordinator_rejects_tools_instructions_approvals(self):
        parts = [{"name": "后端", "prompt": "做后端"}]
        agents = [
            {"name": "x-agent", "tools": ["websearch"]},
            {"name": "x-agent", "instructions": "Do it."},
            {"name": "x-agent", "tool_approvals": {"websearch": "ask"}},
        ]
        for agent in agents:
            status, data = self._create({
                "coordinator_parts": parts, "agent": agent})
            self.assertEqual(status, 400, (agent, data))
            self.assertIn("cannot be used with", data["error"])

    def test_goal_rejects_tools_instructions_approvals(self):
        agents = [
            {"name": "x-agent", "tools": ["websearch"]},
            {"name": "x-agent", "instructions": "Do it."},
            {"name": "x-agent", "tool_approvals": {"websearch": "ask"}},
        ]
        for agent in agents:
            status, data = self._create({
                "goal_mode": 1, "agent": agent})
            self.assertEqual(status, 400, (agent, data))
            self.assertIn("cannot be used with", data["error"])


if __name__ == "__main__":
    unittest.main()