host
9eedee5a3d
_list_tasks 在取附件的同时用 _task_tokens_map 一次查 cost_records 按 task_id 求和（GROUP BY），一屏只多一条 查询，不按任务逐个调成本接口。无记录计 0、表不存在时 全计 0 且列表不失败；_task_to_dict 新增 total_tokens 参数，默认 None（非列表调用方不把“未算”写成 0）。 前端 TaskListSection 的 token 栏位此前已接好 total_tokens， 补上后列表行直接显示真数，不再显示 token —；前端注释 与占位文案同步更新。 新测试 test_task_list_tokens.py 3 项：HTTP 列表聚合 （actual+estimated 加总、无记录为 0）、缺表不炸、 _task_to_dict 默认 None。
123 lines
4.4 KiB
Python
123 lines
4.4 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_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()
|