From b9e0bb33334de35d641e5077d545be1bd2eb9486 Mon Sep 17 00:00:00 2001 From: 0x-copilot-dev <0x-copilot-dev@users.noreply.github.com> Date: Fri, 14 Aug 2026 09:58:43 +0530 Subject: [PATCH] fix(ai-backend): declare the context nobody was declaring MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A live packaged run recorded `undeclared_tokens: 8257` on every model call — 31% of measured input — through a field whose contract says the number is EXPECTED 0 and anything above it is a first-party defect. PR #627 had just stopped journey CB-8 failing on it, and the composer's context meter has been showing users an "N undeclared" notice since #625. It was three unrelated defects, and the largest row was not a tool at all. **Ten library-installed tools, ~3,459 tokens.** The declaration seam stamps tools at `factory._model_visible_tools`, and the PRD-02 AST gate sweeps that same function. But `create_deep_agent` installs the filesystem, todo and subagent middleware tools INSIDE the library, so they reach the provider having passed no append site: no stamp to read, and nothing for the gate to see. A gate scoped to one function cannot fail on contributors that never enter it. `ThirdPartyToolOrigins` now declares them on their author's behalf — §4.3's sanctioned exception, applied to the tool block the way it already was to the system block — resolving the owner from the tool's dispatch callable's `__module__` and injected as `ToolSchemaLedger.measure(fallback_origin=...)`. Resolution is by that callable and never by the tool's class or name. The class is `StructuredTool` whoever wrote the tool, so `type(tool).__module__` would relabel OUR tools as third-party `langchain_core` text. And the allowlist names the libraries we speak for rather than excluding the packages we own, because "unknown means theirs" fails by silent misattribution: `task` looks exactly like a deepagents built-in — same name, same middleware slot — but is `build_atlas_task_tool`, ours, and would have billed ~390 resident tokens of our text to a dependency that does not contain the string. It is declared at its authoring site instead, as is `web_search`. **`web_search`, ~198 tokens.** Declared correctly at the factory and destroyed downstream. `wrap_tools_with_display` copies a `StructuredTool` (the stamp rides along in the instance `__dict__`) but REBUILDS any other `BaseTool` via `from_function`, holding the inner in a closure — no stamp, and no attribute chain for `ContextOriginBinding.of` to walk. Every non-`StructuredTool` on the surface was silently un-declared on its way to the model. The declaration is now carried across the rebuild, so it is still made once, by the code that composed the text. **The system block, 4,798 tokens — 58% of the total, and not a tool.** Every row a packaged run wrote carries `assembly_record_id: None`: F2 prompt assembly does not run on the desktop, so there is no `PromptAssemblyPlan`, and the whole prompt fell through as one unattributed span. Without a plan the block has exactly one first-party contributor — the `system_prompt` the factory hands the builder — so naming it is a more accurate report than calling it undeclared. Keyed strictly on `plan is None`: a plan that exists and matches nothing means the plan and the prompt disagree, which IS a defect and still reports UNDECLARED. What is genuinely lost is the per-fragment breakdown, and that is carried in the segment's `detail` (`system[unassembled]`) rather than smuggled into a field whose contract is about broken declarations. The new tests are built on the real middleware and the real factory-plus-display chain, not on fixture tools — a hand-built tool list is declared by construction and can never notice that the live surface carries objects nobody stamped, which is why 9.5k green tests sat over this. Half the file is the counterweight: an undeclared first-party tool must STAY undeclared, or the fix would be an amnesty that zeroes the number by destroying the only signal the field carries. Verified end-to-end through the recorder against the pinned deepagents 0.7.4: a full desktop-shaped request reports `undeclared_tokens: 0`, with every row naming an owner someone can act on. Co-Authored-By: Claude Opus 5 --- .../middleware/display_metadata.py | 53 +- .../delegation/subagents/atlas_task_tool.py | 44 +- .../context_occupancy_recorder.py | 134 +++- .../observability/context_origin.py | 44 ++ .../observability/context_third_party.py | 171 +++++ .../observability/context_tool_ledger.py | 69 +- .../src/runtime_worker/dependencies.py | 38 +- .../test_context_occupancy_recorder.py | 56 +- .../test_context_tool_declaration_coverage.py | 672 ++++++++++++++++++ .../tests/unit/test_context_origin_gate.py | 23 + 10 files changed, 1266 insertions(+), 38 deletions(-) create mode 100644 services/ai-backend/tests/unit/agent_runtime/observability/test_context_tool_declaration_coverage.py diff --git a/services/ai-backend/src/agent_runtime/capabilities/middleware/display_metadata.py b/services/ai-backend/src/agent_runtime/capabilities/middleware/display_metadata.py index 05e5e3501..2e9174ef1 100644 --- a/services/ai-backend/src/agent_runtime/capabilities/middleware/display_metadata.py +++ b/services/ai-backend/src/agent_runtime/capabilities/middleware/display_metadata.py @@ -3,6 +3,7 @@ from __future__ import annotations import inspect +import logging from collections.abc import Mapping from typing import Annotated, Any, ClassVar, Callable, get_type_hints @@ -12,6 +13,9 @@ from agent_runtime.capabilities.tools.cards import ToolDisplayTemplate + +_LOGGER = logging.getLogger(__name__) + # Wire keys (alias form) for optional agent-supplied display overrides. # The model emits these in tool_call args; the presentation layer reads them # off ``payload.args`` and applies them only when the matched template is @@ -603,6 +607,17 @@ def wrap_tool_with_display(tool: object) -> object: Idempotent: a tool already bearing the ``__display_wrapped__`` schema marker is returned unchanged. Falls back to returning the original tool for unrecognised shapes — the safety contract is "never break a working tool to add display copy." + + The wrapped tool inherits the inner's Context Occupancy declaration. The + ``StructuredTool`` branch keeps it for free (the stamp lives in the instance + ``__dict__``, which ``model_copy`` carries), but the delegation branch + *rebuilds* the tool with ``from_function`` and holds the inner in a closure, + so the new object had neither the stamp nor an attribute chain leading back + to it. That silently un-declared every non-``StructuredTool`` on the model + surface — ``web_search``, which arrives wrapped in ``RetryingTool`` — and + the occupancy report blamed a composition site that had declared it + correctly. Carrying it here rather than re-declaring downstream keeps the + rule that a declaration is made once, by the code that composed the text. """ args_schema = getattr(tool, "args_schema", None) @@ -615,12 +630,46 @@ def wrap_tool_with_display(tool: object) -> object: return tool if isinstance(tool, StructuredTool): - return _wrap_structured_tool(tool, StructuredTool) + return _with_carried_context_origin( + tool, _wrap_structured_tool(tool, StructuredTool) + ) if isinstance(tool, BaseTool): - return _wrap_base_tool_via_delegation(tool, StructuredTool) + return _with_carried_context_origin( + tool, _wrap_base_tool_via_delegation(tool, StructuredTool) + ) return tool +def _with_carried_context_origin(source: object, wrapped: object) -> object: + """Preserve ``source``'s occupancy declaration on the tool that replaces it. + + Deferred import for the same reason ``ModelToolDeclaration`` defers its own: + ``capabilities.middleware`` is imported eagerly from the execution lane, and + a module-scope edge into the observability lane would make importing either + package order-dependent. + + Fail-open, because declaring an origin is observability and this function is + on the graph-build path: a binding that cannot be read or written costs a + label in a report, never a working tool (§6.4). + """ + + if wrapped is source: + return wrapped + try: + from agent_runtime.observability.context_origin import ( # noqa: PLC0415 + carry_context_origin, + ) + + return carry_context_origin(source, wrapped) + except Exception: # noqa: BLE001 — a label is never worth breaking a tool + _LOGGER.debug( + "Could not carry a context origin onto a display-wrapped tool; " + "it will measure as UNDECLARED.", + exc_info=True, + ) + return wrapped + + def wrap_tools_with_display(tools: Any) -> list[object]: """Apply :func:`wrap_tool_with_display` to every tool and return a new list.""" diff --git a/services/ai-backend/src/agent_runtime/delegation/subagents/atlas_task_tool.py b/services/ai-backend/src/agent_runtime/delegation/subagents/atlas_task_tool.py index 36d6ecd28..a6693a62c 100644 --- a/services/ai-backend/src/agent_runtime/delegation/subagents/atlas_task_tool.py +++ b/services/ai-backend/src/agent_runtime/delegation/subagents/atlas_task_tool.py @@ -46,6 +46,13 @@ from langchain_core.tools import StructuredTool from langgraph.types import Command +from agent_runtime.observability.context_origin import ( + ContextLifecycle, + ContextOrigin, + ContextSegmentClass, + declare_context_origin, +) + # Re-import the pieces we mirror so we stay 1:1 with the upstream behavior # except for the metadata injection. from deepagents.middleware.subagents import ( # type: ignore[import-untyped] @@ -361,13 +368,36 @@ async def _dispatch() -> object: return str(result) return _return_command_with_state_update(result, runtime.tool_call_id) - return StructuredTool.from_function( - name="task", - func=task, - coroutine=atask, - description=description, - infer_schema=False, - args_schema=TaskToolSchema, + # Declared here because this is where the text is authored, and there is + # nowhere else it could be. ``task`` does not travel through + # ``factory._model_visible_tools`` — it is handed to the subagent middleware + # and installed by ``create_deep_agent`` — so no append site stamps it and + # the PRD-02 AST gate, which sweeps that one function, cannot see it either. + # It reached the provider undeclared on every model call of every run that + # can delegate, and it is easy to mistake for the library's own ``task``: + # the tool it replaces has the same name and the same slot, so an occupancy + # reader who assumed third-party would go looking in ``deepagents`` for a + # description this repository writes. + return declare_context_origin( + StructuredTool.from_function( + name="task", + func=task, + coroutine=atask, + description=description, + infer_schema=False, + args_schema=TaskToolSchema, + ), + ContextOrigin( + # A literal rather than ``__name__`` so the PRD-02 inventory gate can + # read it statically. That gate is the only review this declaration + # ever gets — it sits outside ``_model_visible_tools``, which is the + # one place the sweep looks — and an owner it cannot resolve is an + # entry that never appears in the pin for anyone to review. + owner="agent_runtime.delegation.subagents.atlas_task_tool", + name="task", + segment_class=ContextSegmentClass.TOOLS, + lifecycle=ContextLifecycle.RESIDENT, + ), ) diff --git a/services/ai-backend/src/agent_runtime/observability/context_occupancy_recorder.py b/services/ai-backend/src/agent_runtime/observability/context_occupancy_recorder.py index 0f354d99c..80acbc41b 100644 --- a/services/ai-backend/src/agent_runtime/observability/context_occupancy_recorder.py +++ b/services/ai-backend/src/agent_runtime/observability/context_occupancy_recorder.py @@ -20,10 +20,14 @@ **Three attribution paths, one vocabulary.** The system block is matched against the typed :class:`~agent_runtime.prompts.assembly.PromptAssemblyPlan` by content digest, so each fragment's own ``source_owner`` labels its span; whatever the -plan does not explain is offered to the pinned third-party adapter (§4.3) and, -failing that, recorded as ``UNDECLARED``. The tool block goes through +plan does not explain is offered to the pinned third-party adapter (§4.3), and +what survives that is ``UNDECLARED`` when a plan existed or the build-time +system prompt when none did. The tool block goes through :class:`~agent_runtime.observability.context_tool_ledger.ToolSchemaLedger`, which -reads the declaration each tool was stamped with at composition. Messages go +reads the declaration each tool was stamped with at composition and, for the +middleware tools the library installs behind our back, falls back to +:class:`~agent_runtime.observability.context_third_party.ThirdPartyToolOrigins`. +Messages go through :class:`~agent_runtime.observability.context_message_classifier.ContextMessageClassifier`, whose structural rules resolve to the same declared origins. Three very @@ -47,7 +51,7 @@ from __future__ import annotations from collections.abc import Callable, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, replace import hashlib import json import logging @@ -71,7 +75,10 @@ ContextSegmentClass, ContextTextWidth, ) -from agent_runtime.observability.context_third_party import ThirdPartyContextOrigins +from agent_runtime.observability.context_third_party import ( + ThirdPartyContextOrigins, + ThirdPartyToolOrigins, +) from agent_runtime.observability.context_token_counter import ( ContextTokenCounter, TokenCounterSource, @@ -146,6 +153,41 @@ class RuntimeContextOrigins: lifecycle=ContextLifecycle.RESIDENT, ) + BUILD_TIME_SYSTEM_PROMPT: Final[ContextOrigin] = ContextOrigin( + owner="agent_runtime.execution.factory", + name="system_prompt", + segment_class=ContextSegmentClass.SYSTEM, + lifecycle=ContextLifecycle.RESIDENT, + ) + """The system prompt handed to ``create_deep_agent`` when F2 assembled none. + + Per-fragment attribution needs a :class:`PromptAssemblyPlan`, and there is + no plan unless the F2 prompt runtime is bound — which it is not on the + desktop, where ``RunControlContext.current()`` yields no control binding. + Every occupancy row a packaged run wrote carried ``assembly_record_id: + None`` and, consequently, a single 4,798-token ``UNDECLARED`` span covering + the whole system block. That is the largest number in the report, it is + constant on every call, and it was being announced to users through the + composer's context meter as a first-party contract defect. + + It is not one. Without a plan the system block has exactly one first-party + contributor — the ``system_prompt`` argument the factory builds and passes + to the builder — so naming it is a *more* accurate report than leaving it + undeclared, not a papering-over. What is genuinely lost is the per-fragment + breakdown, and that loss is carried in the segment's ``detail`` + (``system[unassembled]``) rather than smuggled into ``undeclared_tokens``, + whose contract is about broken declarations and not about which prompt path + a deployment runs. + + Applied **only** when no plan exists at all. A residue left over *after* a + plan was matched is a real attribution gap — the plan and the prompt + disagree — and stays ``UNDECLARED``, which is the case that field is for. + + ``cache_eligibility`` is unset for the same reason it is on the joiner: the + field records a declared intent read from ``PromptFragment`` metadata, and + an unassembled block carries no such metadata to read. + """ + ASSEMBLY_JOINER: Final[ContextOrigin] = ContextOrigin( owner="agent_runtime.prompts", name="assembly_joiner", @@ -328,6 +370,7 @@ class SystemBlockAttributor: ) _UNDECLARED_DETAIL: Final[str] = "system[unattributed]" + _UNASSEMBLED_DETAIL: Final[str] = "system[unassembled]" _JOINER_DETAIL: Final[str] = "system[joiners]" def __init__(self, *, third_party: ThirdPartyPromptIndex) -> None: @@ -343,13 +386,20 @@ def spans( Total: a plan that does not match, a fragment whose owner is not a legal declaration, and a third-party sweep that found nothing all reduce the - result toward one ``UNDECLARED`` span covering the whole block. That is - the honest degradation — the bytes are still reported, they are simply - reported as unexplained. + result toward one span covering the whole block. That is the honest + degradation — the bytes are still reported, they are simply reported at + the coarsest granularity the available evidence supports. + + ``plan is None`` and ``plan matched nothing`` are deliberately *not* the + same outcome. With no plan there was no per-fragment assembly to fail, + so the block is the build-time system prompt and is declared as such. + With a plan that matched nothing, the plan and the prompt disagree — + a real attribution defect — and the residue stays ``UNDECLARED``. """ if not system_text: return () + assembled = plan is not None try: matched = self._plan_spans(system_text, plan=plan) except Exception: # noqa: BLE001 — measurement never fails a run (§6.4) @@ -359,7 +409,7 @@ def spans( exc_info=True, ) matched = () - return self._fill_gaps(system_text, matched=matched) + return self._fill_gaps(system_text, matched=matched, assembled=assembled) def _plan_spans( self, @@ -399,6 +449,7 @@ def _fill_gaps( system_text: str, *, matched: tuple[tuple[int, int, SystemSpan], ...], + assembled: bool, ) -> tuple[SystemSpan, ...]: """Interleave the matched spans with what sits between and after them. @@ -415,11 +466,17 @@ def _fill_gaps( cursor = 0 for start, end, span in matched: spans.extend( - self._unmatched_spans(system_text[cursor:start], joiners=joiners) + self._unmatched_spans( + system_text[cursor:start], joiners=joiners, assembled=assembled + ) ) spans.append(span) cursor = end - spans.extend(self._unmatched_spans(system_text[cursor:], joiners=joiners)) + spans.extend( + self._unmatched_spans( + system_text[cursor:], joiners=joiners, assembled=assembled + ) + ) if joiners: spans.append(self._joiner_span(joiners)) return tuple(spans) @@ -429,6 +486,7 @@ def _unmatched_spans( text: str, *, joiners: list[str], + assembled: bool, ) -> tuple[SystemSpan, ...]: """Attribute leftover system bytes, collecting assembly joiners aside. @@ -437,7 +495,13 @@ def _unmatched_spans( :attr:`RuntimeContextOrigins.ASSEMBLY_JOINER` for why that is a declaration rather than an ``UNDECLARED`` row. Anything else is real text somebody contributed, and is offered to the third-party adapter - before it is admitted as undeclared. + before anything else is decided about it. + + The library's own constants are peeled off first either way — a + ``deepagents`` prompt inside an unassembled block still belongs to + ``deepagents``, and sweeping it into a first-party bucket would trade + one misattribution for another. Only the residue of that split is + affected by ``assembled``. """ if not text: @@ -445,7 +509,28 @@ def _unmatched_spans( if not text.strip(): joiners.append(text) return () - return self._third_party.split(text, undeclared_detail=self._UNDECLARED_DETAIL) + split = self._third_party.split(text, undeclared_detail=self._UNDECLARED_DETAIL) + if assembled: + return split + return tuple(self._as_build_time_prompt(span) for span in split) + + @classmethod + def _as_build_time_prompt(cls, span: SystemSpan) -> SystemSpan: + """Declare an unassembled residue against the build-time system prompt. + + Only an undeclared span is rewritten. A span the third-party index + already claimed keeps its library owner, because that attribution is + evidence-based — the bytes matched a pinned constant — and is strictly + better than the coarse first-party bucket. + """ + + if span.origin is not None: + return span + return replace( + span, + origin=RuntimeContextOrigins.BUILD_TIME_SYSTEM_PROMPT, + detail=cls._UNASSEMBLED_DETAIL, + ) @classmethod def _joiner_span(cls, joiners: Sequence[str]) -> SystemSpan: @@ -739,12 +824,19 @@ def __init__( *, counter: ContextTokenCounter | None = None, third_party: ThirdPartyPromptIndex | None = None, + third_party_tools: ThirdPartyToolOrigins | None = None, builder: SnapshotBuilder | None = None, ) -> None: self._counter = counter or ContextTokenCounter() self._third_party = third_party or ThirdPartyPromptIndex( origins=ThirdPartyContextOrigins() ) + # The tool-block counterpart of ``_third_party``. Separate collaborators + # because they resolve different things by different means — one indexes + # the dependency's prompt *text*, the other reads a tool's authoring + # module — and a test that wants a third-party-free system block still + # wants real tool attribution. + self._third_party_tools = third_party_tools or ThirdPartyToolOrigins() self._builder = builder or SnapshotBuilder() self._system = SystemBlockAttributor(third_party=self._third_party) @@ -1184,12 +1276,26 @@ def _tool_segments( prompt-cache identity cannot drift apart. Its counting is injected here so the numbers come from the real tokenizer chain rather than from the char/4 stand-in it defaults to. + + ``fallback_origin`` is §4.3 applied to the tool block, and it is what + closes the gap that made ``undeclared_tokens`` permanently non-zero. + ``materialized.tools`` is the *library's* final tool list, not ours: + ``create_deep_agent`` installs the filesystem, todo and subagent + middleware tools inside itself, so those nine reach the wire having + never passed the one composition site that stamps a declaration. The + stamp still wins wherever one exists — this only speaks for tools that + have none, and only when their authoring module is outside this + repository. """ if not materialized.tools: return () bridge = _ToolSchemaCounterBridge(counter=self._counter, model=model) - footprints = ToolSchemaLedger.measure(materialized.tools, counter=bridge) + footprints = ToolSchemaLedger.measure( + materialized.tools, + counter=bridge, + fallback_origin=self._third_party_tools.origin_for, + ) sources = bridge.sources_for(footprints) return tuple( self._segment_for_footprint(footprint, source=source) diff --git a/services/ai-backend/src/agent_runtime/observability/context_origin.py b/services/ai-backend/src/agent_runtime/observability/context_origin.py index 7c0f409f1..58baefb72 100644 --- a/services/ai-backend/src/agent_runtime/observability/context_origin.py +++ b/services/ai-backend/src/agent_runtime/observability/context_origin.py @@ -231,6 +231,38 @@ def declare(cls, target: object, origin: ContextOrigin) -> object: ) return target + @classmethod + def carry_over(cls, source: object, target: object) -> object: + """Copy ``source``'s declaration onto ``target`` when ``target`` has none. + + The escape hatch for a re-wrapper that *rebuilds* a tool instead of + wrapping it. :meth:`of` walks the wrapper chain, which covers every + adapter that keeps its inner under one of + :attr:`WRAPPER_ATTRIBUTES` — but a wrapper built with + ``StructuredTool.from_function`` holds its inner in a **closure**, so + there is no attribute to walk and the new object is genuinely + undeclared. + + That is not a hypothetical shape. ``wrap_tools_with_display`` takes + exactly that branch for a tool that is a ``BaseTool`` but not a + ``StructuredTool``, which is what ``web_search`` is once + ``RetryingTool`` has wrapped it: the factory stamped it correctly, the + display decoration returned a fresh tool, and the declaration was gone + by the time the tool block was measured. It read as an undeclared + first-party contributor — the right alarm for the wrong reason, since + the composition site had done its job. + + ``target`` wins when it already carries a declaration: a rebuilt tool + that declared itself knows better than the thing it replaced. + """ + + if cls.of(target) is not None: + return target + declared = cls.of(source) + if declared is None: + return target + return cls.declare(target, declared) + @classmethod def of(cls, target: object) -> ContextOrigin | None: """Return the declaration on ``target`` or its wrappers, or ``None``. @@ -280,6 +312,17 @@ def context_origin_of(target: object) -> ContextOrigin | None: return ContextOriginBinding.of(target) +def carry_context_origin(source: object, target: object) -> object: + """Preserve ``source``'s declaration across a rebuild into ``target``. + + Module-level seam over :meth:`ContextOriginBinding.carry_over`, so a + re-wrapping site outside the observability lane keeps its declaration + handling to one obvious call and does not have to know the binding class. + """ + + return ContextOriginBinding.carry_over(source, target) + + __all__ = ( "UNDECLARED_CONTEXT_LABEL", "ContextLifecycle", @@ -287,6 +330,7 @@ def context_origin_of(target: object) -> ContextOrigin | None: "ContextOriginBinding", "ContextSegmentClass", "ContextTextWidth", + "carry_context_origin", "context_origin_of", "declare_context_origin", ) diff --git a/services/ai-backend/src/agent_runtime/observability/context_third_party.py b/services/ai-backend/src/agent_runtime/observability/context_third_party.py index b26420388..292f28e2f 100644 --- a/services/ai-backend/src/agent_runtime/observability/context_third_party.py +++ b/services/ai-backend/src/agent_runtime/observability/context_third_party.py @@ -670,8 +670,179 @@ def _estimate_tokens(cls, byte_count: int) -> int: return (byte_count + cls.CHARS_PER_TOKEN - 1) // cls.CHARS_PER_TOKEN +class ThirdPartyToolOrigins: + """Declare library-installed tools on their authoring module's behalf (§4.3). + + The tool half of what :class:`ThirdPartyContextOrigins` does for prompt text, + and it exists because the declaration seam has exactly one first-party + composition site. ``ModelToolDeclaration`` stamps an origin at + ``factory._model_visible_tools``, and the PRD-02 AST gate sweeps that same + function — so a tool composed anywhere else is invisible to *both*. That is + not hypothetical: ``create_deep_agent`` installs its own middleware tools + inside the library, and the ``FilesystemMiddleware`` (``ls`` / ``read_file`` + / ``write_file`` / ``edit_file`` / ``delete`` / ``glob`` / ``grep``), the + ``TodoListMiddleware`` (``write_todos``) and the ``SubAgentMiddleware`` + (``task``) never pass an append site. They carried no stamp, the gate could + not see them, and they measured as ``UNDECLARED`` — thousands of estimated + tokens of resident rent per model call, reported through a field whose + contract says "first-party contract defect" about text no first-party edit + could fix. + + **Resolution is by the tool's authoring callable, never by its class and + never by its name.** A name table would be the central registry §3.2 + rejects, stale the moment the dependency adds a tool. The class is worse + than useless: ``StructuredTool`` is a ``langchain_core`` class whoever wrote + the tool, so ``type(tool).__module__`` would relabel *our* tools as + third-party and quietly delete the very signal this module protects. The + function a tool dispatches to is defined in the module that actually wrote + its description and schema, which is both the honest owner and exactly the + dotted shape :attr:`ContextOrigin.owner` wants. + + **Only the pinned libraries named below are declarable, and everything else + resolves to ``None``.** The tempting inversion — declare anything whose + module is not one of this service's own packages — reads as equivalent and + is not. It assumes *unknown means theirs*, so the failure mode is silent + misattribution: a first-party module missing from the list, or a tool built + in a module nobody thought about, gets stamped ``third_party=True`` and the + report sends a reader off to edit a dependency that does not contain the + string. This direction fails the other way, which is the one §4.4 wants — an + unrecognised author lands in ``undeclared_tokens``, where it is loud, + reviewable, and obviously ours to fix. + + That is not a theoretical preference. ``task`` looks exactly like a + ``deepagents`` built-in — same name, same middleware slot — but this + repository substitutes its own + :func:`~agent_runtime.delegation.subagents.atlas_task_tool.build_atlas_task_tool` + into that slot, and a resolver working off "not ours" would have labelled + ~390 tokens of text we wrote as the library's. + + Never raises, like everything else on this path (§6.4). An unreadable + callable, a missing ``__module__`` and a name that cannot form a legal label + all degrade to ``None``, which is the pre-existing ``UNDECLARED`` behaviour. + """ + + class Keys: + """Attribute names read off an untrusted tool object. + + The two callable slots are the ones LangChain fills from + ``StructuredTool.from_function`` and the ``@tool`` decorator, which is + how every library-installed tool in the pinned dependencies is built. + ``func`` is checked first because a tool that has both dispatches the + sync one from the same module. + """ + + FUNC: Final[str] = "func" + COROUTINE: Final[str] = "coroutine" + MODULE: Final[str] = "__module__" + NAME: Final[str] = "name" + + AUTHOR_ATTRIBUTES: Final[tuple[str, ...]] = (Keys.FUNC, Keys.COROUTINE) + + MODULE_SEPARATOR: Final[str] = "." + + DECLARABLE_ROOTS: Final[frozenset[str]] = frozenset( + { + "deepagents", + "langchain", + "langchain_core", + "langgraph", + } + ) + """The pinned packages this module is allowed to speak for. + + §4.3 grants exactly one exception to "contributors declare themselves", and + it is scoped to the dependency: *one module declares on the library's + behalf*. An allowlist is that scope written down. These four are the + packages that can put a tool on the model surface without passing through + any code of ours — ``deepagents`` and ``langchain`` do so today via their + middleware, and the other two are the same lane's lower layers. + + A dependency that moves its tools into a new top-level package makes them + measure as ``UNDECLARED`` until that package is added here. That is the + intended behaviour, not a gap: it is one line to fix, the occupancy report + names the tool, and the alternative — assuming any unfamiliar module is + somebody else's — is how first-party text gets quietly attributed to a + vendor. + """ + + def origin_for(self, tool: object) -> ContextOrigin | None: + """Declare ``tool`` on its authoring module's behalf, or ``None``. + + ``None`` means "this resolver has nothing to say", which leaves the + caller's existing ``UNDECLARED`` path untouched. That is the answer for + a first-party tool, for a tool authored in any package outside + :attr:`DECLARABLE_ROOTS`, for a tool whose authoring callable cannot be + read, and for a tool whose name cannot form a legal label. + """ + + module = self._authoring_module(tool) + if module is None or not self._is_declarable(module): + return None + name = self._tool_name(tool) + if not name: + return None + try: + return ContextOrigin( + owner=module, + name=name, + segment_class=ContextSegmentClass.TOOLS, + lifecycle=ContextLifecycle.RESIDENT, + third_party=True, + ) + except Exception: # noqa: BLE001 — a declaration is never load-bearing + _LOGGER.debug( + "Library-installed tool %r in %s does not form a legal " + "declaration; it will measure as UNDECLARED.", + name, + module, + exc_info=True, + ) + return None + + @classmethod + def _authoring_module(cls, tool: object) -> str | None: + """Return the dotted module that defined ``tool``'s dispatch callable. + + Deliberately no fallback to ``type(tool).__module__``. See the class + docstring: the class is a library class for first-party and third-party + tools alike, so falling back to it would misattribute ours rather than + degrade — and a wrong owner is worse than an absent one, because + ``undeclared_tokens`` would go quiet while the report named the wrong + package to go fix. + """ + + for attribute in cls.AUTHOR_ATTRIBUTES: + try: + callable_ = getattr(tool, attribute, None) + except Exception: # noqa: BLE001 — unreadable is treated as absent + continue + if callable_ is None: + continue + module = getattr(callable_, cls.Keys.MODULE, None) + if isinstance(module, str) and module: + return module + return None + + @classmethod + def _is_declarable(cls, module: str) -> bool: + """Whether ``module`` belongs to a package this adapter speaks for.""" + + root = module.split(cls.MODULE_SEPARATOR, 1)[0] + return root in cls.DECLARABLE_ROOTS + + @staticmethod + def _tool_name(tool: object) -> str: + """Read the tool's model-facing name, or ``""`` when unreadable.""" + + try: + return str(getattr(tool, ThirdPartyToolOrigins.Keys.NAME, "")).strip() + except Exception: # noqa: BLE001 — an unreadable name is an empty name + return "" + + __all__ = ( "ThirdPartyContextOriginRegistry", "ThirdPartyContextOrigins", "ThirdPartyPromptConstant", + "ThirdPartyToolOrigins", ) diff --git a/services/ai-backend/src/agent_runtime/observability/context_tool_ledger.py b/services/ai-backend/src/agent_runtime/observability/context_tool_ledger.py index 52a2764ff..ba5b68936 100644 --- a/services/ai-backend/src/agent_runtime/observability/context_tool_ledger.py +++ b/services/ai-backend/src/agent_runtime/observability/context_tool_ledger.py @@ -70,6 +70,23 @@ """ +ToolOriginResolver = Callable[[object], ContextOrigin | None] +"""Declares a tool that carries no stamp, or returns ``None`` to leave it alone. + +The second and last way a tool can acquire a declaration, and it exists because +the stamp can only be applied where a tool is *composed by us*. Library +middleware installs its own tools inside ``create_deep_agent``, so there is no +append site to stamp and no lexical site for the PRD-02 gate to sweep; §4.3's +"one module declares on the library's behalf" is the sanctioned answer, and this +is the seam it plugs into. + +A plain callable for the same reason the counter is one — the ledger needs +``tool -> declaration | None`` and nothing else — and injected rather than +imported so this module keeps no edge to the third-party adapter and a test can +resolve without depending on whichever ``deepagents`` version is installed. +""" + + class HeuristicToolSchemaTokenCounter: """The char/4 stand-in PRD-04 replaces with the real tokenizer chain. @@ -225,6 +242,7 @@ def measure( model_tools: Sequence[object], *, counter: ToolSchemaTokenCounter | None = None, + fallback_origin: ToolOriginResolver | None = None, ) -> tuple[ToolSchemaFootprint, ...]: """Return one footprint per tool, in composition order. @@ -233,13 +251,23 @@ def measure( gated Wave-1 block sitting last is information about *why* those tools are the ones to defer. - Never raises. A tool that cannot be serialized, a counter that throws, - and a declaration that fails validation all degrade to a recorded row — - occupancy measurement is best-effort and must never fail a run (§6.4). + ``fallback_origin`` is consulted only for a tool that carries no stamp, + so a declaration made at composition always wins over one inferred here + — the composing code knows what it built, and this resolver is reading + a dependency's internals. Omitting it restores the pre-PRD-06 behaviour + exactly: an unstamped tool measures as ``UNDECLARED``. + + Never raises. A tool that cannot be serialized, a counter that throws, a + resolver that throws, and a declaration that fails validation all + degrade to a recorded row — occupancy measurement is best-effort and + must never fail a run (§6.4). """ count = counter or HeuristicToolSchemaTokenCounter.count - return tuple(cls._footprint(tool, counter=count) for tool in model_tools) + return tuple( + cls._footprint(tool, counter=count, fallback_origin=fallback_origin) + for tool in model_tools + ) @classmethod def _footprint( @@ -247,10 +275,13 @@ def _footprint( tool: object, *, counter: ToolSchemaTokenCounter, + fallback_origin: ToolOriginResolver | None = None, ) -> ToolSchemaFootprint: """Measure one tool, degrading to a zero footprint on any failure.""" origin = cls._origin_of(tool) + if origin is None and fallback_origin is not None: + origin = cls._resolved_origin(tool, fallback_origin=fallback_origin) tool_name = cls._tool_name(tool) try: entry = cls.schema_entry(tool) @@ -294,6 +325,35 @@ def _origin_of(cls, tool: object) -> ContextOrigin | None: ) return None + @classmethod + def _resolved_origin( + cls, + tool: object, + *, + fallback_origin: ToolOriginResolver, + ) -> ContextOrigin | None: + """Ask the injected resolver to declare an unstamped tool. + + The resolver is the one collaborator here written outside this module, + so it is the one treated as untrusted: a raise, or anything back that is + not a :class:`ContextOrigin`, is the same answer as "nothing to say" and + leaves the tool ``UNDECLARED``. Trusting the return type would let a + malformed resolver put a non-contract object into ``label`` and + ``lifecycle`` and fail the row's validation instead — on the model-call + path, which §6.4 forbids. + """ + + try: + resolved = fallback_origin(tool) + except Exception: # noqa: BLE001 — a failed resolution is no resolution + _LOGGER.debug( + "Fallback context-origin resolution failed for a composed tool; " + "measuring it as UNDECLARED.", + exc_info=True, + ) + return None + return resolved if isinstance(resolved, ContextOrigin) else None + @classmethod def _count(cls, text: str, *, counter: ToolSchemaTokenCounter) -> int: """Count ``text`` through ``counter``, falling back to the heuristic. @@ -333,6 +393,7 @@ def _tool_name(cls, tool: object) -> str: __all__ = ( "HeuristicToolSchemaTokenCounter", + "ToolOriginResolver", "ToolSchemaFootprint", "ToolSchemaLedger", "ToolSchemaTokenCounter", diff --git a/services/ai-backend/src/runtime_worker/dependencies.py b/services/ai-backend/src/runtime_worker/dependencies.py index b23f3c545..207eb4cc4 100644 --- a/services/ai-backend/src/runtime_worker/dependencies.py +++ b/services/ai-backend/src/runtime_worker/dependencies.py @@ -144,10 +144,27 @@ def _web_search_tool(self) -> object: ``content``, and LangChain — which reads the field off the tool it dispatches — stringified the whole pair into ``ToolMessage.content`` and left ``artifact`` ``None`` on every web search. + + The tool declares its own Context Occupancy origin here, at the site + that authors its description and schema. ``ModelToolDeclaration.declared`` + yields to an existing declaration precisely so an authoring site can do + this, and it has to: the factory composes injected registry tools as one + anonymous group and marks the whole group ``deepagents.middleware`` / + ``third_party=True``. That bucket was honest while nothing could tell + the members apart, but it is wrong about this tool — the ~200 resident + tokens are ours, in this file, and a report that blamed a dependency for + them would send a reader to go edit a package that does not contain the + string. """ from langchain_core.tools import StructuredTool from agent_runtime.capabilities.retrying_tool import RetryingTool + from agent_runtime.observability.context_origin import ( + ContextLifecycle, + ContextOrigin, + ContextSegmentClass, + declare_context_origin, + ) inner = StructuredTool.from_function( func=self._search, @@ -156,11 +173,22 @@ def _web_search_tool(self) -> object: args_schema=self.WebSearchInput, response_format="content_and_artifact", ) - return RetryingTool.wrapping( - inner, - max_attempts=3, - initial_backoff_seconds=1.0, - max_backoff_seconds=8.0, + return declare_context_origin( + RetryingTool.wrapping( + inner, + max_attempts=3, + initial_backoff_seconds=1.0, + max_backoff_seconds=8.0, + ), + ContextOrigin( + # A literal rather than ``__name__`` so the PRD-02 inventory gate + # can read it statically; an owner it cannot resolve never + # reaches the reviewed pin. + owner="runtime_worker.dependencies", + name=self.Values.WEB_SEARCH_TOOL_NAME, + segment_class=ContextSegmentClass.TOOLS, + lifecycle=ContextLifecycle.RESIDENT, + ), ) def _search(self, query: str) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: diff --git a/services/ai-backend/tests/unit/agent_runtime/observability/test_context_occupancy_recorder.py b/services/ai-backend/tests/unit/agent_runtime/observability/test_context_occupancy_recorder.py index d529ef312..879199427 100644 --- a/services/ai-backend/tests/unit/agent_runtime/observability/test_context_occupancy_recorder.py +++ b/services/ai-backend/tests/unit/agent_runtime/observability/test_context_occupancy_recorder.py @@ -92,6 +92,14 @@ _FAKE_LIBRARY_CONSTANT: Final[str] = "SANDBOX_SYSTEM_PROMPT" _FAKE_LIBRARY_TEXT: Final[str] = "Library-owned sandbox guidance the model reads." +_BUILD_TIME_LABEL: Final[str] = RuntimeContextOrigins.BUILD_TIME_SYSTEM_PROMPT.label +"""What an unassembled system block is labelled with. + +Read off the origin rather than restated, so a rename of the declaration moves +these assertions with it instead of leaving them asserting a string that no +longer exists. +""" + class LengthCounter: """A ``TokenCounterPort`` whose count is a deterministic function of the text. @@ -501,11 +509,40 @@ def test_unmatched_library_suffix_is_undeclared_without_a_third_party_index( ) assert snapshot.undeclared_tokens > 0 - def test_without_a_plan_the_whole_system_block_is_undeclared(self) -> None: - # Honest degradation: with no typed plan there is genuinely nothing to - # attribute the system prompt to, and saying so is better than guessing. + def test_without_a_plan_the_system_block_is_the_build_time_prompt(self) -> None: + # With no typed plan there was no per-fragment assembly to fail: the + # block is the ``system_prompt`` the factory handed to the builder, and + # naming that is more accurate than calling it undeclared. This is the + # desktop's permanent state — every packaged run records + # ``assembly_record_id: None`` — so the old UNDECLARED answer reported a + # 4,798-token contract defect on every model call of a healthy system, + # and the composer's context meter showed users the number. snapshot = self.capture(self.request(system_text=self.POLICY_TEXT)) + system = self.of_class(snapshot, ContextSegmentClass.SYSTEM) + assert [segment.label for segment in system] == [ + "agent_runtime.execution.factory:system_prompt" + ] + assert [segment.detail for segment in system] == ["system[unassembled]"] + assert system[0].lifecycle is ContextLifecycle.RESIDENT + assert snapshot.undeclared_tokens == 0 + + def test_a_plan_that_matches_nothing_still_reports_undeclared(self) -> None: + # The distinction the coarse fallback must not erase. A plan existed and + # its fragments are absent from the prompt that was sent, so the two + # disagree — an attribution defect, and exactly what the field is for. + # Keying the fallback on "a plan matched nothing" instead of "there is + # no plan" would have silently swallowed this case. + plan = self.plan( + self.fragment( + fragment_id="00_base_runtime", + source_owner="agent_runtime.prompts", + content="Text that is nowhere in the system block.", + ) + ) + + snapshot = self.capture(self.request(system_text=self.POLICY_TEXT), plan=plan) + system = self.of_class(snapshot, ContextSegmentClass.SYSTEM) assert [segment.label for segment in system] == [UNDECLARED_CONTEXT_LABEL] assert snapshot.undeclared_tokens > 0 @@ -666,7 +703,12 @@ def test_text_around_a_library_constant_stays_separately_attributed( assert sum(segment.byte_count for segment in system) == len( system_text.encode("utf-8") ) - assert sum(1 for segment in system if segment.is_undeclared) == 2 + # Head and tail stay two spans, separate from the library constant. With + # no plan they carry the build-time prompt's label rather than + # UNDECLARED, but the split itself is what this asserts: the library's + # bytes keep the library's owner and are not swallowed into ours. + assert sum(1 for segment in system if segment.label == _BUILD_TIME_LABEL) == 2 + assert any(segment.third_party for segment in system) def test_a_broken_sweep_declares_nothing_rather_than_raising(self) -> None: index = self.index(ExplodingThirdPartyOrigins()) @@ -676,7 +718,9 @@ def test_a_broken_sweep_declares_nothing_rather_than_raising(self) -> None: recorder=self.recorder(third_party=index), ) - assert self.labels(snapshot) == (UNDECLARED_CONTEXT_LABEL,) + # No library label anywhere: the sweep contributed nothing and the block + # falls to the coarse first-party owner instead of raising. + assert self.labels(snapshot) == (_BUILD_TIME_LABEL,) def test_an_unresolvable_constant_is_dropped_from_the_index(self) -> None: # The constant is discovered but its module is not importable, which is @@ -699,7 +743,7 @@ def test_an_unresolvable_constant_is_dropped_from_the_index(self) -> None: recorder=self.recorder(third_party=index), ) - assert self.labels(snapshot) == (UNDECLARED_CONTEXT_LABEL,) + assert self.labels(snapshot) == (_BUILD_TIME_LABEL,) class TestToolBlockMeasurement(RecorderFixtureMixin): diff --git a/services/ai-backend/tests/unit/agent_runtime/observability/test_context_tool_declaration_coverage.py b/services/ai-backend/tests/unit/agent_runtime/observability/test_context_tool_declaration_coverage.py new file mode 100644 index 000000000..87168b189 --- /dev/null +++ b/services/ai-backend/tests/unit/agent_runtime/observability/test_context_tool_declaration_coverage.py @@ -0,0 +1,672 @@ +"""The gate that keeps ``undeclared_tokens`` at zero for the tool block (§4.4). + +``ContextOccupancySnapshot.undeclared_tokens`` is contractually **0** — anything +above it is "a first-party contract defect, not drift", and since PR #625/#626 +the composer's context meter renders it to users as an "N undeclared" notice. A +live packaged run recorded 8,257 undeclared tokens on every model call, 3,459 of +them from the tool block, and every existing test was green over it. + +They were green because the two mechanisms that were supposed to catch it both +looked somewhere else. The PRD-02 AST gate sweeps ``factory._model_visible_tools``, +which is the only place *this repository* composes a tool — but +``create_deep_agent`` installs the ``FilesystemMiddleware``, ``TodoListMiddleware`` +and ``SubAgentMiddleware`` tools inside the library, so ten tools reached the +provider having never passed a site the gate can see. And every unit test that +exercised the ledger built its own fixture tools, which are declared by +construction; a fixture cannot notice that the real surface carries objects +nobody stamped. + +So these tests are deliberately built on the **real** middleware tools and the +**real** composition chain rather than on stand-ins. That is the whole point: +this file has to fail when a dependency bump adds a tool we cannot attribute, +and a hand-built tool list can never do that. It is the reason +:class:`MiddlewareToolSurfaceMixin` builds the actual middleware tools and +:class:`FirstPartyToolSurfaceMixin` runs the actual factory-plus-display +decoration. + +The second half of the file is the counterweight, and it matters as much as the +first. A resolver that declared *everything* would also drive the number to zero, +while destroying the only signal the field carries. So the undeclared path is +asserted to survive intact for a first-party tool that genuinely lost its +declaration — the case that must stay loud. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import pytest +from pydantic import BaseModel, Field + +from agent_runtime.capabilities.middleware.display_metadata import ( + wrap_tools_with_display, +) +from agent_runtime.capabilities.retrying_tool import RetryingTool +from agent_runtime.execution.factory import _model_visible_tools +from agent_runtime.observability.context_occupancy import GraphScope +from agent_runtime.observability.context_occupancy_recorder import ( + ContextOccupancyRecorder, + ThirdPartyPromptIndex, +) +from agent_runtime.observability.context_origin import ( + ContextLifecycle, + ContextOrigin, + ContextSegmentClass, + context_origin_of, + declare_context_origin, +) +from agent_runtime.observability.context_third_party import ThirdPartyToolOrigins +from agent_runtime.observability.context_tool_ledger import ( + ToolSchemaFootprint, + ToolSchemaLedger, +) + + +class SearchArgs(BaseModel): + """A representative typed tool schema, shaped like a real model tool's.""" + + query: str = Field(description="What to look up.") + + +class MiddlewareToolSurfaceMixin: + """Build the tools ``create_deep_agent`` installs without a composition site. + + These are the objects that actually reach ``ModelRequest.tools`` and that no + first-party append site ever touches. Building them for real — rather than + restating the ten tool names — is what makes the assertion a gate instead of + a mirror: a dependency bump that renames ``read_file``, adds a built-in, or + moves a tool into a package outside ``DECLARABLE_ROOTS`` fails here, which is + precisely the event §4.3 wants surfaced at review time. + + ``task`` comes from ``build_atlas_task_tool`` rather than from + ``SubAgentMiddleware`` because that is what a run actually gets: this + repository monkey-patches its own ``task`` over the library's, keeping the + name and the slot. Constructing the library's version would have tested a + tool the product does not ship — and would have needed a fake model and a + full subagent spec to do it, coupling this gate to a constructor signature + that has nothing to do with attribution. + """ + + def middleware_installed_tools(self) -> tuple[object, ...]: + """Return the built-in tools that reach the model with no append site.""" + + from deepagents.middleware import FilesystemMiddleware + from langchain.agents.middleware import TodoListMiddleware + + from agent_runtime.delegation.subagents.atlas_task_tool import ( + build_atlas_task_tool, + ) + + return ( + *FilesystemMiddleware().tools, + *TodoListMiddleware().tools, + build_atlas_task_tool([]), + ) + + +class FirstPartyToolSurfaceMixin: + """Compose the model tool surface exactly as the graph build does. + + ``_model_visible_tools`` then ``wrap_tools_with_display``, in that order, + because the order is where the bug was: the factory declared ``web_search`` + correctly and the display decoration then rebuilt it into a fresh + ``StructuredTool``, dropping the stamp. Asserting on the factory's output + alone passes while the surface the model is shown is undeclared, so the + decoration has to be inside the test. + """ + + def first_party_model_tools(self) -> tuple[object, ...]: + """Return the decorated first-party tool surface for a minimal run.""" + + from runtime_worker.dependencies import WebSearchToolRegistry + + composed = _model_visible_tools( + tools=tuple(WebSearchToolRegistry().list_available_tools(None)), + mcp_registry=None, + skill_registry=None, + prior_tool_result_loader=None, + mcp_discovery_cache=None, + code_mode_tool=None, + sandbox_execute_tool=None, + stage_rowset_write_tool=None, + publish_artifact_tool=None, + revise_artifact_tool=None, + runtime_context=None, + mcp_per_tool=None, + mcp_catalog=None, + ) + return tuple(wrap_tools_with_display(composed)) + + +class MeasurementMixin: + """Measure a tool surface the way the recorder does, and report offenders.""" + + def footprints(self, tools: Sequence[object]) -> tuple[ToolSchemaFootprint, ...]: + """Measure ``tools`` through the ledger with the production resolver.""" + + return ToolSchemaLedger.measure( + tools, + fallback_origin=ThirdPartyToolOrigins().origin_for, + ) + + def undeclared( + self, + footprints: Sequence[ToolSchemaFootprint], + ) -> tuple[ToolSchemaFootprint, ...]: + """Return only the rows that carry no declaration.""" + + return tuple(footprint for footprint in footprints if not footprint.declared) + + def offender_report( + self, + footprints: Sequence[ToolSchemaFootprint], + ) -> str: + """Name the undeclared tools and what they cost, for the failure message. + + A bare count would tell a future reader that something regressed without + telling them which tool to go declare, which is the same unhelpfulness + the ledger exists to fix. + """ + + return ", ".join( + f"{footprint.tool_name} (~{footprint.estimated_tokens} tokens)" + for footprint in self.undeclared(footprints) + ) + + +class StubModelRequest: + """The four parts of a provider request, as the recorder reads them. + + ``MaterializedProviderRequest`` reaches for ``system_message`` / ``tools`` / + ``messages`` / ``response_format`` by attribute, so a plain object with those + four is a faithful stand-in for the library's ``ModelRequest`` without + importing a type whose constructor is not ours. + """ + + class SystemMessage: + """A system block, carried the way LangChain carries one.""" + + def __init__(self, content: str) -> None: + self.content = content + + def __init__( + self, + *, + tools: Sequence[object], + system_text: str = "", + ) -> None: + self.system_message = self.SystemMessage(system_text) if system_text else None + self.tools = tuple(tools) + self.messages: tuple[object, ...] = () + self.response_format = None + + +class StubCallIdentity: + """The one field the recorder reads off a run's call identity.""" + + def __init__(self, model_call_id: str = "model-call:fixture") -> None: + self.model_call_id = model_call_id + + +class RecorderMixin: + """Build a recorder with third-party *prompt* attribution switched off. + + Only the prompt index is disabled, and only so a tool-block assertion cannot + fail for reasons about somebody else's system text. Tool attribution is the + production resolver throughout. + """ + + def recorder(self) -> ContextOccupancyRecorder: + """Return a recorder whose tool-block resolution is the production one.""" + + return ContextOccupancyRecorder(third_party=ThirdPartyPromptIndex.disabled()) + + def snapshot_for(self, tools: Sequence[object], *, system_text: str = ""): + """Capture a snapshot over a request, with no assembly plan. + + ``plan`` is omitted on purpose: that is the desktop's state on every + model call, and it is the configuration the whole defect showed up in. + """ + + return self.recorder().capture( + StubModelRequest(tools=tools, system_text=system_text), + identity=StubCallIdentity(), + attempt_ordinal=1, + graph_scope=GraphScope.ROOT, + provider="anthropic", + model_family="claude-sonnet-5", + context_window_tokens=200_000, + ) + + +class TestBuiltInToolSurfaceIsDeclared( + MiddlewareToolSurfaceMixin, + FirstPartyToolSurfaceMixin, + MeasurementMixin, + RecorderMixin, +): + """Every tool the model is shown must be attributable to somebody.""" + + def test_middleware_installed_tools_contribute_no_undeclared_bytes(self) -> None: + """The ten middleware built-ins all resolve to an authoring module.""" + + footprints = self.footprints(self.middleware_installed_tools()) + + assert footprints, "the pinned library installed no built-in tools" + assert not self.undeclared(footprints), ( + "middleware-installed tools contributed undeclared context bytes: " + f"{self.offender_report(footprints)}" + ) + + def test_library_tools_are_attributed_to_the_package_that_wrote_them( + self, + ) -> None: + """A library tool is owned by its own module and marked third-party. + + The owner is what makes the row actionable, and for text we did not + write the actionable move is a profile exclusion or a dependency change + — which is exactly what ``third_party`` tells a reader. + """ + + by_name = { + footprint.tool_name: footprint + for footprint in self.footprints(self.middleware_installed_tools()) + } + + assert ( + by_name["read_file"].label == "deepagents.middleware.filesystem:read_file" + ) + assert by_name["read_file"].third_party is True + assert ( + by_name["write_todos"].label + == "langchain.agents.middleware.todo:write_todos" + ) + assert by_name["write_todos"].third_party is True + + def test_the_task_tool_is_reported_as_ours_not_the_librarys(self) -> None: + """``task`` occupies a library slot but is authored in this repository. + + ``build_atlas_task_tool`` replaces ``deepagents``' own ``task``, keeping + its name and its middleware slot. Everything visible from the outside + says third-party, and roughly 390 resident tokens of description would + have been billed to a dependency that does not contain the string. This + assertion is the one that stops the resolver guessing by appearance. + """ + + by_name = { + footprint.tool_name: footprint + for footprint in self.footprints(self.middleware_installed_tools()) + } + + assert by_name["task"].third_party is False + assert by_name["task"].label == ( + "agent_runtime.delegation.subagents.atlas_task_tool:task" + ) + + def test_first_party_tool_surface_contributes_no_undeclared_bytes(self) -> None: + """The composed-and-decorated first-party surface is fully declared.""" + + footprints = self.footprints(self.first_party_model_tools()) + + assert footprints, "the factory composed no model tools" + assert not self.undeclared(footprints), ( + "first-party model tools contributed undeclared context bytes: " + f"{self.offender_report(footprints)}" + ) + + def test_web_search_is_owned_by_the_module_that_writes_its_description( + self, + ) -> None: + """``web_search`` is ours, and the report must not blame the library. + + It reaches the factory inside a ``RetryingTool``, which used to make it + both undeclared (the display decoration rebuilt it) and, once declared + by the injected-registry group, attributed to ``deepagents.middleware`` + — a package that does not contain the string being measured. + """ + + by_name = { + footprint.tool_name: footprint + for footprint in self.footprints(self.first_party_model_tools()) + } + + assert by_name["web_search"].label == "runtime_worker.dependencies:web_search" + assert by_name["web_search"].third_party is False + + def test_snapshot_over_the_whole_surface_reports_zero_undeclared_tokens( + self, + ) -> None: + """The wiring, not just the resolver: a captured snapshot reconciles to 0. + + The resolver existing is not the same as the recorder using it — the + third-party *prompt* adapter had shipped for exactly this long while the + tool block never consulted anything — so the assertion that matters is + made on the snapshot the run actually persists. + """ + + tools = (*self.middleware_installed_tools(), *self.first_party_model_tools()) + + snapshot = self.snapshot_for(tools) + + assert snapshot is not None + assert snapshot.estimated_input_tokens > 0, "the surface measured as free" + assert snapshot.undeclared_tokens == 0, ( + "the model tool surface reported undeclared occupancy: " + f"{self.offender_report(self.footprints(tools))}" + ) + + def test_a_whole_desktop_shaped_request_reports_zero_undeclared(self) -> None: + """The end-to-end shape the live run recorded: system block plus tools. + + The packaged run this work started from reported 8,257 undeclared tokens + on every model call — 4,798 from a system block with no assembly plan + and 3,459 from tools nobody had stamped. Both halves are measured here in + one snapshot, with no plan, because that is the configuration the desktop + runs in and asserting the halves separately would let their sum regress + unnoticed. + """ + + tools = (*self.middleware_installed_tools(), *self.first_party_model_tools()) + + snapshot = self.snapshot_for( + tools, + system_text="You are a careful assistant. Follow the workspace rules.", + ) + + assert snapshot is not None + assert snapshot.undeclared_tokens == 0 + assert {segment.label for segment in snapshot.segments} == { + segment.label + for segment in snapshot.segments + if segment.label != "UNDECLARED" + } + + +class TestUndeclaredOccupancyStaysVisible(MeasurementMixin, RecorderMixin): + """The resolver must not become a blanket amnesty. + + Driving ``undeclared_tokens`` to zero is only worth anything if the field can + still reach non-zero. Every test here asserts a case that must *stay* + undeclared, because each one is a real defect a reader needs to see. + """ + + def _first_party_tool(self, name: str = "orphan_tool") -> object: + """Build an undeclared tool whose dispatch callable is defined here.""" + + from langchain_core.tools import StructuredTool + + def _run(query: str) -> str: + return query + + return StructuredTool.from_function( + func=_run, + name=name, + description="A first-party tool nobody declared.", + args_schema=SearchArgs, + ) + + def test_undeclared_first_party_tool_still_measures_undeclared(self) -> None: + """A tool authored in this repository is never declared on its behalf. + + This is the case the whole design turns on. The resolver reads a tool's + authoring module, and text written in this repository has a composition + site that could have declared it — so declaring it here would convert an + actionable defect into a silent one. + """ + + (footprint,) = self.footprints([self._first_party_tool()]) + + assert footprint.declared is False + assert footprint.label == "UNDECLARED" + assert footprint.estimated_tokens > 0 + + def test_snapshot_surfaces_an_undeclared_first_party_tool(self) -> None: + """The undeclared bytes reach ``undeclared_tokens``, not just the row.""" + + snapshot = self.snapshot_for([self._first_party_tool()]) + + assert snapshot is not None + assert snapshot.undeclared_tokens > 0 + + def test_an_existing_declaration_wins_over_the_resolver(self) -> None: + """A stamp made at composition is never overwritten by inference. + + The composing code knows what it built; the resolver is reading a + dependency's internals. If the two ever disagree, the stamp is the one + with standing. + """ + + tool = self._first_party_tool(name="declared_tool") + declare_context_origin( + tool, + ContextOrigin( + owner="agent_runtime.capabilities.tools", + name="declared_tool", + segment_class=ContextSegmentClass.TOOLS, + lifecycle=ContextLifecycle.RESIDENT, + ), + ) + + (footprint,) = ToolSchemaLedger.measure( + [tool], + fallback_origin=lambda _tool: ContextOrigin( + owner="deepagents.middleware", + name="usurper", + segment_class=ContextSegmentClass.TOOLS, + lifecycle=ContextLifecycle.RESIDENT, + third_party=True, + ), + ) + + assert footprint.label == "agent_runtime.capabilities.tools:declared_tool" + assert footprint.third_party is False + + def test_a_resolver_that_raises_degrades_to_undeclared(self) -> None: + """Measurement is on the model-call path and may never raise (§6.4).""" + + def _explode(_tool: object) -> ContextOrigin | None: + raise RuntimeError("resolver failure") + + (footprint,) = ToolSchemaLedger.measure( + [self._first_party_tool()], + fallback_origin=_explode, + ) + + assert footprint.declared is False + assert footprint.label == "UNDECLARED" + + def test_a_resolver_returning_a_non_origin_is_ignored(self) -> None: + """An untyped return is treated as "nothing to say", never trusted. + + Trusting it would put a non-contract object into ``label`` and + ``lifecycle`` and fail the row's own validation instead — on the + model-call path, which is the failure this whole lane is built to avoid. + """ + + (footprint,) = ToolSchemaLedger.measure( + [self._first_party_tool()], + fallback_origin=lambda _tool: "deepagents.middleware:not_an_origin", # type: ignore[return-value] + ) + + assert footprint.declared is False + assert footprint.label == "UNDECLARED" + + def test_omitting_the_resolver_preserves_the_original_behaviour(self) -> None: + """``measure`` without a resolver is exactly what it was before.""" + + (footprint,) = ToolSchemaLedger.measure([self._first_party_tool()]) + + assert footprint.declared is False + assert footprint.label == "UNDECLARED" + + +class TestDisplayWrappingPreservesDeclarations: + """A re-wrap must carry the declaration the composition site made. + + ``wrap_tools_with_display`` has two branches and they used to behave + differently: the ``StructuredTool`` branch copies the instance (stamp + included) while the delegation branch rebuilds the tool with + ``from_function`` and holds the inner in a closure — no stamp, and no + attribute chain for ``ContextOriginBinding.of`` to walk. Every non- + ``StructuredTool`` on the surface was silently un-declared on the way to the + model, which is how a correctly-declared ``web_search`` was reported as a + first-party contract defect. + """ + + def _declared_structured_tool(self, name: str) -> object: + from langchain_core.tools import StructuredTool + + def _run(query: str) -> str: + return query + + tool = StructuredTool.from_function( + func=_run, + name=name, + description="Declared at its composition site.", + args_schema=SearchArgs, + ) + return declare_context_origin( + tool, + ContextOrigin( + owner="agent_runtime.capabilities.tools", + name=name, + segment_class=ContextSegmentClass.TOOLS, + lifecycle=ContextLifecycle.RESIDENT, + ), + ) + + def test_structured_tool_rewrap_keeps_its_declaration(self) -> None: + """The copying branch keeps the stamp, and stays that way.""" + + source = self._declared_structured_tool("structured_tool") + + (wrapped,) = wrap_tools_with_display([source]) + + origin = context_origin_of(wrapped) + assert origin is not None + assert origin.label == "agent_runtime.capabilities.tools:structured_tool" + + def test_delegating_tool_rewrap_keeps_its_declaration(self) -> None: + """The rebuilding branch — the one that dropped it — now carries it over. + + ``RetryingTool`` is the real shape this happens to: it is a ``BaseTool`` + and not a ``StructuredTool``, which is exactly the branch predicate. + """ + + inner = self._declared_structured_tool("retried_tool") + source = RetryingTool.wrapping( + inner, + max_attempts=2, + initial_backoff_seconds=0.1, + max_backoff_seconds=1.0, + ) + assert context_origin_of(source) is not None, "fixture is not declared" + + (wrapped,) = wrap_tools_with_display([source]) + + origin = context_origin_of(wrapped) + assert origin is not None, ( + "the display decoration dropped the declaration made at composition" + ) + assert origin.label == "agent_runtime.capabilities.tools:retried_tool" + + def test_rewrapping_an_undeclared_tool_declares_nothing(self) -> None: + """Carrying a declaration over never invents one.""" + + from langchain_core.tools import StructuredTool + + def _run(query: str) -> str: + return query + + source = RetryingTool.wrapping( + StructuredTool.from_function( + func=_run, + name="undeclared_tool", + description="Nobody declared this.", + args_schema=SearchArgs, + ), + max_attempts=2, + initial_backoff_seconds=0.1, + max_backoff_seconds=1.0, + ) + + (wrapped,) = wrap_tools_with_display([source]) + + assert context_origin_of(wrapped) is None + + +class TestThirdPartyToolOriginResolution: + """The resolver's own rules, asserted directly rather than through a surface.""" + + @pytest.fixture + def origins(self) -> ThirdPartyToolOrigins: + return ThirdPartyToolOrigins() + + def test_a_tool_with_no_readable_callable_is_not_declared( + self, + origins: ThirdPartyToolOrigins, + ) -> None: + """No authoring module means no owner, and an absent owner is honest.""" + + class Opaque: + name = "opaque" + + assert origins.origin_for(Opaque()) is None + + def test_the_tool_class_module_is_never_used_as_the_owner( + self, + origins: ThirdPartyToolOrigins, + ) -> None: + """``StructuredTool`` lives in ``langchain_core`` whoever wrote the tool. + + ``langchain_core`` is a declarable root, so a fallback to + ``type(tool).__module__`` would not merely degrade — it would confidently + stamp every tool on the surface as third-party ``langchain_core`` text + and send a reader to edit a package that contains none of it. + """ + + from langchain_core.tools import StructuredTool + + def _run(query: str) -> str: + return query + + tool = StructuredTool.from_function( + func=_run, + name="first_party", + description="Authored in this test module.", + args_schema=SearchArgs, + ) + # ``func`` is defined here, outside every declarable root. + assert origins.origin_for(tool) is None + + def test_a_module_outside_the_declarable_roots_is_not_declared( + self, + origins: ThirdPartyToolOrigins, + ) -> None: + """An unrecognised author lands in ``undeclared_tokens``, not in a vendor. + + The allowlist direction: "not one of ours" is never taken as evidence of + "therefore theirs". + """ + + class Vendored: + name = "vendored" + + def func(self) -> None: ... + + Vendored.func.__module__ = "some_unpinned_package.tools" + assert origins.origin_for(Vendored()) is None + + def test_an_unnamed_third_party_tool_is_not_declared( + self, + origins: ThirdPartyToolOrigins, + ) -> None: + """A label needs both halves; half a label traces back to nothing.""" + + class Unnamed: + name = " " + + def func(self) -> None: ... + + Unnamed.func.__module__ = "deepagents.middleware.filesystem" + assert origins.origin_for(Unnamed()) is None diff --git a/services/ai-backend/tests/unit/test_context_origin_gate.py b/services/ai-backend/tests/unit/test_context_origin_gate.py index 1e630ef71..f39ca4cde 100644 --- a/services/ai-backend/tests/unit/test_context_origin_gate.py +++ b/services/ai-backend/tests/unit/test_context_origin_gate.py @@ -103,13 +103,36 @@ def test_declared_context_origin_inventory_is_reviewed() -> None: "agent_runtime.conversation:assistant_tool_calls", "agent_runtime.conversation:tool_result", "agent_runtime.conversation:user", + # Declared at its authoring site because it has no other one. ``task`` is + # handed to the subagent middleware and installed by ``create_deep_agent``, + # so it never passes ``_model_visible_tools`` and the site sweep below + # cannot see it. It also *looks* like the library's own ``task`` — same + # name, same slot — which is why leaving it undeclared billed ~390 + # resident tokens of our text to ``deepagents``. + "agent_runtime.delegation.subagents.atlas_task_tool:task", "agent_runtime.execution.factory:10_application_context_boundary", + # The whole system block when no F2 plan was assembled — the desktop's + # permanent state. Coarser than per-fragment attribution and honest about + # it (the segment's detail is ``system[unassembled]``), but far better + # than the alternative it replaced: reporting the entire system prompt as + # a contract defect on every model call. + "agent_runtime.execution.factory:system_prompt", "agent_runtime.execution.model_invocation:response_format", "agent_runtime.execution:state_file", "agent_runtime.prompts.runtime:00_base_runtime", "agent_runtime.prompts:assembly_joiner", "deepagents.middleware:tools", ) + # Not in the pin, and knowingly so: ``runtime_worker.dependencies:web_search`` + # spells its ``name`` as ``Values.WEB_SEARCH_TOOL_NAME`` rather than as a + # literal, and the sweep drops any declaration whose owner or name it cannot + # read statically. Making it literal would duplicate the constant that names + # the tool two lines above it. It is gated behaviourally instead, against the + # real composed surface, in + # ``tests/unit/agent_runtime/observability/test_context_tool_declaration_coverage.py`` + # (``test_web_search_is_owned_by_the_module_that_writes_its_description``), + # which is the stronger of the two checks: it asserts the label the ledger + # actually produces rather than the one the source appears to request. def test_inventory_covers_every_model_tool_composed_by_the_factory() -> None: