@genkit-ai/ai/tests/agent_test.ts

Genkit AI framework generative AI APIs.

/**
 * Copyright 2026 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import { initNodeFeatures } from '@genkit-ai/core/node';
import { Registry } from '@genkit-ai/core/registry';
import { enableTelemetry } from '@genkit-ai/core/tracing';
import { SimpleSpanProcessor } from '@opentelemetry/sdk-trace-base';
import * as assert from 'assert';
import { describe, it } from 'node:test';

import { z } from '@genkit-ai/core';
import { TestSpanExporter } from '../../core/tests/utils.js';
import { AgentError } from '../src/agent-core.js';
import {
  AgentStreamChunk,
  SessionRunner,
  defineAgent,
  defineCustomAgent,
  definePromptAgent,
} from '../src/agent.js';
import { definePrompt } from '../src/prompt.js';
import { InMemorySessionStore } from '../src/session-stores.js';
import {
  Session,
  reserveSnapshotId,
  type SessionSnapshot,
} from '../src/session.js';
import { ToolInterruptError, defineTool, interrupt } from '../src/tool.js';
import {
  defineEchoModel,
  defineProgrammableModel,
  type ProgrammableModel,
} from './helpers.js';

initNodeFeatures();

const spanExporter = new TestSpanExporter();
enableTelemetry({
  spanProcessors: [new SimpleSpanProcessor(spanExporter)],
});

/**
 * Returns a Promise that resolves once the given snapshotId reaches targetStatus
 * in the store. Rejects after timeoutMs if the status is never reached.
 */
function waitForSnapshotStatus<S>(
  store: InMemorySessionStore<S>,
  snapshotId: string,
  targetStatus: NonNullable<SessionSnapshot<S>['status']>,
  timeoutMs = 5000
): Promise<SessionSnapshot<S>> {
  return new Promise((resolve, reject) => {
    const timer = setTimeout(
      () =>
        reject(
          new Error(
            `Timed out waiting for snapshot ${snapshotId} to reach status "${targetStatus}"`
          )
        ),
      timeoutMs
    );

    const unsubscribeFn = store.onSnapshotStateChange(snapshotId, (snap) => {
      if (snap.status === targetStatus) {
        clearTimeout(timer);
        if (typeof unsubscribeFn === 'function') unsubscribeFn();
        resolve(snap);
      }
    });

    // Check in case already at the target status.
    store.getSnapshot({ snapshotId: snapshotId }).then((snap) => {
      if (snap?.status === targetStatus) {
        clearTimeout(timer);
        if (typeof unsubscribeFn === 'function') unsubscribeFn();
        resolve(snap);
      }
    });
  });
}

