From 27d869d3ccc87f069c7fae6b6ab516b52b099a89 Mon Sep 17 00:00:00 2001 From: chengyongru Date: Tue, 16 Jun 2026 10:21:28 +0800 Subject: [PATCH] fix(agent): refresh goal continuation context --- nanobot/agent/loop.py | 22 ++++++++------- nanobot/agent/runner.py | 16 +++++++++-- tests/agent/test_loop_save_turn.py | 35 ++++++++++++++++++++++++ tests/agent/test_runner_goal_continue.py | 34 +++++++++++++++++++++++ 4 files changed, 95 insertions(+), 12 deletions(-) diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index de649601..f9bbefaf 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -65,7 +65,6 @@ from nanobot.utils.image_generation_intent import image_generation_prompt from nanobot.utils.llm_runtime import LLMRuntime from nanobot.utils.runtime import ( EMPTY_FINAL_RESPONSE_MESSAGE, - SUSTAINED_GOAL_CONTINUE_PROMPT, ) if TYPE_CHECKING: @@ -796,15 +795,18 @@ class AgentLoop: file_state_token = bind_file_states(self._file_state_store.for_session(active_session_key)) request_token = bind_request_context(request_ctx) workspace_token = bind_workspace_scope(effective_scope) - # Build continuation message that embeds the active goal objective so - # the LLM can see it even if earlier Runtime Context was truncated. - _goal_lines = goal_state_runtime_lines(session.metadata if session is not None else None) - _goal_continue = ( - "You have an active sustained goal:\n\n" - + "\n".join(_goal_lines) - + "\n\nPlease continue working toward the objective using your tools, " - "or call complete_goal if the work is truly finished." - ) if _goal_lines else SUSTAINED_GOAL_CONTINUE_PROMPT + # Compute lazily because long_task may create goal metadata during this run. + def _goal_continue() -> str | None: + _goal_lines = goal_state_runtime_lines(session.metadata if session is not None else None) + if not _goal_lines: + return None + return ( + "You have an active sustained goal:\n\n" + + "\n".join(_goal_lines) + + "\n\nPlease continue working toward the objective using your tools, " + "or call complete_goal if the work is truly finished." + ) + session_metadata = session.metadata if session is not None else None try: result = await self.runner.run(AgentRunSpec( diff --git a/nanobot/agent/runner.py b/nanobot/agent/runner.py index e5cc579f..e3398ecb 100644 --- a/nanobot/agent/runner.py +++ b/nanobot/agent/runner.py @@ -54,6 +54,8 @@ from nanobot.utils.runtime import ( repeated_workspace_violation_error, ) +GoalContinueMessage = str | Callable[[], str | None] + _DEFAULT_ERROR_MESSAGE = "Sorry, I encountered an error calling the AI model." _ARREARAGE_ERROR_MESSAGE = ( "The AI provider rejected the request because the API key is out of quota or the " @@ -109,7 +111,7 @@ class AgentRunSpec: injection_callback: Any | None = None llm_timeout_s: float | None = None goal_active_predicate: Callable[[], bool] | None = None - goal_continue_message: str | None = None + goal_continue_message: GoalContinueMessage | None = None finalize_on_max_iterations: bool = True @@ -198,7 +200,7 @@ class AgentRunner: if not injections and allow_goal_continue and assistant_message is not None: predicate = spec.goal_active_predicate if predicate is not None and predicate(): - injections = [build_goal_continue_message(spec.goal_continue_message)] + injections = [self._build_goal_continue_message(spec)] if not injections: return False, injection_cycles if real_injection: @@ -227,6 +229,16 @@ class AgentRunner: logger.info("Injected sustained-goal continuation {}", phase) return True, injection_cycles + def _build_goal_continue_message(self, spec: AgentRunSpec) -> dict[str, str]: + custom = spec.goal_continue_message + if callable(custom): + try: + custom = custom() + except Exception: + logger.exception("goal_continue_message callback failed") + custom = None + return build_goal_continue_message(custom) + async def _drain_injections(self, spec: AgentRunSpec) -> list[dict[str, Any]]: """Drain pending user messages via the injection callback. diff --git a/tests/agent/test_loop_save_turn.py b/tests/agent/test_loop_save_turn.py index f314fea4..dfd530c9 100644 --- a/tests/agent/test_loop_save_turn.py +++ b/tests/agent/test_loop_save_turn.py @@ -953,6 +953,41 @@ async def test_process_message_uses_explicit_session_metadata_for_goal_context( assert GOAL_STATE_KEY not in kwargs["session_metadata"] +@pytest.mark.asyncio +async def test_run_agent_loop_goal_continue_message_reads_latest_metadata( + tmp_path: Path, +) -> None: + from nanobot.agent.runner import AgentRunResult + + loop = _make_full_loop(tmp_path) + session = loop.sessions.get_or_create("websocket:late-goal") + seen: dict[str, str | None] = {} + + async def fake_run(spec): + assert callable(spec.goal_continue_message) + session.metadata[GOAL_STATE_KEY] = { + "status": "active", + "objective": "Goal created during this runner call.", + } + seen["goal_continue"] = spec.goal_continue_message() + return AgentRunResult( + final_content="ok", + messages=[{"role": "assistant", "content": "ok"}], + ) + + loop.runner.run = fake_run # type: ignore[method-assign] + + await loop._run_agent_loop( + [], + session=session, + channel="websocket", + chat_id="late-goal", + session_key=session.key, + ) + + assert "Goal created during this runner call." in (seen["goal_continue"] or "") + + def test_set_tool_context_uses_effective_key_for_spawn_tool(tmp_path: Path) -> None: loop = _make_full_loop(tmp_path) spawn_tool = loop.tools.get("spawn") diff --git a/tests/agent/test_runner_goal_continue.py b/tests/agent/test_runner_goal_continue.py index e5aec92f..cca146a0 100644 --- a/tests/agent/test_runner_goal_continue.py +++ b/tests/agent/test_runner_goal_continue.py @@ -210,3 +210,37 @@ async def test_runner_uses_custom_goal_continue_message(): user_msgs = [m for m in result.messages if m.get("role") == "user"] assert any(custom_msg in str(m.get("content", "")) for m in user_msgs) + + +@pytest.mark.asyncio +async def test_runner_resolves_goal_continue_message_lazily(): + """The continuation text can depend on goal metadata created during the run.""" + from nanobot.agent.runner import AgentRunner, AgentRunSpec + + provider = MagicMock(spec=LLMProvider) + provider.chat_with_retry = AsyncMock(return_value=LLMResponse( + content="still working", tool_calls=[], usage={}, + )) + tools = MagicMock() + tools.get_definitions.return_value = [] + calls = {"n": 0} + + def dynamic_msg() -> str: + calls["n"] += 1 + return "Goal (active):\nWrite the article draft." + + runner = AgentRunner(provider) + result = await runner.run(AgentRunSpec( + initial_messages=[{"role": "user", "content": "do task"}], + tools=tools, + model="test-model", + max_iterations=1, + max_tool_result_chars=_MAX_TOOL_RESULT_CHARS, + goal_active_predicate=lambda: True, + goal_continue_message=dynamic_msg, + finalize_on_max_iterations=False, + )) + + user_msgs = [m for m in result.messages if m.get("role") == "user"] + assert calls["n"] == 1 + assert any("Write the article draft." in str(m.get("content", "")) for m in user_msgs)