mirror of
https://github.com/andrewyng/openworker.git
synced 2026-11-04 13:27:11 +00:00
Rohit C Prasad
7db596dd80
The human's ticks on the staffing card are the worker's grant; approvals follow the lead's mode. A lead answers a worker's parked call; team and work-item gates keep their payload in the Inbox.
87 lines
4.1 KiB
Python
87 lines
4.1 KiB
Python
"""Token counting (connectors-across-machines spec §5): per-model totals persisted on the
|
|
session record, exposed on the sessions list, and one content-blind audit row per model
|
|
round-trip so exported token columns are honest. No dollars anywhere."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
|
|
import aisuite as ai
|
|
|
|
from coworker.engine import EventType, TurnEngine
|
|
from coworker.permissions import PermissionEngine
|
|
from coworker.sessions import SessionRecord, usage_totals
|
|
from coworker.tools import ToolRegistry
|
|
from tests.test_persona_connections import _mgr
|
|
from tests.test_token_usage import _UsageProvider
|
|
|
|
|
|
def test_usage_totals_folds_assistant_sidecars_by_model():
|
|
msgs = [
|
|
{"role": "user", "content": "hi"},
|
|
{"role": "assistant", "content": "a", "usage": {"model": "m1", "input": 10, "output": 2, "cache_read": 1, "cache_write": 0}},
|
|
{"role": "assistant", "content": "b", "usage": {"model": "m1", "input": 5, "output": 1, "cache_read": 0, "cache_write": 3}},
|
|
{"role": "assistant", "content": "c", "usage": {"model": "m2", "input": 7, "output": 7}},
|
|
{"role": "assistant", "content": "no sidecar"},
|
|
{"role": "assistant", "content": "junk", "usage": {"model": "m2", "input": "x", "output": -4}},
|
|
]
|
|
t = usage_totals(msgs)
|
|
assert t["m1"] == {"input": 15, "output": 3, "cache_read": 1, "cache_write": 3, "turns": 2}
|
|
assert t["m2"] == {"input": 7, "output": 7, "cache_read": 0, "cache_write": 0, "turns": 2}
|
|
assert usage_totals([]) == {}
|
|
|
|
|
|
def test_save_persists_totals_and_the_list_exposes_them(tmp_path, monkeypatch):
|
|
mgr = _mgr(tmp_path, monkeypatch)
|
|
engine = mgr.get_engine("s-usage", agent="cowork")
|
|
assert engine is not None
|
|
engine.messages.append({"role": "user", "content": "hi"})
|
|
engine.messages.append(
|
|
{"role": "assistant", "content": "a", "usage": {"model": "openai:gpt-5.5", "input": 100, "output": 20, "cache_read": 5, "cache_write": 0}}
|
|
)
|
|
engine.messages.append(
|
|
{"role": "assistant", "content": "b", "usage": {"model": "openai:gpt-5.5", "input": 50, "output": 10, "cache_read": 0, "cache_write": 0}}
|
|
)
|
|
mgr.save("s-usage", engine)
|
|
# Round-trips through the store column…
|
|
loaded = mgr.session_store.load("s-usage")
|
|
assert loaded is not None
|
|
assert loaded.usage == {"openai:gpt-5.5": {"input": 150, "output": 30, "cache_read": 5, "cache_write": 0, "turns": 2}}
|
|
# …and the list row carries it for the Team panel roll-up.
|
|
row = next(r for r in mgr.list_sessions() if r["session_id"] == "s-usage")
|
|
assert row["usage"]["openai:gpt-5.5"]["input"] == 150
|
|
# A re-save never double counts (folded from the transcript, not accumulated).
|
|
mgr.save("s-usage", engine)
|
|
assert mgr.session_store.load("s-usage").usage["openai:gpt-5.5"]["turns"] == 2
|
|
# A record saved without messages keeps whatever it carried.
|
|
mgr.session_store.save(SessionRecord(session_id="s-plain", workspace=str(tmp_path), model="m", mode="interactive"))
|
|
assert mgr.session_store.load("s-plain").usage == {}
|
|
|
|
|
|
def test_engine_audits_one_usage_row_per_round_trip(tmp_path):
|
|
rows: list[dict] = []
|
|
registry = ToolRegistry()
|
|
registry.register_all(ai.toolkits.files(root=str(tmp_path), allow_write=True))
|
|
engine = TurnEngine(
|
|
provider=_UsageProvider(),
|
|
registry=registry,
|
|
permissions=PermissionEngine(workspace_root=tmp_path),
|
|
model="gpt-5.5",
|
|
audit_sink=rows.append,
|
|
)
|
|
engine.audit_context = {"session_id": "s1", "agent": "cowork", "workspace": str(tmp_path)}
|
|
|
|
async def _collect():
|
|
return [ev async for ev in engine.run("hello")]
|
|
|
|
events = asyncio.run(_collect())
|
|
assert any(ev.type == EventType.ASSISTANT_MESSAGE for ev in events)
|
|
usage_rows = [r for r in rows if r.get("stage") == "usage"]
|
|
assert len(usage_rows) == 1
|
|
r = usage_rows[0]
|
|
assert r["session_id"] == "s1" and r["tool"] == ""
|
|
assert r["arguments"] == {"model": "gpt-5.5"}
|
|
assert (r["tokens_in"], r["tokens_out"], r["cache_read"], r["cache_write"]) == (100, 20, 5, 0)
|
|
# Content-blind: no message text rides the row.
|
|
assert "hello" not in repr(r) and "done" not in repr(r)
|