Compare commits

...
Author SHA1 Message Date
MervinPraisonandClaude Fable 5.1 59f69e0ee8 fix(agent): run context compaction on the streaming path
Agent auto-enables a ContextManager whenever tools are present
(auto_compact, threshold 0.8), and the non-streaming chat() path honours
it by calling _apply_context_management between building the messages
and issuing the request. start(stream=True) / iter_stream built the same
messages and sent them unmodified, so a long session streamed its entire
history to the model on every turn while the configured manager sat
idle. The desktop's default agent has tools, so every streamed turn hit
this once the conversation grew.

Both streaming branches now call the same _apply_context_management the
non-streaming path uses: the OpenAI branch on the built messages list
(which the post-tool follow-up request also reuses), and the custom-LLM
branch on the chat_history that is actually handed to
get_response_stream. No new compaction machinery; context=False remains
zero-overhead because the helper returns immediately when no manager is
configured.

The test stubs the transport, seeds 400 long messages (~160k estimated
tokens against a 128k window) and asserts the message list that reaches
the transport is bounded on both branches, plus a context=False control
that pins the full history going out untouched. Removing either call
makes the corresponding test fail.

Closes #4714

Co-Authored-By: Claude Fable 5.1 <[email protected]>
2026-09-02 23:17:19 +01:00
2 changed files with 165 additions and 5 deletions

No files matched your search

@@ -4917,7 +4917,21 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
try:
# Use the new streaming generator from LLM class
response_content = ""
stream_history = self.chat_history
system_prompt_for_stream = self._build_system_prompt(
tool_param,
memory_prefetch_context=memory_prefetch_context,
)
# Apply context management (auto-compaction) to the history
# that is actually sent, the same way the non-streaming
# custom-LLM path does. Without this a configured
# ContextManager -- auto-enabled whenever tools are present
# -- is silently skipped on stream=True (Issue #4714).
# Zero overhead when context=False.
stream_history, _stream_context_result = self._apply_context_management(
messages=self.chat_history,
system_prompt=system_prompt_for_stream,
tools=tool_param,
)
durable_context = self._get_durable_run_context()
if durable_context is not None:
stream_history = durable_context.restore_messages(
@@ -4925,10 +4939,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
)
for chunk in self.llm_instance.get_response_stream(
prompt=actual_prompt,
system_prompt=self._build_system_prompt(
tool_param,
memory_prefetch_context=memory_prefetch_context,
),
system_prompt=system_prompt_for_stream,
chat_history=stream_history,
temperature=kwargs.get('temperature', 1.0),
tools=tool_param,
@@ -5006,6 +5017,18 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
self.chat_history[-1].get("content") == normalized_content):
self._append_to_chat_history({"role": "user", "content": normalized_content})
# Apply context management before the LLM call (auto-compaction),
# exactly as the non-streaming OpenAI path does. Without this a
# configured ContextManager -- auto-enabled whenever tools are
# present -- is silently skipped on stream=True and the full
# history is sent (Issue #4714). Zero overhead when context=False.
system_prompt_content = messages[0].get("content", "") if messages and messages[0].get("role") == "system" else ""
messages, _stream_context_result = self._apply_context_management(
messages=messages,
system_prompt=system_prompt_content,
tools=tool_param,
)
try:
# Check if OpenAI client is available
if self._openai_client is None:
@@ -0,0 +1,137 @@
"""The streaming entry point must run context compaction (Issue #4714).
Agent auto-enables a ContextManager whenever tools are present (auto_compact,
compact_threshold 0.8). The non-streaming ``chat()`` path runs
``_apply_context_management`` between building the messages and issuing the
request; ``start(stream=True)`` built the same messages and sent them
unmodified, so a long session streamed the entire history to the model every
turn while the manager sat idle. The desktop's default agent has tools, so it
hit this on every streamed turn.
These tests stub the transport and assert the EFFECT: the message list that
actually reaches the transport is bounded when a ContextManager is configured,
and untouched when it is not.
"""
import pytest
from praisonaiagents import Agent
SEEDED = 400 # ~160k estimated tokens; well past 80% of gpt-4o-mini's 128k window
FILLER = "lorem ipsum dolor sit amet " * 60
def _add(a: int, b: int) -> int:
"""Add two numbers."""
return a + b
class _RecordingOpenAI:
"""A proxy over the real client wrapper: everything (build_messages, tool
formatting, ...) is delegated; only the transport call is intercepted and
recorded. ``iter([])`` is an empty stream -- no answer, no tool calls."""
def __init__(self, real):
self._real = real
self.calls = []
outer = self
class _Completions:
def create(self, **kw):
outer.calls.append(kw)
return iter([])
class _Chat:
completions = _Completions()
class _Sync:
chat = _Chat()
self.sync_client = _Sync()
def __getattr__(self, name):
return getattr(self._real, name)
def _seed_history(agent, n=SEEDED):
for i in range(n):
agent.chat_history.append({
"role": "user" if i % 2 == 0 else "assistant",
"content": f"msg {i} {FILLER}",
})
def _stub_openai(agent):
from praisonaiagents.llm import get_openai_client
stub = _RecordingOpenAI(get_openai_client(api_key="sk-test"))
agent._Agent__openai_client = stub # the lazily-built client, pre-seeded
return stub
@pytest.fixture(autouse=True)
def _api_key(monkeypatch):
monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
def test_openai_stream_compacts_long_history():
agent = Agent(name="t", instructions="x", llm="gpt-4o-mini", tools=[_add])
assert agent.context_manager is not None, "tools present must auto-enable a ContextManager"
_seed_history(agent)
stub = _stub_openai(agent)
"".join(agent.start("x", stream=True))
assert stub.calls, "the stub transport was never called"
sent = stub.calls[0]["messages"]
# Far fewer than seeded: compaction at threshold 0.8 keeps the window
# under budget, not merely one message under the seed count.
assert len(sent) <= SEEDED * 0.7, (
f"streamed request carried {len(sent)} messages for {SEEDED} seeded; "
"context compaction did not run on the streaming path"
)
# Compaction must not lose the request itself.
assert sent[0]["role"] == "system"
assert sent[-1] == {"role": "user", "content": "x"}
def test_openai_stream_without_context_manager_sends_full_history():
"""Zero-overhead contract: context=False leaves the streamed request alone.
Also pins that the bounded count above is compaction and not some other
trimming in the streaming path."""
agent = Agent(name="t", instructions="x", llm="gpt-4o-mini", tools=[_add], context=False)
assert agent.context_manager is None
_seed_history(agent)
stub = _stub_openai(agent)
"".join(agent.start("x", stream=True))
sent = stub.calls[0]["messages"]
assert len(sent) == SEEDED + 2 # system + seeded history + new user turn
def test_custom_llm_stream_compacts_long_history():
"""The LiteLLM branch passes chat_history rather than a built messages list;
compaction must apply to the history that is actually sent."""
agent = Agent(name="t", instructions="x", llm="openai/gpt-4o-mini", tools=[_add])
assert agent.context_manager is not None
_seed_history(agent)
seen = {}
def fake_stream(**kw):
seen.update(kw)
yield "ok"
agent.llm_instance.get_response_stream = fake_stream
"".join(agent.start("x", stream=True))
assert "chat_history" in seen, "the fake transport was never called"
history = seen["chat_history"]
assert len(history) <= SEEDED * 0.7, (
f"streamed custom-LLM request carried {len(history)} history messages for "
f"{SEEDED} seeded; context compaction did not run on the streaming path"
)
# The turn just appended must survive compaction.
assert history[-1] == {"role": "user", "content": "x"}