feat(tools): tighten patch and session workflows

nanobot/agent/tools/exec_session.py

@@ -16,6 +16,8 @@ from nanobot.agent.tools.schema import BooleanSchema, IntegerSchema, StringSchem
 
 DEFAULT_YIELD_MS = 1000
 MAX_YIELD_MS = 30_000
+DEFAULT_WAIT_FOR_MS = 10_000
+MAX_WAIT_FOR_MS = 120_000
 DEFAULT_MAX_OUTPUT_CHARS = 10_000
 MAX_OUTPUT_CHARS = 50_000
 
@@ -356,6 +358,18 @@ def format_session_poll(session_id: str, poll: _SessionPoll) -> str:
             minimum=0,
             maximum=MAX_YIELD_MS,
         ),
+        wait_for=StringSchema(
+            "Optional text to wait for in output before returning. "
+            "Useful for interactive commands and dev servers.",
+            nullable=True,
+        ),
+        wait_timeout_ms=IntegerSchema(
+            DEFAULT_WAIT_FOR_MS,
+            description="Maximum milliseconds to wait for wait_for text (default 10000, max 120000).",
+            minimum=0,
+            maximum=MAX_WAIT_FOR_MS,
+            nullable=True,
+        ),
         max_output_chars=IntegerSchema(
             DEFAULT_MAX_OUTPUT_CHARS,
             description="Maximum output characters to return from this poll (default 10000, max 50000).",
@@ -412,8 +426,9 @@ class WriteStdinTool(Tool):
         return (
             "Write text to a running exec session and return recent output. "
             "Use chars='' to poll without writing. Set close_stdin=true to send EOF, "
-            "or terminate=true to stop the session. Sessions finish automatically "
-            "when their process exits."
+            "or terminate=true to stop the session. Use wait_for to keep polling "
+            "until expected output appears. Sessions finish automatically when "
+            "their process exits."
         )
 
     async def execute(
@@ -423,6 +438,8 @@ class WriteStdinTool(Tool):
         close_stdin: bool = False,
         terminate: bool = False,
         yield_time_ms: int | None = None,
+        wait_for: str | None = None,
+        wait_timeout_ms: int | None = None,
         max_output_chars: int | None = None,
         max_output_tokens: int | None = None,
         **kwargs: Any,
@@ -430,18 +447,34 @@ class WriteStdinTool(Tool):
         try:
             if max_output_chars is None:
                 max_output_chars = max_output_tokens
+            output_limit = clamp_session_int(
+                max_output_chars,
+                DEFAULT_MAX_OUTPUT_CHARS,
+                1000,
+                MAX_OUTPUT_CHARS,
+            )
+            if wait_for:
+                return await self._wait_for_output(
+                    session_id=session_id,
+                    chars=chars,
+                    close_stdin=close_stdin,
+                    terminate=terminate,
+                    wait_for=wait_for,
+                    wait_timeout_ms=clamp_session_int(
+                        wait_timeout_ms,
+                        DEFAULT_WAIT_FOR_MS,
+                        0,
+                        MAX_WAIT_FOR_MS,
+                    ),
+                    max_output_chars=output_limit,
+                )
             poll = await self._manager.write(
                 session_id=session_id,
                 chars=chars,
                 close_stdin=close_stdin,
                 terminate=terminate,
                 yield_time_ms=clamp_session_int(yield_time_ms, DEFAULT_YIELD_MS, 0, MAX_YIELD_MS),
-                max_output_chars=clamp_session_int(
-                    max_output_chars,
-                    DEFAULT_MAX_OUTPUT_CHARS,
-                    1000,
-                    MAX_OUTPUT_CHARS,
-                ),
+                max_output_chars=output_limit,
             )
             return format_session_poll(session_id, poll)
         except KeyError:
@@ -449,6 +482,47 @@ class WriteStdinTool(Tool):
         except Exception as exc:
             return f"Error writing to exec session: {exc}"
 
+    async def _wait_for_output(
+        self,
+        *,
+        session_id: str,
+        chars: str | None,
+        close_stdin: bool,
+        terminate: bool,
+        wait_for: str,
+        wait_timeout_ms: int,
+        max_output_chars: int,
+    ) -> str:
+        deadline = time.monotonic() + (wait_timeout_ms / 1000)
+        aggregate: list[str] = []
+        first = True
+        poll: _SessionPoll | None = None
+
+        while True:
+            remaining_ms = max(0, int((deadline - time.monotonic()) * 1000))
+            step_ms = min(500, remaining_ms)
+            poll = await self._manager.write(
+                session_id=session_id,
+                chars=chars if first else None,
+                close_stdin=close_stdin if first else False,
+                terminate=terminate if first else False,
+                yield_time_ms=step_ms,
+                max_output_chars=max_output_chars,
+            )
+            first = False
+            if poll.output:
+                aggregate.append(poll.output)
+                joined = "".join(aggregate)
+                if wait_for in joined:
+                    poll.output = joined
+                    return format_session_poll(session_id, poll)
+            if poll.done or remaining_ms <= 0:
+                poll.output = "".join(aggregate)
+                result = format_session_poll(session_id, poll)
+                if wait_for not in poll.output:
+                    result += f"\nWait target not observed: {wait_for!r}"
+                return result
+
 
 @tool_parameters(tool_parameters_schema())
 class ListExecSessionsTool(Tool):