"""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()