code-puppy 0.0.910__tar.gz → 0.0.912__tar.gz

{code_puppy-0.0.910 → code_puppy-0.0.912}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: code-puppy
-Version: 0.0.910
+Version: 0.0.912
 Summary: Code generation agent
 Project-URL: repository, https://github.com/mpfaffenberger/code_puppy
 Project-URL: HomePage, https://github.com/mpfaffenberger/code_puppy

{code_puppy-0.0.910 → code_puppy-0.0.912}/code_puppy/agents/_builder.py

@@ -726,7 +726,7 @@ def build_pydantic_agent(
                 history_compaction,
                 ProcessHistory(steer_processor),
                 build_response_clamp(),
-                build_model_message_transform(logical_agent_name),
+                build_model_message_transform(logical_agent_name, agent),
                 # Sub-agent recursion guards on the wrap_tool_execute seam
                 # (denies invoke_agent calls past the depth caps before the
                 # tool body runs). Sole wrap_tool_execute implementer, so

{code_puppy-0.0.910 → code_puppy-0.0.912}/code_puppy/agents/_compaction.py

@@ -64,6 +64,7 @@ from code_puppy.config import (
     get_protected_token_count,
     get_summarization_model_name,
 )
+from code_puppy.context_accounting import context_tokens, model_names
 from code_puppy.messaging import emit_warning
 
 # Kept importable at module level for the context-indicator plugin, which
@@ -202,7 +203,12 @@ class CodePuppyCompactionStore:
         return self.agent._get_model_context_length()
 
     def context_overhead(self) -> int:
-        return self.agent._estimate_context_overhead()
+        overhead = self.agent._estimate_context_overhead()
+        # Cached so the outbound-request stamp (``_model_message_transform``)
+        # can report the exact overhead this measurement used, instead of
+        # re-deriving a possibly-different value moments later.
+        self.agent._context_overhead = overhead
+        return overhead
 
     def model_name(self) -> Optional[str]:
         if self.agent is None:
@@ -225,6 +231,17 @@ def _take_forced_compaction_request() -> bool:
     return get_pause_controller().take_compaction_request()
 
 
+def _total_tokens(messages, model, overhead) -> int:
+    """Anchor the whole-history total to the live request model's identity.
+
+    Wired as the capability's ``total_estimator`` -- see
+    ``code_puppy.context_accounting.context_tokens`` for the anchor/fallback
+    logic. Kept as a module-level function (not a lambda) so it shows up by
+    name in tracebacks and can be monkeypatched directly in tests.
+    """
+    return context_tokens(messages, model_names(model), overhead)
+
+
 def build_history_compaction(agent: Any) -> PureHistoryCompaction:
     """Build the pure ``HistoryCompaction`` capability wired to Code Puppy.
 
@@ -243,6 +260,7 @@ def build_history_compaction(agent: Any) -> PureHistoryCompaction:
         strategy_name=lambda: get_compaction_strategy(),
         compaction_threshold=lambda: get_compaction_threshold(),
         token_estimator=estimate_tokens_for_message,
+        total_estimator=_total_tokens,
         message_hasher=hash_message,
         history_sanitizer=sanitize_tool_call_ids,
         force_poll=_take_forced_compaction_request,

code_puppy-0.0.912/code_puppy/agents/_model_message_transform.py

@@ -0,0 +1,130 @@
+"""Plugin transforms for final outbound model messages."""
+
+from copy import copy
+from typing import Any
+
+from pydantic_ai import RunContext
+from pydantic_ai.capabilities import Hooks, WrapModelRequestHandler
+from pydantic_ai.messages import ModelResponse
+from pydantic_ai.models import ModelRequestContext
+from pydantic_ai.models.wrapper import WrapperModel
+
+from code_puppy.agents._foreign_thinking import strip_foreign_thinking
+from code_puppy.callbacks import on_transform_model_messages
+
+# Per-request prefix fingerprint, handed from the ``wrap_model_request`` leg
+# to the ``after_model_request`` leg via this attribute on the (mutable,
+# per-request) ``ModelRequestContext``. ``None`` means "this request's
+# outbound messages diverged from the durable history it was built from, so
+# don't stamp an anchor" -- see ``build_model_message_transform``.
+_PREFIX_ATTR = "_context_accounting_prefix"
+
+
+class _ContinuationObserver(WrapperModel):
+    """Observe a suspended->continued chain *before* pydantic-ai merges it.
+
+    ``continuation_delay`` runs on intermediate suspended responses before
+    their usage can merge. This request-local wrapper preserves the real
+    model's delay behavior while marking the final response ineligible for
+    a receipt. The framework's durable-execution capabilities use the same
+    request-model swap seam. Provider details and billing limitations live
+    in ``docs/API_CONTEXT_ACCOUNTING.md``.
+    """
+
+    def __init__(self, wrapped):
+        super().__init__(wrapped)
+        self.continued = False
+
+    def continuation_delay(self, response: ModelResponse) -> float | None:
+        self.continued = True
+        return super().continuation_delay(response)
+
+
+def build_model_message_transform(agent_name: str | None, owner: Any = None) -> Hooks:
+    """Build the request-only plugin transform for an agent.
+
+    ``owner`` is the Code Puppy agent (or agent-shaped object) whose
+    ``_context_overhead`` compaction just computed for this same request --
+    see ``CodePuppyCompactionStore.context_overhead``. Passing it lets the
+    stamped receipt record the *exact* overhead the compaction trigger used,
+    so the next request's ``context_tokens`` call reconstructs the same
+    total instead of silently drifting by whatever overhead changed between
+    the two reads. ``None`` (sub-agents/tests that don't wire one) just
+    stamps ``overhead=0``, which only costs an extra estimate on that one
+    anchor -- never a wrong total, because ``context_tokens`` always folds
+    the caller's *current* overhead back in.
+    """
+
+    async def transform(
+        _ctx: RunContext[Any],
+        *,
+        request_context: ModelRequestContext,
+        handler: WrapModelRequestHandler,
+    ) -> ModelResponse:
+        from code_puppy.context_accounting import fingerprint, invalidate_anchors
+
+        durable_messages = request_context.messages
+        # Snapshot *before* anything can touch the messages. strip_foreign_thinking
+        # is a no-op pass-through for non-Anthropic models (same list AND same
+        # message objects by reference, not a copy) -- so a plugin transform
+        # that mutates in place would otherwise mutate ``durable_messages`` too,
+        # and a fingerprint taken afterwards would never see the divergence.
+        original_fingerprint = fingerprint(durable_messages)
+
+        transformed_context = copy(request_context)
+        transformed_context.messages = strip_foreign_thinking(
+            list(durable_messages), request_context.model
+        )
+        await on_transform_model_messages(agent_name, transformed_context.messages)
+
+        # The anchor must describe what the *durable history* looked like,
+        # because that's what next turn's ``context_tokens`` re-fingerprints
+        # against. If anything between the snapshot above and the handler
+        # call -- the built-in foreign-thinking strip, or a plugin's rewrite,
+        # in place or not -- changed what's actually sent, the response's
+        # usage no longer describes ``durable_messages`` and must not anchor it.
+        prefix = original_fingerprint
+        if fingerprint(transformed_context.messages) != prefix:
+            prefix = None
+
+        # Wrap *after* the fingerprint check above (which needs the real
+        # model for foreign-thinking detection) and *before* the handler
+        # call, so the continuation loop inside `handler` actually drives
+        # this wrapper's `continuation_delay` -- see `_ContinuationObserver`.
+        observer = _ContinuationObserver(transformed_context.model)
+        transformed_context.model = observer
+
+        try:
+            response = await handler(transformed_context)
+        except BaseException:
+            invalidate_anchors(durable_messages)
+            raise
+
+        if observer.continued:
+            # Usage on `response` is a cumulative sum across segments, not
+            # one measured prompt -- decline to anchor it. Billing/usage
+            # itself is untouched; only our own receipt is suppressed.
+            prefix = None
+        setattr(request_context, _PREFIX_ATTR, prefix)
+        return response
+
+    async def stamp(
+        _ctx: RunContext[Any],
+        *,
+        request_context: ModelRequestContext,
+        response: ModelResponse,
+    ) -> ModelResponse:
+        from code_puppy.context_accounting import model_names, record_anchor
+
+        prefix = getattr(request_context, _PREFIX_ATTR, None)
+        if prefix is not None:
+            record_anchor(
+                request_context.messages,
+                response,
+                context_overhead=getattr(owner, "_context_overhead", 0),
+                prefix=prefix,
+                request_model=model_names(request_context.model),
+            )
+        return response
+
+    return Hooks(model_request=transform, after_model_request=stamp)

{code_puppy-0.0.910 → code_puppy-0.0.912}/code_puppy/agents/_runtime.py

@@ -2,9 +2,9 @@
 
 Replaces the monolithic ``BaseAgent.run_with_mcp`` coroutine. Everything here
 is a free function; the agent is passed in explicitly. Integration points
-preserved verbatim:
+include:
 
-- Plugin-supplied async context managers wrap the run (see
+- Plugin-supplied async context managers wrap each physical model attempt (see
   ``on_agent_run_context``); used e.g. by the DBOS plugin to set a workflow
   ID and swap MCP toolsets in/out.
 - SIGINT fallback-handler choice driven by ``sigint_fallback_cancels()``
@@ -25,7 +25,6 @@ import ssl
 import sys
 import threading
 import uuid
-from contextlib import AsyncExitStack
 from typing import Any, Callable, Iterator, List, Optional, Sequence, Type, Union
 
 import httpcore
@@ -82,8 +81,8 @@ from code_puppy.agents._run_signals import (
     sigint_should_cancel,
 )
 from code_puppy.agents.event_stream_handler import event_stream_handler
+from code_puppy.agents.run_invocation import ModelCall, run_with_exception_retry
 from code_puppy.callbacks import (
-    on_agent_exception,
     on_agent_run_cancel,
     on_agent_run_context,
     on_agent_run_end,
@@ -784,10 +783,23 @@ async def _run_with_mcp_impl(
 
     if agent._code_generation_agent is None:
         build_pydantic_agent(agent)
-    pydantic_agent = agent._code_generation_agent
 
     if output_type is not None:
-        pydantic_agent = build_pydantic_agent(agent, output_type=output_type)
+        build_pydantic_agent(agent, output_type=output_type)
+
+    # MCP lifecycle tasks can invalidate the cache during this run. Keep its
+    # built client as a fallback, while still resolving live replacements lazily.
+    built_client = agent._code_generation_agent
+    invocation = ModelCall(
+        lambda: (
+            current
+            if (current := agent._code_generation_agent) is not None
+            else built_client
+        ),
+        lambda current: on_agent_run_context(
+            agent, current, group_id, getattr(agent, "_mcp_servers", None) or []
+        ),
+    )
 
     prompt = _should_prepend_system_prompt(agent, prompt)
     prompt_payload = _build_prompt_payload(prompt, attachments, link_attachments)
@@ -830,42 +842,19 @@ async def _run_with_mcp_impl(
             usage_limits=usage_limits, event_stream_handler=stream_handler, **kwargs
         )
         _call = _main_retry(
-            resumable_call(agent, pydantic_agent, prompt_to_use, **run_options)
+            resumable_call(agent, invocation, prompt_to_use, **run_options)
         )
 
-        async def _call_with_exception_recovery() -> Any:
-            """Run ``_call`` and let plugins request one exception retry."""
-            try:
-                return await _call()
-            except Exception as exc:
-                hook_results = await on_agent_exception(
-                    exc,
-                    agent=agent,
-                    agent_name=agent.name,
-                    model_name=agent.get_model_name(),
-                )
-                retry_req = next(
-                    (r for r in hook_results if isinstance(r, dict) and r.get("retry")),
-                    None,
-                )
-                if not retry_req:
-                    raise
-
-                retry_delay = retry_req.get("delay", 0.0)
-                if retry_delay:
-                    await asyncio.sleep(retry_delay)
-                return await _call()
-
-        result = await _call_with_exception_recovery()
+        result = await run_with_exception_retry(_call, agent=agent)
 
         # ``now``-mode steers are injected by make_steer_history_processor
         # (before every model call); ``queue``-mode ones drain between runs
         # below — additive, won't interrupt in-progress work.
         async def _follow_up_run(follow_up_prompt: Any) -> Any:
             call = _main_retry(
-                resumable_call(agent, pydantic_agent, follow_up_prompt, **run_options)
+                resumable_call(agent, invocation, follow_up_prompt, **run_options)
             )
-            return await call()
+            return await run_with_exception_retry(call, agent=agent)
 
         hook_retries_used = 0
         queued_steers_used = 0
@@ -935,14 +924,7 @@ async def _run_with_mcp_impl(
                 agent._message_history
             )
 
-            mcp_servers = getattr(agent, "_mcp_servers", None) or []
-            run_ctxs = on_agent_run_context(
-                agent, pydantic_agent, group_id, mcp_servers
-            )
-            async with AsyncExitStack() as stack:
-                for cm in run_ctxs:
-                    await stack.enter_async_context(cm)
-                return await _do_run(prompt_payload)
+            return await _do_run(prompt_payload)
         except* UsageLimitExceeded as ule:
             emit_info(f"Usage limit exceeded: {ule}", group_id=group_id)
             emit_info(

{code_puppy-0.0.910 → code_puppy-0.0.912}/code_puppy/agents/retry_checkpoint.py

@@ -32,7 +32,7 @@ class RetryCheckpoint:
         return self.completed
 
 
-def resumable_call(agent, pydantic_agent, prompt, **kwargs):
+def resumable_call(agent, client, prompt, **kwargs):
     """One logical prompt: resend only if it never reached persistent history."""
     checkpointed = False
 
@@ -46,7 +46,7 @@ def resumable_call(agent, pydantic_agent, prompt, **kwargs):
             if isinstance(part, UserPromptPart)
         }
         try:
-            return await pydantic_agent.run(
+            return await client.run(
                 None if checkpointed else prompt,
                 message_history=agent._message_history,
                 **kwargs,

code_puppy-0.0.912/code_puppy/agents/run_invocation.py

@@ -0,0 +1,72 @@
+"""Separate logical-prompt retries from physical model-call resources."""
+
+import asyncio
+from collections.abc import Awaitable, Callable, Sequence
+from contextlib import AsyncExitStack
+from typing import Any, AsyncContextManager
+
+from pydantic_ai.exceptions import RunCancelled
+
+from code_puppy.callbacks import on_agent_exception
+
+
+class ModelCall:
+    """Run one physical attempt against a freshly resolved client.
+
+    ``resolve`` returns the current client, which may have been replaced by
+    recovery. ``contexts(client)`` creates fresh async context managers for
+    that same client. Checkpoint state belongs outside this adapter: pass it
+    to ``resumable_call`` once for each logical prompt, before applying retry.
+    Neither the client nor its context managers are cached between attempts.
+    """
+
+    def __init__(
+        self,
+        resolve: Callable[[], Any],
+        contexts: Callable[[Any], Sequence[AsyncContextManager[Any]]],
+    ) -> None:
+        self.resolve = resolve
+        self.contexts = contexts
+
+    async def run(self, prompt: Any, **kwargs: Any) -> Any:
+        current = self.resolve()
+        async with AsyncExitStack() as stack:
+            for context in self.contexts(current):
+                await stack.enter_async_context(context)
+            return await current.run(prompt, **kwargs)
+
+
+async def run_with_exception_retry(
+    call: Callable[[], Awaitable[Any]], *, agent: Any
+) -> Any:
+    """Let exception callbacks request one retry of the same logical call.
+
+    The callback may replace the model client; ``ModelCall`` resolves that
+    replacement on the next attempt. Retry the existing callable, rather than
+    constructing another checkpoint. Exceptions from the retry propagate and
+    cancellation never enters exception recovery.
+    """
+    try:
+        return await call()
+    except RunCancelled:
+        raise
+    except Exception as exc:
+        if isinstance(exc, ExceptionGroup) and exc.subgroup(RunCancelled) is not None:
+            raise  # Preserve the whole original group, including sibling errors.
+        hook_results = await on_agent_exception(
+            exc,
+            agent=agent,
+            agent_name=agent.name,
+            model_name=agent.get_model_name(),
+        )
+        retry_req = next(
+            (r for r in hook_results if isinstance(r, dict) and r.get("retry")),
+            None,
+        )
+        if not retry_req:
+            raise
+
+        retry_delay = retry_req.get("delay", 0.0)
+        if retry_delay:
+            await asyncio.sleep(retry_delay)
+        return await call()

{code_puppy-0.0.910 → code_puppy-0.0.912}/code_puppy/callbacks.py

@@ -1806,9 +1806,13 @@ def on_wrap_pydantic_agent(
 def on_agent_run_context(agent, pydantic_agent, group_id, mcp_servers) -> List[Any]:
     """Collect async context managers that should wrap the ``pydantic_agent.run()`` call.
 
-    Each callback returns an async CM (with ``__aenter__``/``__aexit__``) or
-    ``None``. The caller composes all non-``None`` results via
-    ``contextlib.AsyncExitStack``.
+    Each callback returns a fresh async CM (with ``__aenter__``/``__aexit__``)
+    or ``None``. Main-agent contexts are created and exited per physical model
+    attempt, including retries and follow-ups, around the currently resolved
+    client. They are not a scope for the whole prompt or retry loop. The caller
+    composes all non-``None`` results via ``contextlib.AsyncExitStack``.
+    In ``tools/subagent_invocation.py``, this same hook retains whole-run
+    scope around the subagent's run and retry loop, not per-attempt scope.
 
     Returns a list of async context managers (may be empty).
     """

{code_puppy-0.0.910 → code_puppy-0.0.912}/code_puppy/capabilities/compaction.py

@@ -24,7 +24,7 @@ from typing import Any, Callable, List, Optional, Protocol, Set, Tuple
 from pydantic_ai import CapabilityEvent
 from pydantic_ai.capabilities import AbstractCapability
 from pydantic_ai.messages import ModelMessage, ModelResponse, ThinkingPart
-from pydantic_ai.models import ModelRequestContext
+from pydantic_ai.models import Model, ModelRequestContext
 from pydantic_ai.tools import RunContext
 from pydantic_ai.usage import RunUsage
 
@@ -255,6 +255,17 @@ class HistoryCompaction(AbstractCapability[Any]):
     message_hasher: Callable[[ModelMessage], str]
     """Stable content hash used for merge dedup + dropped tracking."""
 
+    total_estimator: Optional[Callable[[List[ModelMessage], Model, int], int]] = None
+    """Optional whole-history total, given the resolved request model.
+
+    When set, this replaces ``sum(token_estimator(...)) + context_overhead``
+    for both the pre- and post-compaction measurement. The seam exists so an
+    application can anchor the total to a provider-reported usage figure
+    instead of a per-message heuristic sum, without this module knowing
+    anything about how that anchor is produced or stored. ``None`` keeps the
+    historical sum-based behavior.
+    """
+
     history_sanitizer: Optional[Callable[[List[ModelMessage]], List[ModelMessage]]] = (
         None
     )
@@ -350,6 +361,24 @@ class HistoryCompaction(AbstractCapability[Any]):
         )
         return cleaned
 
+    def _total(
+        self,
+        ctx: RunContext[Any],
+        messages: List[ModelMessage],
+        model_name: Optional[str],
+        context_overhead: int,
+    ) -> int:
+        """Resolve the whole-history total via ``total_estimator`` if wired.
+
+        Falls back to the historical ``sum(token_estimator) + overhead`` when
+        no ``total_estimator`` was injected, so callers that never wire one
+        see unchanged behavior.
+        """
+        if self.total_estimator is not None:
+            return self.total_estimator(messages, ctx.model, context_overhead)
+        message_tokens = sum(self.token_estimator(m, model_name) for m in messages)
+        return message_tokens + context_overhead
+
     async def measure_and_compact(
         self,
         ctx: RunContext[Any],
@@ -366,8 +395,7 @@ class HistoryCompaction(AbstractCapability[Any]):
         always survive a failed compaction.
         """
         model_name = self.store.model_name()
-        message_tokens = sum(self.token_estimator(m, model_name) for m in messages)
-        total_tokens = message_tokens + context_overhead
+        total_tokens = self._total(ctx, messages, model_name, context_overhead)
         proportion_used = total_tokens / model_max if model_max else 0.0
 
         await _safe_emit(
@@ -425,8 +453,12 @@ class HistoryCompaction(AbstractCapability[Any]):
         dropped = [m for m in messages if self.message_hasher(m) not in result_hashes]
 
         # Parity with the historical implementation: the post-compaction
-        # measurement reports message tokens only (no overhead term).
-        final_token_count = sum(self.token_estimator(m, model_name) for m in result)
+        # measurement reports message tokens only (no overhead term). A
+        # ``total_estimator`` naturally falls back to the same "message
+        # tokens only" shape here too -- compaction just removed the
+        # earlier prefix an anchor would need, so the anchor fingerprint
+        # can never match and the estimator sums the (now-shorter) result.
+        final_token_count = self._total(ctx, result, model_name, 0)
         await _safe_emit(
             ctx,
             ContextUsageMeasuredEvent(

code_puppy-0.0.912/code_puppy/context_accounting.py

@@ -0,0 +1,154 @@
+"""Context totals anchored to completed API requests, with a heuristic delta.
+
+A receipt identifies the unchanged measured prefix, response and request
+model; new content and changed overhead remain estimates. Invalid receipts
+fall back rather than reuse stale usage. Continuation-merged responses do
+not anchor because cumulative billing is not a single measured prompt.
+
+Design, provider limitations and rollout semantics are documented in
+``docs/API_CONTEXT_ACCOUNTING.md``. The request hook owns continuation
+observation; this module never alters provider usage or billing.
+"""
+
+import hashlib
+import json
+
+from pydantic_ai.messages import ModelMessagesTypeAdapter, ModelResponse, TextPart
+
+
+def estimate_tokens_for_message(message, model_name):
+    """Import lazily so the shared entry point has no agent-package cycle."""
+    from code_puppy.agents._history import estimate_tokens_for_message as estimate
+
+    return estimate(message, model_name)
+
+
+_ANCHOR = "context_anchor"
+
+
+def fingerprint(messages) -> str:
+    """Include arguments, signatures and binary payloads, excluding receipts.
+
+    Receipts live under ``metadata`` so stamping one does not change the
+    fingerprint of the message it rides on -- otherwise a receipt would
+    invalidate itself the instant it was written.
+    """
+    clean = ModelMessagesTypeAdapter.dump_python(messages, mode="json")
+    payload = [
+        {
+            key: value
+            for key, value in message.items()
+            if key in {"parts", "instructions", "kind"}
+        }
+        for message in clean
+    ]
+    return hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest()
+
+
+def record_anchor(
+    messages, response, *, context_overhead=0, prefix=None, request_model=None
+):
+    """Stamp only a successfully completed response; usage is never fabricated.
+
+    Trusts ``prefix`` and ``response.usage`` as given -- it has no way to
+    tell a continuation-merged (cumulative) usage from an ordinary one on
+    its own. The caller must not call this helper when a continuation was observed
+    (see ``_ContinuationObserver`` in ``agents/_model_message_transform.py``).
+    Here, ``prefix=None`` computes a fingerprint; it does not suppress stamping.
+    """
+    if (
+        response.usage.input_tokens <= 0
+        or not response.model_name
+        or response.state != "complete"
+    ):
+        return
+    response.metadata = dict(response.metadata or {})
+    response.metadata[_ANCHOR] = {
+        "prefix": prefix if prefix is not None else fingerprint(messages),
+        "response": fingerprint([response]),
+        "overhead": context_overhead,
+        "request_models": sorted(
+            request_model
+            if isinstance(request_model, frozenset)
+            else {request_model if request_model is not None else response.model_name}
+        ),
+    }
+
+
+def invalidate_anchors(messages):
+    """Conservatively decline anchors after an unsuccessful request.
+
+    This is a deliberate failed-request policy, not proof the old measurement
+    became invalid. An unchanged prefix could retain its previous receipt,
+    and a retained partial response already forces fallback. We nevertheless
+    choose full estimation until a fresh completion; changing that contract
+    requires reconsidering the failure/cancellation regressions together.
+    """
+    for message in messages:
+        if isinstance(message, ModelResponse) and message.metadata:
+            message.metadata.pop(_ANCHOR, None)
+
+
+def context_tokens(messages, model_name, context_overhead=0) -> int:
+    """One total for compaction, the status bar, and context metadata.
+
+    API input includes system/tool/cache tokens. Plain text output usage is
+    an incremental approximation; it may include hidden reasoning when no
+    ThinkingPart is exposed, conservatively overestimating replay. Visible
+    reasoning/tool replay is estimated instead.
+    Only the latest response may anchor: partial/zero usage must not cause a
+    silent rewind to an older observation.
+    """
+    names = model_name if isinstance(model_name, frozenset) else {model_name}
+    estimate_model = next(iter(names)) if len(names) == 1 else None
+    for index in range(len(messages) - 1, -1, -1):
+        response = messages[index]
+        if not isinstance(response, ModelResponse):
+            continue
+        receipt = (response.metadata or {}).get(_ANCHOR)
+        if (
+            isinstance(receipt, dict)
+            and response.usage.input_tokens > 0
+            and response.state == "complete"
+            and receipt.get("request_models") == sorted(names)
+            and receipt.get("prefix") == fingerprint(messages[:index])
+            and receipt.get("response") == fingerprint([response])
+            and isinstance(receipt.get("overhead"), (int, float))
+        ):
+            replay = (
+                response.usage.output_tokens
+                if response.parts
+                and all(isinstance(p, TextPart) for p in response.parts)
+                and response.usage.output_tokens > 0
+                else estimate_tokens_for_message(response, response.model_name)
+            )
+            return (
+                response.usage.input_tokens
+                + replay
+                + context_overhead
+                - receipt["overhead"]
+                + sum(
+                    estimate_tokens_for_message(m, estimate_model)
+                    for m in messages[index + 1 :]
+                )
+            )
+        break
+    return context_overhead + sum(
+        estimate_tokens_for_message(m, estimate_model) for m in messages
+    )
+
+
+def active_model_name(agent):
+    """Use the request model identity, not an arbitrary config alias."""
+    model = getattr(agent, "cur_model", None)
+    if model is None:
+        model = getattr(getattr(agent, "pydantic_agent", None), "model", None)
+    return model_names(model) if model is not None else agent.get_model_name()
+
+
+def model_names(model):
+    """Identify a request model together with its routing candidates."""
+    candidates = getattr(model, "models", ())
+    if candidates:
+        return frozenset({model.model_name, *(m.model_name for m in candidates)})
+    return model.model_name

{code_puppy-0.0.910 → code_puppy-0.0.912}/code_puppy/model_switching.py

@@ -85,8 +85,11 @@ def _refresh_context_status(agent) -> None:
     try:
         capacity = agent._get_model_context_length()
         history = agent.get_message_history() or []
-        message_tokens = sum(agent.estimate_tokens_for_message(msg) for msg in history)
-        total_tokens = message_tokens + agent._estimate_context_overhead()
+        from code_puppy.context_accounting import active_model_name, context_tokens
+
+        total_tokens = context_tokens(
+            history, active_model_name(agent), agent._estimate_context_overhead()
+        )
         proportion = total_tokens / capacity if capacity else 0.0
         update_spinner_context(format_context_info(total_tokens, capacity, proportion))
     except Exception: