#!/usr/bin/env python3
"""
SSE stream handling component for Klea pages.
The pure event-to-state reducer moved to
:mod:`klea_utils.ui.stream_events` (shared with the terminal client); this
module keeps the UI-driving coroutines (:func:`run_stream`,
:func:`stop_stream`) that consume the backend's ``/query/stream`` events and
update the chat panel, status pane and inspector.
See ``devdocs/system/streams.md`` and ``devdocs/system/api-sse-sequence.md``.
File: klea_utils/ui/web/nicegui/components/stream.py
Copyright 2026 Ankur Sinha
Author: Ankur Sinha <sanjay DOT ankur AT gmail DOT com>
"""
import asyncio
import logging
from datetime import datetime
from typing import Any
import httpx
from nicegui import background_tasks
from klea_utils.api.sse import request_cancel, stream_events
from klea_utils.ui.state import InspectorMarker, ensure_chat
from klea_utils.ui.stream_events import TOOL_ROUND_SECTION_KEY, apply_stream_event
from klea_utils.ui.web.nicegui.components.context import PageContext
logger = logging.getLogger(__name__)
def _retry_stream(ctx: PageContext, query: str, chat_id: str) -> None:
"""Resume the chat's last failed run (Retry action)."""
task = background_tasks.create(run_stream(ctx, query, chat_id, resume=True))
ctx.stream_tasks[f"{ctx.user_id}:{chat_id}"] = task
[docs]
def stop_stream(ctx: PageContext, chat_id: str) -> None:
"""Stop *chat_id*'s active run: cancel locally and ask the server to cancel.
Cancels the chat's ``run_stream`` task (closing the SSE connection) and
fires a best-effort ``/query/cancel`` so the server stops the graph run
even if the disconnect alone would not. Idempotent and safe to call when
the chat is not streaming.
:param ctx: The shared page context.
:param chat_id: Chat whose run should stop.
"""
task = ctx.stream_tasks.get(f"{ctx.user_id}:{chat_id}")
logger.debug("stop_stream(chat=%s, task=%s)", chat_id, task)
if task is not None and not task.done():
task.cancel()
if chat_id:
background_tasks.create(request_cancel(ctx.server_url, chat_id, ctx.user_id))
[docs]
def stop_all_streams(ctx: PageContext) -> None:
"""Stop every active run on this page (used when exiting the frontend).
Cancels each chat's ``run_stream`` task and fires the best-effort cancel
request, so a reused/remote server does not keep running after the page
goes away.
:param ctx: The shared page context.
"""
keys = list(ctx.stream_tasks)
logger.debug(f"stop_all_streams: {keys = }")
for key in keys:
chat_id = key.split(":", 1)[1] if ":" in key else key
stop_stream(ctx, chat_id)
[docs]
async def run_stream(
ctx: PageContext,
query: str,
chat_id: str,
resume: bool = False,
interrupt_response: dict[str, Any] | None = None,
interrupt_id: str | None = None,
interrupt_cancel: bool = False,
) -> None:
"""Stream a query's events into the UI for *chat_id*.
Shows a progress row while streaming, commits the final answer and
inspector data on completion, and surfaces errors inline (with a Retry
action when the run can be resumed from its checkpoint). When the run
pauses for human input it renders the ask (the status region builds the
form) and stops streaming until the user answers or cancels.
:param ctx: The shared page context.
:param query: The user's query text (unused when resuming/answering).
:param chat_id: Chat conversation identifier.
:param resume: Resume the chat's last failed run instead of starting a
new turn (the query is not resent).
:param interrupt_response: Answer a pending HITL interrupt (ADR-0046).
:param interrupt_id: Id of the interrupt being answered.
:param interrupt_cancel: Cancel the pending interrupt instead of answering.
"""
current_chat = ensure_chat(ctx.user_id, chat_id)
is_query = not resume and interrupt_response is None and not interrupt_cancel
logger.debug(
"Streaming chat %s (resume=%s interrupt=%s)",
chat_id,
resume,
"cancel" if interrupt_cancel else interrupt_response is not None,
)
key = f"{ctx.user_id}:{chat_id}"
def _active() -> bool:
"""Whether *chat_id* is the chat currently shown on the page.
Chat data is always updated; the DOM is only touched when the run's
chat is the visible one, so a background run cannot repaint the chat
the user is looking at.
"""
return chat_id == ctx.chat_id
# Per-chat run registry (page-scoped): the callers register the created
# task, so a run is visible for gating before this coroutine starts;
# ``setdefault`` keeps their handle when they raced ahead.
ctx.stream_tasks.setdefault(key, asyncio.current_task())
# Retry resumes this turn; the status region's Retry button invokes it.
# Kept after the run (below) so a failed run's Retry action survives.
ctx.retry_cbs[key] = lambda: _retry_stream(ctx, query, chat_id)
if _active():
ctx.refresh_stream_button()
# Status sections persist across turns: each node updates its own section
# by key, so the pane keeps the last known state (e.g. the live plan)
# instead of blanking until the node re-emits. Only the transient
# tool-round section is dropped: it is scoped to one dispatch, and an
# interrupted turn would otherwise leave a stuck ``Running N tool(s)``
# section behind (its ``running``/``completed`` pair shares one key).
current_chat.setdefault("state_sections", {}).pop(TOOL_ROUND_SECTION_KEY, None)
# The status region is per-chat state (rendered by the chat-area
# component), so it is replaced wholesale on every turn and switching
# chats restores each chat's own status.
current_chat["turn_status"] = {"kind": "progress", "heading": ""}
current_chat.pop("interrupt", None)
if not is_query:
# Continuing the same turn (resume or an interrupt answer/cancel):
# keep the existing inspector section and re-activate it so appended
# entries land in the right place.
if _active():
ctx.refresh_inspector()
else:
# Start a new inspector section for this query. Entries are appended
# live; sections (and their entries) are kept for the session.
marker: InspectorMarker = {
"type": "query",
"text": query,
"stamp": datetime.now().astimezone().strftime("%X"),
}
current_chat["inspector_entries"].append(marker)
ctx.begin_inspector_section(chat_id, marker)
if _active():
ctx.refresh_status_pane()
ctx.refresh_turn_status()
def _reset_streaming_state() -> None:
"""Drop this chat's run entry and refresh the send control/status."""
ctx.stream_tasks.pop(key, None)
if _active():
ctx.refresh_stream_button()
ctx.refresh_status_pane()
try:
async for event in stream_events(
query,
chat_id,
ctx.server_url,
user_id=ctx.user_id,
extra=ctx.query_extra or None,
resume=resume,
interrupt_response=interrupt_response,
interrupt_id=interrupt_id,
interrupt_cancel=interrupt_cancel,
):
t = event.get("type", "?")
logger.debug("chat=%s stream event type=%s", chat_id, t)
if t == "ping":
# Periodic server heartbeat while a long-running node emits no
# events: the connection is alive, nothing to render.
continue
if t == "progress":
heading = (event.get("data") or {}).get("heading") or event.get(
"node", ""
)
# Update the live heading in place (no full re-render).
current_chat["turn_status"]["heading"] = heading
if _active() and ctx.turn_status_label is not None:
ctx.turn_status_label.set_text(heading)
continue
action = apply_stream_event(current_chat, event)
if action == "context":
# Reconcile context controls (mode / access) with the
# checkpointed value the server just reported, e.g. after a
# ``/mode`` command (ADR-0032, ADR-0047).
for hook in ctx.context_hooks:
hook(current_chat)
if action in ("usage", "state", "context"):
if _active():
ctx.refresh_status_pane()
elif action == "inspect":
entry = current_chat["inspector_entries"][-1]
ctx.append_inspector(chat_id, entry)
elif action in ("tool", "notice"):
if _active():
ctx.render_chat_area()
elif action == "interrupt":
# The run paused for HITL input: show the ask and stop
# streaming until the user answers or cancels (ADR-0046).
logger.debug("chat=%s paused for input", chat_id)
current_chat["turn_status"] = {"kind": "awaiting_input"}
if _active():
ctx.refresh_turn_status()
_reset_streaming_state()
break
elif action == "complete":
logger.debug("chat=%s stream complete", chat_id)
current_chat["turn_status"] = {"kind": "idle"}
current_chat.pop("interrupt", None)
if _active():
ctx.render_chat_area()
_reset_streaming_state()
if _active():
ctx.refresh_inspector()
break
elif action == "error":
error_msg = event.get("message", "Unknown error")
resumable = bool(event.get("resumable"))
logger.debug(
"chat=%s stream error: %s (resumable=%s)",
chat_id,
error_msg,
resumable,
)
current_chat["turn_status"] = {
"kind": "error",
"message": error_msg,
"resumable": resumable,
}
if _active():
ctx.refresh_turn_status()
_reset_streaming_state()
break
except asyncio.CancelledError:
# The user pressed Stop: record a "Stopped" status for this chat and
# re-render the region from it. Re-raise so the task ends cancelled
# (NiceGUI's exception handler ignores it).
logger.debug("chat=%s stream cancelled by user", chat_id)
current_chat["turn_status"] = {"kind": "stopped"}
if _active():
ctx.refresh_turn_status()
_reset_streaming_state()
raise
except httpx.HTTPStatusError as e:
# The backend rejected the streaming request (e.g. 409 while a run is
# still active, or 400 for a resume with nothing to continue). Surface
# the status instead of letting it escape the background task.
logger.debug("chat=%s HTTP status error: %s", chat_id, e)
current_chat["turn_status"] = {
"kind": "error",
"message": f"Request failed ({e.response.status_code}): {e}",
"resumable": False,
}
if _active():
ctx.refresh_turn_status()
_reset_streaming_state()
except httpx.RequestError as e:
# Transport failure: a dropped connection or idle timeout. The run is
# checkpointed server-side, so offer a resume.
logger.debug("chat=%s request error: %s", chat_id, e)
current_chat["turn_status"] = {
"kind": "error",
"message": f"Connection error: {e}",
"resumable": True,
}
if _active():
ctx.refresh_turn_status()
_reset_streaming_state()
finally:
# Clear this chat's run entry on every exit path (including an
# unexpected exception) so it never leaks a stale task handle. The
# retry callback is kept so a failed run's Retry action survives.
ctx.stream_tasks.pop(key, None)
# Re-render the chat list so its per-chat run indicator stops
# spinning (or switches to the awaiting-input icon) now that the run
# is no longer registered. Unconditional: the indicator is per chat,
# so it must clear even while a different chat is active.
ctx.refresh_chat_list()