fix(agent): extend sustained goal iteration budget

tests/agent/test_loop_runner_integration.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import asyncio
 import time
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -48,6 +47,30 @@ async def test_loop_max_iterations_message_stays_stable(tmp_path):
     )
 
 
+@pytest.mark.asyncio
+async def test_loop_goal_turn_uses_standard_iteration_budget(tmp_path):
+    loop = _make_loop(tmp_path)
+    loop.provider.chat_with_retry = AsyncMock(return_value=LLMResponse(
+        content="working",
+        tool_calls=[ToolCallRequest(id="call_1", name="list_dir", arguments={})],
+    ))
+    loop.tools.get_definitions = MagicMock(return_value=[])
+    loop.tools.execute = AsyncMock(return_value="ok")
+    loop.max_iterations = 2
+
+    final_content, _, _, stop_reason, _ = await loop._run_agent_loop(
+        [],
+        metadata={"original_command": "/goal"},
+    )
+
+    assert stop_reason == "max_iterations"
+    assert loop.provider.chat_with_retry.await_count == 2
+    assert final_content == (
+        "I reached the maximum number of tool call iterations (2) "
+        "without completing the task. You can try breaking the task into smaller steps."
+    )
+
+
 @pytest.mark.asyncio
 async def test_loop_stream_filter_handles_think_only_prefix_without_crashing(tmp_path):
     loop = _make_loop(tmp_path)

tests/agent/test_loop_save_turn.py

@@ -11,6 +11,10 @@ from nanobot.bus.queue import MessageBus
 from nanobot.providers.base import LLMResponse
 from nanobot.session.goal_state import GOAL_STATE_KEY
 from nanobot.session.manager import Session, SessionManager
