@agent-native/dispatch 0.14.7 → 0.14.8

package/src/server/lib/mcp-gateway.spec.ts

@@ -12,7 +12,9 @@ const mocks = vi.hoisted(() => ({
   managerStop: vi.fn(),
   managerCallTool: vi.fn(),
   managerConstructor: vi.fn(),
-  callAgent: vi.fn(),
+  a2aConstructor: vi.fn(),
+  a2aSend: vi.fn(),
+  a2aGetTask: vi.fn(),
   signA2AToken: vi.fn(),
   getOrgA2ASecret: vi.fn(),
   getOrgDomain: vi.fn(),
@@ -40,7 +42,19 @@ vi.mock("@agent-native/core/server", async (importOriginal) => {
 });
 
 vi.mock("@agent-native/core/a2a", () => ({
-  callAgent: mocks.callAgent,
+  A2AClient: class MockA2AClient {
+    constructor(...args: unknown[]) {
+      mocks.a2aConstructor(...args);
+    }
+
+    send(...args: unknown[]) {
+      return mocks.a2aSend(...args);
+    }
+
+    getTask(...args: unknown[]) {
+      return mocks.a2aGetTask(...args);
+    }
+  },
   signA2AToken: mocks.signA2AToken,
 }));
 
@@ -76,6 +90,7 @@ import { runWithRequestContext } from "@agent-native/core/server";
 import {
   createGrantedDispatchMcpEmbedSession,
   askGrantedDispatchMcpApp,
+  getGrantedDispatchMcpAppTask,
   listGrantedDispatchMcpApps,
   listGrantedDispatchMcpAppOrigins,
   openGrantedDispatchMcpApp,
@@ -91,6 +106,10 @@ const analyticsAgent = {
 };
 
 beforeEach(() => {
+  mocks.a2aConstructor.mockReset();
+  mocks.a2aSend.mockReset();
+  mocks.a2aGetTask.mockReset();
+  mocks.signA2AToken.mockReset();
   mocks.discoverAgents.mockResolvedValue([analyticsAgent]);
   mocks.getUserSetting.mockResolvedValue({ mode: "all-apps" });
   mocks.getOrgSetting.mockResolvedValue({ mode: "all-apps" });
@@ -106,19 +125,29 @@ beforeEach(() => {
       startUrl: "http://localhost:8086/_agent-native/embed/start?ticket=remote",
     },
   });
-  mocks.callAgent.mockResolvedValue("Created the requested dashboard.");
+  mocks.a2aSend.mockResolvedValue({
+    id: "task-1",
+    status: {
+      state: "completed",
+      message: {
+        role: "agent",
+        parts: [{ type: "text", text: "Created the requested dashboard." }],
+      },
+    },
+  });
   mocks.signA2AToken.mockResolvedValue("signed-token");
   mocks.getOrgA2ASecret.mockResolvedValue(null);
   mocks.getOrgDomain.mockResolvedValue(null);
 });
 
 afterEach(() => {
+  vi.useRealTimers();
   vi.unstubAllEnvs();
   vi.clearAllMocks();
 });
 
 describe("Dispatch MCP gateway app discovery", () => {
-  it("defaults to exposing Dispatch only until an owner or admin grants more apps", async () => {
+  it("defaults to exposing every discovered app", async () => {
     mocks.getUserSetting.mockResolvedValue(null);
     const apps = await runWithRequestContext(
       {
@@ -128,7 +157,7 @@ describe("Dispatch MCP gateway app discovery", () => {
       () => listGrantedDispatchMcpApps(),
     );
 
-    expect(apps.map((app) => app.id)).toEqual(["dispatch"]);
+    expect(apps.map((app) => app.id)).toEqual(["dispatch", "analytics"]);
   });
 
   it("includes Dispatch itself so agents can target extension routes", async () => {
@@ -275,20 +304,258 @@ describe("askGrantedDispatchMcpApp", () => {
         ),
     );
 
-    expect(mocks.callAgent).toHaveBeenCalledWith(
+    expect(mocks.a2aConstructor).toHaveBeenCalledWith(
       "http://localhost:8086",
-      "Build a weekly active users dashboard.",
+      "signed-token",
+      { requestTimeoutMs: 10_000 },
+    );
+    expect(mocks.a2aSend).toHaveBeenCalledWith(
+      {
+        role: "user",
+        parts: [
+          { type: "text", text: "Build a weekly active users dashboard." },
+        ],
+      },
+      {
+        async: true,
+        metadata: {
+          userEmail: "owner@example.test",
+          orgDomain: "builder.io",
+          requestOrigin: "http://localhost:8092",
+        },
+      },
+    );
+    expect(result).toMatchObject({
+      app: "analytics",
+      routedVia: "a2a",
+      response: "Created the requested dashboard.",
+      taskId: "task-1",
+      status: "completed",
+    });
+  });
+
+  it("returns a durable polling handle when the downstream task is still working", async () => {
+    mocks.a2aSend.mockResolvedValueOnce({
+      id: "task-working",
+      status: { state: "working" },
+    });
+
+    const result = await runWithRequestContext(
       {
         userEmail: "owner@example.test",
-        orgDomain: "builder.io",
-        orgSecret: "org-specific-secret",
-        timeoutMs: 5 * 60_000,
+        requestOrigin: "http://localhost:8092",
       },
+      () =>
+        askGrantedDispatchMcpApp("analytics", "Build the report.", {
+          async: true,
+        }),
     );
+
     expect(result).toEqual({
       app: "analytics",
       routedVia: "a2a",
-      response: "Created the requested dashboard.",
+      taskId: "task-working",
+      status: "working",
+      pollAfterMs: 1_500,
+      poll: {
+        tool: "ask_app_status",
+        arguments: { app: "analytics", taskId: "task-working" },
+      },
+      message:
+        'ask_app is still working. Call ask_app_status with taskId "task-working" to retrieve the final response.',
+    });
+  });
+
+  it("counts submission and every poll against one inline deadline", async () => {
+    vi.useFakeTimers();
+    vi.setSystemTime(0);
+    mocks.a2aSend.mockImplementationOnce(
+      () =>
+        new Promise((resolve) => {
+          setTimeout(
+            () =>
+              resolve({
+                id: "task-deadline",
+                status: { state: "working" },
+              }),
+            3_000,
+          );
+        }),
+    );
+    mocks.a2aGetTask.mockImplementationOnce(() => new Promise(() => undefined));
+
+    const resultPromise = runWithRequestContext(
+      {
+        userEmail: "owner@example.test",
+        requestOrigin: "http://localhost:8092",
+      },
+      () =>
+        askGrantedDispatchMcpApp("analytics", "Build the report.", {
+          maxWaitMs: 5_000,
+        }),
+    );
+
+    await vi.advanceTimersByTimeAsync(5_000);
+
+    await expect(resultPromise).resolves.toMatchObject({
+      taskId: "task-deadline",
+      status: "working",
+      pollAfterMs: 1_500,
+    });
+    expect(Date.now()).toBe(5_000);
+    expect(mocks.a2aConstructor).toHaveBeenCalledWith(
+      "http://localhost:8086",
+      undefined,
+      { requestTimeoutMs: 5_000 },
+    );
+    expect(mocks.a2aGetTask).toHaveBeenCalledTimes(1);
+  });
+
+  it.each([401, 403, 404])(
+    "surfaces permanent %s status errors without waiting until the deadline",
+    async (statusCode) => {
+      vi.useFakeTimers();
+      vi.setSystemTime(0);
+      mocks.a2aSend.mockResolvedValueOnce({
+        id: "task-missing",
+        status: { state: "working" },
+      });
+      mocks.a2aGetTask.mockRejectedValueOnce(
+        new Error(`A2A request failed (${statusCode}): Task not found`),
+      );
+
+      const resultPromise = runWithRequestContext(
+        {
+          userEmail: "owner@example.test",
+          requestOrigin: "http://localhost:8092",
+        },
+        () =>
+          askGrantedDispatchMcpApp("analytics", "Build the report.", {
+            maxWaitMs: 20_000,
+          }),
+      );
+      const rejection = resultPromise.catch((err) => err);
+
+      await vi.advanceTimersByTimeAsync(0);
+      await vi.advanceTimersByTimeAsync(1_500);
+
+      await expect(rejection).resolves.toEqual(
+        expect.objectContaining({
+          message: expect.stringMatching(
+            new RegExp(`${statusCode}.*not found`, "i"),
+          ),
+        }),
+      );
+      expect(Date.now()).toBe(1_500);
+      expect(mocks.a2aGetTask).toHaveBeenCalledTimes(1);
+    },
+  );
+
+  it("retries transient status failures within the same deadline", async () => {
+    vi.useFakeTimers();
+    vi.setSystemTime(0);
+    mocks.a2aSend.mockResolvedValueOnce({
+      id: "task-transient",
+      status: { state: "working" },
+    });
+    mocks.a2aGetTask
+      .mockRejectedValueOnce(new Error("A2A request failed (503): retry"))
+      .mockResolvedValueOnce({
+        id: "task-transient",
+        status: {
+          state: "completed",
+          message: {
+            role: "agent",
+            parts: [{ type: "text", text: "The report is ready." }],
+          },
+        },
+      });
+
+    const resultPromise = runWithRequestContext(
+      {
+        userEmail: "owner@example.test",
+        requestOrigin: "http://localhost:8092",
+      },
+      () =>
+        askGrantedDispatchMcpApp("analytics", "Build the report.", {
+          maxWaitMs: 20_000,
+        }),
+    );
+
+    await vi.advanceTimersByTimeAsync(0);
+    await vi.advanceTimersByTimeAsync(3_000);
+
+    await expect(resultPromise).resolves.toMatchObject({
+      taskId: "task-transient",
+      status: "completed",
+      response: "The report is ready.",
+    });
+    expect(mocks.a2aGetTask).toHaveBeenCalledTimes(2);
+  });
+
+  it("returns input-required as a terminal handoff instead of a poll loop", async () => {
+    mocks.a2aSend.mockResolvedValueOnce({
+      id: "task-input",
+      status: {
+        state: "input-required",
+        message: {
+          role: "agent",
+          parts: [{ type: "text", text: "Which date range should I use?" }],
+        },
+      },
+    });
+
+    const result = await runWithRequestContext(
+      {
+        userEmail: "owner@example.test",
+        requestOrigin: "http://localhost:8092",
+      },
+      () => askGrantedDispatchMcpApp("analytics", "Build the report."),
+    );
+
+    expect(result).toEqual({
+      app: "analytics",
+      routedVia: "a2a",
+      taskId: "task-input",
+      status: "input-required",
+      response: "Which date range should I use?",
+      inputRequired: "Which date range should I use?",
+      message: "Which date range should I use?",
+    });
+    expect(mocks.a2aGetTask).not.toHaveBeenCalled();
+  });
+
+  it("polls a granted app task through the same authenticated A2A route", async () => {
+    mocks.a2aGetTask.mockResolvedValueOnce({
+      id: "task-working",
+      status: {
+        state: "completed",
+        message: {
+          role: "agent",
+          parts: [{ type: "text", text: "The report is ready." }],
+        },
+      },
+    });
+    mocks.getUserSetting.mockResolvedValue({
+      mode: "selected-apps",
+      selectedAppIds: ["analytics"],
+    });
+
+    const result = await runWithRequestContext(
+      {
+        userEmail: "owner@example.test",
+        requestOrigin: "http://localhost:8092",
+      },
+      () => getGrantedDispatchMcpAppTask("analytics", "task-working"),
+    );
+
+    expect(mocks.a2aGetTask).toHaveBeenCalledWith("task-working");
+    expect(result).toEqual({
+      app: "analytics",
+      routedVia: "a2a",
+      taskId: "task-working",
+      status: "completed",
+      response: "The report is ready.",
     });
   });
 
@@ -307,7 +574,25 @@ describe("askGrantedDispatchMcpApp", () => {
         () => askGrantedDispatchMcpApp("analytics", "Show signups."),
       ),
     ).rejects.toThrow(/not granted/);
-    expect(mocks.callAgent).not.toHaveBeenCalled();
+    expect(mocks.a2aSend).not.toHaveBeenCalled();
+  });
+
+  it("rejects polling a task for an app outside the grant", async () => {
+    mocks.getUserSetting.mockResolvedValue({
+      mode: "selected-apps",
+      selectedAppIds: ["dispatch"],
+    });
+
+    await expect(
+      runWithRequestContext(
+        {
+          userEmail: "owner@example.test",
+          requestOrigin: "http://localhost:8092",
+        },
+        () => getGrantedDispatchMcpAppTask("analytics", "task-working"),
+      ),
+    ).rejects.toThrow(/not granted/);
+    expect(mocks.a2aGetTask).not.toHaveBeenCalled();
   });
 });
 
@@ -547,7 +832,7 @@ describe("createGrantedDispatchMcpEmbedSession", () => {
     expect(mocks.managerConstructor).toHaveBeenCalledWith({
       servers: {
         target: expect.objectContaining({
-          url: "http://localhost:8092/analytics/_agent-native/mcp",
+          url: "http://localhost:8092/analytics/mcp",
         }),
       },
     });
@@ -596,7 +881,7 @@ describe("createGrantedDispatchMcpEmbedSession", () => {
     expect(mocks.managerConstructor).toHaveBeenCalledWith({
       servers: {
         target: expect.objectContaining({
-          url: "https://mail.agent-native.com/_agent-native/mcp",
+          url: "https://mail.agent-native.com/mcp",
         }),
       },
     });

package/src/server/lib/mcp-gateway.ts

@@ -1,4 +1,4 @@
-import { callAgent, signA2AToken } from "@agent-native/core/a2a";
+import { A2AClient, signA2AToken, type Task } from "@agent-native/core/a2a";
 import {
   buildMcpToolName,
   McpClientManager,
@@ -30,6 +30,23 @@ const DISPATCH_DESCRIPTION =
 const DISPATCH_COLOR = "#14B8A6";
 const TARGET_EMBED_SESSION_ATTEMPTS = 3;
 const TARGET_EMBED_SESSION_RETRY_BASE_MS = 250;
+const DISPATCH_ASK_APP_DEFAULT_INLINE_WAIT_MS = 20_000;
+const DISPATCH_ASK_APP_MAX_INLINE_WAIT_MS = 25_000;
+const DISPATCH_ASK_APP_POLL_INTERVAL_MS = 1_500;
+const DISPATCH_A2A_REQUEST_TIMEOUT_MS = 10_000;
+const DISPATCH_ASK_APP_TERMINAL_STATES = new Set([
+  "completed",
+  "failed",
+  "canceled",
+  "input-required",
+]);
+
+class DispatchAskAppInlineDeadlineError extends Error {
+  constructor() {
+    super("ask_app inline wait deadline reached");
+    this.name = "DispatchAskAppInlineDeadlineError";
+  }
+}
 
 export interface DispatchMcpAccessibleApp {
   id: string;
@@ -44,6 +61,204 @@ function normalizeAppId(value: string): string {
   return value.trim().toLowerCase();
 }
 
+function boundedDispatchAskAppWaitMs(raw: unknown): number {
+  if (raw == null || raw === "") {
+    return DISPATCH_ASK_APP_DEFAULT_INLINE_WAIT_MS;
+  }
+  const parsed = Number(raw);
+  if (!Number.isFinite(parsed)) return DISPATCH_ASK_APP_DEFAULT_INLINE_WAIT_MS;
+  return Math.max(
+    0,
+    Math.min(DISPATCH_ASK_APP_MAX_INLINE_WAIT_MS, Math.trunc(parsed)),
+  );
+}
+
+function isTerminalDispatchTask(task: Task): boolean {
+  return DISPATCH_ASK_APP_TERMINAL_STATES.has(String(task.status.state));
+}
+
+function dispatchTaskText(task: Task): string {
+  return (
+    task.status.message?.parts
+      ?.filter(
+        (part): part is { type: "text"; text: string } => part.type === "text",
+      )
+      .map((part) => part.text)
+      .join("\n")
+      .trim() ?? ""
+  );
+}
+
+function dispatchAskAppTaskResult(
+  app: string,
+  task: Task,
+): {
+  app: string;
+  routedVia: "a2a";
+  taskId: string;
+  status: string;
+  response?: string;
+  error?: string;
+  inputRequired?: string;
+  pollAfterMs?: number;
+  poll?: { tool: "ask_app_status"; arguments: { app: string; taskId: string } };
+  message?: string;
+} {
+  const status = String(task.status.state);
+  const response = dispatchTaskText(task);
+  const base = {
+    app,
+    routedVia: "a2a" as const,
+    taskId: task.id,
+    status,
+  };
+
+  if (status === "completed") {
+    return { ...base, response: response || "(no response)" };
+  }
+  if (status === "failed" || status === "canceled") {
+    return {
+      ...base,
+      ...(response ? { response } : {}),
+      error: response || `ask_app task ${status}.`,
+    };
+  }
+  if (status === "input-required") {
+    const inputRequired =
+      response || "The agent needs additional input before it can continue.";
+    return {
+      ...base,
+      response: inputRequired,
+      inputRequired,
+      message: inputRequired,
+    };
+  }
+  return {
+    ...base,
+    pollAfterMs: DISPATCH_ASK_APP_POLL_INTERVAL_MS,
+    poll: {
+      tool: "ask_app_status",
+      arguments: { app, taskId: task.id },
+    },
+    message:
+      `ask_app is still ${status}. Call ask_app_status with ` +
+      `taskId "${task.id}" to retrieve the final response.`,
+  };
+}
+
+async function createDispatchA2AClient(input: {
+  targetUrl: string;
+  userEmail: string;
+  orgDomain?: string;
+  orgSecret?: string;
+  deadline?: number;
+}): Promise<{
+  client: A2AClient;
+  metadata: Record<string, unknown>;
+}> {
+  const apiKeys: string[] = [];
+  const addSignedToken = async (preferGlobalSecret: boolean) => {
+    try {
+      const token = await signA2AToken(
+        input.userEmail,
+        input.orgDomain,
+        input.orgSecret,
+        { preferGlobalSecret },
+      );
+      if (token && !apiKeys.includes(token)) apiKeys.push(token);
+    } catch {
+      // A2A can still be configured for local/dev unauthenticated calls. If
+      // signing is unavailable, let the target return its own auth error.
+    }
+  };
+
+  if (process.env.A2A_SECRET?.trim()) await addSignedToken(true);
+  if (input.orgSecret) await addSignedToken(false);
+
+  const metadata: Record<string, unknown> = {
+    userEmail: input.userEmail,
+    ...(input.orgDomain ? { orgDomain: input.orgDomain } : {}),
+    ...(getRequestContext()?.requestOrigin
+      ? { requestOrigin: getRequestContext()?.requestOrigin }
+      : {}),
+  };
+  const remainingMs =
+    input.deadline == null ? null : input.deadline - Date.now();
+  return {
+    client: new A2AClient(input.targetUrl, apiKeys[0], {
+      requestTimeoutMs:
+        remainingMs == null
+          ? DISPATCH_A2A_REQUEST_TIMEOUT_MS
+          : Math.max(1, Math.min(DISPATCH_A2A_REQUEST_TIMEOUT_MS, remainingMs)),
+      ...(apiKeys.length > 1 ? { fallbackApiKeys: apiKeys.slice(1) } : {}),
+    }),
+    metadata,
+  };
+}
+
+function isTransientDispatchAskAppStatusError(err: unknown): boolean {
+  const message = err instanceof Error ? err.message : String(err ?? "");
+  return (
+    /\bfetch failed\b|failed to fetch|networkerror|socket hang up|econnreset|etimedout|timeout|aborted/i.test(
+      message,
+    ) || /A2A request failed \((?:429|500|502|503|504)\)/i.test(message)
+  );
+}
+
+async function runBeforeDispatchAskAppDeadline<T>(
+  operation: () => Promise<T>,
+  deadline: number,
+): Promise<T> {
+  const remainingMs = deadline - Date.now();
+  if (remainingMs <= 0) throw new DispatchAskAppInlineDeadlineError();
+
+  let timeout: ReturnType<typeof setTimeout> | undefined;
+  try {
+    return await Promise.race([
+      operation(),
+      new Promise<T>((_resolve, reject) => {
+        timeout = setTimeout(
+          () => reject(new DispatchAskAppInlineDeadlineError()),
+          remainingMs,
+        );
+      }),
+    ]);
+  } finally {
+    if (timeout) clearTimeout(timeout);
+  }
+}
+
+async function waitForDispatchA2ATask(
+  client: A2AClient,
+  task: Task,
+  deadline: number | undefined,
+): Promise<Task> {
+  if (deadline == null || isTerminalDispatchTask(task)) return task;
+  let current = task;
+  while (!isTerminalDispatchTask(current)) {
+    const remaining = deadline - Date.now();
+    if (remaining <= 0) return current;
+    await new Promise((resolve) =>
+      setTimeout(
+        resolve,
+        Math.min(DISPATCH_ASK_APP_POLL_INTERVAL_MS, remaining),
+      ),
+    );
+    if (Date.now() >= deadline) return current;
+    try {
+      current = await runBeforeDispatchAskAppDeadline(
+        () => client.getTask(task.id),
+        deadline,
+      );
+    } catch (err) {
+      if (err instanceof DispatchAskAppInlineDeadlineError) return current;
+      if (!isTransientDispatchAskAppStatusError(err)) throw err;
+      if (Date.now() >= deadline) return current;
+    }
+  }
+  return current;
+}
+
 function normalizeBaseUrl(raw: string | undefined | null): string | null {
   const value = raw?.trim();
   if (!value) return null;
@@ -281,7 +496,8 @@ export async function resolveGrantedDispatchMcpApp(
 export async function askGrantedDispatchMcpApp(
   app: string,
   message: string,
-): Promise<{ app: string; routedVia: "a2a"; response: string }> {
+  options?: { async?: boolean; maxWaitMs?: number },
+): Promise<ReturnType<typeof dispatchAskAppTaskResult>> {
   const trimmedMessage = message.trim();
   if (!trimmedMessage) throw new Error("message is required");
   const target = await resolveGrantedDispatchMcpApp(app);
@@ -296,13 +512,55 @@ export async function askGrantedDispatchMcpApp(
       ])
     : [null, null];
 
-  const response = await callAgent(target.url, trimmedMessage, {
+  const inlineWaitMs =
+    options?.async === true
+      ? 0
+      : boundedDispatchAskAppWaitMs(options?.maxWaitMs);
+  const deadline = inlineWaitMs > 0 ? Date.now() + inlineWaitMs : undefined;
+
+  const { client, metadata } = await createDispatchA2AClient({
+    targetUrl: target.url,
+    userEmail,
+    orgDomain: orgDomain ?? undefined,
+    orgSecret: orgSecret ?? undefined,
+    deadline,
+  });
+  const task = await client.send(
+    {
+      role: "user",
+      parts: [{ type: "text", text: trimmedMessage }],
+    },
+    { async: true, metadata },
+  );
+  const finalOrRunning = await waitForDispatchA2ATask(client, task, deadline);
+  return dispatchAskAppTaskResult(target.id, finalOrRunning);
+}
+
+export async function getGrantedDispatchMcpAppTask(
+  app: string,
+  taskId: string,
+): Promise<ReturnType<typeof dispatchAskAppTaskResult>> {
+  const trimmedTaskId = taskId.trim();
+  if (!trimmedTaskId) throw new Error("taskId is required");
+  const target = await resolveGrantedDispatchMcpApp(app);
+  const userEmail = getRequestUserEmail();
+  if (!userEmail) throw new Error("no authenticated user");
+
+  const orgId = getRequestOrgId();
+  const [orgDomain, orgSecret] = orgId
+    ? await Promise.all([
+        getOrgDomain(orgId).catch(() => null),
+        getOrgA2ASecret(orgId).catch(() => null),
+      ])
+    : [null, null];
+  const { client } = await createDispatchA2AClient({
+    targetUrl: target.url,
     userEmail,
     orgDomain: orgDomain ?? undefined,
     orgSecret: orgSecret ?? undefined,
-    timeoutMs: 5 * 60_000,
   });
-  return { app: target.id, routedVia: "a2a", response };
+  const task = await client.getTask(trimmedTaskId);
+  return dispatchAskAppTaskResult(target.id, task);
 }
 
 export async function openGrantedDispatchMcpApp(input: {
@@ -434,7 +692,7 @@ async function callTargetCreateEmbedSession(input: {
       servers: {
         [serverId]: {
           type: "http",
-          url: `${appBaseUrl(input.app)}/_agent-native/mcp`,
+          url: `${appBaseUrl(input.app)}/mcp`,
           headers: {
             Authorization: `Bearer ${input.token}`,
           },