191 lines
6.8 KiB
Python
191 lines
6.8 KiB
Python
#!/usr/bin/env python3
|
|
"""Automatic lifecycle, token, tool and context usage observation.
|
|
|
|
Hook modes: ``--event start|tool|stop``. All events are append-only and bind agent/session,
|
|
workflow, role and exact context-package SHA whenever those identities are available.
|
|
Missing usage fields are recorded honestly rather than estimated as actual usage.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import sys
|
|
import uuid
|
|
from datetime import datetime, timezone
|
|
|
|
import yaml
|
|
|
|
HERE = os.path.dirname(os.path.abspath(__file__))
|
|
ROOT = os.environ.get("CLAUDE_PROJECT_DIR") or os.path.dirname(os.path.dirname(HERE))
|
|
sys.path.insert(0, HERE)
|
|
import _workspace as W # noqa: E402
|
|
|
|
|
|
def _now():
|
|
return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
|
|
|
|
|
|
def _events_path():
|
|
return os.path.join(W.state_dir(), "usage-events.jsonl")
|
|
|
|
|
|
def _append(record):
|
|
path = _events_path()
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
record = {"usage-event-id": "use-" + uuid.uuid4().hex, "observed-at": _now(), **record}
|
|
with open(path, "a", encoding="utf-8") as handle:
|
|
try:
|
|
import fcntl
|
|
fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
|
|
except Exception:
|
|
pass
|
|
handle.write(json.dumps(record, ensure_ascii=False) + "\n")
|
|
handle.flush()
|
|
os.fsync(handle.fileno())
|
|
|
|
|
|
def _agent_id(payload):
|
|
return payload.get("agent_id") or payload.get("agentId") or payload.get("subagent_id")
|
|
|
|
|
|
def _registry(agent_id):
|
|
if not agent_id:
|
|
return None
|
|
path = os.path.join(W.state_dir(), "subagent-registry.jsonl")
|
|
if not os.path.exists(path):
|
|
return None
|
|
found = None
|
|
for line in open(path, encoding="utf-8"):
|
|
try:
|
|
row = json.loads(line)
|
|
except Exception:
|
|
continue
|
|
if row.get("agent_id") == agent_id:
|
|
found = row
|
|
return found
|
|
|
|
|
|
def _package(record):
|
|
if not record or not record.get("context_package"):
|
|
return None
|
|
path = str(record["context_package"])
|
|
path = path if os.path.isabs(path) else os.path.join(ROOT, path)
|
|
try:
|
|
return yaml.safe_load(open(path, encoding="utf-8")) or {}
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def _base(payload):
|
|
agent_id = _agent_id(payload)
|
|
record = _registry(agent_id) or {}
|
|
pkg = _package(record) or {}
|
|
return {
|
|
"agent-id": agent_id,
|
|
"session-id": payload.get("session_id") or payload.get("sessionId"),
|
|
"agent-type": record.get("agent_type") or payload.get("agent_type") or payload.get("agentType"),
|
|
"workflow-id": record.get("workflow_id") or pkg.get("workflow-id"),
|
|
"role-id": record.get("role") or pkg.get("target-role-agent"),
|
|
"tier": pkg.get("tier"),
|
|
"context-package": record.get("context_package"),
|
|
"context-package-sha256": record.get("context_package_sha256"),
|
|
}, pkg
|
|
|
|
|
|
def _usage(payload):
|
|
candidates = [payload.get("usage"), payload.get("token_usage"),
|
|
(payload.get("result") or {}).get("usage") if isinstance(payload.get("result"), dict) else None,
|
|
payload]
|
|
for value in candidates:
|
|
if not isinstance(value, dict):
|
|
continue
|
|
inp = value.get("input_tokens") if value.get("input_tokens") is not None else value.get("inputTokens")
|
|
out = value.get("output_tokens") if value.get("output_tokens") is not None else value.get("outputTokens")
|
|
cached = value.get("cache_read_input_tokens") or value.get("cacheReadInputTokens") or 0
|
|
if inp is not None or out is not None:
|
|
return int(inp or 0), int(out or 0), int(cached or 0)
|
|
return None
|
|
|
|
|
|
def _read_path(payload):
|
|
tool_input = payload.get("tool_input") or payload.get("toolInput") or {}
|
|
for key in ("file_path", "path", "notebook_path"):
|
|
if tool_input.get(key):
|
|
return str(tool_input[key])
|
|
return None
|
|
|
|
|
|
def _context_item(pkg, path):
|
|
if not path:
|
|
return None
|
|
wanted = os.path.realpath(path if os.path.isabs(path) else os.path.join(ROOT, path))
|
|
for item in pkg.get("must-read", []) or []:
|
|
if not isinstance(item, dict):
|
|
continue
|
|
uri = item.get("uri")
|
|
if not uri:
|
|
continue
|
|
resolved = os.path.realpath(uri if os.path.isabs(str(uri)) else os.path.join(ROOT, str(uri)))
|
|
if resolved == wanted:
|
|
estimate = item.get("estimated-tokens")
|
|
if estimate is None:
|
|
try:
|
|
estimate = max(1, os.path.getsize(resolved) // 4)
|
|
except OSError:
|
|
estimate = None
|
|
return {"context-id": item.get("context-id"), "uri": uri,
|
|
"estimated-tokens": estimate, "reason": item.get("reason")}
|
|
return None
|
|
|
|
|
|
def observe(event, payload):
|
|
base, pkg = _base(payload)
|
|
if event == "start":
|
|
_append({"event-type": "SubagentStarted", **base})
|
|
return
|
|
if event == "tool":
|
|
tool_name = payload.get("tool_name") or payload.get("toolName")
|
|
_append({"event-type": "ToolUsageObserved", **base, "tool-name": tool_name})
|
|
if tool_name in {"Read", "Grep", "Glob"}:
|
|
item = _context_item(pkg, _read_path(payload))
|
|
if item:
|
|
_append({"event-type": "ContextItemRead", **base, **item})
|
|
return
|
|
if event == "stop":
|
|
usage = _usage(payload)
|
|
record = {"event-type": "SubagentCompleted", **base, "usage-observed": bool(usage)}
|
|
if usage:
|
|
input_tokens, output_tokens, cached_tokens = usage
|
|
record.update({"input-tokens": input_tokens, "output-tokens": output_tokens,
|
|
"cache-read-input-tokens": cached_tokens,
|
|
"total-tokens": input_tokens + output_tokens})
|
|
_append(record)
|
|
if usage and base.get("workflow-id") and base.get("tier"):
|
|
try:
|
|
import token_ledger
|
|
token_ledger.log(base["workflow-id"], base.get("role-id") or "unknown",
|
|
usage[0] + usage[1], tier=base["tier"],
|
|
usage_source_id=f"subagent:{base.get('agent-id')}")
|
|
except ValueError as exc:
|
|
if "중복 usage-source-id" not in str(exc):
|
|
raise
|
|
|
|
|
|
def main():
|
|
event = "tool"
|
|
if "--event" in sys.argv and sys.argv.index("--event") + 1 < len(sys.argv):
|
|
event = sys.argv[sys.argv.index("--event") + 1]
|
|
try:
|
|
payload = json.loads(sys.stdin.read() or "{}")
|
|
if not isinstance(payload, dict):
|
|
payload = {}
|
|
observe(event, payload)
|
|
except Exception as exc:
|
|
# Observation must not stop delivery; absence stays visible as usage-observed=false/missing.
|
|
sys.stderr.write(f"[usage_observer] observation skipped: {exc}\n")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|