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,