Songhaoz666andClaude Opus 4.8
e4a120b715
回应 @Mem0ried 在 PR #37 的对接问题：除逐条 budget.alert 事件的 cost_phase 外，详情 聚合也需能拆 initial/review_retry。build_runtime_metrics 新增 cost_by_phase 滚动汇总 （按 swarm_id，从各任务 usage.model_cost_usd/model_tokens 按 rework_attributions 归类）， GET …/{id}/metrics 返回 {initial,review_retry}:{cost_usd,model_tokens}。HM 详情聚合 / 客户端 08 用量抽屉可直接展示拆分，无需从事件推导。 测试 test-usage-cost-phase.py 增 /metrics 断言（reworked→review_retry、非 reworked→initial）； usage-billing §5 记录该字段。 Refs #16 Refs #37 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
102 lines
4.9 KiB
Python
102 lines
4.9 KiB
Python
"""Usage review-retry cost attribution test (issue #16).
|
|
|
|
Verifies emit_usage_event tags each usage event with `cost_phase`:
|
|
* "initial" before cross-review has reopened the task for rework;
|
|
* "review_retry" once the task appears in run.metadata["rework_attributions"]
|
|
(i.e. its re-execution cost is attributable as review-retry cost).
|
|
|
|
Hermetic: REDIS_FAKE, no model key, no callback url.
|
|
|
|
Run from agent_swarm_v6 (install deps first):
|
|
pip install -r orchestrator/requirements.txt
|
|
REDIS_FAKE=1 python scripts/test-usage-cost-phase.py
|
|
"""
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
os.environ["REDIS_FAKE"] = "1"
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
|
|
failures = []
|
|
|
|
|
|
def check(name, cond):
|
|
print(("PASS" if cond else "FAIL"), "-", name)
|
|
if not cond:
|
|
failures.append(name)
|
|
|
|
|
|
async def budget_events(swarm_runtime, swarm_id):
|
|
raw = await swarm_runtime.list_events(swarm_id, limit=500)
|
|
return [e for e in raw["events"] if e["event_type"] == "budget.alert"]
|
|
|
|
|
|
async def main():
|
|
from orchestrator.redis_client import redis_client
|
|
from orchestrator.swarm_runtime import swarm_runtime
|
|
from orchestrator.task_queue import task_queue
|
|
from orchestrator import main as orch
|
|
|
|
await redis_client.connect()
|
|
body = {"mode": "swarm", "orchestration_plan": {"objective": "cost phase"},
|
|
"callback": {"url": "", "subscribed_events": []},
|
|
"metadata": {"manager_deployment_id": "m-cost"}}
|
|
run, _ = await swarm_runtime.get_or_create_run(body=body, idempotency_key=None, correlation_id="c-cost")
|
|
t1 = await task_queue.create_task(task_id=f"{run.swarm_id}-t1", description="impl",
|
|
agent_role="impl", required_capabilities=["python"], enqueue=True)
|
|
await swarm_runtime.attach_task(run, t1.task_id)
|
|
|
|
usage_result = {"usage": {"model_id": "gpt-x", "model_tokens": 100, "model_cost_usd": 0.01,
|
|
"billing_source": "newapi"}}
|
|
|
|
# 1) Initial execution: no rework attribution yet → cost_phase=initial.
|
|
await orch.emit_usage_event(run, t1, "agent-1", usage_result)
|
|
evs = await budget_events(swarm_runtime, run.swarm_id)
|
|
check("initial usage tagged cost_phase=initial", evs and evs[-1]["payload"].get("cost_phase") == "initial")
|
|
check("usage payload carries attempt", "attempt" in evs[-1]["payload"])
|
|
|
|
# 2) Cross-review reopened t1 for rework → recorded in rework_attributions.
|
|
run.metadata["rework_attributions"] = [{"target_task_id": t1.task_id, "root_cause": "test"}]
|
|
await swarm_runtime.save_run(run)
|
|
t1.retry_count = 1
|
|
await orch.emit_usage_event(run, t1, "agent-1", usage_result)
|
|
evs = await budget_events(swarm_runtime, run.swarm_id)
|
|
check("redo usage tagged cost_phase=review_retry", evs[-1]["payload"].get("cost_phase") == "review_retry")
|
|
check("review_retry usage attributable by swarm_id + cost_phase",
|
|
any(e["payload"].get("cost_phase") == "review_retry" and e["payload"].get("swarm_id") == run.swarm_id
|
|
for e in evs))
|
|
|
|
# A non-reworked task stays initial even alongside the reworked one.
|
|
t2 = await task_queue.create_task(task_id=f"{run.swarm_id}-t2", description="doc",
|
|
agent_role="documentation", required_capabilities=["technical-writing"], enqueue=True)
|
|
await swarm_runtime.attach_task(run, t2.task_id)
|
|
await orch.emit_usage_event(run, t2, "agent-2", usage_result)
|
|
evs = await budget_events(swarm_runtime, run.swarm_id)
|
|
t2_ev = [e for e in evs if e["payload"].get("task_id") == t2.task_id][-1]
|
|
check("non-reworked task stays cost_phase=initial", t2_ev["payload"].get("cost_phase") == "initial")
|
|
|
|
# --- /metrics run-level cost_by_phase rollup (#37) ---
|
|
# Persist task results with usage so build_runtime_metrics can split them; t1 is a rework target.
|
|
await task_queue.complete_task(t1.task_id, result=json.dumps({"usage": {"model_cost_usd": 0.05, "model_tokens": 500}}))
|
|
await task_queue.complete_task(t2.task_id, result=json.dumps({"usage": {"model_cost_usd": 0.02, "model_tokens": 200}}))
|
|
refreshed = await swarm_runtime.get_run(run.swarm_id)
|
|
metrics = await orch.build_runtime_metrics(refreshed, "15m", "60s")
|
|
cbp = metrics.get("cost_by_phase") or {}
|
|
check("/metrics exposes cost_by_phase", set(cbp.keys()) == {"initial", "review_retry"})
|
|
check("/metrics review_retry cost = reworked task (t1=0.05)", abs(cbp["review_retry"]["cost_usd"] - 0.05) < 1e-9)
|
|
check("/metrics initial cost = non-reworked task (t2=0.02)", abs(cbp["initial"]["cost_usd"] - 0.02) < 1e-9)
|
|
check("/metrics review_retry tokens = 500", cbp["review_retry"]["model_tokens"] == 500)
|
|
|
|
print()
|
|
if failures:
|
|
print(f"{len(failures)} usage cost-phase check(s) FAILED: {failures}")
|
|
sys.exit(1)
|
|
print("all usage cost-phase checks passed")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|