fix: fixed local auth handling for streamed flows

js/flow/src/flow.ts

@@ -344,6 +344,7 @@ export class Flow<
             }) as S extends z.ZodVoid
               ? undefined
               : StreamingCallback<z.infer<S>>,
+            auth: opts?.withLocalAuthContext,
           }
         ).then((s) => s.result)
       )

js/flow/tests/flow_test.ts

@@ -20,6 +20,7 @@ import assert from 'node:assert';
 import { beforeEach, describe, it } from 'node:test';
 import { z } from 'zod';
 import { defineFlow, defineStreamingFlow } from '../src/flow.js';
+import { getFlowAuth } from '../src/utils.js';
 
 function createTestFlow() {
   return defineFlow(
@@ -34,6 +35,42 @@ function createTestFlow() {
   );
 }
 
+function createTestFlowWithAuth() {
+  return defineFlow(
+    {
+      name: 'testFlowWithAuth',
+      inputSchema: z.string(),
+      outputSchema: z.string(),
+      authPolicy: async (auth) => {
+        if (auth != 'open sesame') {
+          throw 'forty thieves!';
+        }
+      },
+    },
+    async (input) => {
+      return `foo ${input}, auth ${JSON.stringify(getFlowAuth())}`;
+    }
+  );
+}
+
+function createTestStreamingFlowWithAuth() {
+  return defineStreamingFlow(
+    {
+      name: 'testFlowWithAuth',
+      inputSchema: z.string(),
+      outputSchema: z.string(),
+      authPolicy: async (auth) => {
+        if (auth != 'open sesame') {
+          throw 'forty thieves!';
+        }
+      },
+    },
+    async (input) => {
+      return `foo ${input}, auth ${JSON.stringify(getFlowAuth())}`;
+    }
+  );
+}
+
 function createTestStreamingFlow() {
   return defineStreamingFlow(
     {
@@ -111,6 +148,28 @@ describe('flow', () => {
         }
       );
     });
+
+    it('should pass auth context all the way', async () => {
+      const testFlow = createTestFlowWithAuth();
+
+      const result = await testFlow('bar', {
+        withLocalAuthContext: 'open sesame',
+      });
+
+      assert.equal(result, 'foo bar, auth "open sesame"');
+    });
+
+    it('should fail auth', async () => {
+      const testFlow = createTestFlowWithAuth();
+
+      await assert.rejects(
+        () =>
+          testFlow('bar', {
+            withLocalAuthContext: 'yolo',
+          }),
+        /forty thieves/
+      );
+    });
   });
 
   describe('streamFlow', () => {
@@ -145,5 +204,24 @@ describe('flow', () => {
         message: 'stream bad happened: foo',
       });
     });
+
+    it('should pass auth context all the way', async () => {
+      const testFlow = createTestStreamingFlowWithAuth();
+
+      const result = testFlow('bar', {
+        withLocalAuthContext: 'open sesame',
+      });
+
+      assert.equal(await result.output, 'foo bar, auth "open sesame"');
+    });
+
+    it('should fail auth', async () => {
+      const testFlow = createTestStreamingFlowWithAuth();
+      const response = testFlow('bar', {
+        withLocalAuthContext: 'yolo',
+      });
+
+      await assert.rejects(() => response.output, /forty thieves/);
+    });
   });
 });

js/plugins/google-cloud/tests/logs_test.ts

@@ -19,10 +19,10 @@ import { defineModel, GenerateResponseData } from '@genkit-ai/ai/model';
 import { configureGenkit } from '@genkit-ai/core';
 import { defineFlow, run } from '@genkit-ai/flow';
 import {
-  googleCloud,
   __addTransportStreamForTesting,
   __forceFlushSpansForTesting,
   __getSpanExporterForTesting,
+  googleCloud,
 } from '@genkit-ai/google-cloud';
 import { ReadableSpan } from '@opentelemetry/sdk-trace-base';
 import assert from 'node:assert';

js/plugins/google-cloud/tests/metrics_test.ts

@@ -19,11 +19,11 @@ import { defineModel } from '@genkit-ai/ai/model';
 import { configureGenkit, defineAction } from '@genkit-ai/core';
 import { defineFlow, run } from '@genkit-ai/flow';
 import {
-  GcpOpenTelemetry,
-  googleCloud,
   __forceFlushSpansForTesting,
   __getMetricExporterForTesting,
   __getSpanExporterForTesting,
+  GcpOpenTelemetry,
+  googleCloud,
 } from '@genkit-ai/google-cloud';
 import {
   DataPoint,