Hide reasoning content after streaming done.

Author: beacoder

python_agent_harness/client.py

@@ -203,6 +203,7 @@ def chat(
             reasoning_effort=reasoning_effort,
         )
         content_parts: list[str] = []
+        reasoning_parts: list[str] = []
         tc_index: dict[int, dict[str, Any]] = {}
         usage = Usage()
 
@@ -238,6 +239,7 @@ def chat(
                             on_delta(delta["content"])
                     if delta.get("reasoning_content"):
                         content_parts.append(delta["reasoning_content"])
+                        reasoning_parts.append(delta["reasoning_content"])
                         if on_delta:
                             on_delta(delta["reasoning_content"])
                     for tc in delta.get("tool_calls") or []:
@@ -268,10 +270,12 @@ def chat(
                 for i in sorted(tc_index)
             ]
         if content or tool_calls:
-            msg = Message(role="assistant", content=content, tool_calls=tool_calls)
+            msg = Message(role="assistant", content=content, tool_calls=tool_calls,
+                          reasoning="".join(reasoning_parts) or None)
             _log_llm_interaction(self.log_path, payload, msg, usage)
             return msg, usage
-        msg = Message(role="assistant", content="")
+        msg = Message(role="assistant", content="",
+                      reasoning="".join(reasoning_parts) or None)
         _log_llm_interaction(self.log_path, payload, msg, usage)
         return msg, usage
 

python_agent_harness/tui.py

@@ -150,6 +150,24 @@ def _strip_final_check(text: str) -> str:
     return "\n".join(lines[:first]).rstrip()
 
 
+def _strip_reasoning(text: str, reasoning: str) -> str:
+    """Remove the leading REASONING block from TEXT, or TEXT unchanged.
+
+    Reasoning content is streamed before the answer, so it forms the
+    leading part of the stored message content.  The TUI collapses it
+    to a marker once the stream is done, so it stops eating the
+    visible-row budget; the stored message is never modified.
+    """
+    if not reasoning:
+        return text
+    if text.startswith(reasoning):
+        return text[len(reasoning):]
+    stripped = text.lstrip()
+    if stripped.startswith(reasoning):
+        return stripped[len(reasoning):]
+    return text
+
+
 def _history_path() -> str:
     d = config.SESSION_DIR / "python-agent-harness"
     d.mkdir(parents=True, exist_ok=True)
@@ -497,7 +515,14 @@ def _build_history_rows(self) -> list[Any]:
                 if body.strip():
                     rows.append(Markdown(f"**user:** {body}"))
             elif m.role == "assistant":
-                body = _tail_lines(_strip_final_check(m.text()), 12)
+                body = m.text()
+                collapsed_reasoning = False
+                if m.reasoning:
+                    stripped = _strip_reasoning(body, m.reasoning)
+                    if stripped != body:
+                        body = stripped
+                        collapsed_reasoning = True
+                body = _tail_lines(_strip_final_check(body), 12)
                 if m.tool_calls:
                     for tc in m.tool_calls:
                         args = tc.arguments
@@ -513,6 +538,11 @@ def _build_history_rows(self) -> list[Any]:
                             params = ""
                         label = f"🤖 {tc.name}({params})" if params else f"🤖 {tc.name}"
                         rows.append(Text(label, style="cyan"))
+                if collapsed_reasoning:
+                    # the reasoning streamed live while it was being
+                    # produced; once it is done it collapses to a marker
+                    # so it doesn't eat the visible-row budget
+                    rows.append(Text("💭 ...", style="dim"))
                 if body.strip():
                     rows.append(Markdown(f"**assistant:** {body}"))
             elif m.role == "tool":

tests/fake_openai_server.py

@@ -12,6 +12,8 @@ def do_POST(self):
         stream = body.get("stream", False)
         if stream:
             chunks = [
+                {"choices": [{"delta": {"role": "assistant", "reasoning_content": "thinking"}}]},
+                {"choices": [{"delta": {"reasoning_content": " hard"}}]},
                 {"choices": [{"delta": {"role": "assistant", "content": "Hello"}}]},
                 {"choices": [{"delta": {"content": " world"}}]},
                 {"choices": [{"delta": {"tool_calls": [

tests/test_client.py

@@ -0,0 +1,39 @@
+"""Client streaming tests against the in-process fake OpenAI server."""
+
+import unittest
+
+from python_agent_harness.client import Client
+
+from fake_openai_server import serve
+
+
+def make_client() -> Client:
+    srv = serve()
+    host, port = srv.server_address
+    c = Client(base_url=f"http://{host}:{port}/v1", api_key="test", model="fake")
+    c._server = srv  # keep the server alive for the test
+    return c
+
+
+class TestClientStreaming(unittest.TestCase):
+    def test_reasoning_content_streamed_and_captured(self):
+        """reasoning_content deltas stream normally (on_delta) AND are
+        captured on the message so the TUI can collapse them later."""
+        c = make_client()
+        deltas: list[str] = []
+        msg, usage = c.chat(
+            [__import__("python_agent_harness.models", fromlist=["Message"]).Message(
+                role="user", content="hi"
+            )],
+            on_delta=deltas.append,
+        )
+        self.assertEqual("".join(deltas), "thinking hardHello world")
+        self.assertEqual(msg.content, "thinking hardHello world")
+        self.assertEqual(msg.reasoning, "thinking hard")
+        self.assertIn("Hello world", msg.content)
+        self.assertEqual(usage.input_tokens, 12)
+        c.close()
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_tui.py

@@ -190,6 +190,71 @@ def test_stream_final_check_hidden(self):
         self.assertIn("Working...", row.plain)
         self.assertNotIn("Status:", row.plain)
 
+    def test_reasoning_streams_normally(self):
+        """While streaming, reasoning content shows up live like any
+        other text — the collapse only happens in the final history."""
+        tui, _ = make_tui()
+        tui.stream_text = "thinking hard...\n\nanswer here"
+        row = tui._stream_row()
+        self.assertIsNotNone(row)
+        self.assertIn("thinking hard...", row.plain)
+        self.assertIn("answer here", row.plain)
+
+    def test_reasoning_collapsed_in_history(self):
+        """Once the stream is done, the stored reasoning collapses to a
+        '...' marker in the history so it stops eating TUI space; the
+        answer stays fully visible.  The stored message is untouched."""
+        tui, buf = make_tui()
+        reasoning = "Let me think: Rayleigh scattering dominates..."
+        tui.session.last_messages = [
+            Message(role="user", content="why is the sky blue?"),
+            Message(
+                role="assistant",
+                content=reasoning + "\n\nThe sky is blue due to Rayleigh scattering.",
+                reasoning=reasoning,
+            ),
+        ]
+        tui.console.print(tui._render_conversation())
+        out = buf.getvalue()
+        self.assertIn("💭 ...", out)
+        self.assertIn("Rayleigh scattering.", out)
+        self.assertNotIn("Let me think", out)
+        # the stored message keeps its reasoning — only the display hides it
+        self.assertEqual(
+            tui.session.last_messages[1].reasoning, reasoning
+        )
+        self.assertIn("Let me think", tui.session.last_messages[1].text())
+
+    def test_reasoning_collapsed_marker_shows_even_without_answer(self):
+        """A reasoning-only assistant message (no answer content) still
+        shows the collapse marker instead of vanishing entirely."""
+        tui, buf = make_tui()
+        tui.session.last_messages = [
+            Message(role="user", content="go"),
+            Message(
+                role="assistant", content="pensive thoughts here",
+                reasoning="pensive thoughts here",
+            ),
+        ]
+        tui.console.print(tui._render_conversation())
+        out = buf.getvalue()
+        self.assertIn("💭 ...", out)
+        self.assertNotIn("pensive thoughts", out)
+
+    def test_strip_reasoning(self):
+        """_strip_reasoning removes the leading reasoning prefix and
+        leaves non-matching text untouched."""
+        from python_agent_harness.tui import _strip_reasoning
+
+        self.assertEqual(
+            _strip_reasoning("ABCanswer", "ABC"), "answer"
+        )
+        self.assertEqual(
+            _strip_reasoning("  ABCanswer", "ABC"), "answer"
+        )
+        self.assertEqual(_strip_reasoning("answer", "ABC"), "answer")
+        self.assertEqual(_strip_reasoning("x", ""), "x")
+
     def test_long_conversation_bounded(self):
         """A long conversation is capped so rendering stays fast."""
         tui, buf = make_tui()