From a634021a807aee982840f06c2185321cd703acaa Mon Sep 17 00:00:00 2001 From: worldmozara Date: Tue, 14 Jul 2026 19:53:34 +0800 Subject: [PATCH] core/cli: interactive continue prompts for loop and PTY limits on_limit raises tool-round allowance; PTY session/budget limits can prompt to raise; MaxRoundsEvent.continued marks extended runs. --- CLAUDE.md | 3 +- src/plyngent/agent/chat.py | 8 +- src/plyngent/agent/events.py | 1 + src/plyngent/agent/loop.py | 100 +++++++++++++--------- src/plyngent/cli/app.py | 2 + src/plyngent/cli/display.py | 8 +- src/plyngent/cli/limits.py | 20 +++++ src/plyngent/cli/state.py | 3 + src/plyngent/tools/process/pty_session.py | 62 +++++++++++--- tests/test_agent/test_loop.py | 45 +++++++++- tests/test_cli/test_limits.py | 32 +++++++ tests/test_tools/test_process.py | 19 ++++ 12 files changed, 248 insertions(+), 55 deletions(-) create mode 100644 src/plyngent/cli/limits.py create mode 100644 tests/test_cli/test_limits.py diff --git a/CLAUDE.md b/CLAUDE.md index ddcad61..4514763 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -54,7 +54,7 @@ Async SQLAlchemy + aiosqlite. `MemoryStore`: schema init, default local user, se - **`ChatClient`** Protocol for `chat_completions`. - **`@tool` / `ToolRegistry`**: decorator infers JSON Schema from type hints; execute tools by name. -- **`run_chat_loop`**: multi-round tool loop, yields `AgentEvent` stream. +- **`run_chat_loop`**: multi-round tool loop, yields `AgentEvent` stream; optional `on_limit` to continue past max rounds. - **`ChatAgent`**: wrapper with optional `MemoryStore` bind (load/append messages). ### Tools (`tools/`) @@ -65,6 +65,7 @@ Module-level `@tool` handlers. Call `set_workspace_root()` before use. - **`file`**: `read_file`, `write_file`, `listdir`, `edit_replace` (first occurrence). - **`process`**: `run_command` (argv, no shell, timeout, optional stdin/env); PTY `open_pty` / `read_pty` / `write_pty` / `close_pty` (**Unix only**: `pty`+`fork`). - PTY: structured status (`alive`/`exit_code`/`data`); `read_pty(..., until=)`; session limit/idle TTL/output budget; close SIGTERM→SIGKILL. +- CLI limit hooks: interactive confirm to raise tool-loop rounds, PTY session cap, or PTY output budget. - **`DEFAULT_TOOLS`**: file + process tool list for a `ToolRegistry`. ### CLI (`cli/`) diff --git a/src/plyngent/agent/chat.py b/src/plyngent/agent/chat.py index 2c5a2c5..ce6a09e 100644 --- a/src/plyngent/agent/chat.py +++ b/src/plyngent/agent/chat.py @@ -7,7 +7,7 @@ from plyngent.lmproto.openai_compatible.model import UserChatMessage from .loop import DEFAULT_MAX_ROUNDS, run_chat_loop if TYPE_CHECKING: - from collections.abc import AsyncIterator, Sequence + from collections.abc import AsyncIterator, Callable, Sequence from plyngent.lmproto.openai_compatible.model import AnyChatMessage from plyngent.memory import MemoryStore @@ -16,6 +16,8 @@ if TYPE_CHECKING: from .events import AgentEvent from .tools import ToolRegistry + type LimitContinueHook = Callable[[str], bool] + class ChatAgent: """Thin wrapper: chat client + optional tools + optional memory bind.""" @@ -27,6 +29,7 @@ class ChatAgent: session_id: int | None max_rounds: int temperature: float | None + on_limit: LimitContinueHook | None messages: list[AnyChatMessage] def __init__( # noqa: PLR0913 @@ -40,6 +43,7 @@ class ChatAgent: max_rounds: int = DEFAULT_MAX_ROUNDS, temperature: float | None = None, messages: Sequence[AnyChatMessage] | None = None, + on_limit: LimitContinueHook | None = None, ) -> None: self.client = client self.model = model @@ -48,6 +52,7 @@ class ChatAgent: self.session_id = session_id self.max_rounds = max_rounds self.temperature = temperature + self.on_limit = on_limit self.messages = list(messages) if messages is not None else [] async def load_history(self) -> None: @@ -84,6 +89,7 @@ class ChatAgent: tools=self.tools, max_rounds=self.max_rounds, temperature=self.temperature, + on_limit=self.on_limit, ): yield event diff --git a/src/plyngent/agent/events.py b/src/plyngent/agent/events.py index 1cd9c55..41d2f4c 100644 --- a/src/plyngent/agent/events.py +++ b/src/plyngent/agent/events.py @@ -25,6 +25,7 @@ class ToolResultEvent(Struct, tag_field="type", tag="tool_result"): class MaxRoundsEvent(Struct, tag_field="type", tag="max_rounds"): rounds: int + continued: bool = False type AgentEvent = TextDeltaEvent | AssistantMessageEvent | ToolCallEvent | ToolResultEvent | MaxRoundsEvent diff --git a/src/plyngent/agent/loop.py b/src/plyngent/agent/loop.py index fa6e1fb..3f81bcc 100644 --- a/src/plyngent/agent/loop.py +++ b/src/plyngent/agent/loop.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING from msgspec import UNSET from plyngent.lmproto.openai_compatible.model import ( + AnyAssistantToolCall, AssistantChatMessage, AssistantFunctionToolCall, ChatCompletionsParam, @@ -21,17 +22,38 @@ from .events import ( ) if TYPE_CHECKING: - from collections.abc import AsyncIterator, Sequence + from collections.abc import AsyncIterator, Callable, Sequence from plyngent.lmproto.openai_compatible.model import AnyChatMessage, AnyToolItem from .client import ChatClient from .tools import ToolRegistry + type LimitContinueHook = Callable[[str], bool] + DEFAULT_MAX_ROUNDS = 32 +async def _execute_tool_calls( + tools: ToolRegistry, + tool_calls: Sequence[AnyAssistantToolCall], + messages: list[AnyChatMessage], +) -> AsyncIterator[AgentEvent]: + for call in tool_calls: + yield ToolCallEvent(tool_call=call) + if isinstance(call, AssistantFunctionToolCall): + result_text = await tools.execute(call.function.name, call.function.arguments) + tool_msg = ToolChatMessage(content=result_text, tool_call_id=call.id) + else: + tool_msg = ToolChatMessage( + content="error: custom tool calls are not supported", + tool_call_id=call.id, + ) + messages.append(tool_msg) + yield ToolResultEvent(message=tool_msg) + + async def run_chat_loop( # noqa: PLR0913 client: ChatClient, messages: list[AnyChatMessage], @@ -40,58 +62,56 @@ async def run_chat_loop( # noqa: PLR0913 tools: ToolRegistry | None = None, max_rounds: int = DEFAULT_MAX_ROUNDS, temperature: float | None = None, + on_limit: LimitContinueHook | None = None, ) -> AsyncIterator[AgentEvent]: """Multi-round chat/tool loop; mutates ``messages`` in place and yields events. Continues until the model returns no tool calls, or ``max_rounds`` is hit. - Streaming is non-stream LLM calls for reliable tool_calls; text is emitted - as a single :class:`TextDeltaEvent` when content is present. + If ``on_limit`` is set and returns True when the cap is reached, another + batch of ``max_rounds`` is granted and the loop continues. """ tool_items: Sequence[AnyToolItem] | None = None if tools is not None and len(tools) > 0: tool_items = tools.tool_items() - for round_idx in range(max_rounds): - param = ChatCompletionsParam( - messages=list(messages), - model=model, - temperature=temperature if temperature is not None else UNSET, - tools=list(tool_items) if tool_items is not None else UNSET, - ) - response = await client.chat_completions(param, stream=False) - if not response.choices: - msg = "chat completion response contained no choices" - raise RuntimeError(msg) - assistant = response.choices[0].message - messages.append(assistant) - yield AssistantMessageEvent(message=assistant) + rounds_used = 0 + allowance = max_rounds - if isinstance(assistant.content, str) and assistant.content: - yield TextDeltaEvent(content=assistant.content) + while True: + while rounds_used < allowance: + rounds_used += 1 + param = ChatCompletionsParam( + messages=list(messages), + model=model, + temperature=temperature if temperature is not None else UNSET, + tools=list(tool_items) if tool_items is not None else UNSET, + ) + response = await client.chat_completions(param, stream=False) + if not response.choices: + msg = "chat completion response contained no choices" + raise RuntimeError(msg) + assistant = response.choices[0].message + messages.append(assistant) + yield AssistantMessageEvent(message=assistant) - tool_calls = assistant.tool_calls - if tool_calls is UNSET or not tool_calls: - return + if isinstance(assistant.content, str) and assistant.content: + yield TextDeltaEvent(content=assistant.content) - if tools is None: - return + tool_calls = assistant.tool_calls + if tool_calls is UNSET or not tool_calls: + return + if tools is None: + return + async for event in _execute_tool_calls(tools, tool_calls, messages): + yield event - for call in tool_calls: - yield ToolCallEvent(tool_call=call) - if isinstance(call, AssistantFunctionToolCall): - result_text = await tools.execute(call.function.name, call.function.arguments) - tool_msg = ToolChatMessage(content=result_text, tool_call_id=call.id) - else: - tool_msg = ToolChatMessage( - content="error: custom tool calls are not supported", - tool_call_id=call.id, - ) - messages.append(tool_msg) - yield ToolResultEvent(message=tool_msg) - - _ = round_idx # used only for loop bound - - yield MaxRoundsEvent(rounds=max_rounds) + reason = f"tool loop reached {allowance} rounds (used {rounds_used})" + if on_limit is not None and on_limit(reason): + yield MaxRoundsEvent(rounds=allowance, continued=True) + allowance += max_rounds + continue + yield MaxRoundsEvent(rounds=allowance, continued=False) + return def collect_assistant_messages(events: Sequence[AgentEvent]) -> list[AssistantChatMessage]: diff --git a/src/plyngent/cli/app.py b/src/plyngent/cli/app.py index 9f73c79..6445501 100644 --- a/src/plyngent/cli/app.py +++ b/src/plyngent/cli/app.py @@ -14,6 +14,7 @@ from plyngent.cli.editor import ( open_in_editor, resolve_config_path, ) +from plyngent.cli.limits import install_cli_limit_hooks from plyngent.cli.repl import run_repl from plyngent.cli.selection import select_model, select_provider from plyngent.cli.state import ReplState @@ -66,6 +67,7 @@ async def _run_chat( # noqa: PLR0913 raise click.ClickException(str(exc)) from exc _ = set_workspace_root(workspace) + install_cli_limit_hooks() memory = await MemoryStore.open(_database_config(store)) try: state = ReplState( diff --git a/src/plyngent/cli/display.py b/src/plyngent/cli/display.py index 3321994..d0aaaab 100644 --- a/src/plyngent/cli/display.py +++ b/src/plyngent/cli/display.py @@ -48,7 +48,13 @@ async def render_events(events: AsyncIterator[AgentEvent]) -> None: ) click.secho(f"[tool result] {preview}", fg="magenta") elif isinstance(event, MaxRoundsEvent): - click.secho(f"\n[max rounds reached: {event.rounds}]", fg="red") + if event.continued: + click.secho( + f"\n[max rounds {event.rounds} reached — continuing with a higher allowance]", + fg="yellow", + ) + else: + click.secho(f"\n[max rounds reached: {event.rounds}]", fg="red") else: # AssistantMessageEvent — text already shown via TextDeltaEvent. _ = event diff --git a/src/plyngent/cli/limits.py b/src/plyngent/cli/limits.py new file mode 100644 index 0000000..02e373f --- /dev/null +++ b/src/plyngent/cli/limits.py @@ -0,0 +1,20 @@ +from __future__ import annotations + +import click + +from plyngent.tools.process.pty_session import PtyManager + + +def prompt_continue_limit(reason: str) -> bool: + """Ask the user whether to raise a limit and continue (TTY).""" + click.echo() + click.secho(f"[limit] {reason}", fg="yellow") + try: + return bool(click.confirm("Raise limit and continue?", default=True)) + except click.Abort: + return False + + +def install_cli_limit_hooks() -> None: + """Register interactive continue hooks for process-global tool limits.""" + PtyManager.set_limit_continue_hook(prompt_continue_limit) diff --git a/src/plyngent/cli/state.py b/src/plyngent/cli/state.py index bdb6ee9..76cc249 100644 --- a/src/plyngent/cli/state.py +++ b/src/plyngent/cli/state.py @@ -43,6 +43,8 @@ class ReplState: return ToolRegistry(list(DEFAULT_TOOLS)) def _make_agent(self) -> ChatAgent: + from plyngent.cli.limits import prompt_continue_limit + return ChatAgent( self.client, model=self.model, @@ -50,6 +52,7 @@ class ReplState: memory=self.memory, session_id=self.session_id, max_rounds=self.max_rounds, + on_limit=prompt_continue_limit, ) def rebuild_client(self) -> None: diff --git a/src/plyngent/tools/process/pty_session.py b/src/plyngent/tools/process/pty_session.py index a018d2d..02ad563 100644 --- a/src/plyngent/tools/process/pty_session.py +++ b/src/plyngent/tools/process/pty_session.py @@ -10,16 +10,23 @@ import signal import time from dataclasses import dataclass, field from threading import Lock -from typing import ClassVar +from typing import TYPE_CHECKING, ClassVar from plyngent.tools.workspace import WorkspaceError, check_command_allowed, resolve_path +if TYPE_CHECKING: + from collections.abc import Callable + + type LimitContinueHook = Callable[[str], bool] + DEFAULT_PTY_READ_BYTES = 8192 DEFAULT_PTY_POLL_TIMEOUT = 0.2 DEFAULT_MAX_SESSIONS = 8 DEFAULT_IDLE_TTL_SECONDS = 600.0 DEFAULT_SESSION_OUTPUT_BUDGET = 256_000 DEFAULT_CLOSE_GRACE_SECONDS = 0.5 +_SESSION_LIMIT_STEP = 4 +_BUDGET_STEP = 256_000 _STDERR_FD = 2 _EXEC_FAIL_MARKER = b"plyngent-pty-exec-failed: " @@ -67,6 +74,7 @@ class PtyManager: max_sessions: ClassVar[int] = DEFAULT_MAX_SESSIONS idle_ttl_seconds: ClassVar[float] = DEFAULT_IDLE_TTL_SECONDS session_output_budget: ClassVar[int] = DEFAULT_SESSION_OUTPUT_BUDGET + _limit_continue: ClassVar[LimitContinueHook | None] = None @classmethod def configure( @@ -83,6 +91,21 @@ class PtyManager: if session_output_budget is not None: cls.session_output_budget = max(1024, session_output_budget) + @classmethod + def set_limit_continue_hook(cls, hook: LimitContinueHook | None) -> None: + """Optional interactive hook: return True to raise a limit and continue.""" + cls._limit_continue = hook + + @classmethod + def _offer_raise(cls, reason: str) -> bool: + hook = cls._limit_continue + if hook is None: + return False + try: + return bool(hook(reason)) + except Exception: # noqa: BLE001 — never break tools on prompt failure + return False + @classmethod def open( cls, @@ -100,8 +123,12 @@ class PtyManager: with cls._lock: alive_count = sum(1 for s in cls._sessions.values() if not s.closed) if alive_count >= cls.max_sessions: - msg = f"PTY session limit reached ({cls.max_sessions}); close idle sessions" - raise WorkspaceError(msg) + reason = f"PTY session limit reached ({cls.max_sessions})" + if cls._offer_raise(f"{reason}; raise by {_SESSION_LIMIT_STEP}?"): + cls.max_sessions += _SESSION_LIMIT_STEP + else: + msg = f"{reason}; close idle sessions or allow a higher limit" + raise WorkspaceError(msg) master_fd, slave_fd = pty.openpty() pid = os.fork() @@ -255,13 +282,19 @@ class PtyManager: raise WorkspaceError(msg) if (cls.session_output_budget - session.bytes_read) <= 0: - return PtyReadResult( - session_id=session_id, - alive=session.alive, - exit_code=session.exit_code, - data="", - budget_exhausted=True, - ) + if cls._offer_raise( + f"PTY output budget exhausted for session {session_id} " + f"({cls.session_output_budget} bytes); raise by {_BUDGET_STEP}?" + ): + cls.session_output_budget += _BUDGET_STEP + else: + return PtyReadResult( + session_id=session_id, + alive=session.alive, + exit_code=session.exit_code, + data="", + budget_exhausted=True, + ) chunks, matched = cls._collect_chunks( session, max_bytes=max_bytes, timeout=timeout, until=until @@ -270,6 +303,13 @@ class PtyManager: truncated = len(data) > max_bytes if truncated: data = data[:max_bytes] + budget_exhausted = (cls.session_output_budget - session.bytes_read) <= 0 + if budget_exhausted and cls._offer_raise( + f"PTY output budget exhausted for session {session_id} " + f"({cls.session_output_budget} bytes); raise by {_BUDGET_STEP}?" + ): + cls.session_output_budget += _BUDGET_STEP + budget_exhausted = False return PtyReadResult( session_id=session_id, alive=session.alive, @@ -277,7 +317,7 @@ class PtyManager: data=data, truncated=truncated, matched=matched, - budget_exhausted=(cls.session_output_budget - session.bytes_read) <= 0, + budget_exhausted=budget_exhausted, ) @classmethod diff --git a/tests/test_agent/test_loop.py b/tests/test_agent/test_loop.py index ac14922..46dbb36 100644 --- a/tests/test_agent/test_loop.py +++ b/tests/test_agent/test_loop.py @@ -155,10 +155,53 @@ async def test_max_rounds() -> None: client = ScriptedClient([forever, forever, forever]) messages: list[AnyChatMessage] = [UserChatMessage(content="x")] events = [e async for e in run_chat_loop(client, messages, model="m", tools=registry, max_rounds=2)] - assert any(isinstance(e, MaxRoundsEvent) and e.rounds == 2 for e in events) # noqa: PLR2004 + assert any(isinstance(e, MaxRoundsEvent) and e.rounds == 2 and not e.continued for e in events) # noqa: PLR2004 assert len(client.calls) == 2 # noqa: PLR2004 +async def test_max_rounds_continue_hook() -> None: + @tool + def ping() -> str: + return "pong" + + registry = ToolRegistry([ping]) + forever = _response( + AssistantChatMessage( + content="", + tool_calls=[ + AssistantFunctionToolCall( + id="c", + function=AssistantFunctionTool(name="ping", arguments="{}"), + ) + ], + ) + ) + final = _response(AssistantChatMessage(content="done")) + client = ScriptedClient([forever, forever, final]) + messages: list[AnyChatMessage] = [UserChatMessage(content="x")] + asks: list[str] = [] + + def on_limit(reason: str) -> bool: + asks.append(reason) + return len(asks) == 1 + + events = [ + e + async for e in run_chat_loop( + client, + messages, + model="m", + tools=registry, + max_rounds=2, + on_limit=on_limit, + ) + ] + assert len(asks) == 1 + assert any(isinstance(e, MaxRoundsEvent) and e.continued for e in events) + assert any(isinstance(e, TextDeltaEvent) and e.content == "done" for e in events) + assert len(client.calls) == 3 # noqa: PLR2004 + + async def test_default_max_rounds_is_generous() -> None: from plyngent.agent.loop import DEFAULT_MAX_ROUNDS diff --git a/tests/test_cli/test_limits.py b/tests/test_cli/test_limits.py new file mode 100644 index 0000000..35bff6f --- /dev/null +++ b/tests/test_cli/test_limits.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +from plyngent.cli.limits import install_cli_limit_hooks, prompt_continue_limit +from plyngent.tools.process.pty_session import PtyManager + +if TYPE_CHECKING: + import pytest + + +def test_prompt_continue_limit_yes(monkeypatch: pytest.MonkeyPatch) -> None: + def _confirm(*_a: object, **_k: object) -> bool: + return True + + monkeypatch.setattr("click.confirm", _confirm) + assert prompt_continue_limit("hit a wall") is True + + +def test_prompt_continue_limit_no(monkeypatch: pytest.MonkeyPatch) -> None: + def _confirm(*_a: object, **_k: object) -> bool: + return False + + monkeypatch.setattr("click.confirm", _confirm) + assert prompt_continue_limit("hit a wall") is False + + +def test_install_cli_limit_hooks() -> None: + install_cli_limit_hooks() + # Hook is installed process-wide for the CLI session. + assert callable(getattr(PtyManager, "_limit_continue", None)) + PtyManager.set_limit_continue_hook(None) diff --git a/tests/test_tools/test_process.py b/tests/test_tools/test_process.py index 60a447f..47400b1 100644 --- a/tests/test_tools/test_process.py +++ b/tests/test_tools/test_process.py @@ -147,6 +147,7 @@ def test_pty_session_limit(workspace: object) -> None: del workspace previous = PtyManager.max_sessions try: + PtyManager.set_limit_continue_hook(None) PtyManager.configure(max_sessions=1) first = call_sync(open_pty, ["sleep", "30"]) assert "session_id=" in first @@ -155,6 +156,24 @@ def test_pty_session_limit(workspace: object) -> None: finally: PtyManager.close_all() PtyManager.configure(max_sessions=previous) + PtyManager.set_limit_continue_hook(None) + + +def test_pty_session_limit_continue(workspace: object) -> None: + del workspace + previous = PtyManager.max_sessions + try: + PtyManager.configure(max_sessions=1) + PtyManager.set_limit_continue_hook(lambda _reason: True) + first = call_sync(open_pty, ["sleep", "30"]) + second = call_sync(open_pty, ["sleep", "30"]) + assert "session_id=" in first + assert "session_id=" in second + assert PtyManager.max_sessions >= 2 # noqa: PLR2004 + finally: + PtyManager.close_all() + PtyManager.configure(max_sessions=previous) + PtyManager.set_limit_continue_hook(None) def test_pty_output_budget(workspace: object) -> None: