From c4f6f2101db1b7cb4285e11aedc65cec233b2d8b Mon Sep 17 00:00:00 2001 From: Elliot Hesp Date: Fri, 27 Feb 2026 21:38:24 +0000 Subject: [PATCH 1/2] feat(js): typed context --- js/ai/src/tool.ts | 123 ++++++++++++++------- js/core/src/action.ts | 7 +- js/core/src/context.ts | 2 +- js/core/src/flow.ts | 30 +++-- js/genkit/src/genkit.ts | 86 +++++++++++---- js/pnpm-lock.yaml | 13 +++ js/testapps/typed-context/README.md | 78 +++++++++++++ js/testapps/typed-context/package.json | 24 ++++ js/testapps/typed-context/src/index.ts | 139 ++++++++++++++++++++++++ js/testapps/typed-context/tsconfig.json | 15 +++ 10 files changed, 438 insertions(+), 79 deletions(-) create mode 100644 js/testapps/typed-context/README.md create mode 100644 js/testapps/typed-context/package.json create mode 100644 js/testapps/typed-context/src/index.ts create mode 100644 js/testapps/typed-context/tsconfig.json diff --git a/js/ai/src/tool.ts b/js/ai/src/tool.ts index 585669ecf0..62cfd19297 100644 --- a/js/ai/src/tool.ts +++ b/js/ai/src/tool.ts @@ -22,7 +22,6 @@ import { stripUndefinedProps, z, type Action, - type ActionContext, type ActionRunOptions, type JSONSchema7, } from '@genkit-ai/core'; @@ -280,38 +279,56 @@ export function toToolDefinition( return out; } -export interface ToolFnOptions extends ActionFnArg { +/** + * Options passed to tool callbacks. Context is typed as C & ActionContext when using defineTool with a typed Genkit instance (genkit()). + */ +export interface ToolFnOptions + extends ActionFnArg { /** * A function that can be called during tool execution that will result in the tool * getting interrupted (immediately) and tool request returned to the upstream caller. */ interrupt: (metadata?: Record) => never; - - context: ActionContext; } -export type ToolFn = ( +export type ToolFn< + I extends z.ZodTypeAny = z.ZodTypeAny, + O extends z.ZodTypeAny = z.ZodTypeAny, + C extends object = object, +> = ( input: z.infer, - ctx: ToolFnOptions & ToolRunOptions + ctx: ToolFnOptions & ToolRunOptions ) => Promise>; -export type MultipartToolFn = ( +export type MultipartToolFn< + I extends z.ZodTypeAny = z.ZodTypeAny, + O extends z.ZodTypeAny = z.ZodTypeAny, + C extends object = object, +> = ( input: z.infer, - ctx: ToolFnOptions & ToolRunOptions + ctx: ToolFnOptions & ToolRunOptions ) => Promise<{ output?: z.infer; content?: Part[]; }>; -export function defineTool( +export function defineTool< + I extends z.ZodTypeAny = z.ZodTypeAny, + O extends z.ZodTypeAny = z.ZodTypeAny, + C extends object = object, +>( registry: Registry, config: { multipart: true } & ToolConfig, - fn?: ToolFn + fn?: ToolFn ): MultipartToolAction; -export function defineTool( +export function defineTool< + I extends z.ZodTypeAny = z.ZodTypeAny, + O extends z.ZodTypeAny = z.ZodTypeAny, + C extends object = object, +>( registry: Registry, config: ToolConfig, - fn?: ToolFn + fn?: ToolFn ): ToolAction; /** @@ -319,17 +336,24 @@ export function defineTool( * * A tool is an action that can be passed to a model to be called automatically if it so chooses. */ -export function defineTool( +export function defineTool< + I extends z.ZodTypeAny = z.ZodTypeAny, + O extends z.ZodTypeAny = z.ZodTypeAny, + C extends object = object, +>( registry: Registry, config: { multipart?: true } & ToolConfig, - fn?: ToolFn | MultipartToolFn + fn?: ToolFn | MultipartToolFn ): ToolAction | MultipartToolAction { const a = tool(config, fn); delete a.__action.metadata.dynamic; registry.registerAction(config.multipart ? 'tool.v2' : 'tool', a); if (!config.multipart) { // For non-multipart tools, we register a v2 tool action as well - registry.registerAction('tool.v2', basicToolV2(config, fn as ToolFn)); + registry.registerAction( + 'tool.v2', + basicToolV2(config, fn as ToolFn) + ); } return a as ToolAction; } @@ -469,30 +493,40 @@ function interruptTool(registry?: Registry) { }; } -export function tool( +export function tool< + I extends z.ZodTypeAny = z.ZodTypeAny, + O extends z.ZodTypeAny = z.ZodTypeAny, + C extends object = object, +>( config: { multipart: true } & ToolConfig, - fn?: ToolFn + fn?: ToolFn ): MultipartToolAction; -export function tool( - config: ToolConfig, - fn?: ToolFn -): ToolAction; +export function tool< + I extends z.ZodTypeAny = z.ZodTypeAny, + O extends z.ZodTypeAny = z.ZodTypeAny, + C extends object = object, +>(config: ToolConfig, fn?: ToolFn): ToolAction; /** * Defines a dynamic tool. Dynamic tools are just like regular tools but will not be registered in the * Genkit registry and can be defined dynamically at runtime. */ -export function tool( +export function tool< + I extends z.ZodTypeAny = z.ZodTypeAny, + O extends z.ZodTypeAny = z.ZodTypeAny, + C extends object = object, +>( config: { multipart?: true } & ToolConfig, - fn?: ToolFn | MultipartToolFn + fn?: ToolFn | MultipartToolFn ): ToolAction | MultipartToolAction { return config.multipart ? multipartTool(config, fn) : basicTool(config, fn); } -function basicTool( - config: ToolConfig, - fn?: ToolFn -): ToolAction { +function basicTool< + I extends z.ZodTypeAny = z.ZodTypeAny, + O extends z.ZodTypeAny = z.ZodTypeAny, + C extends object = object, +>(config: ToolConfig, fn?: ToolFn): ToolAction { const a = action( { ...config, @@ -506,7 +540,7 @@ function basicTool( ...runOptions, context: { ...runOptions.context }, interrupt, - }); + } as unknown as ToolFnOptions & ToolRunOptions); } return interrupt(); } @@ -515,24 +549,32 @@ function basicTool( return a; } -function basicToolV2( - config: ToolConfig, - fn?: ToolFn -): MultipartToolAction { +function basicToolV2< + I extends z.ZodTypeAny = z.ZodTypeAny, + O extends z.ZodTypeAny = z.ZodTypeAny, + C extends object = object, +>(config: ToolConfig, fn?: ToolFn): MultipartToolAction { return multipartTool(config, async (input, ctx) => { if (!fn) { const interrupt = interruptTool(ctx.registry); return interrupt(); } return { - output: await fn(input, ctx), + output: await fn( + input, + ctx as unknown as ToolFnOptions & ToolRunOptions + ), }; }); } -function multipartTool( +function multipartTool< + I extends z.ZodTypeAny = z.ZodTypeAny, + O extends z.ZodTypeAny = z.ZodTypeAny, + C extends object = object, +>( config: ToolConfig, - fn?: MultipartToolFn + fn?: MultipartToolFn ): MultipartToolAction { const a = action( { @@ -553,7 +595,7 @@ function multipartTool( ...runOptions, context: { ...runOptions.context }, interrupt, - }); + } as unknown as ToolFnOptions & ToolRunOptions); } return interrupt() as any; // we cast to any because `interrupt` throws. } @@ -568,10 +610,11 @@ function multipartTool( * * @deprecated renamed to {@link tool}. */ -export function dynamicTool( - config: ToolConfig, - fn?: ToolFn -): DynamicToolAction { +export function dynamicTool< + I extends z.ZodTypeAny = z.ZodTypeAny, + O extends z.ZodTypeAny = z.ZodTypeAny, + C extends object = object, +>(config: ToolConfig, fn?: ToolFn): DynamicToolAction { const t = basicTool(config, fn) as DynamicToolAction; t.attach = (_: Registry) => t; return t; diff --git a/js/core/src/action.ts b/js/core/src/action.ts index b1992e3aa6..13e36e76d9 100644 --- a/js/core/src/action.ts +++ b/js/core/src/action.ts @@ -120,8 +120,9 @@ export interface ActionRunOptions { /** * Options (side channel) data to pass to the model. + * @template C - Your app's context shape; ActionContext (auth, etc.) is merged in automatically, so you only declare your own fields. */ -export interface ActionFnArg { +export interface ActionFnArg { /** * Whether the caller of the action requested streaming. */ @@ -133,9 +134,9 @@ export interface ActionFnArg { sendChunk: StreamingCallback; /** - * Additional runtime context data (ex. auth context data). + * Runtime context (always set when the action runs). Typed as C & ActionContext so auth and other built-in context are always available. */ - context?: ActionContext; + context: C & ActionContext; /** * Trace context containing trace and span IDs. diff --git a/js/core/src/context.ts b/js/core/src/context.ts index a5cfb7c5df..ec0ad63006 100644 --- a/js/core/src/context.ts +++ b/js/core/src/context.ts @@ -21,7 +21,7 @@ import { UserFacingError } from './error.js'; const contextAlsKey = 'core.auth.context'; /** - * Action side channel data, like auth and other invocation context infromation provided by the invoker. + * Action side channel data, like auth and other invocation context information provided by the invoker. */ export interface ActionContext { /** Information about the currently authenticated user if provided. */ diff --git a/js/core/src/flow.ts b/js/core/src/flow.ts index 56c3a3d6b9..95ea68bffa 100644 --- a/js/core/src/flow.ts +++ b/js/core/src/flow.ts @@ -16,6 +16,7 @@ import type { z } from 'zod'; import { ActionFnArg, action, type Action } from './action.js'; +import type { ActionContext } from './context.js'; import { Registry, type HasRegistry } from './registry.js'; import { SPAN_TYPE_ATTR, runInNewSpan } from './tracing.js'; @@ -53,22 +54,25 @@ export interface FlowConfig< * side-channel context data. The context itself is a function, a short-cut * for streaming callback. */ -export interface FlowSideChannel extends ActionFnArg { +export interface FlowSideChannel + extends ActionFnArg { (chunk: S): void; } /** * Function to be executed in the flow. + * @template C - Your app's context shape; ActionContext is merged in automatically in the callback's context. */ export type FlowFn< I extends z.ZodTypeAny = z.ZodTypeAny, O extends z.ZodTypeAny = z.ZodTypeAny, S extends z.ZodTypeAny = z.ZodTypeAny, + C extends object = object, > = ( /** Input to the flow. */ input: z.infer, /** Callback for streaming functions only. */ - streamingCallback: FlowSideChannel> + streamingCallback: FlowSideChannel, C> ) => Promise> | z.infer; /** @@ -78,7 +82,8 @@ export function flow< I extends z.ZodTypeAny = z.ZodTypeAny, O extends z.ZodTypeAny = z.ZodTypeAny, S extends z.ZodTypeAny = z.ZodTypeAny, ->(config: FlowConfig | string, fn: FlowFn): Flow { + C extends object = object, +>(config: FlowConfig | string, fn: FlowFn): Flow { const resolvedConfig: FlowConfig = typeof config === 'string' ? { name: config } : config; @@ -92,10 +97,11 @@ export function defineFlow< I extends z.ZodTypeAny = z.ZodTypeAny, O extends z.ZodTypeAny = z.ZodTypeAny, S extends z.ZodTypeAny = z.ZodTypeAny, + C extends object = object, >( registry: Registry, config: FlowConfig | string, - fn: FlowFn + fn: FlowFn ): Flow { const f = flow(config, fn); @@ -111,7 +117,8 @@ function flowAction< I extends z.ZodTypeAny = z.ZodTypeAny, O extends z.ZodTypeAny = z.ZodTypeAny, S extends z.ZodTypeAny = z.ZodTypeAny, ->(config: FlowConfig, fn: FlowFn): Flow { + C extends object = object, +>(config: FlowConfig, fn: FlowFn): Flow { return action( { actionType: 'flow', @@ -126,13 +133,14 @@ function flowAction< { sendChunk, context, trace, abortSignal, streamingRequested } ) => { const ctx = sendChunk; - (ctx as FlowSideChannel>).sendChunk = sendChunk; - (ctx as FlowSideChannel>).context = context; - (ctx as FlowSideChannel>).trace = trace; - (ctx as FlowSideChannel>).abortSignal = abortSignal; - (ctx as FlowSideChannel>).streamingRequested = + (ctx as FlowSideChannel, C>).sendChunk = sendChunk; + (ctx as FlowSideChannel, C>).context = context as C & + ActionContext; + (ctx as FlowSideChannel, C>).trace = trace; + (ctx as FlowSideChannel, C>).abortSignal = abortSignal; + (ctx as FlowSideChannel, C>).streamingRequested = streamingRequested; - return fn(input, ctx as FlowSideChannel>); + return fn(input, ctx as FlowSideChannel, C>); } ); } diff --git a/js/genkit/src/genkit.ts b/js/genkit/src/genkit.ts index 2f7d4ea005..3b09a4bac9 100644 --- a/js/genkit/src/genkit.ts +++ b/js/genkit/src/genkit.ts @@ -156,16 +156,18 @@ export type PromptFn< /** * Options for initializing Genkit. + * When using genkit(), pass options that extend GenkitOptions so context is type-checked. + * Your C is merged with ActionContext (auth, etc.) in flow/tool callbacks — you only declare your own fields. */ -export interface GenkitOptions { +export interface GenkitOptions { /** List of plugins to load. */ plugins?: (GenkitPlugin | GenkitPluginV2)[]; /** Directory where dotprompts are stored. */ promptDir?: string; /** Default model to use if no model is specified. */ model?: ModelArgument; - /** Additional runtime context data for flows and tools. */ - context?: ActionContext; + /** Additional runtime context data for flows and tools. Must match C when using genkit(). */ + context?: C; /** Display name that will be shown in developer tooling. */ name?: string; /** Additional attribution information to include in the x-goog-api-client header. */ @@ -180,10 +182,12 @@ export interface GenkitOptions { * Registry keeps track of actions, flows, tools, and many other components. Reflection server exposes an API to inspect the registry and trigger executions of actions in the registry. Flow server exposes flows as HTTP endpoints for production use. * * There may be multiple Genkit instances in a single codebase. + * + * @template C - App context type for flow/tool callbacks and {@link currentContext}. Pass when calling {@link genkit}, e.g. `genkit({})`, so `context` is typed in all flows and tools without passing a type param each time. */ -export class Genkit implements HasRegistry { +export class Genkit implements HasRegistry { /** Developer-configured options. */ - readonly options: GenkitOptions; + readonly options: GenkitOptions; /** Registry instance that is exclusively modified by this Genkit instance. */ readonly registry: Registry; /** Reflection server for this registry. May be null if not started. */ @@ -195,7 +199,7 @@ export class Genkit implements HasRegistry { return this.registry.apiStability; } - constructor(options?: GenkitOptions) { + constructor(options?: GenkitOptions) { this.options = options || {}; this.registry = new Registry(); if (this.options.context) { @@ -216,6 +220,7 @@ export class Genkit implements HasRegistry { /** * Defines and registers a flow function. + * Context in the callback is typed from the instance's context type (from `genkit({})`). */ defineFlow< I extends z.ZodTypeAny = z.ZodTypeAny, @@ -223,7 +228,7 @@ export class Genkit implements HasRegistry { S extends z.ZodTypeAny = z.ZodTypeAny, >( config: FlowConfig | string, - fn: FlowFn + fn: FlowFn ): Action { const flow = defineFlow(this.registry, config, fn); this.flows.push(flow); @@ -234,25 +239,27 @@ export class Genkit implements HasRegistry { * Defines and registers a tool that can return multiple parts of content. * * Tools can be passed to models by name or value during `generate` calls to be called automatically based on the prompt and situation. + * Context in the callback is typed from the instance (genkit()). */ defineTool( config: { multipart: true } & ToolConfig, - fn: MultipartToolFn + fn: MultipartToolFn ): MultipartToolAction; /** * Defines and registers a tool. * * Tools can be passed to models by name or value during `generate` calls to be called automatically based on the prompt and situation. + * Context in the callback is typed from the instance (genkit()). */ defineTool( config: ToolConfig, - fn: ToolFn + fn: ToolFn ): ToolAction; defineTool( config: ({ multipart?: true } & ToolConfig) | string, - fn: ToolFn | MultipartToolFn + fn: ToolFn | MultipartToolFn ): ToolAction | MultipartToolAction { return defineTool(this.registry, config as any, fn as any); } @@ -263,7 +270,7 @@ export class Genkit implements HasRegistry { */ dynamicTool( config: ToolConfig, - fn?: ToolFn + fn?: ToolFn ): ToolAction { return dynamicTool(config, fn) as ToolAction; } @@ -298,6 +305,7 @@ export class Genkit implements HasRegistry { /** * Defines a new model and adds it to the registry. + * The runner's options.context is typed from the instance (genkit()). */ defineModel( options: { @@ -305,7 +313,7 @@ export class Genkit implements HasRegistry { } & DefineModelOptions, runner: ( request: GenerateRequest, - options: ActionFnArg + options: ActionFnArg ) => Promise ): ModelAction; @@ -740,8 +748,14 @@ export class Genkit implements HasRegistry { CustomOptions extends z.ZodTypeAny = typeof GenerationCommonConfigSchema, >( opts: - | GenerateOptions - | PromiseLike> + | (Omit, 'context'> & { + context?: C & ActionContext; + }) + | PromiseLike< + Omit, 'context'> & { + context?: C & ActionContext; + } + > ): Promise>>; async generate< @@ -751,8 +765,14 @@ export class Genkit implements HasRegistry { options: | string | Part[] - | GenerateOptions - | PromiseLike> + | (Omit, 'context'> & { + context?: C & ActionContext; + }) + | PromiseLike< + Omit, 'context'> & { + context?: C & ActionContext; + } + > ): Promise>> { let resolvedOptions: GenerateOptions; if (options instanceof Promise) { @@ -844,8 +864,14 @@ export class Genkit implements HasRegistry { CustomOptions extends z.ZodTypeAny = typeof GenerationCommonConfigSchema, >( parts: - | GenerateOptions - | PromiseLike> + | (Omit, 'context'> & { + context?: C & ActionContext; + }) + | PromiseLike< + Omit, 'context'> & { + context?: C & ActionContext; + } + > ): GenerateStreamResponse>; generateStream< @@ -855,8 +881,14 @@ export class Genkit implements HasRegistry { options: | string | Part[] - | GenerateStreamOptions - | PromiseLike> + | (Omit, 'context'> & { + context?: C & ActionContext; + }) + | PromiseLike< + Omit, 'context'> & { + context?: C & ActionContext; + } + > ): GenerateStreamResponse> { if (typeof options === 'string' || Array.isArray(options)) { options = { prompt: options }; @@ -949,8 +981,8 @@ export class Genkit implements HasRegistry { * data set by HTTP server frameworks. If invoked outside of an action (e.g. flow or tool) will * return `undefined`. */ - currentContext(): ActionContext | undefined { - return getContext(); + currentContext(): (C & ActionContext) | undefined { + return getContext() as (C & ActionContext) | undefined; } /** @@ -1080,9 +1112,15 @@ function registerActionV2( * * This will create a new Genkit registry, register the provided plugins, stores, and other configuration. This * should be called before any flows are registered. + * + * Pass your app's context type for typed `context` in flow/tool callbacks and {@link Genkit.currentContext}, e.g. + * `const ai = genkit({})`. If you omit the type argument, context in flows/tools is inferred as + * {@link ActionContext} and properties like context.hello are typed as `any` (no type safety). */ -export function genkit(options: GenkitOptions): Genkit { - return new Genkit(options); +export function genkit( + options?: GenkitOptions +): Genkit { + return new Genkit(options); } const shutdown = async () => { diff --git a/js/pnpm-lock.yaml b/js/pnpm-lock.yaml index 75d3e7d7ec..ccf863f9a5 100644 --- a/js/pnpm-lock.yaml +++ b/js/pnpm-lock.yaml @@ -1886,6 +1886,19 @@ importers: specifier: ^5.3.3 version: 5.8.3 + testapps/typed-context: + dependencies: + genkit: + specifier: workspace:* + version: link:../../genkit + devDependencies: + rimraf: + specifier: ^6.0.1 + version: 6.0.1 + typescript: + specifier: ^5.3.3 + version: 5.8.3 + testapps/vertexai-modelgarden: dependencies: '@genkit-ai/firebase': diff --git a/js/testapps/typed-context/README.md b/js/testapps/typed-context/README.md new file mode 100644 index 0000000000..5a3f69516e --- /dev/null +++ b/js/testapps/typed-context/README.md @@ -0,0 +1,78 @@ +# Express Integration + +Demonstrates integrating Genkit flows with an Express.js server — including +authentication context, streaming responses, and the `expressHandler` utility. + +## Features Demonstrated + +| Feature | Endpoint | Description | +|---------|----------|-------------| +| Flow via `expressHandler` | `POST /jokeFlow` | Genkit flow exposed via Express with auth context | +| Flow handler (no auth) | `POST /jokeHandler` | Flow exposed without auth validation | +| Direct flow invocation | `GET /jokeWithFlow` | Call a flow directly from a route handler | +| Raw streaming | `GET /jokeStream` | Chunked transfer encoding with `ai.generate` | +| Auth context | `Authorization` header | Token-based auth with context validation | +| Context providers | `auth()` factory | Reusable auth context provider pattern | + +## Setup + +### Prerequisites + +- **Node.js** (v18 or higher) +- **pnpm** package manager + +### API Keys + +```bash +export GEMINI_API_KEY='' +``` + +### Build and Install + +From the repo root: + +```bash +pnpm install +pnpm run setup +``` + +## Run the Sample + +```bash +pnpm build && pnpm start +``` + +The Express server starts on port `5000` (or `$PORT`). + +## Testing This Demo + +1. **Test with auth** (requires `Authorization: open sesame` header): + ```bash + curl http://localhost:5000/jokeFlow?stream=true \ + -d '{"data": "banana"}' \ + -H "Content-Type: application/json" \ + -H "Authorization: open sesame" + ``` + +2. **Test without auth**: + ```bash + curl http://localhost:5000/jokeHandler?stream=true \ + -d '{"data": "banana"}' \ + -H "Content-Type: application/json" + ``` + +3. **Test direct flow invocation**: + ```bash + curl "http://localhost:5000/jokeWithFlow?subject=banana" + ``` + +4. **Test raw streaming**: + ```bash + curl "http://localhost:5000/jokeStream?subject=banana" + ``` + +5. **Expected behavior**: + - Authenticated requests (`open sesame`) succeed + - Unauthenticated requests return `PERMISSION_DENIED` + - Streaming endpoints deliver text incrementally + - `?stream=true` enables streaming on `expressHandler` endpoints diff --git a/js/testapps/typed-context/package.json b/js/testapps/typed-context/package.json new file mode 100644 index 0000000000..624ef0f8d8 --- /dev/null +++ b/js/testapps/typed-context/package.json @@ -0,0 +1,24 @@ +{ + "name": "typed-context", + "version": "1.0.0", + "description": "", + "main": "lib/index.js", + "scripts": { + "start": "node lib/index.js", + "compile": "tsc", + "build": "pnpm build:clean && pnpm compile", + "build:clean": "rimraf ./lib", + "build:watch": "tsc --watch", + "build-and-run": "pnpm build && node lib/index.js" + }, + "keywords": [], + "author": "", + "license": "ISC", + "dependencies": { + "genkit": "workspace:*" + }, + "devDependencies": { + "rimraf": "^6.0.1", + "typescript": "^5.3.3" + } +} diff --git a/js/testapps/typed-context/src/index.ts b/js/testapps/typed-context/src/index.ts new file mode 100644 index 0000000000..5a4e9d3247 --- /dev/null +++ b/js/testapps/typed-context/src/index.ts @@ -0,0 +1,139 @@ +/** + * Copyright 2024 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. + */ + +/** + * Typed-context test app: defines an AppContext, creates genkit(), + * then runs flows and a tool that use typed context and currentContext(). + */ + +import { genkit, z } from 'genkit'; +import { logger } from 'genkit/logging'; + +logger.setLogLevel('debug'); + +// --------------------------------------------------------------------------- +// App context type and Genkit instance +// --------------------------------------------------------------------------- + +type AppContext = { + echo: (value: string) => string; + hello: (name: string) => string; + userId?: string; +}; + +const ai = genkit({ + context: { + echo: (value) => value, + hello: (name) => `Hello, ${name}!`, + }, +}); + +// --------------------------------------------------------------------------- +// Flows and tool (context is typed as AppContext & ActionContext) +// --------------------------------------------------------------------------- + +const echoFlow = ai.defineFlow( + { + name: 'echoWithContext', + inputSchema: z.string(), + outputSchema: z.string(), + }, + async (input, { context }) => context.echo(input) +); + +const greetFlow = ai.defineFlow( + { + name: 'greetWithContext', + inputSchema: z.object({ name: z.string() }), + outputSchema: z.string(), + }, + async (input, { context }) => context.hello(input.name) +); + +const flowWithCurrentContext = ai.defineFlow( + { + name: 'flowUsingCurrentContext', + inputSchema: z.string(), + outputSchema: z.string(), + }, + async (input) => { + const ctx = ai.currentContext(); + const userId = ctx?.userId ?? 'anonymous'; + return `${input} (user: ${userId})`; + } +); + +const getGreetingTool = ai.defineTool( + { + name: 'getGreeting', + description: 'Returns a greeting using the current context', + inputSchema: z.object({ name: z.string() }), + outputSchema: z.string(), + }, + async (input, { context }) => context.hello(input.name) +); + +// --------------------------------------------------------------------------- +// Run examples (log output so you can see typed context in action) +// --------------------------------------------------------------------------- + +async function main() { + const echoAction = await ai.registry.lookupAction('/flow/echoWithContext'); + + console.log( + 'Echo with default context:', + (await echoAction.run('foo')).result + ); + + console.log( + 'Echo with overridden context:', + ( + await echoFlow.run('foo', { + context: { + echo: (value: string) => `${value} ${value}`, + hello: (name: string) => `Hi, ${name}`, + }, + }) + ).result + ); + + console.log('Greet:', (await greetFlow.run({ name: 'World' })).result); + + console.log( + 'currentContext() without userId:', + (await flowWithCurrentContext.run('test')).result + ); + + console.log( + 'currentContext() with userId:', + ( + await flowWithCurrentContext.run('test', { + context: { + userId: 'user-123', + echo: (v: string) => v, + hello: (n: string) => `Hey ${n}`, + }, + }) + ).result + ); + + console.log( + 'Tool with typed context:', + (await getGreetingTool.run({ name: 'Tool User' })).result + ); +} + +main().catch(console.error); diff --git a/js/testapps/typed-context/tsconfig.json b/js/testapps/typed-context/tsconfig.json new file mode 100644 index 0000000000..e51f33ae38 --- /dev/null +++ b/js/testapps/typed-context/tsconfig.json @@ -0,0 +1,15 @@ +{ + "compilerOptions": { + "module": "NodeNext", + "noImplicitReturns": true, + "noUnusedLocals": false, + "outDir": "lib", + "sourceMap": true, + "strict": true, + "target": "es2017", + "skipLibCheck": true, + "esModuleInterop": true + }, + "compileOnSave": true, + "include": ["src"] +} From 116358cb278b949519c456503b699550b9485b11 Mon Sep 17 00:00:00 2001 From: Elliot Hesp Date: Sat, 28 Feb 2026 19:20:35 +0000 Subject: [PATCH 2/2] docs: update readme --- js/testapps/typed-context/README.md | 86 ++++++++++------------------- 1 file changed, 29 insertions(+), 57 deletions(-) diff --git a/js/testapps/typed-context/README.md b/js/testapps/typed-context/README.md index 5a3f69516e..387b1f0298 100644 --- a/js/testapps/typed-context/README.md +++ b/js/testapps/typed-context/README.md @@ -1,33 +1,21 @@ -# Express Integration +# Typed Context -Demonstrates integrating Genkit flows with an Express.js server — including -authentication context, streaming responses, and the `expressHandler` utility. +Demonstrates **typed context** in Genkit: define an `AppContext` type, pass it to `genkit({...})`, and get typed `context` in flows, tools, `generate()` options, and `ai.currentContext()`. -## Features Demonstrated +## What this sample shows -| Feature | Endpoint | Description | -|---------|----------|-------------| -| Flow via `expressHandler` | `POST /jokeFlow` | Genkit flow exposed via Express with auth context | -| Flow handler (no auth) | `POST /jokeHandler` | Flow exposed without auth validation | -| Direct flow invocation | `GET /jokeWithFlow` | Call a flow directly from a route handler | -| Raw streaming | `GET /jokeStream` | Chunked transfer encoding with `ai.generate` | -| Auth context | `Authorization` header | Token-based auth with context validation | -| Context providers | `auth()` factory | Reusable auth context provider pattern | +| Feature | Description | +|--------|-------------| +| **AppContext type** | Declare your app’s context shape (e.g. `echo`, `hello`, `userId`). | +| **genkit<AppContext>()** | Create the instance with the type so `context` is typed everywhere. | +| **Flows** | Flow callbacks receive `context: AppContext & ActionContext`. | +| **Tools** | Tool callbacks receive the same typed context. | +| **currentContext()** | `ai.currentContext()` returns `(AppContext & ActionContext) \| undefined`. | +| **Override at run time** | Pass a different `context` when calling a flow or action. | -## Setup - -### Prerequisites - -- **Node.js** (v18 or higher) -- **pnpm** package manager - -### API Keys +If you call `genkit()` without the type argument, `context` in flows/tools is typed as `ActionContext` and properties like `context.hello` are `any`. -```bash -export GEMINI_API_KEY='' -``` - -### Build and Install +## Setup From the repo root: @@ -36,43 +24,27 @@ pnpm install pnpm run setup ``` -## Run the Sample +## Run the sample + +From this directory: ```bash pnpm build && pnpm start ``` -The Express server starts on port `5000` (or `$PORT`). - -## Testing This Demo - -1. **Test with auth** (requires `Authorization: open sesame` header): - ```bash - curl http://localhost:5000/jokeFlow?stream=true \ - -d '{"data": "banana"}' \ - -H "Content-Type: application/json" \ - -H "Authorization: open sesame" - ``` - -2. **Test without auth**: - ```bash - curl http://localhost:5000/jokeHandler?stream=true \ - -d '{"data": "banana"}' \ - -H "Content-Type: application/json" - ``` +Or run the source directly: -3. **Test direct flow invocation**: - ```bash - curl "http://localhost:5000/jokeWithFlow?subject=banana" - ``` +```bash +pnpm exec tsx src/index.ts +``` -4. **Test raw streaming**: - ```bash - curl "http://localhost:5000/jokeStream?subject=banana" - ``` +## Example output -5. **Expected behavior**: - - Authenticated requests (`open sesame`) succeed - - Unauthenticated requests return `PERMISSION_DENIED` - - Streaming endpoints deliver text incrementally - - `?stream=true` enables streaming on `expressHandler` endpoints +``` +Echo with default context: foo +Echo with overridden context: foo foo +Greet: Hello, World! +currentContext() without userId: test (user: anonymous) +currentContext() with userId: test (user: user-123) +Tool with typed context: Hello, Tool User! +```