Initial Commit

tests/test_vm_agent_stream.py

@@ -0,0 +1,133 @@
+#!/usr/bin/env python3
+"""Streaming exists for the person watching, and may not change the turn.
+
+    python3 tests/test_vm_agent_stream.py
+
+WHY
+
+The agent's calls were non-streaming on purpose: it acts on a finished
+message, so nothing in the loop needed tokens as they arrived. That was
+true of the agent and false of the operator, who asked to see what the
+model is saying while it says it - a turn on this model runs tens of
+seconds, and a motionless box makes "still thinking" look like "wedged".
+
+So the reply is now read as it arrives. The rule that makes that safe is
+that it is a DISPLAY change and nothing else: the assembled message handed
+back must be identical in shape and content to the one the non-streaming
+call returned, because the agent is not allowed to care which path
+produced it. These cases are all about that equivalence, plus the one
+about the callback - which belongs to the watcher, so it must never be
+able to fail a turn.
+"""
+
+import json
+import os
+import sys
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, os.path.join(HERE, "..", "vmagent", "app"))
+
+FAILS: list[str] = []
+
+
+def check(label: str, condition: bool, detail: str = "") -> None:
+    if condition:
+        print(f"ok    {label}")
+    else:
+        FAILS.append(label)
+        print(f"FAIL  {label}\n      {detail}")
+
+
+import ollama  # noqa: E402
+
+CALL = {"function": {"name": "run", "arguments": {"command": "uname -a"}}}
+
+
+def chunks(*entries) -> list:
+    return [json.dumps(e).encode() for e in entries]
+
+
+# --- the assembled message ----------------------------------------------
+seen: list[str] = []
+body = ollama._read_stream(chunks(
+    {"message": {"content": "Check"}},
+    {"message": {"content": "ing the"}},
+    {"message": {"content": " kernel."}},
+    {"message": {"tool_calls": [CALL]}},
+    {"done": True, "done_reason": "stop", "prompt_eval_count": 120, "eval_count": 44,
+     "message": {}},
+), seen.append)
+
+check("the text is joined in order", body["message"]["content"] == "Checking the kernel.",
+      body["message"]["content"])
+check("tool calls survive the stream", body["message"]["tool_calls"] == [CALL])
+check("the counts come off the final chunk",
+      body.get("prompt_eval_count") == 120 and body.get("eval_count") == 44)
+check("and so does why it stopped", body.get("done_reason") == "stop")
+check("the role is set, as the non-streaming reply sets it",
+      body["message"]["role"] == "assistant")
+
+# The watcher is handed the text SO FAR, not the delta - it paints a line,
+# it does not append to one.
+check("the callback gets the running text, not the fragment",
+      seen == ["Check", "Checking the", "Checking the kernel."], seen)
+
+# --- thinking ------------------------------------------------------------
+thought: list[str] = []
+body = ollama._read_stream(chunks(
+    {"message": {"thinking": "weighing it up"}},
+    {"message": {"content": "Right, "}},
+    {"message": {"thinking": " and more"}},
+    {"message": {"content": "doing it."}},
+    {"done": True, "message": {}},
+), thought.append)
+check("thinking is kept on the message", body["message"]["thinking"] == "weighing it up and more")
+check("it is shown while there is nothing else to show", thought[0] == "weighing it up",
+      "a model that is thinking has not said anything yet, and a blank minute reads as hung")
+check("but once it speaks, its words are what is shown",
+      thought[-1] == "Right, doing it." and " and more" not in thought[-1],
+      thought)
+
+# --- nothing in a frame may cost the turn --------------------------------
+body = ollama._read_stream(
+    [b'{"message": {"content": "a"}}', b"not json at all", b"",
+     b'{"message": {"content": "b"}, "done": true}'],
+    lambda _t: None)
+check("a malformed frame is skipped, not raised on",
+      body["message"]["content"] == "ab",
+      "losing a whole turn over one bad line would be a poor trade")
+
+
+def explode(_text):
+    raise RuntimeError("the watcher's callback failed")
+
+
+try:
+    body = ollama._read_stream(chunks(
+        {"message": {"content": "still works"}}, {"done": True, "message": {}}), explode)
+    survived = body["message"]["content"] == "still works"
+except Exception as exc:  # noqa: BLE001
+    survived = False
+    why = f"{type(exc).__name__}: {exc}"
+check("a callback that throws cannot fail the turn", survived, locals().get("why", ""))
+
+# --- an empty reply is still a reply --------------------------------------
+body = ollama._read_stream(chunks({"done": True, "message": {}}), seen.append)
+check("a reply with nothing in it assembles rather than crashing",
+      body["message"]["content"] == "" and body["message"]["tool_calls"] == [])
+
+# --- and the old path is untouched ----------------------------------------
+src = open(os.path.join(HERE, "..", "vmagent", "app", "ollama.py"), encoding="utf-8").read()
+check("streaming is off unless somebody is watching",
+      '"stream": on_text is not None' in src,
+      "no caller without a callback may be made to pay for this")
+check("only the agent's own turn streams; the summariser does not",
+      open(os.path.join(HERE, "..", "vmagent", "app", "loop.py"),
+           encoding="utf-8").read().count("on_text=") == 1,
+      "nobody watches a compaction summary")
+
+print()
+if FAILS:
+    print(f"FAILED {len(FAILS)}")
+    raise SystemExit(1)
+print("the watcher sees the model talking, and the turn cannot tell the difference")