"""cloud 规则路由 v0 单测：分类、决策、策略解析、worker 接入。 运行（在仓库根目录下）： python3 -m unittest discover -s cloud/tests -p "test_models_routing.py" -v 说明：只用纯函数与注入的假 runner，不跑真实模型回合、不联网、 不登录任何服务；env 映射也是注入的普通 dict，里面的值只是 占位用的假名，绝不是真实密钥。 """ import json import os import sqlite3 import sys import tempfile import time import unittest from pathlib import Path from types import SimpleNamespace try: from cloud.models import ( ModelRouter, classify_task, decide_routing, is_fallbackable, load_routing_policy, parse_routing_policy, prompt_len_bucket, record_routing, ) from cloud.server import init_db from cloud.worker import Worker except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from cloud.models import ( ModelRouter, classify_task, decide_routing, is_fallbackable, load_routing_policy, parse_routing_policy, prompt_len_bucket, record_routing, ) from cloud.server import init_db from cloud.worker import Worker POLICY = { "version": "routing-v0.1", "mode": "active", "tiers": { "plan": "gpt-6.1-sol", "goal": "gpt-6.1-sol", "default": "gpt-5.6-terra", }, } def _policy(**overrides): policy = json.loads(json.dumps(POLICY)) policy.update(overrides) return parse_routing_policy(policy) def _insert_task(db_path, task_id, prompt="do work", plan=0, model=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)" " VALUES(?, 'p1', ?, 'queued', NULL, '[]', 0, 'tester'," " ?, ?, NULL)", (task_id, prompt, now, now), ) conn.execute( "UPDATE tasks SET plan = ?, model = ? WHERE id = ?", (plan, model, task_id), ) conn.commit() finally: conn.close() 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=json.dumps({"task": name, "status": "completed", "text": "done"}), stderr="") class ClassifyTestCase(unittest.TestCase): def test_plan_beats_goal_and_default(self): self.assertEqual(classify_task({"plan": 1}), "plan") self.assertEqual(classify_task({"plan": 1, "goal_mode": 1}), "plan") self.assertEqual(classify_task({"goal_mode": 1}), "goal") self.assertEqual( classify_task({"verify_command": "make test"}), "goal") self.assertEqual(classify_task({}), "default") self.assertEqual(classify_task({"priority": 9}), "default") def test_prompt_len_bucket_is_logging_only(self): self.assertEqual(prompt_len_bucket("x" * 10), "<500") self.assertEqual(prompt_len_bucket("x" * 500), "500-3000") self.assertEqual(prompt_len_bucket("x" * 3001), ">3000") # 长度再长也不改变分层：分层只看结构信号。 self.assertEqual(classify_task({"prompt": "x" * 5000}), "default") class ParsePolicyTestCase(unittest.TestCase): def test_valid_policy(self): policy = _policy() self.assertEqual(policy.mode, "active") self.assertEqual(policy.version, "routing-v0.1") self.assertEqual(policy.tiers["plan"], "gpt-6.1-sol") def test_invalid_mode_or_missing_tiers_fails_closed(self): self.assertIsNone(parse_routing_policy( {"mode": "bogus", "tiers": {"default": "x"}})) self.assertIsNone(parse_routing_policy({"mode": "active"})) self.assertIsNone(parse_routing_policy("not-a-dict")) self.assertIsNone(parse_routing_policy(None)) def test_empty_tier_model_is_dropped_not_fatal(self): policy = parse_routing_policy( {"mode": "active", "tiers": {"plan": " ", "default": "gpt-5.6-terra"}}) self.assertNotIn("plan", policy.tiers) self.assertEqual(policy.tiers["default"], "gpt-5.6-terra") def test_load_from_json_and_mode_override(self): policy = load_routing_policy({ "CLOUD_ROUTING_POLICY_JSON": json.dumps( {"mode": "shadow", "tiers": {"default": "gpt-6.1-sol"}}), "CLOUD_ROUTING_MODE": "active", }) self.assertEqual(policy.mode, "active") self.assertIsNone(load_routing_policy( {"CLOUD_ROUTING_POLICY_JSON": "bad json"})) def test_load_from_file(self): with tempfile.NamedTemporaryFile( "w", suffix=".json", delete=False) as f: json.dump({"mode": "shadow", "tiers": {"plan": "gpt-6.1-sol"}}, f) path = f.name try: policy = load_routing_policy( {"CLOUD_ROUTING_POLICY_FILE": path}) self.assertEqual(policy.mode, "shadow") self.assertEqual(policy.tiers["plan"], "gpt-6.1-sol") finally: os.unlink(path) class DecideTestCase(unittest.TestCase): def test_explicit_task_model_wins_over_policy(self): decision = decide_routing( {"id": "t", "model": "gpt-6-sol", "plan": 1}, "gpt-5.6-terra", _policy()) self.assertEqual(decision.chosen_model, "gpt-6-sol") self.assertEqual(decision.source, "explicit") def test_worker_cli_model_wins_over_policy(self): decision = decide_routing( {"id": "t", "plan": 1}, "gpt-6.1-sol", _policy(), worker_model_explicit=True) self.assertEqual(decision.chosen_model, "gpt-6.1-sol") self.assertEqual(decision.source, "worker_explicit") def test_shadow_does_not_apply(self): decision = decide_routing( {"id": "t", "plan": 1}, "gpt-5.6-terra", _policy(mode="shadow")) self.assertFalse(decision.applied) self.assertEqual(decision.chosen_model, "gpt-5.6-terra") self.assertEqual(decision.suggested_model, "gpt-6.1-sol") self.assertEqual(decision.source, "shadow") def test_active_applies_tier_model(self): decision = decide_routing( {"id": "t", "plan": 1}, "gpt-5.6-terra", _policy()) self.assertTrue(decision.applied) self.assertEqual(decision.chosen_model, "gpt-6.1-sol") self.assertEqual(decision.source, "policy") self.assertEqual(decision.policy_version, "routing-v0.1") def test_off_or_no_policy_or_missing_tier_falls_back(self): for policy in (None, _policy(mode="off"), _policy(tiers={"plan": "gpt-6.1-sol"})): decision = decide_routing( {"id": "t", "goal_mode": 1}, "gpt-5.6-terra", policy) self.assertEqual(decision.chosen_model, "gpt-5.6-terra") self.assertEqual(decision.source, "worker_default") def test_decide_does_not_mutate_task_or_other_fields(self): task = {"id": "t", "plan": 1, "effort": "low", "approval": "ask", "executor": "docker"} before = dict(task) decide_routing(task, "gpt-5.6-terra", _policy()) self.assertEqual(task, before) def test_decide_is_deterministic(self): task = {"id": "t", "plan": 1} first = decide_routing(task, "gpt-5.6-terra", _policy()).to_event() second = decide_routing(task, "gpt-5.6-terra", _policy()).to_event() self.assertEqual(first, second) def test_record_routing_appends_event(self): record = {"id": "t"} decision = decide_routing( {"id": "t", "plan": 1}, "gpt-5.6-terra", _policy()) event = record_routing(record, decision) self.assertEqual(event["type"], "model_routing") self.assertEqual(record["events"][0], event) def test_fallback_rules_unchanged(self): # 路由 v0 不改事后 fallback：401/400 不触发，404 触发。 self.assertFalse(is_fallbackable(401)) self.assertFalse(is_fallbackable(400)) self.assertTrue(is_fallbackable(404)) router = ModelRouter(["gpt-6.1-sol", "gpt-6-sol"]) self.assertEqual(router.next_after("gpt-6.1-sol", 503), "gpt-6-sol") self.assertIsNone(router.next_after("gpt-6.1-sol", 401)) class WorkerRoutingTestCase(unittest.TestCase): def setUp(self): self._old_push_allow = os.environ.get( "CODEX_PUSH_ALLOW_PRIVATE_TARGETS") os.environ["CODEX_PUSH_ALLOW_PRIVATE_TARGETS"] = "1" self.tmp = tempfile.TemporaryDirectory() self.db_path = os.path.join(self.tmp.name, "control.db") init_db(self.db_path) self.env = {"XCM_API_KEY": "test-placeholder"} self.calls = [] def tearDown(self): if getattr(self, "_old_push_allow", None) is None: os.environ.pop("CODEX_PUSH_ALLOW_PRIVATE_TARGETS", None) else: os.environ["CODEX_PUSH_ALLOW_PRIVATE_TARGETS"] = ( self._old_push_allow) self.tmp.cleanup() def _recording_runner(self): def fake(cmd): self.calls.append(list(cmd)) return _completed_runner(cmd) return fake def _worker(self, **kwargs): kwargs.setdefault("env", self.env) kwargs.setdefault("model_explicit", False) return Worker(self.db_path, "/fake/host.py", runner=self._recording_runner(), **kwargs) def _dispatch_model(self): for cmd in self.calls: if "dispatch" in cmd and "--model" in cmd: return cmd[cmd.index("--model") + 1] return None def test_worker_init_explicit_flag_and_bad_policy(self): # _worker 默认 model_explicit=False（模拟未钉死的 worker）； # 这里显式传 True，验证显式标志与策略解析。 worker = self._worker(model="gpt-6.1-sol", model_explicit=True, routing_policy=POLICY) self.assertTrue(worker.model_explicit) self.assertEqual(worker.routing_policy.mode, "active") worker2 = self._worker( routing_policy={"mode": "bogus", "tiers": {"default": "x"}}) self.assertIsNone(worker2.routing_policy) # 不传 model_explicit 时按是否传 model 推断：传了即显式。 worker3 = Worker(self.db_path, "/fake/host.py", runner=self._recording_runner(), env=self.env, model="gpt-6.1-sol") self.assertTrue(worker3.model_explicit) def test_run_all_active_plan_uses_policy_model(self): _insert_task(self.db_path, "T1", plan=1) worker = self._worker(routing_policy=POLICY) results = worker.run_all() self.assertEqual(results[0]["status"], "completed") self.assertEqual(results[0]["actual_model"], "gpt-6.1-sol") self.assertEqual(self._dispatch_model(), "gpt-6.1-sol") event = results[0]["events"][0] self.assertEqual(event["type"], "model_routing") self.assertEqual(event["tier"], "plan") self.assertEqual(event["source"], "policy") self.assertTrue(event["applied"]) def test_run_all_shadow_does_not_change_dispatch(self): _insert_task(self.db_path, "T1", plan=1) worker = self._worker(routing_policy=_policy(mode="shadow")) results = worker.run_all() self.assertEqual(results[0]["actual_model"], "gpt-5.6-terra") self.assertEqual(self._dispatch_model(), "gpt-5.6-terra") event = results[0]["events"][0] self.assertFalse(event["applied"]) self.assertEqual(event["suggested_model"], "gpt-6.1-sol") def test_run_all_explicit_task_model_has_no_routing_event(self): _insert_task(self.db_path, "T1", plan=1, model="gpt-6-sol") worker = self._worker(routing_policy=POLICY) results = worker.run_all() self.assertEqual(results[0]["actual_model"], "gpt-6-sol") self.assertEqual(self._dispatch_model(), "gpt-6-sol") self.assertEqual( [e for e in results[0]["events"] if e["type"] == "model_routing"], []) if __name__ == "__main__": unittest.main()