host
63823fc2e3
此前每任务只记一条 estimated（字符数/4，金额为空），注释称 host 协议不暴露真实 token 数。本轮实测推翻了该前提： - account/usage/read 在中继上是死的（要求 ChatGPT 账号登录）， 不是来源，已在 cost.py docstring 写明死胡同 - app-server 每上游补全发 thread/tokenUsage/updated，last 是该次 补全的用量；中继 5 调用回合实测 last 之和等于最后 total （input/output/cached/reasoning 四路都对得上） 实现：host.py run_turn 从事件窗口按 threadId+turnId 汇总用量， 重试/goal 各轮累加，dispatch/run JSON 与 goal 用量行透传； sandbox 三路解析（dispatch JSON 解析改用完整 stdout，顺带修 了长 text 下 failed 被误判 completed 的潜 bug）；worker 每任务 记一条 actual（重试/fallback 用量累加进同一条），无用量时回落 estimated；cost.py 价目表（默认 OpenAI 挂牌价，env 可覆盖， 无价模型 usd 记 None、报表进 unpriced_actual_tokens 而非 0 美元）， cost_records 加 cached_input_tokens 列（SQLite/PG 两支迁移）； 报表与按天行新增 cached_input_tokens/unpriced_actual_tokens。 测试：host 32、cloud 483 全过。
274 lines
11 KiB
Python
274 lines
11 KiB
Python
"""Tests for cloud.models (fallback routing) and cloud.cost.
|
|
|
|
No real model turns, no network, no secrets: usage clients are fakes
|
|
and the CostStore uses an in-memory SQLite database.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import unittest
|
|
from datetime import datetime, timezone
|
|
from unittest import mock
|
|
|
|
# cloud/ has no __init__.py by design (the coordinator creates it), so
|
|
# make the project root importable regardless of the discover top level.
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(
|
|
os.path.dirname(os.path.abspath(__file__)))))
|
|
|
|
from cloud.cost import (CostStore, estimate_tokens, fetch_account_usage,
|
|
price_usd_micros)
|
|
from cloud.models import FALLBACKABLE, ModelRouter, is_fallbackable, record_fallback
|
|
|
|
|
|
class FakeError(Exception):
|
|
def __init__(self, status=None, code=None):
|
|
super().__init__(status or code)
|
|
self.status = status
|
|
self.code = code
|
|
|
|
|
|
class TestIsFallbackable(unittest.TestCase):
|
|
def test_fallbackable_status_codes(self):
|
|
for code in (404, 429, 500, 502, 503, 504):
|
|
self.assertTrue(is_fallbackable(code), code)
|
|
self.assertIn(404, FALLBACKABLE)
|
|
self.assertIn(429, FALLBACKABLE)
|
|
self.assertIn("timeout", FALLBACKABLE)
|
|
|
|
def test_timeout_string(self):
|
|
self.assertTrue(is_fallbackable("timeout"))
|
|
|
|
def test_not_fallbackable(self):
|
|
for code in (400, 401, 403, 200):
|
|
self.assertFalse(is_fallbackable(code), code)
|
|
self.assertFalse(is_fallbackable("other"))
|
|
|
|
def test_object_with_status_or_code(self):
|
|
self.assertTrue(is_fallbackable(FakeError(status=429)))
|
|
self.assertTrue(is_fallbackable(FakeError(code=503)))
|
|
self.assertTrue(is_fallbackable(FakeError(status="timeout")))
|
|
self.assertFalse(is_fallbackable(FakeError(status=401)))
|
|
self.assertFalse(is_fallbackable(FakeError(code=400)))
|
|
|
|
|
|
class TestModelRouter(unittest.TestCase):
|
|
def setUp(self):
|
|
self.router = ModelRouter(["gpt-6.1-sol", "gpt-6-sol"])
|
|
|
|
def test_empty_priority_list_rejected(self):
|
|
with self.assertRaises(ValueError):
|
|
ModelRouter([])
|
|
|
|
def test_first(self):
|
|
self.assertEqual(self.router.first(), "gpt-6.1-sol")
|
|
|
|
def test_next_after_in_priority_order(self):
|
|
self.assertEqual(
|
|
self.router.next_after("gpt-6.1-sol", 404), "gpt-6-sol")
|
|
self.assertEqual(
|
|
self.router.next_after("gpt-6.1-sol", "timeout"), "gpt-6-sol")
|
|
|
|
def test_next_after_exhausted_returns_none(self):
|
|
self.assertIsNone(self.router.next_after("gpt-6-sol", 500))
|
|
self.assertIsNone(self.router.next_after("gpt-6-sol", 429))
|
|
|
|
def test_next_after_non_fallbackable_returns_none(self):
|
|
self.assertIsNone(self.router.next_after("gpt-6.1-sol", 401))
|
|
self.assertIsNone(self.router.next_after("gpt-6.1-sol", 400))
|
|
|
|
def test_route_walks_attempts(self):
|
|
self.assertEqual(
|
|
self.router.route([("gpt-6.1-sol", 404)]), "gpt-6-sol")
|
|
self.assertIsNone(self.router.route(
|
|
[("gpt-6.1-sol", 404), ("gpt-6-sol", 500)]))
|
|
self.assertIsNone(self.router.route([("gpt-6.1-sol", 401)]))
|
|
self.assertEqual(self.router.route([]), "gpt-6.1-sol")
|
|
|
|
|
|
class TestRecordFallback(unittest.TestCase):
|
|
def test_event_recorded_on_task_record(self):
|
|
task = {"id": "task-1"}
|
|
event = record_fallback(task, "gpt-6.1-sol", "gpt-6-sol", "429")
|
|
self.assertEqual(event, {
|
|
"type": "model_fallback",
|
|
"task_id": "task-1",
|
|
"from_model": "gpt-6.1-sol",
|
|
"to_model": "gpt-6-sol",
|
|
"reason": "429",
|
|
})
|
|
self.assertIn(event, task["events"])
|
|
|
|
def test_event_appended_to_existing_events(self):
|
|
task = {"id": "task-2", "events": [{"type": "started"}]}
|
|
record_fallback(task, "a", "b", "500")
|
|
self.assertEqual(len(task["events"]), 2)
|
|
self.assertEqual(task["events"][1]["type"], "model_fallback")
|
|
|
|
|
|
class TestFetchAccountUsage(unittest.TestCase):
|
|
def test_success_returns_available_true(self):
|
|
data = {"balance": 123, "currency": "USD"}
|
|
result = fetch_account_usage(lambda: data)
|
|
self.assertEqual(result, {"available": True, "data": data})
|
|
|
|
def test_raising_client_returns_unavailable_not_raise(self):
|
|
def boom():
|
|
raise RuntimeError("no route")
|
|
self.assertEqual(
|
|
fetch_account_usage(boom), {"available": False, "data": None})
|
|
|
|
def test_none_and_empty_dict_return_unavailable(self):
|
|
self.assertEqual(
|
|
fetch_account_usage(lambda: None),
|
|
{"available": False, "data": None})
|
|
self.assertEqual(
|
|
fetch_account_usage(lambda: {}),
|
|
{"available": False, "data": None})
|
|
|
|
|
|
class TestEstimateTokens(unittest.TestCase):
|
|
def test_chars_divided_by_four(self):
|
|
self.assertEqual(estimate_tokens("abcd", "efgh"), 2)
|
|
self.assertEqual(estimate_tokens("a" * 10, "b" * 7), 4) # 17 // 4
|
|
self.assertEqual(estimate_tokens("", ""), 0)
|
|
|
|
|
|
class TestCostStore(unittest.TestCase):
|
|
def setUp(self):
|
|
self.store = CostStore(":memory:")
|
|
|
|
def test_record_computes_total_and_day(self):
|
|
row_id = self.store.record(
|
|
"t1", "p1", "gpt-6.1-sol", input_tokens=100, output_tokens=50,
|
|
usd_micros=2000, source="actual", fallback_count=1)
|
|
self.assertIsInstance(row_id, int)
|
|
report = self.store.report_by_task("t1")
|
|
self.assertEqual(report["total_tokens"], 150)
|
|
self.assertEqual(report["usd_micros"], 2000)
|
|
self.assertEqual(report["turns"], 1)
|
|
self.assertEqual(report["fallback_count"], 1)
|
|
self.assertEqual(report["actual_tokens"], 150)
|
|
self.assertEqual(report["estimated_tokens"], 0)
|
|
|
|
def test_invalid_source_rejected(self):
|
|
with self.assertRaises(ValueError):
|
|
self.store.record("t1", "p1", "m", source="guess")
|
|
|
|
def test_actual_and_estimated_kept_separate(self):
|
|
self.store.record("t1", "p1", "m", input_tokens=100,
|
|
output_tokens=0, source="actual")
|
|
self.store.record_estimated("t1", "p1", "m",
|
|
"x" * 40, "y" * 40) # 80 // 4 = 20
|
|
report = self.store.report_by_task("t1")
|
|
self.assertEqual(report["total_tokens"], 120)
|
|
self.assertEqual(report["actual_tokens"], 100)
|
|
self.assertEqual(report["estimated_tokens"], 20)
|
|
self.assertEqual(report["turns"], 2)
|
|
|
|
def test_record_estimated_shape(self):
|
|
self.store.record_estimated("t9", "p9", "m", "abcd", "efgh")
|
|
report = self.store.report_by_task("t9")
|
|
self.assertEqual(report["estimated_tokens"], 2)
|
|
self.assertEqual(report["actual_tokens"], 0)
|
|
self.assertEqual(report["usd_micros"], 0)
|
|
|
|
def test_report_by_project_and_day(self):
|
|
self.store.record("t1", "p1", "m", input_tokens=10, output_tokens=0)
|
|
self.store.record("t2", "p1", "m", input_tokens=30, output_tokens=0)
|
|
self.store.record("t3", "p2", "m", input_tokens=999, output_tokens=0)
|
|
project = self.store.report_by_project("p1")
|
|
self.assertEqual(project["total_tokens"], 40)
|
|
self.assertEqual(project["turns"], 2)
|
|
today = datetime.now(timezone.utc).strftime("%Y-%m-%d")
|
|
day = self.store.report_by_day(today)
|
|
self.assertEqual(day["total_tokens"], 10 + 30 + 999)
|
|
self.assertEqual(day["turns"], 3)
|
|
self.assertEqual(self.store.report_by_day("2000-01-01")["turns"], 0)
|
|
|
|
def test_empty_report_is_zeroed(self):
|
|
report = self.store.report_by_task("missing")
|
|
self.assertEqual(report, {
|
|
"total_tokens": 0,
|
|
"usd_micros": 0,
|
|
"turns": 0,
|
|
"fallback_count": 0,
|
|
"actual_tokens": 0,
|
|
"estimated_tokens": 0,
|
|
"cached_input_tokens": 0,
|
|
"unpriced_actual_tokens": 0,
|
|
})
|
|
|
|
|
|
class TestModelPrices(unittest.TestCase):
|
|
def test_list_prices(self):
|
|
# 2026-10-26 中继实测 5 调用回合：(58243-45960)*2
|
|
# + 45960*0.10 + 1138*10 = 40542 微美元。
|
|
self.assertEqual(
|
|
price_usd_micros("gpt-6.1-sol", 58243, 45960, 1138), 40542)
|
|
self.assertEqual(
|
|
price_usd_micros("gpt-6-sol", 1000, 200, 100),
|
|
int(round((800 * 2.00 + 200 * 0.20 + 100 * 10.00))))
|
|
|
|
def test_unknown_model_unpriced(self):
|
|
self.assertIsNone(price_usd_micros("gpt-5.6-sol", 100, 0, 100))
|
|
self.assertIsNone(price_usd_micros(None, 100, 0, 100))
|
|
|
|
def test_cached_clamped_to_input(self):
|
|
self.assertEqual(
|
|
price_usd_micros("gpt-6.1-sol", 100, 500, 100),
|
|
price_usd_micros("gpt-6.1-sol", 100, 100, 100))
|
|
|
|
def test_env_override(self):
|
|
with mock.patch.dict(os.environ, {
|
|
"CLOUD_MODEL_PRICES_JSON":
|
|
'{"gpt-6.1-sol": [4.00, 0.20, 20.00]}'}):
|
|
self.assertEqual(
|
|
price_usd_micros("gpt-6.1-sol", 1000, 0, 0), 4000)
|
|
with mock.patch.dict(os.environ, {
|
|
"CLOUD_MODEL_PRICES_JSON": 'not json'}):
|
|
# 坏覆盖不炸，回落默认表。
|
|
self.assertEqual(
|
|
price_usd_micros("gpt-6.1-sol", 1000, 0, 0), 2000)
|
|
|
|
|
|
class TestCachedColumn(unittest.TestCase):
|
|
def test_record_keeps_cached_and_unpriced(self):
|
|
store = CostStore(":memory:")
|
|
store.record("t1", "p1", "gpt-6.1-sol", input_tokens=1000,
|
|
output_tokens=100, cached_input_tokens=400,
|
|
usd_micros=2240, source="actual")
|
|
store.record("t2", "p1", "gpt-5.6-sol", input_tokens=500,
|
|
output_tokens=50, cached_input_tokens=0,
|
|
usd_micros=None, source="actual")
|
|
report = store.report_by_project("p1")
|
|
self.assertEqual(report["cached_input_tokens"], 400)
|
|
self.assertEqual(report["unpriced_actual_tokens"], 550)
|
|
self.assertEqual(report["usd_micros"], 2240)
|
|
|
|
def test_migration_adds_cached_column_to_old_table(self):
|
|
import sqlite3
|
|
import tempfile
|
|
path = os.path.join(tempfile.mkdtemp(), "c.db")
|
|
conn = sqlite3.connect(path)
|
|
conn.execute(
|
|
"CREATE TABLE cost_records (id INTEGER PRIMARY KEY"
|
|
" AUTOINCREMENT, task_id TEXT, project_id TEXT, model TEXT,"
|
|
" input_tokens INTEGER, output_tokens INTEGER,"
|
|
" total_tokens INTEGER, usd_micros INTEGER, source TEXT,"
|
|
" fallback_count INTEGER DEFAULT 0, day TEXT,"
|
|
" created_at REAL)")
|
|
conn.commit()
|
|
conn.close()
|
|
store = CostStore(path)
|
|
cols = {row[1] for row in store._conn.execute(
|
|
"PRAGMA table_info(cost_records)").fetchall()}
|
|
self.assertIn("cached_input_tokens", cols)
|
|
store.record("t1", "p1", "m", input_tokens=10, output_tokens=5,
|
|
cached_input_tokens=3, source="actual")
|
|
self.assertEqual(store.report_by_task("t1")["cached_input_tokens"],
|
|
3)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|