A2A late-reply fold-back to primary session

[rockfordlhotka]

Background

When a subagent invokes A2A and its reply arrives after the subagent terminates, the result is silently dropped (see A2ATaskResultHandler:78-82 and siblings). PendingA2ATask already carries PrimarySessionId — the routing data exists, just unused on the late path.

The existing a2aAwaiter already prevents most late arrivals — subagents block on outstanding A2A before publishing their result (SubagentRunner.cs:283-300). The fold-back triggers when the awaiter times out, the subagent is cancelled, or the primary cancels the subagent's parent task. This is the safety net, not the hot path.

(Tool-profile scoping for this new handler — and the rest of the in-process roles — is split out to #425 and will land after this issue. Until then, the new handler uses the unfiltered tool set, matching today's UserMessageHandler surface.)

Design

In the 4 A2A handlers (A2ATaskResultHandler, A2ATaskStatusHandler, A2ATaskErrorHandler, InputRequiredHandler), after tracker.TryRemove(correlationId) returns a PendingA2ATask:

1. If PendingA2ATask.SubagentSessionId is set AND SubagentManager.IsActive(subagentSessionId) → route to subagent (current path; existing a2aAwaiter handles the still-running case).
2. Else if PendingA2ATask.PrimarySessionId is set → fold back to primary (new path).
3. Else → log Warning and drop (current behavior).

Fold-back path:

• Stash the reply payload to working memory under notifications/a2a/{subagentTaskId}/{kind} (where kind is one of result | status | error | input-required) in the primary's namespace.
• Append a one-line entry to notifications/index so list_working_memory surfaces pending notifications.
• Publish a new LateA2ANotificationMessage:

record LateA2ANotificationMessage(
    string PrimarySessionId,
    string SubagentTaskId,
    string SubagentName,
    string PeerAgent,
    NotificationKind Kind,
    string WorkingMemoryKey);

• LateA2ANotificationHandler (new, registered in RockBot.Agent/Program.cs) receives the message, loads the primary session, and runs AgentLoopRunner.RunAsync (unfiltered tool set; see note above re: Per-invocation tool profiles for in-process roles #425) with a prompt of the form: A late {kind} arrived from A2A peer {peer} for your completed subagent {name}. The payload is in working memory at {key}. Read it, decide whether to act on it, and inform the user if relevant. The prompt requires the model to surface the context to the user.

Why a message + handler rather than a synchronous call: the bus handles queueing, multiple late arrivals serialize into separate primary turns without colliding with an in-progress turn, and the original A2A handler returns promptly.

Primary-agent directive update: add a line telling the primary to check notifications/ on each turn as defense-in-depth.

Implementation order

1. src/RockBot.Subagent/SubagentManager.cs — add IsActive(subagentSessionId) accessor backed by existing SubagentEntry registry
2. src/RockBot.A2A/LateA2ANotificationMessage.cs (new) — record + NotificationKind enum
3. src/RockBot.A2A/LateA2ANotificationHandler.cs (new) + registration in src/RockBot.Agent/Program.cs
4. Fold-back routing in A2ATaskResultHandler.cs, A2ATaskStatusHandler.cs, A2ATaskErrorHandler.cs, InputRequiredHandler.cs: WM stash → publish LateA2ANotificationMessage → return
5. Primary-agent directive update (src/RockBot.Agent/agent/directives.md) — line telling the primary to check notifications/ on each turn
6. Tests:
   • Fold-back routing tests (subagent active vs terminated branches; missing primary session)
   • LateA2ANotificationHandler happy path + missing-WM-key path

Tests

• MSTest + Rocks (per repo convention)
• No ROCKBOT_RABBITMQ_HOST-gated integration test changes needed — this is in-process routing, not transport

Risks

• Existing a2aAwaiter already prevents most late arrivals. This fix is the safety net.
• Notification storms. Several late replies arriving back-to-back create N primary turns. Acceptable for v1; bus serializes them. If noisy in practice, add debounce later (handler reads notifications/index and consolidates pending entries into one turn).
• User-side surprise. The primary speaks without a user prompt. The LateA2ANotificationHandler prompt requires the model to state the source so it is not opaque.
• Cancellation correctness. If PrimarySessionId resolves to nothing (CLI one-shot, disconnected user), the handler drops cleanly with a Warning. Matches existing "session not found" patterns.

Out of scope

• Per-invocation tool profiles — split to Per-invocation tool profiles for in-process roles #425
• Notification debounce/coalescing (add only if storms materialize)
• Surfacing notifications to the user without an intervening primary-agent turn
• Changing AgentLoopRunner.RunAsync's signature

Related

• Per-invocation tool profiles for in-process roles #425 Per-invocation tool profiles — will refine this handler's tool surface once it lands
• Per-server typed bridge tools for MCP (replace generic mcp_invoke_tool) #420 Per-server typed bridge tools for MCP

[rockfordlhotka]

Plan — Issue #424: A2A late-reply fold-back to primary session

Branch: issue-424/a2a-late-reply-fold-back

Context

When a subagent dispatches an A2A call (via a wisp inside its loop), the resulting AgentTaskResult / AgentTaskStatusUpdate / AgentTaskError / InputRequired message carries PendingA2ATask.PrimarySessionId = "subagent/{taskId}" — the subagent's working-memory namespace, not a user-facing session. The existing four A2A handlers stash the payload to WM under that subagent namespace and then short-circuit via IsUserSession(...) → return because subagents are not user-facing. As long as the subagent is still running, A2ATaskAwaiter.WaitForSessionAsync blocks the subagent before it publishes its result so the reply makes it into the subagent's whiteboard. But when the awaiter is interrupted (subagent timeout, explicit cancel, parent-task cancel), the late reply is silently dropped.

This change introduces a fold-back path: when the originating subagent has terminated, the reply is stashed to working memory under the user-facing primary session's namespace, and a LateA2ANotificationMessage is published so a fresh primary-agent turn surfaces it to the user.

Tool-profile scoping for the new handler is deferred to #425. The new handler uses the unfiltered tool surface for now, matching UserMessageHandler today.

Critical Files

Modified:

• src/RockBot.A2A/PendingA2ATask.cs — add fields
• src/RockBot.A2A/InvokeAgentExecutor.cs — populate new fields at dispatch time
• src/RockBot.A2A/A2ATaskResultHandler.cs — fold-back routing on terminal path
• src/RockBot.A2A/A2ATaskStatusHandler.cs — fold-back routing
• src/RockBot.A2A/A2ATaskErrorHandler.cs — fold-back routing
• src/RockBot.A2A/InputRequiredHandler.cs — used by A2ATaskResultHandler's InputRequired branch; new fold-back behavior lives in the result handler (see notes)
• src/RockBot.A2A/A2ACallerServiceCollectionExtensions.cs — register new handler + topic subscription
• src/RockBot.Subagent/SubagentManager.cs + ISubagentManager.cs — add IsActive(subagentSessionId)
• src/RockBot.Agent/agent/directives.md — primary-agent guidance to scan notifications/ per turn

New:

• src/RockBot.A2A/LateA2ANotificationMessage.cs — message + NotificationKind enum
• src/RockBot.A2A/LateA2ANotificationHandler.cs — wakes the primary loop
• src/RockBot.A2A/A2ATopics.cs or extend A2AOptions — topic constant for late-notifications
• tests/RockBot.A2A.Tests/LateA2ANotificationHandlerTests.cs
• tests/RockBot.A2A.Tests/LateReplyFoldBackTests.cs (handler-level fold-back routing)
• tests/RockBot.Subagent.Tests/SubagentManagerIsActiveTests.cs

Approach

1. PendingA2ATask — capture both subagent + user-primary IDs

Add two optional fields without renaming the existing one (rename would ripple through 6+ files; additive is safer):

internal sealed record PendingA2ATask
{
    public required string TaskId { get; init; }
    public required string TargetAgent { get; init; }
    public required string Skill { get; init; }
    public required string PrimarySessionId { get; init; }   // unchanged — the calling session (subagent ns or user session)
    public required DateTimeOffset StartedAt { get; init; }
    public required CancellationTokenSource Cts { get; init; }

    // NEW — set when the call originated from a subagent (e.g. "subagent-abc123def4").
    public string? SubagentSessionId { get; init; }
    // NEW — the user-facing session for fold-back. Equals PrimarySessionId when the
    // call originated from a user session; populated from SubagentEntry when the
    // call originated from a subagent. Null when no user-facing session is known.
    public string? UserPrimarySessionId { get; init; }

    public string? ContextId { get; set; }
    public int InputRequiredRound { get; set; }
    public string? LastInputRequiredQuestion { get; set; }
    public string? LastInputRequiredAnswer { get; set; }
}

Field name choice diverges from the issue text (which writes PendingA2ATask.PrimarySessionId as the user primary). The semantics match — only the names are different — and the divergence is recorded here.

2. InvokeAgentExecutor — populate the new fields

At InvokeAgentExecutor.cs:131-178, after computing var primarySessionId = request.SessionId ?? "unknown";:

• If primarySessionId.StartsWith("subagent/", StringComparison.OrdinalIgnoreCase):
  • Extract taskId = primarySessionId["subagent/".Length..]
  • subagentSessionId = $"subagent-{taskId}"
  • Look up via ISubagentManager (inject the interface) → entry = manager.ListActive().FirstOrDefault(e => e.TaskId == taskId)
  • If found, userPrimary = entry.PrimarySessionId; else userPrimary = null (subagent already gone at dispatch time — unusual but possible)
• Else (assume user-session caller): subagentSessionId = null, userPrimary = primarySessionId

Wisp sessions inside the primary (not inside a subagent) also pass through this code with request.SessionId set by WispRegistryToolFunction to parentSessionId ?? wispId. Those use the user session as parent already, so the else branch is correct for them.

3. ISubagentManager.IsActive(subagentSessionId)

Add to ISubagentManager:

bool IsActive(string subagentSessionId);

Implementation in SubagentManager.cs: scan _active.Values for an entry where entry.SubagentSessionId == subagentSessionId AND !entry.Task.IsCompleted. Skip the cleanup pass — read-only check.

4. LateA2ANotificationMessage + topic

namespace RockBot.A2A;

public enum NotificationKind { Result, Status, Error, InputRequired }

public sealed record LateA2ANotificationMessage
{
    public required string PrimarySessionId { get; init; }   // user-facing
    public required string SubagentTaskId { get; init; }
    public required string SubagentName { get; init; }       // e.g. "subagent-abc123"
    public required string PeerAgent { get; init; }          // TargetAgent of the original call
    public required NotificationKind Kind { get; init; }
    public required string WorkingMemoryKey { get; init; }   // primary-relative key
}

Add A2AOptions.LateNotificationTopic = "agent.notification.late-a2a" (per-agent suffix added at subscribe time, matching CallerResultTopic convention).

5. Fold-back wiring in the 4 handlers

The common branch lives at the "skip" point in each handler (where IsUserSession(...) is false today). Extract a helper to keep the four handlers identical:

// In a new file src/RockBot.A2A/LateA2AFoldBack.cs (or inline static on the message):
internal static class LateA2AFoldBack
{
    public static async Task<bool> TryFoldBackAsync(
        PendingA2ATask pending,
        NotificationKind kind,
        string payloadText,
        IWorkingMemory workingMemory,
        IMessagePublisher publisher,
        ISubagentManager subagentManager,
        AgentIdentity agent,
        A2AOptions options,
        ILogger logger,
        CancellationToken ct)
    {
        // (1) Subagent still active → no fold-back
        if (pending.SubagentSessionId is { } sub && subagentManager.IsActive(sub))
            return false;

        // (2) No user-facing primary recorded → unchanged drop (warning)
        if (pending.UserPrimarySessionId is not { Length: > 0 } userPrimary)
            return false;

        // (3) Fold back
        var kindSlug = kind switch
        {
            NotificationKind.Result => "result",
            NotificationKind.Status => "status",
            NotificationKind.Error => "error",
            NotificationKind.InputRequired => "input-required",
            _ => "unknown"
        };
        var subagentTaskId = pending.SubagentSessionId?.Replace("subagent-", "") ?? "unknown";
        var key = $"{userPrimary}/notifications/a2a/{subagentTaskId}/{kindSlug}";
        await workingMemory.SetAsync(key, payloadText, TimeSpan.FromMinutes(60),
            category: "a2a-late-notification",
            tags: [pending.TargetAgent, subagentTaskId, kindSlug]);

        // Index entry (append; one line)
        var indexKey = $"{userPrimary}/notifications/index";
        var existing = await workingMemory.GetAsync(indexKey) ?? string.Empty;
        var line = $"{DateTimeOffset.UtcNow:O}\ta2a\t{kindSlug}\t{pending.TargetAgent}\t{subagentTaskId}\t{key}\n";
        await workingMemory.SetAsync(indexKey, existing + line, TimeSpan.FromMinutes(60),
            category: "notifications-index");

        var msg = new LateA2ANotificationMessage
        {
            PrimarySessionId = userPrimary,
            SubagentTaskId = subagentTaskId,
            SubagentName = pending.SubagentSessionId ?? "unknown",
            PeerAgent = pending.TargetAgent,
            Kind = kind,
            WorkingMemoryKey = key
        };
        var envelope = msg.ToEnvelope<LateA2ANotificationMessage>(source: agent.Name);
        await publisher.PublishAsync($"{options.LateNotificationTopic}.{agent.Name}", envelope, ct);
        return true;
    }
}

Each handler's existing IsUserSession == false short-circuit becomes:

if (!IsUserSession(pending.PrimarySessionId))
{
    var folded = await LateA2AFoldBack.TryFoldBackAsync(pending, kind, payloadText, ...);
    if (folded)
        logger.LogInformation("Folded late A2A {Kind} to primary session {Session}", kind, pending.UserPrimarySessionId);
    else
        logger.LogInformation("Late A2A {Kind} dropped (no fold-back available) for session {Session}", kind, pending.PrimarySessionId);
    return;
}

• A2ATaskResultHandler (terminal path after WM stash, around line 280): pass resultText plus kind = Result. The WM stash under the subagent namespace already happened above (existing behavior) — fold-back writes a separate user-namespaced entry tailored for the notification.
• A2ATaskStatusHandler (around line 75 onwards — the non-Working path): pass kind = Status and the status text. Working state still relays to the user as today.
• A2ATaskErrorHandler (around line 75): pass kind = Error and the error message.
• InputRequiredHandler is an internal helper used by A2ATaskResultHandler; the late path enters via the result handler's HandleInputRequiredAsync branch, not directly. Add fold-back inside that branch too: if tracker.TryGet(...) succeeds but the subagent is gone, fold back with kind = InputRequired and the question text, then tracker.TryRemove(...) and dispose the CTS instead of sending an autonomous follow-up to the peer.

6. LateA2ANotificationHandler

Mirrors the structure of A2ATaskErrorHandler but builds the synthesis prompt from the WM-stored payload pointer:

• DI: same dependencies as A2ATaskResultHandler minus the InputRequired/tracker
• Subscribes to agent.notification.late-a2a.{agentName} (registered in AddA2ACaller)
• Synthesis turn: "A late {kind} arrived from A2A peer {peer} for your completed subagent {name}. The payload is in working memory at {key}. Read it, decide whether to act on it, and inform the user if relevant."
• Builds chat messages via AgentContextBuilder.BuildAsync(rawSessionId, syntheticUserTurn, ...)
• AgentLoopRunner.RunAsync with no filter — full tool surface (per Per-invocation tool profiles for in-process roles #425 deferral note)
• Publishes an AgentReply to UserProxyTopics.UserResponse.{agent.Name} so the user sees the result

7. DI + subscription registration

In A2ACallerServiceCollectionExtensions.AddA2ACaller:

builder.HandleMessage<LateA2ANotificationMessage, LateA2ANotificationHandler>();
builder.SubscribeTo($"{options.LateNotificationTopic}.{agentName}");

(Registered alongside the existing handlers, not in Program.cs — matches the existing pattern; the issue's "registered in Program.cs" is a small location call-out and the substance is the same.)

8. directives.md

Append one bullet to the operational rhythm section (around the "Task Execution and Planning" area):

Check pending notifications. At the start of each turn, scan list_working_memory for entries under notifications/ — those are late A2A replies or system signals that arrived between sessions. Surface anything that matters to the user.

9. Tests

• SubagentManagerIsActiveTests — returns true for a running entry; false after entry's task completes; false for an unknown session.
• LateReplyFoldBackTests (Result / Status / Error / InputRequired handler variants):
  • Subagent active → no fold-back, no message published, existing skip behavior.
  • Subagent inactive, UserPrimarySessionId set → WM has notification entry + index line; one envelope published to the late-notification topic; tracker.TryRemove cleared the entry.
  • Subagent inactive, no UserPrimarySessionId → behavior unchanged (drop with Warning, no publish).
  • Verify the existing per-namespace WM stash still happens (Result handler) — the fold-back is additional, not a replacement of the subagent-namespace store.
• LateA2ANotificationHandlerTests — happy path: handler reads WM at the provided key, runs the loop runner (mock returns canned content), publishes an AgentReply with that content. Missing-WM-key path: handler logs Warning and exits without publishing.
• Pattern follows A2ATaskResultHandlerTests — MSTest + InMemoryWorkingMemory + TrackingPublisher, with most LLM-side deps as null! since the fold-back path returns before touching them.

10. Open questions

• SubagentSessionId shape — issue text uses subagent-{taskId}; the runtime stores both subagent-{taskId} (SubagentEntry.SubagentSessionId) and subagent/{taskId} (working-memory namespace / PrimarySessionId of the calling A2A task). The plan extracts the taskId from subagent/{taskId} and stores SubagentSessionId = subagent-{taskId} on PendingA2ATask for consistency with SubagentEntry. IsActive compares against that form.
• Per-handler vs shared helper — chose a shared static helper (LateA2AFoldBack) over duplicating logic four times. Trades a tiny amount of indirection for a single regression surface.
• Notification storms — single late replies will produce one primary turn each. The issue explicitly accepts this for v1; debounce is out of scope.
• Existing tests — A2ATaskResultHandlerTests constructs PendingA2ATask without the new fields. Required init properties stay required; new fields are init and string? so they default to null. Existing tests still compile.
• Topic name — picking agent.notification.late-a2a (sub-namespaced for future "late" categories). Open to alternatives.

Verification

1. dotnet build RockBot.slnx
2. dotnet test RockBot.slnx --filter "FullyQualifiedName~A2A|FullyQualifiedName~SubagentManagerIsActive|FullyQualifiedName~LateA2A"
3. dotnet test RockBot.slnx (full suite — no regressions)
4. Smoke test (manual, only if cluster is handy): cancel a subagent that has an outstanding A2A call mid-flight; confirm the late reply arrives and the primary agent surfaces it in the next user reply.