From 16353c21d493d860024ba3ca1448c9a382c66c0a Mon Sep 17 00:00:00 2001 From: Jason-Vaughan <95194903+Jason-Vaughan@users.noreply.github.com> Date: Sat, 4 Jul 2026 00:19:13 -0700 Subject: [PATCH 1/2] Add serve-origin marker + parent-task attribution to usage records (#74) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Records gain an optional origin field (cli|gui|serve) threaded from each entry surface through run_once/run_once_stream; --stats shows the per-origin split, with pre-existing records rolling up as "untagged" rather than being guessed at. The serve endpoint additionally reads an optional X-TangleBrain-Parent-Task header — trimmed, capped at 128 chars — and records it as parent_task_id for cross-system attribution. Both are additive, written-only-when-present record fields; old records and readers unaffected. No routing changes. --- ARCHITECTURE.md | 5 ++- CHANGELOG.md | 11 +++++++ README.md | 7 ++++ tanglebrain/cli.py | 42 ++++++++++++++++++++---- tanglebrain/gui/views.py | 4 ++- tanglebrain/measurement.py | 25 ++++++++++++-- tanglebrain/serve/server.py | 26 ++++++++++++--- tanglebrain/serve/views.py | 49 +++++++++++++++++++++++++--- tests/test_cli.py | 34 +++++++++++++++++++ tests/test_gui.py | 1 + tests/test_measurement.py | 49 ++++++++++++++++++++++++++++ tests/test_serve.py | 65 +++++++++++++++++++++++++++++++++++++ 12 files changed, 298 insertions(+), 20 deletions(-) diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 59e141b..1bfa36f 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -217,7 +217,10 @@ adapters' optional `run_stream` capability): the view primes the pump — the fi pulled before any headers commit, so connect-time failures stay plain JSON errors — then the handler writes one flushed SSE event per delta (close-delimited body). Backends that can't stream (cli kinds, hence the router path) deliver the completed text as a single chunk. The -routing core is untouched — measurement and both paid-API gates behave exactly as for a CLI run. +routing core is untouched — measurement and both paid-API gates behave exactly as for a CLI run, +with two attribution extras (#74): records tag `origin: "serve"` (vs `cli`/`gui`), and the +optional `X-TangleBrain-Parent-Task` request header is sanitized and recorded as +`parent_task_id` — metadata only, never routed on. Like the panel, it binds `127.0.0.1` only and is deliberately keyless: the `Authorization` header is never read (local callers need no credential), and the loopback bind is what keeps an diff --git a/CHANGELOG.md b/CHANGELOG.md index b4e9dbb..58699af 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- **Serve-origin marker + parent-task attribution (#74)** — usage records now carry an + `origin` field (`cli` | `gui` | `serve`) so serve-mode traffic is distinguishable from CLI and + panel runs, and `tanglebrain --stats` shows the per-origin split (records predating the field + roll up as `untagged`, never guessed at). OpenAI-compat callers can additionally send an + optional `X-TangleBrain-Parent-Task` header carrying their own task/session identity — + trimmed, capped at 128 chars, recorded onto the usage record as `parent_task_id` for + cross-system attribution, never routed on. Both are additive record fields; old records and + readers are unaffected. + ## [0.19.0] - 2026-07-04 ### Added diff --git a/README.md b/README.md index 2e982f2..ce015d8 100644 --- a/README.md +++ b/README.md @@ -232,6 +232,13 @@ field reports **which backend actually served**; the requested directive and rou in a `tanglebrain` extension field, and `usage` carries the same `chars/4` estimate the measurement log uses (served requests are metered exactly like CLI runs). +Served requests are attributed in the usage log: each record carries `origin: "serve"` (CLI runs +tag `cli`, panel runs `gui`), and `tanglebrain --stats` shows the per-origin split. A caller can +additionally send an optional `X-TangleBrain-Parent-Task` header carrying its own task/session +identity — trimmed, capped at 128 chars, recorded onto the usage record as `parent_task_id` for +cross-system attribution, and never routed on. The reverse linkage already exists: the response's +completion id is `chatcmpl-`, the same task id the usage record carries. + Caveats, by design: - **Streaming is real where the backend can stream.** `stream: true` delivers incremental diff --git a/tanglebrain/cli.py b/tanglebrain/cli.py index d8b9813..b3dd0f1 100644 --- a/tanglebrain/cli.py +++ b/tanglebrain/cli.py @@ -150,6 +150,8 @@ def run_once( task: str | None = None, return_served: bool = False, gate: bool | None = None, + origin: str = "cli", + parent_task_id: str | None = None, ): """Route a single prompt to a roster tier and return the response text. @@ -180,6 +182,10 @@ def run_once( gate: Override for the classifier gate on the default path. ``None`` (default) uses the ``classifier_gate_enabled`` setting; ``True``/``False`` force the gate on/off for this call. Ignored when ``model`` or ``local`` is set. + origin: Which surface this call entered through, recorded on the usage record (#74) — + ``"cli"`` (default), ``"gui"``, or ``"serve"``. Attribution only; routing unaffected. + parent_task_id: Optional external caller identity recorded on the usage record (#74: + the serve endpoint's ``X-TangleBrain-Parent-Task`` header). Attribution only. Returns: The response text (``str``), or ``(text, served)`` when ``return_served`` is ``True``. @@ -219,12 +225,21 @@ def run_once( text = router.route(prompt, task=task, opts=opts) entry = router.last_served - record_task(path=path, entry=entry, prompt=prompt, response=text, task_id=task_id) + record_task( + path=path, entry=entry, prompt=prompt, response=text, task_id=task_id, + origin=origin, parent_task_id=parent_task_id, + ) return (text, _served(path, entry, task_id)) if return_served else text def _recording_stream( - deltas: Iterator[str], path: str, entry, prompt: str, task_id: str + deltas: Iterator[str], + path: str, + entry, + prompt: str, + task_id: str, + origin: str, + parent_task_id: str | None, ) -> Iterator[str]: """Wrap a delta stream so the task is metered exactly once, however the stream ends. @@ -245,6 +260,8 @@ def _recording_stream( entry: The serving roster entry. prompt: The routed prompt (for the usage estimate). task_id: The task id minted for this run. + origin: The entry surface recorded on the usage record (#74). + parent_task_id: Optional external caller identity recorded on the usage record (#74). Yields: The fragments of ``deltas``, unchanged. @@ -258,7 +275,8 @@ def _record(require_text: bool) -> None: return recorded = True record_task( - path=path, entry=entry, prompt=prompt, response="".join(pieces), task_id=task_id + path=path, entry=entry, prompt=prompt, response="".join(pieces), task_id=task_id, + origin=origin, parent_task_id=parent_task_id, ) try: @@ -284,6 +302,8 @@ def run_once_stream( local: bool = False, task: str | None = None, gate: bool | None = None, + origin: str = "cli", + parent_task_id: str | None = None, ) -> tuple[Iterator[str], dict | None]: """Route a single prompt like :func:`run_once`, delivering the response as a delta stream. @@ -310,6 +330,8 @@ def run_once_stream( local: Force the free local tier instead of the frontier-first router. task: Optional task-fit hint for the router (a ``good_at`` tag). gate: Classifier-gate override for the default path, as in :func:`run_once`. + origin: Which surface this call entered through, recorded on the usage record (#74). + parent_task_id: Optional external caller identity recorded on the usage record (#74). Returns: ``(deltas, served)`` — ``deltas`` yields response text fragments in order (joined, they @@ -344,7 +366,10 @@ def run_once_stream( router = Router(roster) text = router.route(prompt, task=task, opts=opts) entry = router.last_served - record_task(path="router", entry=entry, prompt=prompt, response=text, task_id=task_id) + record_task( + path="router", entry=entry, prompt=prompt, response=text, task_id=task_id, + origin=origin, parent_task_id=parent_task_id, + ) return iter([text]), _served("router", entry, task_id) adapter = build_adapter(entry) @@ -352,10 +377,15 @@ def run_once_stream( if run_stream is None: # Per-backend emulation: no streaming capability — run blocking, frame as one delta. text = adapter.run(prompt, opts) - record_task(path=path, entry=entry, prompt=prompt, response=text, task_id=task_id) + record_task( + path=path, entry=entry, prompt=prompt, response=text, task_id=task_id, + origin=origin, parent_task_id=parent_task_id, + ) return iter([text]), _served(path, entry, task_id) - deltas = _recording_stream(run_stream(prompt, opts), path, entry, prompt, task_id) + deltas = _recording_stream( + run_stream(prompt, opts), path, entry, prompt, task_id, origin, parent_task_id + ) return deltas, _served(path, entry, task_id) diff --git a/tanglebrain/gui/views.py b/tanglebrain/gui/views.py index 2cba0e4..e8d9118 100644 --- a/tanglebrain/gui/views.py +++ b/tanglebrain/gui/views.py @@ -128,7 +128,9 @@ def run_prompt(payload: dict) -> dict: try: # return_served gives us the served tier/model directly — no usage-log re-read, no race. - text, served = run_once(str(prompt), model=model, local=local, task=task, return_served=True) + text, served = run_once( + str(prompt), model=model, local=local, task=task, return_served=True, origin="gui" + ) except _RUN_ERRORS as exc: return {"ok": False, "error": str(exc)} diff --git a/tanglebrain/measurement.py b/tanglebrain/measurement.py index dd7da44..fc06f86 100644 --- a/tanglebrain/measurement.py +++ b/tanglebrain/measurement.py @@ -303,6 +303,7 @@ def record_task( kind: str = "task", task_id: str | None = None, parent_task_id: str | None = None, + origin: str | None = None, log_path: str | os.PathLike[str] | None = None, pricing: Pricing | None = None, ) -> None: @@ -325,7 +326,13 @@ def record_task( can be linked back to it). Omitted from the record when ``None``. parent_task_id: For a delegated sub-call, the id of the top-level task that spawned it (read from :data:`PARENT_TASK_ID_ENV`). Omitted from the record when ``None`` — e.g. a delegate - invoked outside a propagated task, which rolls up as ``unlinked``. + invoked outside a propagated task, which rolls up as ``unlinked``. For a top-level task, + an external caller's own task/session identity (#74: the serve endpoint's + ``X-TangleBrain-Parent-Task`` header) — pure attribution metadata; the delegate tree's + ``by_parent`` rollup reads it only off ``delegate`` records. + origin: Which surface the work entered through — ``"cli"`` | ``"gui"`` | ``"serve"`` + (#74). Omitted from the record when ``None``; records without it roll up as + ``untagged`` (pre-#74 history is never guessed at). log_path: Override the usage-log path (tests inject a temp path). Defaults to :func:`default_log_path`. pricing: Override the pricing. Defaults to :func:`load_pricing`. @@ -359,6 +366,8 @@ def record_task( record["task_id"] = str(task_id) if parent_task_id is not None: record["parent_task_id"] = str(parent_task_id) + if origin is not None: + record["origin"] = str(origin) target = Path(log_path) if log_path is not None else default_log_path() target.parent.mkdir(parents=True, exist_ok=True) with _LOG_LOCK: @@ -420,8 +429,10 @@ def rollup(records: list[dict]) -> dict: records: The records from :func:`read_records`. Returns: - A dict with: ``tasks`` (int), ``by_tier`` (tier → count), ``in_tokens_est`` / - ``out_tokens_est`` (summed estimates), and ``cloud_equiv_usd`` / ``spend_avoided_usd`` + A dict with: ``tasks`` (int), ``by_tier`` (tier → count), ``by_origin`` (origin → count, + where a record without an ``origin`` field counts as ``untagged`` — pre-#74 history is + never guessed at), ``in_tokens_est`` / ``out_tokens_est`` (summed estimates), and + ``cloud_equiv_usd`` / ``spend_avoided_usd`` (summed dollars) — all over **top-level tasks only** — plus ``delegates``, a separate sub-rollup of delegated sub-calls ``{count, by_backend: {model: {count, in_tokens_est, out_tokens_est}}, by_parent: {parent_task_id: {count, by_backend: {model: count}}}, @@ -434,6 +445,7 @@ def rollup(records: list[dict]) -> dict: summary: dict = { "tasks": 0, "by_tier": {}, + "by_origin": {}, "in_tokens_est": 0, "out_tokens_est": 0, "cloud_equiv_usd": 0.0, @@ -473,6 +485,8 @@ def rollup(records: list[dict]) -> dict: summary["tasks"] += 1 tier = str(r.get("tier", "unknown")) summary["by_tier"][tier] = summary["by_tier"].get(tier, 0) + 1 + origin = str(r.get("origin") or "untagged") + summary["by_origin"][origin] = summary["by_origin"].get(origin, 0) + 1 summary["in_tokens_est"] += in_tok summary["out_tokens_est"] += out_tok summary["cloud_equiv_usd"] += _as_float(r.get("cloud_equiv_usd")) @@ -503,6 +517,11 @@ def format_rollup(summary: dict, pricing: Pricing) -> str: if by_tier: tiers = ", ".join(f"{k} {v}" for k, v in sorted(by_tier.items())) lines.append(f" By tier: {tiers}") + by_origin = summary.get("by_origin") or {} + # Show the origin split only once it says something — all-untagged history adds no signal. + if any(k != "untagged" for k in by_origin): + origins = ", ".join(f"{k} {v}" for k, v in sorted(by_origin.items())) + lines.append(f" By origin: {origins}") lines.append( f" Est. tokens: in {summary.get('in_tokens_est', 0):,} / " f"out {summary.get('out_tokens_est', 0):,}" diff --git a/tanglebrain/serve/server.py b/tanglebrain/serve/server.py index 03be88e..218a6ec 100644 --- a/tanglebrain/serve/server.py +++ b/tanglebrain/serve/server.py @@ -21,10 +21,12 @@ from tanglebrain.serve.views import ( DEFAULT_PORT, + PARENT_TASK_HEADER, error_envelope, handle_chat_completion, handle_chat_completion_stream, list_models, + sanitize_parent_task, wants_stream, ) @@ -38,7 +40,11 @@ def _json_response(status: int, obj: object) -> tuple[int, str, bytes]: def dispatch( - method: str, path: str, body: bytes = b"", content_type: str = "application/json" + method: str, + path: str, + body: bytes = b"", + content_type: str = "application/json", + parent_task: str | None = None, ) -> tuple[int, str, bytes | Iterator[bytes]]: """Route one request to a view and return ``(status, content_type, body)``. @@ -61,6 +67,8 @@ def dispatch( body: Raw request body bytes (for ``POST``). content_type: The request's ``Content-Type`` header value (POST only; defaults to JSON so socket-free tests needn't supply it). + parent_task: The raw ``X-TangleBrain-Parent-Task`` header value (or ``None``) — sanitized + here and recorded onto the usage record for cross-system attribution (#74). Returns: ``(status_code, content_type, body_bytes)``. @@ -94,13 +102,14 @@ def dispatch( return _json_response( 400, error_envelope("request body must be a JSON object", "invalid_request_error") ) + caller_task = sanitize_parent_task(parent_task) try: if wants_stream(payload): - status, result = handle_chat_completion_stream(payload) + status, result = handle_chat_completion_stream(payload, caller_task) if status == 200: return 200, _SSE, result # Iterator[bytes] — pump already primed return _json_response(status, result) - status, obj = handle_chat_completion(payload) + status, obj = handle_chat_completion(payload, caller_task) except Exception as exc: # noqa: BLE001 — any escape must be clean JSON, never a # dropped connection (e.g. a malformed settings.yaml raising SettingsError on the # auto path). Typed, expected failures are already mapped inside the handlers, @@ -119,7 +128,8 @@ class Handler(BaseHTTPRequestHandler): ``Authorization`` is deliberately never consulted: local callers need no key, and any dummy bearer a client insists on sending is simply ignored. The only headers read are the framing ones — ``Content-Length`` and ``Content-Type`` (see :func:`dispatch` for why the latter is - enforced). + enforced) — plus the optional ``X-TangleBrain-Parent-Task`` attribution header (#74), which + is recorded onto the usage record and never routed on. """ def do_GET(self) -> None: # noqa: N802 (stdlib naming) @@ -138,7 +148,13 @@ def do_POST(self) -> None: # noqa: N802 (stdlib naming) ) return body = self.rfile.read(length) if length else b"" - self._respond(*dispatch("POST", self.path, body, self.headers.get("Content-Type", ""))) + self._respond( + *dispatch( + "POST", self.path, body, + self.headers.get("Content-Type", ""), + self.headers.get(PARENT_TASK_HEADER), + ) + ) def _respond(self, status: int, content_type: str, body: bytes | "Iterator[bytes]") -> None: """Write a complete HTTP response — buffered bytes, or a streamed body. diff --git a/tanglebrain/serve/views.py b/tanglebrain/serve/views.py index f3212d5..2b7cd17 100644 --- a/tanglebrain/serve/views.py +++ b/tanglebrain/serve/views.py @@ -49,9 +49,38 @@ # The model-param alias that engages the full router. AUTO_ALIAS = "auto" +# Optional request header carrying the caller's own task/session identity (#74). Recorded onto +# the usage record verbatim (after sanitizing) for cross-system attribution; never routed on. +PARENT_TASK_HEADER = "X-TangleBrain-Parent-Task" + +# Length cap for the recorded header value — attribution metadata, not a payload channel. +_PARENT_TASK_MAX_LEN = 128 + _OWNED_BY = "tanglebrain" +def sanitize_parent_task(value: object) -> str | None: + """Sanitize a raw ``X-TangleBrain-Parent-Task`` header value for recording (#74). + + Attribution metadata only, so the stance is trim-don't-reject: whitespace is stripped, an + empty/absent/non-string value becomes ``None`` (field omitted from the record), and anything + longer than 128 chars is truncated — a caller cannot stuff arbitrary payloads into the + usage log through this header. + + Args: + value: The raw header value (or ``None`` when the header is absent). + + Returns: + The cleaned identity string, or ``None`` when there is nothing worth recording. + """ + if not isinstance(value, str): + return None + cleaned = value.strip() + if not cleaned: + return None + return cleaned[:_PARENT_TASK_MAX_LEN] + + class BadRequestError(ValueError): """An invalid request payload — maps to HTTP 400 with an OpenAI-style error body.""" @@ -339,7 +368,9 @@ def _parse_chat_request(payload: dict) -> tuple[str, str, int | None]: return model, prompt, max_tokens -def handle_chat_completion_stream(payload: dict) -> tuple[int, dict | Iterator[bytes]]: +def handle_chat_completion_stream( + payload: dict, parent_task: str | None = None +) -> tuple[int, dict | Iterator[bytes]]: """Handle one ``stream: true`` ``POST /v1/chat/completions`` request body. Same validation and error mapping as :func:`handle_chat_completion`, but the request runs @@ -350,6 +381,8 @@ def handle_chat_completion_stream(payload: dict) -> tuple[int, dict | Iterator[b Args: payload: The parsed JSON request body (a dict). + parent_task: The sanitized ``X-TangleBrain-Parent-Task`` header value (or ``None``), + recorded onto the usage record for cross-system attribution (#74). Returns: ``(status, body)`` where a non-200 ``body`` is an OpenAI-style error dict (serialize as @@ -363,7 +396,10 @@ def handle_chat_completion_stream(payload: dict) -> tuple[int, dict | Iterator[b pinned = None if model == AUTO_ALIAS else model try: - deltas, served = run_once_stream(prompt, model=pinned, max_tokens=max_tokens) + deltas, served = run_once_stream( + prompt, model=pinned, max_tokens=max_tokens, + origin="serve", parent_task_id=parent_task, + ) first = next(deltas) except StopIteration: # Defensive: adapters raise on zero-content streams, and emulated paths always carry @@ -382,7 +418,7 @@ def handle_chat_completion_stream(payload: dict) -> tuple[int, dict | Iterator[b return 200, sse_stream_events(first, deltas, served, model, prompt) -def handle_chat_completion(payload: dict) -> tuple[int, dict]: +def handle_chat_completion(payload: dict, parent_task: str | None = None) -> tuple[int, dict]: """Handle one ``POST /v1/chat/completions`` request body. Resolves the model directive, flattens the messages, runs the request through @@ -397,6 +433,8 @@ def handle_chat_completion(payload: dict) -> tuple[int, dict]: Args: payload: The parsed JSON request body (a dict). + parent_task: The sanitized ``X-TangleBrain-Parent-Task`` header value (or ``None``), + recorded onto the usage record for cross-system attribution (#74). Returns: ``(status_code, body_dict)``. This is the non-streaming handler — a ``stream: true`` @@ -409,7 +447,10 @@ def handle_chat_completion(payload: dict) -> tuple[int, dict]: pinned = None if model == AUTO_ALIAS else model try: - text, served = run_once(prompt, model=pinned, max_tokens=max_tokens, return_served=True) + text, served = run_once( + prompt, model=pinned, max_tokens=max_tokens, return_served=True, + origin="serve", parent_task_id=parent_task, + ) except SelectionError as exc: if pinned is not None: # The only selection to fail on the pinned path is the id lookup itself. diff --git a/tests/test_cli.py b/tests/test_cli.py index 6c0fef3..fe00941 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -295,6 +295,19 @@ def test_gate_none_uses_setting(self): self.assertEqual(served["path"], "gate-local") RouterCls.assert_not_called() + def test_origin_and_parent_task_recorded(self): + # #74: origin defaults to "cli"; an explicit origin + caller identity land on the record. + fake_adapter = MagicMock() + fake_adapter.run.return_value = "x" + with patch("tanglebrain.cli.build_adapter", return_value=fake_adapter): + run_once("hi", local=True) + run_once("hi", local=True, origin="serve", parent_task_id="tc-42") + default, tagged = self._records() + self.assertEqual(default["origin"], "cli") + self.assertNotIn("parent_task_id", default) + self.assertEqual(tagged["origin"], "serve") + self.assertEqual(tagged["parent_task_id"], "tc-42") + def test_router_path_records_served_entry(self): # The router surfaces last_served; run_once records that tier/model. served = MagicMock() @@ -457,6 +470,27 @@ def test_unknown_model_raises_selection_error_at_call_time(self): with self.assertRaises(SelectionError): run_once_stream("hi", model="no-such-model") + def test_origin_and_parent_task_recorded_on_streamed_and_emulated_paths(self): + # #74: attribution rides every run_once_stream recording site. + streaming = self._streaming_adapter("a") + with patch("tanglebrain.cli.build_adapter", return_value=streaming): + deltas, _ = run_once_stream( + "hi", model="claude", roster_path=_pinned_roster(self), + origin="serve", parent_task_id="tc-42", + ) + list(deltas) + blocking = MagicMock(spec=["run"]) + blocking.run.return_value = "whole" + with patch("tanglebrain.cli.build_adapter", return_value=blocking): + deltas, _ = run_once_stream( + "hi", model="claude", roster_path=_pinned_roster(self), + origin="serve", parent_task_id="tc-42", + ) + list(deltas) + for record in self._records(): + self.assertEqual(record["origin"], "serve") + self.assertEqual(record["parent_task_id"], "tc-42") + def test_gate_not_consulted_for_model_or_local(self): fake = self._streaming_adapter("x") roster = _pinned_roster(self) diff --git a/tests/test_gui.py b/tests/test_gui.py index f6be222..ad7efb9 100644 --- a/tests/test_gui.py +++ b/tests/test_gui.py @@ -153,6 +153,7 @@ def test_happy_path_reports_served(self): self.assertEqual(out["served"]["model"], "claude") self.assertEqual(run.call_args.kwargs["task"], "code") self.assertTrue(run.call_args.kwargs["return_served"]) # uses the returned meta, no log re-read + self.assertEqual(run.call_args.kwargs["origin"], "gui") # #74 attribution def test_does_not_reread_log(self): # The race fix: run_prompt must NOT call read_records (served comes from run_once). diff --git a/tests/test_measurement.py b/tests/test_measurement.py index c7d1661..3c4526c 100644 --- a/tests/test_measurement.py +++ b/tests/test_measurement.py @@ -285,6 +285,55 @@ def test_adversarial_values_roundtrip(self): self.assertEqual(back.output_per_mtok, 1e20) +class OriginAttributionTest(unittest.TestCase): + """#74: the origin field on records, its rollup bucket, and the --stats line.""" + + def setUp(self): + self.tmp = tempfile.mkdtemp() + self.log = Path(self.tmp) / "usage.jsonl" + + def _record(self, **kwargs): + record_task( + path="model", entry=FakeEntry("m", "local"), prompt="p", response="r", + log_path=self.log, pricing=FIXED, **kwargs, + ) + + def test_record_writes_origin_when_given_and_omits_when_absent(self): + self._record(origin="serve") + self._record() + tagged, untagged = read_records(self.log) + self.assertEqual(tagged["origin"], "serve") + self.assertNotIn("origin", untagged) + + def test_record_writes_parent_task_id_on_task_records(self): + # #74: an external caller's identity (the serve header) rides on a kind="task" record. + self._record(parent_task_id="tc-session-42") + record = read_records(self.log)[0] + self.assertEqual(record["kind"], "task") + self.assertEqual(record["parent_task_id"], "tc-session-42") + + def test_rollup_buckets_by_origin_with_untagged_sentinel(self): + summary = rollup([ + {"tier": "local", "origin": "serve"}, + {"tier": "local", "origin": "serve"}, + {"tier": "sub", "origin": "cli"}, + {"tier": "sub"}, # pre-#74 record — never guessed at + {"kind": "delegate", "model": "m", "origin": "serve"}, # delegates stay out + ]) + self.assertEqual(summary["by_origin"], {"serve": 2, "cli": 1, "untagged": 1}) + + def test_format_rollup_shows_origin_split_only_when_tagged(self): + tagged = format_rollup( + rollup([{"tier": "local", "origin": "serve"}, {"tier": "local"}]), FIXED + ) + self.assertIn("By origin:", tagged) + self.assertIn("serve 1", tagged) + self.assertIn("untagged 1", tagged) + # All-untagged history says nothing — the line stays hidden. + untagged_only = format_rollup(rollup([{"tier": "local"}]), FIXED) + self.assertNotIn("By origin:", untagged_only) + + class FormatRollupTest(unittest.TestCase): def test_renders_figures(self): s = rollup([{"tier": "local", "in_tokens_est": 10, "out_tokens_est": 20, diff --git a/tests/test_serve.py b/tests/test_serve.py index 12bfdd9..6ca43c3 100644 --- a/tests/test_serve.py +++ b/tests/test_serve.py @@ -22,12 +22,14 @@ from tanglebrain.serve.server import Handler, dispatch from tanglebrain.serve.views import ( AUTO_ALIAS, + PARENT_TASK_HEADER, BadRequestError, completion_envelope, flatten_messages, handle_chat_completion, handle_chat_completion_stream, list_models, + sanitize_parent_task, sse_stream_events, wants_stream, ) @@ -414,6 +416,54 @@ def test_empty_stream_is_defensive_502(self): self.assertIn("before any content", body["error"]["message"]) +class ParentTaskAttributionTest(unittest.TestCase): + """#74: origin + X-TangleBrain-Parent-Task threading from transport to run_once*.""" + + def test_sanitize_parent_task(self): + self.assertIsNone(sanitize_parent_task(None)) + self.assertIsNone(sanitize_parent_task("")) + self.assertIsNone(sanitize_parent_task(" ")) + self.assertIsNone(sanitize_parent_task(42)) + self.assertEqual(sanitize_parent_task(" tc-42 "), "tc-42") + self.assertEqual(len(sanitize_parent_task("x" * 5000)), 128) # length-capped + + def test_plain_handler_threads_origin_and_parent_task(self): + run = MagicMock(return_value=("t", dict(_SERVED))) + with patch("tanglebrain.serve.views.run_once", run): + status, _ = handle_chat_completion(_chat_payload(), "tc-session-42") + self.assertEqual(status, 200) + self.assertEqual(run.call_args.kwargs["origin"], "serve") + self.assertEqual(run.call_args.kwargs["parent_task_id"], "tc-session-42") + + def test_plain_handler_defaults_parent_task_to_none(self): + run = MagicMock(return_value=("t", dict(_SERVED))) + with patch("tanglebrain.serve.views.run_once", run): + handle_chat_completion(_chat_payload()) + self.assertEqual(run.call_args.kwargs["origin"], "serve") + self.assertIsNone(run.call_args.kwargs["parent_task_id"]) + + def test_stream_handler_threads_origin_and_parent_task(self): + stream = MagicMock(return_value=(iter(["a"]), dict(_SERVED))) + with patch("tanglebrain.serve.views.run_once_stream", stream): + status, body = handle_chat_completion_stream( + _chat_payload(stream=True), "tc-session-42" + ) + list(body) + self.assertEqual(status, 200) + self.assertEqual(stream.call_args.kwargs["origin"], "serve") + self.assertEqual(stream.call_args.kwargs["parent_task_id"], "tc-session-42") + + def test_dispatch_sanitizes_the_raw_header_value(self): + run = MagicMock(return_value=("t", dict(_SERVED))) + payload = json.dumps(_chat_payload()).encode("utf-8") + with patch("tanglebrain.serve.views.run_once", run): + dispatch("POST", "/v1/chat/completions", payload, parent_task=" tc-42 ") + dispatch("POST", "/v1/chat/completions", payload, parent_task=" ") + first, second = run.call_args_list + self.assertEqual(first.kwargs["parent_task_id"], "tc-42") + self.assertIsNone(second.kwargs["parent_task_id"]) + + class WantsStreamTest(unittest.TestCase): def test_wants_stream_only_on_json_true(self): self.assertTrue(wants_stream({"stream": True})) @@ -587,6 +637,21 @@ def test_request_with_garbage_authorization_succeeds(self): self.assertEqual(body["choices"][0]["message"]["content"], "pong") self.assertEqual(body["model"], "claude") + def test_parent_task_header_reaches_routing_over_the_socket(self): + # #74 end-to-end: the real Handler reads X-TangleBrain-Parent-Task off the wire. + run = MagicMock(return_value=("pong", dict(_SERVED))) + request = urllib.request.Request( + f"http://127.0.0.1:{self.port}/v1/chat/completions", + data=json.dumps(_chat_payload()).encode("utf-8"), + headers={"Content-Type": "application/json", PARENT_TASK_HEADER: "tc-session-42"}, + method="POST", + ) + with patch("tanglebrain.serve.views.run_once", run): + with urllib.request.urlopen(request, timeout=5) as response: + self.assertEqual(response.status, 200) + self.assertEqual(run.call_args.kwargs["parent_task_id"], "tc-session-42") + self.assertEqual(run.call_args.kwargs["origin"], "serve") + def test_streaming_deltas_arrive_incrementally_over_the_socket(self): # The point of c13: the first content chunk must be readable while the backend is still # generating. The second delta is gated on an Event the test only sets AFTER it has read From 7299ba8f742fdac6c923ccd809028aa5804713a2 Mon Sep 17 00:00:00 2001 From: Jason-Vaughan <95194903+Jason-Vaughan@users.noreply.github.com> Date: Sat, 4 Jul 2026 00:23:53 -0700 Subject: [PATCH 2/2] Address Independent Critic findings on #74 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Panel now shows the per-origin split (the issue names the panel as a contract surface), hidden while all history is untagged — same stance as the CLI's format_rollup. - sanitize_parent_task drops non-printable characters (folded-header \r\n, ANSI escapes) so no control bytes ever reach usage records; tested. - Router-path recording site in run_once_stream now has the same attribution assertion as the other two sites. --- tanglebrain/gui/static/index.html | 6 ++++++ tanglebrain/serve/views.py | 12 +++++++----- tests/test_cli.py | 14 +++++++++++++- tests/test_serve.py | 5 +++++ 4 files changed, 31 insertions(+), 6 deletions(-) diff --git a/tanglebrain/gui/static/index.html b/tanglebrain/gui/static/index.html index 53cd300..3e3cc27 100644 --- a/tanglebrain/gui/static/index.html +++ b/tanglebrain/gui/static/index.html @@ -109,10 +109,16 @@

Pricing reference

const d = await getJSON("/api/stats"); const s = d.summary || {}; const tiers = Object.entries(s.by_tier || {}).map(([k, v]) => `${esc(k)} ${v}`).join(", ") || "—"; + // Origin split (#74): shown only once some record is actually tagged — all-untagged history + // adds no signal (mirrors the CLI's format_rollup stance). + const byOrigin = s.by_origin || {}; + const hasOrigin = Object.keys(byOrigin).some((k) => k !== "untagged"); + const origins = Object.entries(byOrigin).map(([k, v]) => `${esc(k)} ${v}`).join(", "); let html = `
Tasks routed
${s.tasks || 0}
Spend avoided
${money(s.spend_avoided_usd)}
By tier
${tiers}
+ ${hasOrigin ? `
By origin
${origins}
` : ""}
Est. tokens (in / out)
${(s.in_tokens_est||0).toLocaleString()} / ${(s.out_tokens_est||0).toLocaleString()}
`; if (d.is_placeholder) html += `
⚠ pricing: PLACEHOLDER — figures illustrative until the anchor is ratified.
`; diff --git a/tanglebrain/serve/views.py b/tanglebrain/serve/views.py index 2b7cd17..84e0aa8 100644 --- a/tanglebrain/serve/views.py +++ b/tanglebrain/serve/views.py @@ -62,10 +62,12 @@ def sanitize_parent_task(value: object) -> str | None: """Sanitize a raw ``X-TangleBrain-Parent-Task`` header value for recording (#74). - Attribution metadata only, so the stance is trim-don't-reject: whitespace is stripped, an - empty/absent/non-string value becomes ``None`` (field omitted from the record), and anything - longer than 128 chars is truncated — a caller cannot stuff arbitrary payloads into the - usage log through this header. + Attribution metadata only, so the stance is trim-don't-reject: whitespace is stripped, + non-printable characters are dropped (the stdlib header parser accepts folded values, so raw + ``\\r\\n`` — and ANSI escapes — could otherwise ride into records and bite any future + consumer that prints the field to a terminal), an empty/absent/non-string value becomes + ``None`` (field omitted from the record), and anything longer than 128 chars is truncated — + a caller cannot stuff arbitrary payloads into the usage log through this header. Args: value: The raw header value (or ``None`` when the header is absent). @@ -75,7 +77,7 @@ def sanitize_parent_task(value: object) -> str | None: """ if not isinstance(value, str): return None - cleaned = value.strip() + cleaned = "".join(ch for ch in value.strip() if ch.isprintable()) if not cleaned: return None return cleaned[:_PARENT_TASK_MAX_LEN] diff --git a/tests/test_cli.py b/tests/test_cli.py index fe00941..cdf9902 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -487,7 +487,19 @@ def test_origin_and_parent_task_recorded_on_streamed_and_emulated_paths(self): origin="serve", parent_task_id="tc-42", ) list(deltas) - for record in self._records(): + served_entry = MagicMock(); served_entry.tier = "sub"; served_entry.id = "codex" + fake_router = MagicMock() + fake_router.route.return_value = "routed" + fake_router.last_served = served_entry + with patch("tanglebrain.cli.load_roster"), \ + patch("tanglebrain.cli.Router", return_value=fake_router): + deltas, _ = run_once_stream( + "hi", gate=False, origin="serve", parent_task_id="tc-42" + ) + list(deltas) + records = self._records() + self.assertEqual(len(records), 3) # streamed pin + emulated pin + emulated router + for record in records: self.assertEqual(record["origin"], "serve") self.assertEqual(record["parent_task_id"], "tc-42") diff --git a/tests/test_serve.py b/tests/test_serve.py index 6ca43c3..f2fd58d 100644 --- a/tests/test_serve.py +++ b/tests/test_serve.py @@ -426,6 +426,11 @@ def test_sanitize_parent_task(self): self.assertIsNone(sanitize_parent_task(42)) self.assertEqual(sanitize_parent_task(" tc-42 "), "tc-42") self.assertEqual(len(sanitize_parent_task("x" * 5000)), 128) # length-capped + # Control chars are dropped, not recorded: header folding can smuggle \r\n through the + # stdlib parser, and ANSI escapes would bite any consumer printing the field raw. + self.assertEqual(sanitize_parent_task("a\r\nb"), "ab") + self.assertEqual(sanitize_parent_task("\x1b[31mevil\x1b[0m"), "[31mevil[0m") + self.assertIsNone(sanitize_parent_task("\x1b\x00\x07")) # nothing printable left def test_plain_handler_threads_origin_and_parent_task(self): run = MagicMock(return_value=("t", dict(_SERVED)))