@genkit-ai/ai
Version:
Genkit AI framework generative AI APIs.
944 lines • 32.6 kB
JavaScript
import {
GenkitError,
deepEqual,
defineAction,
defineBidiAction,
getContext,
run,
z
} from "@genkit-ai/core";
import { Channel } from "@genkit-ai/core/async";
import {
createAgentAPI
} from "./agent-core.mjs";
import { parseSchema, toJsonSchema } from "@genkit-ai/core/schema";
import {
setCustomMetadataAttribute,
setCustomMetadataAttributes
} from "@genkit-ai/core/tracing";
import {
AgentAbortRequestSchema,
AgentAbortResponseSchema,
AgentInitSchema,
AgentInputSchema,
AgentOutputSchema,
AgentStreamChunkSchema,
GetSnapshotRequestSchema
} from "./agent-types.mjs";
import { generateStream } from "./generate.mjs";
import { diff } from "./json-patch.mjs";
import {
definePrompt
} from "./prompt.mjs";
import { InMemorySessionStore } from "./session-stores.mjs";
import {
Session,
SessionSnapshotSchema,
reserveSnapshotId,
runWithSession
} from "./session.mjs";
import {
AgentAbortRequestSchema as AgentAbortRequestSchema2,
AgentAbortResponseSchema as AgentAbortResponseSchema2,
AgentInitSchema as AgentInitSchema2,
AgentInputSchema as AgentInputSchema2,
AgentOutputSchema as AgentOutputSchema2,
AgentResultSchema,
AgentStreamChunkSchema as AgentStreamChunkSchema2,
GetSnapshotRequestSchema as GetSnapshotRequestSchema2,
JsonPatchOperationSchema,
JsonPatchSchema,
TurnEndSchema
} from "./agent-types.mjs";
const DEFAULT_HEARTBEAT_INTERVAL_MS = 3e4;
const DEFAULT_HEARTBEAT_TIMEOUT_MS = 6e4;
function isHeartbeatExpired(snapshot, timeoutMs = DEFAULT_HEARTBEAT_TIMEOUT_MS) {
if (snapshot.status !== "pending" || !snapshot.heartbeatAt) {
return false;
}
const last = Date.parse(snapshot.heartbeatAt);
if (Number.isNaN(last)) {
return false;
}
return Date.now() - last > timeoutMs;
}
function toErrorDetails(e) {
return {
status: e?.status || "INTERNAL",
message: e?.message || "Internal failure",
// Only surface explicitly-provided structured details. Never fall back to
// the raw thrown value: it is serialized over the wire (AgentOutput.error)
// and persisted into snapshots, so leaking it could expose stack traces /
// internal state, or break JSON.stringify on circular error objects.
details: e?.detail ?? e?.details
};
}
function abortAwareMutator(input) {
return (current) => current?.status === "aborted" ? null : input;
}
function requireStore(store, operation, agentName) {
if (!store) {
throw new GenkitError({
status: "FAILED_PRECONDITION",
message: `${operation} requires a persistent store. Provide a 'store' when defining '${agentName}'.`
});
}
}
async function abortSnapshotInStore(store, snapshotId, options) {
let previousStatus;
await store.saveSnapshot(
snapshotId,
(current) => {
if (!current) return null;
previousStatus = current.status;
if (current.status === "completed" || current.status === "failed" || current.status === "aborted") {
return null;
}
return { ...current, status: "aborted" };
},
options
);
return previousStatus;
}
class SessionRunner {
session;
inputCh;
turnIndex = 0;
onEndTurn;
onDetach;
newSnapshotId;
/** The finish reason of the most recently completed turn. */
lastTurnFinishReason;
/**
* Error details of the most recent failed turn. Set when a turn throws and
* the runner resolves gracefully instead of propagating the exception.
*/
lastTurnError;
/**
* The state the most recently *successful* turn left behind. On a failed
* turn this is the state the failed turn started with - the last-good state
* returned to the caller (for client-managed agents).
*/
lastGoodState;
/**
* The snapshotId of the most recently *successful* (persisted, `done`) turn.
* On a failed turn this is the last-good snapshot the caller resumes from;
* `undefined` when no turn has succeeded yet (e.g. a first-turn failure).
*/
lastGoodSnapshotId;
lastSnapshot;
lastSnapshotVersion = 0;
store;
isDetached = false;
/**
* Aborts in-flight turns. When set and aborted, a turn that rejects out of
* `generate` is reported as `aborted` (not `failed`) and its failed snapshot
* write is skipped (the abort path already persisted the `aborted` status).
*/
abortSignal;
/**
* True until the first `customPatch` chunk of the current turn has been
* emitted. The first patch of every turn is a whole-document replace
* (re-basing clients that may not share the server's baseline); reset to
* `true` at the start of each turn.
*/
firstCustomPatchInTurn = true;
constructor(session, inputCh, options) {
this.session = session;
this.inputCh = inputCh;
this.lastSnapshot = options?.lastSnapshot;
this.store = options?.store;
this.abortSignal = options?.abortSignal;
this.onEndTurn = options?.onEndTurn;
this.onDetach = options?.onDetach;
this.lastGoodState = this.session.getState();
this.lastGoodSnapshotId = options?.lastSnapshot?.snapshotId;
}
// ── Session delegate methods ────────────────────────────────────────
// These forward to `this.session` so callers can write `sess.addMessages()`
// instead of the verbose `sess.session.addMessages()`.
/** Returns a deep copy of the current session state. */
getState() {
return this.session.getState();
}
/** Retrieves all messages associated with the session. */
getMessages() {
return this.session.getMessages();
}
/** Appends messages to the session. */
addMessages(messages) {
this.session.addMessages(messages);
}
/** Overwrites the session messages. */
setMessages(messages) {
this.session.setMessages(messages);
}
/** Retrieves the custom state of the session. */
getCustom() {
return this.session.getCustom();
}
/** Updates the custom state using a mutator function. */
updateCustom(fn) {
this.session.updateCustom(fn);
}
/** Retrieves the list of artifacts generated during the session. */
getArtifacts() {
return this.session.getArtifacts();
}
/** Adds artifacts to the session, deduplicating by name. */
addArtifacts(artifacts) {
this.session.addArtifacts(artifacts);
}
/** Invokes the end-of-turn callback, absorbing errors from a closed stream. */
notifyEndTurn(snapshotId, finishReason) {
try {
this.onEndTurn?.(snapshotId, finishReason);
} catch {
}
}
/**
* Executes the flow handler against incoming input messages sequentially.
*
* The handler receives the turn's {@link AgentInput} and a {@link TurnContext}
* whose `snapshotId` is *reserved up front* - it is the id the snapshot
* persisted at turn end will reuse. This lets a handler set up external,
* snapshot-correlated state (e.g. a git branch/worktree named after the
* snapshot) before generating, then commit it under that id.
*
* The handler may return a {@link TurnResult} carrying an explicit
* `finishReason` for the just-completed turn. When omitted, no per-turn
* reason is reported. Failures always report `failed`.
*/
async run(fn) {
for await (const input of this.inputCh) {
if (input.message) {
this.session.addMessages([input.message]);
}
this.firstCustomPatchInTurn = true;
const parentSnapshotId = this.lastSnapshot?.snapshotId;
if (this.store && !this.newSnapshotId) {
this.newSnapshotId = reserveSnapshotId();
}
const turnSnapshotId = this.newSnapshotId;
this.newSnapshotId = void 0;
const turnContext = {
snapshotId: turnSnapshotId,
parentSnapshotId,
turnIndex: this.turnIndex
};
try {
await run(`runTurn-${this.turnIndex + 1}`, input, async () => {
const turnResult = await fn(input, turnContext);
const finishReason = turnResult?.finishReason;
this.lastTurnFinishReason = finishReason;
this.lastTurnError = void 0;
const snapshotId = await this.maybeSnapshot(
"completed",
void 0,
turnSnapshotId,
finishReason
);
this.lastGoodState = this.session.getState();
this.lastGoodSnapshotId = snapshotId;
if (snapshotId) {
setCustomMetadataAttribute("agent:snapshotId", snapshotId);
}
this.notifyEndTurn(snapshotId, finishReason);
return { state: this.session.getState() };
});
this.turnIndex++;
} catch (e) {
if (this.abortSignal?.aborted) {
this.lastTurnFinishReason = "aborted";
this.lastTurnError = void 0;
this.notifyEndTurn(this.lastSnapshot?.snapshotId, "aborted");
break;
}
this.lastTurnFinishReason = "failed";
this.lastTurnError = toErrorDetails(e);
const snapshotId = await this.maybeSnapshot(
"failed",
this.lastTurnError,
turnSnapshotId,
"failed"
);
this.notifyEndTurn(snapshotId, "failed");
break;
}
}
}
/**
* Saves a snapshot of the current session state to the persistent store.
*
* When a store is configured every turn is persisted (snapshotting is no
* longer opt-out). Uses the mutator-based `saveSnapshot` to atomically check
* that the snapshot has not been concurrently aborted before writing -
* preventing a race where a "done" write could overwrite a concurrent
* "aborted" status.
*/
async maybeSnapshot(status, error, snapshotId, finishReason) {
if (!this.store || this.isDetached && snapshotId !== this.lastSnapshot?.snapshotId)
return this.lastSnapshot?.snapshotId;
const currentVersion = this.session.getVersion();
if (currentVersion === this.lastSnapshotVersion && !status) {
return this.lastSnapshot?.snapshotId;
}
const currentState = this.session.getState();
const snapshotInput = {
...snapshotId || this.newSnapshotId ? { snapshotId: snapshotId || this.newSnapshotId } : {},
// Stamp the session id onto every snapshot in the chain so callers can
// resolve a snapshot's session without reaching into its state.
sessionId: this.session.sessionId,
createdAt: (/* @__PURE__ */ new Date()).toISOString(),
updatedAt: (/* @__PURE__ */ new Date()).toISOString(),
state: currentState,
parentId: this.lastSnapshot?.snapshotId,
// Default to a resumable `completed` status. The only caller that omits a
// status is the post-invocation write (which fires when the handler
// mutates state after the last turn); persisting it as `completed` keeps
// it a valid resume target under the "only `completed` is resumable" rule.
status: status ?? "completed",
// Stamp an initial heartbeat on a `pending` (detached, in-flight)
// snapshot. A background heartbeat loop refreshes it; if it goes stale the
// snapshot is reported as `expired` on read (the worker is presumed dead).
...status === "pending" && { heartbeatAt: (/* @__PURE__ */ new Date()).toISOString() },
...finishReason && { finishReason },
error
};
const effectiveId = snapshotId || this.newSnapshotId;
const assignedId = await this.store.saveSnapshot(
effectiveId,
abortAwareMutator(snapshotInput),
{ context: getContext() }
);
if (assignedId === null) {
return effectiveId;
}
this.lastSnapshot = { ...snapshotInput, snapshotId: assignedId };
this.lastSnapshotVersion = currentVersion;
return assignedId;
}
}
const GetSnapshotDataInputSchema = z.object({
snapshotId: z.string().optional(),
sessionId: z.string().optional()
});
class AgentInitError extends GenkitError {
}
function assertInitMatchesStateManagement(config, init) {
if ((init?.snapshotId || init?.sessionId) && !config.store) {
throw new AgentInitError({
status: "FAILED_PRECONDITION",
message: `Cannot use '${init.snapshotId ? "snapshotId" : "sessionId"}' with agent '${config.name}': this agent has no store configured (client-managed state). Send 'state' instead.`
});
}
if (init?.state && config.store) {
throw new AgentInitError({
status: "FAILED_PRECONDITION",
message: `Cannot send 'state' to agent '${config.name}': this agent uses a server-managed store. Send 'snapshotId' or 'sessionId' instead.`
});
}
}
async function resolveSession(config, store, init, validateCustomState) {
if (init?.snapshotId) {
const snapshot = await store.getSnapshot({
snapshotId: init.snapshotId,
context: getContext()
});
if (!snapshot) {
throw new GenkitError({
status: "NOT_FOUND",
message: `Snapshot ${init.snapshotId} not found`
});
}
const snapshotSessionId = snapshot.sessionId ?? snapshot.state?.sessionId;
if (init.sessionId && snapshotSessionId !== init.sessionId) {
throw new AgentInitError({
status: "INVALID_ARGUMENT",
message: `Snapshot ${init.snapshotId} does not belong to session ${init.sessionId} (it belongs to ${snapshotSessionId ?? "an unknown session"}).`
});
}
if (snapshot.status !== "completed") {
throw new GenkitError({
status: "INVALID_ARGUMENT",
message: `Snapshot ${init.snapshotId} is not resumable (status: ${snapshot.status ?? "unknown"}). Only 'completed' snapshots can be resumed.`
});
}
validateCustomState(snapshot.state?.custom);
return {
snapshot,
session: new Session(snapshot.state)
};
}
if (init?.sessionId) {
let snapshot = await store.getSnapshot({
sessionId: init.sessionId,
context: getContext()
});
const visited = /* @__PURE__ */ new Set();
while (snapshot && snapshot.status !== "completed") {
if (visited.has(snapshot.snapshotId)) {
throw new GenkitError({
status: "FAILED_PRECONDITION",
message: `Session '${init.sessionId}' has a cyclic snapshot parent chain (snapshot '${snapshot.snapshotId}' was visited twice). Resume by snapshotId instead.`
});
}
visited.add(snapshot.snapshotId);
snapshot = snapshot.parentId ? await store.getSnapshot({
snapshotId: snapshot.parentId,
context: getContext()
}) : void 0;
}
if (snapshot) {
validateCustomState(snapshot.state?.custom);
return {
snapshot,
session: new Session(snapshot.state)
};
}
return {
session: new Session({
custom: void 0,
artifacts: [],
messages: [],
sessionId: init.sessionId
})
};
}
if (init?.state && !config.store) {
validateCustomState(init.state.custom);
return {
session: new Session(init.state)
};
}
return {
session: new Session({
custom: void 0,
artifacts: [],
messages: []
})
};
}
function pipeInputWithDetach(inputStream, target, getRunner, storeEnabled, rejectDetach) {
(async () => {
try {
for await (const input of inputStream) {
if (input.detach) {
if (!storeEnabled) {
rejectDetach(
new GenkitError({
status: "FAILED_PRECONDITION",
message: "Detach is only supported when a session store is provided."
})
);
} else {
const runner = getRunner();
const turnSnapshotId = runner.newSnapshotId || reserveSnapshotId();
runner.newSnapshotId = turnSnapshotId;
await runner.maybeSnapshot("pending", void 0, turnSnapshotId);
runner.isDetached = true;
if (runner.onDetach) {
runner.onDetach(turnSnapshotId);
}
}
const hasPayload = !!(input.message || input.resume?.restart?.length || input.resume?.respond?.length);
if (hasPayload) {
target.send(input);
}
} else {
target.send(input);
}
}
target.close();
} catch (e) {
target.error(e);
}
})();
}
function defineCustomAgent(registry, config, fn) {
const toClientState = (state) => {
if (config.clientTransform?.state) {
return config.clientTransform.state(state);
}
return state;
};
const stateJsonSchema = config.stateSchema ? toJsonSchema({ schema: config.stateSchema }) : void 0;
const validateCustomState = (custom) => {
if (config.stateSchema && custom !== void 0) {
parseSchema(custom, { schema: config.stateSchema });
}
};
const primaryAction = defineBidiAction(
registry,
{
name: config.name,
description: config.description,
actionType: "agent",
inputSchema: AgentInputSchema,
outputSchema: AgentOutputSchema,
streamSchema: AgentStreamChunkSchema,
initSchema: AgentInitSchema,
metadata: {
agent: {
stateManagement: config.store ? "server" : "client",
abortable: !!config.store?.onSnapshotStateChange,
...stateJsonSchema && { stateSchema: stateJsonSchema }
}
}
},
async function* (arg) {
const init = arg.init;
const store = config.store || new InMemorySessionStore();
assertInitMatchesStateManagement(config, init);
let session;
let snapshot;
try {
({ session, snapshot } = await resolveSession(
config,
store,
init,
validateCustomState
));
} catch (e) {
if (e instanceof AgentInitError) {
throw e;
}
return {
finishReason: "failed",
error: toErrorDetails(e),
...!config.store && init?.state && { state: init.state }
};
}
setCustomMetadataAttributes({
"agent:sessionId": session.sessionId
});
let detachedSnapshotId;
let resolveDetach;
let rejectDetach;
const detachPromise = new Promise((resolve, reject) => {
resolveDetach = resolve;
rejectDetach = reject;
});
const abortController = new AbortController();
let unsubscribe = void 0;
let heartbeatTimer;
const stopHeartbeat = () => {
if (heartbeatTimer) {
clearInterval(heartbeatTimer);
heartbeatTimer = void 0;
}
};
let runner;
const emitChunk = (chunk) => {
try {
let toSend = chunk;
if (config.clientTransform?.chunk) {
toSend = config.clientTransform.chunk(chunk);
}
if (!toSend) return;
arg.sendChunk(toSend);
} catch {
}
};
const runnerInputChannel = new Channel();
pipeInputWithDetach(
arg.inputStream,
runnerInputChannel,
() => runner,
!!config.store,
(reason) => rejectDetach?.(reason)
);
runner = new SessionRunner(session, runnerInputChannel, {
store,
lastSnapshot: snapshot,
abortSignal: abortController.signal,
onDetach: (snapshotId) => {
detachedSnapshotId = snapshotId;
if (resolveDetach) {
resolveDetach();
}
const ctx = getContext();
heartbeatTimer = setInterval(() => {
void store.saveSnapshot(
snapshotId,
(current) => current?.status === "pending" ? { ...current, heartbeatAt: (/* @__PURE__ */ new Date()).toISOString() } : null,
{ context: ctx }
).catch(() => {
});
}, DEFAULT_HEARTBEAT_INTERVAL_MS);
heartbeatTimer.unref?.();
if (store.onSnapshotStateChange) {
unsubscribe = store.onSnapshotStateChange(
snapshotId,
(snap) => {
if (snap.status === "aborted") {
stopHeartbeat();
abortController.abort();
if (unsubscribe) unsubscribe();
}
},
{ context: getContext() }
);
}
},
onEndTurn: (snapshotId, finishReason2) => {
if (!runner.isDetached) {
emitChunk({
turnEnd: {
...config.store && { snapshotId },
...finishReason2 && { finishReason: finishReason2 }
}
});
}
}
});
const sendArtifactChunk = (a) => {
if (!runner.isDetached) {
emitChunk({ artifact: a });
}
};
session.on("artifactAdded", sendArtifactChunk);
session.on("artifactUpdated", sendArtifactChunk);
let lastSentCustom;
const sendCustomPatch = () => {
if (runner.isDetached) return;
const transformed = toClientState(session.getState())?.custom;
let patch;
if (runner.firstCustomPatchInTurn) {
patch = [
{ op: "replace", path: "", value: structuredClone(transformed) }
];
runner.firstCustomPatchInTurn = false;
} else {
patch = diff(lastSentCustom, transformed);
}
lastSentCustom = structuredClone(transformed);
if (patch.length) {
emitChunk({ customPatch: patch });
}
};
session.on("customChanged", sendCustomPatch);
const sendChunk = (chunk) => {
if (!runner.isDetached) {
emitChunk(chunk);
}
};
const flowPromise = (async () => {
try {
const result2 = await runWithSession(
registry,
session,
() => fn(runner, {
sendChunk,
abortSignal: abortController.signal,
context: getContext()
})
);
const finalSnapshotId2 = await runner.maybeSnapshot();
return { result: result2, finalSnapshotId: finalSnapshotId2 };
} finally {
stopHeartbeat();
if (unsubscribe) unsubscribe();
session.off("artifactAdded", sendArtifactChunk);
session.off("artifactUpdated", sendArtifactChunk);
session.off("customChanged", sendCustomPatch);
}
})();
const outcome = await Promise.race([
flowPromise,
detachPromise.then(() => "detached")
]);
if (outcome === "detached") {
return {
sessionId: session.sessionId,
snapshotId: detachedSnapshotId,
finishReason: "detached",
...!config.store && { state: toClientState(session.getState()) }
};
}
const { result, finalSnapshotId } = outcome;
if (runner.lastTurnFinishReason === "failed" && runner.lastTurnError) {
const lastGood = runner.lastGoodState ?? session.getState();
const lastGoodMessages = lastGood.messages;
return {
sessionId: session.sessionId,
finishReason: "failed",
error: runner.lastTurnError,
...result.artifacts?.length && { artifacts: result.artifacts },
...lastGoodMessages?.length && {
message: lastGoodMessages[lastGoodMessages.length - 1]
},
// Server-managed: the last successful turn is already persisted (every
// turn is snapshotted), so point at its snapshot. The failed turn's
// own snapshot is persisted too but is not resumable - `sessionId`
// resume skips it back to this last-good `done` snapshot. Undefined on
// a first-turn failure (no successful turn yet; client holds the seed).
...config.store && { snapshotId: runner.lastGoodSnapshotId },
// Client-managed: return the last-good state directly.
...!config.store && { state: toClientState(lastGood) }
};
}
const finishReason = result.finishReason ?? runner.lastTurnFinishReason;
return {
sessionId: session.sessionId,
...result.artifacts?.length && { artifacts: result.artifacts },
...result.message && { message: result.message },
...finishReason && { finishReason },
...config.store && { snapshotId: finalSnapshotId },
...!config.store && { state: toClientState(session.getState()) }
};
}
);
const toClientSnapshot = (snapshot) => {
if (!config.clientTransform?.state || !snapshot.state) {
return snapshot;
}
return {
...snapshot,
state: config.clientTransform.state(snapshot.state)
};
};
const resolveSnapshot = async (lookup) => {
requireStore(config.store, "getSnapshotData", config.name);
const snapshot = await config.store.getSnapshot(lookup);
if (!snapshot) return void 0;
const effective = isHeartbeatExpired(snapshot) ? { ...snapshot, status: "expired" } : snapshot;
return toClientSnapshot(effective);
};
const runAbort = (snapshotId, options) => {
requireStore(config.store, "abort", config.name);
return abortSnapshotInStore(config.store, snapshotId, options);
};
const getSnapshotDataAction = defineAction(
registry,
{
name: config.name,
description: `Gets snapshot data for ${config.name} by snapshotId or sessionId`,
actionType: "agent-snapshot",
inputSchema: GetSnapshotRequestSchema,
outputSchema: SessionSnapshotSchema.optional()
},
async (lookup) => resolveSnapshot({ ...lookup, context: getContext() })
);
const abortAgentAction = defineAction(
registry,
{
name: config.name,
description: `Aborts ${config.name} agent by snapshotId. Returns the snapshot id and its status after the abort attempt.`,
actionType: "agent-abort",
inputSchema: AgentAbortRequestSchema,
outputSchema: AgentAbortResponseSchema
},
async ({ snapshotId }) => {
const status = await runAbort(snapshotId, { context: getContext() });
return { snapshotId, status };
}
);
const composite = Object.assign(primaryAction, {
getSnapshotData: (opts) => resolveSnapshot(opts),
abort: (snapshotId, options) => runAbort(snapshotId, options),
getSnapshotDataAction,
abortAgentAction
});
const startBidi = (input, init, opts) => {
const bidi = primaryAction.streamBidi(init, {
abortSignal: opts.abortSignal
});
bidi.send(input);
bidi.close();
return bidi;
};
const transport = {
stateManagement: config.store ? "server" : "client",
runTurn(input, init, opts) {
const bidi = startBidi(input, init, opts);
return { stream: bidi.stream, output: bidi.output };
},
async getSnapshot(lookup) {
return composite.getSnapshotData(lookup);
},
abort(snapshotId) {
return composite.abort(snapshotId);
}
};
const agentApi = createAgentAPI(transport);
Object.assign(composite, {
chat: agentApi.chat,
loadChat: agentApi.loadChat,
getSnapshot: agentApi.getSnapshot
});
return composite;
}
function definePromptAgent(registry, config) {
let cachedPromptAction;
const fn = async (sess, { sendChunk, abortSignal }) => {
await sess.run(async (input) => {
const promptInput = config.promptInput ?? {};
if (!cachedPromptAction) {
cachedPromptAction = await registry.lookupAction(
`/prompt/${config.promptName}`
);
if (!cachedPromptAction) {
throw new Error(
`Prompt '${config.promptName}' not found. Ensure it is defined before the agent is invoked.`
);
}
}
const historyTag = "_genkit_history";
const promptTag = "agentPreamble";
const history = (sess.getMessages() || []).map((m) => ({
...m,
metadata: { ...m.metadata, [historyTag]: true }
}));
const genOpts = await cachedPromptAction.__executablePrompt.render(
promptInput,
{ messages: history }
);
if (genOpts.messages) {
genOpts.messages = genOpts.messages.map((m) => {
if (m.metadata?.[historyTag]) {
const { [historyTag]: _, ...restMeta } = m.metadata;
return {
...m,
metadata: Object.keys(restMeta).length ? restMeta : void 0
};
}
return { ...m, metadata: { ...m.metadata, [promptTag]: true } };
});
}
if (input.resume) {
validateResumeAgainstHistory(input.resume, sess.getMessages());
genOpts.resume = {
...input.resume.restart?.length && {
restart: input.resume.restart
},
...input.resume.respond?.length && {
respond: input.resume.respond
}
};
}
const result = generateStream(registry, { ...genOpts, abortSignal });
for await (const chunk of result.stream) {
sendChunk({ modelChunk: chunk });
}
const res = await result.response;
if (res.request?.messages) {
const msgs2 = res.request.messages.filter(
(m) => !m.metadata?.[promptTag]
);
if (res.message) {
msgs2.push(res.message);
}
sess.setMessages(msgs2);
} else if (res.message) {
sess.addMessages([res.message]);
}
if (res.finishReason === "interrupted") {
const parts = res.message?.content?.filter((p) => !!p.toolRequest) || [];
if (parts.length > 0) {
sendChunk({
modelChunk: {
role: "tool",
content: parts
}
});
}
}
return { finishReason: res.finishReason };
});
const msgs = sess.getMessages();
return {
artifacts: sess.getArtifacts(),
message: msgs.length > 0 ? msgs[msgs.length - 1] : void 0,
...sess.lastTurnFinishReason && {
finishReason: sess.lastTurnFinishReason
}
};
};
return defineCustomAgent(
registry,
{
name: config.promptName,
description: config.description,
stateSchema: config.stateSchema,
store: config.store,
clientTransform: config.clientTransform
},
fn
);
}
function validateResumeAgainstHistory(resume, history) {
const allToolRequests = [];
for (const msg of history) {
if (msg.role === "model") {
for (const part of msg.content) {
if (part.toolRequest) {
allToolRequests.push(part.toolRequest);
}
}
}
}
for (const restart of resume.restart || []) {
const { name, ref, input } = restart.toolRequest;
const match = allToolRequests.find(
(tr) => tr.name === name && tr.ref === ref
);
if (!match) {
throw new GenkitError({
status: "INVALID_ARGUMENT",
message: `resume.restart references tool '${name}'` + (ref ? ` (ref: ${ref})` : "") + ` which was not found in session history.`
});
}
if (!deepEqual(input, match.input)) {
throw new GenkitError({
status: "INVALID_ARGUMENT",
message: `resume.restart for tool '${name}'` + (ref ? ` (ref: ${ref})` : "") + ` has modified inputs that do not match the original tool request in session history. Restart inputs must exactly match the interrupted tool request.`
});
}
}
for (const respond of resume.respond || []) {
const { name, ref } = respond.toolResponse;
const match = allToolRequests.find(
(tr) => tr.name === name && tr.ref === ref
);
if (!match) {
throw new GenkitError({
status: "INVALID_ARGUMENT",
message: `resume.respond references tool '${name}'` + (ref ? ` (ref: ${ref})` : "") + ` which was not found in session history.`
});
}
}
}
function defineAgent(registry, config) {
const { stateSchema, store, clientTransform, promptInput, ...promptConfig } = config;
definePrompt(registry, promptConfig);
return definePromptAgent(registry, {
promptName: promptConfig.name,
description: promptConfig.description,
promptInput,
stateSchema,
store,
clientTransform
});
}
export {
AgentAbortRequestSchema2 as AgentAbortRequestSchema,
AgentAbortResponseSchema2 as AgentAbortResponseSchema,
AgentInitError,
AgentInitSchema2 as AgentInitSchema,
AgentInputSchema2 as AgentInputSchema,
AgentOutputSchema2 as AgentOutputSchema,
AgentResultSchema,
AgentStreamChunkSchema2 as AgentStreamChunkSchema,
GetSnapshotDataInputSchema,
GetSnapshotRequestSchema2 as GetSnapshotRequestSchema,
JsonPatchOperationSchema,
JsonPatchSchema,
SessionRunner,
TurnEndSchema,
defineAgent,
defineCustomAgent,
definePromptAgent,
validateResumeAgainstHistory
};
//# sourceMappingURL=agent.mjs.map