feat(patch): opt-in SSE streaming for /v1/patch_sweep grids#236
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #235 (feat/patch-sweep-spans).
What
Adds opt-in
stream: bool = FalsetoPatchSweepRequest. When true,/v1/patch_sweepreturns atext/event-stream(SSE) response that emits results as the grid runs instead of holding one JSON response open until the whole sweep finishes:startevent (grid shape / resolved axes / auto_captured) so a consumer can size a live heatmap up front.cellevent per cell as it completes (completion order, no ordering promise), each carryinghook(forward-compatible with multi-hook sweeps),layer,position,value. A cell that voids mid-sweep re-emits withvalue: null+error, mirroring how voided cells land inskipped.summaryevent carrying the exact non-streamingPatchSweepResponsepayload, then[DONE].The fan-out + response assembly are factored into a single async generator both paths consume: streaming serializes each event to SSE; non-streaming drains it and returns the summary. The non-streaming contract is unchanged. Cells still run concurrently (
asyncio.as_completed); on client disconnect the outstanding cell tasks are cancelled, aborting their engine requests best-effort.Pre-fan-out errors (bad hook/layers, span/alignment failure, missing source without
clean_prompt, source validation) still return a plain JSON 400 — the stream only starts once the fan-out begins.Client:
PatchStudy.sweep_layers_positions(..., server_side=True, on_cell=fn)sendsstream: true, invokeson_cellper cell event, and builds the returnedSweepResultfrom the summary event (identical to the non-streaming result). Withouton_cell, behavior is unchanged.Why
A large grid (all_prompt × many layers on a real prompt = 1000+ cells) holds a single HTTP response open for minutes with zero progress signal; a proxy/client timeout loses the whole nearly-done sweep. Streaming surfaces incremental progress and lets clients drive a live heatmap.
Expected conflicts
Overlaps with sibling branch feat/patch-multihook-lifecycle (multi-hook
hooksfield + source-run lifecycle) in api_router.py's fan-out / response-assembly sections, protocol.py, and docs/features/activation_patching.md. The cell-event schema already includeshookso the event contract does not need a breaking change when multi-hook lands.