describe('Agent', () => {
  describe('Session', () => {
    it('should maintain custom state', () => {
      const session = new Session<{ foo: string }>({ custom: { foo: 'bar' } });
      assert.strictEqual(session.getCustom()?.foo, 'bar');

      session.updateCustom((c) => ({ ...c!, foo: 'baz' }));
      assert.strictEqual(session.getCustom()?.foo, 'baz');
    });

    it('should add and set messages', () => {
      const session = new Session({});
      session.addMessages([{ role: 'user', content: [{ text: 'hi' }] }]);
      assert.strictEqual(session.getMessages().length, 1);
      assert.strictEqual(session.getMessages()[0].role, 'user');

      session.setMessages([{ role: 'model', content: [{ text: 'hello' }] }]);
      assert.strictEqual(session.getMessages().length, 1);
      assert.strictEqual(session.getMessages()[0].role, 'model');
    });

    it('should add and deduplicate artifacts', () => {
      const session = new Session({});
      session.addArtifacts([{ name: 'art1', parts: [{ text: 'content1' }] }]);
      assert.strictEqual(session.getArtifacts().length, 1);

      // Add with same name should replace
      session.addArtifacts([{ name: 'art1', parts: [{ text: 'content2' }] }]);
      assert.strictEqual(session.getArtifacts().length, 1);
      assert.deepStrictEqual(session.getArtifacts()[0].parts, [
        { text: 'content2' },
      ]);

      // Add with different name should append
      session.addArtifacts([{ name: 'art2', parts: [{ text: 'content3' }] }]);
      assert.strictEqual(session.getArtifacts().length, 2);
    });

    it('should process all artifacts in a batch without dropping any', () => {
      const session = new Session({});
      session.addArtifacts([{ name: 'art1', parts: [{ text: 'v1' }] }]);

      // Replace art1 and add art2 and art3 in the same batch.
      session.addArtifacts([
        { name: 'art1', parts: [{ text: 'v2' }] },
        { name: 'art2', parts: [{ text: 'new' }] },
        { name: 'art3', parts: [{ text: 'another' }] },
      ]);

      const arts = session.getArtifacts();
      assert.strictEqual(arts.length, 3);
      assert.strictEqual(
        arts.find((a) => a.name === 'art1')?.parts[0].text,
        'v2'
      );
      assert.strictEqual(
        arts.find((a) => a.name === 'art2')?.parts[0].text,
        'new'
      );
      assert.strictEqual(
        arts.find((a) => a.name === 'art3')?.parts[0].text,
        'another'
      );
    });

    it('should emit artifactAdded for new and artifactUpdated for replaced', () => {
      const session = new Session({});
      const added: string[] = [];
      const updated: string[] = [];
      session.on('artifactAdded', (a: { name?: string }) =>
        added.push(a.name ?? '')
      );
      session.on('artifactUpdated', (a: { name?: string }) =>
        updated.push(a.name ?? '')
      );

      session.addArtifacts([{ name: 'art1', parts: [] }]);
      session.addArtifacts([
        { name: 'art1', parts: [] }, // replace
        { name: 'art2', parts: [] }, // new
      ]);

      assert.deepStrictEqual(added, ['art1', 'art2']);
      assert.deepStrictEqual(updated, ['art1']);
    });

    it('should increment version on mutation', () => {
      const session = new Session({});
      const v0 = session.getVersion();

      session.addMessages([{ role: 'user', content: [{ text: 'hi' }] }]);
      const v1 = session.getVersion();
      assert.ok(v1 > v0);

      session.updateCustom((c) => c);
      const v2 = session.getVersion();
      assert.ok(v2 > v1);

      session.addArtifacts([{ name: 'a', parts: [] }]);
      const v3 = session.getVersion();
      assert.ok(v3 > v2);
    });
  });

  describe('SessionRunner', () => {
    it('should loop over inputs and call handler', async () => {
      const session = new Session({});
      const inputs = [
        { message: { role: 'user' as const, content: [{ text: 'hi' }] } },
        { message: { role: 'user' as const, content: [{ text: 'bye' }] } },
      ];

      async function* inputGen() {
        for (const input of inputs) {
          yield input;
        }
      }

      const runner = new SessionRunner(session, inputGen());
      let turns = 0;
      const seenInputs: any[] = [];

      await runner.run(async (input) => {
        turns++;
        seenInputs.push(input);
      });

      assert.strictEqual(turns, 2);
      assert.deepStrictEqual(seenInputs, inputs);
      assert.strictEqual(session.getMessages().length, 2);
    });

    it('should trigger snapshots if store is present', async () => {
      const store = new InMemorySessionStore();
      const session = new Session({});
      const inputs = [
        { message: { role: 'user' as const, content: [{ text: 'hi' }] } },
      ];

      async function* inputGen() {
        for (const input of inputs) {
          yield input;
        }
      }

      let turnEnded = false;
      let turnSnapshotId: string | undefined;

      const runner = new SessionRunner(session, inputGen(), {
        store,
        onEndTurn: (snapshotId) => {
          turnEnded = true;
          turnSnapshotId = snapshotId;
        },
      });

      await runner.run(async () => {});

      assert.ok(turnEnded);
      assert.ok(turnSnapshotId);

      const saved = await store.getSnapshot({ snapshotId: turnSnapshotId! });
      assert.ok(saved);
      assert.strictEqual(saved?.snapshotId, turnSnapshotId);
    });

    it('reserves the turn snapshotId at turn start and persists under it', async () => {
      const store = new InMemorySessionStore();
      const session = new Session({});

      async function* inputGen() {
        yield {
          message: { role: 'user' as const, content: [{ text: 'hi' }] },
        };
      }

      let ctxSnapshotId: string | undefined;
      let ctxParentSnapshotId: string | undefined;
      let ctxTurnIndex: number | undefined;
      let endTurnSnapshotId: string | undefined;

      const runner = new SessionRunner(session, inputGen(), {
        store,
        onEndTurn: (snapshotId) => {
          endTurnSnapshotId = snapshotId;
        },
      });

      await runner.run(async (_input, ctx) => {
        ctxSnapshotId = ctx.snapshotId;
        ctxParentSnapshotId = ctx.parentSnapshotId;
        ctxTurnIndex = ctx.turnIndex;
      });

      // The id is reserved at turn start and made available to the handler.
      assert.ok(ctxSnapshotId, 'handler should receive a reserved snapshotId');
      // First turn of a fresh session has no parent.
      assert.strictEqual(ctxParentSnapshotId, undefined);
      assert.strictEqual(ctxTurnIndex, 0);

      // The snapshot persisted at turn end reuses the reserved id.
      assert.strictEqual(endTurnSnapshotId, ctxSnapshotId);
      const saved = await store.getSnapshot({ snapshotId: ctxSnapshotId! });
      assert.ok(saved);
      assert.strictEqual(saved?.snapshotId, ctxSnapshotId);
    });

    it('passes the prior snapshotId as parentSnapshotId on subsequent turns', async () => {
      const store = new InMemorySessionStore();
      const session = new Session({});

      async function* inputGen() {
        yield {
          message: { role: 'user' as const, content: [{ text: 'one' }] },
        };
        yield {
          message: { role: 'user' as const, content: [{ text: 'two' }] },
        };
      }

      const snapshotIds: string[] = [];
      const parentIds: Array<string | undefined> = [];

      const runner = new SessionRunner(session, inputGen(), { store });

      await runner.run(async (_input, ctx) => {
        snapshotIds.push(ctx.snapshotId);
        parentIds.push(ctx.parentSnapshotId);
      });

      assert.strictEqual(snapshotIds.length, 2);
      // First turn: no parent. Second turn: parent is the first turn's snapshot.
      assert.strictEqual(parentIds[0], undefined);
      assert.strictEqual(parentIds[1], snapshotIds[0]);
    });

    it('does not reserve a snapshotId when no store is configured', async () => {
      const session = new Session({});

      async function* inputGen() {
        yield {
          message: { role: 'user' as const, content: [{ text: 'hi' }] },
        };
      }

      let ctxSnapshotId: string | undefined = 'sentinel';
      const runner = new SessionRunner(session, inputGen());
      await runner.run(async (_input, ctx) => {
        ctxSnapshotId = ctx.snapshotId;
      });

      // Without a store there is nothing to reserve, so snapshotId is undefined.
      assert.strictEqual(ctxSnapshotId, undefined);
    });
  });

  describe('reserveSnapshotId', () => {
    const UUID =
      /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i;

    it('mints a plain UUID snapshotId', () => {
      const id = reserveSnapshotId();
      assert.match(id, UUID, `id should be a UUID: ${id}`);
    });

    it('produces unique ids on successive calls', () => {
      const a = reserveSnapshotId();
      const b = reserveSnapshotId();
      assert.notStrictEqual(a, b);
    });
  });

  describe('defineCustomAgent', () => {
    it('should set client stateManagement and abortable=false when no store is provided', () => {
      const registry = new Registry();
      const agent = defineCustomAgent(
        registry,
        { name: 'noStoreMetadataTest' },
        async () => ({ artifacts: [] })
      );
      assert.strictEqual(
        agent.__action.metadata?.agent?.stateManagement,
        'client'
      );
      assert.strictEqual(agent.__action.metadata?.agent?.abortable, false);
    });

    it('should set server stateManagement and abortable=true when store with onSnapshotStateChange is provided', () => {
      const registry = new Registry();
      const store = new InMemorySessionStore();
      const agent = defineCustomAgent(
        registry,
        { name: 'fullStoreMetadataTest', store },
        async () => ({ artifacts: [] })
      );
      assert.strictEqual(
        agent.__action.metadata?.agent?.stateManagement,
        'server'
      );
      assert.strictEqual(agent.__action.metadata?.agent?.abortable, true);
    });

    it('should reject init.state for server-managed agents (store is set)', async () => {
      const registry = new Registry();
      const store = new InMemorySessionStore<{ foo: string }>();

      const flow = defineCustomAgent<{ foo: string }>(
        registry,
        { name: 'rejectInitStateTest', store },
        async (sess) => {
          await sess.run(async () => {});
          return {
            artifacts: [],
            message: { role: 'model', content: [{ text: 'done' }] },
          };
        }
      );

      // Pass init.state - this is API misuse for a server-managed agent and
      // must throw a FAILED_PRECONDITION error (rather than resolving with a
      // graceful finishReason 'failed' output).
      const session = flow.streamBidi({
        state: {
          custom: { foo: 'should-be-rejected' },
          messages: [{ role: 'user', content: [{ text: 'stale history' }] }],
          artifacts: [],
        },
      });
      session.send({
        message: { role: 'user', content: [{ text: 'hello' }] },
      });
      session.close();

      let thrown: any;
      try {
        for await (const _ of session.stream) {
        }
        await session.output;
      } catch (e: any) {
        thrown = e;
      }

      assert.ok(thrown, 'Expected the turn to throw an error');
      assert.strictEqual(thrown.status, 'FAILED_PRECONDITION');
      assert.ok(
        thrown.message.includes("Cannot send 'state' to agent"),
        `Expected FAILED_PRECONDITION message, got: ${thrown.message}`
      );
    });

    it('should use init.state for client-managed agents (no store)', async () => {
      const registry = new Registry();

      const flow = defineCustomAgent<{ foo: string }>(
        registry,
        { name: 'useInitStateTest' },
        async (sess) => {
          await sess.run(async () => {});
          return {
            artifacts: [],
            message: { role: 'model', content: [{ text: 'done' }] },
          };
        }
      );

      // Pass init.state - it should be used because no store is set
      const session = flow.streamBidi({
        state: {
          custom: { foo: 'seeded' },
          messages: [{ role: 'user', content: [{ text: 'prior msg' }] }],
          artifacts: [],
        },
      });
      session.send({
        message: { role: 'user', content: [{ text: 'hello' }] },
      });
      session.close();

      for await (const _ of session.stream) {
      }
      const output = await session.output;

      // State should include the seeded state plus the new message
      assert.ok(output.state);
      assert.strictEqual((output.state!.custom as any).foo, 'seeded');
      // Messages: 1 from init.state + 1 from input
      assert.strictEqual(output.state!.messages!.length, 2);
      assert.strictEqual(
        output.state!.messages![0].content[0].text,
        'prior msg'
      );
      assert.strictEqual(output.state!.messages![1].content[0].text, 'hello');
    });

    it('should set server stateManagement and abortable=false when store lacks onSnapshotStateChange', () => {
      const registry = new Registry();
      const store: any = {
        getSnapshot: async () => undefined,
        saveSnapshot: async () => {},
        // no onSnapshotStateChange
      };
      const agent = defineCustomAgent(
        registry,
        { name: 'noAbortStoreMetadataTest', store },
        async () => ({ artifacts: [] })
      );
      assert.strictEqual(
        agent.__action.metadata?.agent?.stateManagement,
        'server'
      );
      assert.strictEqual(agent.__action.metadata?.agent?.abortable, false);
    });

    it('should register and execute agent', async () => {
      const registry = new Registry();

      const flow = defineCustomAgent(
        registry,
        { name: 'testFlow' },
        async (sess, { sendChunk }) => {
          let receivedInput = false;
          await sess.run(async (input) => {
            receivedInput = true;
            assert.strictEqual(input.message?.role, 'user');
          });
          assert.ok(receivedInput);
          return { message: { role: 'model', content: [{ text: 'done' }] } };
        }
      );

      const session = flow.streamBidi({});

      session.send({
        message: { role: 'user' as const, content: [{ text: 'hi' }] },
      });
      session.close();

      const chunks: AgentStreamChunk[] = [];
      for await (const chunk of session.stream) {
        chunks.push(chunk);
      }

      const output = await session.output;
      assert.strictEqual(output.message?.role, 'model');
      assert.strictEqual(output.message?.content[0].text, 'done');
    });

    it('should automatically stream artifacts added via Session.addArtifacts()', async () => {
      const registry = new Registry();

      const flow = defineCustomAgent(
        registry,
        { name: 'testEventFlow' },
        async (sess, { sendChunk }) => {
          await sess.run(async (input) => {
            sess.session.addArtifacts([
              { name: 'testArt', parts: [{ text: 'testPart' }] },
            ]);
          });
          return { message: { role: 'model', content: [{ text: 'done' }] } };
        }
      );

      const session = flow.streamBidi({});
      session.send({
        message: { role: 'user' as const, content: [{ text: 'hi' }] },
      });
      session.close();

      const chunks: AgentStreamChunk[] = [];
      for await (const chunk of session.stream) {
        chunks.push(chunk);
      }

      const artChunks = chunks.filter((c) => !!c.artifact);
      assert.strictEqual(artChunks.length, 1);
      assert.strictEqual(artChunks[0].artifact?.name, 'testArt');
    });

    it('should stream artifactUpdated chunks when an artifact is replaced', async () => {
      const registry = new Registry();

      const flow = defineCustomAgent(
        registry,
        { name: 'testArtifactUpdateFlow' },
        async (sess) => {
          await sess.run(async () => {
            sess.session.addArtifacts([{ name: 'a', parts: [{ text: 'v1' }] }]);
            sess.session.addArtifacts([{ name: 'a', parts: [{ text: 'v2' }] }]);
          });
          return {};
        }
      );

      const session = flow.streamBidi({});
      session.send({ message: { role: 'user', content: [{ text: 'go' }] } });
      session.close();

      const chunks: AgentStreamChunk[] = [];
      for await (const chunk of session.stream) {
        chunks.push(chunk);
      }

      const artChunks = chunks.filter((c) => !!c.artifact);
      assert.strictEqual(artChunks.length, 2);
      assert.strictEqual(artChunks[0].artifact?.parts[0].text, 'v1');
      assert.strictEqual(artChunks[1].artifact?.parts[0].text, 'v2');
    });

    it('records the snapshotId and state on the turn span (server-managed)', async () => {
      spanExporter.exportedSpans = [];
      const store = new InMemorySessionStore<{ count: number }>();

      const flow = defineCustomAgent<{ count: number }>(
        new Registry(),
        { name: 'turnSpanServerTest', store },
        async (sess) => {
          await sess.run(async () => {
            sess.updateCustom((c) => ({ count: (c?.count ?? 0) + 1 }));
            return { finishReason: 'stop' as const };
          });
          return {
            message: { role: 'model', content: [{ text: 'done' }] },
            finishReason: 'stop' as const,
          };
        }
      );

      const session = flow.streamBidi({});
      session.send({
        message: { role: 'user' as const, content: [{ text: 'hi' }] },
      });
      session.close();

      for await (const _ of session.stream) {
      }
      const output = await session.output;
      assert.ok(output.snapshotId);

      const turnSpan = spanExporter.exportedSpans.find(
        (s) => s.displayName === 'runTurn-1'
      );
      assert.ok(turnSpan, 'expected a runTurn-1 span to be exported');

      // The turn span carries the snapshotId this turn persisted under.
      assert.strictEqual(
        turnSpan.attributes['genkit:metadata:agent:snapshotId'],
        output.snapshotId
      );

      // The turn span's output is the session state this turn produced.
      const out = JSON.parse(turnSpan.attributes['genkit:output']);
      assert.deepStrictEqual(out.state.custom, { count: 1 });
    });

    it('records state on the turn span for a client-managed agent', async () => {
      spanExporter.exportedSpans = [];
      const registry = new Registry();

      const flow = defineCustomAgent<{ count: number }>(
        registry,
        { name: 'turnSpanClientTest' },
        async (sess) => {
          await sess.run(async () => {
            sess.updateCustom((c) => ({ count: (c?.count ?? 0) + 1 }));
            return { finishReason: 'stop' as const };
          });
          return {
            message: { role: 'model', content: [{ text: 'done' }] },
            finishReason: 'stop' as const,
          };
        }
      );

      const session = flow.streamBidi({});
      session.send({
        message: { role: 'user' as const, content: [{ text: 'hi' }] },
      });
      session.close();

      for await (const _ of session.stream) {
      }
      await session.output;

      const turnSpan = spanExporter.exportedSpans.find(
        (s) => s.displayName === 'runTurn-1'
      );
      assert.ok(turnSpan, 'expected a runTurn-1 span to be exported');

      // The turn span's output carries the session state this turn produced,
      // for client-managed agents too.
      const out = JSON.parse(turnSpan.attributes['genkit:output']);
      assert.deepStrictEqual(out.state.custom, { count: 1 });
    });
  });

  describe('sessionId', () => {
    it('should generate sessionId for a fresh client-managed agent', async () => {
      const registry = new Registry();

      const flow = defineCustomAgent(
        registry,
        { name: 'sessionIdFreshClient' },
        async (sess) => {
          await sess.run(async () => {});
          return { message: { role: 'model', content: [{ text: 'done' }] } };
        }
      );

      const session = flow.streamBidi({});
      session.send({
        message: { role: 'user' as const, content: [{ text: 'hi' }] },
      });
      session.close();

      for await (const _ of session.stream) {
      }
      const output = await session.output;

      // Client-managed agents return state, which should contain a sessionId
      assert.ok(output.state, 'output.state should be present');
      assert.ok(
        output.state!.sessionId,
        'sessionId should be generated for a fresh session'
      );
      // Should be a valid UUID
      assert.match(
        output.state!.sessionId!,
        /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i
      );
    });

    it('should preserve sessionId across turns for client-managed agents', async () => {
      const registry = new Registry();

      const flow = defineCustomAgent(
        registry,
        { name: 'sessionIdPreserveClient' },
        async (sess) => {
          await sess.run(async () => {});
          return { message: { role: 'model', content: [{ text: 'done' }] } };
        }
      );

      // Turn 1: fresh session
      const session1 = flow.streamBidi({});
      session1.send({
        message: { role: 'user' as const, content: [{ text: 'hi' }] },
      });
      session1.close();
      for await (const _ of session1.stream) {
      }
      const output1 = await session1.output;
      const firstSessionId = output1.state!.sessionId!;
      assert.ok(firstSessionId, 'First turn should have sessionId');

      // Turn 2: pass state back (client-managed)
      const session2 = flow.streamBidi({ state: output1.state });
      session2.send({
        message: { role: 'user' as const, content: [{ text: 'bye' }] },
      });
      session2.close();
      for await (const _ of session2.stream) {
      }
      const output2 = await session2.output;

      assert.strictEqual(
        output2.state!.sessionId,
        firstSessionId,
        'sessionId should be preserved across turns'
      );
    });

    it('should generate sessionId for a fresh server-managed agent and persist in snapshot', async () => {
      const registry = new Registry();
      const store = new InMemorySessionStore();

      const flow = defineCustomAgent(
        registry,
        { name: 'sessionIdServerManaged', store },
        async (sess) => {
          await sess.run(async () => {});
          return { message: { role: 'model', content: [{ text: 'done' }] } };
        }
      );

      const session = flow.streamBidi({});
      session.send({
        message: { role: 'user' as const, content: [{ text: 'hi' }] },
      });
      session.close();

      for await (const _ of session.stream) {
      }
      const output = await session.output;

      assert.ok(output.snapshotId, 'should have snapshotId');

      // Read snapshot and verify sessionId is persisted in the state
      const snapshot = await store.getSnapshot({
        snapshotId: output.snapshotId!,
      });
      assert.ok(snapshot, 'snapshot should exist');
      assert.ok(
        snapshot!.state.sessionId,
        'sessionId should be persisted in snapshot state'
      );
      assert.match(
        snapshot!.state.sessionId!,
        /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i
      );
    });

    it('should preserve sessionId from snapshot for server-managed agents', async () => {
      const registry = new Registry();
      const store = new InMemorySessionStore();

      const flow = defineCustomAgent(
        registry,
        { name: 'sessionIdServerPreserve', store },
        async (sess) => {
          await sess.run(async () => {});
          return { message: { role: 'model', content: [{ text: 'done' }] } };
        }
      );

      // Turn 1
      const session1 = flow.streamBidi({});
      session1.send({
        message: { role: 'user' as const, content: [{ text: 'hi' }] },
      });
      session1.close();
      for await (const _ of session1.stream) {
      }
      const output1 = await session1.output;
      const firstSnapshotId = output1.snapshotId!;

      const snapshot1 = await store.getSnapshot({
        snapshotId: firstSnapshotId,
      });
      const firstSessionId = snapshot1!.state.sessionId!;
      assert.ok(firstSessionId, 'First turn should have sessionId');

      // Turn 2: resume from snapshot
      const session2 = flow.streamBidi({ snapshotId: firstSnapshotId });
      session2.send({
        message: { role: 'user' as const, content: [{ text: 'bye' }] },
      });
      session2.close();
      for await (const _ of session2.stream) {
      }
      const output2 = await session2.output;

      const snapshot2 = await store.getSnapshot({
        snapshotId: output2.snapshotId!,
      });
      assert.strictEqual(
        snapshot2!.state.sessionId,
        firstSessionId,
        'sessionId should be preserved across turns via snapshot'
      );
    });

    it('resumes a snapshot when both snapshotId and a matching sessionId are provided', async () => {
      const registry = new Registry();
      const store = new InMemorySessionStore();

      const flow = defineCustomAgent(
        registry,
        { name: 'snapshotAndMatchingSession', store },
        async (sess) => {
          await sess.run(async () => {});
          return { message: { role: 'model', content: [{ text: 'done' }] } };
        }
      );

      // Turn 1
      const session1 = flow.streamBidi({});
      session1.send({
        message: { role: 'user' as const, content: [{ text: 'hi' }] },
      });
      session1.close();
      for await (const _ of session1.stream) {
      }
      const output1 = await session1.output;
      const firstSnapshotId = output1.snapshotId!;
      const sessionId = output1.sessionId!;
      assert.ok(sessionId, 'First turn should have sessionId');

      // Turn 2: resume passing BOTH the snapshotId and its owning sessionId.
      const session2 = flow.streamBidi({
        snapshotId: firstSnapshotId,
        sessionId,
      });
      session2.send({
        message: { role: 'user' as const, content: [{ text: 'bye' }] },
      });
      session2.close();
      for await (const _ of session2.stream) {
      }
      const output2 = await session2.output;

      assert.notStrictEqual(
        output2.finishReason,
        'failed',
        `Expected a successful resume, got error: ${output2.error?.message}`
      );
      assert.strictEqual(
        output2.sessionId,
        sessionId,
        'sessionId should be preserved when resuming by snapshotId + sessionId'
      );
    });

    it('rejects when snapshotId belongs to a different session than the provided sessionId', async () => {
      const registry = new Registry();
      const store = new InMemorySessionStore();

      const flow = defineCustomAgent(
        registry,
        { name: 'snapshotSessionMismatch', store },
        async (sess) => {
          await sess.run(async () => {});
          return { message: { role: 'model', content: [{ text: 'done' }] } };
        }
      );

      // Create a snapshot that belongs to session A.
      const session1 = flow.streamBidi({});
      session1.send({
        message: { role: 'user' as const, content: [{ text: 'hi' }] },
      });
      session1.close();
      for await (const _ of session1.stream) {
      }
      const output1 = await session1.output;
      const snapshotId = output1.snapshotId!;

      // Resume that snapshot but claim it belongs to a different session. This
      // is API misuse and must throw (rather than resolving gracefully).
      const session2 = flow.streamBidi({
        snapshotId,
        sessionId: 'a-different-session-id',
      });
      session2.send({
        message: { role: 'user' as const, content: [{ text: 'bye' }] },
      });
      session2.close();

      let thrown: any;
      try {
        for await (const _ of session2.stream) {
        }
        await session2.output;
      } catch (e: any) {
        thrown = e;
      }

      assert.ok(thrown, 'Expected the turn to throw an error');
      assert.strictEqual(thrown.status, 'INVALID_ARGUMENT');
      assert.ok(
        thrown.message.includes('does not belong to session'),
        `Expected an ownership-mismatch error, got: ${thrown.message}`
      );
    });
  });

  describe('definePromptAgent', () => {
    it('should register and execute agent from prompt', async () => {
      const registry = new Registry();
      defineEchoModel(registry);
      definePrompt(registry, {
        name: 'agent',
        model: 'echoModel',
        config: { temperature: 1 },
        system: 'hello from template',
      });

      const flow = definePromptAgent(registry, {
        promptName: 'agent',
      });

      const session = flow.streamBidi({});
      session.send({
        message: { role: 'user' as const, content: [{ text: 'hi' }] },
      });
      session.close();

      const chunks: AgentStreamChunk[] = [];
      for await (const chunk of session.stream) {
        chunks.push(chunk);
      }

      const output = await session.output;
      assert.strictEqual(output.message?.role, 'model');
    });

    it('should pass promptInput to the prompt template', async () => {
      const registry = new Registry();
      defineEchoModel(registry);
      definePrompt(registry, {
        name: 'personaAgentPrompt',
        model: 'echoModel',
        input: { schema: z.object({ persona: z.string() }) },
        system: 'You are a {{persona}}.',
      });

      const flow = definePromptAgent<unknown, z.ZodTypeAny>(registry, {
        promptName: 'personaAgentPrompt',
        promptInput: { persona: 'pirate' },
      });

      const session = flow.streamBidi({});
      session.send({
        message: { role: 'user' as const, content: [{ text: 'hi' }] },
      });
      session.close();

      for await (const _ of session.stream) {
        // drain
      }

      const output = await session.output;
      // The echo model echoes the rendered system message, which interpolated
      // the supplied promptInput.
      assert.ok(
        output.message?.content[0].text?.includes('You are a pirate.'),
        `expected rendered prompt to include the promptInput, got: ${output.message?.content[0].text}`
      );
    });

    it('should let agents in separate registries reuse one prompt definition with different promptInput', async () => {
      // The same prompt definition (config) reused to build two differently
      // customized agents. Each lives in its own registry because an agent is
      // keyed by its promptName.
      const promptConfig = {
        name: 'sharedPersonaPrompt',
        model: 'echoModel',
        input: { schema: z.object({ persona: z.string() }) },
        system: 'You are a {{persona}}.',
      };

      const buildAgent = (persona: string) => {
        const registry = new Registry();
        defineEchoModel(registry);
        definePrompt(registry, promptConfig);
        return definePromptAgent<unknown, z.ZodTypeAny>(registry, {
          promptName: 'sharedPersonaPrompt',
          promptInput: { persona },
        });
      };

      const runAgent = async (flow: ReturnType<typeof buildAgent>) => {
        const session = flow.streamBidi({});
        session.send({
          message: { role: 'user' as const, content: [{ text: 'hi' }] },
        });
        session.close();
        for await (const _ of session.stream) {
          // drain
        }
        return session.output;
      };

      const pirateOut = await runAgent(buildAgent('pirate'));
      const ninjaOut = await runAgent(buildAgent('ninja'));

      assert.ok(pirateOut.message?.content[0].text?.includes('pirate'));
      assert.ok(ninjaOut.message?.content[0].text?.includes('ninja'));
    });

    it('should detach asynchronously and continue execution in the background', async () => {
      const store = new InMemorySessionStore<{ foo: string }>();
      let resolvePromise: () => void = () => {};
      const releasePromise = new Promise<void>((resolve) => {
        resolvePromise = resolve;
      });

      const flow = defineCustomAgent<{ foo: string }>(
        new Registry(),
        {
          name: 'detachTest',
          store,
        },
        async (sess, { sendChunk }) => {
          await sess.run(async () => {
            await releasePromise;
          });
          return {
            artifacts: [],
            message: { role: 'model', content: [{ text: 'hi' }] },
          };
        }
      );

      const session = flow.streamBidi({});
      session.send({
        message: { role: 'user' as const, content: [{ text: 'hi' }] },
        detach: true,
      });

      const output = await session.output;
      const snapshotId = output.snapshotId;
      assert.ok(snapshotId);

      const snapPending = await store.getSnapshot({ snapshotId: snapshotId! });
      assert.strictEqual(snapPending?.status, 'pending');

      resolvePromise();
      session.close();

      const snapDone = await waitForSnapshotStatus(
        store,
        snapshotId!,
        'completed'
      );
      assert.strictEqual(snapDone.status, 'completed');
    });

    it('should abort a detached agent', async () => {
      const store = new InMemorySessionStore<{ foo: string }>();
      let aborted = false;

      const flow = defineCustomAgent<{ foo: string }>(
        new Registry(),
        {
          name: 'abortTest',
          store,
        },
        async (sess, { abortSignal }) => {
          if (abortSignal) {
            abortSignal.onabort = () => {
              aborted = true;
            };
          }
          await sess.run(async () => {
            await new Promise((resolve) => setTimeout(resolve, 5000));
          });
          return {
            artifacts: [],
            message: { role: 'model', content: [{ text: 'hi' }] },
          };
        }
      );

      const session = flow.streamBidi({});
      session.send({
        message: { role: 'user' as const, content: [{ text: 'hi' }] },
        detach: true,
      });

      const output = await session.output;
      const snapshotId = output.snapshotId;
      assert.ok(snapshotId);

      const previousStatus = await flow.abort(snapshotId!);

      assert.strictEqual(previousStatus, 'pending');
      const snapAborted = await store.getSnapshot({ snapshotId: snapshotId! });
      assert.strictEqual(snapAborted?.status, 'aborted');
      // AbortController.abort() fires onabort synchronously, so no delay needed.
      assert.strictEqual(aborted, true);
    });

    it('should stamp a heartbeat on the pending detached snapshot', async () => {
      const store = new InMemorySessionStore<{ foo: string }>();
      let resolvePromise: () => void = () => {};
      const releasePromise = new Promise<void>((resolve) => {
        resolvePromise = resolve;
      });

      const flow = defineCustomAgent<{ foo: string }>(
        new Registry(),
        {
          name: 'heartbeatStampTest',
          store,
        },
        async (sess) => {
          await sess.run(async () => {
            await releasePromise;
          });
          return {
            artifacts: [],
            message: { role: 'model', content: [{ text: 'hi' }] },
          };
        }
      );

      const session = flow.streamBidi({});
      session.send({
        message: { role: 'user' as const, content: [{ text: 'hi' }] },
        detach: true,
      });

      const output = await session.output;
      const snapshotId = output.snapshotId!;
      assert.ok(snapshotId);

      const snapPending = await store.getSnapshot({ snapshotId });
      assert.strictEqual(snapPending?.status, 'pending');
      assert.ok(
        snapPending?.heartbeatAt,
        'pending detached snapshot should carry a heartbeatAt'
      );

      resolvePromise();
      session.close();
      await waitForSnapshotStatus(store, snapshotId, 'completed');
    });

    it('reports a pending snapshot with a stale heartbeat as expired', async () => {
      const store = new InMemorySessionStore<{ foo: string }>();

      const flow = defineCustomAgent<{ foo: string }>(
        new Registry(),
        {
          name: 'heartbeatExpiredTest',
          store,
        },
        async (sess) => {
          await sess.run(async () => {});
          return { artifacts: [] };
        }
      );

      // Simulate an orphaned detached snapshot: pending, with a heartbeat that
      // is older than the expiry timeout (default 60s).
      const stale = new Date(Date.now() - 120_000).toISOString();
      const snapshotId = await store.saveSnapshot(undefined, () => ({
        createdAt: stale,
        updatedAt: stale,
        heartbeatAt: stale,
        status: 'pending',
        state: { sessionId: 'sess-expired', custom: { foo: 'bar' } },
      }));
      assert.ok(snapshotId);

      // The raw store still has it as pending (compute-on-read does not write
      // back), but getSnapshotData surfaces it as expired.
      const raw = await store.getSnapshot({ snapshotId: snapshotId! });
      assert.strictEqual(raw?.status, 'pending');

      const viaData = await flow.getSnapshotData({ snapshotId: snapshotId! });
      assert.strictEqual(viaData?.status, 'expired');
    });

    it('keeps a pending snapshot with a fresh heartbeat as pending', async () => {
      const store = new InMemorySessionStore<{ foo: string }>();

      const flow = defineCustomAgent<{ foo: string }>(
        new Registry(),
        {
          name: 'heartbeatFreshTest',
          store,
        },
        async (sess) => {
          await sess.run(async () => {});
          return { artifacts: [] };
        }
      );

      const now = new Date().toISOString();
      const snapshotId = await store.saveSnapshot(undefined, () => ({
        createdAt: now,
        updatedAt: now,
        heartbeatAt: now,
        status: 'pending',
        state: { sessionId: 'sess-fresh', custom: { foo: 'bar' } },
      }));
      assert.ok(snapshotId);

      const viaData = await flow.getSnapshotData({ snapshotId: snapshotId! });
      assert.strictEqual(viaData?.status, 'pending');
    });

    it('does not expire a pending snapshot that has no heartbeat yet', async () => {
      const store = new InMemorySessionStore<{ foo: string }>();

      const flow = defineCustomAgent<{ foo: string }>(
        new Registry(),
        {
          name: 'heartbeatNoneTest',
          store,
        },
        async (sess) => {
          await sess.run(async () => {});
          return { artifacts: [] };
        }
      );

      const old = new Date(Date.now() - 120_000).toISOString();
      const snapshotId = await store.saveSnapshot(undefined, () => ({
        createdAt: old,
        updatedAt: old,
        status: 'pending',
        state: { sessionId: 'sess-noheartbeat', custom: { foo: 'bar' } },
      }));
      assert.ok(snapshotId);

      const viaData = await flow.getSnapshotData({ snapshotId: snapshotId! });
      assert.strictEqual(viaData?.status, 'pending');
    });

    it('should not override terminal status when aborting an already-completed flow', async () => {
      const store = new InMemorySessionStore<{ foo: string }>();

      const flow = defineCustomAgent<{ foo: string }>(
        new Registry(),
        {
          name: 'abortDoneTest',
          store,
        },
        async (sess) => {
          await sess.run(async () => {});
          return {
            artifacts: [],
            message: { role: 'model', content: [{ text: 'hi' }] },
          };
        }
      );

      const session = flow.streamBidi({});
      session.send({
        message: { role: 'user' as const, content: [{ text: 'hi' }] },
      });
      session.close();
      const output = await session.output;
      assert.ok(output.snapshotId);

      // Snapshot should be 'done' now
      const snapBefore = await store.getSnapshot({
        snapshotId: output.snapshotId!,
      });
      assert.strictEqual(snapBefore?.status, 'completed');

      // Abort returns the previous status but does not override terminal states
      const previousStatus = await flow.abort(output.snapshotId!);
      assert.strictEqual(previousStatus, 'completed');

      // Snapshot should still be 'done' - the mutator skips terminal states
      const snapAfter = await store.getSnapshot({
        snapshotId: output.snapshotId!,
      });
      assert.strictEqual(snapAfter?.status, 'completed');
    });

    it('should return undefined when aborting a non-existent snapshot', async () => {
      const store = new InMemorySessionStore<{ foo: string }>();

      const flow = defineCustomAgent<{ foo: string }>(
        new Registry(),
        {
          name: 'abortMissingTest',
          store,
        },
        async (sess) => {
          await sess.run(async () => {});
          return {
            artifacts: [],
            message: { role: 'model', content: [{ text: 'hi' }] },
          };
        }
      );

      const previousStatus = await flow.abort('non-existent-id');
      assert.strictEqual(previousStatus, undefined);
    });

    it('should throw error when detach is requested without session store', async () => {
      const flow = defineCustomAgent<{ foo: string }>(
        new Registry(),
        {
          name: 'noStoreTest',
        },
        async (sess) => {
          await sess.run(async () => {});
          return {
            artifacts: [],
            message: { role: 'model', content: [{ text: 'hi' }] },
          };
        }
      );

      const session = flow.streamBidi({});
      session.send({
        message: { role: 'user' as const, content: [{ text: 'hi' }] },
        detach: true,
      });

      try {
        await session.output;
        assert.fail('Should have thrown error');
      } catch (e: any) {
        assert.strictEqual(
          e.message,
          'FAILED_PRECONDITION: Detach is only supported when a session store is provided.'
        );
      }
    });

    it('should save failed snapshot if detached flow throws', async () => {
      const store = new InMemorySessionStore<{ foo: string }>();
      let resolvePromise: () => void = () => {};
      const releasePromise = new Promise<void>((resolve) => {
        resolvePromise = resolve;
      });

      const flow = defineCustomAgent<{ foo: string }>(
        new Registry(),
        {
          name: 'detachErrorTest',
          store,
        },
        async (sess, { sendChunk }) => {
          await sess.run(async () => {
            await releasePromise;
            throw new Error('intentional background failure');
          });
          return {
            artifacts: [],
            message: { role: 'model', content: [{ text: 'hi' }] },
          };
        }
      );

      const session = flow.streamBidi({});
      session.send({
        message: { role: 'user' as const, content: [{ text: 'hi' }] },
        detach: true,
      });

      const output = await session.output;
      const snapshotId = output.snapshotId;
      assert.ok(snapshotId);

      resolvePromise();
      session.close();

      const snapFailed = await waitForSnapshotStatus(
        store,
        snapshotId!,
        'failed'
      );
      assert.strictEqual(snapFailed.status, 'failed');
      assert.strictEqual(
        snapFailed.error?.message,
        'intentional background failure'
      );
    });

    it('should mark snapshot aborted even without subscription support', async () => {
      const baseStore = new InMemorySessionStore();
      const store = Object.assign(Object.create(baseStore), {
        onSnapshotStateChange: undefined,
        getSnapshot: baseStore.getSnapshot.bind(baseStore),
        saveSnapshot: baseStore.saveSnapshot.bind(baseStore),
      }) as InMemorySessionStore<any>;

      let resolveBlock: () => void = () => {};
      const blockPromise = new Promise<void>((resolve) => {
        resolveBlock = resolve;
      });

      const flow = defineCustomAgent<{ foo: string }>(
        new Registry(),
        {
          name: 'legacyStoreTest',
          store,
        },
        async (sess, { sendChunk }) => {
          await sess.run(async () => {
            await blockPromise; // Keep flow pending until abort is called
          });
          return {
            artifacts: [],
            message: { role: 'model', content: [{ text: 'hi' }] },
          };
        }
      );

      const session = flow.streamBidi({});
      session.send({
        message: { role: 'user' as const, content: [{ text: 'hi' }] },
        detach: true,
      });

      const output = await session.output;
      const snapshotId = output.snapshotId;

      // Snapshot should be 'pending' since the flow is still blocked
      const snapBefore = await store.getSnapshot({ snapshotId: snapshotId! });
      assert.strictEqual(snapBefore?.status, 'pending');

      await flow.abort(snapshotId!);

      const snapshot = await store.getSnapshot({ snapshotId: snapshotId! });
      assert.strictEqual(snapshot?.status, 'aborted');

      // Release the flow so it doesn't hang
      resolveBlock();
      session.close();
    });

    it('should fetch snapshot data via companion action', async () => {
      const store = new InMemorySessionStore<{ foo: string }>();
      const flow = defineCustomAgent<{ foo: string }>(
        new Registry(),
        {
          name: 'companionActionFlow',
          store,
        },
        async (sess) => {
          // Mutate session state so a snapshot is persisted on turn end.
          await sess.run(async () => {});
          return {
            artifacts: [],
            message: { role: 'model', content: [{ text: 'hi' }] },
          };
        }
      );

      const session = flow.streamBidi({});
      session.send({
        message: { role: 'user' as const, content: [{ text: 'hi' }] },
      });
      session.close();
      const output = await session.output;

      assert.ok(output.snapshotId, 'should have a snapshotId');
      const snapData = await flow.getSnapshotData({
        snapshotId: output.snapshotId!,
      });
      assert.strictEqual(snapData?.snapshotId, output.snapshotId);
    });

    it('should chain parentId properly across session snapshots', async () => {
      const store = new InMemorySessionStore<{ foo: string }>();
      const flow = defineCustomAgent<{ foo: string }>(
        new Registry(),
        {
          name: 'lineageTest',
          store,
        },
        async (sess) => {
          await sess.run(async () => {});
          return {
            artifacts: [],
            message: { role: 'model', content: [{ text: 'hi' }] },
          };
        }
      );

      const session1 = flow.streamBidi({});
      session1.send({
        message: { role: 'user' as const, content: [{ text: 'first' }] },