+from nanobot.session.turn_continuation import (
+    INTERNAL_CONTINUATION_META,
+    INTERNAL_CONTINUATION_RUN_STARTED_AT_META,
+)
 from nanobot.session.webui_turns import (
     TITLE_GENERATION_MAX_TOKENS,
     TITLE_GENERATION_REASONING_EFFORT,
@@ -560,6 +564,226 @@ async def test_process_message_does_not_duplicate_early_persisted_user_message(t
     assert AgentLoop._PENDING_USER_TURN_KEY not in session.metadata
 
 
+@pytest.mark.asyncio
+async def test_internal_continuation_queues_turn_without_fake_user_history(
+    tmp_path: Path,
+) -> None:
+    loop = _make_full_loop(tmp_path)
+    loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False)  # type: ignore[method-assign]
+    session = loop.sessions.get_or_create("feishu:c-auto")
+    session.metadata[GOAL_STATE_KEY] = {
+        "status": "active",
+        "objective": "Finish the long goal.",
+    }
+    loop.sessions.save(session)
+
+    calls: list[dict] = []
+
+    async def fake_run_agent_loop(initial_messages, *, metadata=None, **_kwargs):
+        calls.append({"initial_messages": initial_messages, "metadata": metadata})
+        if len(calls) == 1:
+            return (
+                "paused",
+                [],
+                [*initial_messages, {"role": "assistant", "content": "paused"}],
+                "max_iterations",
+                False,
+            )
+        return (
+            "done",
+            [],
+            [*initial_messages, {"role": "assistant", "content": "done"}],
+            "completed",
+            False,
+        )
+
+    loop._run_agent_loop = fake_run_agent_loop  # type: ignore[method-assign]
+    pending: asyncio.Queue[InboundMessage] = asyncio.Queue()
+
+    first = await loop._process_message(
+        InboundMessage(
+            channel="feishu",
+            sender_id="u1",
+            chat_id="c-auto",
+            content="start the goal",
+        ),
+        pending_queue=pending,
+    )
+
+    assert first is None
+    queued = pending.get_nowait()
+    assert queued.sender_id == "system:continuation"
+    assert queued.metadata[INTERNAL_CONTINUATION_META] is True
+    assert "Finish the long goal." in queued.content
+
+    session = loop.sessions.get_or_create("feishu:c-auto")
+    assert [
+        {k: v for k, v in m.items() if k in {"role", "content"}}
+        for m in session.messages
+    ] == [{"role": "user", "content": "start the goal"}]
+
+    second = await loop._process_message(queued, pending_queue=asyncio.Queue())
+
+    assert second is not None
+    assert second.content == "done"
+    session = loop.sessions.get_or_create("feishu:c-auto")
+    assert [
+        {k: v for k, v in m.items() if k in {"role", "content"}}
+        for m in session.messages
+    ] == [
+        {"role": "user", "content": "start the goal"},
+        {"role": "assistant", "content": "done"},
+    ]
+
+
+@pytest.mark.asyncio
+async def test_internal_continuation_preserves_streaming_route_metadata(
+    tmp_path: Path,
+) -> None:
+    loop = _make_full_loop(tmp_path)
+    loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False)  # type: ignore[method-assign]
+    session = loop.sessions.get_or_create("feishu:c-stream")
+    session.metadata[GOAL_STATE_KEY] = {
+        "status": "active",
+        "objective": "Finish the streamed long goal.",
+    }
+    loop.sessions.save(session)
+
+    calls = 0
+
+    async def fake_run_agent_loop(initial_messages, *, on_stream=None, on_stream_end=None, **_kwargs):
+        nonlocal calls
+        calls += 1
+        if calls == 1:
+            return (
+                "paused",
+                [],
+                [*initial_messages, {"role": "assistant", "content": "paused"}],
+                "max_iterations",
+                False,
+            )
+        assert on_stream is not None
+        assert on_stream_end is not None
+        await on_stream("done")
+        await on_stream_end(resuming=False)
+        return (
+            "done",
+            [],
+            [*initial_messages, {"role": "assistant", "content": "done"}],
+            "completed",
+            False,
+        )
+
+    loop._run_agent_loop = fake_run_agent_loop  # type: ignore[method-assign]
+
+    await loop._dispatch(InboundMessage(
+        channel="feishu",
+        sender_id="u1",
+        chat_id="c-stream",
+        content="start the goal",
+        metadata={
+            "_wants_stream": True,
+            "message_id": "om_001",
+            "origin_message_id": "root_001",
+            "_stream_id": "old-stream",
+        },
+    ))
+
+    assert loop.bus.outbound_size == 0
+    queued = await asyncio.wait_for(loop.bus.consume_inbound(), timeout=0.5)
+    assert queued.metadata[INTERNAL_CONTINUATION_META] is True
+    assert queued.metadata["_wants_stream"] is True
+    assert queued.metadata["message_id"] == "om_001"
+    assert queued.metadata["origin_message_id"] == "root_001"
+    assert "_stream_id" not in queued.metadata
+
+    await loop._dispatch(queued)
+
+    outbound = []
+    while loop.bus.outbound_size:
+        outbound.append(await loop.bus.consume_outbound())
+    deltas = [m for m in outbound if m.metadata.get("_stream_delta")]
+    ends = [m for m in outbound if m.metadata.get("_stream_end")]
+    streamed_markers = [m for m in outbound if m.metadata.get("_streamed")]
+
+    assert [m.content for m in deltas] == ["done"]
+    assert len(ends) == 1
+    assert ends[0].metadata["_resuming"] is False
+    assert ends[0].metadata["message_id"] == "om_001"
+    assert ends[0].metadata["origin_message_id"] == "root_001"
+    assert isinstance(ends[0].metadata.get("_stream_id"), str)
+    assert streamed_markers and streamed_markers[-1].content == "done"
+
+
+@pytest.mark.asyncio
+async def test_websocket_internal_continuation_keeps_single_visible_run(
+    tmp_path: Path,
+) -> None:
+    loop = _make_full_loop(tmp_path)
+    loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False)  # type: ignore[method-assign]
+    session = loop.sessions.get_or_create("websocket:c-auto")
+    session.metadata[GOAL_STATE_KEY] = {
+        "status": "active",
+        "objective": "Finish the long goal.",
+    }
+    loop.sessions.save(session)
+
+    calls = 0
+
+    async def fake_run_agent_loop(initial_messages, **_kwargs):
+        nonlocal calls
+        calls += 1
+        if calls == 1:
+            return (
+                "paused",
+                [],
+                [*initial_messages, {"role": "assistant", "content": "paused"}],
+                "max_iterations",
+                False,
+            )
+        return (
+            "done",
+            [],
+            [*initial_messages, {"role": "assistant", "content": "done"}],
+            "completed",
+            False,
+        )
+
+    loop._run_agent_loop = fake_run_agent_loop  # type: ignore[method-assign]
+
+    await loop._dispatch(InboundMessage(
+        channel="websocket",
+        sender_id="u1",
+        chat_id="c-auto",
+        content="start the goal",
+        metadata={"webui": True},
+    ))
+
+    first_outbound = []
+    while loop.bus.outbound_size:
+        first_outbound.append(await loop.bus.consume_outbound())
+    first_statuses = [m.metadata for m in first_outbound if m.metadata.get("_goal_status")]
+    assert [m["goal_status"] for m in first_statuses] == ["running"]
+    assert not [m for m in first_outbound if m.metadata.get("_turn_end")]
+    started_at = first_statuses[0]["started_at"]
+
+    queued = await asyncio.wait_for(loop.bus.consume_inbound(), timeout=0.5)
+    assert queued.metadata[INTERNAL_CONTINUATION_META] is True
+    assert queued.metadata[INTERNAL_CONTINUATION_RUN_STARTED_AT_META] == started_at
+
+    await loop._dispatch(queued)
+
+    second_outbound = []
+    while loop.bus.outbound_size:
+        second_outbound.append(await loop.bus.consume_outbound())
+    second_statuses = [m.metadata for m in second_outbound if m.metadata.get("_goal_status")]
+    assert [m["goal_status"] for m in second_statuses] == ["running", "idle"]
+    assert second_statuses[0]["started_at"] == started_at
+    turn_end = [m for m in second_outbound if m.metadata.get("_turn_end")]
+    assert len(turn_end) == 1
+    assert isinstance(turn_end[0].metadata.get("latency_ms"), int)
+
+
 @pytest.mark.asyncio
 async def test_process_message_uses_context_chat_id_for_runtime_prompt(tmp_path: Path) -> None:
     loop = _make_full_loop(tmp_path)

