Skip to content

Feat/streaming preemption uvloop - #6611

Open
blasscoc wants to merge 4 commits into
google:mainfrom
blasscoc:feat/streaming-preemption-uvloop
Open

Feat/streaming preemption uvloop#6611
blasscoc wants to merge 4 commits into
google:mainfrom
blasscoc:feat/streaming-preemption-uvloop

Conversation

@blasscoc

@blasscoc blasscoc commented Aug 6, 2026

Copy link
Copy Markdown

Problem: Workflow nodes today wait for each LLM turn to fully complete before the graph advances, and fan-out always waits for every branch. For search/retrieval-style work this wastes both output tokens (the model keeps generating a long answer we've already decided we don't need) and wall-clock time (downstream work can't start until generation ends, and losing branches keep running). Agents used in these pipelines often "waffle" or aren't perfectly aligned, but they're still perfectly usable if we can act on the useful part of the stream and drop the rest.

Solution: Three opt-in workflow primitives (plus one runtime switch) that let a graph react to a stream as it arrives. None change existing behavior unless you use them.

StreamingRouterNode — mid-stream preemption. A monitor callback inspects the SSE token stream and can commit an output and cancel the rest of generation the moment the answer is known (e.g. stop after a VERDICT: line instead of generating a 300-word summary you'll discard). Cancellation propagates cooperatively via aclosing.
SpeculativeRouterNode — speculative dispatch/overlap. As the agent streams a structured call, it repairs the still-truncated JSON (repair_json), dispatches a downstream target immediately, and keeps generating; when the finalized call arrives it verifies (keep on hit, cancel + re-run on miss). A combine hook lets the agent's full text remain a returned deliverable, so the overlap win is real and not from throwing output away. Use only for idempotent/cancel-safe targets.
FirstMatchNode — first-answer-wins fan-out. Races several branches, returns the first whose output satisfies a match predicate, and cancels the losers (search Recall@k).
enable_uvloop() — one-line runtime switch to run agents on the libuv event loop.
These are aimed squarely at search/fan-out tasks: read N sources in parallel, take the first good answer, don't keep summarizing the ones you're discarding, and overlap dependent steps.

Testing Plan
Unit Tests:

I have added or updated unit tests for my change.

All unit tests pass locally.
New/updated suites: test_streaming_router.py, test_speculative_router.py (incl. repair_json + combine), test_first_match_node.py, test_event_loop.py.

uv run pytest -q tests/unittests/workflow/test_streaming_router.py
tests/unittests/workflow/test_speculative_router.py
tests/unittests/workflow/test_first_match_node.py
tests/unittests/utils/test_event_loop.py
29 passed, 1 warning in 6.88s
Unit tests use scripted agents (no network) and cover: preemption cut points, speculation hit/miss + rollback, JSON repair edge cases, combine returning both agent text and target result, cross-branch cancellation, no-match fallback, and max_parallel.

Manual End-to-End (E2E) Tests:

Real-LLM integration tests (Vertex, gemini-3.5-flash-lite, whole documents, no chunking). They read config from the environment and skip unless GOOGLE_CLOUD_PROJECT is set:

uv run pytest -s -p no:cacheprovider
tests/integration/test_streaming_router_preemption_timing.py
tests/integration/test_speculative_router_chained_llm.py
tests/integration/test_streaming_router_tsla_10k.py
Both integration tests are written to be honest — the feature's output is a required deliverable, so a naive baseline gets no shortcut:

Preemption (5 arXiv papers in parallel): the relevant paper's full summary is kept and asserted in both runs; only the 4 irrelevant summaries are skipped. Result: 6.31s → 3.36s (1.88× faster), output tokens 3531 → 673 (5.25× fewer), cost −10.4% (−48.3% with context caching). Relevant summary preserved (465 vs 413 words).
Speculative chained LLM (planner → worker): planner rationale is a required output; both paths return {plan, answer}. Result: 6.11s → 3.52s, 2.59s saved (~72% of the worker call hidden behind the required rationale tail).
Checklist

I have read the CONTRIBUTING.md document.

I have performed a self-review of my own code.

I have commented my code, particularly in hard-to-understand areas.

I have added tests that prove my fix is effective or that my feature works.

New and existing unit tests pass locally with my changes.

I have manually tested my changes end-to-end.

Any dependent changes have been merged and published in downstream modules. (N/A — no downstream dependencies.)
Additional context
All features are opt-in and additive; no existing node behavior changes. Samples are included under contributing/samples/workflows/ (streaming_route, fan_out_preempt, search_fanout_first_answer, speculative_tool) and a guide at docs/guides/workflow/streaming_preemption/. Speculative dispatch must be used with idempotent/cancel-safe targets only.

Network-bound agent workloads spend most of their time awaiting the event
loop; swapping CPython's default loop for uvloop (libuv) is a cheap win but
previously required application code to manage loop policy itself.

Add google.adk.enable_uvloop(), a one-line, idempotent switch to install the
uvloop policy process-wide before the first Runner.run/asyncio.run. Deployments
can opt in without code changes via ADK_UVLOOP=1, which the sync Runner.run path
honours. uvloop is declared as an optional extra ("google-adk[uvloop]"; no
Windows wheels). uvloop only accelerates code that awaits on the loop, so it is
the last 10%, not a 10x on its own.
The LlmAgent-as-node wrapper only commits a node's output on the final,
non-partial event, so the graph always advances at turn granularity: the model
finishes generating, then the scheduler moves on. There was no way to act on a
decision the model has already made mid-turn.

StreamingRouterNode runs a wrapped agent in SSE mode and hands every streamed
delta to a caller-supplied monitor. When the monitor returns a StreamDecision,
the node commits the route/output and (by default) closes the model stream;
closing propagates GeneratorExit down the aclosing chain, cooperatively
cancelling the in-flight model call so the scheduler advances immediately.

Measured on real Vertex gemini-3.5-flash-lite reading five whole arXiv papers in
parallel ("is this an AI paper?"): preempting once the verdict streams in is
~3.5x faster and generates ~28x fewer output tokens than streaming each answer
to completion, while producing identical classifications. Preemption saves
generation, not the input prefill (the whole doc is still read); with context
caching the combined cost drop is ~58%.

Includes unit tests, a real-LLM timing/token/cost integration test, runnable
samples (streaming_route, fan_out_preempt), and a workflow guide.
Big-document sibling of the arXiv timing test. Pulls the latest Tesla annual
report (Form 10-K, ~100k+ input tokens) live from SEC EDGAR and hands it to
Gemini whole in one streaming call -- no chunking. Compares plain SSE streaming
against SSE + mid-stream preemption: the monitor cancels generation the instant
the VERDICT line streams in, so the long analysis is never decoded. On a filing
this size the input is paid once (prefill), making preemption a large wall-clock
win and, with context caching, a large cost win. Skips without Vertex creds or
SEC network access.
Adds two workflow primitives for latency- and cost-sensitive fan-out, plus
honest real-LLM integration tests that make the feature's output a required
deliverable so a naive baseline gets no shortcut.

- SpeculativeRouterNode: dispatches a downstream target from a partial,
  repaired tool call while the agent keeps streaming, then verifies against the
  finalized call (keep on hit, cancel + re-run on miss). A new `combine` hook
  lets the agent's full streamed text (e.g. a planner rationale) be a returned
  deliverable alongside the target result, so overlap wins are real, not from
  discarded output. Ships repair_json and make_marker_extractor helpers.
- FirstMatchNode: races several branches, returns the first whose output
  satisfies a match predicate, and cancels the losers (search Recall@k).
- Honest tests: the chained-LLM speculation test requires the planner rationale
  as output, and the arXiv preemption test keeps (and asserts) the relevant
  paper's full summary while only skipping the summaries it discards.
@adk-bot adk-bot added the core [Component] This issue is related to the core interface and implementation label Aug 6, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

core [Component] This issue is related to the core interface and implementation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants