Custom orchestration

Most Genkit agents should use the standard prompt-backed loop. Custom orchestration is for cases where the application must control turn processing directly while still using the Agents API for sessions, snapshots, streaming, HTTP transport, and background execution. If you need complete ownership of the backend contract instead, a Genkit flow with direct generate() calls may be a better fit.

When to use a custom agent

Use ai.defineCustomAgent() when the workflow needs one of these behaviors:

Several model calls in one user turn.
Dynamic model selection or custom stopping rules.
Planner and executor loops with application decisions between calls.
Manual message history management.
Custom state and artifact updates before the final response.
Custom stream chunks that do not come directly from a single model call.

Runtime contract

A custom agent receives a SessionRunner and an options object:

async (sess, { sendChunk, abortSignal, context }) => {
  // Custom loop.
};

Use sess.run(async (input, turnContext) => {}) to process each input turn. The runner adds input.message to history before the callback runs. For server-managed agents, turnContext.snapshotId is reserved before the turn starts and reused when the turn snapshot is saved.

The runner exposes helpers for the parts of session state you are most likely to need. Use getState() for the full state, getMessages() and addMessages() for conversation history, getCustom() and updateCustom(fn) for typed application state, and getArtifacts() and addArtifacts(artifacts) for generated outputs.

Use updateCustom(fn) for progress and control data that the UI should react to during a turn. Use addArtifacts() when the agent produces a named output the user may inspect later, such as a report, patch, or JSON result. Custom state should stay compact because it is part of the conversation state that gets snapshotted or returned to the client.

Multi-step example

export const researchAgent = ai.defineCustomAgent(
  {
    name: 'researchAgent',
    description:
      'Plans research, answers subquestions, and synthesizes results.',
    stateSchema: ResearchStateSchema,
    store,
  },
  async (sess, { sendChunk, abortSignal }) => {
    let finalMessage;

    await sess.run(async (input, turnContext) => {
      const userText =
        input.message?.content.find((part) => part.text)?.text ?? '';
      const priorMessages = sess.getMessages();

      sess.updateCustom((state) => ({
        ...state,
        status: 'Decomposing question',
        turn: turnContext.turnIndex,
      }));

      const plan = await ai.generate({
        model: liteModel,
        prompt: `Break this question into three subquestions:\n${userText}`,
        output: { format: 'json', schema: z.array(z.string()).length(3) },
        abortSignal,
      });

      const subQuestions = plan.output ?? [userText];

      sess.updateCustom((state) => ({
        ...state,
        subQuestions,
        status: 'Researching',
      }));

      const answers = [];
      for (const question of subQuestions) {
        const answer = await ai.generate({
          prompt: `Answer in two paragraphs:\n${question}`,
          abortSignal,
        });
        answers.push({ question, answer: answer.text });
      }

      sess.updateCustom((state) => ({
        ...state,
        answers,
        status: 'Synthesizing',
      }));

      const stream = ai.generateStream({
        messages: priorMessages,
        prompt: `Synthesize these findings:\n${JSON.stringify(answers)}`,
        abortSignal,
      });

      for await (const chunk of stream.stream) {
        sendChunk({ modelChunk: chunk });
      }

      const response = await stream.response;
      finalMessage = response.message;

      if (response.message) {
        sess.addMessages([response.message]);
      }

      sess.addArtifacts([
        {
          name: `research-${turnContext.snapshotId}.json`,
          parts: [{ text: JSON.stringify(answers) }],
        },
      ]);

      return { finishReason: response.finishReason };
    });

    return {
      message: finalMessage,
      artifacts: sess.getArtifacts(),
      finishReason: sess.lastTurnFinishReason,
    };
  },
);

Use input.message for the current user message. The custom handler is not passed an input.messages array. Read history from sess.getMessages().

Failure and recovery

If the per-turn callback throws, the runtime ends the turn, emits its turn end, and resolves the invocation with finishReason: 'failed', or 'aborted' when the caller stopped the run. By default the turn rolls back. For client-managed agents, the response carries the last-good state. For server-managed agents, the response carries the last-good snapshot ID. This lets clients retry without preserving partial failed-turn mutations.

To keep the work a failed turn already did, throw CommittedTurnError from genkit/beta instead. The runner commits the session as the callback left it and reports it as the resume point. With a session store, it saves that state as a failed snapshot (aborted when the caller stopped the run) with the error on it; without one, the response carries that state. Throw it only when the session is at a point the turn can continue from.

A client retries a committed turn by sending an input with no message and no resume; see Resume after a failure or a stop. The runner passes that input to the callback as is, so handle it by running the turn again on the conversation already in the session:

import { CommittedTurnError } from 'genkit/beta';

await sess.run(async (input) => {
  // No message and no resume: a retry of a committed turn.
  const retry = !input.message && !input.resume;
  try {
    await runSteps(sess, { retry, abortSignal });
  } catch (e) {
    // Commit the steps that finished so the client can retry from here.
    throw new CommittedTurnError(e);
  }
});

Streaming custom data

sendChunk({ modelChunk }) forwards model chunks. sess.updateCustom() emits a customPatch chunk. sess.addArtifacts() records artifacts, while sendChunk({ artifact }) can stream an artifact chunk explicitly when the UI needs immediate visibility.

When to use a custom agent

Use genkitx.DefineCustomAgent when the prompt-backed loop does not fit:

The agent must call multiple models in one turn.
The workflow chooses models or tools dynamically.
You need custom retry, planning, or validation around each turn.
The agent emits artifacts or state updates outside a normal model stream.
You need the reserved turn snapshot ID before work starts.

Runtime contract

The custom function receives:

ctx is the invocation context, which is not the transport connection. Before a detach, a client disconnect cancels it. After a detach it stays live and the turn keeps running; only Agent.Abort or the process exiting cancels it. Every genkit.GenerateStream call made with this context follows the same rule.
resp aix.Responder streams model chunks and artifacts to the client.
sess *aix.SessionRunner[State] manages turns, messages, custom state, artifacts, and snapshots.

SessionRunner methods

SessionRunner[State] adds two methods of its own and embeds *aix.Session[State] for everything else:

Method	Purpose
`Run(ctx, fn) error`	Loops over the invocation’s inputs, calling `fn` once per turn.
`Result() *AgentResult`	The last message and artifacts currently recorded.
`State() *SessionState[State]`	The whole session state.
`SessionID() string`	The session this invocation belongs to.
`Messages() []*ai.Message`	Conversation history.
`AddMessages(...*ai.Message)`	Appends messages.
`SetMessages([]*ai.Message)`	Replaces history wholesale.
`UpdateMessages(func([]ai.Message) []ai.Message)`	Atomic read-modify-write on history.
`Custom() State`	Typed custom state.
`UpdateCustom(func(State) State)`	Atomic update; emits a custom patch chunk.
`Artifacts() []*Artifact`	Artifacts recorded so far.
`AddArtifacts(...*Artifact)`	Appends artifacts.
`UpdateArtifacts(func([]Artifact) []Artifact)`	Atomic read-modify-write on artifacts.

The three Update* callbacks run while the session lock is held. Do not call another Session method or send on a Responder from inside one.

Turn loop semantics

sess.Run(ctx, fn) loops over the invocation’s input channel, calling fn once per turn, and returns only when the invocation ends or fn returns an error. It is not a single turn, so put per-turn timeouts and retries inside fn, not around Run. Each turn runs in its own trace span. The runner adds the user message to the session before fn, then emits a TurnEnd chunk and writes a snapshot when a store exists.

fn returns (*aix.TurnResult, error). TurnResult has one field, FinishReason. Returning nil reports no finish reason and the framework infers nothing.

When fn returns a bare error, Run rolls the turn back: it emits TurnEnd with aix.AgentFinishReasonFailed, stops looping, and returns the error. It still writes a failed snapshot when a store exists, holding the messages through the last committed turn, the error, and the session’s usage, since the turn’s tokens were spent. That snapshot is the resume point, and it holds the same conversation as the turn before it. Return a non-nil TurnResult together with the error to commit the turn instead. The session as fn left it then persists as the failed snapshot; do that only when the messages you added end at a turn seam, with every tool request answered, since that is what the next model call needs. A turn that ends because the invocation’s context was cancelled lands as aborted on the same terms. Either way, return the error to resolve the invocation, or call Run again to keep serving inputs on the same connection.

A context that ends between turns ends the loop the way a closed input channel does: Run returns nil before the next queued input starts, and the inputs that never started are dropped.

aix.AgentResult, the agent function’s return value, has three fields: Message, Artifacts, and FinishReason. Leave FinishReason empty to accept the last turn’s reason. Returning an AgentResult completes the invocation on the function’s own terms, even after a failed turn. Returning nil, nil after a failed turn does not: the invocation then takes the failed turn’s outcome and keeps its error.

