diff --git a/js/ai/src/generate.ts b/js/ai/src/generate.ts
index 2cea52ebd1..8affb2aec7 100755
--- a/js/ai/src/generate.ts
+++ b/js/ai/src/generate.ts
@@ -156,7 +156,7 @@ export interface GenerateOptions<
   /** Maximum number of tool call iterations that can be performed in a single generate call (default 5). */
   maxTurns?: number;
   /** When provided, models supporting streaming will call the provided callback with chunks as generation progresses. */
-  onChunk?: StreamingCallback<GenerateResponseChunk>;
+  onChunk?: StreamingCallback<GenerateResponseChunk<z.infer<O>>>;
   /**
    * When provided, models supporting streaming will call the provided callback with chunks as generation progresses.
    *
@@ -298,7 +298,10 @@ async function toolsToActionRefs(
   return tools;
 }
 
-function messagesFromOptions(options: GenerateOptions): MessageData[] {
+function messagesFromOptions<
+  O extends z.ZodTypeAny = z.ZodTypeAny,
+  CustomOptions extends z.ZodTypeAny = z.ZodType,
+>(options: GenerateOptions<O, CustomOptions>): MessageData[] {
   const messages: MessageData[] = [];
   if (options.system) {
     messages.push({
@@ -549,7 +552,7 @@ export type GenerateStreamOptions<
 > = Omit<GenerateOptions<O, CustomOptions>, 'streamingCallback'>;
 
 export interface GenerateStreamResponse<O extends z.ZodTypeAny = z.ZodTypeAny> {
-  get stream(): AsyncIterable<GenerateResponseChunk>;
+  get stream(): AsyncIterable<GenerateResponseChunk<O>>;
   get response(): Promise<GenerateResponse<O>>;
 }
 
@@ -562,7 +565,7 @@ export function generateStream<
     | GenerateOptions<O, CustomOptions>
     | PromiseLike<GenerateOptions<O, CustomOptions>>
 ): GenerateStreamResponse<O> {
-  const channel = new Channel<GenerateResponseChunk>();
+  const channel = new Channel<GenerateResponseChunk<O>>();
 
   const generated = Promise.resolve(options).then((resolvedOptions) =>
     generate<O, CustomOptions>(registry, {
diff --git a/js/genkit/src/genkit.ts b/js/genkit/src/genkit.ts
index 77abd1e531..f8def4622f 100644
--- a/js/genkit/src/genkit.ts
+++ b/js/genkit/src/genkit.ts
@@ -358,7 +358,7 @@ export class Genkit implements HasRegistry {
       input?: I,
       opts?: PromptGenerateOptions<O, CustomOptions>
     ): GenerateStreamResponse<O> => {
-      let channel = new Channel<GenerateResponseChunk>();
+      let channel = new Channel<GenerateResponseChunk<O>>();
 
       const generated = runInNewSpan(
         this.registry,