diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index 6f6913b71d0..77c79ed3abb 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -309,6 +309,7 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID - **`todos_remaining(*, looping_modes=None)`** / **`todos_remaining_message`** - Helper factories for todo-driven loops (the Python counterpart of .NET's `TodoCompletionLoopEvaluator`), designed for `create_harness_agent` but usable with any agent that registers a `TodoProvider` via `context_providers`. They resolve the `TodoProvider`/`AgentModeProvider` from the *running agent* (`agent.context_providers`, via `_resolve_context_provider`) rather than taking the provider as an argument, so they can be wired directly into `loop_should_continue`/`loop_next_message`. `todos_remaining` returns a `should_continue` predicate that loops while any todo is open; pass `looping_modes=[...]` to gate looping to specific operating modes (case-insensitive; honors the `AgentModeProvider`'s `source_id`/`available_modes`), `looping_modes=None` (default) applies in every mode, and an empty sequence raises `ValueError`. `todos_remaining_message` is a `next_message` callable that lists the still-open todo titles and tells the agent to finish them, returning `None` when the session/agent/provider is unavailable or nothing is open (in which case the middleware's default `None` handling applies: reuse the previous iteration's messages verbatim under the default `fresh_context=False`, or `DEFAULT_NEXT_MESSAGE` only when `fresh_context=True`). - **`background_tasks_running()`** / **`background_tasks_running_message`** - Helper factories for background-agent-driven loops, mirroring the `todos_remaining` pair. They resolve the `BackgroundAgentsProvider` from the *running agent* (`agent.context_providers`, via `_resolve_context_provider`) rather than taking the provider as an argument, so they can be wired directly into `create_harness_agent`'s `loop_should_continue`/`loop_next_message`. `background_tasks_running` returns a `should_continue` predicate that loops while the provider's persisted state shows any task with `status == RUNNING` (pair it with `max_iterations` so the loop is bounded even if a task's persisted status is never refreshed). `background_tasks_running_message` is a `next_message` callable that lists the still-running tasks (`# (): `) and tells the agent to wait for them to finish and retrieve their results, returning `None` when the session/agent/provider is unavailable or no task is running. - **Approval escape hatch** - `_has_pending_approval_request(result)` checks whether an iteration's response carries a pending tool-approval request (any content with `type == "function_approval_request"`). Both the streaming and non-streaming loops stop and return that response to the caller *before* evaluating `should_continue`/`max_iterations` or injecting `next_message`, so the loop is HITL-safe even when wrapped outermost around a `ToolApprovalMiddleware` (mirrors the C# `LoopAgent`'s `HasPendingApprovalRequests`). + - **Signal protocol** (`_harness/_loop.py`, alongside `AgentLoopMiddleware`) - `SignalParser` and `signal_should_continue()` provide a drop-in `should_continue` predicate that stops the loop when the agent emits `TASK_COMPLETE: ` (sets exit reason `"completed"`) or `NEED_INPUT: ` (sets exit reason `"need_input"`). A token counts only at the start of a line (leading whitespace allowed); mid-sentence or backtick-quoted mentions are ignored. `SignalParser` and `get_loop_exit_reason` share one detection helper (`_find_signal`). Only assistant-role messages are scanned; user-role nudges injected between iterations are ignored. `LoopExitReason` holds the typed string constants; the tokens are exposed as `SignalParser.TASK_COMPLETE_TOKEN` / `SignalParser.NEED_INPUT_TOKEN`. All four public symbols are `@experimental(feature_id=ExperimentalFeature.HARNESS)`. `get_loop_exit_reason(response)` scans assistant text first, then falls back to `response.additional_properties["loop_exit_reason"]` (unconditionally overwritten by the middleware when the iteration cap fires; non-string values are ignored). Signal tokens always win over the cap. The middleware stamps the cap reason using the shared `LOOP_EXIT_REASON_KEY` / `LoopExitReason.iteration_cap_reached` constants defined in the same module. - **Harness integration** - `create_harness_agent` enables the loop when a `loop_should_continue` callable is passed; it prepends `AgentLoopMiddleware(loop_should_continue, max_iterations=loop_max_iterations, next_message=loop_next_message)` ahead of `ToolApprovalMiddleware` so the loop is the outermost middleware (each iteration is a full agent run including tool approval, and the escape hatch hands pending approvals back to the caller). `loop_next_message` and `loop_max_iterations` only take effect together with `loop_should_continue` (with no `loop_should_continue` there is no loop, so they are ignored); `loop_max_iterations` defaults to the loop's default cap (`None` → unbounded). ### Workflows (`_workflows/`) diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index 3953187578e..43e2ab8c27d 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/python/packages/core/agent_framework/__init__.py @@ -134,8 +134,12 @@ "._harness._loop": ( "AgentLoopMiddleware", "JudgeVerdict", + "LoopExitReason", + "SignalParser", "background_tasks_running", "background_tasks_running_message", + "get_loop_exit_reason", + "signal_should_continue", "todos_remaining", "todos_remaining_message", ), @@ -554,6 +558,7 @@ "InlineSkillScript", "JudgeVerdict", "LocalEvaluator", + "LoopExitReason", "MCPSkill", "MCPSkillResource", "MCPSkillsSource", @@ -596,6 +601,7 @@ "ServiceSessionId", "SessionContext", "SessionStore", + "SignalParser", "SingleEdgeGroup", "Skill", "SkillFrontmatter", @@ -696,6 +702,7 @@ "executor", "function_middleware", "get_agent_mode", + "get_loop_exit_reason", "get_run_context", "handler", "included_messages", @@ -715,6 +722,7 @@ "resolve_agent_id", "response_handler", "set_agent_mode", + "signal_should_continue", "step", "todos_remaining", "todos_remaining_message", diff --git a/python/packages/core/agent_framework/__init__.pyi b/python/packages/core/agent_framework/__init__.pyi index e8d1bc6e411..862b42059a0 100644 --- a/python/packages/core/agent_framework/__init__.pyi +++ b/python/packages/core/agent_framework/__init__.pyi @@ -92,8 +92,12 @@ from ._harness._file_memory import DEFAULT_FILE_MEMORY_INSTRUCTIONS, DEFAULT_FIL from ._harness._loop import ( AgentLoopMiddleware, JudgeVerdict, + LoopExitReason, + SignalParser, background_tasks_running, background_tasks_running_message, + get_loop_exit_reason, + signal_should_continue, todos_remaining, todos_remaining_message, ) @@ -509,6 +513,7 @@ __all__ = [ "InlineSkillScript", "JudgeVerdict", "LocalEvaluator", + "LoopExitReason", "MCPSkill", "MCPSkillResource", "MCPSkillsSource", @@ -551,6 +556,7 @@ __all__ = [ "ServiceSessionId", "SessionContext", "SessionStore", + "SignalParser", "SingleEdgeGroup", "Skill", "SkillFrontmatter", @@ -651,6 +657,7 @@ __all__ = [ "executor", "function_middleware", "get_agent_mode", + "get_loop_exit_reason", "get_run_context", "handler", "included_messages", @@ -670,6 +677,7 @@ __all__ = [ "resolve_agent_id", "response_handler", "set_agent_mode", + "signal_should_continue", "step", "todos_remaining", "todos_remaining_message", diff --git a/python/packages/core/agent_framework/_harness/_loop.py b/python/packages/core/agent_framework/_harness/_loop.py index b4617a23e01..a0eec04271e 100644 --- a/python/packages/core/agent_framework/_harness/_loop.py +++ b/python/packages/core/agent_framework/_harness/_loop.py @@ -22,6 +22,16 @@ structured output; provider-specific formats can be supplied with a parser that converts the response into ``JudgeVerdict``. The loop continues while the answer is "no". This is a convenience wrapper that builds an async ``should_continue`` predicate, so it is a special case of (1). +3. Signal protocol (via :func:`~agent_framework.signal_should_continue`) - the + agent emits ``TASK_COMPLETE: `` or ``NEED_INPUT: `` at the start of a line in + its response text to signal terminal intent (leading whitespace is allowed; mid-sentence or + backtick-quoted mentions are ignored). Only assistant-role messages are scanned; user-role + nudges injected between iterations do not trigger false positives. ``TASK_COMPLETE:`` is + checked first. Read the result with :func:`~agent_framework.get_loop_exit_reason`, which returns a + :class:`~agent_framework.LoopExitReason` string (``"completed"``, ``"iteration_cap_reached"``, + ``"need_input"``). When ``max_iterations`` fires the loop stamps + ``additional_properties["loop_exit_reason"] = "iteration_cap_reached"``; signal tokens always win + over the cap. In every case, the input for the next iteration is controlled by the ``next_message`` callable. """ @@ -29,6 +39,7 @@ from __future__ import annotations import inspect +import re from collections.abc import Awaitable, Callable, Mapping, Sequence from typing import TYPE_CHECKING, Any, TypeAlias @@ -59,8 +70,12 @@ __all__ = [ "AgentLoopMiddleware", "JudgeVerdict", + "LoopExitReason", + "SignalParser", "background_tasks_running", "background_tasks_running_message", + "get_loop_exit_reason", + "signal_should_continue", "todos_remaining", "todos_remaining_message", ] @@ -131,6 +146,58 @@ class JudgeVerdict(BaseModel): DEFAULT_JUDGE_MAX_ITERATIONS = 5 +# Text tokens the agent emits in its response to signal loop state. +# Case-sensitive: the agent is instructed to emit the exact tokens. +TASK_COMPLETE_TOKEN: str = "TASK_COMPLETE:" # ruff: ignore[hardcoded-password-string] +NEED_INPUT_TOKEN: str = "NEED_INPUT:" # ruff: ignore[hardcoded-password-string] + +# A token only counts when it starts a line (leading spaces/tabs allowed), so a mid-sentence or +# backtick-quoted mention such as "I'll emit `TASK_COMPLETE:` when done" is not a signal. +# The ``text`` group captures the rest of the token's line. +_TASK_COMPLETE_PATTERN = re.compile(rf"^[ \t]*{re.escape(TASK_COMPLETE_TOKEN)}(?P[^\n]*)", re.MULTILINE) +_NEED_INPUT_PATTERN = re.compile(rf"^[ \t]*{re.escape(NEED_INPUT_TOKEN)}(?P[^\n]*)", re.MULTILINE) + +# Key used in AgentResponse.additional_properties to surface the loop exit reason. +LOOP_EXIT_REASON_KEY: str = "loop_exit_reason" + + +@experimental(feature_id=ExperimentalFeature.HARNESS) +class LoopExitReason: + """Typed exit reason strings for AgentLoopMiddleware runs. + + Values are plain strings (not an enum) to match the ``FinishReason`` + convention in this codebase and remain JSON-transparent. + + Examples:: + + from agent_framework import get_loop_exit_reason, LoopExitReason + + exit_reason = get_loop_exit_reason(response) + if exit_reason == LoopExitReason.completed: + print("Task done!") + elif exit_reason == LoopExitReason.need_input: + # question is in response.text after NEED_INPUT: + """ + + completed: str = "completed" + """``TASK_COMPLETE:`` signal detected in the agent's response.""" + + iteration_cap_reached: str = "iteration_cap_reached" + """``max_iterations`` cap fired before any terminal signal was emitted.""" + + need_input: str = "need_input" + """``NEED_INPUT:`` signal detected; loop stopped to surface a question.""" + + cancelled: str = "cancelled" + """External cancellation (reserved for future use).""" + + token_budget_exceeded: str = "token_budget_exceeded" # ruff: ignore[hardcoded-password-string] + """Token budget (``max_tokens``) exhausted before the loop completed.""" + + time_budget_exceeded: str = "time_budget_exceeded" + """Wall-clock time budget (``max_duration``) exhausted before the loop completed.""" + + # A callable invoked between iterations. It always receives the loop keyword arguments # (``iteration``, ``last_result``, ``messages``, ``original_messages``, ``session``, ``agent``, # ``progress``, ``feedback``). Callers declare only the keywords they need plus ``**kwargs`` to @@ -538,6 +605,7 @@ async def _process_non_streaming( snapshot: dict[str, Any] | None, ) -> None: iteration = 0 + cap_fired = False work_iterations = 0 progress: list[str] = [] # Aggregated transcript across iterations: each iteration's response messages plus the @@ -584,6 +652,8 @@ async def _process_non_streaming( # Decide whether to stop and capture any feedback from should_continue first, so the # feedback is available to both the progress and next-message callables this iteration. stop, feedback = await self._evaluate_stop(loop_kwargs, work_iterations) + if stop and self.max_iterations is not None and work_iterations >= self.max_iterations: + cap_fired = True loop_kwargs = self._build_loop_kwargs( context=context, iteration=iteration, @@ -617,6 +687,9 @@ async def _process_non_streaming( finally: context.options.pop(_LOOP_ITERATION_TOKEN_KEY, None) + if cap_fired: + final_result.additional_properties[LOOP_EXIT_REASON_KEY] = LoopExitReason.iteration_cap_reached + if not self.return_final_only: context.result = self._aggregate_response(final_result, aggregated, aggregated_usage) await self._fire_turn_scoped_after_providers( @@ -638,6 +711,7 @@ def _process_streaming( async def _generator() -> Any: iteration = 0 + cap_fired = False work_iterations = 0 progress: list[str] = [] stamped_options = dict(context.options) if context.options is not None else {} @@ -686,6 +760,8 @@ async def _generator() -> Any: # Decide whether to stop and capture any feedback from should_continue first, so the # feedback is available to both the progress and next-message callables this iteration. stop, feedback = await self._evaluate_stop(loop_kwargs, work_iterations) + if stop and self.max_iterations is not None and work_iterations >= self.max_iterations: + cap_fired = True loop_kwargs = self._build_loop_kwargs( context=context, iteration=iteration, @@ -721,6 +797,8 @@ async def _generator() -> Any: yield self._message_to_update(message) finally: context.options.pop(_LOOP_ITERATION_TOKEN_KEY, None) + if cap_fired and holder["final"] is not None: + holder["final"].additional_properties[LOOP_EXIT_REASON_KEY] = LoopExitReason.iteration_cap_reached await self._fire_turn_scoped_after_providers(context, holder["final"], original_messages) def _finalize(updates: Sequence[AgentResponseUpdate]) -> AgentResponse: @@ -1046,3 +1124,129 @@ async def todos_remaining_message(*, session: Any = None, agent: Any = None, **k "Continue working through them now. Mark each todo complete as you finish it, and only stop " "once every todo item is complete." ) + + +@experimental(feature_id=ExperimentalFeature.HARNESS) +class SignalParser: + """A ``should_continue`` predicate that scans for signal tokens. + + Scans the agent's latest response text for :attr:`TASK_COMPLETE_TOKEN` + and :attr:`NEED_INPUT_TOKEN` and returns the appropriate continuation + decision for :class:`~agent_framework.AgentLoopMiddleware`. + + A token only counts as a signal when it starts a line in an assistant + message (leading spaces/tabs are allowed). A mention mid-sentence or after + other text on the same line, including markdown backticks, is ignored. + + ``TASK_COMPLETE:`` is checked before ``NEED_INPUT:``; if an agent emits + both, the loop treats the run as completed. + + The extracted text after the token (up to the first newline) is returned + as the feedback string so ``next_message`` and ``record_feedback`` + callbacks can reference it. + + Examples:: + + from agent_framework import AgentLoopMiddleware, signal_should_continue + + agent = Agent( + client=client, + middleware=[AgentLoopMiddleware(signal_should_continue())], + ) + """ + + TASK_COMPLETE_TOKEN: str = TASK_COMPLETE_TOKEN + """Token the agent emits to signal that the task is complete.""" + + NEED_INPUT_TOKEN: str = NEED_INPUT_TOKEN + """Token the agent emits to signal that it needs human input.""" + + def __call__(self, *, last_result: AgentResponse, **kwargs: Any) -> tuple[bool, str | None]: + """Return ``(False, extracted_text)`` on a terminal signal, ``(True, None)`` otherwise.""" + signal = _find_signal(last_result) + if signal is None: + return (True, None) + return (False, signal[1] or None) + + +def _find_signal(response: AgentResponse) -> tuple[str, str] | None: + """Find the terminal signal in the assistant-role messages of ``response``. + + This is the single detection rule shared by :class:`SignalParser` and + :func:`get_loop_exit_reason`: a token counts only when it starts a line (after + optional spaces/tabs) in an assistant-role message. ``TASK_COMPLETE:`` takes + precedence over ``NEED_INPUT:``. + + Returns: + ``(exit_reason, text_after_token)`` where the text is stripped and cut at the + end of the token's line, or ``None`` when no signal is present. + """ + text = "\n".join(m.text or "" for m in response.messages if m.role == "assistant") + for pattern, reason in ( + (_TASK_COMPLETE_PATTERN, LoopExitReason.completed), + (_NEED_INPUT_PATTERN, LoopExitReason.need_input), + ): + match = pattern.search(text) + if match is not None: + return (reason, match.group("text").strip()) + return None + + +@experimental(feature_id=ExperimentalFeature.HARNESS) +def signal_should_continue() -> SignalParser: + """Return a :class:`SignalParser` instance ready for :class:`~agent_framework.AgentLoopMiddleware`. + + Pass the returned callable as ``should_continue``:: + + middleware = AgentLoopMiddleware(signal_should_continue()) + """ + return SignalParser() + + +@experimental(feature_id=ExperimentalFeature.HARNESS) +def get_loop_exit_reason(response: AgentResponse) -> str | None: + """Return the typed exit reason for a completed loop run, or ``None`` if unknown. + + Scans the agent's own (assistant-role) response messages for signal tokens first + (a token counts only at the start of a line, the same rule :class:`SignalParser` uses), + then falls back to ``response.additional_properties["loop_exit_reason"]`` + (set by :class:`~agent_framework.AgentLoopMiddleware` when the iteration cap fires). + + Signal tokens take priority: an agent that emitted ``TASK_COMPLETE:`` on the final + allowed iteration is treated as completed, not cap-reached. + + Args: + response: The :class:`~agent_framework.AgentResponse` returned by the loop. + + Returns: + A :class:`LoopExitReason` string, or ``None`` when the reason cannot + be determined (e.g. the loop was stopped by a custom predicate with no + signal tokens in the response). + + Examples: + + .. code-block:: python + + from agent_framework import get_loop_exit_reason, LoopExitReason + + response = await agent.run(task) + match get_loop_exit_reason(response): + case LoopExitReason.completed: + print("Done:", response.text) + case LoopExitReason.need_input: + # Question is in response.text after NEED_INPUT: + print("Agent needs input; see response.text for the question.") + case LoopExitReason.iteration_cap_reached: + print("Hit cap - partial result") + case _: + print("Unknown exit reason") + """ + # Scan only assistant messages - the agent's own output. User-role nudges injected + # by the loop between iterations must not trigger a false completion signal. + signal = _find_signal(response) + if signal is not None: + return signal[0] + # Fall back to additional_properties for framework-set reasons (e.g. iteration cap). + # additional_properties accepts arbitrary values, so only return non-empty strings. + reason = response.additional_properties.get(LOOP_EXIT_REASON_KEY) + return reason if isinstance(reason, str) and reason else None diff --git a/python/packages/core/tests/core/test_harness_loop.py b/python/packages/core/tests/core/test_harness_loop.py index 3b87544d804..21a7bce0947 100644 --- a/python/packages/core/tests/core/test_harness_loop.py +++ b/python/packages/core/tests/core/test_harness_loop.py @@ -24,14 +24,18 @@ Content, ContextProvider, JudgeVerdict, + LoopExitReason, Message, MiddlewareTermination, ResponseStream, + SignalParser, TodoItem, TodoProvider, background_tasks_running, background_tasks_running_message, + get_loop_exit_reason, set_agent_mode, + signal_should_continue, todos_remaining, todos_remaining_message, tool, @@ -40,6 +44,7 @@ DEFAULT_JUDGE_MAX_ITERATIONS, DEFAULT_MAX_ITERATIONS, DEFAULT_NEXT_MESSAGE, + LOOP_EXIT_REASON_KEY, AgentLoopMiddleware, ) @@ -1697,3 +1702,331 @@ async def test_concurrent_same_agent_run_during_stream_pause_is_not_suppressed() # once for the loop boundary, once for the concurrent run's own boundary assert turn_scoped.after_calls == 2 + + +# --- Signal Protocol Tests --- + + +def _make_signal_response(text: str, additional_properties: dict[str, Any] | None = None) -> AgentResponse: + return AgentResponse( + messages=[Message(role="assistant", contents=[text])], + additional_properties=additional_properties, + ) + + +def test_signal_parser_task_complete_stops_and_extracts_summary() -> None: + response = _make_signal_response("TASK_COMPLETE: finished the job") + assert SignalParser()(last_result=response) == (False, "finished the job") + + +def test_signal_parser_need_input_stops_and_extracts_question() -> None: + response = _make_signal_response("NEED_INPUT: which file should I edit?") + assert SignalParser()(last_result=response) == (False, "which file should I edit?") + + +def test_signal_parser_no_signal_continues() -> None: + response = _make_signal_response("I am still working on it.") + assert SignalParser()(last_result=response) == (True, None) + + +def test_signal_parser_task_complete_wins_over_need_input() -> None: + response = _make_signal_response("TASK_COMPLETE: done\nNEED_INPUT: are you sure?") + assert SignalParser()(last_result=response) == (False, "done") + + +def test_signal_parser_extracts_only_first_line_after_token() -> None: + response = _make_signal_response("Some preamble\nTASK_COMPLETE: line1\nmore stuff after") + assert SignalParser()(last_result=response) == (False, "line1") + + +def test_signal_parser_empty_response_continues() -> None: + response = AgentResponse(messages=[]) + assert SignalParser()(last_result=response) == (True, None) + + +def test_signal_parser_token_with_no_following_text_returns_none_feedback() -> None: + response = _make_signal_response("TASK_COMPLETE:") + assert SignalParser()(last_result=response) == (False, None) + + +def test_get_loop_exit_reason_completed_from_text() -> None: + response = _make_signal_response("TASK_COMPLETE: done") + assert get_loop_exit_reason(response) == LoopExitReason.completed + + +def test_get_loop_exit_reason_need_input_from_text() -> None: + response = _make_signal_response("NEED_INPUT: which approach?") + assert get_loop_exit_reason(response) == LoopExitReason.need_input + + +def test_get_loop_exit_reason_cap_from_additional_properties() -> None: + response = _make_signal_response( + "still going", additional_properties={LOOP_EXIT_REASON_KEY: "iteration_cap_reached"} + ) + assert get_loop_exit_reason(response) == LoopExitReason.iteration_cap_reached + + +def test_get_loop_exit_reason_signal_wins_over_additional_properties() -> None: + response = _make_signal_response( + "TASK_COMPLETE: done", additional_properties={LOOP_EXIT_REASON_KEY: "iteration_cap_reached"} + ) + assert get_loop_exit_reason(response) == LoopExitReason.completed + + +def test_get_loop_exit_reason_unknown_returns_none() -> None: + response = _make_signal_response("nothing special") + assert get_loop_exit_reason(response) is None + + +@pytest.mark.parametrize("value", [1, True, ["iteration_cap_reached"], {"reason": "x"}, ""]) +def test_get_loop_exit_reason_ignores_non_string_additional_property(value: Any) -> None: + response = _make_signal_response("nothing special", additional_properties={LOOP_EXIT_REASON_KEY: value}) + assert get_loop_exit_reason(response) is None + + +def test_signal_parser_exposes_tokens() -> None: + assert SignalParser.TASK_COMPLETE_TOKEN == "TASK_COMPLETE:" + assert SignalParser.NEED_INPUT_TOKEN == "NEED_INPUT:" + + +@pytest.mark.parametrize("obj", [LoopExitReason, SignalParser, signal_should_continue, get_loop_exit_reason]) +def test_signal_protocol_symbols_are_experimental(obj: Any) -> None: + assert obj.__feature_stage__ == "experimental" + assert obj.__feature_id__ == "HARNESS" + + +def test_signal_should_continue_factory_returns_signal_parser() -> None: + assert isinstance(signal_should_continue(), SignalParser) + + +async def test_loop_stops_on_task_complete_signal() -> None: + client = RecordingChatClient(texts=["TASK_COMPLETE: all done"]) + agent = Agent(client=client, middleware=[AgentLoopMiddleware(signal_should_continue(), max_iterations=10)]) + + response = await agent.run("task") + + assert client.call_count == 1 + assert get_loop_exit_reason(response) == LoopExitReason.completed + assert LOOP_EXIT_REASON_KEY not in response.additional_properties + + +async def test_loop_stops_on_need_input_signal() -> None: + client = RecordingChatClient(texts=["NEED_INPUT: which approach?"]) + agent = Agent(client=client, middleware=[AgentLoopMiddleware(signal_should_continue(), max_iterations=10)]) + + response = await agent.run("task") + + assert client.call_count == 1 + assert get_loop_exit_reason(response) == LoopExitReason.need_input + assert LOOP_EXIT_REASON_KEY not in response.additional_properties + + +async def test_loop_sets_iteration_cap_reason_when_no_signal() -> None: + client = RecordingChatClient() + agent = Agent(client=client, middleware=[AgentLoopMiddleware(signal_should_continue(), max_iterations=2)]) + + response = await agent.run("task") + + assert client.call_count == 2 + assert response.additional_properties.get(LOOP_EXIT_REASON_KEY) == LoopExitReason.iteration_cap_reached + assert get_loop_exit_reason(response) == LoopExitReason.iteration_cap_reached + + +async def test_loop_sets_iteration_cap_reason_streaming() -> None: + client = RecordingChatClient() + agent = Agent(client=client, middleware=[AgentLoopMiddleware(signal_should_continue(), max_iterations=2)]) + + stream = agent.run("task", stream=True) + async for _ in stream: + pass + response = await stream.get_final_response() + + assert client.call_count == 2 + assert response.additional_properties.get(LOOP_EXIT_REASON_KEY) == LoopExitReason.iteration_cap_reached + + +class StaleExitReasonChatClient(RecordingChatClient): + """Stamps a stale ``loop_exit_reason`` on every response, emulating an inner provider or nested loop.""" + + def _inner_get_response( + self, + *, + messages: Sequence[Message], + stream: bool = False, + options: Mapping[str, Any], + **kwargs: Any, + ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: + inner = super()._inner_get_response(messages=messages, stream=stream, options=options, **kwargs) + if isinstance(inner, ResponseStream): + + async def _gen() -> AsyncIterable[ChatResponseUpdate]: + async for update in inner: + update.additional_properties = {LOOP_EXIT_REASON_KEY: "stale_reason"} + yield update + + return ResponseStream(_gen(), finalizer=ChatResponse.from_updates) + + async def _get() -> ChatResponse: + response = await inner + response.additional_properties = {LOOP_EXIT_REASON_KEY: "stale_reason"} + return response + + return _get() + + +async def test_loop_cap_overwrites_pre_existing_exit_reason() -> None: + client = StaleExitReasonChatClient() + agent = Agent(client=client, middleware=[AgentLoopMiddleware(signal_should_continue(), max_iterations=2)]) + + response = await agent.run("task") + + assert client.call_count == 2 + assert response.additional_properties.get(LOOP_EXIT_REASON_KEY) == LoopExitReason.iteration_cap_reached + assert get_loop_exit_reason(response) == LoopExitReason.iteration_cap_reached + + +async def test_loop_cap_overwrites_pre_existing_exit_reason_streaming() -> None: + client = StaleExitReasonChatClient() + agent = Agent(client=client, middleware=[AgentLoopMiddleware(signal_should_continue(), max_iterations=2)]) + + stream = agent.run("task", stream=True) + async for _ in stream: + pass + response = await stream.get_final_response() + + assert client.call_count == 2 + assert response.additional_properties.get(LOOP_EXIT_REASON_KEY) == LoopExitReason.iteration_cap_reached + assert get_loop_exit_reason(response) == LoopExitReason.iteration_cap_reached + + +async def test_get_loop_exit_reason_ignores_nudge_text_with_token() -> None: + # A next_message nudge that mentions TASK_COMPLETE: must not trigger a false positive. + # Scripted texts keep the agent from echoing the nudge back in its own messages. + client = RecordingChatClient(texts=["working...", "still working..."]) + agent = Agent( + client=client, + middleware=[ + AgentLoopMiddleware( + signal_should_continue(), + max_iterations=2, + next_message=lambda **kwargs: "Reminder: emit TASK_COMPLETE: when done.", + ) + ], + ) + + response = await agent.run("task") + + assert client.call_count == 2 + assert response.additional_properties.get(LOOP_EXIT_REASON_KEY) == LoopExitReason.iteration_cap_reached + assert get_loop_exit_reason(response) == LoopExitReason.iteration_cap_reached + + +async def test_loop_continues_when_no_signal() -> None: + client = RecordingChatClient(texts=["working...", "still working...", "TASK_COMPLETE: finally done"]) + agent = Agent(client=client, middleware=[AgentLoopMiddleware(signal_should_continue(), max_iterations=10)]) + + response = await agent.run("task") + + assert client.call_count == 3 + assert get_loop_exit_reason(response) == LoopExitReason.completed + + +# Signal tokens only count at the start of a line (leading whitespace allowed). +_NON_SIGNAL_MENTIONS = [ + "I'll emit TASK_COMPLETE: when done.", + "I'll emit `TASK_COMPLETE:` when done.", + "I'll ask with NEED_INPUT: if I get stuck.", + "I'll emit `NEED_INPUT:` if I get stuck.", + "Working...\nWhen finished I will print `TASK_COMPLETE: summary`.", +] + + +@pytest.mark.parametrize("text", _NON_SIGNAL_MENTIONS) +def test_signal_parser_ignores_mid_line_mentions(text: str) -> None: + response = _make_signal_response(text) + assert SignalParser()(last_result=response) == (True, None) + assert get_loop_exit_reason(response) is None + + +@pytest.mark.parametrize( + ("text", "expected_reason", "expected_feedback"), + [ + ("Here is my analysis.\nTASK_COMPLETE: wrote the report", LoopExitReason.completed, "wrote the report"), + ("I looked at both options.\nNEED_INPUT: which one?", LoopExitReason.need_input, "which one?"), + ("Mentioning `TASK_COMPLETE:` here.\nTASK_COMPLETE: real", LoopExitReason.completed, "real"), + ], +) +def test_signal_parser_detects_token_at_start_of_later_line( + text: str, expected_reason: str, expected_feedback: str +) -> None: + response = _make_signal_response(text) + assert SignalParser()(last_result=response) == (False, expected_feedback) + assert get_loop_exit_reason(response) == expected_reason + + +@pytest.mark.parametrize( + ("text", "expected_reason", "expected_feedback"), + [ + (" TASK_COMPLETE: indented", LoopExitReason.completed, "indented"), + ("prose\n\tNEED_INPUT: tabbed", LoopExitReason.need_input, "tabbed"), + ], +) +def test_signal_parser_allows_leading_whitespace(text: str, expected_reason: str, expected_feedback: str) -> None: + response = _make_signal_response(text) + assert SignalParser()(last_result=response) == (False, expected_feedback) + assert get_loop_exit_reason(response) == expected_reason + + +def test_signal_parser_mid_line_task_complete_does_not_shadow_need_input() -> None: + response = _make_signal_response("I'll say `TASK_COMPLETE:` later.\nNEED_INPUT: which file?") + assert SignalParser()(last_result=response) == (False, "which file?") + assert get_loop_exit_reason(response) == LoopExitReason.need_input + + +@pytest.mark.parametrize("text", _NON_SIGNAL_MENTIONS) +async def test_loop_ignores_mid_line_signal_mentions(text: str) -> None: + client = RecordingChatClient(texts=[text, text]) + agent = Agent(client=client, middleware=[AgentLoopMiddleware(signal_should_continue(), max_iterations=2)]) + + response = await agent.run("task") + + assert client.call_count == 2 + assert get_loop_exit_reason(response) == LoopExitReason.iteration_cap_reached + + +@pytest.mark.parametrize("text", _NON_SIGNAL_MENTIONS) +async def test_loop_ignores_mid_line_signal_mentions_streaming(text: str) -> None: + client = RecordingChatClient(texts=[text, text]) + agent = Agent(client=client, middleware=[AgentLoopMiddleware(signal_should_continue(), max_iterations=2)]) + + stream = agent.run("task", stream=True) + async for _ in stream: + pass + response = await stream.get_final_response() + + assert client.call_count == 2 + assert get_loop_exit_reason(response) == LoopExitReason.iteration_cap_reached + + +async def test_loop_stops_on_task_complete_at_start_of_later_line() -> None: + client = RecordingChatClient(texts=["Summary of work.\n TASK_COMPLETE: done"]) + agent = Agent(client=client, middleware=[AgentLoopMiddleware(signal_should_continue(), max_iterations=10)]) + + response = await agent.run("task") + + assert client.call_count == 1 + assert get_loop_exit_reason(response) == LoopExitReason.completed + assert LOOP_EXIT_REASON_KEY not in response.additional_properties + + +async def test_loop_stops_on_need_input_at_start_of_later_line_streaming() -> None: + client = RecordingChatClient(texts=["I compared the options.\nNEED_INPUT: which one?"]) + agent = Agent(client=client, middleware=[AgentLoopMiddleware(signal_should_continue(), max_iterations=10)]) + + stream = agent.run("task", stream=True) + async for _ in stream: + pass + response = await stream.get_final_response() + + assert client.call_count == 1 + assert get_loop_exit_reason(response) == LoopExitReason.need_input