Usage follows the same split. The model calls made inside fn, including those of a tool that calls genkit.Generate itself, count toward the turn. A call the agent function makes between turns counts toward the invocation and the session, but not toward any turn. See Token usage.

Custom agent example

import (
    "context"
    "fmt"

    "github.com/firebase/genkit/go/ai"
    aix "github.com/firebase/genkit/go/ai/exp"
    "github.com/firebase/genkit/go/genkit"
    genkitx "github.com/firebase/genkit/go/genkit/exp"
)

coder := genkitx.DefineCustomAgent(g, "coder",
    func(ctx context.Context, resp aix.Responder, sess *aix.SessionRunner[CoderState]) (*aix.AgentResult, error) {
        err := sess.Run(ctx, func(ctx context.Context, input *aix.AgentInput) (*aix.TurnResult, error) {
            turn := aix.TurnContextFromContext(ctx)

            sess.UpdateCustom(func(state CoderState) CoderState {
                state.Status = "Generating answer"
                state.LastSnapshotID = turn.SnapshotID
                return state
            })

            for chunk, err := range genkit.GenerateStream(ctx, g,
                ai.WithModelName("googleai/gemini-flash-latest"),
                ai.WithSystem("You are a concise coding assistant."),
                ai.WithMessages(sess.Messages()...),
            ) {
                if err != nil {
                    // A bare error rolls the turn back. Return a TurnResult with
                    // it to commit the messages added so far as a failed snapshot.
                    return nil, fmt.Errorf("stream model: %w", err)
                }
                if chunk.Done {
                    sess.AddMessages(chunk.Response.Message)
                    return &aix.TurnResult{
                        FinishReason: aix.AgentFinishReason(chunk.Response.FinishReason),
                    }, nil
                }
                resp.SendModelChunk(chunk.Chunk)
            }

            // No TurnResult: report no finish reason. The framework infers
            // nothing from a nil result.
            return nil, nil
        })
        if err != nil {
            // sess.Run stopped looping on the first callback error. Returning it
            // resolves the invocation as failed with the last-good state; calling
            // sess.Run again would keep serving inputs instead.
            return nil, fmt.Errorf("run turn: %w", err)
        }
        return sess.Result(), nil
    },
    aix.WithSessionStore(store),
    aix.WithDescription[CoderState]("Concise coding assistant"),
)

sess.Result() is a convenience that returns the last message and artifacts currently recorded in the session.

The coder agent in the basic-agents sample is this shape running alongside five agents built the ordinary way, so the extra work a custom agent takes on is easy to compare.

Turn context

aix.TurnContextFromContext(ctx) returns read-only turn metadata:

SnapshotID is the snapshot ID reserved before this turn runs. It is empty for client-managed agents.
ParentSnapshotID is the snapshot this turn continues from.
TurnIndex is the zero-based turn number within the invocation.

Use this when external resources need to line up with the snapshot that will later be saved.

Responder behavior

Responder.SendModelChunk(chunk) streams token-level model output. Responder.SendArtifact(artifact) streams an artifact and records it in the session. Send methods return promptly when the work context is canceled. Their session side effects are applied before they return, so snapshots and sess.Result() observe them.

Failure and detach behavior

If the per-turn callback returns an error, the invocation resolves as a failed AgentOutput with structured error details and a resume point: the failed turn’s snapshot, holding what the callback committed, or the messages through the last committed turn when it committed nothing. Every chunk the turn produced, its TurnEnd included, reaches the client before the output resolves. When a client detaches, chunks after detach are not forwarded, but session side effects such as artifacts still apply to the final snapshot.

Next steps

Sessions and state covers custom state, artifacts, and state transforms.
Background execution covers detach, pending snapshots, and abort.
Agent error handling covers the failure channels a custom agent resolves into.

When to use a custom agent

Use defineCustomAgent() when you need full control over the turn logic:

Running multiple model calls sequentially inside a single user turn.
Choosing models, prompts, or tools dynamically at runtime.
Implementing custom loops (e.g. planner-executor or self-correction).
Emitting status updates or artifacts during a turn.

Runtime contract

A custom agent receives a SessionRunner and an AgentFnOptions object:

import 'package:genkit/genkit.dart';
import 'package:genkit/experimental.dart';

final customAgent = ai.defineCustomAgent(
  name: 'myAgent',
  fn: (sess, options) async {
    // Custom logic...
  },
);

