Phase 3a: A2A message/stream + tasks/resubscribe via SSE - #50
Conversation
📝 WalkthroughWalkthroughThe changes add comprehensive Server-Sent Events (SSE) streaming support to the A2A JSON-RPC endpoint. A new Changes
Sequence DiagramsequenceDiagram
participant Client as Client (SSE)
participant Handler as Handler<br/>(handle_stream)
participant Registry as StreamRegistry
participant Task as Task<br/>(persistence)
participant Function as Function<br/>(trigger)
Client->>Handler: POST /a2a<br/>(message/stream)
Handler->>Task: Load or create task<br/>(Working state)
Handler->>Registry: subscribe(task_id,<br/>writer)
Registry-->>Handler: subscriber_id
Handler->>Client: [SSE] Submitted frame
Handler->>Client: [SSE] Working frame
Handler->>Handler: Validate function
Handler->>Function: Spawn trigger
par Function Execution
Function->>Function: Execute
and Task Broadcast
Function-->>Registry: broadcast(artifact)
Registry->>Client: [SSE] artifact-update
end
Function-->>Handler: Complete/Error
Handler->>Registry: broadcast(Completed<br/>or Failed, final=true)
Registry->>Client: [SSE] Final status frame
Handler->>Registry: close_task(task_id)
Handler->>Client: Close stream
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (6)
a2a/src/streaming.rs (4)
209-230: SSE preamble is missingX-Accel-Buffering: noandConnection: keep-alive.Many reverse proxies (nginx, several CDNs) buffer
text/event-streamresponses by default unlessX-Accel-Buffering: nois set, which delays — sometimes indefinitely — frames reaching the client until the response body is large enough to flush. Likewise some HTTP/1.1 stacks benefit from an explicitConnection: keep-alive. Adding both to theset_headersmap is cheap insurance for production deployments behind any L7 proxy.🛡️ Suggested headers
&serde_json::to_string(&json!({ "type": "set_headers", "headers": { "content-type": "text/event-stream", "cache-control": "no-cache", + "connection": "keep-alive", + "x-accel-buffering": "no", } }))🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@a2a/src/streaming.rs` around lines 209 - 230, The SSE preamble in send_sse_preamble currently sets content-type and cache-control but omits buffering/keep-alive signals; update the set_headers payload sent via ChannelWriter in send_sse_preamble to include "X-Accel-Buffering": "no" and "Connection": "keep-alive" so reverse proxies and HTTP/1.1 stacks do not buffer the event-stream; locate the set_headers JSON sent in send_sse_preamble and add those two header entries to the headers map before calling writer.send_message.
282-322:handle_streamignoresparams.message.task_idfor idempotent re-stream after terminal.If a client resends a streaming
message/sendwith an existing terminaltask_id(the documented idempotency contract for syncmessage/send), this path correctly emits a single terminalstatus-updateand closes. But unlike the sync path it doesn't return the storedtaskback to the caller in a JSON-RPC envelope — that's expected on a streaming socket, just confirming that's intentional. Worth a one-line comment near the early return so the parallel withhandle_send's terminal-idempotency branch is obvious.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@a2a/src/streaming.rs` around lines 282 - 322, handle_stream currently detects an existing terminal task via load_task and emits a terminal "status-update" then closes the writer, but unlike the sync handle_send path it does not return the stored task in a JSON-RPC envelope; add a one-line comment immediately before the early return in the existing-terminal branch (the block that writes the frame, closes the writer and returns) noting that this is intentional idempotent re-stream behavior — we emit a single terminal status-update and close the SSE stream instead of returning the full task object (to mirror handle_send’s idempotency contract for sync calls).
96-152: Broadcast pruning skips writers whose tasks panic, and a slow socket still blocksclose_task.Two minor points on the JoinSet fan-out:
The
Err(join_err)arm at 137-143 logs but doesn't capture the subscriber id — if the spawned task panics, that subscriber stays on the bus until the next failed write. Given each broadcast spawns one task per subscriber, a writer whosewrite()future panics consistently will be retried on every broadcast. Consider tagging each spawned task with itssub_id(e.g.,join_set.spawn(async move { (sub_id, writer.write(&frame).await) })) so the JoinError arm can also push todead.
broadcastawaits every spawned write before returning, so one stuck writer (TCP push-back, dead peer keeping the half-open socket) holds up the call. Sinceclose_taskruns strictly after the final broadcast inhandle_stream/handle_send, a single slow client can delay terminal notification to all other subscribers waiting onterminal_watch. A bounded write timeout per task (e.g., wrapwriter.writeintokio::time::timeout) would put a ceiling on the worst-case fan-out latency.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@a2a/src/streaming.rs` around lines 96 - 152, The broadcast loop in broadcast spawns a task per subscriber but neither associates the spawned task with its subscriber id nor bounds each write, so panics or stuck writes leave dead writers on the bus and a single slow client can block the whole fan-out; modify the spawn to return the subscriber id with the write result (e.g., spawn an async move that yields (sub_id, writer.write(&frame).await)) so the Err(join_err) arm can mark that sub_id as dead, and wrap the writer.write future in a tokio::time::timeout with a reasonable duration so timed-out writes return an error which you convert into marking the sub_id dead; ensure you still collect dead ids into the dead HashSet and call g.subscribers.retain(|s| !dead.contains(&s.id)) on the bus after join_set completes.
660-727: Streaming module's#[cfg(test)]block duplicates the integration tests ina2a/tests/streaming.rs.
task_status_update_serializes_camel_case_with_final,task_artifact_update_serializes_camel_case,resubscribe_params_round_trip, andsse_frame_formathere cover the same wire-shape assertions astask_status_update_uses_final_not_final_event/task_artifact_update_serializes_camel_case/resubscribe_params_round_trip/sse_frame_matches_a2a_layoutina2a/tests/streaming.rs. Pick one location (the integration test file is fine since it's already exercising the public surface) and drop the other to avoid drift.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@a2a/src/streaming.rs` around lines 660 - 727, The tests inside the #[cfg(test)] mod in streaming.rs duplicate integration tests in a2a/tests/streaming.rs; remove the duplicated unit tests here to avoid drift. Specifically, delete or comment out the test functions sse_frame_format, task_status_update_serializes_camel_case_with_final, task_artifact_update_serializes_camel_case, and resubscribe_params_round_trip (the entire #[cfg(test)] mod) so the project relies on the existing integration tests (e.g., task_status_update_uses_final_not_final_event, task_artifact_update_serializes_camel_case, resubscribe_params_round_trip, sse_frame_matches_a2a_layout) instead. Ensure no other non-test code in that module is removed.a2a/src/handler.rs (2)
421-444: Duplicated payload builders acrosshandler.rsandstreaming.rs.
status_update_payload/artifact_update_payloadhere mirrortask_status_update_payload/task_artifact_update_payloadina2a/src/streaming.rs(lines 232–257). Both serialize the sameTaskStatusUpdateEvent/TaskArtifactUpdateEventshapes with the same fallback-to-Value::Null. Consolidate into a single pair ofpub(crate)helpers instreaming.rs(ortypes.rs) and reuse from both call sites so the wire shape can't drift between sync and streaming flows.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@a2a/src/handler.rs` around lines 421 - 444, status_update_payload and artifact_update_payload duplicate task_status_update_payload/task_artifact_update_payload in streaming.rs; move the shared serializers into a single pub(crate) helper pair (e.g., task_status_update_payload and task_artifact_update_payload) placed in streaming.rs or a common types.rs, keep the exact serialized shapes (fields like append: None and last_chunk: Some(true) for artifacts and final_event for status) and the same unwrap_or(Value::Null) fallback, make them pub(crate) and replace calls in handler.rs to call the consolidated helpers so both sync and streaming code reuse the same functions.
596-603: Mid-flight cancel path drops the trigger result silently — confirm that's intentional.When
tasks/cancellands during theiii.trigger(...).await, this branch returns the canceled task and discardsresult. That matches the documented "result is discarded and the task keeps itsCanceledstate" semantic in the README, so behaviorally fine. Worth a single-line tracing log though, since otherwise a function that produced a real artifact gets dropped with no breadcrumb in the worker logs.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@a2a/src/handler.rs` around lines 596 - 603, The cancel branch in the Ok(result) path currently returns the canceled task and silently drops result; add a one-line tracing log before the return so there's a breadcrumb when a trigger-produced artifact is discarded. In the Ok(result) handling (around the load_task(...) / matches!(t.status.state, TaskState::Canceled) block), emit a concise trace/info via the tracing/logger used in the file (e.g., mentioning task_id and that the trigger result was discarded) immediately before the A2AResponse::success(...) return so callers can see that the trigger result was intentionally dropped.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@a2a/src/handler.rs`:
- Around line 187-190: The code calls iii_inner.address() which doesn't exist in
iii-sdk 0.11.3; change the call site to accept and use the engine WebSocket URL
passed into the handler (e.g., the engine_ws_base parameter) instead of trying
to call III::address(). Specifically, update the place creating the
ChannelWriter (the Arc::new(iii_sdk::ChannelWriter::new(...))) to pass the
engine_ws_base (or equivalent function parameter) as the first argument rather
than iii_inner.address(), and ensure the handler signature/callers provide that
engine_ws_base value like image-resize's handler does.
In `@a2a/src/streaming.rs`:
- Around line 367-384: store_task wrote the task as Working before emitting a
synthesized Submitted frame, so persisted state never shows Submitted; either
persist Submitted first, broadcast it, then transition the stored record to
Working and persist again (use store_task with the cloned submitted Task, call
registry.broadcast with task_status_update_payload(&submitted, false), then
update the original task to Working and call store_task + registry.broadcast for
Working), or if the wire-only emission is intentional add a one-line comment
near store_task/Submitted explaining the deliberate mismatch; refer to
store_task, Submitted, Working, task.clone, task_status_update_payload,
registry.broadcast and task_id to locate and implement the change.
- Around line 600-634: Race occurs when task becomes terminal between load_task
and registry.subscribe causing a new empty TaskBus to be created and the handler
to wait forever; after calling registry.subscribe (and after
reserve_event_id/replay), immediately re-load the task state via load_task (or
the same storage fetch used earlier) and if the reloaded task is terminal, build
and write the final frame with build_sse_frame(task_status_update_payload(...,
true)) and close the writer instead of entering the terminal_watch loop; ensure
you release/unsubscribe the subscriber_id if your registry has an
unsubscribe/remove method and avoid waiting on terminal_watch for a
newly-created bus when the persistent task is already terminal.
- Around line 202-207: The function writer_ref_from_input currently uses
iii_sdk::extract_channel_refs(...) and .find(...) which silently returns the
first ChannelDirection::Write it sees; change this to explicitly verify the SDK
contract by collecting all writable channels from extract_channel_refs(input)
(filtering by r.direction == ChannelDirection::Write) and then: if more than one
writable channel is found, emit a debug/error via your logger or panic/assert in
non-production builds to surface the contract violation; if exactly one, return
that StreamChannelRef; if zero, return None. Update writer_ref_from_input and
reference iii_sdk::extract_channel_refs and the
StreamChannelRef/ChannelDirection symbols in the fix and add a short comment
documenting the expected single-writable-channel contract or that the code
defends against violations.
---
Nitpick comments:
In `@a2a/src/handler.rs`:
- Around line 421-444: status_update_payload and artifact_update_payload
duplicate task_status_update_payload/task_artifact_update_payload in
streaming.rs; move the shared serializers into a single pub(crate) helper pair
(e.g., task_status_update_payload and task_artifact_update_payload) placed in
streaming.rs or a common types.rs, keep the exact serialized shapes (fields like
append: None and last_chunk: Some(true) for artifacts and final_event for
status) and the same unwrap_or(Value::Null) fallback, make them pub(crate) and
replace calls in handler.rs to call the consolidated helpers so both sync and
streaming code reuse the same functions.
- Around line 596-603: The cancel branch in the Ok(result) path currently
returns the canceled task and silently drops result; add a one-line tracing log
before the return so there's a breadcrumb when a trigger-produced artifact is
discarded. In the Ok(result) handling (around the load_task(...) /
matches!(t.status.state, TaskState::Canceled) block), emit a concise trace/info
via the tracing/logger used in the file (e.g., mentioning task_id and that the
trigger result was discarded) immediately before the A2AResponse::success(...)
return so callers can see that the trigger result was intentionally dropped.
In `@a2a/src/streaming.rs`:
- Around line 209-230: The SSE preamble in send_sse_preamble currently sets
content-type and cache-control but omits buffering/keep-alive signals; update
the set_headers payload sent via ChannelWriter in send_sse_preamble to include
"X-Accel-Buffering": "no" and "Connection": "keep-alive" so reverse proxies and
HTTP/1.1 stacks do not buffer the event-stream; locate the set_headers JSON sent
in send_sse_preamble and add those two header entries to the headers map before
calling writer.send_message.
- Around line 282-322: handle_stream currently detects an existing terminal task
via load_task and emits a terminal "status-update" then closes the writer, but
unlike the sync handle_send path it does not return the stored task in a
JSON-RPC envelope; add a one-line comment immediately before the early return in
the existing-terminal branch (the block that writes the frame, closes the writer
and returns) noting that this is intentional idempotent re-stream behavior — we
emit a single terminal status-update and close the SSE stream instead of
returning the full task object (to mirror handle_send’s idempotency contract for
sync calls).
- Around line 96-152: The broadcast loop in broadcast spawns a task per
subscriber but neither associates the spawned task with its subscriber id nor
bounds each write, so panics or stuck writes leave dead writers on the bus and a
single slow client can block the whole fan-out; modify the spawn to return the
subscriber id with the write result (e.g., spawn an async move that yields
(sub_id, writer.write(&frame).await)) so the Err(join_err) arm can mark that
sub_id as dead, and wrap the writer.write future in a tokio::time::timeout with
a reasonable duration so timed-out writes return an error which you convert into
marking the sub_id dead; ensure you still collect dead ids into the dead HashSet
and call g.subscribers.retain(|s| !dead.contains(&s.id)) on the bus after
join_set completes.
- Around line 660-727: The tests inside the #[cfg(test)] mod in streaming.rs
duplicate integration tests in a2a/tests/streaming.rs; remove the duplicated
unit tests here to avoid drift. Specifically, delete or comment out the test
functions sse_frame_format, task_status_update_serializes_camel_case_with_final,
task_artifact_update_serializes_camel_case, and resubscribe_params_round_trip
(the entire #[cfg(test)] mod) so the project relies on the existing integration
tests (e.g., task_status_update_uses_final_not_final_event,
task_artifact_update_serializes_camel_case, resubscribe_params_round_trip,
sse_frame_matches_a2a_layout) instead. Ensure no other non-test code in that
module is removed.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: b3271454-a18f-41d7-887b-68545d69d701
⛔ Files ignored due to path filters (1)
a2a/Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (8)
a2a/Cargo.tomla2a/README.mda2a/src/handler.rsa2a/src/lib.rsa2a/src/main.rsa2a/src/streaming.rsa2a/src/types.rsa2a/tests/streaming.rs
| let writer = Arc::new(iii_sdk::ChannelWriter::new( | ||
| iii_inner.address(), | ||
| &writer_ref, | ||
| )); |
There was a problem hiding this comment.
🧩 Analysis chain
🌐 Web query:
iii-sdk Rust crate 0.11.3 III struct address method
💡 Result:
The iii-sdk Rust crate version 0.11.3 exists (released 2026-04-22 per crates.io version history). It provides the core III struct (re-exported from the iii module), which is the main SDK client instance obtained via register_worker(address: &str, options: InitOptions) -> Result. The address parameter to register_worker is a WebSocket URL string (e.g., "ws://127.0.0.1:49134") for connecting to the III Engine. No evidence of an address method on the III struct itself. The III struct exposes methods like register_function, register_trigger, and trigger for registering/invoking functions and managing triggers. Related types include WorkerInfo and WorkerMetadata (re-exported from iii module), which may contain worker address/metadata, but no specific address method details found. Official documentation: https://docs.rs/iii-sdk/latest/iii_sdk/ (shows re-exports including III). Source repository: https://github.com/iii-hq/iii/tree/main/sdk (Rust SDK in packages/rust or similar). For full API, check docs.rs for exact version or GitHub source.
Citations:
- 1: https://docs.rs/crate/iii-sdk/latest
- 2: https://crates.io/crates/iii-sdk/0.11.0-next.4
- 3: https://crates.io/crates/iii-sdk/0.11.0-next.6
- 4: https://docs.rs/iii-sdk/latest/iii_sdk/
- 5: https://github.com/iii-hq/iii/blob/06a6580d46649804e7635bb3160d8f9709247e74/sdk/README.md
- 6: https://github.com/iii-hq/iii/blob/f997e40e35304e4aa7364cfa84fce40905ff1dfe/skills/iii-rust-sdk/SKILL.md
🏁 Script executed:
#!/bin/bash
# Check how other workers in this repo construct ChannelWriter to see the canonical engine_ws_base source.
rg -nP --type=rust -C3 'ChannelWriter::new\s*\('
# And how III exposes its address.
rg -nP --type=rust -C2 '\.address\s*\(\)'Repository: iii-hq/workers
Length of output: 1724
🏁 Script executed:
# Find where iii_inner is created in a2a
rg -nP --type=rust 'iii_inner' a2a/src/handler.rs | head -20Repository: iii-hq/workers
Length of output: 481
🏁 Script executed:
# Search for III struct definition or address method usage in the codebase
rg -nP --type=rust 'impl.*III|fn address\s*\(' | head -30Repository: iii-hq/workers
Length of output: 3723
🏁 Script executed:
# Check the full function signature of the streaming method to understand iii_inner's type
sed -n '170,200p' a2a/src/handler.rsRepository: iii-hq/workers
Length of output: 1569
🏁 Script executed:
# Search for iii_card and iii_rpc definitions
rg -nP --type=rust 'iii_card|iii_rpc' a2a/src/handler.rs | head -40Repository: iii-hq/workers
Length of output: 228
🏁 Script executed:
# Look at the register_worker call and surrounding context
sed -n '1,130p' a2a/src/handler.rsRepository: iii-hq/workers
Length of output: 4263
🏁 Script executed:
# Search for any trait implementations on III
rg -nP --type=rust 'impl.*III|fn address' a2a/src/Repository: iii-hq/workers
Length of output: 40
🏁 Script executed:
# Also check the Cargo.toml to see the exact iii-sdk version and check if there are any patches
cat a2a/Cargo.tomlRepository: iii-hq/workers
Length of output: 935
🏁 Script executed:
# Search for address method definition anywhere in the codebase
rg -nP --type=rust 'fn address\s*\(' --type-listRepository: iii-hq/workers
Length of output: 448
🏁 Script executed:
# Search the entire codebase for address() to see all usages
rg -nP --type=rust 'address\s*\('Repository: iii-hq/workers
Length of output: 126
🏁 Script executed:
# Check if there's a wrapper type around III or a trait
rg -nP --type=rust 'struct.*III|trait' a2a/src/ | grep -i addrRepository: iii-hq/workers
Length of output: 40
🏁 Script executed:
# Let's try a web search for iii-sdk documentation on III::addressRepository: iii-hq/workers
Length of output: 40
III::address() method does not exist in iii-sdk 0.11.3; code at line 187-190 will not compile.
The III struct in iii-sdk 0.11.3 has no .address() method. Web search and official documentation (docs.rs, crates.io) show only methods like register_function, register_trigger, and trigger; no address getter is exposed. This is the only place in the codebase calling .address().
Compare with image-resize/src/handler.rs:36, which correctly receives engine_ws_base as a function parameter and passes it directly to ChannelWriter::new(). Pass the engine WebSocket URL the same way instead of attempting to retrieve it via a non-existent method.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@a2a/src/handler.rs` around lines 187 - 190, The code calls
iii_inner.address() which doesn't exist in iii-sdk 0.11.3; change the call site
to accept and use the engine WebSocket URL passed into the handler (e.g., the
engine_ws_base parameter) instead of trying to call III::address().
Specifically, update the place creating the ChannelWriter (the
Arc::new(iii_sdk::ChannelWriter::new(...))) to pass the engine_ws_base (or
equivalent function parameter) as the first argument rather than
iii_inner.address(), and ensure the handler signature/callers provide that
engine_ws_base value like image-resize's handler does.
| pub fn writer_ref_from_input(input: &Value) -> Option<StreamChannelRef> { | ||
| iii_sdk::extract_channel_refs(input) | ||
| .into_iter() | ||
| .find(|(_, r)| matches!(r.direction, ChannelDirection::Write)) | ||
| .map(|(_, r)| r) | ||
| } |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Find canonical usage of extract_channel_refs in this repo and any docs.
rg -nP --type=rust -C5 'extract_channel_refs'
fd -e md | xargs rg -n -C2 'extract_channel_refs|writer.?ref|StreamChannelRef' 2>/dev/nullRepository: iii-hq/workers
Length of output: 3270
🏁 Script executed:
# Check for HTTP trigger specifications or tests
fd -e rs -e md | xargs rg -l 'http.*trigger|trigger.*http|HTTP' 2>/dev/null | head -20Repository: iii-hq/workers
Length of output: 508
🏁 Script executed:
# Look for tests related to extract_channel_refs or writer_ref_from_input
rg -n 'test.*extract_channel_refs|test.*writer_ref|#\[test\]' a2a/src/streaming.rs -B2 -A10Repository: iii-hq/workers
Length of output: 1673
🏁 Script executed:
# Search for iii-sdk version specification and any contracts/docs
fd -e toml | xargs rg -n 'iii.sdk|iii_sdk' 2>/dev/nullRepository: iii-hq/workers
Length of output: 716
🏁 Script executed:
# Check for any documentation about channel constraints
fd -e md | xargs rg -n 'one.*writable|single.*writable|writable.*channel.*one' 2>/dev/nullRepository: iii-hq/workers
Length of output: 42
🏁 Script executed:
# Find all tests in a2a/src/streaming.rs
rg -n '#\[test\]' a2a/src/streaming.rs -A20Repository: iii-hq/workers
Length of output: 2449
🏁 Script executed:
# Check if there are HTTP trigger tests or docs in a2a
fd -e rs -e md a2a/ | xargs rg -n 'writer_ref_from_input|HTTP.*trigger|http_trigger' 2>/dev/nullRepository: iii-hq/workers
Length of output: 3796
🏁 Script executed:
# Look for the full function context and any assertions about channel count
rg -n 'writer_ref|extract_channel_refs' a2a/src/ -B3 -A3Repository: iii-hq/workers
Length of output: 1947
🏁 Script executed:
# Search for any validation or assertions about multiple writable channels
rg -n 'writable|Write|multiple.*channel|one.*channel' a2a/src/ --type=rust | head -30Repository: iii-hq/workers
Length of output: 1169
🏁 Script executed:
# Check if there's any documentation about HTTP trigger format/structure
fd . a2a/ -type f \( -name '*.md' -o -name '*.rs' \) | xargs rg -n 'trigger.*input|input.*envelope|channel.*ref' 2>/dev/null | head -20Repository: iii-hq/workers
Length of output: 1721
🏁 Script executed:
# Check if there are integration tests or examples showing HTTP trigger payloads
fd . a2a/ -type f -name '*.rs' | xargs rg -n 'json!' a2a/src/handler.rs | head -10Repository: iii-hq/workers
Length of output: 855
🏁 Script executed:
# Look for any documentation about the iii-sdk extract_channel_refs behavior
rg -n 'extract_channel_refs' --type=rust -B2 -A8 image-resize/src/handler.rsRepository: iii-hq/workers
Length of output: 1334
Add defensive code or verify SDK contract for single writable channel.
The code assumes HTTP triggers surface exactly one writable channel per request (documented in the comment), but this assumption is neither verified against the iii-sdk 0.11.3 contract nor defended in code. Using .find() to pick the first writable channel will silently select the wrong one if the assumption is violated.
Either:
- Confirm in SDK 0.11.3 documentation that HTTP triggers guarantee exactly one writable channel per request, and document that guarantee in code, or
- Add a debug-log or assert when multiple writable channels are found to catch contract violations early.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@a2a/src/streaming.rs` around lines 202 - 207, The function
writer_ref_from_input currently uses iii_sdk::extract_channel_refs(...) and
.find(...) which silently returns the first ChannelDirection::Write it sees;
change this to explicitly verify the SDK contract by collecting all writable
channels from extract_channel_refs(input) (filtering by r.direction ==
ChannelDirection::Write) and then: if more than one writable channel is found,
emit a debug/error via your logger or panic/assert in non-production builds to
surface the contract violation; if exactly one, return that StreamChannelRef; if
zero, return None. Update writer_ref_from_input and reference
iii_sdk::extract_channel_refs and the StreamChannelRef/ChannelDirection symbols
in the fix and add a short comment documenting the expected
single-writable-channel contract or that the code defends against violations.
Adds A2A v0.3 streaming for the a2a worker. Both SSE endpoints emit TaskStatusUpdateEvent and TaskArtifactUpdateEvent through a shared process-local StreamRegistry, so a sync `message/send` or `tasks/cancel` broadcasts to concurrent stream subscribers without polling. - New `streaming` module with StreamRegistry (stable u64 subscriber ids, per-subscriber event counters, JoinSet-based fan-out, prune-by-id). - `message/stream` walks Submitted (wire-only) -> Working -> artifact -> Completed/Failed. - `tasks/resubscribe` replays the current state then forwards broadcasts until terminal. Race guard re-loads the task post-subscribe and synthesizes a final frame if the producer transitioned in between. - handler::register now plumbs an Arc<StreamRegistry>; the JSON-RPC dispatch closure snapshots the writable channel ref before stripping the body so streaming methods can hand it to ChannelWriter::new. - AgentCapabilities.streaming flipped to true. - Cargo.toml bumped to 0.3.4. Tests: 13 unit tests (frame layout, A2A v0.3 wire fields, registry mechanics) and 3 #[ignore]'d e2e tests documenting the live-engine contract.
54222f5 to
f5d7762
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
a2a/src/handler.rs (2)
165-166: Nit: the intermediateregistrybinding is unused.
registryon line 165 is only used to produceregistry_rpc; nothing else inregister()references it. Collapsing to a single binding is a one-line cleanup.- let registry = Arc::new(StreamRegistry::new()); - let registry_rpc = registry.clone(); + let registry_rpc = Arc::new(StreamRegistry::new());🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@a2a/src/handler.rs` around lines 165 - 166, The local variable `registry` is unused and only cloned immediately to produce `registry_rpc` inside the `register()` function; remove the intermediate binding and create `registry_rpc` directly as Arc::new(StreamRegistry::new()). Update the code to eliminate `let registry = ...` and replace `let registry_rpc = registry.clone();` with a single `let registry_rpc = Arc::new(StreamRegistry::new());` so only the used symbol (`registry_rpc`) remains.
209-245: Optional: extract the "no writable channel" error envelope.The same 8-line
A2AResponse::errorenvelope is repeated formessage/streamandtasks/resubscribe. A small helper (or a single check before the match) would tighten this and make it harder for the two arms to drift apart in error code/message.♻️ Sketch
+ let no_writer_response = |id: Option<Value>| { + json!({ + "status_code": 200, + "headers": { "content-type": "application/json" }, + "body": A2AResponse::error( + id, + -32004, + "Streaming not supported on this transport (no writable channel)" + ) + }) + }; + match request.method.as_str() { "message/stream" | "SendStreamingMessage" => { let Some(writer_ref) = writer_ref else { - return Ok(json!({ - "status_code": 200, - "headers": { "content-type": "application/json" }, - "body": A2AResponse::error( - request.id, - -32004, - "Streaming not supported on this transport (no writable channel)" - ) - })); + return Ok(no_writer_response(request.id)); }; dispatch_stream(&iii_inner, request.params, writer_ref, registry, cfg) .await; return Ok(Value::Null); } "tasks/resubscribe" | "SubscribeToTask" => { let Some(writer_ref) = writer_ref else { - return Ok(json!({ ... duplicate ... })); + return Ok(no_writer_response(request.id)); }; ... } _ => {} }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@a2a/src/handler.rs` around lines 209 - 245, Extract the repeated "no writable channel" error envelope into a small helper or pre-check to avoid duplication: add a function or closure (e.g., ensure_writable_channel) that takes request.id and returns the standardized A2AResponse::error Value used when writer_ref is None, then call that helper from both the "message/stream" / "SendStreamingMessage" and "tasks/resubscribe" / "SubscribeToTask" arms (before calling dispatch_stream or dispatch_resubscribe) or perform the check once before the match and reuse the same error Value; reference writer_ref, A2AResponse::error, dispatch_stream, dispatch_resubscribe, and request.id when implementing the change.a2a/src/streaming.rs (1)
757-769: Deadtouch_taskhelper — drop or wire it up.
touch_taskis#[allow(dead_code)]and never called;store_taskinhandler.rsalready does the equivalent state write. The "kept for symmetry with future store-side uses" comment will rot — easier to delete and re-add when the future use materialises.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@a2a/src/streaming.rs` around lines 757 - 769, The helper touch_task (async fn touch_task(iii: &III, task: &Task)) is dead code and should be removed to avoid rot since store_task in handler.rs already performs the same state write; delete the touch_task function, its #[allow(dead_code)] and the accompanying comment. Alternatively, if you intend to keep the behavior, wire it up by invoking touch_task from the existing store_task or the place where tasks are persisted (passing the III and Task), instead of leaving it unused.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@a2a/src/streaming.rs`:
- Around line 495-548: This branch can overwrite a concurrently-set Canceled
state (TOCTOU); before calling store_task(&iii_run, &completed) and broadcasting
the Completed status, reload the task (using load_task(&iii_run,
&task_id_run).await) and if its status.state is now TaskState::Canceled (or any
terminal non-Completed state) skip the store_task and final
registry_run.broadcast of EVENT_STATUS; alternatively, if you intend to accept
this race, add a one-line comment next to the existing load_task check
referencing the accepted design tradeoff and why no CAS is available.
- Around line 477-593: The spawned tokio::spawn future running iii.trigger can
panic or be dropped and currently never guarantees registry.close_task is
called, leaving terminal_watch waiting and the SSE open; fix by retaining the
JoinHandle returned by tokio::spawn (instead of dropping it) and concurrently
select!ing the JoinHandle against the registry.terminal_watch() watcher in the
parent task: if the JoinHandle returns Err(JoinError) or finishes unexpectedly,
call registry.close_task(&task_id).await and break the wait loop so
writer.close() runs; ensure the same cleanup is invoked when the cancel-branch
returns early (or alternatively add a small drop-guard tied to the spawned task
that invokes registry_run.close_task(&task_id_run) if the spawned future is
dropped) so registry.close_task is always executed even on panics in
iii.trigger, serde_json::to_string_pretty, or load_task.
---
Nitpick comments:
In `@a2a/src/handler.rs`:
- Around line 165-166: The local variable `registry` is unused and only cloned
immediately to produce `registry_rpc` inside the `register()` function; remove
the intermediate binding and create `registry_rpc` directly as
Arc::new(StreamRegistry::new()). Update the code to eliminate `let registry =
...` and replace `let registry_rpc = registry.clone();` with a single `let
registry_rpc = Arc::new(StreamRegistry::new());` so only the used symbol
(`registry_rpc`) remains.
- Around line 209-245: Extract the repeated "no writable channel" error envelope
into a small helper or pre-check to avoid duplication: add a function or closure
(e.g., ensure_writable_channel) that takes request.id and returns the
standardized A2AResponse::error Value used when writer_ref is None, then call
that helper from both the "message/stream" / "SendStreamingMessage" and
"tasks/resubscribe" / "SubscribeToTask" arms (before calling dispatch_stream or
dispatch_resubscribe) or perform the check once before the match and reuse the
same error Value; reference writer_ref, A2AResponse::error, dispatch_stream,
dispatch_resubscribe, and request.id when implementing the change.
In `@a2a/src/streaming.rs`:
- Around line 757-769: The helper touch_task (async fn touch_task(iii: &III,
task: &Task)) is dead code and should be removed to avoid rot since store_task
in handler.rs already performs the same state write; delete the touch_task
function, its #[allow(dead_code)] and the accompanying comment. Alternatively,
if you intend to keep the behavior, wire it up by invoking touch_task from the
existing store_task or the place where tasks are persisted (passing the III and
Task), instead of leaving it unused.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 8176bba1-9b35-467b-8595-3b142072ea8e
⛔ Files ignored due to path filters (1)
a2a/Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (7)
a2a/Cargo.tomla2a/README.mda2a/src/handler.rsa2a/src/lib.rsa2a/src/streaming.rsa2a/src/types.rsa2a/tests/streaming.rs
✅ Files skipped from review due to trivial changes (1)
- a2a/Cargo.toml
🚧 Files skipped from review as they are similar to previous changes (3)
- a2a/src/lib.rs
- a2a/README.md
- a2a/tests/streaming.rs
| // Spawn the actual function call. Runs concurrently with any further | ||
| // broadcasts (e.g. from a sibling `tasks/cancel`). | ||
| let iii_run = iii.clone(); | ||
| let registry_run = registry.clone(); | ||
| let task_id_run = task_id.clone(); | ||
| let context_id_run = task.context_id.clone(); | ||
| let metadata_run = task.metadata.clone(); | ||
| let history_run = task.history.clone(); | ||
| tokio::spawn(async move { | ||
| let result = iii_run | ||
| .trigger(TriggerRequest { | ||
| function_id, | ||
| payload, | ||
| action: None, | ||
| timeout_ms: Some(30000), | ||
| }) | ||
| .await; | ||
|
|
||
| // Cancel-while-running: if the task is now Canceled, don't | ||
| // overwrite that with Completed. The cancel path has already | ||
| // broadcast its own final frame. | ||
| let fresh = load_task(&iii_run, &task_id_run).await; | ||
| if let Some(ref t) = fresh | ||
| && matches!(t.status.state, TaskState::Canceled) | ||
| { | ||
| return; | ||
| } | ||
|
|
||
| match result { | ||
| Ok(value) => { | ||
| let result_text = | ||
| serde_json::to_string_pretty(&value).unwrap_or_else(|_| value.to_string()); | ||
| let artifact = Artifact { | ||
| artifact_id: uuid::Uuid::new_v4().to_string(), | ||
| parts: vec![Part { | ||
| text: Some(result_text), | ||
| data: None, | ||
| url: None, | ||
| raw: None, | ||
| media_type: Some("application/json".to_string()), | ||
| }], | ||
| name: Some(fn_name), | ||
| metadata: None, | ||
| }; | ||
|
|
||
| let artifact_payload = artifact_event_payload( | ||
| &task_id_run, | ||
| context_id_run.as_deref(), | ||
| artifact.clone(), | ||
| ); | ||
| registry_run | ||
| .broadcast(&task_id_run, EVENT_ARTIFACT, &artifact_payload) | ||
| .await; | ||
|
|
||
| let completed = Task { | ||
| id: task_id_run.clone(), | ||
| context_id: context_id_run.clone(), | ||
| status: TaskStatus { | ||
| state: TaskState::Completed, | ||
| message: None, | ||
| timestamp: Some(iso_now()), | ||
| }, | ||
| artifacts: Some(vec![artifact]), | ||
| history: history_run.clone(), | ||
| metadata: metadata_run.clone(), | ||
| }; | ||
| store_task(&iii_run, &completed).await; | ||
| let final_payload = status_event_payload(&completed, TaskState::Completed, true); | ||
| registry_run | ||
| .broadcast(&task_id_run, EVENT_STATUS, &final_payload) | ||
| .await; | ||
| } | ||
| Err(err) => { | ||
| let failed = Task { | ||
| id: task_id_run.clone(), | ||
| context_id: context_id_run.clone(), | ||
| status: TaskStatus { | ||
| state: TaskState::Failed, | ||
| message: Some(Message { | ||
| message_id: msg_id(), | ||
| role: MessageRole::Agent, | ||
| parts: vec![text_part(format!("Error: {}", err))], | ||
| task_id: None, | ||
| context_id: None, | ||
| metadata: None, | ||
| }), | ||
| timestamp: Some(iso_now()), | ||
| }, | ||
| artifacts: None, | ||
| history: history_run.clone(), | ||
| metadata: metadata_run.clone(), | ||
| }; | ||
| store_task(&iii_run, &failed).await; | ||
| let final_payload = status_event_payload(&failed, TaskState::Failed, true); | ||
| registry_run | ||
| .broadcast(&task_id_run, EVENT_STATUS, &final_payload) | ||
| .await; | ||
| } | ||
| } | ||
|
|
||
| registry_run.close_task(&task_id_run).await; | ||
| }); | ||
|
|
||
| // Wait for the producer to signal terminal, then close our writer so | ||
| // the SSE response completes cleanly. | ||
| if let Some(mut rx) = registry.terminal_watch(&task_id).await { | ||
| loop { | ||
| if *rx.borrow() { | ||
| break; | ||
| } | ||
| if rx.changed().await.is_err() { | ||
| break; | ||
| } | ||
| } | ||
| } | ||
| let _ = writer.close().await; | ||
| } |
There was a problem hiding this comment.
Spawned trigger task can hang the SSE response on panic.
The trigger driver is tokio::spawn'd (line 485) and the JoinHandle is dropped, so a panic anywhere inside the future — including serde_json::to_string_pretty, load_task deserialization, or any future panic in iii.trigger — never reaches close_task at line 577. The bus is left in buses with terminal_tx alive, and the parent's loop at 582-591 (rx.changed().await) blocks indefinitely because the watch sender is never sent and never dropped. The SSE socket then stays open with no further frames.
Two low-cost mitigations:
- Hold the
JoinHandleandselect!it againstterminal_watch; onErr(JoinError)(panic), forceregistry.close_task(...)and break. - Or use a small drop-guard that calls
close_taskon the registry when the task future is dropped (covers panic and explicit early returns like the cancel branch at 502, which today also skipsclose_task).
The cancel-branch early return at line 502 has the same shape: it relies on the cancel path having already called close_task, but if cancel raced in a way that skipped it (e.g., cancel saw the task still Working and didn't close because of any future code change), the same hang reappears. A drop-guard would make this self-healing.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@a2a/src/streaming.rs` around lines 477 - 593, The spawned tokio::spawn future
running iii.trigger can panic or be dropped and currently never guarantees
registry.close_task is called, leaving terminal_watch waiting and the SSE open;
fix by retaining the JoinHandle returned by tokio::spawn (instead of dropping
it) and concurrently select!ing the JoinHandle against the
registry.terminal_watch() watcher in the parent task: if the JoinHandle returns
Err(JoinError) or finishes unexpectedly, call
registry.close_task(&task_id).await and break the wait loop so writer.close()
runs; ensure the same cleanup is invoked when the cancel-branch returns early
(or alternatively add a small drop-guard tied to the spawned task that invokes
registry_run.close_task(&task_id_run) if the spawned future is dropped) so
registry.close_task is always executed even on panics in iii.trigger,
serde_json::to_string_pretty, or load_task.
| // Cancel-while-running: if the task is now Canceled, don't | ||
| // overwrite that with Completed. The cancel path has already | ||
| // broadcast its own final frame. | ||
| let fresh = load_task(&iii_run, &task_id_run).await; | ||
| if let Some(ref t) = fresh | ||
| && matches!(t.status.state, TaskState::Canceled) | ||
| { | ||
| return; | ||
| } | ||
|
|
||
| match result { | ||
| Ok(value) => { | ||
| let result_text = | ||
| serde_json::to_string_pretty(&value).unwrap_or_else(|_| value.to_string()); | ||
| let artifact = Artifact { | ||
| artifact_id: uuid::Uuid::new_v4().to_string(), | ||
| parts: vec![Part { | ||
| text: Some(result_text), | ||
| data: None, | ||
| url: None, | ||
| raw: None, | ||
| media_type: Some("application/json".to_string()), | ||
| }], | ||
| name: Some(fn_name), | ||
| metadata: None, | ||
| }; | ||
|
|
||
| let artifact_payload = artifact_event_payload( | ||
| &task_id_run, | ||
| context_id_run.as_deref(), | ||
| artifact.clone(), | ||
| ); | ||
| registry_run | ||
| .broadcast(&task_id_run, EVENT_ARTIFACT, &artifact_payload) | ||
| .await; | ||
|
|
||
| let completed = Task { | ||
| id: task_id_run.clone(), | ||
| context_id: context_id_run.clone(), | ||
| status: TaskStatus { | ||
| state: TaskState::Completed, | ||
| message: None, | ||
| timestamp: Some(iso_now()), | ||
| }, | ||
| artifacts: Some(vec![artifact]), | ||
| history: history_run.clone(), | ||
| metadata: metadata_run.clone(), | ||
| }; | ||
| store_task(&iii_run, &completed).await; | ||
| let final_payload = status_event_payload(&completed, TaskState::Completed, true); | ||
| registry_run | ||
| .broadcast(&task_id_run, EVENT_STATUS, &final_payload) | ||
| .await; | ||
| } |
There was a problem hiding this comment.
Cancel-during-running TOCTOU can clobber a Canceled state with Completed.
Sequence:
- Trigger completes successfully (line 493).
load_taskat 498 returnsWorking(cancel hasn't fired yet).- Concurrent
tasks/cancelruns: storesCanceled, broadcasts the final frame, callsclose_task. - This branch proceeds to
store_task(completed)at 543, silently overwriting the persistedCanceledstate, and broadcasts aCompletedframe to a now-empty bus.
End user sees a canceled final frame on the original stream (correct), but tasks/get later returns completed (incorrect) and any race-losing resubscriber loads completed from storage.
Per prior repo learnings, the iii engine's state API has no atomic CAS, so this exact TOCTOU pattern has been accepted elsewhere (sandbox capacity check). If you want to keep that posture here, a one-line code comment near line 498 explicitly noting the accepted race would prevent it being read as a bug later. Otherwise, the cleanest fix is to make the post-trigger update conditional: re-check state immediately before store_task and skip the write if it's now terminal.
Based on learnings: "the iii engine has no atomic reserve/CAS API for global enforcement across worker instances ... this is documented in the code by comment and is a known, accepted design tradeoff."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@a2a/src/streaming.rs` around lines 495 - 548, This branch can overwrite a
concurrently-set Canceled state (TOCTOU); before calling store_task(&iii_run,
&completed) and broadcasting the Completed status, reload the task (using
load_task(&iii_run, &task_id_run).await) and if its status.state is now
TaskState::Canceled (or any terminal non-Completed state) skip the store_task
and final registry_run.broadcast of EVENT_STATUS; alternatively, if you intend
to accept this race, add a one-line comment next to the existing load_task check
referencing the accepted design tradeoff and why no CAS is available.
1 participant
Tracks #45 — Phase 3a of MCP+A2A overhaul.
Summary
A2A v0.3 streaming over SSE using iii-sdk `ChannelWriter`:
Reference SDK pattern: `iii/sdk/packages/rust/iii/tests/api_triggers.rs:498`. Note: SDK exposes `ChannelWriter::new(addr, &ref)` (lazy connect), not `connect()` — README documents this.
Test plan
Open issues (must fix before merge)
Sequencing
Needs rebase onto phase1 once it lands — uses `ExposureConfig` / `is_always_hidden_pub` / `is_exposed_pub` / `resolve_function_pub` from `crate::handler`, all of which Phase 1 removes.
Summary by CodeRabbit
Release Notes v0.3.4
New Features
Documentation
Chores