tests/session/test_turn_continuation.py

@@ -0,0 +1,127 @@
+"""Tests for internal turn continuation policy."""
+
+from __future__ import annotations
+
+import asyncio
+from types import SimpleNamespace
+
+import pytest
+
+from nanobot.bus.events import InboundMessage
+from nanobot.session.goal_state import GOAL_STATE_KEY
+from nanobot.session.turn_continuation import (
+    INTERNAL_CONTINUATION_KIND_META,
+    INTERNAL_CONTINUATION_META,
+    INTERNAL_CONTINUATION_PENDING_META,
+    INTERNAL_CONTINUATION_RUN_STARTED_AT_META,
+    internal_continuation_pending,
+    internal_continuation_run_started_at,
+    maybe_continue_turn,
+    should_stream_budget_response,
+)
+
+
+@pytest.mark.asyncio
+async def test_maybe_continue_turn_queues_internal_message():
+    meta = {
+        GOAL_STATE_KEY: {
+            "status": "active",
+            "objective": "Finish the migration.",
+            "ui_summary": "migration",
+        },
+    }
+    messages = [
+        {"role": "system", "content": "system"},
+        {"role": "user", "content": "start"},
+        {"role": "assistant", "content": "paused"},
+    ]
+    pending: asyncio.Queue[InboundMessage] = asyncio.Queue()
+    ctx = SimpleNamespace(
+        session=SimpleNamespace(metadata=meta),
+        msg=InboundMessage(
+            channel="feishu",
+            sender_id="u1",
+            chat_id="c1",
+            content="start",
+            metadata={
+                "message_id": "msg-1",
+                "origin_message_id": "msg-0",
+                "_wants_stream": True,
+                "_stream_id": "stream-1",
+                "_stream_delta": True,
+                "_stream_end": True,
+                "_resuming": True,
+                "webui": True,
+            },
+        ),
+        session_key="feishu:c1",
+        pending_queue=pending,
+        stop_reason="max_iterations",
+        final_content="paused",
+        all_messages=messages,
+        suppress_response=False,
+        visible_run_started_at=1234.5,
+    )
+
+    assert await maybe_continue_turn(ctx) is True
+
+    queued = pending.get_nowait()
+    assert queued.sender_id == "system:continuation"
+    assert queued.metadata[INTERNAL_CONTINUATION_META] is True
+    assert queued.metadata[INTERNAL_CONTINUATION_KIND_META] == "sustained_goal"
+    assert queued.metadata[INTERNAL_CONTINUATION_RUN_STARTED_AT_META] == 1234.5
+    assert internal_continuation_run_started_at(queued.metadata) == 1234.5
+    assert internal_continuation_pending(ctx.msg.metadata)
+    assert queued.metadata["webui"] is True
+    assert queued.metadata["message_id"] == "msg-1"
+    assert queued.metadata["origin_message_id"] == "msg-0"
+    assert queued.metadata["_wants_stream"] is True
+    assert "_stream_id" not in queued.metadata
+    assert "_stream_delta" not in queued.metadata
+    assert "_stream_end" not in queued.metadata
+    assert "_resuming" not in queued.metadata
+    assert "Finish the migration." in queued.content
+    assert ctx.all_messages == messages[:-1]
+    assert ctx.final_content == ""
+    assert ctx.suppress_response is True
+    assert ctx.msg.metadata[INTERNAL_CONTINUATION_PENDING_META] is True
+    assert meta["_sustained_goal_continuation_rounds"] == 1
+
+
+@pytest.mark.asyncio
+async def test_internal_continuation_respects_round_limit():
+    meta = {
+        GOAL_STATE_KEY: {"status": "active", "objective": "x"},
+        "_sustained_goal_continuation_rounds": 12,
+    }
+    ctx = SimpleNamespace(
+        session=SimpleNamespace(metadata=meta),
+        msg=InboundMessage(channel="feishu", sender_id="u1", chat_id="c1", content="start"),
+        session_key="feishu:c1",
+        pending_queue=asyncio.Queue(),
+        stop_reason="max_iterations",
+        final_content="paused",
+        all_messages=[],
+    )
+
+    assert should_stream_budget_response(
+        stop_reason="max_iterations",
+        pending_queue_available=True,
+        session_metadata=meta,
+    )
+    assert await maybe_continue_turn(ctx) is False
+
+
+def test_internal_continuation_requires_budget_boundary_and_queue():
+    meta = {GOAL_STATE_KEY: {"status": "active", "objective": "x"}}
+
+    assert should_stream_budget_response(
+        stop_reason="completed",
+        pending_queue_available=True,
+        session_metadata=meta,
+    )
+    assert should_stream_budget_response(
+        stop_reason="max_iterations",
+        pending_queue_available=False,
+        session_metadata=meta,
+    )

tests/utils/test_webui_turn_helpers.py

@@ -31,6 +31,19 @@ async def test_publish_turn_run_status_running_records_wall_clock() -> None:
     assert call.metadata.get("started_at") == t0
 
 
+@pytest.mark.asyncio
+async def test_publish_turn_run_status_reuses_explicit_wall_clock() -> None:
+    bus = MagicMock()
+    bus.publish_outbound = AsyncMock()
+    msg = InboundMessage(channel="websocket", sender_id="u", chat_id="chat-a", content="hi")
+
+    await wth.publish_turn_run_status(bus, msg, "running", started_at=1234.5)
+
+    assert wth.websocket_turn_wall_started_at("chat-a") == 1234.5
+    call = bus.publish_outbound.await_args[0][0]
+    assert call.metadata.get("started_at") == 1234.5
+
+
 @pytest.mark.asyncio
 async def test_publish_turn_run_status_idle_clears_wall_clock() -> None:
     bus = MagicMock()