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

"""Tests for cloud/auto_memory.py + worker injection + HTTP settings.

Scope, stated so a future reader does not over-read it: v1
deterministic extraction only. No test here calls a model for
extraction, touches host.py, CODEX_HOME or any Codex thread; the
store, the worker hook and the HTTP endpoints only read/write the
control-plane database, and injection only changes the in-memory
task copy the worker hands to its executor.

Run from the repo root:
    python3 -m unittest discover -s cloud/tests -p "test_auto_memory.py" -v
"""

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

try:
    from cloud.auth import AuthStore
    from cloud.auto_memory import AutoMemoryStore
    from cloud.auto_memory import build_injection_prefix
    from cloud.auto_memory import process_due_tasks
    from cloud.auto_memory import redact_text
    from cloud.db import POSTGRES_SCHEMA
    from cloud.identity import IdentityStore
    from cloud.migrations import MIGRATION_TABLES
    from cloud.server import create_server
    from cloud.server import init_db
    from cloud.worker import Worker
except ImportError:
    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
    from cloud.auth import AuthStore
    from cloud.auto_memory import AutoMemoryStore
    from cloud.auto_memory import build_injection_prefix
    from cloud.auto_memory import process_due_tasks
    from cloud.auto_memory import redact_text
    from cloud.db import POSTGRES_SCHEMA
    from cloud.identity import IdentityStore
    from cloud.migrations import MIGRATION_TABLES
    from cloud.server import create_server
    from cloud.server import init_db
    from cloud.worker import Worker


def _tmp_db():
    fd, path = tempfile.mkstemp(suffix=".db")
    os.close(fd)
    os.unlink(path)
    return path


def _insert_completed(db_path, task_id, owner="u1", project="p1",
                      prompt="做登录页", result=None, started_ago=2000,
                      duration=120, goal=None):
    now = time.time()
    finished = now - 1000
    started = finished - duration
    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, goal,"
            " started_at, finished_at, result_text)"
            " VALUES(?, ?, ?, 'completed', NULL, '[]', 0, ?, ?, ?,"
            " NULL, ?, ?, ?, ?)",
            (task_id, project, prompt, owner, now, finished,
             goal, started, finished,
             result if result is not None else "登录页已完成。" + "x" * 300),
        )
        conn.commit()
    finally:
        conn.close()
    return finished


class RedactTest(unittest.TestCase):
    def test_key_value_and_bearer_redacted(self):
        out = redact_text("api_key=abc123 and password: hunter2")
        self.assertNotIn("abc123", out)
        self.assertNotIn("hunter2", out)
        self.assertIn("[REDACTED]", out)

    def test_plain_text_untouched(self):
        self.assertEqual(redact_text("hello world"), "hello world")
        self.assertEqual(redact_text(None), "")


class StoreTest(unittest.TestCase):
    def setUp(self):
        self.db = _tmp_db()
        self.store = AutoMemoryStore(self.db)

    def tearDown(self):
        try:
            self.store.close()
        finally:
            if os.path.exists(self.db):
                os.unlink(self.db)

    def test_defaults_off(self):
        settings = self.store.get_settings("u1")
        self.assertEqual(settings["generate_memories"], 0)
        self.assertEqual(settings["use_memories"], 0)
        self.assertEqual(settings["disable_on_external_context"], 1)
        self.assertEqual(settings["min_rate_limit_remaining_percent"], 20)
        self.assertEqual(settings["idle_seconds"], 900)

    def test_upsert_validates_ranges(self):
        with self.assertRaises(ValueError):
            self.store.upsert_settings("u1",
                                       min_rate_limit_remaining_percent=101)
        with self.assertRaises(ValueError):
            self.store.upsert_settings("u1", idle_seconds=-1)
        with self.assertRaises(ValueError):
            self.store.upsert_settings("u1", bogus_field=1)
        got = self.store.upsert_settings(
            "u1", generate_memories=True, use_memories=1,
            idle_seconds=0, min_duration_seconds=0, min_result_chars=0)
        self.assertEqual(got["generate_memories"], 1)
        self.assertEqual(got["idle_seconds"], 0)

    def test_task_flag_overrides_settings(self):
        settings = self.store.upsert_settings("u1", generate_memories=False)
        self.assertFalse(AutoMemoryStore.effective(settings, None,
                                                    "generate_memories"))
        self.store.set_task_flags("t1", generate_memories=True)
        flags = self.store.get_task_flags("t1")
        self.assertTrue(AutoMemoryStore.effective(settings, flags,
                                                  "generate_memories"))
        with self.assertRaises(ValueError):
            self.store.set_task_flags("t2")

    def test_pg_schema_and_migration_lists_include_new_tables(self):
        schema_text = "\n".join(POSTGRES_SCHEMA)
        for table in ("memory_settings", "task_memory_flags",
                      "memory_extraction_log"):
            self.assertIn("CREATE TABLE IF NOT EXISTS %s" % table,
                          schema_text)
            self.assertIn(table, MIGRATION_TABLES)


