feat(js): implemented defineAgent#5251
Conversation
Add a new test case in `session-flow_test.ts` to ensure that session flows correctly handle native tool interrupts and can successfully resume using the provided tool responses. Additionally, update the agents testapp index to import, log, and export the new `interrupt-agent` for further testing and validation of interrupt flows.
Change `abort()` to return the snapshot's previous status before setting it to 'aborted', or `undefined` if the snapshot was not found. This allows callers to distinguish between aborting a pending, done, or non-existent session. Update action output schema from `z.void()` to `z.string().optional()` and add tests covering all abort return value scenarios.
Wrap the flow function invocation in `runWithSession` to ensure the session is available in the async context during execution. This imports `runWithSession` from `session.js` and applies it around the `fn(runner, ...)` call in `defineSessionFlow`.
Add AgentInit, AgentInput, and AgentOutput type exports from @genkit-ai/ai to the beta public API surface.
Wrap the turn execution logic in `SessionRunner.runTurn` with a
`run()` call to produce proper tracing spans for each turn. This
improves observability by labeling each turn as `runTurn-{index}`.
Also includes formatting cleanups and updated tests.
Expose metadata on custom agent actions indicating whether state is managed on the server or client, and whether the agent supports abort via `onSnapshotStateChange`. This allows consumers to inspect agent capabilities without needing access to the original config.
Introduce a new 'agent-abort' action type to distinguish abort actions from regular agent actions. Previously, the abort action was registered with actionType 'agent', which made it indistinguishable from normal agent actions in the registry.
Introduce a `ClientStateTransform` type and `clientStateTransform` config option for `defineCustomAgent`, `definePromptAgent`, and `defineAgent`. When provided, the transform is applied to session state before it is returned to the client (in `AgentOutput.state`, `getSnapshotData`, etc.), allowing agents to redact sensitive fields or reshape state for external consumers.
The `getSnapshotDataAction` and `abortAgentAction` no longer append `__getSnapshotData` and `__abort` suffixes to `config.name`. These actions already have distinct `actionType` values (`agent-snapshot` and `agent-abort`), making the suffixes unnecessary for disambiguation.
Refactor prompt agent to pass session history into the prompt's
render() call instead of blindly appending it after rendered messages.
This lets prompts (e.g. dotprompt {{history}}) control where history
is placed. Tag history and preamble messages separately so prompt
template messages are correctly stripped from stored history after
generation while preserving tool-loop and model response messages.
Add comprehensive tests covering system-only, system+user, and
multi-turn prompt rendering to verify preamble messages don't
accumulate in stored history across turns.
github-actions
Bot
added
docs
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a unified Agent primitive for Genkit JS/TS, replacing the existing beta Session and SessionStore APIs with a snapshot-based model to support multi-turn conversations and state persistence. The changes include the addition of a SessionRunner for execution, new Agent definition APIs, and a file-system backed store. Review feedback identified several critical and medium-severity issues: the hardcoding of prompt inputs, a breach of encapsulation when accessing internal prompt properties, Node.js version compatibility concerns regarding crypto.randomUUID(), a lack of error handling robustness in the new EventEmitter utility, and the use of permissive z.any() types in action schemas which compromises type safety.
| { sendChunk, abortSignal } | ||
| ) => { | ||
| await sess.run(async (input) => { | ||
| const promptInput = {}; |
There was a problem hiding this comment.
The promptInput is hardcoded to an empty object {}. This is a critical issue because any agent defined from a prompt that requires input variables (e.g., {{name}} in a template) will fail to render correctly, as those variables are never provided.
The agent implementation should allow for initial input variables to be passed (likely via AgentInit) and persisted in the session state so they can be used across multiple turns.
| // Let the prompt control where history is placed (e.g. dotprompt | ||
| // {{history}}). When the prompt has no explicit `messages` config | ||
| // the render helper simply appends history after system/user. | ||
| const genOpts = await cachedPromptAction.__executablePrompt.render( |
There was a problem hiding this comment.
Accessing __executablePrompt breaks encapsulation as it is an internal/private property of PromptAction. This makes the agent implementation fragile and dependent on internal details of the prompt system.
Consider exposing a public render method on PromptAction or ExecutablePrompt that supports rendering with a provided history, or refactoring how agents interact with prompts to avoid relying on internal properties.
| this.session.addMessages(input.messages); | ||
| } | ||
|
|
||
| const turnSnapshotId = this.newSnapshotId || crypto.randomUUID(); |
There was a problem hiding this comment.
crypto.randomUUID() is not available globally in Node.js versions prior to 19.0.0 (including Node 18 LTS, which Genkit typically supports). In Node 18, crypto must be imported from the node:crypto module or accessed via globalThis.crypto (which was only added in later 18.x releases).
This will cause a ReferenceError in environments where it's not global. Please ensure compatibility by importing randomUUID or using a cross-platform utility.
|
|
||
| emit(event: string, ...args: any[]) { | ||
| if (!this.listeners[event]) return; | ||
| this.listeners[event].forEach((l) => l(...args)); |
There was a problem hiding this comment.
The custom EventEmitter implementation lacks robustness in its emit method. If a listener throws an error, it will stop the execution of subsequent listeners for the same event and propagate the error to the caller, which might not be expected.
Consider wrapping listener calls in a try-catch block to ensure that one failing listener doesn't prevent others from running.
References
- Avoid making changes for minor optimizations if the existing code is functionally correct, but ensure robustness in core utilities.
| description: `Gets snapshot data for ${config.name} by snapshotId`, | ||
| actionType: 'agent-snapshot', | ||
| inputSchema: z.string(), | ||
| outputSchema: z.any(), // SessionSnapshot Schema |
There was a problem hiding this comment.
Using z.any() for the outputSchema of getSnapshotDataAction reduces type safety and discoverability in the Genkit registry. Since the action is expected to return a SessionSnapshot, consider defining and using a proper Zod schema for it. Stricter validation is preferable to maintain consistency across API versions.
References
- Stricter validation is preferable to maintain consistency across API versions and improve type safety, even if it requires replacing permissive types like z.any().
Move `stateManagement` and `abortable` properties into a nested `agent` object within action metadata to avoid polluting the top-level metadata namespace. Update tests accordingly.
Use conditional spread syntax to only include fields like snapshotId, artifacts, message, and state when they have meaningful values, reducing noise in response objects.
Refactor AgentInputSchema to replace the flat `toolRestarts` field with a structured `resume` object containing `restart` and `respond` arrays. This enables both re-executing interrupted tools and directly providing tool responses when resuming a session. Add `validateResumeAgainstHistory` to verify that all resume entries reference tool requests present in session history and that restart inputs have not been tampered with, improving safety for interrupted generation flows.
Make `snapshotId` optional in `saveSnapshot` input and return the assigned ID from the store. This allows store implementations to encode grouping or routing information into snapshot IDs rather than requiring callers to generate UUIDs upfront. - Add `SessionSnapshotInput` type with optional `snapshotId` - Change `SessionStore.saveSnapshot` return type from `void` to `string` - Refactor `FileSessionStore` to use composite IDs (`convoId/suffix`) with a conversation-based directory structure - Update `InMemorySessionStore` to generate IDs when not provided - Update `SessionRunner` to use store-assigned snapshot IDs - Add helpers for parsing, building, and validating composite IDs
| */ | ||
| export const AgentInitSchema = z.object({ | ||
| snapshotId: z.string().optional(), | ||
| newSnapshotId: z.string().optional(), |
| const previousStatus = snapshot.status; | ||
| snapshot.status = 'aborted'; | ||
| await config.store.saveSnapshot(snapshot, { context: getContext() }); | ||
| return previousStatus; |
There was a problem hiding this comment.
I'm not sure this logic is helpful. It should return aborted unless it has reached a different terminal state in which case that status should be returned instead.
Remove the `newSnapshotId` optional field from `AgentInitSchema` as it is no longer needed in the agent initialization schema.
Export AgentInitSchema, AgentInputSchema, AgentOutputSchema, and AgentStreamChunkSchema alongside existing type exports from @genkit-ai/ai to allow runtime validation of agent types.
Refactor `saveSnapshot` to accept a snapshot ID and a mutator function instead of a pre-built snapshot object. The mutator receives the current snapshot and returns the updated value (or null to skip the write), enabling atomic read-check-write semantics. This prevents a race condition where a "done" write could overwrite a concurrent "aborted" status. The abort logic is also updated to skip writes on already-terminal snapshots (done, failed, aborted). - Add `SnapshotMutator` type to `SessionStore` interface - Update `InMemorySessionStore` and all callers to new signature - Remove separate `getSnapshot` + `saveSnapshot` pattern in abort flows - Add JSDoc explaining the concurrency-safety rationale
2 participants
Parent PR: #5248