@assistant-ui/react-google-adk 0.0.34 → 0.0.36

package/src/adkThreadState.ts

@@ -0,0 +1,124 @@
+import type { AppendMessage } from "@assistant-ui/core";
+import type {
+  AdkAuthRequest,
+  AdkMessage,
+  AdkMessageMetadata,
+  AdkThreadSnapshot,
+  AdkToolConfirmation,
+} from "./types";
+
+export type AdkStagedEntry = {
+  message: AdkMessage & { id: string };
+  runConfig: AppendMessage["runConfig"];
+};
+
+export type AdkThreadState = {
+  messages: AdkMessage[];
+  stateDelta: Record<string, unknown>;
+  agentInfo: { name?: string | undefined; branch?: string | undefined };
+  longRunningToolIds: string[];
+  artifactDelta: Record<string, number>;
+  toolConfirmations: AdkToolConfirmation[];
+  authRequests: AdkAuthRequest[];
+  escalated: boolean;
+  messageMetadata: Map<string, AdkMessageMetadata>;
+  stagedEntries: ReadonlyMap<string, AdkStagedEntry>;
+};
+
+export type AdkThreadAction =
+  | { type: "event.published"; state: Omit<AdkThreadState, "stagedEntries"> }
+  | { type: "snapshot.applied"; snapshot: AdkThreadSnapshot }
+  | { type: "messages.replaced"; messages: AdkMessage[] }
+  | { type: "messages.set"; messages: AdkMessage[] }
+  | { type: "longRunningToolIds.set"; ids: string[] }
+  | { type: "staged.stage"; entry: AdkStagedEntry }
+  | { type: "staged.unstage"; ids: readonly string[] }
+  | {
+      type: "run.started";
+      messages: AdkMessage[];
+      longRunningToolIds: string[];
+      toolConfirmations: AdkToolConfirmation[];
+      authRequests: AdkAuthRequest[];
+    };
+
+export const createAdkThreadState = (): AdkThreadState => ({
+  messages: [],
+  stateDelta: {},
+  agentInfo: {},
+  longRunningToolIds: [],
+  artifactDelta: {},
+  toolConfirmations: [],
+  authRequests: [],
+  escalated: false,
+  messageMetadata: new Map(),
+  stagedEntries: new Map(),
+});
+
+export const reduceAdkThreadState = (
+  state: AdkThreadState,
+  action: AdkThreadAction,
+): AdkThreadState => {
+  switch (action.type) {
+    case "event.published": {
+      const next = action.state;
+      return {
+        ...state,
+        ...next,
+        stateDelta: { ...state.stateDelta, ...next.stateDelta },
+        artifactDelta: { ...state.artifactDelta, ...next.artifactDelta },
+        messageMetadata:
+          next.messageMetadata.size > 0
+            ? new Map([...state.messageMetadata, ...next.messageMetadata])
+            : state.messageMetadata,
+      };
+    }
+    case "snapshot.applied": {
+      const snapshot = action.snapshot;
+      return {
+        ...state,
+        messages: snapshot.messages,
+        stateDelta: snapshot.stateDelta ?? {},
+        agentInfo: snapshot.agentInfo ?? {},
+        longRunningToolIds: snapshot.longRunningToolIds ?? [],
+        artifactDelta: snapshot.artifactDelta ?? {},
+        toolConfirmations: snapshot.toolConfirmations ?? [],
+        authRequests: snapshot.authRequests ?? [],
+        escalated: snapshot.escalated ?? false,
+        messageMetadata: snapshot.messageMetadata ?? new Map(),
+      };
+    }
+    case "messages.replaced":
+      return {
+        ...state,
+        messages: action.messages,
+        longRunningToolIds: [],
+        toolConfirmations: [],
+        authRequests: [],
+        escalated: false,
+        messageMetadata: new Map(),
+      };
+    case "messages.set":
+      return { ...state, messages: action.messages };
+    case "longRunningToolIds.set":
+      return { ...state, longRunningToolIds: action.ids };
+    case "staged.stage": {
+      const stagedEntries = new Map(state.stagedEntries);
+      stagedEntries.set(action.entry.message.id, action.entry);
+      return { ...state, stagedEntries };
+    }
+    case "staged.unstage": {
+      if (!action.ids.some((id) => state.stagedEntries.has(id))) return state;
+      const stagedEntries = new Map(state.stagedEntries);
+      for (const id of action.ids) stagedEntries.delete(id);
+      return { ...state, stagedEntries };
+    }
+    case "run.started":
+      return {
+        ...state,
+        messages: action.messages,
+        longRunningToolIds: action.longRunningToolIds,
+        toolConfirmations: action.toolConfirmations,
+        authRequests: action.authRequests,
+      };
+  }
+};

