diff --git a/python/packages/core/agent_framework/_harness/_agent.py b/python/packages/core/agent_framework/_harness/_agent.py index d4dd8d0ae1..603d822b4d 100644 --- a/python/packages/core/agent_framework/_harness/_agent.py +++ b/python/packages/core/agent_framework/_harness/_agent.py @@ -18,7 +18,7 @@ from .._agents import Agent, SupportsAgentRun from .._clients import SupportsShellTool, SupportsWebSearchTool -from .._compaction import CompactionProvider, ContextWindowCompactionStrategy, ToolResultCompactionStrategy +from .._compaction import CompactionProvider, ContextWindowCompactionStrategy from .._feature_stage import ExperimentalFeature, experimental from .._sessions import ContextProvider, HistoryProvider, InMemoryHistoryProvider, MessageInjectionMiddleware from .._skills import SkillsProvider @@ -78,7 +78,7 @@ def _assemble_instructions( return f"{harness}\n\n{agent_instructions or ''}".strip() or None -def _assemble_compaction_provider( +def _assemble_compaction( *, disable_compaction: bool, max_context_window_tokens: int | None, @@ -87,46 +87,58 @@ def _assemble_compaction_provider( before_compaction_strategy: CompactionStrategy | None, after_compaction_strategy: CompactionStrategy | None, tokenizer: TokenizerProtocol | None, -) -> CompactionProvider | None: - """Build the compaction provider from parameters or defaults. - - The token-budget defaults (``ContextWindowCompactionStrategy`` for the before phase and - ``ToolResultCompactionStrategy`` for the after phase) are only applied when the token - params are provided. Caller-supplied strategies are always honored. Either phase may end - up ``None``, which ``CompactionProvider`` interprets as "skip that phase". - - Returns None when compaction is explicitly disabled, or when neither phase has a strategy - (no custom strategies and no token budget to build the defaults). +) -> tuple[CompactionStrategy | None, CompactionProvider | None]: + """Resolve the harness compaction strategies into their execution sites. + + The token-budget default (``ContextWindowCompactionStrategy``) is used for **both** the before + and after phases, and is only applied when the token params are provided. A single shared + instance is reused across both phases (it is stateless). Caller-supplied strategies always win + per phase. + + Because the harness enables per-service-call history persistence, running the before-phase + compaction as a ``CompactionProvider.before_run`` hook is a no-op: the agent skips + ``HistoryProvider.before_run``, so the provider only ever sees an empty context. + Instead the before-strategy is returned separately so the caller can wire it as the agent's + ``compaction_strategy`` chat option. That runs it inside ``BaseChatClient.get_response`` — + per model call, inner of ``PerServiceCallHistoryPersistingMiddleware`` (which has already + loaded the full history into the outgoing messages) and outer of the leaf client. + + The after-strategy stays on a ``CompactionProvider`` (its ``after_run`` compacts the persisted + session history in place and works correctly). + + Returns a ``(before_strategy, after_provider)`` tuple. Both elements are ``None`` when + compaction is disabled or when the corresponding phase has no strategy (no custom strategy + and no token budget to build the default). """ if disable_compaction: - return None + return None, None - # Resolve the before-strategy: custom strategy wins; otherwise fall back to the - # token-budget-aware default when token params are available. - before_strategy = before_compaction_strategy - if before_strategy is None and max_context_window_tokens is not None and max_output_tokens is not None: - before_strategy = ContextWindowCompactionStrategy( + # Token-budget default, shared across both phases when token params are available. + default_strategy: CompactionStrategy | None = None + if max_context_window_tokens is not None and max_output_tokens is not None: + default_strategy = ContextWindowCompactionStrategy( max_context_window_tokens=max_context_window_tokens, max_output_tokens=max_output_tokens, tokenizer=tokenizer, ) - # Resolve the after-strategy: custom strategy wins; otherwise fall back to the default - # when token params are available. - after_strategy = after_compaction_strategy - if after_strategy is None and max_context_window_tokens is not None and max_output_tokens is not None: - after_strategy = ToolResultCompactionStrategy(keep_last_tool_call_groups=2) + # Resolve each phase: caller-supplied strategy wins; otherwise fall back to the shared default. + before_strategy = before_compaction_strategy if before_compaction_strategy is not None else default_strategy + after_strategy = after_compaction_strategy if after_compaction_strategy is not None else default_strategy - # Nothing to compact in either phase: skip the provider entirely. - if before_strategy is None and after_strategy is None: - return None - - return CompactionProvider( - before_strategy=before_strategy, - after_strategy=after_strategy, - tokenizer=tokenizer, - history_source_id=history_source_id, + # The after-strategy runs post-turn against the persisted history via a CompactionProvider; + # the before-strategy runs per model call via the agent's compaction_strategy option. + after_provider = ( + CompactionProvider( + before_strategy=None, + after_strategy=after_strategy, + tokenizer=tokenizer, + history_source_id=history_source_id, + ) + if after_strategy is not None + else None ) + return before_strategy, after_provider def _assemble_context_providers( @@ -383,8 +395,9 @@ def create_harness_agent( compaction runs even if token params are omitted. Defaults to ContextWindowCompactionStrategy (token-budget aware) when token params are provided. after_compaction_strategy: Custom after-run compaction strategy. When provided, - compaction runs even if token params are omitted. Defaults to - ToolResultCompactionStrategy when token params are provided. + compaction runs even if token params are omitted. Defaults to the same + ContextWindowCompactionStrategy used for the before phase when token params are + provided. tokenizer: Custom tokenizer for compaction strategies. disable_todo: When True, skip the TodoProvider. todo_provider: Custom TodoProvider instance. Ignored when disable_todo is True. @@ -496,8 +509,10 @@ def create_harness_agent( # Build history provider. resolved_history = history_provider or InMemoryHistoryProvider() - # Build compaction provider. - compaction_provider = _assemble_compaction_provider( + # Build compaction. The before-strategy is wired as the agent's compaction_strategy option + # (runs per model call, inner of per-service-call persistence); the after-strategy stays on a + # CompactionProvider that compacts the persisted history post-turn. See issue #7011. + before_compaction, compaction_provider = _assemble_compaction( disable_compaction=disable_compaction, max_context_window_tokens=max_context_window_tokens, max_output_tokens=max_output_tokens, @@ -600,6 +615,8 @@ def create_harness_agent( default_options=default_opts, # type: ignore[arg-type] context_providers=assembled_providers, middleware=assembled_middleware or None, + compaction_strategy=before_compaction, + tokenizer=tokenizer, require_per_service_call_history_persistence=True, ) diff --git a/python/packages/core/tests/core/test_harness_agent.py b/python/packages/core/tests/core/test_harness_agent.py index 9cb6d94a37..19cb5a006b 100644 --- a/python/packages/core/tests/core/test_harness_agent.py +++ b/python/packages/core/tests/core/test_harness_agent.py @@ -6,6 +6,7 @@ from collections.abc import AsyncIterable, Awaitable, Mapping, Sequence from pathlib import Path from typing import Any +from unittest.mock import patch import pytest from agent_framework_tools.shell import ShellResult @@ -13,17 +14,20 @@ from agent_framework import ( AgentSession, BaseChatClient, + CharacterEstimatorTokenizer, ChatOptions, ChatResponse, ChatResponseUpdate, CompactionProvider, Content, + ContextWindowCompactionStrategy, FileAccessProvider, FileMemoryProvider, FileSystemAgentFileStore, InMemoryAgentFileStore, InMemoryHistoryProvider, Message, + MessageInjectionMiddleware, ResponseStream, ServiceSessionId, SkillsProvider, @@ -32,7 +36,8 @@ ) from agent_framework._harness._agent import DEFAULT_HARNESS_INSTRUCTIONS, _assemble_instructions from agent_framework._harness._mode import AgentModeProvider -from agent_framework._sessions import ContextProvider +from agent_framework._sessions import ContextProvider, PerServiceCallHistoryPersistingMiddleware +from agent_framework._tools import FunctionInvocationLayer class _FakeChatClient(BaseChatClient[ChatOptions[Any]]): @@ -295,16 +300,19 @@ def test_create_harness_agent_no_token_params_skips_max_tokens_option() -> None: assert agent.default_options.get("max_tokens") is None -def test_create_harness_agent_custom_before_strategy_enables_compaction_without_tokens() -> None: - """A custom before_compaction_strategy enables compaction even when token params are omitted.""" +def test_create_harness_agent_custom_before_strategy_wires_compaction_strategy_without_tokens() -> None: + """A custom before_compaction_strategy is wired as the agent compaction_strategy, even without tokens.""" from agent_framework import ToolResultCompactionStrategy + before_strategy = ToolResultCompactionStrategy() agent = create_harness_agent( client=_FakeChatClient(), - before_compaction_strategy=ToolResultCompactionStrategy(), + before_compaction_strategy=before_strategy, ) - provider_types = [type(p) for p in agent.context_providers] - assert CompactionProvider in provider_types + # The before-strategy runs per model call via the agent compaction_strategy option, not a provider. + assert agent.compaction_strategy is before_strategy + # No after-strategy and no token budget, so no CompactionProvider is added. + assert CompactionProvider not in [type(p) for p in agent.context_providers] def test_create_harness_agent_disable_compaction_overrides_custom_before_strategy() -> None: @@ -318,6 +326,7 @@ def test_create_harness_agent_disable_compaction_overrides_custom_before_strateg ) provider_types = [type(p) for p in agent.context_providers] assert CompactionProvider not in provider_types + assert agent.compaction_strategy is None def test_create_harness_agent_custom_after_strategy_enables_compaction_without_tokens() -> None: @@ -333,6 +342,207 @@ def test_create_harness_agent_custom_after_strategy_enables_compaction_without_t # Before phase is skipped (no token budget, no custom before strategy), after phase is set. assert compaction_providers[0].before_strategy is None assert compaction_providers[0].after_strategy is not None + # An after-only strategy must not wire anything as the agent-level (per-call) compaction. + assert agent.compaction_strategy is None + + +def test_create_harness_agent_default_tokens_split_compaction_phases() -> None: + """With token params, the before-strategy is the agent compaction_strategy; after stays on the provider.""" + agent = create_harness_agent( + client=_FakeChatClient(), + max_context_window_tokens=128_000, + max_output_tokens=16_384, + ) + # Before-strategy runs per model call via the agent option (issue #7011). + assert isinstance(agent.compaction_strategy, ContextWindowCompactionStrategy) + # After-strategy runs post-turn via the provider, which no longer carries a before-strategy. + compaction_providers = [p for p in agent.context_providers if isinstance(p, CompactionProvider)] + assert len(compaction_providers) == 1 + assert compaction_providers[0].before_strategy is None + # Both phases default to the same shared ContextWindowCompactionStrategy instance. + assert isinstance(compaction_providers[0].after_strategy, ContextWindowCompactionStrategy) + assert compaction_providers[0].after_strategy is agent.compaction_strategy + + +async def test_before_strategy_compaction_fires_on_loaded_history_under_per_service_call_persistence( + chat_client_base: Any, +) -> None: + """Regression for #7011: before-strategy compaction runs per model call on the loaded history. + + Because the harness enables per-service-call persistence, wiring the before-strategy as a + ``CompactionProvider.before_run`` hook was a no-op (empty context). Wiring it as the agent + ``compaction_strategy`` option makes it run inside ``BaseChatClient.get_response`` on the full + history that ``PerServiceCallHistoryPersistingMiddleware`` loads into the outgoing messages. + """ + captured: dict[str, list[Message]] = {} + + async def fake_get_response(*, messages: Sequence[Message], options: dict[str, Any], **kwargs: Any) -> ChatResponse: + captured["messages"] = list(messages) + return ChatResponse(messages=Message(role="assistant", contents=["ok"])) + + def build_agent(*, disable_compaction: bool) -> Any: + return create_harness_agent( + client=chat_client_base, + max_context_window_tokens=2_000, + max_output_tokens=500, + tokenizer=CharacterEstimatorTokenizer(), + disable_compaction=disable_compaction, + disable_todo=True, + disable_mode=True, + disable_file_memory=True, + disable_file_access=True, + ) + + def seeded_session(agent: Any) -> tuple[Any, int]: + session = agent.create_session() + long_history: list[Message] = [] + for i in range(20): + long_history.append(Message(role="user", contents=[f"u{i} " * 50])) + long_history.append(Message(role="assistant", contents=[f"a{i} " * 50])) + session.state[InMemoryHistoryProvider.DEFAULT_SOURCE_ID] = {"messages": long_history} + return session, len(long_history) + + with patch.object(chat_client_base, "_get_non_streaming_response", side_effect=fake_get_response): + # Control: compaction disabled -> the full loaded history reaches the leaf client. + control_agent = build_agent(disable_compaction=True) + control_session, history_len = seeded_session(control_agent) + assert control_agent.compaction_strategy is None + await control_agent.run("final question", session=control_session) + baseline_count = len(captured["messages"]) + assert baseline_count == history_len + 1 # loaded history + the input message + + # Compaction enabled -> the before-strategy truncates the loaded history per model call. + compacting_agent = build_agent(disable_compaction=False) + compacting_session, _ = seeded_session(compacting_agent) + assert isinstance(compacting_agent.compaction_strategy, ContextWindowCompactionStrategy) + await compacting_agent.run("final question", session=compacting_session) + compacted_count = len(captured["messages"]) + + # Truncation dropped older messages from the model input, proving the before-strategy fired + # on the history loaded by per-service-call persistence. + assert compacted_count < baseline_count + + +async def test_before_strategy_compaction_multi_turn_keeps_persisted_history_coherent( + chat_client_base: Any, +) -> None: + """The per-call before-strategy must not corrupt or lose persisted history across turns. + + The before-strategy runs on a per-call shallow copy that shares ``Message`` objects with the + stored history, so its ``_excluded`` annotations leak onto those shared objects. With the + default ``skip_excluded=False`` the provider reloads every message each turn, so stored history + must keep growing by exactly the persisted input/output of each turn (no loss, no duplication, + no leaked per-call summaries). + """ + per_call_counts: list[int] = [] + + async def fake_get_response(*, messages: Sequence[Message], options: dict[str, Any], **kwargs: Any) -> ChatResponse: + per_call_counts.append(len(messages)) + return ChatResponse(messages=Message(role="assistant", contents=["ok"])) + + agent = create_harness_agent( + client=chat_client_base, + max_context_window_tokens=2_000, + max_output_tokens=500, + tokenizer=CharacterEstimatorTokenizer(), + disable_todo=True, + disable_mode=True, + disable_file_memory=True, + disable_file_access=True, + ) + + session = agent.create_session() + long_history: list[Message] = [] + for i in range(15): + long_history.append(Message(role="user", contents=[f"u{i} " * 50])) + long_history.append(Message(role="assistant", contents=[f"a{i} " * 50])) + history_len = len(long_history) + session.state[InMemoryHistoryProvider.DEFAULT_SOURCE_ID] = {"messages": long_history} + + with patch.object(chat_client_base, "_get_non_streaming_response", side_effect=fake_get_response): + await agent.run("first question", session=session) + await agent.run("second question", session=session) + + # Compaction fired on both turns: each model call saw far fewer messages than were persisted. + assert len(per_call_counts) == 2 + assert all(0 < count < history_len for count in per_call_counts) + + # Persisted history stays coherent: a flat list of Message objects that grew by exactly the + # input + output of each turn (2 turns * 2 messages), with no leaked per-call summaries. + stored = session.state[InMemoryHistoryProvider.DEFAULT_SOURCE_ID]["messages"] + assert all(isinstance(m, Message) for m in stored) + assert len(stored) == history_len + 4 + + +async def test_harness_chat_client_middleware_execution_order( + chat_client_base: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The harness chat-client pipeline must execute in a specific outer-to-inner order. + + Injected messages must be persisted (injection outer of per-service-call persistence) and + compaction must run on the full persisted history (inner of persistence), all inside the + function-invocation loop so injection can happen within the tool-calling loop. See issue #7011. + + Expected order (outermost first): + 1. Function Invocation Loop + 2. MessageInjectionMiddleware + 3. PerServiceCallHistoryPersistingMiddleware + 4. Compaction (agent ``compaction_strategy`` option, per model call) + 5. Leaf ChatClient + """ + order: list[str] = [] + + original_function_loop = FunctionInvocationLayer.get_response + + def recording_function_loop(self: Any, *args: Any, **kwargs: Any) -> Any: + # Fires once at the top of the run; the internal ``super().get_response`` is a different + # bound method, so this does not re-fire per model iteration. + order.append("function_loop") + return original_function_loop(self, *args, **kwargs) + + original_injection = MessageInjectionMiddleware.process + + async def recording_injection(self: Any, context: Any, call_next: Any) -> None: + order.append("message_injection") + await original_injection(self, context, call_next) + + original_persistence = PerServiceCallHistoryPersistingMiddleware.process + + async def recording_persistence(self: Any, context: Any, call_next: Any) -> None: + order.append("per_service_call") + await original_persistence(self, context, call_next) + + monkeypatch.setattr(FunctionInvocationLayer, "get_response", recording_function_loop) + monkeypatch.setattr(MessageInjectionMiddleware, "process", recording_injection) + monkeypatch.setattr(PerServiceCallHistoryPersistingMiddleware, "process", recording_persistence) + + class _RecordingCompactionStrategy: + async def __call__(self, messages: list[Message]) -> bool: + order.append("compaction") + return False + + async def recording_leaf(*, messages: Sequence[Message], options: dict[str, Any], **kwargs: Any) -> ChatResponse: + order.append("leaf") + return ChatResponse(messages=Message(role="assistant", contents=["ok"])) + + agent = create_harness_agent( + client=chat_client_base, + max_context_window_tokens=2_000, + max_output_tokens=500, + tokenizer=CharacterEstimatorTokenizer(), + before_compaction_strategy=_RecordingCompactionStrategy(), + disable_todo=True, + disable_mode=True, + disable_file_memory=True, + disable_file_access=True, + ) + + session = agent.create_session() + with patch.object(chat_client_base, "_get_non_streaming_response", side_effect=recording_leaf): + await agent.run("hello", session=session) + + assert order == ["function_loop", "message_injection", "per_service_call", "compaction", "leaf"] # --- Validation Tests ---