host
0f7162934b
列表行与详情共用同一口径：新抽 api/taskMeta.ts（首字/
总耗时/token 计算），TaskListSection 改 import 它（逻辑
原样不动），TaskDetailPanel 在状态行下方加一行：首字
方块 + 耗时 + token，title 里写口径。
详情的数据缺口一并补上：单任务接口 _get_task 此前走
_task_to_dict 默认 total_tokens=None，详情 token 只能
显示占位；现与列表一样用 _task_tokens_map 一次聚合
（单任务只多一条查询），详情直接出真数。
types.ts 里过期的“列表不回 token”注释同步更新。
新测试：单任务 GET /tasks/{id} 聚合回 total_tokens。
141 lines
5.2 KiB
Python
141 lines
5.2 KiB
Python
"""任务列表 total_tokens：一次聚合，不逐个调成本接口。
|
|
|
|
用真实 HTTP 服务 + 临时 SQLite：建两个任务，给其中一个记
|
|
两条成本记录（actual + estimated），列表应一次回两行的
|
|
total_tokens（有记录的加总、无记录的 0），且不回 404/500；
|
|
_task_tokens_map 在 cost_records 表不存在时也不炸列表。
|
|
"""
|
|
|
|
import http.client
|
|
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
try:
|
|
from cloud.auth import AuthStore
|
|
from cloud.cost import CostStore
|
|
from cloud.server import (
|
|
_task_to_dict, _task_tokens_map, create_server)
|
|
except ImportError:
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
|
from cloud.auth import AuthStore
|
|
from cloud.cost import CostStore
|
|
from cloud.server import (
|
|
_task_to_dict, _task_tokens_map, create_server)
|
|
|
|
|
|
class ListTokensTestCase(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_list_returns_aggregated_total_tokens(self):
|
|
pid = self.project["id"]
|
|
_, t1 = self._req("POST", "/tasks",
|
|
{"project_id": pid, "prompt": "有成本"},
|
|
self.owner_key)
|
|
_, t2 = self._req("POST", "/tasks",
|
|
{"project_id": pid, "prompt": "无成本"},
|
|
self.owner_key)
|
|
store = CostStore(self.db_path)
|
|
try:
|
|
store.record(t1["id"], pid, "m", input_tokens=100,
|
|
output_tokens=50, source="actual")
|
|
store.record(t1["id"], pid, "m", input_tokens=20,
|
|
output_tokens=0, source="estimated")
|
|
finally:
|
|
store._conn.close()
|
|
status, body = self._req(
|
|
"GET", "/tasks?project_id=%s" % pid, key=self.owner_key)
|
|
self.assertEqual(status, 200, body)
|
|
by_id = {t["id"]: t for t in body["tasks"]}
|
|
self.assertEqual(by_id[t1["id"]]["total_tokens"], 170)
|
|
self.assertEqual(by_id[t2["id"]]["total_tokens"], 0)
|
|
|
|
def test_single_task_returns_aggregated_total_tokens(self):
|
|
# 详情接口（_get_task）与列表口径一致：单任务也聚合回
|
|
# total_tokens，否则详情里的 token 栏只能显示占位。
|
|
pid = self.project["id"]
|
|
_, t1 = self._req("POST", "/tasks",
|
|
{"project_id": pid, "prompt": "详情成本"},
|
|
self.owner_key)
|
|
store = CostStore(self.db_path)
|
|
try:
|
|
store.record(t1["id"], pid, "m", input_tokens=60,
|
|
output_tokens=40, source="actual")
|
|
finally:
|
|
store._conn.close()
|
|
status, body = self._req(
|
|
"GET", "/tasks/%s" % t1["id"], key=self.owner_key)
|
|
self.assertEqual(status, 200, body)
|
|
self.assertEqual(body["total_tokens"], 100)
|
|
|
|
def test_tokens_map_missing_table_is_zero_not_error(self):
|
|
import sqlite3
|
|
conn = sqlite3.connect(":memory:")
|
|
try:
|
|
self.assertEqual(
|
|
_task_tokens_map(conn, ["t1", "t2"]),
|
|
{"t1": 0, "t2": 0})
|
|
finally:
|
|
conn.close()
|
|
|
|
def test_task_to_dict_default_total_tokens_is_none(self):
|
|
row = {"id": "t1", "project_id": "p1", "prompt": "p",
|
|
"status": "completed", "idempotency_key": None,
|
|
"depends_on": "[]", "priority": 0, "created_by": "u",
|
|
"created_at": 1.0, "updated_at": 2.0,
|
|
"last_heartbeat": None}
|
|
self.assertIsNone(_task_to_dict(row, [])["total_tokens"])
|
|
self.assertEqual(
|
|
_task_to_dict(row, [], 42)["total_tokens"], 42)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|