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

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(

package/src/server/parseAdkRequest.test.ts

@@ -28,11 +28,18 @@ describe("parseAdkRequest", () => {
     });
   });
 
-  it("parses a message request with stateDelta", async () => {
+  it("parses a message request with sessionId and stateDelta", async () => {
     const result = await parseAdkRequest(
-      makeRequest({ message: "Hello", stateDelta: { count: 1 } }),
+      makeRequest({
+        message: "Hello",
+        sessionId: "session-1",
+        stateDelta: { count: 1 },
+      }),
     );
-    expect(result).toMatchObject({ stateDelta: { count: 1 } });
+    expect(result).toMatchObject({
+      sessionId: "session-1",
+      stateDelta: { count: 1 },
+    });
   });
 
   it("parses a message request with parts (multimodal)", async () => {
@@ -130,6 +137,7 @@ describe("parseAdkRequest", () => {
     [{ parts: [null] }, 'field "parts"'],
     [{ type: "unknown", message: "hello" }, 'field "type"'],
     [{ message: "hello", checkpointId: 42 }, 'field "checkpointId"'],
+    [{ message: "hello", sessionId: 42 }, 'field "sessionId"'],
     [{ message: "hello", stateDelta: [] }, 'field "stateDelta"'],
   ])("rejects malformed message requests %#", async (body, error) => {
     await expect(parseAdkRequest(makeRequest(body))).rejects.toThrow(error);

package/src/server/parseAdkRequest.ts

@@ -7,6 +7,7 @@ type ParsedAdkRequest =
       type: "message";
       text: string;
       parts?: Array<Record<string, unknown>> | undefined;
+      sessionId?: string | undefined;
       config: AdkSendMessageConfig;
       stateDelta?: Record<string, unknown> | undefined;
     }
@@ -16,6 +17,7 @@ type ParsedAdkRequest =
       toolName: string;
       result: unknown;
       isError: boolean;
+      sessionId?: string | undefined;
       config: AdkSendMessageConfig;
       stateDelta?: Record<string, unknown> | undefined;
     };
@@ -162,6 +164,7 @@ export const parseAdkRequest = async (
   if (body.runConfig !== undefined) config.runConfig = body.runConfig;
   const checkpointId = readOptionalString(body, "checkpointId");
   if (checkpointId !== undefined) config.checkpointId = checkpointId;
+  const sessionId = readOptionalString(body, "sessionId");
 
   const stateDelta = body.stateDelta;
   if (stateDelta !== undefined && !isRecord(stateDelta)) {
@@ -181,6 +184,7 @@ export const parseAdkRequest = async (
       toolName: readString(body, "toolName"),
       result: body.result,
       isError: body.isError ?? false,
+      ...(sessionId !== undefined && { sessionId }),
       config,
       ...(stateDelta != null && { stateDelta }),
     };
@@ -213,6 +217,7 @@ export const parseAdkRequest = async (
     type: "message",
     text: text ?? "",
     ...(parts !== undefined && { parts }),
+    ...(sessionId !== undefined && { sessionId }),
     config,
     ...(stateDelta != null && { stateDelta }),
   };
@@ -226,7 +231,8 @@ export const parseAdkRequest = async (
  * ```ts
  * const parsed = await parseAdkRequest(req);
  * const newMessage = toAdkContent(parsed);
- * const events = runner.runAsync({ userId, sessionId, newMessage, stateDelta: parsed.stateDelta });
+ * const stateDelta = validateSessionState(parsed.stateDelta);
+ * const events = runner.runAsync({ userId, sessionId, newMessage, stateDelta });
  * return adkEventStream(events);
  * ```
  */

package/src/useAdkMessages.test.ts

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