diff --git a/docs/source-en/rst_source/usage/configure_planner.rst b/docs/source-en/rst_source/usage/configure_planner.rst index 5a495a41..acc60c6c 100644 --- a/docs/source-en/rst_source/usage/configure_planner.rst +++ b/docs/source-en/rst_source/usage/configure_planner.rst @@ -172,6 +172,11 @@ Notes: the model configured as the Codex SDK default. - ``--planner-timeout-s`` limits the Codex run. Its default is ``CODEX_TIMEOUT_S``, then ``CELL_TIMEOUT_S``, then ``1200`` seconds. +- ``--max-turns`` limits completed model responses, counted from the SDK's + cumulative usage updates. Reasoning/tool-only responses count too; multiple + text or tool items in one response count once. Duplicate usage updates do + not count again. CLI and Dashboard use the same budget and request an + interrupt when it is reached; a recorded ``finish`` is preserved. - By default, the Codex SDK reuses existing Codex authentication. For a custom Responses-compatible endpoint, set ``CODEX_BASE_URL`` and ``CODEX_API_KEY``. This backend does not read ``OPENAI_BASE_URL`` or diff --git a/docs/source-zh/rst_source/usage/configure_planner.rst b/docs/source-zh/rst_source/usage/configure_planner.rst index 3ca3196f..efd1ce80 100644 --- a/docs/source-zh/rst_source/usage/configure_planner.rst +++ b/docs/source-zh/rst_source/usage/configure_planner.rst @@ -153,6 +153,10 @@ Codex 规划会话不会自动加载仓库的 ``AGENTS.md`` 和 ``.agents/skills 配置的默认模型。 - ``--planner-timeout-s`` 限制 Codex 运行时间。默认依次读取 ``CODEX_TIMEOUT_S``、``CELL_TIMEOUT_S``,均未设置时为 ``1200`` 秒。 +- ``--max-turns`` 根据 SDK 的累计 usage 更新限制已完成的模型响应次数。 + 只有推理和工具调用的响应也计入预算;同一次响应中的多个文本或工具条目 + 只计一次,重复 usage 通知不重复计数。CLI 和 Dashboard 使用相同预算, + 达到上限时请求中断,并保留已经记录的 ``finish``。 - 默认情况下,Codex SDK 会复用已有的 Codex 认证。若要接入自定义的 Responses API 兼容端点,请设置 ``CODEX_BASE_URL`` 和 ``CODEX_API_KEY``;这里不读取 ``OPENAI_BASE_URL`` 或 diff --git a/rpent/planner/codex.py b/rpent/planner/codex.py index 14eb87ac..d28be491 100644 --- a/rpent/planner/codex.py +++ b/rpent/planner/codex.py @@ -326,6 +326,7 @@ def _steer() -> None: daemon=True, ).start() + limit_reached = False try: for event in turn.stream(): _write_jsonl(raw_f, _message_to_json(event)) @@ -335,6 +336,14 @@ def _steer() -> None: out_f.write(rendered) out_f.flush() logger.info(rendered.strip()) + if ( + str(_get(event, "method", "")) != "turn/completed" + and not limit_reached + and recorder.finish_result is None + and recorder.turns >= recorder.max_turns + ): + limit_reached = True + turn.interrupt() finally: if stop_steer is not None: stop_steer.set() @@ -618,6 +627,7 @@ class _Recorder: max_turns: int dashboard_events: DashboardEventSink turns: int = 0 + _seen_usage: set[tuple[int, ...]] = field(default_factory=set) tool_calls: int = 0 usage: dict[str, int] = field( default_factory=lambda: { @@ -643,7 +653,8 @@ def observe(self, event: Any) -> str: if method == "item/completed": return self._render_item(_get(payload, "item")) if method == "thread/tokenUsage/updated": - self._set_usage(_get(payload, "token_usage")) + if self._set_usage(_get(payload, "token_usage")): + return f"\n[agent] === turn {self.turns}/{self.max_turns} ===\n" return "" if method == "turn/completed": return self._render_turn_completed(_get(payload, "turn")) @@ -672,12 +683,8 @@ def _render_item(self, item: Any) -> str: if not text: return "" self.final_response = text - self.turns += 1 self.dashboard_events.emit(TranscriptEvent({"type": "text", "text": text})) - return ( - f"\n[agent] === turn {self.turns}/{self.max_turns} ===\n" - f"[codex] {text}\n" - ) + return f"\n[codex] {text}\n" if item_type == "reasoning": text = _extract_text(_get(item, "summary") or _get(item, "content")) @@ -731,11 +738,18 @@ def _render_turn_completed(self, turn: Any) -> str: # -- helpers ----------------------------------------------------------- - def _set_usage(self, usage: Any) -> None: + def _set_usage(self, usage: Any) -> bool: + """Count completed model responses, including reasoning/tool-only ones. + + The SDK updates cumulative token usage after each model response. Text + and tool items within that response do not consume additional turns. + Repeated notifications (including context-window-only updates) do not + count again or overwrite newer usage totals. + """ if usage is None: - return + return False total = _get(usage, "total", usage) - self.usage = { + updated = { "total_input_tokens": _int_attr(total, "input_tokens"), "total_cached_input_tokens": _int_attr(total, "cached_input_tokens"), "total_output_tokens": _int_attr(total, "output_tokens"), @@ -743,6 +757,12 @@ def _set_usage(self, usage: Any) -> None: total, "reasoning_output_tokens" ), } + key = tuple(updated.values()) + if not any(key) or key in self._seen_usage: + return False + self._seen_usage.add(key) + self.usage = updated + self.turns += 1 self.dashboard_events.emit( UsageEvent( inp=self.usage["total_input_tokens"], @@ -750,6 +770,7 @@ def _set_usage(self, usage: Any) -> None: tool_calls=self.tool_calls, ) ) + return True def _maybe_capture_finish(self, name: str, item: Any) -> None: if self.finish_result is not None: diff --git a/tests/unit_tests/rpent/planner/test_codex_contracts.py b/tests/unit_tests/rpent/planner/test_codex_contracts.py index 44221ae7..6b8ddd8e 100644 --- a/tests/unit_tests/rpent/planner/test_codex_contracts.py +++ b/tests/unit_tests/rpent/planner/test_codex_contracts.py @@ -14,6 +14,7 @@ from __future__ import annotations +import asyncio import json import os import queue @@ -454,6 +455,184 @@ def test_successful_fake_codex_lifecycle_uses_fake_mcp_and_accounts_events( assert any(isinstance(event, UsageEvent) for event in sink.events) +def _model_response_events(index: int, *, text: bool = False, finish: bool = False): + yield { + "method": "item/completed", + "payload": {"item": {"type": "reasoning", "summary": []}}, + } + for message in ("working", "still working") if text else ("",): + yield { + "method": "item/completed", + "payload": {"item": {"type": "agentMessage", "text": message}}, + } + for tool in ("read_text_file", "finish" if finish else "list_dir"): + yield { + "method": "item/completed", + "payload": { + "item": { + "type": "mcpToolCall", + "tool": tool, + "status": "completed", + "arguments": {"status": "stuck", "summary": "done"} + if tool == "finish" + else {}, + "result": "accepted", + } + }, + } + usage = { + "method": "thread/tokenUsage/updated", + "payload": { + "token_usage": { + "total": {"input_tokens": 10 * index, "output_tokens": 2 * index} + } + }, + } + yield usage + yield usage # A repeated SDK notification is not another model response. + + +@pytest.mark.parametrize("max_turns", [1, 2]) +@pytest.mark.parametrize("text", [False, True]) +def test_cli_limits_model_responses_with_or_without_text( + tmp_path, monkeypatch, max_turns, text +): + install_fake_backend(monkeypatch) + + def stream(self): + for index in range(1, 6): + if self.interrupt_calls: + break + yield from _model_response_events(index, text=text) + yield { + "method": "turn/completed", + "payload": {"turn": {"status": "interrupted"}}, + } + + monkeypatch.setattr(FakeTurn, "stream", stream) + result = make_planner(tmp_path, RecordingSink()).solve( + system_prompt="system", + user_message="task", + toolkit=FakeToolkit(), + max_turns=max_turns, + ) + assert result.stats["turns_used"] == max_turns + assert result.stats["tool_calls"] == 2 * max_turns + assert result.stats["total_input_tokens"] == 10 * max_turns + assert FakeCodex.instances[0].thread.fake_turn.interrupt_calls == 1 + assert FakeCodex.instances[0].closed + assert FakeMcpServer.instances[0].stopped + assert result.error is None + + +def test_cli_keeps_finish_at_the_response_budget(tmp_path, monkeypatch): + install_fake_backend(monkeypatch) + FakeCodex.events = [ + *_model_response_events(1, finish=True), + {"method": "turn/completed", "payload": {"turn": {"status": "completed"}}}, + ] + result = make_planner(tmp_path, RecordingSink()).solve( + system_prompt="system", + user_message="task", + toolkit=FakeToolkit(), + max_turns=1, + ) + assert result.stats["turns_used"] == 1 + assert result.finish_result["status"] == "stuck" + assert FakeCodex.instances[0].thread.fake_turn.interrupt_calls == 0 + assert FakeCodex.instances[0].closed + assert FakeMcpServer.instances[0].stopped + assert result.error is None + + +@pytest.mark.parametrize("max_turns", [1, 2]) +@pytest.mark.parametrize("finish", [False, True]) +def test_dashboard_uses_the_same_response_budget_and_closes(max_turns, finish): + from rpent.planner.codex import _CodexDashboardSession, _Recorder + + async def run(): + class Turn: + interrupt_calls = 0 + + async def stream(self): + for index in range(1, 6): + if self.interrupt_calls: + break + yield_events = _model_response_events(index, finish=finish) + for event in yield_events: + yield event + if finish: + break + yield { + "method": "turn/completed", + "payload": { + "turn": {"status": "completed" if finish else "interrupted"} + }, + } + + async def interrupt(self): + self.interrupt_calls += 1 + + class Control: + ended = False + closed = False + + async def tool_completed(self, session): + pass + + async def complete(self, session): + raise AssertionError("a finished or budget-limited session must end") + + def end(self): + self.ended = True + + async def close(self): + self.closed = True + + recorder = _Recorder(max_turns=max_turns, dashboard_events=RecordingSink()) + control = Control() + session = _CodexDashboardSession( + config=None, + thread_options={}, + turn_options={}, + recorder=recorder, + emit_event=recorder.observe, + control=control, + ) + turn = Turn() + done = asyncio.Event() + session._codex = control + session._turn = turn + session._turn_done = done + await session._consume_turn(turn, done) + await session.close() + assert done.is_set() and control.ended and control.closed + assert session.error is None + assert turn.interrupt_calls == (0 if finish else 1) + assert recorder.turns == (1 if finish else max_turns) + assert (recorder.finish_result is not None) == finish + + asyncio.run(run()) + + +def test_usage_replays_do_not_count_again_or_regress_totals(): + from rpent.planner.codex import _Recorder + + recorder = _Recorder(max_turns=10, dashboard_events=RecordingSink()) + for index in (1, 1, 2, 1): + for event in _model_response_events(index): + recorder.observe(event) + assert recorder.turns == 2 + assert recorder.stats()["total_input_tokens"] == 20 + recorder.observe( + { + "method": "thread/tokenUsage/updated", + "payload": {"token_usage": {"total": {}}}, + } + ) + assert recorder.turns == 2 + + def test_rejected_finish_item_is_not_promoted() -> None: from rpent.planner.codex import _Recorder