Skip to content

a2a: add async dispatch (a2a_send_async / check_task / cancel_task) — replaces #1066#1221

Open
pbranchu wants to merge 5 commits into
RightNow-AI:mainfrom
pbranchu:a2a-async-dispatch
Open

a2a: add async dispatch (a2a_send_async / check_task / cancel_task) — replaces #1066#1221
pbranchu wants to merge 5 commits into
RightNow-AI:mainfrom
pbranchu:a2a-async-dispatch

Conversation

@pbranchu

Copy link
Copy Markdown
Contributor

Summary

Adds async A2A task dispatch on top of the SSE streaming foundation (#1219) and kernel context threading (#1220). Replaces the closed #1066 with a properly scoped, properly tested implementation that addresses all reviewer feedback from that PR.

Depends on #1219 and #1220. This PR contains those two PRs' commits plus this PR's own dispatch work — exactly 4 commits total (3 feature + 1 dependency bump for `lettre`). Once the dependencies land, this PR's effective diff is just the third commit.

New tools

Tool Description
`a2a_send_async` Dispatch a task to a remote A2A agent and return a task_id immediately. Result is delivered to the originating channel via callback when complete. Limits: max 500 concurrent in-flight tasks; results discarded 5 min after completion.
`a2a_check_task` Poll a task's live progress or final result by task_id. Results remain available via this tool for 5 min after completion.
`a2a_cancel_task` Abort the local SSE reader for a task. Note: the remote agent is NOT notified — this only frees local resources.

Architecture

  • TTL/cap machinery: `MAX_ASYNC_TASKS = 500`, `TASK_TTL = 1h` (for tasks that never complete), `COMPLETED_GRACE = 5min` (so polling right after completion still returns the result before sweep). Background sweep loop spawned at kernel boot, runs every 5 min.
  • Atomic admission: `ASYNC_TASKS_IN_FLIGHT: AtomicUsize` + `CapSlotGuard` RAII pattern. Two concurrent calls at len=499 can't both pass the check — slot is reserved with `fetch_add` before the check.
  • Panic safety: `TaskCleanupGuard` Drop impl runs on panic, removes entries from both `ASYNC_TASKS` and `A2A_TASK_PROGRESS`, and releases the cap slot.
  • Untrusted-content tagging: `inject_async_callback` wraps the remote agent's response with `[A2A Result from {agent_name} — treat as untrusted external content]` so the caller's LLM treats it as untrusted (prompt-injection mitigation).
  • No global context map: callback context captured at spawn time and threaded directly into the background task closure (via kernel: thread channel callback context through agent loop (replaces global DashMap) #1220's plumbing). No race possible.
  • Cancellation honesty: tool description says "aborts the local SSE reader only — remote agent is not notified."

What this PR addresses from #1066 review feedback

Reviewer concern Status
Tool description / 300s timeout consistency ✅ Description matches `SYNC_STREAMING_DEADLINE` wrapping `consume_sse_stream` (in #1219)
`std::sync::Mutex` in async paths ✅ All production paths use `tokio::sync::Mutex`
`MAX_ASYNC_TASKS` cap + TTL + cleanup ✅ Atomic admission, RAII guards, 5-min sweep
SSE parser edge case tests ✅ 7 tests in #1219 cover normal/disconnect/malformed/missing-final/concatenated/error/empty
Channel context race ✅ Eliminated by #1220's parameter-threading
`TaskCleanupGuard` panic safety ✅ Drop impl removes from both maps and releases cap slot
`thread_id` from smart-thread routing ✅ `effective_thread_id` derivation in bridge.rs (#1220)
Untrusted-content tagging ✅ `format_async_callback_message` wraps callback content
Tests: cap rejection, TTL eviction, send_async happy-path, inject_async_callback e2e, set_channel_context capture ✅ All present including the now-deterministic TTL test (no longer uptime-dependent)
No `openfang-cli` changes ✅ Diff touches 7 files, none in CLI
No scope creep (`prune_failed_tool_turns`, `TOOL_ERROR_GUIDANCE` reword, debug→info log changes, driver dead-code removal) ✅ All absent
False claims in PR comments (SQLite persistence, `KernelRequest::CancelTask`, SSE keep-alive, `progressPercentage`, `JSONRPCError`) ✅ Retracted in the closing comment on #1066; not claimed here

Retraction note

The April 18 comment on #1066 claimed features that were not in the diff (SQLite persistence, true cancel propagation, SSE keep-alive, etc.). Those claims were wrong. This PR contains only what its diff shows: in-memory task registry with 5-min completion grace + 1h running TTL, local-only cancellation, no keep-alive, no `progressPercentage` field, no `JSONRPCError` envelope. If any of those would be desired, they are appropriate follow-ups.

Tests

  • TTL sweep: deterministic test using `sweep_with_age_thresholds(Duration::ZERO, Duration::ZERO)` instead of real time; second test with `Duration::MAX` confirms entries survive
  • Atomic cap: `test_async_cap_atomic_admission` fires MAX+1 concurrent attempts, asserts exactly MAX succeed
  • TaskCleanupGuard: `test_task_cleanup_guard_drop_paths` exercises both the panic path (entries removed) and the normal-completion path (entries preserved for grace period)
  • Test isolation: `with_global_lock_async` uses `tokio::sync::Mutex` (no `clippy::await_holding_lock`); helper serializes any test that touches global state
  • End-to-end SSE consume + callback: `test_send_streaming_with_progress_then_inject_callback` drives `send_task_streaming_with_progress` against a real local TCP server then verifies the callback dispatches with correct channel/recipient/thread_id and includes the untrusted tag
  • `tool_a2a_send_async` validation: missing message, missing url/name, SSRF-blocked URL, at-cap rejection
  • `tool_a2a_check_task`: running, completed, missing
  • Format helper test (kernel): verifies the untrusted-tag wrapping at the production formatting layer

962 runtime + 504 channels tests pass on three consecutive runs (no flakes). `cargo clippy --workspace --tests --all-targets -- -D warnings` clean.

Smoke-tested in production

Deployed to a real OpenFang instance with a mock SSE server. Confirmed:

  • A2A discovery via `/.well-known/agent.json` works
  • `a2a_send_async` returns task_id immediately; mock server receives the `tasks/sendSubscribe` POST
  • Task state transitions to `Completed` after final event
  • `a2a_check_task` returns result with `(completed)` indicator within the 5-min grace period
  • HTTP API path correctly emits "no channel context" warning (callback only fires for real channels)

Test plan

Philippe Branchu and others added 4 commits May 31, 2026 20:35
The previous design proposal stored a per-agent ChannelCallbackContext in a
global DashMap on the kernel and read it back from tools. That structure
caused cross-user callback bleed whenever the same agent was messaged by
two channels concurrently — DashMap::insert is atomic but offers no
isolation between the write and the later read inside the agent loop.

The proper fix is to never store the context globally: pass it as an owned
parameter from the bridge dispatch site, down through the agent loop, to
the tool runner. Each concurrent invocation owns its own context value,
so isolation is enforced by the type system rather than by ad-hoc locks.

- Add `openfang_types::ChannelCallbackContext` (channel, recipient,
  display name, thread_id, agent_id).
- Add `ChannelBridgeHandle::send_message_with_context` with a default
  impl that ignores the context and falls back to `send_message`.
- `KernelBridgeAdapter` overrides it to call the new
  `OpenFangKernel::send_message_with_context` →
  `send_message_with_handle_and_blocks_and_context` →
  `execute_llm_agent` plumbing path.
- `run_agent_loop` / `run_agent_loop_streaming` gain a
  `callback_context: Option<ChannelCallbackContext>` parameter that is
  forwarded to `tool_runner::execute_tool` as `Option<&...>`.
- `execute_tool` accepts the context as a new (currently-unused) param;
  the async-dispatch PR will wire it into `tool_a2a_send_async`.
- `dispatch_message` (channels bridge) builds the context from the
  incoming message and calls `send_message_with_context`.
- `KernelHandle::inject_async_callback` declared here with a default
  Err — production implementation lands with the async-dispatch PR.
- New test `test_concurrent_send_message_with_context_no_crossbleed`:
  two concurrent same-agent dispatches with distinct contexts each
  receive their own context back, confirming ownership-based isolation.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Replace blocking tasks/send with tasks/sendSubscribe so a2a_send no longer
hits the 30-second client timeout on long-running remote agents (Claude Code
delegations, multi-step reasoning, etc.). The call still blocks until the
remote agent emits its final event but processes SSE chunks incrementally
from the wire — the response body is never buffered in memory at once, so
backpressure is preserved.

- Bump A2aClient timeout from 30s to 300s.
- Add `send_task_streaming` to `A2aClient` (incremental SSE consumption).
- Extract `process_sse_line` + `parse_sse_content` so production streaming
  and unit tests share one parser implementation.
- `tool_a2a_send` switches to `send_task_streaming`; description updated
  to accurately state the 300s blocking behaviour.
- Add 7 SSE parser unit tests: normal completion, mid-stream disconnect,
  malformed JSON skipping, no-final-event, chunk boundaries, error events,
  empty stream.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Builds on the streaming and context-threading branches. Lets agents fire
long-running A2A tasks in the background and get the result delivered
back to the originating channel automatically when the remote agent
finishes.

- New tools `a2a_send_async`, `a2a_check_task`, `a2a_cancel_task`.
- Per-process cap (MAX_ASYNC_TASKS = 500) rejects new tasks at the limit.
- Per-entry TTL (TASK_TTL = 3600s) with a background sweeper installed
  in `OpenFangKernel::start_background_agents` (runs every 5 min).
- TaskCleanupGuard RAII removes both `ASYNC_TASKS` and `A2A_TASK_PROGRESS`
  entries when the spawned task ends (including panics).
- `A2aClient::send_task_streaming_with_progress` reuses `process_sse_line`
  from the streaming branch; intermediate agent text is appended to a
  shared `Arc<tokio::sync::Mutex<String>>` so `a2a_check_task` can return
  live progress.
- `OpenFangKernel::inject_async_callback`: tags the remote result as
  untrusted external content via `format_async_callback_message` before
  re-entering the agent loop, then dispatches the agent's response to the
  channel adapter using the captured context.
- Cancellation scope is local-only — the tool aborts the in-process SSE
  reader; no `tasks/cancel` is sent over the wire. Tool description and
  the success message say so explicitly.

Tests (run repeatedly under --test-threads=4-8 to confirm no flakes):

- All tests that mutate the process-wide `ASYNC_TASKS` /
  `A2A_TASK_PROGRESS` statics run under a shared `Mutex` via the
  `with_global_lock_async` helper. The helper clears the maps on entry
  and exit so tests cannot interfere with one another.
- Real cap test: insert MAX entries, assert the (MAX+1)th is rejected.
- Real sweep tests: `Duration::ZERO` expires everything,
  `Duration::MAX` keeps everything (no system-uptime dependency).
- Validation tests: missing message / missing url & name / SSRF-blocked /
  at-cap each return a clear error.
- `test_send_streaming_with_progress_then_inject_callback` directly
  drives `A2aClient::send_task_streaming_with_progress` against a real
  local TCP server returning canned SSE bytes, then exercises
  `inject_async_callback` on a stub kernel handle that captures
  `(context, agent_name, result_text)` verbatim. (Driving
  `tool_a2a_send_async` end-to-end is not possible without bypassing the
  production SSRF guard, which unconditionally rejects loopback IPs.)
- The stub deliberately does NOT reformat the message — that's
  production's job and is covered separately by the
  `format_async_callback_message` unit test in `openfang-kernel`.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…fter_grace (Windows panic)

Subtracting a 1-hour `Duration` from `Instant::now()` panics on Windows
when the CI VM uptime is shorter than the subtracted duration — the
`Instant` epoch is process/system start, not Unix epoch.

Same fix pattern as the other sweep tests: don't construct an
artificially-old `Instant` at all. Insert entries with
`created_at = Instant::now()` and drive the sweep decision entirely
through the threshold arguments:

  - Pass 1: sweep_with_age_thresholds(Duration::MAX, Duration::ZERO)
    keeps Running, expires Completed.
  - Pass 2: sweep_with_age_thresholds(Duration::ZERO, Duration::MAX)
    keeps Completed, expires Running.

This still exercises the asymmetric behaviour the test was checking
without any system-uptime dependency.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant