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: