host
cbc90f9d30
开关默认关（generate/use），设置接口 GET/PUT /users/me/auto-memory-settings，任务级覆盖 POST /tasks/{id}/auto-memory-flags，手动触发 POST /users/me/auto-memory/run，日志 GET /users/me/auto-memory-log。
提取：worker run_loop 轮次边界扫 completed 任务，闲置/时长/长度/外部上下文/配额剩余百分比五道门，脱敏后写 project 级 identity_memories，同一任务去重、同一目标 supersede 合并；只写 project 级，不跨项目泄漏。
注入：worker run_one 首轮把最新各 5 条用户/项目记忆拼进内存 prompt 副本（goal 任务同步改内存 goal），库原 prompt 不动；续跑/断点/协调员不注入，注入失败 fail-open。
v1 诚实边界：extract/consolidation 模型只存设置不调模型；MCP/web 线程内调用暂无 tasks 表标记，外部上下文门只认任务间消息与协调员子任务；无控制台 UI。
483 lines
19 KiB
Python
483 lines
19 KiB
Python
"""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()
|