class ExtractionTest(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)
        store = AutoMemoryStore(self.db_path)
        store.upsert_settings(
            "u1", generate_memories=True, use_memories=True,
            idle_seconds=900, min_duration_seconds=30,
            min_result_chars=200, min_rate_limit_remaining_percent=0)
        store.close()
        # Unlimited quota: explicit NULL row, so the quota gate passes.
        conn = sqlite3.connect(self.db_path)
        try:
            conn.execute(
                "INSERT INTO user_quotas(user_id, quota_usd_micros,"
                " updated_at) VALUES('u1', NULL, ?)", (time.time(),))
            conn.commit()
        finally:
            conn.close()

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

    def _memories(self):
        store = IdentityStore(self.db_path)
        try:
            return store.list_project_memories("p1")
        finally:
            store.conn.close()

    def test_happy_path_creates_project_memory_and_is_idempotent(self):
        _insert_completed(self.db_path, "t1")
        results = process_due_tasks(self.db_path)
        self.assertEqual(results[0]["status"], "extracted")
        self.assertEqual(results[0]["reason"], "created")
        memories = self._memories()
        self.assertEqual(len(memories), 1)
        self.assertEqual(memories[0]["scope"], "project")
        self.assertIn("做登录页", memories[0]["content"])
        # Second sweep: no duplicate, same memory reused.
        results2 = process_due_tasks(self.db_path)
        self.assertEqual(results2[0]["reason"], "already_extracted")
        self.assertEqual(len(self._memories()), 1)

    def test_secret_in_result_is_redacted(self):
        _insert_completed(self.db_path, "t1",
                          result="完成，api_key=should-not-leak " + "y" * 300)
        process_due_tasks(self.db_path)
        self.assertNotIn("should-not-leak", self._memories()[0]["content"])

    def test_idle_gate_and_force(self):
        now = time.time()
        conn = sqlite3.connect(self.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, started_at,"
                " finished_at, result_text)"
                " VALUES('t1', 'p1', '做登录页', 'completed', NULL,"
                " '[]', 0, 'u1', ?, ?, NULL, ?, ?, ?)",
                (now, now, now - 120, now, "z" * 300))
            conn.commit()
        finally:
            conn.close()
        results = process_due_tasks(self.db_path)
        self.assertEqual(results[0]["reason"], "idle_not_reached")
        self.assertEqual(self._memories(), [])
        results = process_due_tasks(self.db_path, force=True)
        self.assertEqual(results[0]["status"], "extracted")

    def test_too_short_skipped(self):
        _insert_completed(self.db_path, "t1", duration=5)
        results = process_due_tasks(self.db_path)
        self.assertEqual(results[0]["reason"], "too_short")
        self.assertEqual(self._memories(), [])

    def test_generate_disabled_skipped(self):
        store = AutoMemoryStore(self.db_path)
        store.upsert_settings("u1", generate_memories=False)
        store.close()
        _insert_completed(self.db_path, "t1")
        results = process_due_tasks(self.db_path)
        self.assertEqual(results[0]["reason"], "generate_disabled")
        self.assertEqual(self._memories(), [])

    def test_external_context_gate(self):
        _insert_completed(self.db_path, "t1")
        conn = sqlite3.connect(self.db_path)
        try:
            conn.execute(
                "CREATE TABLE IF NOT EXISTS task_followups("
                "id TEXT PRIMARY KEY, task_id TEXT, project_id TEXT,"
                " body TEXT, status TEXT, created_at REAL,"
                " sender_task_id TEXT)")
            conn.execute(
                "INSERT INTO task_followups(id, task_id, project_id,"
                " body, status, created_at, sender_task_id)"
                " VALUES('f1', 't1', 'p1', 'hi', 'done', ?, 'other')",
                (time.time(),))
            conn.commit()
        finally:
            conn.close()
        results = process_due_tasks(self.db_path)
        self.assertEqual(results[0]["reason"], "external_context")
        store = AutoMemoryStore(self.db_path)
        store.upsert_settings("u1", disable_on_external_context=False)
        store.close()
        results = process_due_tasks(self.db_path)
        self.assertEqual(results[0]["status"], "extracted")

    def test_quota_gate(self):
        import datetime
        month = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m")
        conn = sqlite3.connect(self.db_path)
        try:
            conn.execute(
                "UPDATE user_quotas SET quota_usd_micros = 10000000"
                " WHERE user_id = 'u1'")
            conn.execute(
                "CREATE TABLE IF NOT EXISTS cost_records("
                "id INTEGER PRIMARY KEY, task_id TEXT, project_id TEXT,"
                " model TEXT, input_tokens INTEGER, output_tokens INTEGER,"
                " cached_input_tokens INTEGER, total_tokens INTEGER,"
                " usd_micros INTEGER, source TEXT, fallback_count INTEGER,"
                " day TEXT, created_at REAL)")
            conn.execute(
                "INSERT INTO cost_records(task_id, project_id, model,"
                " input_tokens, output_tokens, cached_input_tokens,"
                " total_tokens, usd_micros, source, fallback_count, day,"
                " created_at) VALUES('t1', 'p1', 'm', 0, 0, 0, 0, 9500000,"
                " 'actual', 0, ?, ?)", (month + "-01", time.time()))
            conn.commit()
        finally:
            conn.close()
        store = AutoMemoryStore(self.db_path)
        store.upsert_settings(
            "u1", min_rate_limit_remaining_percent=20)
        store.close()
        _insert_completed(self.db_path, "t1")
        results = process_due_tasks(self.db_path)
        self.assertEqual(results[0]["reason"], "quota_low")
        self.assertEqual(self._memories(), [])

    def test_same_goal_consolidates_instead_of_duplicating(self):
        _insert_completed(self.db_path, "t1", result="第一版完成。" + "a" * 300)
        process_due_tasks(self.db_path)
        self.assertEqual(len(self._memories()), 1)
        _insert_completed(self.db_path, "t2", result="第二版完成。" + "b" * 300)
        results = process_due_tasks(self.db_path)
        extracted = [r for r in results if r["task_id"] == "t2"][0]
        self.assertEqual(extracted["reason"], "consolidated")
        memories = self._memories()
        self.assertEqual(len(memories), 1)
        self.assertIn("第二版完成", memories[0]["content"])


