core/agent: soft context compact and cooperative cancel points

Shrink older tool dumps for model requests without mutating history;
check task cancellation between tool/model steps.

src/plyngent/agent/budget.py

@@ -1,6 +1,24 @@
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
+import msgspec
+from msgspec import UNSET
+
+from plyngent.lmproto.openai_compatible.model import (
+    AssistantChatMessage,
+    ToolChatMessage,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from plyngent.lmproto.openai_compatible.model import AnyChatMessage
+
 DEFAULT_TOOL_RESULT_MAX_CHARS = 32_000
+DEFAULT_CONTEXT_MAX_CHARS = 200_000
+DEFAULT_OLD_TOOL_RESULT_CHARS = 800
+DEFAULT_RECENT_TOOL_RESULTS = 4
 
 
 def truncate_tool_result(text: str, max_chars: int = DEFAULT_TOOL_RESULT_MAX_CHARS) -> str:
@@ -11,3 +29,117 @@ def truncate_tool_result(text: str, max_chars: int = DEFAULT_TOOL_RESULT_MAX_CHA
         return text
     omitted = len(text) - max_chars
     return f"{text[:max_chars]}\n...[truncated {omitted} characters]"
+
+
+def estimate_message_chars(message: AnyChatMessage) -> int:
+    """Rough character estimate for soft context budgeting (not tokenizer-accurate)."""
+    total = 0
+    content = getattr(message, "content", None)
+    if isinstance(content, str):
+        total += len(content)
+    if isinstance(message, AssistantChatMessage):
+        tool_calls = message.tool_calls
+        if tool_calls is not UNSET and tool_calls:
+            for call in tool_calls:
+                total += len(getattr(call, "id", "") or "")
+                function = getattr(call, "function", None)
+                if function is not None:
+                    total += len(getattr(function, "name", "") or "")
+                    total += len(getattr(function, "arguments", "") or "")
+                custom = getattr(call, "custom", None)
+                if custom is not None:
+                    total += len(getattr(custom, "name", "") or "")
+                    total += len(getattr(custom, "input", "") or "")
+    if isinstance(message, ToolChatMessage):
+        total += len(message.tool_call_id)
+    return total
+
+
+def estimate_messages_chars(messages: Sequence[AnyChatMessage]) -> int:
+    return sum(estimate_message_chars(m) for m in messages)
+
+
+def _shrink_tool(message: ToolChatMessage, max_chars: int) -> ToolChatMessage:
+    if len(message.content) <= max_chars:
+        return message
+    return msgspec.structs.replace(
+        message,
+        content=truncate_tool_result(message.content, max_chars),
+    )
+
+
+def _tool_indices(messages: Sequence[AnyChatMessage]) -> list[int]:
+    return [i for i, m in enumerate(messages) if isinstance(m, ToolChatMessage)]
+
+
+def _protect_indices(tool_indices: Sequence[int], keep_recent: int) -> set[int]:
+    if keep_recent <= 0:
+        return set()
+    return set(tool_indices[-keep_recent:])
+
+
+def _shrink_except(
+    messages: list[AnyChatMessage],
+    tool_indices: Sequence[int],
+    protect: set[int],
+    max_chars: int,
+) -> None:
+    for idx in tool_indices:
+        if idx in protect:
+            continue
+        tool_msg = messages[idx]
+        if isinstance(tool_msg, ToolChatMessage):
+            messages[idx] = _shrink_tool(tool_msg, max_chars)
+
+
+def _shrink_largest(
+    messages: list[AnyChatMessage],
+    tool_indices: Sequence[int],
+    *,
+    max_chars: int,
+    shrink_cap: int,
+) -> None:
+    def tool_len(i: int) -> int:
+        msg = messages[i]
+        return len(msg.content) if isinstance(msg, ToolChatMessage) else 0
+
+    for idx in sorted(tool_indices, key=tool_len, reverse=True):
+        if estimate_messages_chars(messages) <= max_chars:
+            return
+        tool_msg = messages[idx]
+        if isinstance(tool_msg, ToolChatMessage):
+            messages[idx] = _shrink_tool(tool_msg, shrink_cap)
+
+
+def compact_messages_for_request(
+    messages: Sequence[AnyChatMessage],
+    *,
+    max_chars: int = DEFAULT_CONTEXT_MAX_CHARS,
+    old_tool_result_chars: int = DEFAULT_OLD_TOOL_RESULT_CHARS,
+    keep_recent_tool_results: int = DEFAULT_RECENT_TOOL_RESULTS,
+) -> list[AnyChatMessage]:
+    """Return a request-time copy with older tool dumps shrunk if over budget.
+
+    Does not mutate the original history (full results stay for persistence/UI).
+    ``max_chars < 1`` disables compacting.
+    """
+    out: list[AnyChatMessage] = list(messages)
+    if max_chars < 1 or estimate_messages_chars(out) <= max_chars:
+        return out
+
+    indices = _tool_indices(out)
+    if not indices:
+        return out
+
+    protect = _protect_indices(indices, keep_recent_tool_results)
+    _shrink_except(out, indices, protect, old_tool_result_chars)
+    if estimate_messages_chars(out) <= max_chars:
+        return out
+
+    _shrink_largest(
+        out,
+        indices,
+        max_chars=max_chars,
+        shrink_cap=max(64, old_tool_result_chars // 2),
+    )
+    return out

src/plyngent/agent/loop.py

@@ -16,7 +16,12 @@ from plyngent.lmproto.openai_compatible.model import (
 )
 from plyngent.typedef import Unset  # noqa: TC001
 
-from .budget import DEFAULT_TOOL_RESULT_MAX_CHARS, truncate_tool_result
+from .budget import (
+    DEFAULT_CONTEXT_MAX_CHARS,
+    DEFAULT_TOOL_RESULT_MAX_CHARS,
+    compact_messages_for_request,
+    truncate_tool_result,
+)
 from .events import (
     AgentEvent,
     AssistantMessageEvent,
@@ -41,18 +46,28 @@ if TYPE_CHECKING:
 DEFAULT_MAX_ROUNDS = 32
 
 
+def _raise_if_cancelled() -> None:
+    """Cooperative cancel points between model/tool steps."""
+    task = asyncio.current_task()
+    if task is not None and task.cancelling():
+        raise asyncio.CancelledError
+
+
 async def _run_one_tool(
     tools: ToolRegistry,
     call: AnyAssistantToolCall,
     *,
     max_result_chars: int,
 ) -> tuple[ToolChatMessage, ErrorEvent | None]:
+    _raise_if_cancelled()
     if isinstance(call, AssistantFunctionToolCall):
         try:
             result_text = await tools.execute(call.function.name, call.function.arguments)
+        except asyncio.CancelledError:
+            raise
         except Exception as exc:  # noqa: BLE001
             result_text = f"error: tool {call.function.name!r} failed: {exc}"
-            err = ErrorEvent(message=result_text)
+            err = ErrorEvent(message=result_text, retryable=True, source="tool")
             truncated = truncate_tool_result(result_text, max_result_chars)
             return ToolChatMessage(content=truncated, tool_call_id=call.id), err
         truncated = truncate_tool_result(result_text, max_result_chars)
@@ -61,7 +76,11 @@ async def _run_one_tool(
         content="error: custom tool calls are not supported",
         tool_call_id=call.id,
     )
-    return msg, None
+    return msg, ErrorEvent(
+        message=msg.content,
+        retryable=False,
+        source="tool",
+    )
 
 
 async def _execute_tool_calls(
@@ -75,6 +94,7 @@ async def _execute_tool_calls(
     for call in tool_calls:
         yield ToolCallEvent(tool_call=call)
 
+    _raise_if_cancelled()
     if parallel and len(tool_calls) > 1:
         results = await asyncio.gather(
             *[_run_one_tool(tools, call, max_result_chars=max_result_chars) for call in tool_calls]
@@ -158,12 +178,14 @@ async def run_chat_loop(
     stream: bool = True,
     max_tool_result_chars: int = DEFAULT_TOOL_RESULT_MAX_CHARS,
     parallel_tools: bool = True,
+    max_context_chars: int = DEFAULT_CONTEXT_MAX_CHARS,
 ) -> AsyncIterator[AgentEvent]:
     """Multi-round chat/tool loop; mutates ``messages`` in place and yields events.
 
     When ``stream=True``, uses ``chat_completions(..., stream=True)`` and yields
     text deltas as chunks arrive; tool calls are merged from stream deltas.
     Multiple tool calls in one round run in parallel when ``parallel_tools``.
+    Request payloads may shrink older tool results when over ``max_context_chars``.
     """
     tool_items: Sequence[AnyToolItem] | None = None
     if tools is not None and len(tools) > 0:
@@ -175,8 +197,13 @@ async def run_chat_loop(
     while True:
         while rounds_used < allowance:
             rounds_used += 1
+            _raise_if_cancelled()
+            request_messages = compact_messages_for_request(
+                messages,
+                max_chars=max_context_chars,
+            )
             param = ChatCompletionsParam(
-                messages=list(messages),
+                messages=request_messages,
                 model=model,
                 temperature=temperature if temperature is not None else UNSET,
                 tools=list(tool_items) if tool_items is not None else UNSET,

tests/test_agent/test_budget.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
-from plyngent.agent.budget import truncate_tool_result
+from plyngent.agent.budget import estimate_message_chars, truncate_tool_result
+from plyngent.lmproto.openai_compatible.model import ToolChatMessage, UserChatMessage
 
 
 def test_truncate_tool_result_short() -> None:
@@ -13,3 +14,9 @@ def test_truncate_tool_result_long() -> None:
     assert out.startswith("a" * 20)
     assert "truncated" in out
     assert "30" in out
+
+
+def test_estimate_message_chars() -> None:
+    assert estimate_message_chars(UserChatMessage(content="hello")) == 5  # noqa: PLR2004
+    tool = ToolChatMessage(content="abc", tool_call_id="id1")
+    assert estimate_message_chars(tool) == 6  # noqa: PLR2004

tests/test_agent/test_compact.py

@@ -0,0 +1,224 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal, overload
+
+from msgspec import UNSET
+
+from plyngent.agent.budget import (
+    compact_messages_for_request,
+    estimate_messages_chars,
+    truncate_tool_result,
+)
+from plyngent.agent.loop import run_chat_loop
+from plyngent.lmproto.openai_compatible.model import (
+    AssistantChatMessage,
+    AssistantFunctionTool,
+    AssistantFunctionToolCall,
+    ChatCompletionChoice,
+    ChatCompletionChunk,
+    ChatCompletionResponse,
+    ChatCompletionsParam,
+    ChunkChoice,
+    DeltaMessage,
+    StreamFunctionDelta,
+    StreamToolCallDelta,
+    ToolChatMessage,
+    UserChatMessage,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+
+    from plyngent.lmproto.openai_compatible.model import AnyChatMessage
+
+
+def test_truncate_tool_result_short() -> None:
+    assert truncate_tool_result("hello", 100) == "hello"
+
+
+def test_truncate_tool_result_long() -> None:
+    text = "a" * 50
+    out = truncate_tool_result(text, 20)
+    assert out.startswith("a" * 20)
+    assert "truncated" in out
+    assert "30" in out
+
+
+def test_compact_shrinks_old_tool_results() -> None:
+    messages: list[AnyChatMessage] = [
+        UserChatMessage(content="start"),
+        AssistantChatMessage(
+            content="",
+            tool_calls=[
+                AssistantFunctionToolCall(
+                    id="1",
+                    function=AssistantFunctionTool(name="t", arguments="{}"),
+                )
+            ],
+        ),
+        ToolChatMessage(content="OLD" * 200, tool_call_id="1"),
+        UserChatMessage(content="again"),
+        AssistantChatMessage(
+            content="",
+            tool_calls=[
+                AssistantFunctionToolCall(
+                    id="2",
+                    function=AssistantFunctionTool(name="t", arguments="{}"),
+                )
+            ],
+        ),
+        ToolChatMessage(content="NEW" * 50, tool_call_id="2"),
+    ]
+    original_old = messages[2]
+    assert isinstance(original_old, ToolChatMessage)
+    original_len = len(original_old.content)
+
+    compacted = compact_messages_for_request(
+        messages,
+        max_chars=estimate_messages_chars(messages) - 1,
+        old_tool_result_chars=40,
+        keep_recent_tool_results=1,
+    )
+    assert isinstance(compacted[2], ToolChatMessage)
+    assert len(compacted[2].content) < original_len
+    assert "truncated" in compacted[2].content
+    # Full history unchanged
+    assert isinstance(messages[2], ToolChatMessage)
+    assert len(messages[2].content) == original_len
+    # Recent tool kept
+    assert isinstance(compacted[5], ToolChatMessage)
+    assert compacted[5].content == "NEW" * 50
+
+
+def test_compact_disabled_when_max_chars_zero() -> None:
+    messages: list[AnyChatMessage] = [
+        ToolChatMessage(content="x" * 500, tool_call_id="1"),
+    ]
+    out = compact_messages_for_request(messages, max_chars=0)
+    assert out[0] is messages[0] or (
+        isinstance(out[0], ToolChatMessage) and out[0].content == "x" * 500
+    )
+
+
+def _response(message: AssistantChatMessage) -> ChatCompletionResponse:
+    return ChatCompletionResponse(
+        id="1",
+        object="chat.completion",
+        created=0,
+        model="t",
+        choices=[ChatCompletionChoice(index=0, message=message, logprobs={}, finish_reason="stop")],
+        system_fingerprint="",
+        usage={},
+    )
+
+
+class CaptureClient:
+    _responses: list[ChatCompletionResponse]
+    calls: list[ChatCompletionsParam]
+
+    def __init__(self, responses: list[ChatCompletionResponse]) -> None:
+        self._responses = list(responses)
+        self.calls = []
+
+    @overload
+    async def chat_completions(
+        self, param: ChatCompletionsParam, *, stream: Literal[False] = False
+    ) -> ChatCompletionResponse: ...
+
+    @overload
+    async def chat_completions(
+        self, param: ChatCompletionsParam, *, stream: Literal[True]
+    ) -> AsyncIterator[ChatCompletionChunk]: ...
+
+    async def chat_completions(
+        self, param: ChatCompletionsParam, *, stream: bool = False
+    ) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]:
+        self.calls.append(param)
+        response = self._responses.pop(0)
+        if stream:
+
+            async def as_stream() -> AsyncIterator[ChatCompletionChunk]:
+                message = response.choices[0].message
+                if isinstance(message.content, str) and message.content:
+                    yield ChatCompletionChunk(
+                        id="1",
+                        object="chat.completion.chunk",
+                        created=0,
+                        model="t",
+                        choices=[
+                            ChunkChoice(
+                                index=0,
+                                delta=DeltaMessage(content=message.content),
+                                finish_reason=None,
+                            )
+                        ],
+                    )
+                tool_calls = message.tool_calls
+                if tool_calls is not UNSET and tool_calls:
+                    deltas: list[StreamToolCallDelta] = []
+                    for i, call in enumerate(tool_calls):
+                        if isinstance(call, AssistantFunctionToolCall):
+                            deltas.append(
+                                StreamToolCallDelta(
+                                    index=i,
+                                    id=call.id,
+                                    type="function",
+                                    function=StreamFunctionDelta(
+                                        name=call.function.name,
+                                        arguments=call.function.arguments,
+                                    ),
+                                )
+                            )
+                    yield ChatCompletionChunk(
+                        id="1",
+                        object="chat.completion.chunk",
+                        created=0,
+                        model="t",
+                        choices=[
+                            ChunkChoice(
+                                index=0,
+                                delta=DeltaMessage(tool_calls=deltas),
+                                finish_reason="tool_calls",
+                            )
+                        ],
+                    )
+
+            return as_stream()
+        return response
+
+
+async def test_loop_sends_compacted_request_not_history() -> None:
+    big = "Z" * 500
+    history: list[AnyChatMessage] = [
+        UserChatMessage(content="u"),
+        AssistantChatMessage(
+            content="",
+            tool_calls=[
+                AssistantFunctionToolCall(
+                    id="1",
+                    function=AssistantFunctionTool(name="t", arguments="{}"),
+                )
+            ],
+        ),
+        ToolChatMessage(content=big, tool_call_id="1"),
+        UserChatMessage(content="next"),
+    ]
+    client = CaptureClient([_response(AssistantChatMessage(content="ok"))])
+    _ = [
+        e
+        async for e in run_chat_loop(
+            client,
+            history,
+            model="m",
+            stream=False,
+            max_context_chars=200,
+            max_tool_result_chars=50,
+        )
+    ]
+    assert client.calls
+    sent = client.calls[0].messages
+    tool_sent = next(m for m in sent if isinstance(m, ToolChatMessage))
+    assert len(tool_sent.content) < len(big)
+    # In-memory history still full for the old tool result
+    assert isinstance(history[2], ToolChatMessage)
+    assert history[2].content == big