core/tools: two-line wait prompt with reason input and colors

src/plyngent/tools/chat/wait.py

@@ -1,20 +1,44 @@
 from __future__ import annotations
 
 import asyncio
+import re
+import sys
+
+from click import style
 
 from plyngent.agent import ToolTag, tool
 from plyngent.prompting import get_prompt_backend, read_line_with_timeout, run_prompt_async
 
-WAIT_PROMPT = "Waiting for {duration}s. Press Enter to disturb. Feel free to ship any reason before pressing Enter. "
+_ANSI_ESCAPE_RE = re.compile(r"\x1b\[[0-9;]*m")
+
+
+def _wait_prompt(duration: int, *, reason: str | None) -> str:
+    """Two-line wait prompt: status line, then the optional-reason input prompt.
+
+    ``read_line_with_timeout`` writes the prompt verbatim to stdout (no click
+    echo ANSI stripping), so styling is applied here and stripped when stdout
+    is not a TTY. On a terminal: cyan status, dimmed detail, yellow input
+    prompt; plain text otherwise (tests, redirected output).
+    """
+    status = style(f"Waiting for {duration}s", fg="cyan")
+    if reason:
+        status += style(f" ({reason})", fg="bright_black")
+    status += style(". Press Enter to disturb.", fg="bright_black")
+    input_prompt = style("Reason (optional): ", fg="yellow")
+    rendered = f"{status}\n{input_prompt}"
+    if not sys.stdout.isatty():
+        return _ANSI_ESCAPE_RE.sub("", rendered)
+    return rendered
 
 
 @tool(name="wait", tags=ToolTag.LOCAL | ToolTag.READ_ONLY)
 async def wait(duration: int, *, reason: str | None = None) -> str:
     """Wait ``duration`` seconds before continuing.
 
-    Interactive sessions show a line prompt: pressing Enter (optionally after
+    Interactive sessions show a two-line prompt: a status line, then an
-    typing a reason) "disturbs" the wait so the turn continues immediately.
+    optional-reason input — pressing Enter (optionally after typing a reason)
-    Non-interactive runs simply sleep the full duration.
+    "disturbs" the wait so the turn continues immediately. Non-interactive runs
+    simply sleep the full duration.
     """
     if duration < 0:
         return "error: duration must not be negative"
@@ -24,12 +48,7 @@ async def wait(duration: int, *, reason: str | None = None) -> str:
     if not backend.is_interactive():
         await asyncio.sleep(duration)
         return f"waited {duration}s"
-    prompt = WAIT_PROMPT.format(duration=duration)
+    prompt = _wait_prompt(duration, reason=reason)
-    if reason:
-        prompt = (
-            f"Waiting for {duration}s (reason: {reason}). Press Enter to disturb. "
-            "Feel free to ship any reason before pressing Enter. "
-        )
     line = await run_prompt_async(read_line_with_timeout, prompt, float(duration))
     if line is None:
         return f"waited {duration}s"

tests/test_tools/test_chat_tools.py

@@ -144,3 +144,51 @@ async def test_wait_tool_times_out(monkeypatch) -> None:
         registry = ToolRegistry([wait])
         out = await registry.execute("wait", '{"duration": 3}')
     assert out == "waited 3s"
+
+
+async def test_wait_tool_prompt_two_lines_with_reason(monkeypatch) -> None:
+    prompts: list[str] = []
+    monkeypatch.setattr(
+        wait_module,
+        "read_line_with_timeout",
+        lambda prompt, timeout: prompts.append(prompt) or None,
+    )
+    with temporary_backend(ScriptedBackend([])):
+        registry = ToolRegistry([wait])
+        out = await registry.execute("wait", '{"duration": 5, "reason": "tests are slow"}')
+    assert out == "waited 5s"
+    assert prompts == ["Waiting for 5s (tests are slow). Press Enter to disturb.\nReason (optional): "]
+
+
+async def test_wait_tool_prompt_two_lines_no_reason(monkeypatch) -> None:
+    prompts: list[str] = []
+    monkeypatch.setattr(
+        wait_module,
+        "read_line_with_timeout",
+        lambda prompt, timeout: prompts.append(prompt) or None,
+    )
+    with temporary_backend(ScriptedBackend([])):
+        registry = ToolRegistry([wait])
+        out = await registry.execute("wait", '{"duration": 5}')
+    assert out == "waited 5s"
+    assert prompts == ["Waiting for 5s. Press Enter to disturb.\nReason (optional): "]
+
+
+def test_wait_prompt_plain_when_not_a_tty() -> None:
+    """Prompt styling is stripped when stdout is not a TTY (pytest capture)."""
+    prompt = wait_module._wait_prompt(5, reason="tests")
+    assert "\x1b[" not in prompt
+    assert prompt == "Waiting for 5s (tests). Press Enter to disturb.\nReason (optional): "
+
+
+def test_wait_prompt_colored_on_tty(monkeypatch) -> None:
+    """On a TTY the prompt is styled: cyan status, yellow input prompt."""
+
+    class _FakeStdout:
+        def isatty(self) -> bool:
+            return True
+
+    monkeypatch.setattr(wait_module.sys, "stdout", _FakeStdout())
+    prompt = wait_module._wait_prompt(5, reason="tests")
+    assert "\x1b[36m" in prompt  # cyan status
+    assert "\x1b[33m" in prompt  # yellow input prompt