class InjectionTest(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)
        identity = IdentityStore(self.db_path)
        identity.add_memory("u1", "project", "convention", "提交前先跑测试",
                            project_id="p1")
        identity.conn.close()
        store = AutoMemoryStore(self.db_path)
        store.upsert_settings("u1", use_memories=True)
        store.close()
        self.env = {"XCM_API_KEY": "test-placeholder"}
        self.calls = []

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

    def _completed_runner(self, cmd):
        if "dispatch" in cmd:
            name = cmd[cmd.index("--name") + 1]
        else:
            name = cmd[cmd.index("--task") + 1]
        return SimpleNamespace(
            returncode=0,
            stdout=json.dumps({"task": name, "status": "completed",
                               "text": "done"}),
            stderr="")

    def _recording_runner(self, inner):
        def fake(cmd):
            self.calls.append(list(cmd))
            return inner(cmd)
        return fake

    def _insert_queued(self, task_id, prompt="做新功能", owner="u1"):
        now = time.time()
        conn = sqlite3.connect(self.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)"
                " VALUES(?, 'p1', ?, 'queued', NULL, '[]', 0, ?, ?, ?,"
                " NULL)",
                (task_id, prompt, owner, now, now),
            )
            conn.commit()
        finally:
            conn.close()

    def test_worker_prepends_memory_to_first_round_only(self):
        self._insert_queued("n1")
        worker = Worker(self.db_path, "/fake/host.py",
                        runner=self._recording_runner(self._completed_runner),
                        env=self.env)
        results = worker.run_all()
        self.assertEqual([r["status"] for r in results], ["completed"])
        dispatch = [c for c in self.calls if "dispatch" in c][0]
        task_arg = dispatch[dispatch.index("--task") + 1]
        self.assertIn("提交前先跑测试", task_arg)
        self.assertIn("做新功能", task_arg)
        # DB prompt is unchanged: injection only touched the copy.
        conn = sqlite3.connect(self.db_path)
        try:
            row = conn.execute(
                "SELECT prompt FROM tasks WHERE id = 'n1'").fetchone()
        finally:
            conn.close()
        self.assertEqual(row[0], "做新功能")

    def test_injection_off_by_default(self):
        store = AutoMemoryStore(self.db_path)
        store.upsert_settings("u1", use_memories=False)
        store.close()
        self.assertEqual(
            build_injection_prefix(self.db_path, "u1", "p1"), "")


