fix: stream provider deltas through SSE

Connect provider streaming to the SSE endpoint while preserving tool receipts, usage, memory, completion, and legacy payload fields. Add a timing regression and make the precision test use a fixed peak timestamp.\n\nFixes #86

Author: EvanProgramming

docs/self-evolution.md

@@ -9,7 +9,7 @@ Historical verification snapshot: `51be33361422e55e1f2f00c33a0e0f8c56132a91`
 (the post-#54 `main` revision, captured before this #55 documentation-only
 update). Snapshot date: 2026-09-04.
 
-Current repository test count at this snapshot: **160 unittest cases**.
+Current repository test count at this snapshot: **161 unittest cases**.
 
 ## Verified surface
 
@@ -106,8 +106,11 @@ curl -sS http://127.0.0.1:8000/api/chat \
   -d '{"message":"Write a sourced migration note", "profile":"researcher", "session_id":"demo", "speaker":"authenticated", "audience":"team", "channel":"chat"}'
 ```
 
-The streaming endpoint `/api/chat/stream` accepts the same body and emits the
-receipt as an SSE record when one is available.
+The streaming endpoint `/api/chat/stream` accepts the same body and forwards
+provider content deltas as they arrive. Control responses are buffered only
+long enough to parse actions; durable tool receipts, task progress, usage,
+memory, and completion are emitted as typed SSE records. The existing `chunk`,
+`cost`, `memory_receipt`, `error`, and `[DONE]` payloads remain compatible.
 
 ### Inspect learning evidence
 

main.py

@@ -508,8 +508,20 @@ def _retrieve_failure(query: str, n: int = 3) -> list[str]:
 _last_prompt_tokens: int = 0
 _last_completion_tokens: int = 0
 _active_usage_run_id: ContextVar[str | None] = ContextVar("active_usage_run_id", default=None)
+_stream_event_callback: ContextVar[Any] = ContextVar("stream_event_callback", default=None)
 _turn_cost_log: list[dict] = []  # {"tokens":int, "time":float, "tool_calls":int}
 
+
+def _emit_stream_event(event: dict[str, Any]) -> None:
+    """Forward typed stream progress without letting a client affect the turn."""
+    callback = _stream_event_callback.get()
+    if not callable(callback):
+        return
+    try:
+        callback(event)
+    except Exception:
+        pass
+
 def _track_tool_performance(action: str, result: str, elapsed: float) -> None:
     stats = _tool_stats.setdefault(action, {"calls":0,"successes":0,"total_time":0.0})
     stats["calls"] += 1
@@ -997,6 +1009,13 @@ def _spinner_worker(stop_event: threading.Event) -> None:
 
 def _call_llm_with_spinner(messages: list[dict], model: str | None = None) -> str:
     global _SPINNER_STOP, _SPINNER_THREAD
+    streaming = callable(_stream_event_callback.get())
+    if streaming:
+        return _get_llm_response(
+            messages, model=model, stream=True,
+            on_chunk=lambda chunk: _emit_stream_event({"event": "content", "chunk": str(chunk)}),
+            on_stream_end=lambda: _emit_stream_event({"event": "model_complete"}),
+        )
     _SPINNER_STOP.clear()
     _SPINNER_THREAD = threading.Thread(target=_spinner_worker, args=(_SPINNER_STOP,), daemon=True)
     _SPINNER_THREAD.start()
@@ -3303,12 +3322,18 @@ def _record_turn_receipt(receipt: ExecutionReceipt) -> dict[str, Any]:
         "execution.receipt", receipt.as_dict(), user_id=tasks.user_id,
         workspace_id=tasks.workspace_id, session_id=tasks.session_id, task_id=task_id,
     )
-    return {
+    result = {
         "receipt_id": receipt.receipt_id, "operation_id": receipt.operation_id,
         "action": receipt.action, "args": receipt.args, "result": receipt.result,
         "success": receipt.success, "authorized": receipt.authorized,
         "acceptance": receipt.acceptance, "failure": receipt.failure,
     }
+    _emit_stream_event({"event": "tool_receipt", "tool_receipt": result})
+    _emit_stream_event({"event": "tasks", "tasks": [
+        {"id": item["id"], "description": item["description"], "status": item["status"]}
+        for item in tasks.tasks
+    ]})
+    return result
 
 
 def _execute_durable_task(task: dict[str, Any]) -> dict[str, Any]:
@@ -3449,7 +3474,8 @@ def run() -> None:
     return value
 
 
-def _get_llm_response(messages: list[dict[str, str]], model: str | None = None, stream: bool = False) -> str:
+def _get_llm_response(messages: list[dict[str, str]], model: str | None = None, stream: bool = False,
+                      on_chunk: Any = None, on_stream_end: Any = None) -> str:
     global _last_prompt_tokens, _last_completion_tokens, _total_prompt_tokens, _total_completion_tokens
     if llm_provider is None:
         return "[Error] LLM provider not initialised"
@@ -3459,15 +3485,29 @@ def _get_llm_response(messages: list[dict[str, str]], model: str | None = None,
                 workspace_id=memory_bank.workspace_id, session_id=memory_bank.session_id,
                 run_id=_active_usage_run_id.get(), surface=_EXECUTION_SURFACE):
             if stream and hasattr(llm_provider, 'chat_stream'):
-                # Streaming usage frames are persisted by the streaming implementation.
+                before = memory_bank.store.usage_totals(
+                    user_id=memory_bank.user_id, workspace_id=memory_bank.workspace_id,
+                    session_id=memory_bank.session_id, run_id=_active_usage_run_id.get(),
+                )
                 collected: list[str] = []
                 for chunk in llm_provider.chat_stream(messages, model or DEEPSEEK_MODEL):
                     collected.append(chunk)
-                    sys.stdout.write(chunk)
-                    sys.stdout.flush()
+                    if on_chunk:
+                        on_chunk(chunk)
+                    else:
+                        sys.stdout.write(chunk)
+                        sys.stdout.flush()
                 text = "".join(collected).strip()
-                _last_prompt_tokens = 0
-                _last_completion_tokens = len(text) // 4  # presentation-only until a usage frame arrives
+                after = memory_bank.store.usage_totals(
+                    user_id=memory_bank.user_id, workspace_id=memory_bank.workspace_id,
+                    session_id=memory_bank.session_id, run_id=_active_usage_run_id.get(),
+                )
+                _last_prompt_tokens = max(0, int(after["prompt_tokens"]) - int(before["prompt_tokens"]))
+                _last_completion_tokens = max(0, int(after["completion_tokens"]) - int(before["completion_tokens"]))
+                _total_prompt_tokens += _last_prompt_tokens
+                _total_completion_tokens += _last_completion_tokens
+                if on_stream_end:
+                    on_stream_end()
             else:
                 text, usage_dict = _bounded_provider_call(
                     lambda: llm_provider.chat(messages, model or DEEPSEEK_MODEL)

server.py

@@ -16,6 +16,7 @@
 import copy
 import re
 import asyncio
+import queue
 from pathlib import Path
 from typing import Any
 from urllib.parse import urlparse
@@ -778,26 +779,89 @@ async def api_chat_stream(request: Request):
     _set_memory_context(session, body)
     _audit("CHAT_STREAM", f"user={session['user_id']} msg={msg[:80]}", session["user_id"])
 
-    async def generate():
+    class StreamProjection:
+        """Pass plain deltas through while holding model control prefixes."""
+
+        prefixes = ("Thought:", "Plan:", "TaskList:", "TaskDone:", "Action:", "DefineTool:")
+
+        def __init__(self, sink):
+            self.sink = sink
+            self.buffer = ""
+
+        def __call__(self, event: dict[str, Any]) -> None:
+            kind = event.get("event")
+            if kind == "content":
+                self.buffer += str(event.get("chunk", ""))
+                candidate = self.buffer.lstrip()
+                lowered = candidate.lower()
+                if candidate and (
+                    any(prefix.lower().startswith(lowered) for prefix in self.prefixes)
+                    or any(lowered.startswith(prefix.lower()) for prefix in self.prefixes)
+                ):
+                    return
+                if self.buffer:
+                    self.sink({"event": "content", "chunk": self.buffer})
+                    self.buffer = ""
+                return
+            if kind == "model_complete":
+                self._flush()
+                return
+            self.sink(event)
+
+        def _flush(self) -> None:
+            if not self.buffer:
+                return
+            text = self.buffer
+            self.buffer = ""
+            if _agent._collect_tool_calls(text) or any(
+                    text.lstrip().lower().startswith(prefix.lower()) for prefix in self.prefixes
+            ):
+                text = _agent._clean_final_response(text)
+            if text:
+                self.sink({"event": "content", "chunk": text})
+
+    events: queue.Queue[dict[str, Any]] = queue.Queue()
+    projection = StreamProjection(events.put)
+
+    def run_streaming_turn() -> None:
+        callback_token = _agent._stream_event_callback.set(projection)
         try:
-            # The legacy agent is synchronous and may perform network and disk
-            # I/O. Keep it off the ASGI event loop so one chat cannot stall
-            # health checks or unrelated requests.
-            reply = await asyncio.to_thread(_run_session_chat, session, msg)
-            # Send chunks (simulated streaming for non-streaming providers)
-            chunk_size = 20
-            for i in range(0, len(reply), chunk_size):
-                chunk = reply[i:i+chunk_size]
-                yield f"data: {json.dumps({'chunk': chunk})}\n\n"
-                await asyncio_sleep(0.01)
-            yield f"data: {json.dumps({'cost': _cost_summary()})}\n\n"
-            if session.get("last_memory_receipt"):
-                yield f"data: {json.dumps({'memory_receipt': session['last_memory_receipt']})}\n\n"
-            yield "data: [DONE]\n\n"
-            _emit_chat_completed(session, reply, streamed=True)
-            _audit("REPLY_STREAM", f"len={len(reply)}", session["user_id"])
-        except Exception as e:
-            yield f"data: {json.dumps({'error': str(e)})}\n\n"
+            reply = _run_session_chat(session, msg)
+            if str(reply).startswith("[LLM Error]"):
+                events.put({"event": "error", "error": str(reply)})
+            else:
+                events.put({"event": "complete", "reply": reply})
+        except Exception as exc:
+            events.put({"event": "error", "error": str(exc)})
+        finally:
+            _agent._stream_event_callback.reset(callback_token)
+
+    async def generate():
+        # The synchronous agent runs in a dedicated worker and safely drains if
+        # a client disconnects; it never blocks the ASGI event loop.
+        threading.Thread(target=run_streaming_turn, daemon=True).start()
+        while True:
+            event = await asyncio.to_thread(events.get)
+            kind = event.get("event")
+            if kind == "content":
+                yield f"data: {json.dumps({'event': 'content', 'chunk': event.get('chunk', '')}, ensure_ascii=False)}\n\n"
+            elif kind == "tool_receipt":
+                yield f"data: {json.dumps({'event': 'tool_receipt', 'tool_receipt': event.get('tool_receipt')}, ensure_ascii=False)}\n\n"
+            elif kind == "tasks":
+                yield f"data: {json.dumps({'event': 'tasks', 'tasks': event.get('tasks', [])}, ensure_ascii=False)}\n\n"
+            elif kind == "error":
+                yield f"data: {json.dumps({'event': 'error', 'error': event.get('error', 'stream failed')}, ensure_ascii=False)}\n\n"
+                break
+            elif kind == "complete":
+                reply = str(event.get("reply", ""))
+                yield f"data: {json.dumps({'event': 'usage', 'cost': _cost_summary()})}\n\n"
+                if session.get("last_memory_receipt"):
+                    yield f"data: {json.dumps({'event': 'memory_receipt', 'memory_receipt': session['last_memory_receipt']}, ensure_ascii=False)}\n\n"
+                yield f"data: {json.dumps({'event': 'completion', 'status': 'completed'})}\n\n"
+                yield "data: [DONE]\n\n"
+                _emit_chat_completed(session, reply, streamed=True)
+                _audit("REPLY_STREAM", f"len={len(reply)}", session["user_id"])
+                break
 
     return StreamingResponse(generate(), media_type="text/event-stream")
 
@@ -1660,11 +1724,6 @@ async def pwa_manifest():
     }
 
 
-# Async sleep helper
-async def asyncio_sleep(seconds: float):
-    await asyncio.sleep(seconds)
-
-
 # ---------------------------------------------------------------------------
 # Plugin system
 # ---------------------------------------------------------------------------

tests/test_streaming.py

@@ -0,0 +1,92 @@
+import asyncio
+import json
+import tempfile
+import threading
+import time
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+import main
+import server
+from memory import MemoryBank
+from providers import ProviderConfig
+
+
+class _Request:
+    def __init__(self, body):
+        self.body = body
+
+    async def json(self):
+        return self.body
+
+
+class _DelayedProvider:
+    def __init__(self):
+        self.config = ProviderConfig(provider="deepseek", model_simple="deepseek-v4-flash")
+        self.release = threading.Event()
+        self.completed = threading.Event()
+
+    def chat_stream(self, _messages, _model=None):
+        yield "FIRST"
+        self.release.wait(timeout=3)
+        yield " SECOND"
+        self.completed.set()
+
+
+class StreamingEndpointTests(unittest.TestCase):
+    def test_first_sse_content_arrives_before_provider_stream_completes(self):
+        session_id = "stream-timing-regression"
+        provider = _DelayedProvider()
+
+        async def exercise():
+            response = await server.api_chat_stream(_Request({
+                "message": "hello", "session_id": session_id,
+            }))
+            iterator = response.body_iterator.__aiter__()
+            started = time.monotonic()
+            first = await asyncio.wait_for(anext(iterator), timeout=1)
+            first_elapsed = time.monotonic() - started
+            provider_complete_before_release = provider.completed.is_set()
+            provider.release.set()
+            remaining = []
+            try:
+                while True:
+                    remaining.append(await asyncio.wait_for(anext(iterator), timeout=2))
+            except StopAsyncIteration:
+                pass
+            return first, first_elapsed, provider_complete_before_release, remaining
+
+        with tempfile.TemporaryDirectory(prefix="openkyrozen-stream-") as directory:
+            memory = MemoryBank(Path(directory) / "state.sqlite3", workspace_id="stream-test")
+            original_sessions = server._sessions
+            server._sessions = {}
+            try:
+                with (patch.object(server._agent, "memory_bank", memory),
+                      patch.object(server._agent, "llm_provider", provider),
+                      patch.object(server._agent, "DEEPSEEK_MODEL", "deepseek-v4-flash"),
+                      patch.object(server._agent, "_chat_turn", side_effect=lambda message, **_: (
+                          main._call_llm_with_spinner([{"role": "user", "content": message}])
+                      ))):
+                    first, elapsed, provider_complete_before_release, remaining = asyncio.run(exercise())
+            finally:
+                server._sessions = original_sessions
+
+        def text(item):
+            return item.decode() if isinstance(item, bytes) else item
+
+        first_payload = json.loads(text(first).split("data: ", 1)[1].splitlines()[0])
+        payloads = [json.loads(text(item).split("data: ", 1)[1].splitlines()[0])
+                    for item in remaining
+                    if text(item).startswith("data: ") and text(item) != "data: [DONE]\n\n"]
+        self.assertLess(elapsed, 1)
+        self.assertFalse(provider_complete_before_release)
+        self.assertEqual(first_payload, {"event": "content", "chunk": "FIRST"})
+        self.assertEqual([item.get("event") for item in payloads[:3]], ["content", "usage", "completion"])
+        self.assertEqual(payloads[0]["chunk"], " SECOND")
+        self.assertEqual(sum(text(item) == "data: [DONE]\n\n" for item in remaining), 1)
+        self.assertTrue(provider.completed.is_set())
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_usage_ledger.py

@@ -121,20 +121,21 @@ def test_reset_requires_explicit_confirmation_and_preserves_installation_ledger(
     def test_short_and_mixed_calls_aggregate_before_display_rounding(self):
         with tempfile.TemporaryDirectory(prefix="openkyrozen-usage-precision-") as directory:
             store = EventStore(Path(directory) / "state.sqlite3")
+            peak_time = datetime(2026, 9, 7, 2, tzinfo=timezone.utc)
             with usage_scope(store=store, workspace_id="project", session_id="short"):
                 for _ in range(1000):
                     _track_cost("deepseek", {"prompt_tokens": 0, "completion_tokens": 100},
-                                model="deepseek-chat")
+                                model="deepseek-chat", occurred_at=peak_time)
             with usage_scope(store=store, workspace_id="project", session_id="mixed"):
                 for _ in range(1000):
                     _track_cost("deepseek", {"prompt_tokens": 100, "completion_tokens": 100},
-                                model="deepseek-chat")
+                                model="deepseek-chat", occurred_at=peak_time)
             with usage_scope(store=store, workspace_id="project", session_id="combined-short"):
                 _track_cost("deepseek", {"prompt_tokens": 0, "completion_tokens": 100_000},
-                            model="deepseek-chat")
+                            model="deepseek-chat", occurred_at=peak_time)
             with usage_scope(store=store, workspace_id="project", session_id="combined-mixed"):
                 _track_cost("deepseek", {"prompt_tokens": 100_000, "completion_tokens": 100_000},
-                            model="deepseek-chat")
+                            model="deepseek-chat", occurred_at=peak_time)
 
             short = store.usage_totals(workspace_id="project", session_id="short", user_id="local")
             mixed = store.usage_totals(workspace_id="project", session_id="mixed", user_id="local")