One streamed request per run, instead of start / poll / collect

[JadenFiotto-Kaufman]

A tool run took three legs: POST /start for a job id, the browser polling NDIF's /response/{id} until COMPLETED, then POST /results/{id} to collect and shape it. It is now one POST that stays open and answers in Server-Sent Events — each NDIF status as it lands, then the finished payload.

event: status   {"status":"QUEUED","description":"Added to Queue at position 1."}
event: status   {"status":"DISPATCHED", ...}
event: status   {"status":"RUNNING", ...}
event: status   {"status":"COMPLETED", ...}
event: data     {"meta":{...},"layers":[...],"topk":[...]}

How

nnsight 0.8 ships AsyncRemoteBackend, which submits on the trace's exit and then hands back the raw status updates instead of consuming them. That is the whole mechanism. _api/sse.py turns those updates into status frames, calls the tool's to_data_obj when the saved values arrive, and emits one data frame.

There is no custom backend subclass. An earlier attempt at this — feat/sse-remote-backend, written against 0.7 — needed 102 lines of one, because 0.7 had no async path.

The three tool routes (logit lens, j-lens, activation patching) differ only in which tool and which arguments, so they now say exactly that and nothing else:

@router.post("/run")
async def run_logit_lens(req, state, user_email):
    return stream_tool(state, logit_lens, state[req.model], req.prompt, top_k=req.topk)

What it buys

• The browser no longer talks to NDIF at all. config.ts has no NDIF URL and needs none — one origin, no second host to reach through a tunnel or an ingress. The hardcoded localhost:5001 goes with it; NDIF's API has been on 8001 since long before 0.8.
• Status is pushed, not sampled — QUEUED position and the RUNNING→COMPLETED transition land when they happen rather than up to a second later.
• No window between legs. A collect step could find its model gone; causal_mediation carried a 503 for exactly that case. One connection, no window.

What it costs

A run lives and dies with its connection. Polling a job id survived a reload; this does not. NDIF still finishes the job and the result is still written by the caller's mutation, so what is lost is the ability to rejoin a run in progress. Runs are seconds to a couple of minutes.

Telemetry's job_id correlation is gone too — it belonged to the poll-and-collect flow and the async backend never surfaces one. Telemetry is currently disabled (# TelemetryClient.init(self)); if it comes back and the correlation matters, take the id from the first status update.

Details worth a look

• Encoding. The routes lost their response_model, so payloads are encoded by hand — through jsonable_encoder, because a payload is often a plain dict with pydantic models nested inside it (a generation's completion is a list of Token). Plain json.dumps refuses those, which is how this first failed.
• No heartbeat. Nothing fills the silence between updates. A cold 70B deploy can sit quiet for minutes, and the nginx ingress in front of the preview deployments cuts a connection idle for 60s (proxy-read-timeout, not overridden in deploy/preview/values.yaml). If that bites, the annotation is a better lever than a comment frame; stream_backend's docstring says so. deployApi used to poll on a 20-minute ceiling and now just holds the stream open, which is the flow most exposed to this.
• Errors split by when they happen. Before the stream opens — a 403 for a model the caller cannot use — they stay ordinary HTTP failures. After, they can only be an error frame, because the status line is long gone.
• saves["results"]. NDIF keys returned values by the name of the variable the tool saved, which for every nnsightful tool is results. That is a real coupling to the tool's internals; it is the same one the old collect routes had (backend()["results"]).

Testing

Verified end to end against the self-hosted 0.8 NDIF: the logit lens streams the four statuses above and returns Paris for "The Eiffel Tower is in the city of"; generation returns "…is in the city of Paris, and the E". tsc --noEmit and eslint report the same errors as the base commit — none new (34 before, 34 after; the diff is line numbers only).

Not exercised against a live model: activation patching, j-lens, causal mediation, and lens v1, which share the same two helpers but have not been run.

🤖 Generated with Claude Code

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
workbench	Ready	Preview	Aug 13, 2026 8:28pm

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 446280e9-709b-4b27-b2f0-7c855e1c9591

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[github-actions]

🧹 Preview for PR #141 torn down.