package/src/convertToAdkMessages.test.ts

@@ -300,6 +300,25 @@ describe("getMessageContent", () => {
     ]);
   });
 
+  it("uses the binary data URL default for a media-less file", () => {
+    const result = getMessageContent(
+      makeAppendMessage([
+        {
+          type: "file",
+          mimeType: "",
+          data: "data:;base64,SGVsbG8=",
+        },
+      ]),
+    );
+    expect(result).toEqual([
+      {
+        type: "file",
+        mimeType: "application/octet-stream",
+        data: "SGVsbG8=",
+      },
+    ]);
+  });
+
   it("emits a file_url part for file parts with sourceType url", () => {
     const result = getMessageContent(
       makeAppendMessage([

package/src/convertToAdkMessages.ts

@@ -52,7 +52,7 @@ export const getMessageContent = (msg: AppendMessage) => {
         }
         return {
           type: "file" as const,
-          mimeType: part.mimeType,
+          mimeType: source.mimeType,
           // Lands in Gemini `inlineData.data`, which takes bare base64, so a
           // data URL envelope is stripped rather than forwarded.
           data: source.data,

package/src/hooks.test.tsx

@@ -1,3 +1,4 @@
+/** @vitest-environment jsdom */
 import { renderHook } from "@testing-library/react";
 import { describe, expect, it, vi } from "vitest";
 

package/src/server/createAdkApiRoute.controls.test.ts

@@ -0,0 +1,66 @@
+import { describe, expect, it, vi } from "vitest";
+import { createAdkApiRoute } from "./createAdkApiRoute";
+
+const makeRequest = (body: unknown) =>
+  new Request("https://example.test/api/adk", {
+    method: "POST",
+    body: JSON.stringify(body),
+  });
+
+describe("createAdkApiRoute request controls", () => {
+  it("does not trust client run configuration or state by default", async () => {
+    const runAsync = vi.fn(async function* () {});
+    const handler = createAdkApiRoute({
+      runner: { runAsync },
+      userId: "user-1",
+      sessionId: "session-1",
+    });
+
+    await handler(
+      makeRequest({
+        message: "Hello",
+        runConfig: { maxLlmCalls: 1_000_000 },
+        stateDelta: { "app:role": "admin" },
+      }),
+    );
+
+    expect(runAsync).toHaveBeenCalledWith({
+      userId: "user-1",
+      sessionId: "session-1",
+      newMessage: { role: "user", parts: [{ text: "Hello" }] },
+    });
+  });
+
+  it("forwards only values returned by server resolvers", async () => {
+    const runAsync = vi.fn(async function* () {});
+    const resolveRunConfig = vi.fn(() => ({ maxLlmCalls: 10 }));
+    const resolveStateDelta = vi.fn(() => ({ taskId: "server-task" }));
+    const handler = createAdkApiRoute({
+      runner: { runAsync },
+      userId: "user-1",
+      sessionId: "session-1",
+      resolveRunConfig,
+      resolveStateDelta,
+    });
+    const request = makeRequest({
+      message: "Hello",
+      runConfig: { maxLlmCalls: 1_000_000 },
+      stateDelta: { "app:role": "admin" },
+    });
+
+    await handler(request);
+
+    expect(resolveRunConfig).toHaveBeenCalledWith(request, {
+      maxLlmCalls: 1_000_000,
+    });
+    expect(resolveStateDelta).toHaveBeenCalledWith(request, {
+      "app:role": "admin",
+    });
+    expect(runAsync).toHaveBeenCalledWith(
+      expect.objectContaining({
+        runConfig: { maxLlmCalls: 10 },
+        stateDelta: { taskId: "server-task" },
+      }),
+    );
+  });
+});

package/src/server/createAdkApiRoute.test.ts

@@ -0,0 +1,282 @@
+import { describe, expect, it, vi } from "vitest";
+import { createAdkApiRoute } from "./createAdkApiRoute";
+
+describe("createAdkApiRoute", () => {
+  it("creates the client thread session before the first run", async () => {
+    const sessions = new Set<string>();
+    const getSession = vi.fn(async ({ sessionId }: { sessionId: string }) =>
+      sessions.has(sessionId) ? { id: sessionId } : undefined,
+    );
+    const createSession = vi.fn(
+      async ({ sessionId }: { sessionId: string }) => {
+        sessions.add(sessionId);
+        return { id: sessionId };
+      },
+    );
+    const runAsync = vi.fn((options: Record<string, unknown>) => {
+      expect(sessions.has(options.sessionId as string)).toBe(true);
+      return (async function* () {})();
+    });
+    const sessionId = vi.fn(
+      (_request: Request, clientSessionId: string | undefined) => {
+        if (!clientSessionId) throw new Error("Missing session ID");
+        return `scoped-${clientSessionId}`;
+      },
+    );
+    const handler = createAdkApiRoute({
+      runner: {
+        appName: "test-app",
+        sessionService: { getSession, createSession },
+        runAsync,
+      },
+      userId: "user-1",
+      sessionId,
+    });
+
+    await handler(
+      new Request("https://example.test/api/adk", {
+        method: "POST",
+        body: JSON.stringify({ message: "Hello", sessionId: "thread-1" }),
+      }),
+    );
+
+    expect(sessionId).toHaveBeenCalledWith(expect.any(Request), "thread-1");
+    expect(createSession).toHaveBeenCalledWith({
+      appName: "test-app",
+      userId: "user-1",
+      sessionId: "scoped-thread-1",
+    });
+    expect(runAsync).toHaveBeenCalledWith(
+      expect.objectContaining({
+        userId: "user-1",
+        sessionId: "scoped-thread-1",
+      }),
+    );
+  });
+
+  it("uses a separate session for each thread of the same user", async () => {
+    const getSession = vi.fn(async () => undefined);
+    const createSession = vi.fn(
+      async ({ sessionId }: { sessionId: string }) => ({
+        id: sessionId,
+      }),
+    );
+    const runAsync = vi.fn(async function* (
+      _options: Record<string, unknown>,
+    ) {});
+    const handler = createAdkApiRoute({
+      runner: {
+        appName: "test-app",
+        sessionService: { getSession, createSession },
+        runAsync,
+      },
+      userId: "user-1",
+      sessionId: (_request, clientSessionId) => {
+        if (!clientSessionId) throw new Error("Missing session ID");
+        return clientSessionId;
+      },
+    });
+
+    for (const threadId of ["thread-1", "thread-2"]) {
+      await handler(
+        new Request("https://example.test/api/adk", {
+          method: "POST",
+          body: JSON.stringify({ message: "Hello", sessionId: threadId }),
+        }),
+      );
+    }
+
+    expect(getSession).toHaveBeenCalledTimes(2);
+    expect(getSession).toHaveBeenNthCalledWith(1, {
+      appName: "test-app",
+      userId: "user-1",
+      sessionId: "thread-1",
+    });
+    expect(getSession).toHaveBeenNthCalledWith(2, {
+      appName: "test-app",
+      userId: "user-1",
+      sessionId: "thread-2",
+    });
+    expect(createSession).toHaveBeenCalledTimes(2);
+    expect(createSession).toHaveBeenNthCalledWith(1, {
+      appName: "test-app",
+      userId: "user-1",
+      sessionId: "thread-1",
+    });
+    expect(createSession).toHaveBeenNthCalledWith(2, {
+      appName: "test-app",
+      userId: "user-1",
+      sessionId: "thread-2",
+    });
+    expect(runAsync).toHaveBeenCalledTimes(2);
+    expect(runAsync.mock.calls[0]?.[0]).toMatchObject({
+      userId: "user-1",
+      sessionId: "thread-1",
+    });
+    expect(runAsync.mock.calls[1]?.[0]).toMatchObject({
+      userId: "user-1",
+      sessionId: "thread-2",
+    });
+  });
+
+  it("resolves a supplied session ID only within the authenticated user", async () => {
+    const getSession = vi.fn(async ({ userId }: { userId: string }) =>
+      userId === "user-a" ? { id: "user-a-session" } : undefined,
+    );
+    const createSession = vi.fn(
+      async ({ sessionId }: { sessionId: string }) => ({
+        id: sessionId,
+      }),
+    );
+    const runAsync = vi.fn(async function* () {});
+    const authenticatedUsers = new WeakMap<Request, string>();
+    const handler = createAdkApiRoute({
+      runner: {
+        appName: "test-app",
+        sessionService: { getSession, createSession },
+        runAsync,
+      },
+      userId: (request) => {
+        const userId = authenticatedUsers.get(request);
+        if (!userId) throw new Error("Unauthenticated");
+        return userId;
+      },
+      sessionId: (_request, clientSessionId) => {
+        if (!clientSessionId) throw new Error("Missing session ID");
+        return clientSessionId;
+      },
+    });
+    const request = new Request("https://example.test/api/adk", {
+      method: "POST",
+      body: JSON.stringify({ message: "Hello", sessionId: "user-a-session" }),
+    });
+    authenticatedUsers.set(request, "user-b");
+
+    await handler(request);
+
+    expect(getSession).toHaveBeenCalledExactlyOnceWith({
+      appName: "test-app",
+      userId: "user-b",
+      sessionId: "user-a-session",
+    });
+    expect(createSession).toHaveBeenCalledExactlyOnceWith({
+      appName: "test-app",
+      userId: "user-b",
+      sessionId: "user-a-session",
+    });
+    expect(runAsync).toHaveBeenCalledWith(
+      expect.objectContaining({
+        userId: "user-b",
+        sessionId: "user-a-session",
+      }),
+    );
+  });
+
+  it("shares session creation across concurrent first requests", async () => {
+    const getSession = vi.fn(async () => undefined);
+    let releaseCreation!: () => void;
+    let markCreationStarted!: () => void;
+    const creationStarted = new Promise<void>((resolve) => {
+      markCreationStarted = resolve;
+    });
+    const createSession = vi.fn(async () => {
+      markCreationStarted();
+      await new Promise<void>((resolve) => {
+        releaseCreation = resolve;
+      });
+      return { id: "thread-1" };
+    });
+    const runAsync = vi.fn(async function* () {});
+    const sessionId = vi.fn(
+      (_request: Request, clientSessionId: string | undefined) =>
+        clientSessionId ?? "missing",
+    );
+    const handler = createAdkApiRoute({
+      runner: {
+        appName: "test-app",
+        sessionService: { getSession, createSession },
+        runAsync,
+      },
+      userId: "user-1",
+      sessionId,
+    });
+    const makeRequest = () =>
+      new Request("https://example.test/api/adk", {
+        method: "POST",
+        body: JSON.stringify({ message: "Hello", sessionId: "thread-1" }),
+      });
+
+    const firstRequest = handler(makeRequest());
+    await creationStarted;
+    const secondRequest = handler(makeRequest());
+    await vi.waitFor(() => expect(sessionId).toHaveBeenCalledTimes(2));
+    await Promise.resolve();
+    releaseCreation();
+    await Promise.all([firstRequest, secondRequest]);
+
+    expect(getSession).toHaveBeenCalledTimes(1);
+    expect(createSession).toHaveBeenCalledTimes(1);
+    expect(runAsync).toHaveBeenCalledTimes(2);
+  });
+
+  it("continues when another process creates the session first", async () => {
+    let exists = false;
+    const getSession = vi.fn(async () =>
+      exists ? { id: "thread-1" } : undefined,
+    );
+    const createSession = vi.fn(async () => {
+      exists = true;
+      throw new Error("Session already exists");
+    });
+    const runAsync = vi.fn(async function* () {});
+    const handler = createAdkApiRoute({
+      runner: {
+        appName: "test-app",
+        sessionService: { getSession, createSession },
+        runAsync,
+      },
+      userId: "user-1",
+      sessionId: "thread-1",
+    });
+
+    await handler(
+      new Request("https://example.test/api/adk", {
+        method: "POST",
+        body: JSON.stringify({ message: "Hello" }),
+      }),
+    );
+
+    expect(getSession).toHaveBeenCalledTimes(2);
+    expect(createSession).toHaveBeenCalledOnce();
+    expect(runAsync).toHaveBeenCalledOnce();
+  });
+
+  it("preserves create failures when the session still does not exist", async () => {
+    const failure = new Error("database unavailable");
+    const getSession = vi.fn(async () => undefined);
+    const createSession = vi.fn(async () => {
+      throw failure;
+    });
+    const runAsync = vi.fn(async function* () {});
+    const handler = createAdkApiRoute({
+      runner: {
+        appName: "test-app",
+        sessionService: { getSession, createSession },
+        runAsync,
+      },
+      userId: "user-1",
+      sessionId: "thread-1",
+    });
+
+    await expect(
+      handler(
+        new Request("https://example.test/api/adk", {
+          method: "POST",
+          body: JSON.stringify({ message: "Hello" }),
+        }),
+      ),
+    ).rejects.toBe(failure);
+    expect(getSession).toHaveBeenCalledTimes(2);
+    expect(runAsync).not.toHaveBeenCalled();
+  });
+});

package/src/server/createAdkApiRoute.ts

@@ -6,11 +6,77 @@ import { adkEventStream, type AdkEventStreamOptions } from "./adkEventStream";
  * Avoids requiring `@google/adk` as a dependency.
  */
 type AdkRunner = {
+  readonly appName?: string;
+  readonly sessionService?: {
+    getSession(options: {
+      appName: string;
+      userId: string;
+      sessionId: string;
+    }): Promise<unknown | undefined>;
+    createSession(options: {
+      appName: string;
+      userId: string;
+      sessionId: string;
+    }): Promise<unknown>;
+  };
   runAsync(
     options: Record<string, unknown>,
   ): AsyncGenerator<any, void, undefined>;
 };
 
+type AdkSessionService = NonNullable<AdkRunner["sessionService"]>;
+
+const pendingSessions = new WeakMap<
+  AdkSessionService,
+  Map<string, Promise<void>>
+>();
+
+const ensureRunnerSession = async (
+  runner: AdkRunner,
+  userId: string,
+  sessionId: string,
+) => {
+  const { appName, sessionService } = runner;
+  if (!appName || !sessionService) return;
+
+  let serviceSessions = pendingSessions.get(sessionService);
+  if (!serviceSessions) {
+    serviceSessions = new Map();
+    pendingSessions.set(sessionService, serviceSessions);
+  }
+
+  const key = JSON.stringify([appName, userId, sessionId]);
+  let pending = serviceSessions.get(key);
+  if (!pending) {
+    pending = (async () => {
+      const session = await sessionService.getSession({
+        appName,
+        userId,
+        sessionId,
+      });
+      if (!session) {
+        try {
+          await sessionService.createSession({ appName, userId, sessionId });
+        } catch (error) {
+          const existing = await sessionService.getSession({
+            appName,
+            userId,
+            sessionId,
+          });
+          if (!existing) throw error;
+        }
+      }
+    })();
+    serviceSessions.set(key, pending);
+  }
+
+  try {
+    await pending;
+  } finally {
+    if (serviceSessions.get(key) === pending) serviceSessions.delete(key);
+  }
+};
+
 export type CreateAdkApiRouteOptions = {
   /**
    * ADK Runner instance.
@@ -18,16 +84,47 @@ export type CreateAdkApiRouteOptions = {
   runner: AdkRunner;
 
   /**
-   * User ID to use for the ADK session. Can be a static string
-   * or a function that extracts it from the request.
+   * User ID to use for the ADK session. Production routes should resolve this
+   * from the authenticated request. A static value is suitable only for a
+   * single-user development route.
    */
   userId: string | ((req: Request) => string | Promise<string>);
 
   /**
    * Session ID to use. Can be a static string or a function
-   * that extracts it from the request (e.g. from query params or headers).
+   * that validates or transforms the client thread ID sent by
+   * `createAdkStream`. The client value is an identifier, not authorization;
+   * scope access with an authenticated `userId`.
    */
-  sessionId: string | ((req: Request) => string | Promise<string>);
+  sessionId:
+    | string
+    | ((
+        req: Request,
+        clientSessionId: string | undefined,
+      ) => string | Promise<string>);
+
+  /**
+   * Validates or replaces the client-provided ADK run configuration.
+   * Client values are ignored unless this resolver is provided.
+   */
+  resolveRunConfig?:
+    | ((req: Request, runConfig: unknown) => unknown | Promise<unknown>)
+    | undefined;
+
+  /**
+   * Validates or replaces the client-provided ADK state delta.
+   * Client values are ignored unless this resolver is provided. In particular,
+   * `app:` and `user:` keys affect state beyond the current session.
+   */
+  resolveStateDelta?:
+    | ((
+        req: Request,
+        stateDelta: Record<string, unknown> | undefined,
+      ) =>
+        | Record<string, unknown>
+        | undefined
+        | Promise<Record<string, unknown> | undefined>)
+    | undefined;
 
   /**
    * Error handler for stream errors.
@@ -43,11 +140,15 @@ export type CreateAdkApiRouteOptions = {
  * ```ts
  * import { createAdkApiRoute } from '@assistant-ui/react-google-adk/server';
  * import { runner } from './agent';
+ * import { requireUser } from './auth';
  *
  * export const POST = createAdkApiRoute({
  *   runner,
- *   userId: "default-user",
- *   sessionId: (req) => new URL(req.url).searchParams.get("sessionId") ?? "default",
+ *   userId: async (req) => (await requireUser(req)).id,
+ *   sessionId: (_req, clientSessionId) => {
+ *     if (!clientSessionId) throw new Error("Missing ADK session ID");
+ *     return clientSessionId;
+ *   },
  * });
  * ```
  */
@@ -65,17 +166,24 @@ export function createAdkApiRoute(
 
     const sessionId =
       typeof options.sessionId === "function"
-        ? await options.sessionId(req)
+        ? await options.sessionId(req, parsed.sessionId)
         : options.sessionId;
 
+    const runConfig = options.resolveRunConfig
+      ? await options.resolveRunConfig(req, parsed.config.runConfig)
+      : undefined;
+    const stateDelta = options.resolveStateDelta
+      ? await options.resolveStateDelta(req, parsed.stateDelta)
+      : undefined;
+
+    await ensureRunnerSession(options.runner, userId, sessionId);
+
     const events = options.runner.runAsync({
       userId,
       sessionId,
       newMessage,
-      ...(parsed.stateDelta != null && { stateDelta: parsed.stateDelta }),
-      ...(parsed.config.runConfig != null && {
-        runConfig: parsed.config.runConfig,
-      }),
+      ...(stateDelta != null && { stateDelta }),
+      ...(runConfig != null && { runConfig }),
     });
 
     return adkEventStream(