host
e364970b46
Pre-dispatch routing only fills in tasks with no explicit model, classified by structural signals only (plan > goal > default). Explicit task model and worker CLI --model win over policy; policy via CLOUD_ROUTING_POLICY_JSON/FILE + CLOUD_ROUTING_MODE. Shadow mode records a suggestion without changing dispatch. Only policy/shadow decisions emit a model_routing event, so the existing events contract (fallback tests) is unchanged. Fallback rules, effort/approval/sandbox/quota and provider are untouched. Tests: test_models_routing.py (20 tests).
327 lines
12 KiB
Python
327 lines
12 KiB
Python
"""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()
|