sess (the SessionRunner) manages the session’s active turns, messages, custom state, and artifacts.
options (the AgentFnOptions) provides sendChunk(chunk) to stream model chunks back to the client, along with a cancel token and ambient request context.

Call sess.run((input, turnContext) async => TurnResult) to process each input turn. The runner appends the user input to the message history automatically. For server-managed agents, turnContext.snapshotId is reserved beforehand so external state matches the snapshot.

Multi-step example

This simplified multi-step researcher is modeled directly on the research_agent.dart sample:

import 'package:genkit/genkit.dart';
import 'package:genkit/experimental.dart';
import 'package:schemantic/schemantic.dart';

final researchAgent = ai.defineCustomAgent(
  name: 'researchAgent',
  // A loose JSON map is a good fit for ad-hoc status/progress state.
  stateSchema: SchemanticType.map(
    SchemanticType.string(),
    SchemanticType.dynamicSchema(),
  ),
  fn: (sess, options) async {
    Message? lastMessage;

    await sess.run((input, turnContext) async {
      final userText = input.message?.content.firstOrNull?.text ?? '';

      // Step 1: Decompose the question into subquestions.
      // The updater receives the typed state (Map<String, dynamic>? here).
      sess.updateCustom((state) {
        final s = state ?? <String, dynamic>{};
        s['status'] = 'Decomposing question...';
        return s;
      });

      final decompose = await ai.generate(
        model: liteModel,
        prompt: 'Break this question into two sub-questions:\n$userText',
        outputFormat: 'json',
        outputSchema: SchemanticType.list(SchemanticType.string()),
      );
      final subQuestions = (decompose.output ?? [userText])
          .map((q) => q.toString())
          .toList();

      // Step 2: Research each subquestion
      final answers = [];
      for (final q in subQuestions) {
        sess.updateCustom((state) {
          final s = state ?? <String, dynamic>{};
          s['status'] = 'Researching: $q';
          return s;
        });

        final research = await ai.generate(prompt: q);
        answers.add({'question': q, 'answer': research.text});
      }

      // Step 3: Synthesize and stream the final response
      sess.updateCustom((state) {
        final s = state ?? <String, dynamic>{};
        s['status'] = 'Synthesizing final response...';
        return s;
      });

      final synthesis = ai.generateStream(
        prompt: 'Synthesize these findings:\n$answers',
      );

      await for (final chunk in synthesis) {
        options.sendChunk(AgentStreamChunk(modelChunk: chunk.modelChunk));
      }

      final finalRes = await synthesis.onResult;
      lastMessage = finalRes.message;
      if (lastMessage != null) {
        sess.addMessages([lastMessage!]);
      }

      return null;
    });

    return AgentResult(
      message:
          lastMessage ??
          Message(
            role: Role.model,
            content: [TextPart(text: 'Research complete.')],
          ),
    );
  },
  store: InMemorySessionStore(),
);

Failure and recovery

If the per-turn callback throws an exception, the runtime marks the turn as failed and resolves the action with finishReason: 'failed'. The response carries the last-good state or snapshot ID, allowing the client to safely retry without continuing from broken or corrupted partial states.

When to use a custom agent

Use ai.define_custom_agent() on genkit.exp.Genkit when the prompt-backed loop does not fit:

Several model calls in one user turn, such as a planner followed by an answerer.
Models, prompts, or tools chosen per turn by application code.
Status updates or artifacts the UI should see between steps.

Runtime contract

A custom agent function takes a SessionRunner and an AgentContext, and returns an AgentResult. All three come from genkit.exp.agent:

from genkit.exp import Genkit
from genkit.exp.agent import AgentContext, AgentInput, AgentResult, SessionRunner, TurnContext, TurnResult
from genkit_google_genai import GoogleAI

ai = Genkit(plugins=[GoogleAI()])


async def triage_fn(sess: SessionRunner, ctx: AgentContext) -> AgentResult:
    async def handle_turn(inp: AgentInput, turn_ctx: TurnContext) -> TurnResult | None:
        ...  # One user turn: call models, update state, stream chunks.

    await sess.run(handle_turn)
    return await sess.result()


triage_agent = ai.define_custom_agent(name='triageAgent', fn=triage_fn)

