From 8023c2802e40f1bc635f83273bcbd01b1cd2cf8d Mon Sep 17 00:00:00 2001 From: Shashank Shekhar Singh <123410790+Shashankss1205@users.noreply.github.com> Date: Fri, 9 Oct 2026 22:59:06 +0530 Subject: [PATCH] feat(observe): attribute recorded runs and costs to explicit tenants Propagate optional tenant labels through runtime, planner, agent and server traces. Add tenant cost aggregation and CLI reports while preserving unlabelled legacy traces, separating recorded prices from estimates, and counting planner spend. --- README.md | 2 + ROADMAP.md | 15 +- docs/cookbook/06-serving-and-ops.md | 77 ++++- docs/deep-dive.md | 8 +- grapharc/cli/agent.py | 2 +- grapharc/cli/graphrun.py | 6 +- grapharc/cli/main.py | 57 ++++ grapharc/cli/plan.py | 6 +- grapharc/harness/agent.py | 10 +- grapharc/observe/__init__.py | 4 + grapharc/observe/cost.py | 95 +++++- grapharc/observe/replay.py | 2 +- grapharc/observe/trace.py | 19 +- grapharc/planner/admission.py | 1 + grapharc/planner/loop.py | 10 +- grapharc/planner/proposal.py | 1 + grapharc/runtime/graph.py | 34 ++- grapharc/server/runtime.py | 8 +- tests/test_cli_style.py | 1 + tests/test_tenant_cost.py | 447 ++++++++++++++++++++++++++++ 20 files changed, 759 insertions(+), 46 deletions(-) create mode 100644 tests/test_tenant_cost.py diff --git a/README.md b/README.md index 3852a7d..3a8bd80 100644 --- a/README.md +++ b/README.md @@ -70,6 +70,8 @@ print(g.compile().invoke({"question": "meaning of life"})) Typed state, declared writes, a budget — none optional. The rest of the surface (agents, tools, memory, sessions, policy documents, cost attribution, OTel) is in the [deep dive](docs/deep-dive.md) and the [cookbook](docs/cookbook/). +Cost attribution also supports explicit tenants: pass `tenant="team-a"` to a compiled graph or governed loop, then read its bill with `attribute_tenant(trace, "team-a", rates=...)` or `grapharc cost trace.jsonl --tenant team-a --json`. `run`, `plan`, `go` and policy-governed `agent` runs carry an explicit `--tenant`, `GRAPHARC_TENANT` or config value into the trace. Recorded USD and rate-card estimates stay separate; old unlabelled runs remain unassigned. See [the cost guide](docs/cookbook/06-serving-and-ops.md#where-did-the-tokens-go) for the boundaries. + ## The admission gate You cannot pre-author a graph for "investigate this incident" — the shape is discovered while working. So the graph is proposed at runtime, and a deterministic checker stands between proposing and running: registry, policy, remaining budget, depth, acyclicity, all on every proposal. A rejection is structured feedback the planner replans against; work discovered mid-run re-enters the same gate. Watch it refuse (free, scripted, the first proposal names a policy-denied kind): diff --git a/ROADMAP.md b/ROADMAP.md index 93ca052..90cc39a 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -451,7 +451,7 @@ allowances; the agent records document-driven denials and approval decisions. importing the module needs no OTel. Previously documented as unverified against the real SDK — verified during this pass against `opentelemetry-sdk` 1.44.0, with spans reaching an `InMemorySpanExporter`. -- [~] **10.4 — Cost attribution** per run, thread (session) and node. The +- [x] **10.4 — Cost attribution** per run, thread (session), node and explicit tenant. The price is now recorded, not guessed: every gateway publishes a `cost_usd` through the same `llm_output` envelope, the runtime's usage callback accumulates it per node and writes it onto the `end` event, and @@ -461,9 +461,16 @@ allowances; the agent records document-driven denials and approval decisions. figure wins outright rather than being averaged with a guess. Tokens are counted from the same events `metrics.summarize` uses — node `end` events *plus* work outside any node span, which is what a `grapharc agent` run - consists of entirely — and the suite asserts the two agree. **One gap - left:** there is no tenant on a trace event, so tenant attribution is not - offered rather than being approximated. + consists of entirely — and the suite asserts the two agree. An optional + tenant on `RunContext` labels every native graph, planner and agent event; + the loop carries it through round sub-runs. `attribute_tenant` and + `grapharc cost --tenant NAME` select whole labelled runs, refuse conflicting + labels, and leave old unlabelled runs unassigned. Planning tokens without + a recorded price are estimated or marked unpriced, including in tenant + totals. **Producer limit:** planner turns and failed node terminal events + still omit the provider price ([#138](https://github.com/CodeGraphContext/GraphARC/issues/138)); + a rate-card estimate is not a recorded bill. + Tenant labels provide attribution, not authorization. - [ ] **10.2 — Rollback** and versioned graph/prompt configs. - [ ] **10.5 — Alerting** on budget, failure, and verifier-drift. diff --git a/docs/cookbook/06-serving-and-ops.md b/docs/cookbook/06-serving-and-ops.md index a6d9cf8..332d4df 100644 --- a/docs/cookbook/06-serving-and-ops.md +++ b/docs/cookbook/06-serving-and-ops.md @@ -1570,7 +1570,9 @@ from pathlib import Path from langchain_core.messages import HumanMessage -from grapharc.observe import RateCard, TraceRecorder, attribute, summarize, to_mermaid +from grapharc.observe import ( + RateCard, TraceRecorder, attribute, attribute_tenant, summarize, to_mermaid, +) from grapharc.runtime.graph import END, START, GraphARC from grapharc.runtime.state import GraphARCState from grapharc.testing import ScriptedChatModel @@ -1612,6 +1614,12 @@ metrics = summarize(trace, "r1") print("metrics :", metrics.nodes_executed, "nodes,", metrics.tokens, "tokens,", metrics.per_node) print() print(to_mermaid(trace, "r1")) + +# A second run has an explicit tenant; the old unlabelled run stays excluded. +g.compile().invoke({"question": "budgets?"}, run_id="r2", tenant="team-a") +bill = attribute_tenant(trace, "team-a", rates=RateCard(default=3.0)) +assert [run.run_id for run in bill.runs] == ["r2"] +assert bill.tokens == attribute(trace, "r2").tokens ``` ``` @@ -1634,24 +1642,25 @@ flowchart TD ``` `RunCost.tokens` and `RunMetrics.tokens` agree by construction — both count the -`end` events of node executions, and the test suite asserts they match, because +`end` and `error` events plus work outside node spans, and the suite asserts they match, because a cost report and an audit trail that disagree are worse than either alone. **Three honest limits:** -- **Nothing in GraphARC writes `cost_usd` onto a trace event today.** The field - exists and `TraceRecorder.event` accepts it, but the kernel does not pass it, - so `recorded_cost_usd` is always `None` on a trace from today's runtime. - Everything you see above is `estimated_cost_usd` — tokens × your rate card, - reported in its own field so nobody mistakes an estimate for an invoice. +- **Recorded cost depends on the producer.** The kernel writes provider prices + onto successful node terminal events, and agents record their per-call prices. + Planner turns and failed node terminal events currently omit the provider + price. The example above uses a backend that reports no price, so its money + is `estimated_cost_usd` — tokens × your rate card. Recorded and estimated + figures remain separate; a recorded figure takes precedence over an estimate. - **`complete` is what tells you the total is a total.** Tokens with neither a recorded cost nor a matching rate are counted in `unpriced_tokens`, and a non-zero count means `cost_usd` is a lower bound. A `RateCard` with no `default` and no matching model prices nothing. -- **A node that raised has no token count at node level.** No `end` event was - written. Where it was an `AgentNode`, its per-call sub-events still hold the - spend, reported as `tokens_before_error` — kept out of the total so the total - keeps matching `metrics`. +- **Older error events may have no token count.** Current node wrappers record + tokens on both `end` and `error`. When an older or hand-built trace omits the + failed node's count, an agent's per-call sub-events may still hold the spend, + reported as `tokens_before_error` and kept out of the total that matches `metrics`. `to_mermaid` renders the graph's *declared topology* — the `topology` event every run now writes — with execution status overlaid per node: `done`, `running`, @@ -1666,11 +1675,54 @@ renderer that speaks Mermaid. `attribute_thread(trace, thread_id)` is the same for a whole session across resumes, and `by_node(trace)` ranks every node in a file by cost. +For tenant attribution, supply one label at the run entry point, as the example +above does with `tenant="team-a"`. `attribute_tenant` selects its second run, +`r2`, while the first unlabelled run stays excluded. + +Every compiled execution entry (`invoke`, `stream`, `ainvoke`, `astream`, +`astream_events`) accepts `tenant=`, as does `GovernedLoop.run`. The loop carries +it through planning, admission, approval, round execution and stop events; +`AgentNode` inherits its supplied `RunContext.tenant`. Custom producers must +label their own events consistently. Attribution is a caller-supplied label, +not an authorization check; it does not apply a policy on its own. + +`attribute_tenant(trace, name, rates=...)` reads one snapshot and returns +`TenantCost`: the matching runs, tokens, duration, recorded USD, estimated USD, +combined known USD, unpriced tokens and `complete`. It prices the same whole +runs as `attribute`, so model substeps do not add a second copy of node spend. +An unknown tenant raises `ReplayError`. Any run with inconsistent labels, +including a mixture of labelled and unlabelled events, makes the reader refuse +the report rather than guess a bill. Legacy runs with no tenant remain +unassigned, and writing an unlabelled run adds no tenant key. + +The CLI uses the same reader: + +```bash +grapharc cost trace.jsonl --tenant team-a --json +grapharc cost trace.jsonl --run-id r2 +``` + +Choose exactly one scope. JSON includes separate `recorded_cost_usd` and +`estimated_cost_usd`, `unpriced_tokens`, `complete`, and per-run detail for a +tenant report. The CLI supplies no rate card, so tokens without recorded prices +remain unpriced; use the Python API for estimates. Missing or unreadable files +exit 2, and missing selections or inconsistent labels exit 1, including in JSON. + +`run`, `plan`, `go` and policy-governed `agent` runs carry a tenant supplied by +`--tenant`, `GRAPHARC_TENANT` +or `grapharc.toml` into their traces. The policy layer's implicit `"default"` +does not label a trace; an explicit `--tenant default` does. Executing a saved +plan uses the current invocation's explicit tenant, rather than inheriting a +label from the planning run. The standalone `grapharc agent` command +currently requires `--policy` when selecting a tenant; its trace uses the explicit +selection, while the document's implicit default tenant stays unlabelled. + + --- ## The CLI tour -Twelve commands. Every one takes `--json`, which prints the same payload as one +The commands below all take `--json`, which prints the same payload as one document on stdout — including failures, which become the document rather than a line on stderr. @@ -1687,6 +1739,7 @@ line on stderr. | `grapharc diff ` | compare two runs in one trace | | `grapharc trace ` | pretty-print a trace file | | `grapharc metrics ` | summarize one run | +| `grapharc cost --tenant NAME` | attribute whole labelled runs; alternatively select `--run-id ID` | | `grapharc viz ` | render the executed path as Mermaid | `run --check-only` prints a `topology` digest for comparing the same graph diff --git a/docs/deep-dive.md b/docs/deep-dive.md index d34c1e5..cfba7da 100644 --- a/docs/deep-dive.md +++ b/docs/deep-dive.md @@ -94,7 +94,7 @@ OpenRouter also carries routing: model-level `fallback_models` chains, provider **Cost ceilings are enforced, on both sides of a call.** A `SpendMeter` refuses before a call once the ceiling is reached, and after a call it charges the cost and *then* raises if that crossed the line — so overspend is bounded by the single call that crossed it rather than discovered a node later. The limit worth stating: a call whose cost the provider does not report cannot be charged, and those land in `unpriced_calls` rather than being guessed at. `unpriced_calls > 0` means the ceiling saw less than the whole bill. -Caveats each backend accepts openly. **Claude CLI:** `bind_tools` and `with_structured_output` raise `NotImplementedError` — the adapter implements neither, so what you get is LangChain's `BaseChatModel` default, and `claude -p` in the tool-free mode GraphARC drives it in offers nothing to implement them with. GraphARC also has no cache control on this path — the CLI decides and the usage envelope reports what it did — and calls spend subscription quota. **OpenRouter:** credit is reserved against `max_tokens`, so the default is deliberately modest. **OpenAI:** the API returns token counts and no price, so `cost_usd` is `None` and a dollar ceiling counts calls instead of enforcing — pass `price_per_million=` or price the trace afterwards with `observe.cost.RateCard`; token budgets are unaffected. **Ollama:** free by definition, so calls are charged `0.0` rather than counted as unpriced, and whether tool-calling works depends on the model you pulled rather than on the adapter. **All of them:** the per-call `cost_usd` is captured, budgeted against, *and* written onto the trace, so `observe.cost` reports a recorded figure rather than an estimate whenever the provider gave one. See [ROADMAP.md](../ROADMAP.md) §10.4 for what is still missing (a tenant on the event). +Caveats each backend accepts openly. **Claude CLI:** `bind_tools` and `with_structured_output` raise `NotImplementedError` — the adapter implements neither, so what you get is LangChain's `BaseChatModel` default, and `claude -p` in the tool-free mode GraphARC drives it in offers nothing to implement them with. GraphARC also has no cache control on this path — the CLI decides and the usage envelope reports what it did — and calls spend subscription quota. **OpenRouter:** credit is reserved against `max_tokens`, so the default is deliberately modest. **OpenAI:** the API returns token counts and no price, so `cost_usd` is `None` and a dollar ceiling counts calls instead of enforcing — pass `price_per_million=` or price the trace afterwards with `observe.cost.RateCard`; token budgets are unaffected. **Ollama:** free by definition, so calls are charged `0.0` rather than counted as unpriced, and whether tool-calling works depends on the model you pulled rather than on the adapter. **All of them:** a reported per-call `cost_usd` is captured and budgeted against. Successful node events and agent model events carry it into the trace; planner turns and failed node terminal events currently omit it. See [ROADMAP.md](../ROADMAP.md) §10.4 for explicit tenant attribution and producer pricing limits. ## Tools and the harness @@ -172,7 +172,7 @@ grapharc metrics trace.jsonl # tokens, retries, termination rea **Spans are optional by construction.** One root span per run, one child per node execution, with `AgentNode` sub-steps parented by inference — and a sub-step whose parent cannot be identified is parented to the run span rather than to a guess. The OpenTelemetry dependency is confined behind a Protocol, so importing the module needs no OTel installed. This was carried as unverified against the real SDK for a while; it has now been run against `opentelemetry-sdk` 1.44.0 with spans arriving at a real exporter. -**Cost attribution is per run, thread and node**, and it distinguishes what was measured from what was guessed. Tokens are counted from the same events `metrics` uses — node `end` events plus work that happened outside any node span, which is what a `grapharc agent` run is entirely made of — and the suite asserts the two agree, so a cost report and an audit trail cannot drift apart. The provider's own `cost_usd` is written onto the trace, so `recorded_cost_usd` holds a real figure when the backend reported one; a backend that reports none falls back to tokens priced against a `RateCard` you supply, and the two never mix. There is still no tenant on a trace event, so per-tenant attribution is not offered rather than being approximated. +**Cost attribution is per run, thread, node and explicit tenant**, and it distinguishes what was measured from what was guessed. Tokens are counted from the same events `metrics` uses — node `end` events plus work that happened outside any node span, which is what a `grapharc agent` run is entirely made of — and the suite asserts the two agree, so a cost report and an audit trail cannot drift apart. Where a producer records the provider's own `cost_usd` on the trace, `recorded_cost_usd` holds that real figure; a backend that reports none falls back to tokens priced against a `RateCard` you supply, and the two never mix. `RunContext.tenant` explicitly labels graph, planner and agent events, including round sub-runs. `attribute_tenant` selects whole runs from the trace and refuses conflicting labels; old unlabelled traces stay unassigned. The CLI exposes the same selection through `grapharc cost --tenant NAME --json`. Planner tokens without a recorded price are estimated from a supplied rate card or marked unpriced. ## Configuration, and the zero-config path @@ -248,7 +248,7 @@ A stable system is not one that claims to have no edges — it is one whose edge - **`run_command` is not confined.** Argv-only and never a shell, but the child is an ordinary process with your privileges. - **`interrupt()` suspends but cannot be resumed.** LangGraph's native interrupt stops the graph and shows on `get_state`, and there is no supported resume path — resuming means passing a `Command` as *input*, which is closed by design. Use the session layer's approval gate for human-in-the-loop. - **Still unwrapped from LangGraph:** `retry_policy`, `cache_policy`, `durability`, subgraphs. `.inner` reaches them, but execution entry points there fail closed, so `.inner` is an inspection escape hatch and not a way to run the graph. -- **Cost is recorded when a backend reports one, estimated when it does not.** Both gateways publish the provider's `cost_usd` through the same `llm_output` envelope, the runtime's usage callback writes it onto the node's `end` event, and an agent's `model` events carry the per-call breakdown. A backend that reports no price still falls back to a `RateCard` estimate, and the two figures stay apart — `recorded_cost_usd` is never a guess. Still missing: no tenant on a trace event, so per-tenant attribution is not offered. +- **Cost is recorded when a backend reports one, estimated when it does not.** Both gateways publish the provider's `cost_usd` through the same `llm_output` envelope, the runtime's usage callback writes it onto the node's `end` event, and an agent's `model` events carry the per-call breakdown. A backend that reports no price still falls back to a `RateCard` estimate, and the two figures stay apart — `recorded_cost_usd` is never a guess. Explicit tenants are supported by `attribute_tenant` and `grapharc cost`; legacy unlabelled runs are excluded from named tenant bills. Planner turns and failed node terminal events still omit the provider price, so their tokens need a rate card or remain unpriced. - **A node's tokens are its own, not the run's movement while it ran.** Worth stating because it was the other way round: an `end` event carried the difference between two readings of the run's *shared* meter, so under fan-out the workers' windows overlapped and each was credited with its siblings' concurrent spend. Three workers costing 8 tokens each traced as 24/16/8, and `metrics` and `cost` agreed on 48 for 24 tokens of real work — doubling the estimated bill purely because the work ran in parallel. Attribution now comes from a per-node scope on the meter, so the same work costs the same serially and in parallel; a hand charge the usage callback never saw still lands on the node that made it. - **A planning round is an envelope, not a measurement.** A `round` event used to carry the planner's `tokens` and the round's `duration_ms`, both of which `metrics`, `cost` and `replay` add on top of node totals — and the planner's spend was already reported by its own `plan` event, so it was counted twice, and a round's duration encloses the plan plus every node it ran. Neither is on the event now; both are on its `state_delta` as `round_tokens` / `round_iterations` / `round_duration_ms`, where no reader sums them. `RoundRecord.iterations` also holds a figure now rather than always `0`. - **The Claude CLI backend is completion-only, and an agent node on it is *delegated* rather than governed.** The CLI has no tool-calling wire format, so GraphARC cannot run its own gated loop over it. Rather than refuse, `AgentNode` hands the whole loop to Claude Code's headless agent, in one of two named tiers. `allowlist`, the default, pre-approves exactly the Claude Code twins of the node's own registered tools (`read_file`→`Read`, … `run_command`→`Bash`), so one operator declaration governs both the governed loop and the delegated one — but the enforcement is Claude Code's own gating, not this graph's per-call policy, there are no per-tool trace events, and anything unlisted falls to headless default gating, which fails closed. `bypass`, explicit opt-in only, runs `bypassPermissions`: every tool Claude Code has, no checks at all. In either tier the calls are not confined by the sandbox executor and the token figure is the sub-agent's own rather than one GraphARC metered call by call; a `--max-tokens` the delegated path cannot enforce is refused rather than silently unapplied. The workspace boundary and the wall-clock ceiling still hold — the CLI runs in its own session, and the deadline kills the whole process group, not just the direct child. It warns on `DelegatedToolUseWarning` at construction, naming the tier, and marks every trace event `executor=delegated` with its `delegated_mode`; filter that warning to an error to get the old refusal back. Structured output still needs an OpenAI-wire backend: `openrouter`, `openai`, or a local `ollama`. @@ -256,7 +256,7 @@ A stable system is not one that claims to have no edges — it is one whose edge - **`.env` and `grapharc.toml` follow the same discovery rule: the working directory, and nowhere else.** Neither searches parent directories — a run must not be governed by a file you did not know about, and must not be *billed* to one either. **This is a behaviour change:** the credential loader used to walk up to `/`, so a `.env` in an ancestor directory (a `$HOME` one on a shared box, a client project one above a demo checkout) was picked up silently. If you relied on that, move the file into the directory you run from, `export` the variable, or pass `env_file=` to name it explicitly. A real environment variable still beats any file. - **`grapharc run` has no budget unless you give it one.** Set any of `--max-tokens`, `--max-iterations`, `--max-seconds`, or `--max-concurrency`; without them each dimension is unlimited and the gate admits a topology of any worst-case cost. -**Verified this pass:** `pytest` → green, 2,305 selected and 13 deselected (the live ones); `ruff check .` clean; all eight `grapharc demo` stages green, plus the `trace` / `metrics` / `viz` / `replay` tour against a freshly recorded demo trace; the wheel builds and imports all submodules in a clean virtualenv with `[all]`, and `0.1.8` on PyPI is that wheel. The counts are a snapshot, not a property of the project — `pytest` re-derives them in one command, which is the only reason they are quoted, and `tests/test_deep_dive.py` fails this line rather than letting it drift. +**Verified this pass:** `pytest` → green, 2,354 selected and 13 deselected (the live ones); `ruff check .` clean; all eight `grapharc demo` stages green, plus the `trace` / `metrics` / `viz` / `replay` tour against a freshly recorded demo trace; the wheel builds and imports all submodules in a clean virtualenv with `[all]`, and `0.1.8` on PyPI is that wheel. The counts are a snapshot, not a property of the project — `pytest` re-derives them in one command, which is the only reason they are quoted, and `tests/test_deep_dive.py` fails this line rather than letting it drift. [ROADMAP.md](../ROADMAP.md) tracks what is built and what is not, item by item. diff --git a/grapharc/cli/agent.py b/grapharc/cli/agent.py index b033ae0..65e6360 100644 --- a/grapharc/cli/agent.py +++ b/grapharc/cli/agent.py @@ -502,7 +502,7 @@ def run_agent( max_seconds=max_seconds, ) ) - ctx = RunContext(run_id=run_id, graph="cli-agent", meter=meter) + ctx = RunContext(run_id=run_id, graph="cli-agent", meter=meter, tenant=tenant) node = AgentNode( model=model, harness=harness, diff --git a/grapharc/cli/graphrun.py b/grapharc/cli/graphrun.py index 49bd6c6..3d08a89 100644 --- a/grapharc/cli/graphrun.py +++ b/grapharc/cli/graphrun.py @@ -169,6 +169,7 @@ def run_graph( registry_target = settings.resolve("registry", registry_target, DEFAULT_REGISTRY) policy_path = settings.resolve_path("policy", policy_path) tenant = settings.resolve("tenant", tenant, "default") + trace_tenant = tenant if settings.sources["tenant"] != "default" else None max_tokens = settings.resolve("max_tokens", max_tokens) max_iterations = settings.resolve("max_iterations", max_iterations) max_seconds = settings.resolve("max_seconds", max_seconds) @@ -221,7 +222,8 @@ def run_graph( # `grapharc metrics ` finds nothing under that id. rid = run_id or uuid.uuid4().hex[:12] verdict = checker.check( - proposal, meter=meter, ctx=RunContext(run_id=rid, graph="cli-run", meter=meter) + proposal, meter=meter, + ctx=RunContext(run_id=rid, graph="cli-run", meter=meter, tenant=trace_tenant), ) common = { @@ -322,7 +324,7 @@ def run_graph( ) return EXIT_OK - state = compiled.invoke(schema().model_dump(), run_id=rid) + state = compiled.invoke(schema().model_dump(), run_id=rid, tenant=trace_tenant) payload = {"ok": True, "checked_only": False, **common, "state": state} lines = [ diff --git a/grapharc/cli/main.py b/grapharc/cli/main.py index 75980ab..c815f92 100644 --- a/grapharc/cli/main.py +++ b/grapharc/cli/main.py @@ -677,6 +677,54 @@ def _cmd_metrics(args: argparse.Namespace) -> int: return EXIT_OK +def _cmd_cost(args: argparse.Namespace) -> int: + from grapharc.observe.cost import attribute, attribute_tenant + + recorder = _existing_trace(args.path, command="cost", as_json=args.json) + if isinstance(recorder, int): + return recorder + try: + report = ( + attribute_tenant(recorder, args.tenant) + if args.tenant is not None + else attribute(recorder, args.run_id) + ) + except TraceReadError as exc: + return fail(str(exc), as_json=args.json, command="cost") + except (OSError, UnicodeDecodeError) as exc: + return fail( + f"unreadable trace file: {args.path}: {exc}", as_json=args.json, command="cost" + ) + except ReplayError as exc: + return fail(str(exc), as_json=args.json, command="cost", code=EXIT_FAILED) + totals = { + "tokens": report.tokens, + "duration_ms": report.duration_ms, + "recorded_cost_usd": report.recorded_cost_usd, + "estimated_cost_usd": report.estimated_cost_usd, + "cost_usd": report.cost_usd, + "unpriced_tokens": report.unpriced_tokens, + "complete": report.complete, + } + data = {**report.model_dump(), **totals} + lines = [] + if args.tenant is not None: + lines.append(f"tenant {args.tenant} · {len(report.runs)} runs") + else: + lines.append(report.format()) + for key, value in totals.items(): + text = ( + ("unknown" if value is None else f"${value:.6f}") + if key.endswith("cost_usd") else str(value) + ) + lines.append(style.kv(key, text)) + emit( + {"ok": True, "command": "cost", "path": str(args.path), **data}, + lines, as_json=args.json, + ) + return EXIT_OK + + def _cmd_viz(args: argparse.Namespace) -> int: recorder = _existing_trace(args.path, command="viz", as_json=args.json) if isinstance(recorder, int): @@ -1229,6 +1277,15 @@ def build_parser() -> argparse.ArgumentParser: mx.add_argument("run_id") mx.set_defaults(handler=_cmd_metrics) + cost = sub.add_parser( + "cost", parents=[common], help="attribute recorded spend by run or tenant" + ) + cost.add_argument("path", type=Path) + scope = cost.add_mutually_exclusive_group(required=True) + scope.add_argument("--run-id", help="one recorded run") + scope.add_argument("--tenant", help="all runs explicitly labelled with this tenant") + cost.set_defaults(handler=_cmd_cost) + vz = sub.add_parser("viz", parents=[common], help="render a run's executed path as Mermaid") vz.add_argument("path", type=Path) vz.add_argument("run_id") diff --git a/grapharc/cli/plan.py b/grapharc/cli/plan.py index b9b2f75..b1ddc4d 100644 --- a/grapharc/cli/plan.py +++ b/grapharc/cli/plan.py @@ -451,6 +451,7 @@ def execute_plan( model_spec = settings.resolve("model", model_spec, record.get("model")) policy_path = settings.resolve_path("policy", policy_path) tenant = settings.resolve("tenant", tenant, "default") + trace_tenant = tenant if settings.sources["tenant"] != "default" else None max_tokens = settings.resolve("max_tokens", max_tokens, 100_000) model_args = _parse_model_args(model_arg_pairs) # The registry gets the real model (agent-backed kinds need one); the @@ -514,7 +515,7 @@ def execute_plan( ), ) initial = schema(goal=goal) if "goal" in schema.model_fields else schema() - result = loop.run(goal, initial, run_id=run_id) + result = loop.run(goal, initial, run_id=run_id, tenant=trace_tenant) executed = any(r.executed for r in result.rounds) if executed: @@ -950,6 +951,7 @@ def plan( ) policy_path = settings.resolve_path("policy", policy_path) tenant = settings.resolve("tenant", tenant, "default") + trace_tenant = tenant if settings.sources["tenant"] != "default" else None max_rounds = settings.resolve("max_rounds", max_rounds, 8) max_tokens = settings.resolve("max_tokens", max_tokens, 100_000) max_planning_failures = settings.resolve( @@ -1018,7 +1020,7 @@ def plan( # `plan` plans; `go` (and `plan --go`) executes. The attribute rather # than a ctor param keeps every registry module's build_loop signature. loop.plan_only = command == "plan" and not go_after - result = loop.run(goal, initial, run_id=run_id) + result = loop.run(goal, initial, run_id=run_id, tenant=trace_tenant) # The admitted shape as data, computed before the plan file so the file # can carry the same `mutating` verdict the payload does. An external diff --git a/grapharc/harness/agent.py b/grapharc/harness/agent.py index a24338c..78fcf32 100644 --- a/grapharc/harness/agent.py +++ b/grapharc/harness/agent.py @@ -667,7 +667,8 @@ def _run_delegated(self, prompt: str, ctx: RunContext) -> AgentResult: if self.trace is not None: self.trace.event( run_id=ctx.run_id, graph=ctx.graph, node=self.name, phase="model", - step=step, thread_id=ctx.thread_id, attempt=ctx.attempt, + step=step, thread_id=ctx.thread_id, + tenant=ctx.tenant, attempt=ctx.attempt, state_delta={"executor": "delegated", "delegated_mode": self.delegated_mode, "tools": ("all of Claude Code's" if bypass @@ -705,7 +706,8 @@ def _run_delegated(self, prompt: str, ctx: RunContext) -> AgentResult: if self.trace is not None: self.trace.event( run_id=ctx.run_id, graph=ctx.graph, node=self.name, phase="stop", - step=step, thread_id=ctx.thread_id, attempt=ctx.attempt, + step=step, thread_id=ctx.thread_id, + tenant=ctx.tenant, attempt=ctx.attempt, state_delta={"executor": "delegated", "delegated_mode": self.delegated_mode, "termination_reason": exc.reason}, @@ -724,7 +726,8 @@ def _run_delegated(self, prompt: str, ctx: RunContext) -> AgentResult: if self.trace is not None: self.trace.event( run_id=ctx.run_id, graph=ctx.graph, node=self.name, phase="stop", - step=step, thread_id=ctx.thread_id, attempt=ctx.attempt, + step=step, thread_id=ctx.thread_id, + tenant=ctx.tenant, attempt=ctx.attempt, tokens=run.tokens_reported or None, cost_usd=run.cost_usd, state_delta={"executor": "delegated", @@ -947,6 +950,7 @@ def _emit(self, ctx: RunContext, phase: str, *, node: str | None = None, **field self.trace.event( run_id=ctx.run_id, thread_id=ctx.thread_id, + tenant=ctx.tenant, attempt=ctx.attempt, graph=ctx.graph, node=node or f"{self.name}:{phase}", diff --git a/grapharc/observe/__init__.py b/grapharc/observe/__init__.py index a39ef7e..bd97a1d 100644 --- a/grapharc/observe/__init__.py +++ b/grapharc/observe/__init__.py @@ -10,8 +10,10 @@ NodeCost, RateCard, RunCost, + TenantCost, ThreadCost, attribute, + attribute_tenant, attribute_thread, by_node, tokens_by_model, @@ -62,11 +64,13 @@ "RunMetrics", "Span", "SpanExporter", + "TenantCost", "ThreadCost", "TraceEvent", "TraceReadError", "TraceRecorder", "attribute", + "attribute_tenant", "attribute_thread", "by_node", "diff_runs", diff --git a/grapharc/observe/cost.py b/grapharc/observe/cost.py index 690a28d..ce1ae93 100644 --- a/grapharc/observe/cost.py +++ b/grapharc/observe/cost.py @@ -29,14 +29,14 @@ Where the node was an `AgentNode`, its per-call `"model"` sub-events still hold them, and they are reported as `tokens_before_error` rather than folded into the total that must match `metrics`. -- There is no tenant on a trace event, so tenant attribution is not offered - here. Run, thread (session) and node are what the format supports today. +- Tenant attribution needs an explicit, consistent label on every event in + a run. Older unlabelled runs are not assigned to a tenant retrospectively. """ from __future__ import annotations from collections import defaultdict -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from pathlib import Path from pydantic import BaseModel, Field @@ -48,7 +48,7 @@ replay, replay_thread, ) -from grapharc.observe.trace import TraceRecorder +from grapharc.observe.trace import TraceEvent, TraceRecorder, load_events class RateCard(BaseModel): @@ -133,6 +133,7 @@ class RunCost(BaseModel): run_id: str graph: str thread_id: str | None = None + tenant: str | None = None executions: int = 0 errors: int = 0 tokens: int = 0 @@ -166,6 +167,8 @@ def node(self, name: str) -> NodeCost | None: def format(self) -> str: lines = [f"run {self.run_id} · graph {self.graph}"] + if self.tenant is not None: + lines.append(f"tenant {self.tenant}") for entry in self.per_node: cost = "—" if entry.cost_usd is None else f"${entry.cost_usd:.6f}" lines.append( @@ -220,6 +223,52 @@ def cost_usd(self) -> float | None: return _combine(self.recorded_cost_usd, self.estimated_cost_usd) +class TenantCost(BaseModel): + """Whole runs explicitly attributed to one tenant, in first-seen order. + + Runs retain their per-node and per-model detail. Recorded and estimated + USD stay separate, and missing prices remain visible as unpriced tokens. + """ + + tenant: str + runs: list[RunCost] = Field(default_factory=list) + + @property + def tokens(self) -> int: + return sum(r.tokens for r in self.runs) + + @property + def duration_ms(self) -> float: + return round(sum(r.duration_ms for r in self.runs), 2) + + @property + def unpriced_tokens(self) -> int: + return sum(r.unpriced_tokens for r in self.runs) + + @property + def recorded_cost_usd(self) -> float | None: + return _sum_optional(r.recorded_cost_usd for r in self.runs) + + @property + def estimated_cost_usd(self) -> float | None: + return _sum_optional(r.estimated_cost_usd for r in self.runs) + + @property + def cost_usd(self) -> float | None: + return _combine(self.recorded_cost_usd, self.estimated_cost_usd) + + @property + def complete(self) -> bool: + return self.unpriced_tokens == 0 + + +def _run_tenant(events: Sequence[TraceEvent]) -> str | None: + labels = {event.tenant for event in events} + if len(labels) > 1: + raise ReplayError(f"inconsistent tenant labels for run_id {events[0].run_id!r}") + return next(iter(labels), None) + + def _combine(*values: float | None) -> float | None: """Sum what is known; None only when nothing is.""" known = [v for v in values if v is not None] @@ -365,6 +414,15 @@ def _price_run(run: ReplayedRun, rates: RateCard | None) -> RunCost: if sub.phase != "model": if sub.cost_usd is not None: recorded_total.append(sub.cost_usd) + elif sub.tokens: + # Planner turns also spend tokens outside node spans. A + # non-model phase must not make that spend disappear from + # pricing or leave an incomplete report marked complete. + amount = card.price(sub.tokens, sub.model) + if amount is None: + unpriced += sub.tokens + else: + estimated_total.append(amount) continue sub_estimate = ( None if sub.cost_usd is not None else card.price(sub.tokens or 0, sub.model) @@ -394,6 +452,7 @@ def _price_run(run: ReplayedRun, rates: RateCard | None) -> RunCost: entry.duration_ms = round(entry.duration_ms, 2) return RunCost( run_id=run.run_id, + tenant=_run_tenant(run.events), graph=run.graph, thread_id=run.thread_id, executions=sum(n.executions for n in per_node), @@ -472,6 +531,32 @@ def attribute_thread( ) +def attribute_tenant( + source: TraceRecorder | str | Path, + tenant: str, + *, + rates: RateCard | None = None, +) -> TenantCost: + """Attribute all whole runs labelled with `tenant`, from the trace alone. + + Read one snapshot so appended events cannot change labels between + selection and pricing. Refuse any run whose events disagree on tenant, + including a mixture of labelled and unlabelled events. Legacy runs stay + unassigned. Raises `ReplayError` for inconsistent labels or no matching + runs, as `attribute_thread` does for a missing thread. + """ + grouped: dict[str, list[TraceEvent]] = {} + for event in load_events(source): + grouped.setdefault(event.run_id, []).append(event) + selected = [] + for run_id, events in grouped.items(): + if _run_tenant(events) == tenant: + selected.append(_price_run(replay(events, run_id), rates)) + if not selected: + raise ReplayError(f"no events for tenant {tenant!r} in the trace") + return TenantCost(tenant=tenant, runs=selected) + + def by_node(source: TraceRecorder | str | Path, *, rates: RateCard | None = None) -> list[NodeCost]: """Every node's cost across every run in a trace file, most expensive first. @@ -520,8 +605,10 @@ def tokens_by_model(source: TraceRecorder | str | Path) -> dict[str, int]: "NodeCost", "RateCard", "RunCost", + "TenantCost", "ThreadCost", "attribute", + "attribute_tenant", "attribute_thread", "by_node", "tokens_by_model", diff --git a/grapharc/observe/replay.py b/grapharc/observe/replay.py index e347b6c..6e6581f 100644 --- a/grapharc/observe/replay.py +++ b/grapharc/observe/replay.py @@ -286,7 +286,7 @@ def _build_executions( return order, unattached -def replay(source: TraceRecorder | str | Path, run_id: str) -> ReplayedRun: +def replay(source: TraceRecorder | str | Path | Sequence[TraceEvent], run_id: str) -> ReplayedRun: """Reconstruct one run from its trace. Raises `ReplayError` if the trace holds no events for `run_id` — an empty diff --git a/grapharc/observe/trace.py b/grapharc/observe/trace.py index 5051c29..bc85a07 100644 --- a/grapharc/observe/trace.py +++ b/grapharc/observe/trace.py @@ -16,6 +16,7 @@ import json import threading +from collections.abc import Sequence from datetime import UTC, datetime from pathlib import Path from typing import Any @@ -62,6 +63,9 @@ class TraceEvent(BaseModel): ts: str run_id: str thread_id: str | None = None + # An explicit attribution label, not an authorization boundary. Omitted + # when unset so legacy trace output and parsing remain unchanged. + tenant: str | None = None attempt: int = 1 graph: str node: str @@ -124,7 +128,12 @@ def __init__(self, path: str | Path) -> None: self._thread_max: dict[str, tuple[int, int]] = {} def record(self, event: TraceEvent) -> None: - line = json.dumps(_jsonable(event.model_dump(exclude_none=True)), ensure_ascii=False) + payload = _jsonable(event.model_dump(exclude_none=True)) + if event.tenant is not None: + # Attribution identifiers must survive exactly; truncation is + # useful for state text, but would silently rename this tenant. + payload["tenant"] = event.tenant + line = json.dumps(payload, ensure_ascii=False) with self._lock, self.path.open("a", encoding="utf-8") as f: f.write(line + "\n") @@ -137,6 +146,7 @@ def event( phase: str, step: int, thread_id: str | None = None, + tenant: str | None = None, attempt: int = 1, state_delta: dict[str, Any] | None = None, duration_ms: float | None = None, @@ -150,6 +160,7 @@ def event( ts=datetime.now(UTC).isoformat(timespec="milliseconds"), run_id=run_id, thread_id=thread_id, + tenant=tenant, attempt=attempt, graph=graph, node=node, @@ -315,9 +326,11 @@ def read_tail(self, run_id: str | None = None) -> tuple[list[TraceEvent], int]: def load_events( - source: TraceRecorder | str | Path, run_id: str | None = None + source: TraceRecorder | str | Path | Sequence[TraceEvent], run_id: str | None = None ) -> list[TraceEvent]: - """Read a trace from a recorder or a path, so callers can pass either.""" + """Read a recorder, path, or an already-read snapshot of trace events.""" + if not isinstance(source, (TraceRecorder, str, Path)): + return [event for event in source if run_id is None or event.run_id == run_id] recorder = source if isinstance(source, TraceRecorder) else TraceRecorder(source) return recorder.read_events(run_id) diff --git a/grapharc/planner/admission.py b/grapharc/planner/admission.py index 75f15b5..d859f31 100644 --- a/grapharc/planner/admission.py +++ b/grapharc/planner/admission.py @@ -1085,6 +1085,7 @@ def _emit( # proposal's own id stands in for the run id. run_id=ctx.run_id if ctx is not None else proposal.proposal_id, thread_id=ctx.thread_id if ctx is not None else None, + tenant=ctx.tenant if ctx is not None else None, attempt=ctx.attempt if ctx is not None else 1, graph=ctx.graph if ctx is not None else self.name, node=f"{self.name}:{proposal.proposal_id}", diff --git a/grapharc/planner/loop.py b/grapharc/planner/loop.py index 192599d..ba2289a 100644 --- a/grapharc/planner/loop.py +++ b/grapharc/planner/loop.py @@ -366,17 +366,20 @@ def run( *, run_id: str | None = None, thread_id: str | None = None, + tenant: str | None = None, ) -> LoopResult: """Carry `goal` from first proposal to recorded stop. `state` may be a `state_schema` instance, a dict, or None for the schema's defaults. The returned `LoopResult.state` is the last state any - executed round produced. + executed round produced. `tenant` labels planning, admission and every + executed round with the same explicit attribution label. """ rid = run_id or uuid.uuid4().hex[:12] meter = BudgetMeter(self.budget) ctx = RunContext( - run_id=rid, graph=self.name, meter=meter, thread_id=thread_id or rid, attempt=1 + run_id=rid, graph=self.name, meter=meter, + thread_id=thread_id or rid, tenant=tenant, attempt=1 ) current = self._initial_state(state) rounds: list[RoundRecord] = [] @@ -701,6 +704,7 @@ def _execute( self.trace.event( run_id=ctx.run_id, thread_id=ctx.thread_id, + tenant=ctx.tenant, attempt=ctx.attempt, graph=compiled.arc.name, node="topology", @@ -772,6 +776,7 @@ def _execute( state, run_id=ctx.run_id, thread_id=f"{ctx.thread_id}/r{round_number}", + tenant=ctx.tenant, budget=budget, ) except BudgetExceeded as exc: @@ -949,6 +954,7 @@ def _emit(self, ctx: RunContext, *, node: str, phase: str, **fields: Any) -> Non self.trace.event( run_id=ctx.run_id, thread_id=ctx.thread_id, + tenant=ctx.tenant, attempt=ctx.attempt, graph=self.name, node=node, diff --git a/grapharc/planner/proposal.py b/grapharc/planner/proposal.py index 4c59bff..40954c4 100644 --- a/grapharc/planner/proposal.py +++ b/grapharc/planner/proposal.py @@ -823,6 +823,7 @@ def _emit(self, ctx: RunContext, **fields: Any) -> None: self.trace.event( run_id=ctx.run_id, thread_id=ctx.thread_id, + tenant=ctx.tenant, attempt=ctx.attempt, graph=ctx.graph, node=f"{self.name}:plan", diff --git a/grapharc/runtime/graph.py b/grapharc/runtime/graph.py index a7b722f..02b7c7a 100644 --- a/grapharc/runtime/graph.py +++ b/grapharc/runtime/graph.py @@ -141,7 +141,11 @@ class MissingRunContextError(Exception): class RunContext: - """Per-run bookkeeping shared with every node via the runnable config.""" + """Per-run bookkeeping shared with every node via the runnable config. + + `tenant` is an explicit trace attribution label, not a policy decision or + an authorization boundary. None keeps the run unlabelled. + """ def __init__( self, @@ -150,6 +154,7 @@ def __init__( graph: str, meter: BudgetMeter, thread_id: str | None = None, + tenant: str | None = None, attempt: int = 1, step_seed: int = 0, ) -> None: @@ -157,6 +162,7 @@ def __init__( self.graph = graph self.meter = meter self.thread_id = thread_id + self.tenant = tenant self.attempt = attempt self._step = step_seed self._lock = threading.Lock() @@ -802,6 +808,7 @@ def emit(phase: str, **kw: Any) -> None: trace.event( run_id=ctx.run_id, thread_id=ctx.thread_id, + tenant=ctx.tenant, attempt=ctx.attempt, graph=self.name, node=name, @@ -1012,7 +1019,11 @@ def __init__(self, inner: Any, arc: GraphARC) -> None: self.last_run: RunContext | None = None def _run_config( - self, thread_id: str | None, run_id: str | None, budget: Budget | None + self, + thread_id: str | None, + run_id: str | None, + budget: Budget | None, + tenant: str | None, ) -> dict[str, Any]: rid = run_id or uuid.uuid4().hex[:12] thread = thread_id or rid @@ -1025,6 +1036,7 @@ def _run_config( graph=self.arc.name, meter=BudgetMeter(budget or self.arc.budget or _NOOP_BUDGET), thread_id=thread, + tenant=tenant, attempt=attempt, step_seed=step_seed, ) @@ -1044,6 +1056,7 @@ def _run_config( phase="topology", step=0, thread_id=thread, + tenant=tenant, attempt=attempt, state_delta=topology_delta(self.arc), ) @@ -1102,6 +1115,7 @@ def invoke( *, thread_id: str | None = None, run_id: str | None = None, + tenant: str | None = None, budget: Budget | None = None, ) -> Any: """Run the graph. Pass `input=None` with a previous `thread_id` to resume @@ -1110,10 +1124,12 @@ def invoke( Notes: the budget meter is fresh per invoke — limits bound one attempt, not the lifetime of a thread. Trace step/attempt counters continue from the thread's history so replay points stay unique across resumes. + `tenant` labels this attempt's events; it is never inferred from state + or inherited from a previous checkpoint. """ self._reject_async_nodes("invoke") input = self._checked_input("invoke", input) - return self.inner.invoke(input, self._run_config(thread_id, run_id, budget)) + return self.inner.invoke(input, self._run_config(thread_id, run_id, budget, tenant)) def stream( self, @@ -1121,6 +1137,7 @@ def stream( *, thread_id: str | None = None, run_id: str | None = None, + tenant: str | None = None, budget: Budget | None = None, **stream_kwargs: Any, ) -> Iterator[Any]: @@ -1132,7 +1149,7 @@ def stream( self._reject_async_nodes("stream") input = self._checked_input("stream", input) yield from self.inner.stream( - input, self._run_config(thread_id, run_id, budget), **stream_kwargs + input, self._run_config(thread_id, run_id, budget, tenant), **stream_kwargs ) async def ainvoke( @@ -1141,6 +1158,7 @@ async def ainvoke( *, thread_id: str | None = None, run_id: str | None = None, + tenant: str | None = None, budget: Budget | None = None, ) -> Any: """Async twin of `invoke()`, with the same discipline and the same notes. @@ -1150,7 +1168,7 @@ async def ainvoke( nodes is how `max_seconds` is delivered: see `_async_deadline`. """ input = self._checked_input("ainvoke", input) - return await self.inner.ainvoke(input, self._run_config(thread_id, run_id, budget)) + return await self.inner.ainvoke(input, self._run_config(thread_id, run_id, budget, tenant)) async def astream( self, @@ -1158,13 +1176,14 @@ async def astream( *, thread_id: str | None = None, run_id: str | None = None, + tenant: str | None = None, budget: Budget | None = None, **stream_kwargs: Any, ) -> AsyncIterator[Any]: """Async twin of `stream()`; `.inner.astream()` fails closed.""" input = self._checked_input("astream", input) async for chunk in self.inner.astream( - input, self._run_config(thread_id, run_id, budget), **stream_kwargs + input, self._run_config(thread_id, run_id, budget, tenant), **stream_kwargs ): yield chunk @@ -1174,6 +1193,7 @@ async def astream_events( *, thread_id: str | None = None, run_id: str | None = None, + tenant: str | None = None, budget: Budget | None = None, version: Literal["v1", "v2"] = "v2", **kwargs: Any, @@ -1188,7 +1208,7 @@ async def astream_events( raise ValueError(f"astream_events supports version 'v1' or 'v2', got {version!r}") input = self._checked_input("astream_events", input) async for event in self.inner.astream_events( - input, self._run_config(thread_id, run_id, budget), version=version, **kwargs + input, self._run_config(thread_id, run_id, budget, tenant), version=version, **kwargs ): yield event diff --git a/grapharc/server/runtime.py b/grapharc/server/runtime.py index f29fa6d..eef847f 100644 --- a/grapharc/server/runtime.py +++ b/grapharc/server/runtime.py @@ -248,7 +248,11 @@ def _trace_payload(event: TraceEvent) -> dict[str, Any]: # `_trace_shape` reprs it below — which is what the trace format does # with such a value anyway. dumped = event.model_dump(exclude_none=True) - return _trace_shape(dumped) + payload = _trace_shape(dumped) + if event.tenant is not None: + # Match TraceRecorder: attribution identifiers are never clipped. + payload["tenant"] = event.tenant + return payload class BroadcastRecorder(TraceRecorder): @@ -295,6 +299,7 @@ def event( phase: str, step: int, thread_id: str | None = None, + tenant: str | None = None, attempt: int = 1, state_delta: dict[str, Any] | None = None, duration_ms: float | None = None, @@ -314,6 +319,7 @@ def event( ts=_now(), run_id=run_id, thread_id=thread_id, + tenant=tenant, attempt=attempt, graph=graph, node=node, diff --git a/tests/test_cli_style.py b/tests/test_cli_style.py index 6f30f2f..bbba558 100644 --- a/tests/test_cli_style.py +++ b/tests/test_cli_style.py @@ -66,6 +66,7 @@ def build(_workspace: dict) -> list[str]: # A painted row per trace event: dim, cell, accent and err in one output. pytest.param(lambda w: ["trace", str(w["trace"])], id="trace"), pytest.param(lambda w: ["metrics", str(w["trace"]), w["run_id"]], id="metrics"), + pytest.param(lambda w: ["cost", str(w["trace"]), "--run-id", w["run_id"]], id="cost"), ] # The subset whose output is reproducible enough to compare byte-for-byte across diff --git a/tests/test_tenant_cost.py b/tests/test_tenant_cost.py new file mode 100644 index 0000000..e4c3e7e --- /dev/null +++ b/tests/test_tenant_cost.py @@ -0,0 +1,447 @@ +"""Explicit run tenants must survive every producer and trace-only cost reader.""" + +from __future__ import annotations + +import json + +import pytest + +from grapharc.cli.main import main +from grapharc.examples import plan_incident, plan_research +from grapharc.harness import AgentNode, Harness, LocalExecutor, PermissionPolicy, ToolRegistry +from grapharc.observe import RateCard, ReplayError, TraceEvent, TraceRecorder, attribute +from grapharc.runtime.budget import Budget, BudgetMeter +from grapharc.runtime.graph import END, START, GraphARC, RunContext +from grapharc.runtime.state import GraphARCState +from grapharc.testing import ScriptedChatModel + + +class State(GraphARCState): + answer: str = "" + + +def priced_graph(trace, *, fail=False): + model = ScriptedChatModel(responses=["done"], cost_usd=0.002, model_name="acme/fast") + graph = GraphARC(State, name="priced", trace=trace) + + def body(state): + answer = model.invoke("work").content + if fail: + raise ValueError("failed after spending") + return {"answer": answer} + + graph.add_node("think", body, writes={"answer"}) + graph.add_edge(START, "think") + graph.add_edge("think", END) + return graph.compile() + + +def tenant_cost(source, tenant, **kwargs): + # Lazy import keeps pre-feature collection working: the new tests can + # demonstrate the missing behavior rather than aborting collection. + from grapharc.observe import attribute_tenant + + return attribute_tenant(source, tenant, **kwargs) + + +def event(trace, run, tenant, phase="end", step=1, **kwargs): + trace.event( + run_id=run, tenant=tenant, graph="g", node="worker", phase=phase, step=step, **kwargs + ) + + +def test_legacy_trace_bytes_and_parsing_are_unchanged(tmp_path): + trace = TraceRecorder(tmp_path / "old.jsonl") + old = { + "ts": "2026-01-01", + "run_id": "old", + "attempt": 1, + "graph": "g", + "node": "n", + "phase": "end", + "step": 1, + "tokens": 7, + } + trace.record(TraceEvent(**old)) + assert trace.path.read_text() == json.dumps(old) + "\n" + parsed = trace.read_events()[0] + assert parsed.tenant is None + assert "tenant" not in parsed.model_dump(exclude_none=True) + assert attribute(trace, "old").tokens == 7 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("entry", ["invoke", "stream", "ainvoke", "astream", "astream_events"]) +async def test_every_execution_entry_carries_the_tenant(tmp_path, entry): + trace = TraceRecorder(tmp_path / "trace.jsonl") + compiled = priced_graph(trace) + if entry == "invoke": + compiled.invoke({}, tenant="alpha", run_id="run") + elif entry == "stream": + list(compiled.stream({}, tenant="alpha", run_id="run")) + elif entry == "ainvoke": + await compiled.ainvoke({}, tenant="alpha", run_id="run") + else: + _ = [item async for item in getattr(compiled, entry)({}, tenant="alpha", run_id="run")] + events = trace.read_events() + assert {e.phase for e in events} >= {"topology", "start", "end"} + assert {e.tenant for e in events} == {"alpha"} + assert compiled.last_run.tenant == "alpha" + cost = tenant_cost(trace, "alpha") + assert cost.recorded_cost_usd == pytest.approx(0.002) + assert cost.tokens == attribute(trace, "run").tokens > 0 + + +def test_error_events_keep_the_tenant_and_spend(tmp_path): + trace = TraceRecorder(tmp_path / "trace.jsonl") + with pytest.raises(ValueError, match="failed after spending"): + priced_graph(trace, fail=True).invoke({}, tenant="alpha", run_id="failed") + assert {e.tenant for e in trace.read_events()} == {"alpha"} + cost = tenant_cost(trace, "alpha") + assert cost.runs[0].errors == 1 + assert cost.recorded_cost_usd == attribute(trace, "failed").recorded_cost_usd + assert cost.tokens > 0 + assert cost.unpriced_tokens == cost.tokens + + +def test_tenants_partition_runs_without_double_counting_substeps(tmp_path): + trace = TraceRecorder(tmp_path / "trace.jsonl") + event(trace, "a1", "alpha", "start") + event(trace, "b1", "beta", tokens=200, cost_usd=0.02) + event(trace, "a1", "alpha", "model", step=2, tokens=100, cost_usd=0.01) + event(trace, "a1", "alpha", tokens=100, cost_usd=0.01) + event(trace, "a2", "alpha", "start") + event(trace, "a2", "alpha", "model", step=2, tokens=50, model="acme/fast") + event(trace, "a2", "alpha", tokens=50) + # Legacy spend cannot silently enter a named tenant's bill. + event(trace, "old", None, tokens=9, cost_usd=1.0) + rates = RateCard(per_model={"acme/": 2.0}) + alpha = tenant_cost(trace.path, "alpha", rates=rates) + beta = tenant_cost(trace, "beta", rates=rates) + assert [r.run_id for r in alpha.runs] == ["a1", "a2"] + assert alpha.tokens == 150 + assert beta.tokens == 200 + assert alpha.recorded_cost_usd == pytest.approx(0.01) + assert alpha.estimated_cost_usd == pytest.approx(0.1) + assert alpha.cost_usd == pytest.approx(0.11) + assert alpha.unpriced_tokens == 0 + assert alpha.tokens + beta.tokens == sum(attribute(trace, r).tokens for r in ("a1", "a2", "b1")) + assert alpha.recorded_cost_usd + beta.recorded_cost_usd == pytest.approx(0.03) + assert tenant_cost(trace, "alpha").unpriced_tokens == 50 + assert tenant_cost(trace, "alpha").estimated_cost_usd is None + + +@pytest.mark.parametrize("labels", [("alpha", "beta"), ("alpha", None)]) +def test_a_run_with_conflicting_labels_cannot_be_billed(tmp_path, labels): + trace = TraceRecorder(tmp_path / "trace.jsonl") + for step, label in enumerate(labels, start=1): + event(trace, "mixed", label, step=step, tokens=10) + with pytest.raises(ReplayError, match="inconsistent tenant"): + tenant_cost(trace, "alpha") + + +def test_an_unknown_tenant_is_distinct_from_an_empty_bill(tmp_path): + trace = TraceRecorder(tmp_path / "trace.jsonl") + event(trace, "old", None, tokens=9) + with pytest.raises(ReplayError, match="no events for tenant"): + tenant_cost(trace, "alpha") + + +@pytest.mark.parametrize("label", ["研究:team/one", "a" * 2500], ids=["unicode", "long"]) +def test_tenant_identifiers_are_preserved_exactly(tmp_path, label): + trace = TraceRecorder(tmp_path / "trace.jsonl") + event(trace, "r", label, tokens=5) + assert trace.read_events()[0].tenant == label + assert tenant_cost(trace, label).tokens == 5 + + +@pytest.mark.parametrize("label", [None, "alpha", "a" * 2500], ids=["unlabelled", "named", "long"]) +def test_broadcast_recorder_matches_the_file_contract(tmp_path, label): + pytest.importorskip("fastapi") + from grapharc.server.runtime import BroadcastRecorder + + seen = [] + trace = BroadcastRecorder(tmp_path / "broadcast.jsonl", seen.append) + priced_graph(trace).invoke({}, run_id="run", tenant=label) + on_disk = [json.loads(line) for line in trace.path.read_text().splitlines()] + assert seen == on_disk + assert {e.tenant for e in trace.read_events()} == {label} + if label is None: + assert all("tenant" not in item for item in seen) + else: + assert tenant_cost(trace, label).recorded_cost_usd == pytest.approx(0.002) + + +def test_tenant_attribution_reads_one_snapshot(tmp_path): + class CountingRecorder(TraceRecorder): + reads = 0 + + def read_events(self, run_id=None): + self.reads += 1 + return super().read_events(run_id) + + trace = CountingRecorder(tmp_path / "trace.jsonl") + event(trace, "a", "alpha", tokens=10) + event(trace, "b", "alpha", tokens=20) + assert tenant_cost(trace, "alpha").tokens == 30 + assert trace.reads == 1 + + +class ToolModel(ScriptedChatModel): + def bind_tools(self, tools, **kwargs): + return self + + +def test_an_agent_inherits_its_run_context_tenant(tmp_path): + trace = TraceRecorder(tmp_path / "trace.jsonl") + model = ToolModel(responses=["done"], cost_usd=0.004) + harness = Harness(ToolRegistry(), PermissionPolicy(), executor=LocalExecutor()) + node = AgentNode(model, harness, trace=trace) + ctx = RunContext(run_id="agent", graph="g", meter=BudgetMeter(Budget()), tenant="alpha") + node.run("finish", ctx) + assert {e.phase for e in trace.read_events()} == {"model", "stop"} + assert {e.tenant for e in trace.read_events()} == {"alpha"} + assert tenant_cost(trace, "alpha").recorded_cost_usd == pytest.approx(0.004) + + +@pytest.mark.parametrize("explicit", [False, True]) +def test_policy_governed_agent_cli_records_only_an_explicit_tenant( + tmp_path, monkeypatch, capsys, explicit +): + monkeypatch.chdir(tmp_path) + monkeypatch.delenv("GRAPHARC_TENANT", raising=False) + monkeypatch.delenv("GRAPHARC_CONFIG", raising=False) + policy = tmp_path / "policy.toml" + policy.write_text( + 'version = "1"\ndefault = "deny"\ntenants = ["default", "acme"]\n' + '[[rule]]\nid = "tools"\nresource = "tool"\nmatch = "*"\neffect = "allow"\n' + ) + model = ToolModel(responses=["done"], cost_usd=0.004) + monkeypatch.setattr("grapharc.gateway.get_model", lambda *a, **kw: model) + path = tmp_path / "trace.jsonl" + flags = ["--tenant", "acme"] if explicit else [] + assert ( + main( + [ + "agent", + "finish", + "--model", + "test-double", + "--policy", + str(policy), + "--workspace", + str(tmp_path / "work"), + "--trace", + str(path), + *flags, + "--json", + ] + ) + == 0 + ) + payload = json.loads(capsys.readouterr().out) + events = TraceRecorder(path).read_events() + assert events + assert {e.tenant for e in events} == ({"acme"} if explicit else {None}) + if explicit: + assert tenant_cost(path, "acme").recorded_cost_usd == pytest.approx(0.004) + else: + assert attribute(path, payload["run_id"]).tenant is None + assert all("tenant" not in e.model_dump(exclude_none=True) for e in events) + + +@pytest.mark.parametrize("failed", [False, True]) +def test_delegated_agent_events_keep_the_tenant(tmp_path, monkeypatch, failed): + from grapharc.cli.delegate import DelegatedRun, DelegationError + + def delegate(*args, **kwargs): + if failed: + raise DelegationError("failed", reason="test_failure") + return DelegatedRun(True, "done", "done", 1, 15, 0.02, None, [], []) + + monkeypatch.setattr("grapharc.cli.delegate.delegate_task", delegate) + trace = TraceRecorder(tmp_path / "trace.jsonl") + harness = Harness( + ToolRegistry(), PermissionPolicy(), executor=LocalExecutor(), workspace=str(tmp_path) + ) + node = AgentNode(ToolModel(responses=[]), harness, trace=trace) + ctx = RunContext(run_id="delegated", graph="g", meter=BudgetMeter(Budget()), tenant="alpha") + node._run_delegated("finish", ctx) + assert {e.tenant for e in trace.read_events()} == {"alpha"} + cost = tenant_cost(trace, "alpha") + assert cost.tokens == (0 if failed else 15) + assert cost.recorded_cost_usd == (None if failed else pytest.approx(0.02)) + + +@pytest.mark.parametrize("label", ["alpha", None]) +def test_saved_plan_execution_uses_the_current_explicit_tenant( + tmp_path, monkeypatch, capsys, label +): + monkeypatch.chdir(tmp_path) + monkeypatch.delenv("GRAPHARC_TENANT", raising=False) + monkeypatch.delenv("GRAPHARC_CONFIG", raising=False) + path = tmp_path / "trace.jsonl" + assert ( + main( + [ + "plan", + "fix incident", + "--scripted", + "--trace", + str(path), + "--tenant", + "planning", + "--json", + ] + ) + == 0 + ) + planned = json.loads(capsys.readouterr().out) + old_runs = {e.run_id for e in TraceRecorder(path).read_events()} + flags = ["--tenant", label] if label is not None else [] + assert main(["go", planned["plan_file"], *flags, "--json"]) == 0 + _ = capsys.readouterr() + events = TraceRecorder(path).read_events() + executed = [e for e in events if e.run_id not in old_runs] + assert any(e.phase == "end" for e in executed) + assert {e.tenant for e in executed} == {label} + assert {e.tenant for e in events if e.run_id in old_runs} == {"planning"} + + +def test_governed_planning_and_rounds_share_one_tenant_bill(tmp_path): + trace = TraceRecorder(tmp_path / "trace.jsonl") + model = ScriptedChatModel(responses=plan_research.scripted_planner_replies(), cost_usd=0.003) + loop = plan_research.build_loop(model, trace=trace) + result = loop.run("explain elevated latency", run_id="loop", tenant="research") + events = trace.read_events() + assert {e.phase for e in events} >= { + "plan", + "admission", + "topology", + "start", + "end", + "round", + "stop", + } + assert len({e.thread_id for e in events}) > 1 + assert {e.tenant for e in events} == {"research"} + assert any(r.executed for r in result.rounds) + cost = tenant_cost(trace, "research", rates=RateCard(default=2.0)) + assert [r.run_id for r in cost.runs] == ["loop"] + # Planner events currently carry tokens, not a provider price. That + # missing price must remain an estimate, even when the model has one. + assert cost.recorded_cost_usd is None + assert cost.tokens == attribute(trace, "loop").tokens > 0 + assert cost.estimated_cost_usd == pytest.approx(cost.tokens * 2.0 / 1000) + unpriced = tenant_cost(trace, "research") + assert unpriced.unpriced_tokens == unpriced.tokens + assert unpriced.complete is False + + +def test_unlabelled_planner_spend_is_also_priced_or_marked_incomplete(tmp_path): + trace = TraceRecorder(tmp_path / "legacy.jsonl") + trace.event(run_id="old", graph="g", node="planner", phase="plan", step=1, tokens=50) + unpriced = attribute(trace, "old") + assert unpriced.tokens == 50 + assert unpriced.unpriced_tokens == 50 + assert unpriced.complete is False + estimated = attribute(trace, "old", rates=RateCard(default=2.0)) + assert estimated.estimated_cost_usd == pytest.approx(0.1) + assert estimated.recorded_cost_usd is None + assert estimated.complete is True + + +@pytest.mark.parametrize("command", ["run", "plan", "go"]) +@pytest.mark.parametrize("source", ["flag", "env", "config", "default"]) +def test_cli_propagates_only_explicit_tenants(tmp_path, monkeypatch, capsys, command, source): + monkeypatch.chdir(tmp_path) + monkeypatch.delenv("GRAPHARC_TENANT", raising=False) + monkeypatch.delenv("GRAPHARC_CONFIG", raising=False) + trace = tmp_path / "trace.jsonl" + extra = [] + if source == "flag": + extra = ["--tenant", "alpha"] + elif source == "env": + monkeypatch.setenv("GRAPHARC_TENANT", "alpha") + elif source == "config": + (tmp_path / "grapharc.toml").write_text('tenant = "alpha"\n') + if command == "run": + graph = tmp_path / "graph.json" + graph.write_text( + json.dumps( + { + "nodes": [{"name": "triage"}], + "edges": [ + {"source": START, "target": "triage"}, + {"source": "triage", "target": END}, + ], + } + ) + ) + args = ["run", str(graph)] + else: + model = ScriptedChatModel(responses=plan_incident.scripted_planner_replies()) + monkeypatch.setattr("grapharc.gateway.get_model", lambda *a, **kw: model) + model_flags = ["--scripted"] if command == "plan" else ["--model", "test-double"] + args = [ + command, + "fix incident", + "--registry", + "grapharc.examples.plan_incident:build_registry", + *model_flags, + ] + main([*args, *extra, "--trace", str(trace), "--json"]) + payload = json.loads(capsys.readouterr().out) + assert payload["ok"] is True + events = TraceRecorder(trace).read_events() + assert events + expected = None if source == "default" else "alpha" + assert {e.tenant for e in events} == {expected} + if source == "default": + assert all("tenant" not in e.model_dump(exclude_none=True) for e in events) + + +@pytest.mark.parametrize("scope", ["tenant", "run-id"]) +@pytest.mark.parametrize("as_json", [True, False]) +def test_cli_cost_scopes_text_and_json(tmp_path, capsys, scope, as_json): + trace = TraceRecorder(tmp_path / "trace.jsonl") + event(trace, "a", "alpha", tokens=10, cost_usd=0.01) + event(trace, "b", "beta", tokens=50, cost_usd=0.05) + value = "alpha" if scope == "tenant" else "a" + code = main(["cost", str(trace.path), f"--{scope}", value, *(["--json"] if as_json else [])]) + out = capsys.readouterr().out + assert code == 0 + if as_json: + payload = json.loads(out) + assert payload["tokens"] == 10 + assert payload["recorded_cost_usd"] == pytest.approx(0.01) + assert payload["estimated_cost_usd"] is None + assert payload["unpriced_tokens"] == 0 + assert payload["tenant"] == "alpha" + else: + assert "alpha" in out + assert "0.010000" in out + assert "beta" not in out + + +@pytest.mark.parametrize( + "problem,expected", + [("missing", 2), ("malformed", 2), ("utf8", 2), ("unknown", 1), ("mixed", 1)], +) +def test_cli_cost_failures_are_structured(tmp_path, capsys, problem, expected): + path = tmp_path / "trace.jsonl" + if problem == "malformed": + path.write_text("not json\n") + elif problem == "utf8": + path.write_bytes(b"\xff\n") + elif problem in ("unknown", "mixed"): + trace = TraceRecorder(path) + event(trace, "r", "beta", tokens=5) + if problem == "mixed": + event(trace, "r", "alpha", step=2, tokens=10) + code = main(["cost", str(path), "--tenant", "alpha", "--json"]) + payload = json.loads(capsys.readouterr().out) + assert code == expected + assert payload["ok"] is False + assert payload["command"] == "cost" + assert payload["error"]