fix(ai-backend): declare the context nobody was declaring

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."""
 

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,
+        ),
     )
 
 

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,15 +495,42 @@ 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:
             return ()
         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)