From e5a71e4a2f88128752b0437369da6325fdb407d1 Mon Sep 17 00:00:00 2001 From: zhanghanduo Date: Wed, 2 Sep 2026 13:38:25 +0800 Subject: [PATCH 1/6] fix(observers): count the token usage the loop actually produces MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `BudgetObserver` read `ctx.usage["input_tokens"]` / `["output_tokens"]`, but the loop normalises every provider payload to `prompt_tokens` / `completion_tokens` (`extract_usage`). `tokens_used` therefore stayed 0 for the whole run and neither the `budget_exhausted` intervention nor the `warn_ratio` warning could ever fire. The alias-only tests passed because they fed raw `input_tokens`. `ContextSizeGuard` carries a comment describing having fixed exactly this for itself, and `InputTokenGauge` open-codes the same two-spelling read — three consumers, three chances to get it wrong. Add `usage_input_tokens` / `usage_output_tokens` beside `extract_usage` and route all three through them, so the next consumer inherits the knowledge instead of rediscovering it. Co-Authored-By: Claude Opus 5 (1M context) --- .../components/observers/budget_observer.py | 16 ++++- .../observers/context_size_guard.py | 6 +- agent_core/runtime/loop/__init__.py | 4 ++ agent_core/runtime/loop/_response.py | 24 +++++++ agent_core/runtime/loop/llm_client.py | 4 ++ agent_core/runtime/loop/tiered_compact.py | 3 +- tests/test_budget_observer.py | 53 ++++++++++++++++ tests/test_usage_token_accessors.py | 62 +++++++++++++++++++ 8 files changed, 165 insertions(+), 7 deletions(-) create mode 100644 tests/test_usage_token_accessors.py diff --git a/agent_core/components/observers/budget_observer.py b/agent_core/components/observers/budget_observer.py index 0452dbc..a9b7af8 100644 --- a/agent_core/components/observers/budget_observer.py +++ b/agent_core/components/observers/budget_observer.py @@ -5,6 +5,7 @@ from __future__ import annotations from agent_core.loop_types import Intervention, LoopConfig, TurnContext +from agent_core.runtime.loop import usage_input_tokens, usage_output_tokens class BudgetObserver: @@ -26,10 +27,21 @@ async def on_loop_start(self, config: LoopConfig) -> None: self._warned = False async def on_llm_response(self, ctx: TurnContext) -> None: - """Accumulate token counts from the LLM response usage dict.""" + """Accumulate token counts from the LLM response usage dict. + + Reads both spellings through the shared accessors: the loop + normalises usage to ``prompt_tokens`` / ``completion_tokens`` + (``extract_usage``), while hosts and test doubles may hand in the + ``input_tokens`` / ``output_tokens`` aliases. Reading only the + aliases left ``tokens_used`` at 0 against normalised usage, so + neither the ``budget_exhausted`` intervention nor the warning at + ``warn_ratio`` could ever fire in production. + """ if ctx.usage is None: return - self.tokens_used += ctx.usage.get("input_tokens", 0) + ctx.usage.get("output_tokens", 0) + self.tokens_used += ( + usage_input_tokens(ctx.usage) + usage_output_tokens(ctx.usage) + ) ctx.metadata["budget_tokens_used"] = self.tokens_used ctx.metadata["budget_tokens_limit"] = self.max_tokens diff --git a/agent_core/components/observers/context_size_guard.py b/agent_core/components/observers/context_size_guard.py index 2b1ee13..51d2d19 100644 --- a/agent_core/components/observers/context_size_guard.py +++ b/agent_core/components/observers/context_size_guard.py @@ -10,6 +10,7 @@ LoopConfig, TurnContext, ) +from agent_core.runtime.loop import usage_input_tokens from agent_core.runtime.loop.compact import ( COMPACTION_SEQ_KEY, FORCE_COMPACTION_KEY, @@ -51,10 +52,7 @@ async def on_loop_start(self, config: LoopConfig) -> None: async def on_llm_response(self, ctx: TurnContext) -> Intervention | None: if self._tripped or ctx.usage is None: return None - # Normalised usage (infra/usage.py) exposes ``prompt_tokens``; raw - # Anthropic payloads use ``input_tokens``. Read prompt_tokens first — - # reading only ``input_tokens`` meant this guard saw 0 and never tripped. - used = int(ctx.usage.get("prompt_tokens") or ctx.usage.get("input_tokens") or 0) + used = usage_input_tokens(ctx.usage) if used <= self._limit: self._requested_at_seq = None return None diff --git a/agent_core/runtime/loop/__init__.py b/agent_core/runtime/loop/__init__.py index 6653ac9..7ab7d38 100644 --- a/agent_core/runtime/loop/__init__.py +++ b/agent_core/runtime/loop/__init__.py @@ -28,6 +28,8 @@ extract_model_name, extract_usage, is_truncated_with_text, + usage_input_tokens, + usage_output_tokens, ) from agent_core.runtime.loop.message_trimmer import ( MessageTrimmer, @@ -101,4 +103,6 @@ "extract_usage", "is_truncated_with_text", "run_agent_loop", + "usage_input_tokens", + "usage_output_tokens", ] diff --git a/agent_core/runtime/loop/_response.py b/agent_core/runtime/loop/_response.py index 826c210..03f96e7 100644 --- a/agent_core/runtime/loop/_response.py +++ b/agent_core/runtime/loop/_response.py @@ -3,6 +3,7 @@ import logging import re +from collections.abc import Mapping from typing import Any from agent_core.llm import LLMResponse @@ -119,6 +120,29 @@ def _pick_int(*candidates: Any) -> int: return 0 +def usage_input_tokens(usage: Mapping[str, Any] | None) -> int: + """Prompt-side token count from any usage mapping the loop may carry. + + ``extract_usage`` normalises to ``prompt_tokens`` / ``completion_tokens``, + but ``TurnContext.usage`` is deliberately an open mapping (see + ``UsageMetadata``): hosts hand in raw provider payloads and test doubles + that use the ``input_tokens`` / ``output_tokens`` aliases instead. A + consumer that reads only one spelling silently sees 0 — which is how the + context-size guard and the budget observer each ended up inert against + normalised usage. Read both, here, once. + """ + if usage is None: + return 0 + return int(usage.get("prompt_tokens") or usage.get("input_tokens") or 0) + + +def usage_output_tokens(usage: Mapping[str, Any] | None) -> int: + """Completion-side counterpart of :func:`usage_input_tokens`.""" + if usage is None: + return 0 + return int(usage.get("completion_tokens") or usage.get("output_tokens") or 0) + + def extract_usage(response: Any) -> UsageMetadata | None: """Extract token usage from an LLM response, normalized to OpenAI shape. diff --git a/agent_core/runtime/loop/llm_client.py b/agent_core/runtime/loop/llm_client.py index 9b88d1c..3bf6489 100644 --- a/agent_core/runtime/loop/llm_client.py +++ b/agent_core/runtime/loop/llm_client.py @@ -30,6 +30,8 @@ extract_leaked_reasoning, extract_model_name, extract_usage, + usage_input_tokens, + usage_output_tokens, ) from agent_core.runtime.loop._runaway import ( RUNAWAY_STATE_KEY, @@ -62,4 +64,6 @@ "extract_model_name", "extract_usage", "is_truncated_with_text", + "usage_input_tokens", + "usage_output_tokens", ] diff --git a/agent_core/runtime/loop/tiered_compact.py b/agent_core/runtime/loop/tiered_compact.py index fe52308..6ae7acd 100644 --- a/agent_core/runtime/loop/tiered_compact.py +++ b/agent_core/runtime/loop/tiered_compact.py @@ -13,6 +13,7 @@ TurnContext, ) from agent_core.messages import Message, is_tool_msg, text_of, user_msg +from agent_core.runtime.loop._response import usage_input_tokens from agent_core.runtime.loop.compact import ( INPUT_ESTIMATE_KEY, KeepLastNToolResultsCompactor, @@ -109,7 +110,7 @@ async def on_llm_response(self, ctx: TurnContext) -> None: u = ctx.usage or {} # Normalised usage uses ``prompt_tokens``; accept raw Anthropic # ``input_tokens`` as a fallback. - self.tokens = int(u.get("prompt_tokens") or u.get("input_tokens") or 0) + self.tokens = usage_input_tokens(u) # Both sides of the scale must describe the SAME request. Estimating # ``ctx.messages`` here cannot: the loop appends this turn's assistant # reply before building the context, and the per-call system addendum diff --git a/tests/test_budget_observer.py b/tests/test_budget_observer.py index 69dc156..7edcfed 100644 --- a/tests/test_budget_observer.py +++ b/tests/test_budget_observer.py @@ -115,3 +115,56 @@ async def test_exposes_usage_in_metadata() -> None: await obs.on_llm_response(ctx) assert ctx.metadata["budget_tokens_used"] == 1500 assert ctx.metadata["budget_tokens_limit"] == 50_000 + + +# --------------------------------------------------------------------------- +# Normalised usage — the shape the loop actually produces +# --------------------------------------------------------------------------- +# +# ``extract_usage`` normalises every provider payload to ``prompt_tokens`` / +# ``completion_tokens``. Reading only the ``input_tokens`` / ``output_tokens`` +# aliases left ``tokens_used`` at 0 for the whole run, so neither intervention +# below could fire in production while the alias-based tests above passed. + + +@pytest.mark.asyncio +async def test_counts_normalised_usage_keys() -> None: + obs = BudgetObserver(max_tokens=1000) + await obs.on_llm_response(_make_ctx({"prompt_tokens": 300, "completion_tokens": 200})) + assert obs.tokens_used == 500 + + +@pytest.mark.asyncio +async def test_budget_exhausted_fires_on_normalised_usage() -> None: + obs = BudgetObserver(max_tokens=1000) + ctx = _make_ctx({"prompt_tokens": 900, "completion_tokens": 200}) + await obs.on_llm_response(ctx) + + intervention = await obs.on_turn_end(ctx) + + assert isinstance(intervention, Intervention) + assert intervention.stop_reason == "budget_exhausted" + + +@pytest.mark.asyncio +async def test_warning_fires_on_normalised_usage() -> None: + obs = BudgetObserver(max_tokens=1000, warn_ratio=0.8) + ctx = _make_ctx({"prompt_tokens": 800, "completion_tokens": 50}) + await obs.on_llm_response(ctx) + + intervention = await obs.on_turn_end(ctx) + + assert isinstance(intervention, Intervention) + assert intervention.stop_reason is None + assert "85%" in intervention.inject_messages[0] + + +@pytest.mark.asyncio +async def test_alias_and_normalised_keys_are_not_double_counted() -> None: + """A payload carrying both spellings must count each side once.""" + obs = BudgetObserver(max_tokens=1000) + await obs.on_llm_response(_make_ctx({ + "prompt_tokens": 300, "input_tokens": 300, + "completion_tokens": 200, "output_tokens": 200, + })) + assert obs.tokens_used == 500 diff --git a/tests/test_usage_token_accessors.py b/tests/test_usage_token_accessors.py new file mode 100644 index 0000000..71fcf66 --- /dev/null +++ b/tests/test_usage_token_accessors.py @@ -0,0 +1,62 @@ +"""One place to read token counts out of a usage mapping. + +``TurnContext.usage`` is deliberately an open ``Mapping`` (see the +``UsageMetadata`` docstring): ``extract_usage`` normalises to +``prompt_tokens`` / ``completion_tokens``, but hosts hand in raw provider +payloads and test doubles use the ``input_tokens`` / ``output_tokens`` +aliases. Every consumer that read only one spelling has been silently inert at +some point — the context-size guard and the budget observer both were — so the +accessors are the single place that knows about both. +""" + +from __future__ import annotations + +import pytest + +from agent_core.runtime.loop import usage_input_tokens, usage_output_tokens + + +@pytest.mark.parametrize( + ("usage", "expected"), + [ + ({"prompt_tokens": 120}, 120), + ({"input_tokens": 120}, 120), + ({"prompt_tokens": 120, "input_tokens": 999}, 120), + ({"prompt_tokens": 0, "input_tokens": 120}, 120), + ({}, 0), + (None, 0), + ({"prompt_tokens": None}, 0), + ], +) +def test_input_tokens(usage, expected): + assert usage_input_tokens(usage) == expected + + +@pytest.mark.parametrize( + ("usage", "expected"), + [ + ({"completion_tokens": 40}, 40), + ({"output_tokens": 40}, 40), + ({"completion_tokens": 40, "output_tokens": 999}, 40), + ({"completion_tokens": 0, "output_tokens": 40}, 40), + ({}, 0), + (None, 0), + ], +) +def test_output_tokens(usage, expected): + assert usage_output_tokens(usage) == expected + + +def test_reads_what_extract_usage_produces(): + """Pins the accessors against the real producer, not a hand-written dict.""" + from agent_core.llm import LLMResponse + from agent_core.runtime.loop import extract_usage + + usage = extract_usage(LLMResponse( + content="hi", + usage={"prompt_tokens": 31, "completion_tokens": 7}, + )) + + assert usage is not None + assert usage_input_tokens(usage) == 31 + assert usage_output_tokens(usage) == 7 From 871fff32df39969f88eda248e9762118bcf510a1 Mon Sep 17 00:00:00 2001 From: zhanghanduo Date: Wed, 2 Sep 2026 13:42:25 +0800 Subject: [PATCH 2/6] fix(middleware): orphan tool results, duplicate after_llm, shared-scope trace ids MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three defects in the LLM middleware layer. `SummarizationMiddleware` advanced the split point past leading tool messages to avoid an orphan tool result, but bounded the loop at `len(rest) - 1` and so stopped one element short. With `keep_recent=3` over a tail of `[assistant(tool_calls), tool, tool, tool]` the kept window was `[tool]` — precisely the orphan the loop exists to prevent, and the Azure `400 No tool call found for function call output` it was written to stop. Running off the end is the safe outcome: everything is summarised and the kept window is empty. `LLMProxy.stream` emitted `run_after` from a `finally` nested inside the retry `while True`, so a stream failing before its first chunk called every middleware's `after_llm` with an empty `LLMResponse` and then again with the real one. Cost/usage accounting and history-appending middleware double-counted, and the phantom response read as a genuine empty answer. Hoist the block so it runs once per call. `LLMProxy._make_ctx` wrote its freshly minted `step_id` / `prompt_id` back into the shared `ExecutionScope.metadata`. Two concurrent `chat` / `stream` calls on one scope — which `_next_call_index` documents as supported — then clobbered each other and `_log_llm_exception` attributed a failure to the other call's step. Only `session_id` is scope-stable, so only it is published back. Each fix has a regression test verified to fail against the previous code. Co-Authored-By: Claude Opus 5 (1M context) --- agent_core/components/middleware/llm/proxy.py | 158 +++++++------- .../middleware/llm/summarization.py | 10 +- tests/test_llm_proxy_stream_lifecycle.py | 196 ++++++++++++++++++ tests/test_summarization_orphan_tool.py | 122 +++++++++++ 4 files changed, 414 insertions(+), 72 deletions(-) create mode 100644 tests/test_llm_proxy_stream_lifecycle.py create mode 100644 tests/test_summarization_orphan_tool.py diff --git a/agent_core/components/middleware/llm/proxy.py b/agent_core/components/middleware/llm/proxy.py index 3ddcbcb..037f751 100644 --- a/agent_core/components/middleware/llm/proxy.py +++ b/agent_core/components/middleware/llm/proxy.py @@ -106,7 +106,13 @@ def _make_ctx(self, call_index: int) -> LLMCallContext: refresh_prompt_id=True, ) metadata["step_id"] = default_step_id - scope.metadata.update(metadata) + # Only the scope-stable identifier goes back to the scope. + # ``step_id`` / ``prompt_id`` identify THIS call, and an + # execution scope is shared by concurrent ``chat`` / ``stream`` + # invocations (see ``_next_call_index``) — writing them onto the + # scope let two calls overwrite each other, so a failure was + # reported against the other call's step. + scope.metadata.setdefault("session_id", metadata["session_id"]) # Stash the inner model id so a cost-accounting middleware can record # spend even when the provider's response_metadata doesn't echo # model_name (some models behind an OpenAI-compat gateway don't). @@ -176,79 +182,89 @@ async def stream( ctx = self._make_ctx(self._next_call_index()) messages = await self.chain.run_before(ctx, messages) attempt = 0 - while True: - start_time = time.time() - full_content = "" - full_reasoning = "" - stream_error: Exception | None = None - any_chunk_yielded = False - try: - async for delta in self.inner.stream( - messages, - tools=tools, - temperature=temperature, - max_tokens=max_tokens, - extra_headers=extra_headers, - timeout=timeout, - ): - full_content += delta.content or "" - full_reasoning += delta.reasoning_content or "" - any_chunk_yielded = True - # Per-chunk middleware hook. A middleware returning - # True (e.g. StreamRepetitionDetector noticing a - # degenerate loop) tells us to stop consuming the - # inner stream and exit cleanly — the partial - # response still flows through ``after_llm`` in the - # finally block so observers see the truncated - # content rather than nothing. The delta that - # triggered the abort IS still yielded so the - # consumer's accumulator stays consistent with the - # LLMResponse we'll synthesise. - yield delta - if await self.chain.run_on_chunk( - ctx, - delta, - full_content, + start_time = time.time() + full_content = "" + full_reasoning = "" + stream_error: Exception | None = None + try: + while True: + start_time = time.time() + full_content = "" + full_reasoning = "" + stream_error = None + any_chunk_yielded = False + try: + async for delta in self.inner.stream( + messages, + tools=tools, + temperature=temperature, + max_tokens=max_tokens, + extra_headers=extra_headers, + timeout=timeout, ): - ctx.metadata["stream_aborted_by_middleware"] = True - logger.info( - "LLMProxy: stream aborted by middleware after %d chars", - len(full_content), - ) - break - break - except Exception as e: - stream_error = e - await _log_llm_exception(ctx, e) - if not any_chunk_yielded: - should_retry = await self.chain.run_on_llm_error( - ctx, - e, - attempt, - ) - if should_retry: - attempt += 1 - logger.info( - "LLMProxy: retrying stream after attempt %d (%s)", + full_content += delta.content or "" + full_reasoning += delta.reasoning_content or "" + any_chunk_yielded = True + # Per-chunk middleware hook. A middleware returning + # True (e.g. StreamRepetitionDetector noticing a + # degenerate loop) tells us to stop consuming the + # inner stream and exit cleanly — the partial + # response still flows through ``after_llm`` in the + # finally block so observers see the truncated + # content rather than nothing. The delta that + # triggered the abort IS still yielded so the + # consumer's accumulator stays consistent with the + # LLMResponse we'll synthesise. + yield delta + if await self.chain.run_on_chunk( + ctx, + delta, + full_content, + ): + ctx.metadata["stream_aborted_by_middleware"] = True + logger.info( + "LLMProxy: stream aborted by middleware after %d chars", + len(full_content), + ) + break + break + except Exception as e: + stream_error = e + await _log_llm_exception(ctx, e) + if not any_chunk_yielded: + should_retry = await self.chain.run_on_llm_error( + ctx, + e, attempt, - type(e).__name__, ) - continue - raise - finally: - duration_ms = int( - (time.time() - start_time) * 1000, - ) - ctx.metadata["duration_ms"] = duration_ms - if stream_error: - ctx.metadata["error"] = str(stream_error) - await self.chain.run_after( - ctx, - LLMResponse( - content=full_content, - reasoning_content=full_reasoning, - ), - ) + if should_retry: + attempt += 1 + logger.info( + "LLMProxy: retrying stream after attempt %d (%s)", + attempt, + type(e).__name__, + ) + continue + raise + finally: + # Once per ``stream()`` call, not once per retry attempt. This + # block used to sit inside ``while True``, so a stream that failed + # before its first chunk and retried handed every middleware an + # empty ``LLMResponse`` and then the real one — any usage/cost + # accounting or history-appending middleware double-counted. + duration_ms = int( + (time.time() - start_time) * 1000, + ) + ctx.metadata["duration_ms"] = duration_ms + if stream_error: + ctx.metadata["error"] = str(stream_error) + await self.chain.run_after( + ctx, + LLMResponse( + content=full_content, + reasoning_content=full_reasoning, + ), + ) def __getattr__(self, name: str) -> Any: # Defer unknown attributes to the inner client so callers that diff --git a/agent_core/components/middleware/llm/summarization.py b/agent_core/components/middleware/llm/summarization.py index 42efb1f..c3641fa 100644 --- a/agent_core/components/middleware/llm/summarization.py +++ b/agent_core/components/middleware/llm/summarization.py @@ -111,7 +111,15 @@ async def before_llm(self, ctx: LLMCallContext, messages: list[Message]) -> list # return HTTP 400 "No tool call found for function call output # with call_id ...". split_idx = len(rest) - keep_recent - while split_idx < len(rest) - 1 and rest[split_idx].get("role") == "tool": + # ``len(rest)``, not ``len(rest) - 1``: stopping one short left the + # window starting on a tool message whenever the tail was all tool + # results — e.g. keep_recent=3 over + # ``[assistant(tool_calls), tool, tool, tool]`` kept ``[tool]``, the + # exact orphan this loop exists to prevent. Running off the end is + # the safe outcome: everything gets summarised and the kept window is + # empty, rather than shipping a tool result whose ``tool_calls`` was + # summarised away. + while split_idx < len(rest) and rest[split_idx].get("role") == "tool": split_idx += 1 to_summarize = rest[:split_idx] keep = rest[split_idx:] diff --git a/tests/test_llm_proxy_stream_lifecycle.py b/tests/test_llm_proxy_stream_lifecycle.py new file mode 100644 index 0000000..f836b3a --- /dev/null +++ b/tests/test_llm_proxy_stream_lifecycle.py @@ -0,0 +1,196 @@ +"""``LLMProxy.stream`` must call ``after_llm`` exactly once per call. + +The ``finally`` that emits ``run_after`` used to sit *inside* the retry +``while True``, so a stream that failed before its first chunk and was retried +handed every middleware an empty ``LLMResponse`` and then the real one. Any +middleware that accumulates — cost/usage accounting, history appending — +double-counted, and the phantom empty response looked like a real (empty) +answer to anything reading content. + +Also pinned here: an execution scope is shared by concurrent ``chat`` / +``stream`` calls, so per-call trace ids must stay on the call's own context +instead of being written back onto the scope. +""" + +from __future__ import annotations + +import asyncio + +import pytest + +from agent_core.components.middleware.llm.base import ( + LLMCallContext, + LLMMiddleware, + LLMMiddlewareChain, +) +from agent_core.components.middleware.llm.proxy import LLMProxy +from agent_core.execution_context import ( + build_execution_scope, + reset_current_execution_scope, + set_current_execution_scope, +) +from agent_core.llm import LLMResponse, StreamDelta +from agent_core.messages import user_msg + + +class _Recorder(LLMMiddleware): + """Records every after_llm, and can demand one retry.""" + + name = "recorder" + + def __init__(self, *, retries: int = 0) -> None: + super().__init__() + self.after_calls: list[str] = [] + self.error_calls: int = 0 + self._retries_left = retries + + async def after_llm( + self, ctx: LLMCallContext, response: LLMResponse, + ) -> LLMResponse: + self.after_calls.append(response.content) + return response + + async def on_llm_error( + self, ctx: LLMCallContext, error: Exception, attempt: int, + ) -> bool: + self.error_calls += 1 + if self._retries_left > 0: + self._retries_left -= 1 + return True + return False + + +class _Inner: + """Fails the first ``fail_times`` streams before its first chunk.""" + + model = "test-model" + + def __init__(self, fail_times: int = 0) -> None: + self.fail_times = fail_times + self.stream_calls = 0 + + async def stream(self, messages, **kwargs): + self.stream_calls += 1 + if self.stream_calls <= self.fail_times: + raise RuntimeError("upstream reset") + for piece in ("he", "llo"): + yield StreamDelta(content=piece) + + +def _chain(mw: LLMMiddleware) -> LLMMiddlewareChain: + chain = LLMMiddlewareChain() + chain.add(mw) + return chain + + +async def _drain(proxy: LLMProxy) -> str: + return "".join([d.content or "" async for d in proxy.stream([user_msg("q")])]) + + +@pytest.mark.asyncio +async def test_after_llm_fires_once_on_a_clean_stream(): + rec = _Recorder() + proxy = LLMProxy(inner=_Inner(), chain=_chain(rec), role_id="r") + + assert await _drain(proxy) == "hello" + assert rec.after_calls == ["hello"] + + +@pytest.mark.asyncio +async def test_after_llm_fires_once_across_a_retry(): + """The empty first attempt must NOT reach after_llm.""" + rec = _Recorder(retries=1) + inner = _Inner(fail_times=1) + proxy = LLMProxy(inner=inner, chain=_chain(rec), role_id="r") + + assert await _drain(proxy) == "hello" + assert inner.stream_calls == 2 + assert rec.error_calls == 1 + assert rec.after_calls == ["hello"] + + +@pytest.mark.asyncio +async def test_after_llm_still_fires_once_when_every_attempt_fails(): + rec = _Recorder(retries=1) + inner = _Inner(fail_times=99) + proxy = LLMProxy(inner=inner, chain=_chain(rec), role_id="r") + + with pytest.raises(RuntimeError, match="upstream reset"): + await _drain(proxy) + + assert inner.stream_calls == 2 + assert rec.after_calls == [""] + + +@pytest.mark.asyncio +async def test_failure_metadata_describes_the_last_attempt(): + rec = _Recorder() + proxy = LLMProxy(inner=_Inner(fail_times=1), chain=_chain(rec), role_id="r") + seen: dict[str, object] = {} + + class _Capture(LLMMiddleware): + name = "capture" + + async def after_llm(self, ctx, response): + seen.update(ctx.metadata) + return response + + proxy.chain.add(_Capture()) + with pytest.raises(RuntimeError): + await _drain(proxy) + + assert seen.get("error") == "upstream reset" + assert "duration_ms" in seen + + +# ── per-call trace ids stay off the shared scope ────────────────────────── + + +@pytest.mark.asyncio +async def test_concurrent_calls_do_not_clobber_each_others_step_id(): + scope = build_execution_scope( + task_id="t1", phase_id="research", role_id="r", state={}, + ) + token = set_current_execution_scope(scope) + try: + proxy = LLMProxy(inner=_Inner(), chain=LLMMiddlewareChain(), role_id="r") + ctx_a = proxy._make_ctx(proxy._next_call_index()) + ctx_b = proxy._make_ctx(proxy._next_call_index()) + + assert ctx_a.metadata["step_id"] != ctx_b.metadata["step_id"] + assert ctx_a.metadata["prompt_id"] != ctx_b.metadata["prompt_id"] + # The second call must not have rewritten the first one's identity. + assert ctx_a.metadata["step_id"] == "research:llm:1" + assert ctx_b.metadata["step_id"] == "research:llm:2" + # ...and neither leaked a per-call id onto the shared scope. + assert "step_id" not in scope.metadata + assert "prompt_id" not in scope.metadata + # The scope-stable identifier is still published once. + assert scope.metadata["session_id"] == ctx_a.metadata["session_id"] + finally: + reset_current_execution_scope(token) + + +@pytest.mark.asyncio +async def test_interleaved_streams_keep_their_own_step_id(): + """The race the scope write created, exercised through the public API.""" + scope = build_execution_scope( + task_id="t1", phase_id="research", role_id="r", state={}, + ) + token = set_current_execution_scope(scope) + try: + steps: list[str] = [] + + class _Step(LLMMiddleware): + name = "step" + + async def after_llm(self, ctx, response): + steps.append(str(ctx.metadata.get("step_id"))) + return response + + proxy = LLMProxy(inner=_Inner(), chain=_chain(_Step()), role_id="r") + await asyncio.gather(_drain(proxy), _drain(proxy)) + + assert sorted(steps) == ["research:llm:1", "research:llm:2"] + finally: + reset_current_execution_scope(token) diff --git a/tests/test_summarization_orphan_tool.py b/tests/test_summarization_orphan_tool.py new file mode 100644 index 0000000..c9c965a --- /dev/null +++ b/tests/test_summarization_orphan_tool.py @@ -0,0 +1,122 @@ +"""The kept window after summarization must never start on a tool message. + +An orphan tool result — one whose assistant ``tool_calls`` was compressed into +the summary — makes Azure reject the request outright: +``400 No tool call found for function call output with call_id ...``. The +middleware already advanced the split point past leading tool messages, but the +loop bound stopped one element short, so a tail made entirely of tool results +still left exactly one orphan behind. +""" + +from __future__ import annotations + +import pytest + +from agent_core.components.middleware.llm.base import LLMCallContext +from agent_core.components.middleware.llm.summarization import ( + SummarizationMiddleware, +) + + +def _ctx() -> LLMCallContext: + return LLMCallContext(task_id="t1", role_id="r", phase_id="p", call_index=1) + + +def _sys(text: str = "sys") -> dict: + return {"role": "system", "content": text} + + +def _tool(text: str) -> dict: + return {"role": "tool", "content": text, "tool_call_id": text} + + +def _assistant_with_calls(text: str = "calling") -> dict: + return { + "role": "assistant", "content": text, + "tool_calls": [{"id": "c1", "function": {"name": "web_search"}}], + } + + +def _big(text: str, *, repeat: int = 4000) -> str: + return text * repeat + + +async def _compress(messages: list[dict], *, keep_recent: int) -> list[dict]: + mw = SummarizationMiddleware(threshold=10, keep_recent=keep_recent) + return await mw.before_llm(_ctx(), messages) + + +def _assert_no_orphan_tool(result: list[dict]) -> None: + """No tool message may appear before the assistant that requested it.""" + seen_tool_calls = False + for msg in result: + if msg.get("role") == "assistant" and msg.get("tool_calls"): + seen_tool_calls = True + if msg.get("role") == "tool": + assert seen_tool_calls, f"orphan tool message in {[m['role'] for m in result]}" + + +@pytest.mark.asyncio +async def test_all_tool_tail_leaves_no_orphan(): + """keep_recent=3 over [assistant(tool_calls), tool, tool, tool].""" + messages = [ + _sys(), + {"role": "user", "content": _big("question ")}, + _assistant_with_calls(), + _tool(_big("r1 ")), + _tool(_big("r2 ")), + _tool(_big("r3 ")), + ] + + result = await _compress(messages, keep_recent=3) + + assert len(result) < len(messages), "should have compressed" + _assert_no_orphan_tool(result) + assert result[-1].get("role") != "tool" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("keep_recent", [1, 2, 3, 4, 5]) +async def test_no_orphan_for_any_keep_recent_over_a_tool_tail(keep_recent): + messages = [ + _sys(), + {"role": "user", "content": _big("q ")}, + _assistant_with_calls(), + _tool(_big("r1 ")), + _tool(_big("r2 ")), + _tool(_big("r3 ")), + _tool(_big("r4 ")), + ] + + result = await _compress(messages, keep_recent=keep_recent) + + _assert_no_orphan_tool(result) + + +@pytest.mark.asyncio +async def test_split_still_keeps_a_healthy_tail_intact(): + """Guard against over-advancing: a non-tool tail must be kept as-is.""" + messages = [ + _sys(), + {"role": "user", "content": _big("q ")}, + _assistant_with_calls(), + _tool(_big("r1 ")), + {"role": "assistant", "content": "answer A"}, + {"role": "user", "content": "follow up"}, + {"role": "assistant", "content": "answer B"}, + ] + + result = await _compress(messages, keep_recent=3) + + assert [m["content"] for m in result[-3:]] == [ + "answer A", "follow up", "answer B", + ] + _assert_no_orphan_tool(result) + + +@pytest.mark.asyncio +async def test_below_threshold_is_untouched(): + messages = [_sys(), {"role": "user", "content": "hi"}, _tool("r")] + mw = SummarizationMiddleware(threshold=10_000, keep_recent=2) + + assert await mw.before_llm(_ctx(), messages) is messages From 7eb7aa2b909f96d7654f45abacbb160117709946 Mon Sep 17 00:00:00 2001 From: zhanghanduo Date: Wed, 2 Sep 2026 13:45:46 +0800 Subject: [PATCH 3/6] fix(dag): honour max_retries, allow None deltas, keep task_id out of the filter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three defects in the node-wrapper stack `DynamicGraphBuilder.build` assembles. Each is invisible in the default configuration and only surfaces once a second feature is switched on, which is why the suite missed all three. `metadata["max_retries"]` was gated on `err_result is None`, i.e. on a middleware *suppressing* the error. With the default `ExecutionMiddleware.on_error` — which returns the error — a node declaring retries raised on attempt 0 and was never retried. Retry now depends on attempts remaining, which is what the field means. The suppression branch keeps its documented behaviour (advance on a synthetic delta) but logs the swallowed exception, so a genuine crash is no longer traceless. `MiniDAGRunner` explicitly accepts a `None` delta as "no state change", but the middleware wrapper passed it straight into `run_after_phase` and then `result.get(...)`. A node that worked with no middleware registered raised `AttributeError: 'NoneType' has no attribute 'get'` the moment a chain was added — and the surrounding `except` reported the contract violation as a middleware error. Normalise `None` to `{}` before the hook. The NodeContext wrapper was applied *inside* the context filter, so `task_id_getter` read the already-filtered state. Any node whose `ContextPolicy.include_fields` omitted `task_id` — the normal use of the policy — saw `ctx.task_id == ""`, and every task-scoped lookup or telemetry emission from it was attributed to the empty task id. The filter now sits innermost; it forwards extra args so the context wrapper can sit outside it, and `build` decides the arity question from the original node signature rather than from a wrapper's. Co-Authored-By: Claude Opus 5 (1M context) --- agent_core/runtime/dag/graph_builder.py | 59 ++++- tests/test_graph_builder_node_contract.py | 269 ++++++++++++++++++++++ 2 files changed, 317 insertions(+), 11 deletions(-) create mode 100644 tests/test_graph_builder_node_contract.py diff --git a/agent_core/runtime/dag/graph_builder.py b/agent_core/runtime/dag/graph_builder.py index 0bc9741..6e542b0 100644 --- a/agent_core/runtime/dag/graph_builder.py +++ b/agent_core/runtime/dag/graph_builder.py @@ -77,12 +77,20 @@ def build(self, spec: PipelineSpec, state_type: type | None = None) -> MiniDAG: # Use resolved_nodes (works with both nodes and phases via bridge) for node_def in spec.resolved_nodes: node_fn = self._resolve_node_function_from_def(node_def) - # Wrap with NodeContext injection (before context filter) + # Decided from the original signature: every wrapper below takes + # ``(state, ...)`` and would fool the arity probe. + needs_ctx = _needs_context(node_fn) + # Context filter goes innermost so the NodeContext built outside it + # still sees the *unfiltered* state. ``task_id_getter`` reads + # ``state["task_id"]``, which a ``ContextPolicy.include_fields`` + # list normally omits — wrapping the other way round gave those + # nodes ``ctx.task_id == ""`` and attributed every task-scoped + # lookup and telemetry emission to the empty task id. + node_fn = _wrap_with_context_filter(node_def, node_fn) node_fn = _wrap_with_node_context( - node_def, node_fn, self._node_context_factory + node_def, node_fn, self._node_context_factory, + needs_context=needs_ctx, ) - # Wrap with context filter - node_fn = _wrap_with_context_filter(node_def, node_fn) # Wrap with field truncation if node_def.compression and node_def.compression.max_field_tokens: node_fn = _wrap_with_field_truncation(node_def, node_fn) @@ -189,6 +197,8 @@ def _wrap_with_node_context( node_def: NodeDefinition, fn: Any, node_context_factory: NodeContextFactory, + *, + needs_context: bool | None = None, ) -> Any: """Wrap node function to inject the node context as second arg. @@ -198,8 +208,13 @@ def _wrap_with_node_context( If the function has a single-param ``(state)`` signature (legacy), it passes through unchanged — backward compatible with solver specs that predate the NodeContext facade (e.g. solvers/gaia). + + ``needs_context`` lets the caller decide from the *original* node + signature. :meth:`DynamicGraphBuilder.build` must pass it, because it + applies this wrapper on top of the context filter, whose own + ``(state, *extra)`` signature says nothing about what the node wants. """ - if not _needs_context(fn): + if not (needs_context if needs_context is not None else _needs_context(fn)): return fn # legacy (state) signature — no injection async def wrapped(state: dict[str, Any]) -> dict[str, Any]: @@ -266,9 +281,11 @@ def _wrap_with_context_filter( if policy.filter_fn is None and policy.include_fields is None: return node_fn # full state, skip wrapping - async def wrapped(state: dict[str, Any]) -> dict[str, Any]: + async def wrapped( + state: dict[str, Any], *extra: Any, + ) -> dict[str, Any]: filtered = apply_context_filter(policy, state) - return await node_fn(filtered) + return await node_fn(filtered, *extra) wrapped.__name__ = getattr( node_fn, "__name__", f"node_{node_def.node_id}", @@ -356,6 +373,16 @@ async def wrapped(state: dict[str, Any]) -> dict[str, Any]: for attempt in range(1 + max_retries): try: result = await node_fn(state) + # ``MiniDAGRunner`` treats a ``None`` delta as "no state + # change" (``if delta is None: delta = {}``), so a node is + # allowed to return it. Normalise here or the middleware + # hook below is handed None and ``result.get`` raises — a + # node that worked with no middleware registered would + # start failing the moment a chain was added, and the + # surrounding ``except`` would misreport the contract + # violation as a middleware error. + if result is None: + result = {} result = await middleware_chain.run_after_phase( ctx, result, ) @@ -377,10 +404,15 @@ async def wrapped(state: dict[str, Any]) -> dict[str, Any]: err_result = await middleware_chain.run_on_error( ctx, e, ) - if err_result is None and attempt < max_retries: + # Retry on any failure while attempts remain. Gating this + # on ``err_result is None`` made retrying depend on a + # middleware *suppressing* the error, so with the default + # ``ExecutionMiddleware.on_error`` (which returns the + # error) a node declaring ``metadata["max_retries"]`` was + # never retried at all. + if attempt < max_retries: logger.warning( - "Middleware suppressed error in '%s', " - "retrying (%d/%d): %s", + "Node '%s' failed, retrying (%d/%d): %s", node_def.node_id, attempt + 1, max_retries, @@ -388,10 +420,15 @@ async def wrapped(state: dict[str, Any]) -> dict[str, Any]: ) continue if err_result is None: + # Middleware returned None = "handled, do not raise". + # The DAG then advances on a synthetic delta, so log + # the swallowed exception itself — otherwise a genuine + # node crash leaves no trace anywhere. logger.warning( "Middleware suppressed error in '%s' " - "(no more retries)", + "(retries exhausted); advancing the DAG", node_def.node_id, + exc_info=e, ) return { "current_phase": node_def.node_id, diff --git a/tests/test_graph_builder_node_contract.py b/tests/test_graph_builder_node_contract.py new file mode 100644 index 0000000..093a7ed --- /dev/null +++ b/tests/test_graph_builder_node_contract.py @@ -0,0 +1,269 @@ +"""Node-wrapper contracts in ``DynamicGraphBuilder.build``. + +Three defects lived in the wrapper stack, all invisible in the default +configuration and all surfacing only once a *second* feature was switched on: + +* ``metadata["max_retries"]`` only retried when a middleware *suppressed* the + error, so with the default ``ExecutionMiddleware.on_error`` (which returns + the error) it never retried at all; +* a node returning ``None`` — which ``MiniDAGRunner`` explicitly supports — + crashed the middleware wrapper with ``AttributeError``, so registering a + middleware chain broke nodes that had always worked; +* the NodeContext wrapper sat *inside* the context filter, so + ``ctx.task_id`` read the filtered state and came back empty for any node + whose ``ContextPolicy.include_fields`` omitted ``task_id`` — the normal use + of the policy. +""" + +from __future__ import annotations + +from typing import Any + +import pytest + +from agent_core.models.pipeline_spec import ( + ContextPolicy, + NodeDefinition, + PipelineSpec, + TransitionSpec, +) +from agent_core.protocols import ExecutionMiddleware, PhaseContext +from agent_core.runtime.dag.graph_builder import DynamicGraphBuilder + +# ── node functions ──────────────────────────────────────────────────────── +# +# Injected through ``_resolve_node_function_from_def`` rather than a dotted +# path: importing this module by name gives pytest's copy and importlib's copy +# separate module objects, so a call counter at module level would be counted +# in the one the test cannot see. + + +def _counting_node(*, fails_before: int) -> tuple[Any, list[int]]: + """Node that fails its first ``fails_before - 1`` attempts.""" + calls: list[int] = [] + + async def node(state): + calls.append(1) + if len(calls) < fails_before: + raise RuntimeError(f"attempt {len(calls)} failed") + return {"current_phase": "n", "attempts": len(calls)} + + return node, calls + + +def _always_failing_node() -> tuple[Any, list[int]]: + calls: list[int] = [] + + async def node(state): + calls.append(1) + raise RuntimeError("permanently broken") + + return node, calls + + +async def none_node(state): + """A node that reports no state change — legal per MiniDAGRunner.""" + return None + + +async def task_id_node(state, ctx): + return {"current_phase": "tid", "seen_task_id": ctx.task_id} + + +class _Chain: + """Minimal ``PhaseMiddlewareChain`` over real ``ExecutionMiddleware``.""" + + def __init__(self, *mws: ExecutionMiddleware) -> None: + self._mws = list(mws) + + @property + def middlewares(self) -> list[ExecutionMiddleware]: + return list(self._mws) + + async def run_before_phase(self, ctx: PhaseContext) -> PhaseContext: + for mw in self._mws: + ctx = await mw.before_phase(ctx) + return ctx + + async def run_after_phase( + self, ctx: PhaseContext, result: dict[str, Any], + ) -> dict[str, Any]: + for mw in self._mws: + result = await mw.after_phase(ctx, result) + return result + + async def run_on_error( + self, ctx: PhaseContext, error: Exception, + ) -> Exception | None: + out: Exception | None = error + for mw in self._mws: + out = await mw.on_error(ctx, error) + return out + + +class _Suppressing(ExecutionMiddleware): + async def on_error(self, ctx: PhaseContext, error: Exception) -> None: + return None + + +def _builder(chain: Any | None, node_fn: Any = None) -> DynamicGraphBuilder: + builder = DynamicGraphBuilder() + builder._get_middleware_chain = lambda: chain # type: ignore[method-assign] + if node_fn is not None: + builder._resolve_node_function_from_def = ( # type: ignore[method-assign] + lambda node_def: node_fn + ) + return builder + + +def _spec(node: NodeDefinition) -> PipelineSpec: + return PipelineSpec( + pipeline_id="t", + nodes=[node], + transitions=[TransitionSpec(from_phase=node.node_id, to_phase="__END__")], + entry_point=node.node_id, + terminal_nodes=[node.node_id], + ) + + +def _node(fn: str, **kw: Any) -> NodeDefinition: + return NodeDefinition( + node_id=kw.pop("node_id", "n"), role_id="researcher", + node_function=f"tests.test_graph_builder_node_contract.{fn}", **kw, + ) + + +# ── max_retries ─────────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_max_retries_works_with_the_default_middleware(): + """The default ``on_error`` returns the error; that must not disable retry.""" + node, calls = _counting_node(fails_before=3) + dag = _builder(_Chain(ExecutionMiddleware()), node).build( + _spec(_node("x", node_id="n", metadata={"max_retries": 2})), + ) + + result = await dag.nodes["n"]({"task_id": "t1"}) + + assert len(calls) == 3 + assert result["attempts"] == 3 + + +@pytest.mark.asyncio +async def test_retries_are_bounded_and_the_error_surfaces(): + node, calls = _always_failing_node() + dag = _builder(_Chain(ExecutionMiddleware()), node).build( + _spec(_node("x", node_id="n", metadata={"max_retries": 2})), + ) + + with pytest.raises(RuntimeError, match="permanently broken"): + await dag.nodes["n"]({"task_id": "t1"}) + + assert len(calls) == 3 # 1 initial + 2 retries + + +@pytest.mark.asyncio +async def test_no_retries_declared_means_one_attempt(): + node, calls = _always_failing_node() + dag = _builder(_Chain(ExecutionMiddleware()), node).build( + _spec(_node("x", node_id="n")), + ) + + with pytest.raises(RuntimeError): + await dag.nodes["n"]({"task_id": "t1"}) + + assert len(calls) == 1 + + +@pytest.mark.asyncio +async def test_suppressing_middleware_still_retries_then_advances(): + node, calls = _always_failing_node() + dag = _builder(_Chain(_Suppressing()), node).build( + _spec(_node("x", node_id="n", metadata={"max_retries": 1})), + ) + + result = await dag.nodes["n"]({"task_id": "t1"}) + + assert len(calls) == 2 + assert result["current_phase"] == "n" + + +# ── a node may return None ──────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_none_result_survives_a_registered_middleware_chain(): + """Worked with no chain; must not start raising once one is registered.""" + dag = _builder(_Chain(ExecutionMiddleware())).build( + _spec(_node("none_node", node_id="none")), + ) + + result = await dag.nodes["none"]({"task_id": "t1"}) + + assert isinstance(result, dict) + assert "execution_context" in result + + +@pytest.mark.asyncio +async def test_none_result_without_middleware_is_unchanged(): + dag = _builder(None).build(_spec(_node("none_node", node_id="none"))) + + assert await dag.nodes["none"]({"task_id": "t1"}) is None + + +# ── NodeContext sees the unfiltered state ───────────────────────────────── + + +@pytest.mark.asyncio +async def test_task_id_survives_a_context_policy_that_omits_it(): + dag = _builder(None).build(_spec(_node( + "task_id_node", node_id="tid", + context_policy=ContextPolicy(include_fields=["original_question"]), + ))) + + result = await dag.nodes["tid"]({ + "task_id": "task-42", "original_question": "q", "secret": "hidden", + }) + + assert result["seen_task_id"] == "task-42" + + +@pytest.mark.asyncio +async def test_context_policy_still_filters_the_state_the_node_receives(): + """Guard against the reorder leaking unfiltered state into the node.""" + seen: dict[str, Any] = {} + + async def spy(state, ctx): + seen["keys"] = sorted(state.keys()) + return {"current_phase": "spy"} + + builder = _builder(None) + builder._resolve_node_function_from_def = lambda node_def: spy # type: ignore[method-assign] + dag = builder.build(_spec(_node( + "task_id_node", node_id="spy", + context_policy=ContextPolicy(include_fields=["original_question"]), + ))) + + await dag.nodes["spy"]({ + "task_id": "task-42", "original_question": "q", "secret": "hidden", + }) + + assert seen["keys"] == ["original_question"] + + +@pytest.mark.asyncio +async def test_legacy_single_param_nodes_still_pass_through(): + async def legacy(state): + return {"current_phase": "legacy", "keys": sorted(state.keys())} + + builder = _builder(None) + builder._resolve_node_function_from_def = lambda node_def: legacy # type: ignore[method-assign] + dag = builder.build(_spec(_node( + "task_id_node", node_id="legacy", + context_policy=ContextPolicy(include_fields=["original_question"]), + ))) + + result = await dag.nodes["legacy"]({"task_id": "t", "original_question": "q"}) + + assert result["keys"] == ["original_question"] From e4fa1966363bfbef4daaff6ac5893104e3807adb Mon Sep 17 00:00:00 2001 From: zhanghanduo Date: Wed, 2 Sep 2026 13:47:24 +0800 Subject: [PATCH 4/6] fix(skills): reload the config file the loader actually came from MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `get_enabled_skills` auto-reloads when the extensions config changes on disk, and `reload()` does the same on demand. Both called `ExtensionsConfig.from_file()` with no argument, discarding the path the config was loaded from and restarting the cwd/env search. For a loader built with a config loaded from an explicit path outside those locations the search finds nothing and returns empty defaults. Since `is_skill_enabled` defaults to True for anything unlisted, an operator disabling one skill got every skill re-enabled — the decision inverted, not merely lost. The empty config also carries no `_file_path`, so `has_changed` returned False from then on and the state could never recover. `ExtensionsConfig` now exposes `source_path`, and the loader re-reads from it. A config that has become unreadable keeps the previously loaded state rather than failing open to "everything enabled", which is the wrong direction for a gate. Separately, `toggle_skill` was the only accessor that did not discover lazily: on a freshly built loader it read an empty `_skills` and returned False for a skill that exists on disk. The existing test called `discover()` first, which hid it. Co-Authored-By: Claude Opus 5 (1M context) --- .../components/skills/extensions_config.py | 10 ++ .../components/skills/file_system_loader.py | 35 ++++- tests/test_skills_loader_reload.py | 130 ++++++++++++++++++ 3 files changed, 172 insertions(+), 3 deletions(-) create mode 100644 tests/test_skills_loader_reload.py diff --git a/agent_core/components/skills/extensions_config.py b/agent_core/components/skills/extensions_config.py index c22df0c..a60dbdf 100644 --- a/agent_core/components/skills/extensions_config.py +++ b/agent_core/components/skills/extensions_config.py @@ -83,6 +83,16 @@ def from_file(cls, config_path: str | Path | None = None) -> ExtensionsConfig: logger.warning("Failed to load extensions config %s: %s", resolved, e) return cls() + @property + def source_path(self) -> Path | None: + """File this config was loaded from; ``None`` for in-memory defaults. + + A reloader must pass this back to :meth:`from_file` — calling it with + no argument restarts the cwd/env search, which silently returns empty + defaults for a config that was loaded from an explicit path. + """ + return self._file_path + def has_changed(self) -> bool: """Return True if the backing file has been modified since load.""" if self._file_path is None or not self._file_path.is_file(): diff --git a/agent_core/components/skills/file_system_loader.py b/agent_core/components/skills/file_system_loader.py index 9a4bd53..b6f381d 100644 --- a/agent_core/components/skills/file_system_loader.py +++ b/agent_core/components/skills/file_system_loader.py @@ -188,7 +188,7 @@ def get_enabled_skills(self) -> list[SkillConfig]: """ if self._extensions_config.has_changed(): logger.info("Extensions config changed on disk — reloading skill state") - self._extensions_config = ExtensionsConfig.from_file() + self._extensions_config = self._reload_extensions_config() # Re-apply enabled state to already-discovered skills for skill in self._skills.values(): skill.enabled = self._extensions_config.is_skill_enabled(skill.skill_id) @@ -197,7 +197,14 @@ def get_enabled_skills(self) -> list[SkillConfig]: # ── Mutations ───────────────────────────────────────────────────── def toggle_skill(self, skill_id: str, enabled: bool) -> bool: - """Enable or disable a skill. Returns True if skill exists.""" + """Enable or disable a skill. Returns True if skill exists. + + Discovers lazily like every other accessor: without this, toggling on + a freshly built loader read an empty ``_skills`` and reported False for + a skill that exists on disk — a silent no-op. + """ + if not self._discovered: + self.discover() skill = self._skills.get(skill_id) if skill is None: return False @@ -208,5 +215,27 @@ def reload(self) -> None: """Force re-discovery of skills.""" self._skills.clear() self._discovered = False - self._extensions_config = ExtensionsConfig.from_file() + self._extensions_config = self._reload_extensions_config() self.discover() + + def _reload_extensions_config(self) -> ExtensionsConfig: + """Re-read the config from the file it came from. + + Calling ``from_file()`` with no argument restarts the cwd/env search. + For a loader built with a config loaded from an explicit path outside + those locations, that search finds nothing and returns empty defaults — + and because ``is_skill_enabled`` defaults to True for anything unlisted, + an operator disabling a skill got every skill re-enabled instead. The + empty config also carries no ``_file_path``, so ``has_changed`` returned + False from then on and the state could never recover. + """ + source = self._extensions_config.source_path + reloaded = ExtensionsConfig.from_file(source) + if source is not None and reloaded.source_path is None: + logger.warning( + "Extensions config %s could not be re-read; keeping the " + "previously loaded skill state", + source, + ) + return self._extensions_config + return reloaded diff --git a/tests/test_skills_loader_reload.py b/tests/test_skills_loader_reload.py new file mode 100644 index 0000000..c932dca --- /dev/null +++ b/tests/test_skills_loader_reload.py @@ -0,0 +1,130 @@ +"""Reloading skill state must not invert an operator's decision. + +``get_enabled_skills`` auto-reloads when the extensions config changes on +disk. It called ``ExtensionsConfig.from_file()`` with no argument, which +restarts the cwd/env search instead of re-reading the file the config actually +came from. For a config loaded from an explicit path outside those locations +the search finds nothing and returns empty defaults — and because +``is_skill_enabled`` defaults to True for anything unlisted, "disable one +skill" was applied as "enable everything". The empty config also carries no +``_file_path``, so ``has_changed`` returned False from then on and the state +could never recover. +""" + +from __future__ import annotations + +import json +import os + +import pytest + +from agent_core.components.skills import ExtensionsConfig, FileSystemSkillLoader + + +@pytest.fixture() +def skill_dir(tmp_path): + skills = tmp_path / "skills" + skills.mkdir() + for name in ("code-review", "debug"): + d = skills / name + d.mkdir() + (d / "SKILL.md").write_text( + f"---\nname: {name}\ndescription: d\nversion: 1.0.0\n---\n# Body\n", + ) + return skills + + +@pytest.fixture() +def config_file(tmp_path): + """A config in a directory the cwd/env search will never look in.""" + path = tmp_path / "elsewhere" / "extensions.json" + path.parent.mkdir() + path.write_text(json.dumps({"skills": {}})) + return path + + +def _write(path, skills: dict[str, bool]) -> None: + path.write_text(json.dumps({ + "skills": {k: {"enabled": v} for k, v in skills.items()}, + })) + # Ensure the mtime moves even on a coarse-grained filesystem clock. + stat = path.stat() + os.utime(path, (stat.st_atime, stat.st_mtime + 10)) + + +def _loader(skill_dir, config_file) -> FileSystemSkillLoader: + return FileSystemSkillLoader( + skill_dirs=[skill_dir], + extensions_config=ExtensionsConfig.from_file(config_file), + ) + + +def test_source_path_is_exposed(): + """The reloader needs to know where the config came from.""" + assert ExtensionsConfig().source_path is None + + +def test_source_path_records_an_explicit_load(config_file): + config = ExtensionsConfig.from_file(config_file) + assert config.source_path == config_file + + +def test_disable_survives_the_auto_reload(skill_dir, config_file): + loader = _loader(skill_dir, config_file) + assert {s.skill_id for s in loader.get_enabled_skills()} == { + "code-review", "debug", + } + + _write(config_file, {"debug": False}) + + assert {s.skill_id for s in loader.get_enabled_skills()} == {"code-review"} + + +def test_reload_reads_the_same_file(skill_dir, config_file): + loader = _loader(skill_dir, config_file) + _write(config_file, {"debug": False}) + + loader.reload() + + assert {s.skill_id for s in loader.get_enabled_skills()} == {"code-review"} + + +def test_auto_reload_keeps_working_after_the_first_change(skill_dir, config_file): + """The old failure mode wedged ``has_changed`` off permanently.""" + loader = _loader(skill_dir, config_file) + _write(config_file, {"debug": False}) + assert {s.skill_id for s in loader.get_enabled_skills()} == {"code-review"} + + _write(config_file, {"code-review": False}) + + assert {s.skill_id for s in loader.get_enabled_skills()} == {"debug"} + + +def test_unreadable_config_keeps_the_previous_state(skill_dir, config_file): + """Failing open to "everything enabled" is the wrong direction.""" + loader = _loader(skill_dir, config_file) + _write(config_file, {"debug": False}) + assert {s.skill_id for s in loader.get_enabled_skills()} == {"code-review"} + + config_file.unlink() + loader.reload() + + assert {s.skill_id for s in loader.get_enabled_skills()} == {"code-review"} + + +# ── toggle_skill discovers lazily, like every other accessor ─────────────── + + +def test_toggle_on_a_fresh_loader_is_not_a_silent_no_op(skill_dir, config_file): + loader = _loader(skill_dir, config_file) + + assert loader.toggle_skill("debug", False) is True + + skill = loader.get_skill("debug") + assert skill is not None + assert skill.enabled is False + + +def test_toggle_still_reports_a_genuinely_missing_skill(skill_dir, config_file): + loader = _loader(skill_dir, config_file) + assert loader.toggle_skill("no-such-skill", False) is False From 58de97c7f6b31de27e96cb60d1566e99c5f047fa Mon Sep 17 00:00:00 2001 From: zhanghanduo Date: Wed, 2 Sep 2026 13:50:05 +0800 Subject: [PATCH 5/6] fix(observers): never replay a tool-less turn; honour disabled coalescing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `LeakedToolCallRetryObserver` returns `continue_to_next_turn=True` to replay a turn whose tool call leaked into prose. On the landing turn that is wrong: `LastTurnForcer` strips tools by planting `_llm_strip_tools`, which `_prepare_llm_request` consumes one-shot, so the replayed turn is re-bound *with* tools — the one turn the policy guaranteed would have none. The existing `blocked_tool_calls` guard cannot see this, because with no tools bound nothing was parsed and that list is empty, while a prose answer merely mentioning `finalize_answer(...)` still trips `_looks_like_leak`. `_prepare_llm_request` now records `_llm_tools_stripped` for the turn and the observer declines to replay. `TrajectoryFileObserver` read a coalescing bound of 0 as "never due" rather than "bound disabled", contradicting the module header. With `SWARM_TRAJECTORY_COALESCE_N=0` and the default `COALESCE_MS=0` the JSON snapshot was written once and then frozen until `on_loop_end` forced a final flush — so a process killed before that (OOM, hard timeout) left a file holding only the `start` state, the exact case the snapshot exists for. With every bound disabled there is nothing to coalesce, so every event flushes. Both regression tests verified to fail against the previous code. Co-Authored-By: Claude Opus 5 (1M context) --- .../observers/leaked_tool_call_retry.py | 12 ++ agent_core/components/observers/trajectory.py | 10 +- agent_core/runtime/loop/agent_loop.py | 6 + tests/test_leaked_tool_call_landing_turn.py | 118 ++++++++++++++++++ tests/test_trajectory_coalescing.py | 105 ++++++++++++++++ 5 files changed, 250 insertions(+), 1 deletion(-) create mode 100644 tests/test_leaked_tool_call_landing_turn.py create mode 100644 tests/test_trajectory_coalescing.py diff --git a/agent_core/components/observers/leaked_tool_call_retry.py b/agent_core/components/observers/leaked_tool_call_retry.py index 4cf26d5..3b7b425 100644 --- a/agent_core/components/observers/leaked_tool_call_retry.py +++ b/agent_core/components/observers/leaked_tool_call_retry.py @@ -151,6 +151,18 @@ async def on_llm_response(self, ctx: TurnContext) -> Intervention | None: ctx.metadata.pop(_METADATA_TEMP_KEY, None) return None + # This turn ran with no tools bound at all — ``LastTurnForcer`` strips + # them for the landing turn. Nothing could be parsed, so + # ``blocked_tool_calls`` is empty and the guard above does not cover + # it, yet a prose answer that merely mentions ``finalize_answer(...)`` + # trips ``_looks_like_leak``. Replaying would hand tools back on the + # one turn that was supposed to have none, because + # ``_llm_strip_tools`` is one-shot and already consumed. + if ctx.metadata.get("_llm_tools_stripped"): + ctx.metadata.pop(_METADATA_COUNT_KEY, None) + ctx.metadata.pop(_METADATA_TEMP_KEY, None) + return None + # A structured tool call was parsed — nothing to retry. Also reset # the escalation counter so a future leak starts at 0.3 again. if ctx.tool_calls: diff --git a/agent_core/components/observers/trajectory.py b/agent_core/components/observers/trajectory.py index 71d0e80..b11649b 100644 --- a/agent_core/components/observers/trajectory.py +++ b/agent_core/components/observers/trajectory.py @@ -346,8 +346,16 @@ def _flush_json(self, *, force: bool = False) -> None: if not force: self._pending_flushes += 1 now = time.monotonic() + # "0 disables the respective bound" (module header). With every + # bound disabled there is nothing left to coalesce, so every event + # flushes — treating a disabled bound as "never due" instead meant + # ``SWARM_TRAJECTORY_COALESCE_N=0`` froze the JSON snapshot after + # the first flush, and a process killed before ``on_loop_end`` + # left only the ``start`` state on disk. + coalescing = _COALESCE_N > 0 or _COALESCE_MS > 0 due = ( - self._last_flush_at == 0.0 + not coalescing + or self._last_flush_at == 0.0 or ( _COALESCE_N > 0 and self._pending_flushes >= _COALESCE_N diff --git a/agent_core/runtime/loop/agent_loop.py b/agent_core/runtime/loop/agent_loop.py index e4646f6..29befdf 100644 --- a/agent_core/runtime/loop/agent_loop.py +++ b/agent_core/runtime/loop/agent_loop.py @@ -542,6 +542,12 @@ async def _prepare_llm_request( ) -> tuple[Any, list[Message], bool, str]: temp_override = metadata.pop("_llm_temp_override", None) strip_tools = metadata.pop("_llm_strip_tools", False) + # ``_llm_strip_tools`` is one-shot and consumed right here, so an observer + # running later in the same turn cannot tell whether tools were bound. + # Record it for the turn: an observer that replays the turn + # (``continue_to_next_turn``) would otherwise re-bind tools on the landing + # turn that was deliberately given none. + metadata["_llm_tools_stripped"] = strip_tools llm_base = llm_with_session if strip_tools else llm_with_tools llm_for_turn = ( bind_temperature(llm_base, temp_override) diff --git a/tests/test_leaked_tool_call_landing_turn.py b/tests/test_leaked_tool_call_landing_turn.py new file mode 100644 index 0000000..c99eff9 --- /dev/null +++ b/tests/test_leaked_tool_call_landing_turn.py @@ -0,0 +1,118 @@ +"""A turn that ran without tools must never be replayed with tools. + +``LastTurnForcer`` strips tools for the landing turn by planting +``_llm_strip_tools``, which ``_prepare_llm_request`` consumes one-shot. On that +turn nothing can be parsed, so ``ctx.tool_calls`` and ``ctx.blocked_tool_calls`` +are both empty — while a prose answer that merely mentions ``finalize_answer`` +still trips ``_looks_like_leak``. The observer then asked for +``continue_to_next_turn``, and because the strip flag was already gone the +replayed turn was bound *with* tools: the one turn the policy guaranteed would +have none. +""" + +from __future__ import annotations + +import pytest + +from agent_core.components.observers.last_turn_forcer import LastTurnForcer +from agent_core.components.observers.leaked_tool_call_retry import ( + LeakedToolCallRetryObserver, +) +from agent_core.loop_types import TurnContext + +TOOLS = ("finalize_answer", "web_search") +LEAK_TEXT = "I will call finalize_answer with the complete answer now." + + +def _ctx(text: str, *, metadata: dict | None = None, turn: int = 3) -> TurnContext: + return TurnContext( + turn=turn, + max_turns=4, + task_id="t1", + role_id="solver", + ai_text=text, + thinking="", + tool_calls=[], + messages=[], + usage=None, + metadata=metadata if metadata is not None else {}, + ) + + +def _observer() -> LeakedToolCallRetryObserver: + return LeakedToolCallRetryObserver(TOOLS) + + +@pytest.mark.asyncio +async def test_a_genuine_leak_on_a_normal_turn_still_retries(): + """The observer's whole reason to exist must keep working.""" + ctx = _ctx(LEAK_TEXT, metadata={"_llm_tools_stripped": False}) + + intervention = await _observer().on_llm_response(ctx) + + assert intervention is not None + assert intervention.continue_to_next_turn is True + assert ctx.metadata["_llm_temp_override"] == 0.3 + + +@pytest.mark.asyncio +async def test_no_replay_when_the_turn_ran_without_tools(): + ctx = _ctx(LEAK_TEXT, metadata={"_llm_tools_stripped": True}) + + assert await _observer().on_llm_response(ctx) is None + assert "_llm_temp_override" not in ctx.metadata + assert "_leak_retry_count" not in ctx.metadata + + +@pytest.mark.asyncio +async def test_a_pending_retry_is_abandoned_on_the_stripped_turn(): + """State carried in from an earlier nudge must not resurrect the replay.""" + ctx = _ctx(LEAK_TEXT, metadata={ + "_llm_tools_stripped": True, + "_leak_retry_count": 1, + "_llm_temp_override": 0.6, + }) + + assert await _observer().on_llm_response(ctx) is None + assert "_llm_temp_override" not in ctx.metadata + assert "_leak_retry_count" not in ctx.metadata + + +@pytest.mark.asyncio +async def test_blocked_tool_calls_guard_is_unaffected(): + ctx = _ctx(LEAK_TEXT, metadata={"_llm_tools_stripped": False}) + ctx.blocked_tool_calls = [{"name": "web_search"}] + + assert await _observer().on_llm_response(ctx) is None + + +@pytest.mark.asyncio +async def test_end_to_end_with_last_turn_forcer(): + """The two observers must agree about the landing turn. + + LastTurnForcer plants the flag at the end of the penultimate turn; + ``_prepare_llm_request`` converts it into the per-turn marker; the retry + observer then declines to replay. + """ + from agent_core.loop_types import LoopConfig + from agent_core.runtime.loop.agent_loop import _prepare_llm_request + + metadata: dict = {} + forcer = LastTurnForcer(terminal_tool="finalize_answer") + penultimate = _ctx("...", metadata=metadata, turn=3) + assert await forcer.on_turn_end(penultimate) is not None + assert metadata["_llm_strip_tools"] is True + + with_tools = object() + without_tools = object() + llm_for_turn, _messages, _stripped, _ = await _prepare_llm_request( + LoopConfig(task_id="t1", role_id="solver", max_turns=4), + [], without_tools, with_tools, [], metadata, 4, + ) + + assert llm_for_turn is without_tools, "landing turn must run without tools" + assert metadata["_llm_tools_stripped"] is True + assert "_llm_strip_tools" not in metadata + + landing = _ctx(LEAK_TEXT, metadata=metadata, turn=4) + assert await _observer().on_llm_response(landing) is None diff --git a/tests/test_trajectory_coalescing.py b/tests/test_trajectory_coalescing.py new file mode 100644 index 0000000..5f2f31a --- /dev/null +++ b/tests/test_trajectory_coalescing.py @@ -0,0 +1,105 @@ +"""``0`` disables a coalescing bound; it must not mean "never flush". + +The module header states "0 disables the respective bound". The predicate read +a disabled bound as *never due* instead, so with ``COALESCE_N=0`` (and the +default ``COALESCE_MS=0``) the JSON snapshot was written once and then frozen +until ``on_loop_end`` forced a final flush. A process killed before that — an +OOM kill, a hard timeout — left a file containing only the ``start`` state, +which is exactly the situation the snapshot exists for. +""" + +from __future__ import annotations + +import importlib +import json + +import pytest + +from agent_core.loop_types import LoopConfig + + +def _observer(tmp_path, monkeypatch, *, coalesce_n: int, coalesce_ms: int = 0): + """Re-import the module so its env-derived constants are re-read. + + Flushes are counted by spying on ``_write_envelope``: the coalescing + predicate decides *whether* a rewrite happens, which is the behaviour under + test — the document's contents are a separate concern. + """ + monkeypatch.setenv("SWARM_TRAJECTORY_COALESCE_N", str(coalesce_n)) + monkeypatch.setenv("SWARM_TRAJECTORY_COALESCE_MS", str(coalesce_ms)) + import agent_core.components.observers.trajectory as traj + + traj = importlib.reload(traj) + obs = traj.TrajectoryFileObserver(tmp_path, filename="t", formats=["json"]) + writes: list[int] = [] + original = obs._write_envelope + def _spy(path): + writes.append(1) + return original(path) + monkeypatch.setattr(obs, "_write_envelope", _spy) + return obs, writes + + +async def _flush_n_times(obs, n: int) -> None: + """Drive ``_flush_json`` exactly ``n`` times through the non-forced path.""" + for _ in range(n): + obs._flush_json() + + +@pytest.mark.asyncio +async def test_disabled_bounds_flush_every_event(tmp_path, monkeypatch): + obs, writes = _observer(tmp_path, monkeypatch, coalesce_n=0, coalesce_ms=0) + + await _flush_n_times(obs, 5) + + assert len(writes) == 5 + + +@pytest.mark.asyncio +async def test_count_bound_coalesces_then_flushes(tmp_path, monkeypatch): + obs, writes = _observer(tmp_path, monkeypatch, coalesce_n=3, coalesce_ms=0) + + await _flush_n_times(obs, 1) + assert len(writes) == 1, "the first flush always lands" + + await _flush_n_times(obs, 2) + assert len(writes) == 1, "held back until the batch fills" + + await _flush_n_times(obs, 1) + assert len(writes) == 2, "batch of 3 reached" + + +@pytest.mark.asyncio +async def test_time_bound_alone_still_coalesces(tmp_path, monkeypatch): + """N disabled, MS set — the remaining bound must still apply.""" + obs, writes = _observer(tmp_path, monkeypatch, coalesce_n=0, coalesce_ms=60_000) + + await _flush_n_times(obs, 5) + + assert len(writes) == 1 + + +@pytest.mark.asyncio +async def test_force_always_writes(tmp_path, monkeypatch): + """``on_loop_end`` relies on this regardless of the bounds.""" + obs, writes = _observer(tmp_path, monkeypatch, coalesce_n=1000, coalesce_ms=0) + + obs._flush_json() + obs._flush_json() + obs._flush_json(force=True) + + assert len(writes) == 2 # first (always) + forced + + +@pytest.mark.asyncio +async def test_snapshot_is_readable_after_every_event(tmp_path, monkeypatch): + """The point of flushing often: a kill -9 must leave a usable file.""" + obs, _writes = _observer(tmp_path, monkeypatch, coalesce_n=0, coalesce_ms=0) + await obs.on_loop_start( + LoopConfig(task_id="t", role_id="r", max_turns=5), + ) + + obs._flush_json() + + doc = json.loads((tmp_path / "t.json").read_text()) + assert isinstance(doc, dict) From 78fff87b2af02cb38a5f7ec18ecac34c7d3d46e6 Mon Sep 17 00:00:00 2001 From: zhanghanduo Date: Wed, 2 Sep 2026 14:10:46 +0800 Subject: [PATCH 6/6] fix(skills): stabilize relative config reload paths --- .../components/skills/extensions_config.py | 11 ++++++++- tests/test_skills_loader_reload.py | 24 +++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/agent_core/components/skills/extensions_config.py b/agent_core/components/skills/extensions_config.py index a60dbdf..49abe6c 100644 --- a/agent_core/components/skills/extensions_config.py +++ b/agent_core/components/skills/extensions_config.py @@ -63,7 +63,16 @@ class ExtensionsConfig(BaseModel): @classmethod def from_file(cls, config_path: str | Path | None = None) -> ExtensionsConfig: """Load config from JSON file with environment variable resolution.""" - resolved = Path(config_path) if config_path else _find_config_file() + if config_path: + resolved = Path(config_path) + # Keep the source stable across later ``chdir`` calls. Do not use + # ``resolve()`` here: retaining a configured symlink path lets an + # operator atomically repoint that symlink and have reload observe + # the new target. + if not resolved.is_absolute(): + resolved = Path.cwd() / resolved + else: + resolved = _find_config_file() if resolved is None or not resolved.is_file(): logger.debug("No extensions config found — using empty defaults") diff --git a/tests/test_skills_loader_reload.py b/tests/test_skills_loader_reload.py index c932dca..38b799d 100644 --- a/tests/test_skills_loader_reload.py +++ b/tests/test_skills_loader_reload.py @@ -89,6 +89,30 @@ def test_reload_reads_the_same_file(skill_dir, config_file): assert {s.skill_id for s in loader.get_enabled_skills()} == {"code-review"} +def test_relative_source_path_survives_cwd_change( + skill_dir, config_file, monkeypatch, +): + """Reload must keep tracking the originally opened relative path.""" + monkeypatch.chdir(config_file.parent) + config = ExtensionsConfig.from_file(config_file.name) + assert config.source_path == config_file + loader = FileSystemSkillLoader( + skill_dirs=[skill_dir], + extensions_config=config, + ) + + other = config_file.parent / "other" + other.mkdir() + monkeypatch.chdir(other) + _write(config_file, {"debug": False}) + + assert {s.skill_id for s in loader.get_enabled_skills()} == {"code-review"} + + _write(config_file, {"code-review": False}) + loader.reload() + assert {s.skill_id for s in loader.get_enabled_skills()} == {"debug"} + + def test_auto_reload_keeps_working_after_the_first_change(skill_dir, config_file): """The old failure mode wedged ``has_changed`` off permanently.""" loader = _loader(skill_dir, config_file)