class EndpointsTest(unittest.TestCase):
    def setUp(self):
        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()
        status, self.project = self._req(
            "POST", "/projects", {"name": "demo"}, self.owner_key)
        self.assertEqual(status, 201)

    def tearDown(self):
        self.server.shutdown()
        self.server.server_close()
        self.thread.join(timeout=5)
        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 test_settings_get_put_and_validation(self):
        status, body = self._req(
            "GET", "/users/me/auto-memory-settings", key=self.owner_key)
        self.assertEqual(status, 200)
        self.assertEqual(body["generate_memories"], 0)
        status, body = self._req(
            "PUT", "/users/me/auto-memory-settings",
            {"generate_memories": True, "use_memories": True},
            self.owner_key)
        self.assertEqual(status, 200, body)
        self.assertEqual(body["generate_memories"], 1)
        status, _ = self._req(
            "PUT", "/users/me/auto-memory-settings",
            {"idle_seconds": -5}, self.owner_key)
        self.assertEqual(status, 400)

    def test_run_and_log_and_task_flags(self):
        status, body = self._req(
            "POST", "/users/me/auto-memory/run", {"force": True},
            self.owner_key)
        self.assertEqual(status, 200, body)
        self.assertEqual(body["results"], [])
        status, body = self._req(
            "GET", "/users/me/auto-memory-log", key=self.owner_key)
        self.assertEqual(status, 200)
        self.assertEqual(body["log"], [])
        # Task flags: unknown task 404, real task 200.
        status, _ = self._req(
            "POST", "/tasks/nope/auto-memory-flags",
            {"use_memories": False}, self.owner_key)
        self.assertEqual(status, 404)
        status, task = self._req(
            "POST", "/tasks",
            {"project_id": self.project["id"], "prompt": "做事"},
            self.owner_key)
        self.assertEqual(status, 201, task)
        status, body = self._req(
            "POST", "/tasks/%s/auto-memory-flags" % task["id"],
            {"use_memories": False}, self.owner_key)
        self.assertEqual(status, 200, body)
        self.assertEqual(body["flags"]["use_memories"], 0)
        status, _ = self._req(
            "POST", "/tasks/%s/auto-memory-flags" % task["id"],
            {}, self.owner_key)
        self.assertEqual(status, 400)


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