sess runs the turn loop (run, result) and exposes history (get_messages, add_messages, set_messages), typed custom state (get_custom, update_custom), and artifacts (get_artifacts, add_artifacts).
ctx streams chunks with ctx.send_chunk(AgentStreamChunk(...)) and carries request metadata in ctx.context. There is no abort signal; cancellation arrives as asyncio.CancelledError at the next await.
turn_ctx carries turn_index, parent_snapshot_id, and snapshot_id. With a store, the runner reserves snapshot_id before handle_turn runs and saves the turn’s snapshot under the same ID, so external resources can key on it.

await sess.run(handle_turn) calls handle_turn once per input and returns when the inputs end or a turn fails. The runner appends inp.message to history before each call, so read the conversation from sess.get_messages(). sess.result() returns the last message, the artifacts, and the last turn’s finish reason.

sess.update_custom(...) streams a custom_patch chunk and sess.add_artifacts(...) streams an artifact chunk, so the client sees both before the turn ends.

Multi-step example

This research agent plans with one structured call, streams the answer from a second call, and reports progress through custom state:

from pydantic import BaseModel

from genkit import FinishReason, GenkitError, Part
from genkit.exp import Genkit
from genkit.exp.agent import (
    AgentContext,
    AgentFinishReason,
    AgentInput,
    AgentResult,
    AgentStreamChunk,
    Artifact,
    InMemorySessionStore,
    SessionRunner,
    TurnContext,
    TurnResult,
)
from genkit_google_genai import GoogleAI

ai = Genkit(plugins=[GoogleAI()])
model = GoogleAI.gemini('gemini-flash-latest')


class ResearchPlan(BaseModel):
    subquestions: list[str]


class ResearchState(BaseModel):
    status: str = ''
    subquestions: list[str] = []


async def research_fn(sess: SessionRunner[ResearchState], ctx: AgentContext) -> AgentResult:
    async def handle_turn(inp: AgentInput, turn_ctx: TurnContext) -> TurnResult:
        # 1. The runner already appended the user message to history.
        history = await sess.get_messages()
        question = history[-1].text

        # 2. Plan: one structured call splits the question.
        await sess.update_custom(lambda _: ResearchState(status='Planning'))
        plan = await ai.generate(
            model=model,
            prompt=f'Split this question into three short subquestions: {question}',
            output_schema=ResearchPlan,
        )
        subquestions = plan.output.subquestions if plan.output else [question]
        await sess.update_custom(lambda _: ResearchState(status='Answering', subquestions=subquestions))

        # 3. Answer: stream the final reply to the client.
        stream = ai.generate_stream(
            model=model,
            system='Answer the last question. Cover these subquestions in order: ' + '; '.join(subquestions),
            messages=history,
        )
        async for chunk in stream.stream:
            ctx.send_chunk(AgentStreamChunk(model_chunk=chunk))
        res = await stream.response
        if res.error:
            # generate returns model failures as values. Raise to fail the turn.
            raise GenkitError(status=res.error.status, message=res.error.message)

        # 4. Record the reply, a named artifact, and the final status.
        if res.message:
            await sess.add_messages([res.message])
        await sess.add_artifacts([
            Artifact(name=f'research-{turn_ctx.snapshot_id}.md', parts=[Part.from_text(res.text)]),
        ])
        await sess.update_custom(lambda _: ResearchState(status='Done', subquestions=subquestions))

        stopped = res.finish_reason == FinishReason.STOP
        return TurnResult(finish_reason=AgentFinishReason.STOP if stopped else AgentFinishReason.UNKNOWN)

    await sess.run(handle_turn)
    return await sess.result()


research_agent = ai.define_custom_agent(
    name='researchAgent',
    fn=research_fn,
    state_schema=ResearchState,
    store=InMemorySessionStore(),
)

Call it like any other agent. state_schema types res.state, and the artifact name matches the saved snapshot:

chat = research_agent.chat()
res = await chat.send('Why do sourdough starters need regular feeding?')

print(res.state.status)
# => Done
print(res.artifacts[0].name == f'research-{res.snapshot_id}.md')
# => True

Failure and recovery

If handle_turn raises, the turn ends with finish_reason == AgentFinishReason.FAILED and the error on res.error. chat.send() returns instead of raising, and the chat stays on the last completed turn, so the caller can retry. With a store, the failed turn is saved as a failed snapshot for inspection.

chat.send() raises GenkitError only when no turn result came back, such as a dropped connection or init misuse. See Agent error handling for both channels.

Next steps

Sessions and state covers custom state, artifacts, and state transforms.
Background execution covers detach, pending snapshots, and abort.
Agent error handling covers the failure channels a custom agent resolves into.