Stream LLM responses

The SDK tracks streaming responses correctly — every chunk is forwarded to your caller in real time, and the cost is reported from the final chunk (which carries the usage block).

The pattern

streaming_agent.py

import nullrun
from openai import AsyncOpenAI
from nullrun import init_or_die, protect

init_or_die()
client = AsyncOpenAI()


@protect
async def stream_answer(prompt: str):
    stream = await client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        stream=True,
    )
    async for chunk in stream:
        yield chunk.choices[0].delta.content or ""

The transport hook reads the final usage block before emitting /track, while forwarding chunks to your caller in real time.

Long streams and soft mode

A long stream that exceeds the chain idle TTL (300s) will be killed mid-chunk. The SDK sends a wall-clock heartbeat every 30 seconds per policy (configurable in [10s, 120s]) — not per chunk. For multi-minute responses, use a chain context to keep the gate alive:

@protect
def long_stream(prompt: str):
    with chain("my-long-stream", op="start"):
        stream = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": prompt}],
            stream=True,
        )
        for chunk in stream:
            yield chunk.choices[0].delta.content or ""

For budget headroom, set enforcement_mode = "Soft" on the policy. See Chain context.

Chain heartbeat

The SDK keeps the chain alive with a wall-clock heartbeat every 30 seconds by default (configurable per policy in [10s, 120s]). The interval is time-based, not chunk-based: a slow stream with one chunk per minute still gets a heartbeat; a fast stream does not spam them.

If the chain dies (idle TTL expired, max duration exceeded, or op="end"), the SDK raises WorkflowKilledInterrupt at the next yield boundary.

Kill signal mid-stream

An operator hit on Kill raises WorkflowKilledInterrupt (alias NullRunWorkflowKilledError) at the next yield boundary. It is a NullRunError subclass — caught by except Exception: like every other SDK error. If you want kill-specific handling (close the stream, flush state), catch the typed alias explicitly first and re-raise after.

Cancellation latency

The kill signal typically arrives at the SDK within ~100 ms of the operator clicking Kill (WebSocket push path). If the WebSocket is unavailable and polling fallback is active, latency rises to the poll interval (default 1 s) plus the next /gate boundary. For long-running streams, keep the WS connection healthy — set NULLRUN_TRANSPORT=ws (the default) and avoid restrictive outbound firewalls on the SDK host.

from nullrun import WorkflowKilledInterrupt

@protect
async def stream_kill_safe(prompt: str):
    stream = await client.chat.completions.create(
        model="gpt-4o-mini", messages=[{"role": "user", "content": prompt}],
        stream=True,
    )
    try:
        async for chunk in stream:
            yield chunk.choices[0].delta.content or ""
    except WorkflowKilledInterrupt:
        await stream.close()
        raise

Tracking without auto-instrumentation

If the SDK's httpx transport hook can't see your custom streaming client (a vendor SDK that bypasses httpx), call track_llm manually after the stream ends. Use stream_options={"include_usage": True} so the final chunk carries the usage block; otherwise you have to estimate. See OpenAI streaming reference.

from nullrun import init_or_die, protect, track_llm


@protect
def custom_stream(prompt: str):
    stream = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        stream=True,
        stream_options={"include_usage": True},
    )
    final = None
    for chunk in stream:
        final = chunk
        yield chunk.choices[0].delta.content or ""
    if final and getattr(final, "usage", None):
        track_llm(
            input_tokens=final.usage.prompt_tokens,
            output_tokens=final.usage.completion_tokens,
            model="gpt-4o-mini",
        )

Without track_llm() the budget counter is never credited and the next /gate may reject the next call based on stale spend.

Common pitfalls

Pitfall	Symptom	Fix
Heartbeat every N chunks	Chain dies silently during slow streams	Heartbeat on a wall-clock timer (30s default)
`await stream.close()` after kill	Half-written chunks can leak to the caller	Wrap the stream in `try/finally`, always close
Catching only `Exception` around the loop with no kill handler	Kill still propagates but with no cleanup hook	Catch `NullRunWorkflowKilledError` explicitly first to close the stream
Forgetting `track_llm()` after a manual stream	Dashboard shows zero cost, budget never decremented	Always report final usage, even via estimation