Working with Agent State (Beta)

Beta / preview API. Read agents.md first.

Beyond message history, an agent session can hold typed custom state — your own structured data (a task list, a workflow status, counters, etc.). Tools read and mutate it during a turn, and it is automatically synced to the remoteAgent client.

Declare the state shape

Pass a stateSchema (Zod) to defineAgent. State is inferred from it and validated when a snapshot is loaded.

import { z } from 'genkit';
import { ai } from './genkit.js';

const TaskItem = z.object({
  id: z.number(),
  title: z.string(),
  done: z.boolean(),
});

const TaskState = z.object({
  tasks: z.array(TaskItem),
  nextId: z.number(),
});

export const taskAgent = ai.defineAgent({
  name: 'taskAgent',
  stateSchema: TaskState,
  system: 'You manage the user\'s task list. Use the tools to modify it.',
  tools: [
    /* addTask, toggleTask, removeTask — below */
  ],
});

Read & mutate state inside tools

Tools call ai.currentSession<S>() to access the live session, then session.getCustom() / session.updateCustom(mutator). updateCustom takes (custom?: S) => S and returns the new state.

type TaskState = z.infer<typeof TaskState>;

const addTask = ai.defineTool(
  {
    name: 'addTask',
    description: 'Add a new task. Returns the created task.',
    inputSchema: z.object({ title: z.string() }),
    outputSchema: z.object({
      id: z.number(),
      title: z.string(),
      done: z.boolean(),
    }),
  },
  async (input) => {
    const session = ai.currentSession<TaskState>();
    let created!: { id: number; title: string; done: boolean };
    session.updateCustom((state) => {
      const s = state ?? { tasks: [], nextId: 1 };
      created = { id: s.nextId, title: input.title, done: false };
      s.tasks.push(created);
      s.nextId++;
      return s;
    });
    return created;
  }
);

ai.currentSession() throws if called outside an active session (e.g. a tool invoked without a running agent turn), so only use it inside agent tools.

Seed and read state (server-side)

Seed initial custom state when opening a chat. The state argument is a SessionState<S>: custom data goes under .custom (alongside messages and artifacts).

const chat = taskAgent.chat({
  state: {
    custom: { tasks: [], nextId: 1 },
    messages: [],
    artifacts: [],
  },
});

const res = await chat.send('Add a task: buy groceries');
console.log(res.state); // res.state returns the custom state directly

Auto-sync to the remoteAgent client

When you talk to the agent over HTTP, the remoteAgent client tracks custom state for you. Parameterize the client with your state type, seed it the same way (state.custom), and after a turn read it off chat.state.

Important: on the client, custom fields are flattened directly onto chat.state (e.g. chat.state.tasks) — not under chat.state.custom.

import { remoteAgent, type AgentChat } from 'genkit/beta/client';

interface TaskState {
  tasks: { id: number; title: string; done: boolean }[];
  nextId: number;
}

const agent = remoteAgent<TaskState>({ url: '/api/taskAgent' });
const chat: AgentChat<TaskState> = agent.chat({
  state: { custom: { tasks: [], nextId: 1 }, messages: [], artifacts: [] },
});

const turn = chat.sendStream('Add buy groceries, then mark it done');
for await (const chunk of turn.stream) {
  /* render chunk.text */
}
await turn.response;

// Custom state is flattened onto chat.state:
console.log(chat.state?.tasks);

There is no onStateChange subscription — state updates ride on the streamed chunks. Read the authoritative state from chat.state after await turn.response. (For live mid-stream status updates from a custom agent, see advanced custom agents, which emit customPatch chunks as state changes.)