@genkit-ai/express 1.24.0 → 1.26.0-rc.0

package/README.md

@@ -59,6 +59,44 @@ app.post(
 );
 ```
 
+### Durable Streaming (Beta)
+
+You can configure flows to use a `StreamManager` to persist their state. This allows clients to disconnect and reconnect to a stream without losing its state.
+
+To enable durable streaming, provide a `streamManager` in the `expressHandler` options. The `InMemoryStreamManager` is useful for development and testing:
+
+```ts
+import { InMemoryStreamManager } from 'genkit/beta';
+
+// ...
+
+app.post(
+  '/myDurableFlow',
+  expressHandler(myFlow, {
+    streamManager: new InMemoryStreamManager(),
+  })
+);
+```
+
+For production environments, you should use a durable `StreamManager` implementation, such as `FirestoreStreamManager` or `RtdbStreamManager` from the `@genkit-ai/firebase` plugin, or a custom implementation.
+
+Clients can then connect and reconnect to the stream using the `streamId`:
+
+```ts
+// Start a new stream
+const result = streamFlow({
+  url: `http://localhost:8080/myDurableFlow`,
+  input: 'tell me a long story',
+});
+const streamId = await result.streamId; // Save this ID
+
+// ... later, reconnect if needed ...
+const reconnectedResult = streamFlow({
+  url: `http://localhost:8080/myDurableFlow`,
+  streamId: streamId,
+});
+```
+
 Flows and actions exposed using the `expressHandler` function can be accessed using `genkit/beta/client` library:
 
 ```ts

package/lib/index.d.mts

@@ -1,7 +1,7 @@
 import bodyParser from 'body-parser';
 import { CorsOptions } from 'cors';
 import express from 'express';
-import { ActionContext, z, Action, Flow } from 'genkit';
+import { ActionContext, z, Action, StreamManager, Flow } from 'genkit/beta';
 import { ContextProvider } from 'genkit/context';
 
 /**
@@ -25,24 +25,46 @@ import { ContextProvider } from 'genkit/context';
  */
 declare function expressHandler<C extends ActionContext = ActionContext, I extends z.ZodTypeAny = z.ZodTypeAny, O extends z.ZodTypeAny = z.ZodTypeAny, S extends z.ZodTypeAny = z.ZodTypeAny>(action: Action<I, O, S>, opts?: {
     contextProvider?: ContextProvider<C, I>;
+    streamManager?: StreamManager;
 }): express.RequestHandler;
 /**
  * A wrapper object containing a flow with its associated auth policy.
+ * @deprecated Use `withFlowOptions` instead.
  */
 type FlowWithContextProvider<C extends ActionContext = ActionContext, I extends z.ZodTypeAny = z.ZodTypeAny, O extends z.ZodTypeAny = z.ZodTypeAny, S extends z.ZodTypeAny = z.ZodTypeAny> = {
     flow: Flow<I, O, S>;
     context: ContextProvider<C, I>;
 };
+/**
+ * A wrapper object containing a flow with its associated options.
+ */
+type FlowWithOptions<I extends z.ZodTypeAny = z.ZodTypeAny, O extends z.ZodTypeAny = z.ZodTypeAny, S extends z.ZodTypeAny = z.ZodTypeAny> = {
+    flow: Flow<I, O, S>;
+    options: {
+        contextProvider?: ContextProvider<any, I>;
+        streamManager?: StreamManager;
+        path?: string;
+    };
+};
 /**
  * Adds an auth policy to the flow.
+ * @deprecated Use `withFlowOptions` instead.
  */
 declare function withContextProvider<C extends ActionContext = ActionContext, I extends z.ZodTypeAny = z.ZodTypeAny, O extends z.ZodTypeAny = z.ZodTypeAny, S extends z.ZodTypeAny = z.ZodTypeAny>(flow: Flow<I, O, S>, context: ContextProvider<C, I>): FlowWithContextProvider<C, I, O, S>;
+/**
+ * Adds an auth policy to the flow.
+ */
+declare function withFlowOptions<I extends z.ZodTypeAny, O extends z.ZodTypeAny, S extends z.ZodTypeAny>(flow: Flow<I, O, S>, options: {
+    contextProvider?: ContextProvider<any, I>;
+    streamManager?: StreamManager;
+    path?: string;
+}): FlowWithOptions<I, O, S>;
 /**
  * Options to configure the flow server.
  */
 interface FlowServerOptions {
     /** List of flows to expose via the flow server. */
-    flows: (Flow<any, any, any> | FlowWithContextProvider<any, any, any>)[];
+    flows: (Flow<any, any, any> | FlowWithContextProvider<any, any, any> | FlowWithOptions<any, any, any>)[];
     /** Port to run the server on. Defaults to env.PORT or 3400. */
     port?: number;
     /** CORS options for the server. */
@@ -87,4 +109,4 @@ declare class FlowServer {
     static stopAll(): Promise<void[]>;
 }
 
-export { FlowServer, type FlowServerOptions, type FlowWithContextProvider, expressHandler, startFlowServer, withContextProvider };
+export { FlowServer, type FlowServerOptions, type FlowWithContextProvider, type FlowWithOptions, expressHandler, startFlowServer, withContextProvider, withFlowOptions };

package/lib/index.d.ts

@@ -1,7 +1,7 @@
 import bodyParser from 'body-parser';
 import { CorsOptions } from 'cors';
 import express from 'express';
-import { ActionContext, z, Action, Flow } from 'genkit';
+import { ActionContext, z, Action, StreamManager, Flow } from 'genkit/beta';
 import { ContextProvider } from 'genkit/context';
 
 /**
@@ -25,24 +25,46 @@ import { ContextProvider } from 'genkit/context';
  */
 declare function expressHandler<C extends ActionContext = ActionContext, I extends z.ZodTypeAny = z.ZodTypeAny, O extends z.ZodTypeAny = z.ZodTypeAny, S extends z.ZodTypeAny = z.ZodTypeAny>(action: Action<I, O, S>, opts?: {
     contextProvider?: ContextProvider<C, I>;
+    streamManager?: StreamManager;
 }): express.RequestHandler;
 /**
  * A wrapper object containing a flow with its associated auth policy.
+ * @deprecated Use `withFlowOptions` instead.
  */
 type FlowWithContextProvider<C extends ActionContext = ActionContext, I extends z.ZodTypeAny = z.ZodTypeAny, O extends z.ZodTypeAny = z.ZodTypeAny, S extends z.ZodTypeAny = z.ZodTypeAny> = {
     flow: Flow<I, O, S>;
     context: ContextProvider<C, I>;
 };
+/**
+ * A wrapper object containing a flow with its associated options.
+ */
+type FlowWithOptions<I extends z.ZodTypeAny = z.ZodTypeAny, O extends z.ZodTypeAny = z.ZodTypeAny, S extends z.ZodTypeAny = z.ZodTypeAny> = {
+    flow: Flow<I, O, S>;
+    options: {
+        contextProvider?: ContextProvider<any, I>;
+        streamManager?: StreamManager;
+        path?: string;
+    };
+};
 /**
  * Adds an auth policy to the flow.
+ * @deprecated Use `withFlowOptions` instead.
  */
 declare function withContextProvider<C extends ActionContext = ActionContext, I extends z.ZodTypeAny = z.ZodTypeAny, O extends z.ZodTypeAny = z.ZodTypeAny, S extends z.ZodTypeAny = z.ZodTypeAny>(flow: Flow<I, O, S>, context: ContextProvider<C, I>): FlowWithContextProvider<C, I, O, S>;
+/**
+ * Adds an auth policy to the flow.
+ */
+declare function withFlowOptions<I extends z.ZodTypeAny, O extends z.ZodTypeAny, S extends z.ZodTypeAny>(flow: Flow<I, O, S>, options: {
+    contextProvider?: ContextProvider<any, I>;
+    streamManager?: StreamManager;
+    path?: string;
+}): FlowWithOptions<I, O, S>;
 /**
  * Options to configure the flow server.
  */
 interface FlowServerOptions {
     /** List of flows to expose via the flow server. */
-    flows: (Flow<any, any, any> | FlowWithContextProvider<any, any, any>)[];
+    flows: (Flow<any, any, any> | FlowWithContextProvider<any, any, any> | FlowWithOptions<any, any, any>)[];
     /** Port to run the server on. Defaults to env.PORT or 3400. */
     port?: number;
     /** CORS options for the server. */
@@ -87,4 +109,4 @@ declare class FlowServer {
     static stopAll(): Promise<void[]>;
 }
 
-export { FlowServer, type FlowServerOptions, type FlowWithContextProvider, expressHandler, startFlowServer, withContextProvider };
+export { FlowServer, type FlowServerOptions, type FlowWithContextProvider, type FlowWithOptions, expressHandler, startFlowServer, withContextProvider, withFlowOptions };

package/lib/index.js

@@ -31,18 +31,23 @@ __export(index_exports, {
   FlowServer: () => FlowServer,
   expressHandler: () => expressHandler,
   startFlowServer: () => startFlowServer,
-  withContextProvider: () => withContextProvider
+  withContextProvider: () => withContextProvider,
+  withFlowOptions: () => withFlowOptions
 });
 module.exports = __toCommonJS(index_exports);
 var import_body_parser = __toESM(require("body-parser"));
 var import_cors = __toESM(require("cors"));
+var import_crypto = require("crypto");
 var import_express = __toESM(require("express"));
+var import_beta = require("genkit/beta");
 var import_context = require("genkit/context");
 var import_logging = require("genkit/logging");
 const streamDelimiter = "\n\n";
 function expressHandler(action, opts) {
   return async (request, response) => {
     const { stream } = request.query;
+    const streamIdHeader = request.headers["x-genkit-stream-id"];
+    const streamId = Array.isArray(streamIdHeader) ? streamIdHeader[0] : streamIdHeader;
     if (!request.body) {
       const errMsg = `Error: request.body is undefined. Possible reasons: missing 'content-type: application/json' in request headers or misconfigured JSON middleware ('app.use(express.json()')? `;
       import_logging.logger.error(errMsg);
@@ -78,35 +83,29 @@ ${e.stack}`
       abortController.abort();
     });
     if (request.get("Accept") === "text/event-stream" || stream === "true") {
-      response.writeHead(200, {
+      const streamManager = opts?.streamManager;
+      if (streamManager && streamId) {
+        await subscribeToStream(streamManager, streamId, response);
+        return;
+      }
+      const streamIdToUse = (0, import_crypto.randomUUID)();
+      const headers = {
         "Content-Type": "text/plain",
         "Transfer-Encoding": "chunked"
-      });
-      try {
-        const onChunk = (chunk) => {
-          response.write(
-            "data: " + JSON.stringify({ message: chunk }) + streamDelimiter
-          );
-        };
-        const result = await action.run(input, {
-          onChunk,
-          context,
-          abortSignal: abortController.signal
-        });
-        response.write(
-          "data: " + JSON.stringify({ result: result.result }) + streamDelimiter
-        );
-        response.end();
-      } catch (e) {
-        import_logging.logger.error(
-          `Streaming request failed with error: ${e.message}
-${e.stack}`
-        );
-        response.write(
-          `error: ${JSON.stringify({ error: (0, import_context.getCallableJSON)(e) })}${streamDelimiter}`
-        );
-        response.end();
+      };
+      if (streamManager) {
+        headers["x-genkit-stream-id"] = streamIdToUse;
       }
+      response.writeHead(200, headers);
+      runActionWithDurableStreaming(
+        action,
+        streamManager,
+        streamIdToUse,
+        input,
+        context,
+        response,
+        abortController.signal
+      );
     } else {
       try {
         const result = await action.run(input, {
@@ -128,12 +127,112 @@ ${e.stack}`
     }
   };
 }
+async function runActionWithDurableStreaming(action, streamManager, streamId, input, context, response, abortSignal) {
+  let taskQueue;
+  let durableStream;
+  if (streamManager) {
+    taskQueue = new import_beta.AsyncTaskQueue();
+    durableStream = await streamManager.open(streamId);
+  }
+  try {
+    let onChunk = (chunk) => {
+      response.write(
+        "data: " + JSON.stringify({ message: chunk }) + streamDelimiter
+      );
+    };
+    if (streamManager) {
+      const originalOnChunk = onChunk;
+      onChunk = (chunk) => {
+        originalOnChunk(chunk);
+        taskQueue.enqueue(() => durableStream.write(chunk));
+      };
+    }
+    const result = await action.run(input, {
+      onChunk,
+      context,
+      abortSignal
+    });
+    if (streamManager) {
+      taskQueue.enqueue(() => durableStream.done(result.result));
+      await taskQueue.merge();
+    }
+    response.write(
+      "data: " + JSON.stringify({ result: result.result }) + streamDelimiter
+    );
+    response.end();
+  } catch (e) {
+    if (durableStream) {
+      taskQueue.enqueue(() => durableStream.error(e));
+      await taskQueue.merge();
+    }
+    import_logging.logger.error(
+      `Streaming request failed with error: ${e.message}
+${e.stack}`
+    );
+    response.write(
+      `error: ${JSON.stringify({
+        error: (0, import_context.getCallableJSON)(e)
+      })}${streamDelimiter}`
+    );
+    response.end();
+  }
+}
+async function subscribeToStream(streamManager, streamId, response) {
+  try {
+    await streamManager.subscribe(streamId, {
+      onChunk: (chunk) => {
+        response.write(
+          "data: " + JSON.stringify({ message: chunk }) + streamDelimiter
+        );
+      },
+      onDone: (output) => {
+        response.write(
+          "data: " + JSON.stringify({ result: output }) + streamDelimiter
+        );
+        response.end();
+      },
+      onError: (err) => {
+        import_logging.logger.error(
+          `Streaming request failed with error: ${err.message}
+${err.stack}`
+        );
+        response.write(
+          `error: ${JSON.stringify({
+            error: (0, import_context.getCallableJSON)(err)
+          })}${streamDelimiter}`
+        );
+        response.end();
+      }
+    });
+  } catch (e) {
+    if (e instanceof import_beta.StreamNotFoundError) {
+      response.status(204).end();
+      return;
+    }
+    if (e.status === "DEADLINE_EXCEEDED") {
+      response.write(
+        `error: ${JSON.stringify({
+          error: (0, import_context.getCallableJSON)(e)
+        })}${streamDelimiter}`
+      );
+      response.end();
+      return;
+    }
+    throw e;
+  }
+}
 function withContextProvider(flow, context) {
   return {
     flow,
     context
   };
 }
+function withFlowOptions(flow, options) {
+  return {
+    flow,
+    options
+  };
+}
 function startFlowServer(options) {
   const server = new FlowServer(options);
   server.start();
@@ -163,13 +262,11 @@ class FlowServer {
     import_logging.logger.debug("Running flow server with flow paths:");
     const pathPrefix = this.options.pathPrefix ?? "";
     this.options.flows?.forEach((flow) => {
-      if ("context" in flow) {
-        const flowPath = `/${pathPrefix}${flow.flow.__action.name}`;
+      if ("flow" in flow) {
+        const flowPath = `/${pathPrefix}${"options" in flow && flow.options.path || flow.flow.__action.name}`;
         import_logging.logger.debug(` - ${flowPath}`);
-        server.post(
-          flowPath,
-          expressHandler(flow.flow, { contextProvider: flow.context })
-        );
+        const options = "options" in flow ? flow.options : { contextProvider: flow.context };
+        server.post(flowPath, expressHandler(flow.flow, options));
       } else {
         const flowPath = `/${pathPrefix}${flow.__action.name}`;
         import_logging.logger.debug(` - ${flowPath}`);
@@ -224,6 +321,7 @@ class FlowServer {
   FlowServer,
   expressHandler,
   startFlowServer,
-  withContextProvider
+  withContextProvider,
+  withFlowOptions
 });
 //# sourceMappingURL=index.js.map

package/lib/index.js.map

@@ -1 +1 @@
-{"version":3,"sources":["../src/index.ts"],"sourcesContent":["/**\n * Copyright 2024 Google LLC\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\nimport bodyParser from 'body-parser';\nimport cors, { type CorsOptions } from 'cors';\nimport express from 'express';\nimport { type Action, type ActionContext, type Flow, type z } from 'genkit';\nimport {\n  getCallableJSON,\n  getHttpStatus,\n  type ContextProvider,\n  type RequestData,\n} from 'genkit/context';\nimport { logger } from 'genkit/logging';\nimport type { Server } from 'http';\n\nconst streamDelimiter = '\\n\\n';\n\n/**\n * Exposes provided flow or an action as express handler.\n */\nexport function expressHandler<\n  C extends ActionContext = ActionContext,\n  I extends z.ZodTypeAny = z.ZodTypeAny,\n  O extends z.ZodTypeAny = z.ZodTypeAny,\n  S extends z.ZodTypeAny = z.ZodTypeAny,\n>(\n  action: Action<I, O, S>,\n  opts?: {\n    contextProvider?: ContextProvider<C, I>;\n  }\n): express.RequestHandler {\n  return async (\n    request: express.Request,\n    response: express.Response\n  ): Promise<void> => {\n    const { stream } = request.query;\n    if (!request.body) {\n      const errMsg =\n        `Error: request.body is undefined. ` +\n        `Possible reasons: missing 'content-type: application/json' in request ` +\n        `headers or misconfigured JSON middleware ('app.use(express.json()')? `;\n      logger.error(errMsg);\n      response\n        .status(400)\n        .json({ message: errMsg, status: 'INVALID ARGUMENT' })\n        .end();\n      return;\n    }\n\n    const input = request.body.data as z.infer<I>;\n    let context: Record<string, any>;\n\n    try {\n      context =\n        (await opts?.contextProvider?.({\n          method: request.method as RequestData['method'],\n          headers: Object.fromEntries(\n            Object.entries(request.headers).map(([key, value]) => [\n              key.toLowerCase(),\n              Array.isArray(value) ? value.join(' ') : String(value),\n            ])\n          ),\n          input,\n        })) || {};\n    } catch (e: any) {\n      logger.error(\n        `Auth policy failed with error: ${(e as Error).message}\\n${(e as Error).stack}`\n      );\n      response.status(getHttpStatus(e)).json(getCallableJSON(e)).end();\n      return;\n    }\n\n    const abortController = new AbortController();\n    request.on('close', () => {\n      abortController.abort();\n    });\n    // when/if using timeout middleware, it will emit 'timeout' event.\n    request.on('timeout', () => {\n      abortController.abort();\n    });\n\n    if (request.get('Accept') === 'text/event-stream' || stream === 'true') {\n      response.writeHead(200, {\n        'Content-Type': 'text/plain',\n        'Transfer-Encoding': 'chunked',\n      });\n      try {\n        const onChunk = (chunk: z.infer<S>) => {\n          response.write(\n            'data: ' + JSON.stringify({ message: chunk }) + streamDelimiter\n          );\n        };\n        const result = await action.run(input, {\n          onChunk,\n          context,\n          abortSignal: abortController.signal,\n        });\n        response.write(\n          'data: ' + JSON.stringify({ result: result.result }) + streamDelimiter\n        );\n        response.end();\n      } catch (e) {\n        logger.error(\n          `Streaming request failed with error: ${(e as Error).message}\\n${(e as Error).stack}`\n        );\n        response.write(\n          `error: ${JSON.stringify({ error: getCallableJSON(e) })}${streamDelimiter}`\n        );\n        response.end();\n      }\n    } else {\n      try {\n        const result = await action.run(input, {\n          context,\n          abortSignal: abortController.signal,\n        });\n        response.setHeader('x-genkit-trace-id', result.telemetry.traceId);\n        response.setHeader('x-genkit-span-id', result.telemetry.spanId);\n        // Responses for non-streaming flows are passed back with the flow result stored in a field called \"result.\"\n        response\n          .status(200)\n          .json({\n            result: result.result,\n          })\n          .end();\n      } catch (e) {\n        // Errors for non-streaming flows are passed back as standard API errors.\n        logger.error(\n          `Non-streaming request failed with error: ${(e as Error).message}\\n${(e as Error).stack}`\n        );\n        response.status(getHttpStatus(e)).json(getCallableJSON(e)).end();\n      }\n    }\n  };\n}\n\n/**\n * A wrapper object containing a flow with its associated auth policy.\n */\nexport type FlowWithContextProvider<\n  C extends ActionContext = ActionContext,\n  I extends z.ZodTypeAny = z.ZodTypeAny,\n  O extends z.ZodTypeAny = z.ZodTypeAny,\n  S extends z.ZodTypeAny = z.ZodTypeAny,\n> = {\n  flow: Flow<I, O, S>;\n  context: ContextProvider<C, I>;\n};\n\n/**\n * Adds an auth policy to the flow.\n */\nexport function withContextProvider<\n  C extends ActionContext = ActionContext,\n  I extends z.ZodTypeAny = z.ZodTypeAny,\n  O extends z.ZodTypeAny = z.ZodTypeAny,\n  S extends z.ZodTypeAny = z.ZodTypeAny,\n>(\n  flow: Flow<I, O, S>,\n  context: ContextProvider<C, I>\n): FlowWithContextProvider<C, I, O, S> {\n  return {\n    flow,\n    context,\n  };\n}\n\n/**\n * Options to configure the flow server.\n */\nexport interface FlowServerOptions {\n  /** List of flows to expose via the flow server. */\n  flows: (Flow<any, any, any> | FlowWithContextProvider<any, any, any>)[];\n  /** Port to run the server on. Defaults to env.PORT or 3400. */\n  port?: number;\n  /** CORS options for the server. */\n  cors?: CorsOptions;\n  /** HTTP method path prefix for the exposed flows. */\n  pathPrefix?: string;\n  /** JSON body parser options. */\n  jsonParserOptions?: bodyParser.OptionsJson;\n}\n\n/**\n * Starts an express server with the provided flows and options.\n */\nexport function startFlowServer(options: FlowServerOptions): FlowServer {\n  const server = new FlowServer(options);\n  server.start();\n  return server;\n}\n\n/**\n * Flow server exposes registered flows as HTTP endpoints.\n *\n * This is for use in production environments.\n *\n * @hidden\n */\nexport class FlowServer {\n  /** List of all running servers needed to be cleaned up on process exit. */\n  private static RUNNING_SERVERS: FlowServer[] = [];\n\n  /** Options for the flow server configured by the developer. */\n  private options: FlowServerOptions;\n  /** Port the server is actually running on. This may differ from `options.port` if the original was occupied. Null is server is not running. */\n  private port: number | null = null;\n  /** Express server instance. Null if server is not running. */\n  private server: Server | null = null;\n\n  constructor(options: FlowServerOptions) {\n    this.options = {\n      ...options,\n    };\n  }\n\n  /**\n   * Starts the server and adds it to the list of running servers to clean up on exit.\n   */\n  async start() {\n    const server = express();\n\n    server.use(bodyParser.json(this.options.jsonParserOptions));\n    server.use(cors(this.options.cors));\n\n    logger.debug('Running flow server with flow paths:');\n    const pathPrefix = this.options.pathPrefix ?? '';\n    this.options.flows?.forEach((flow) => {\n      if ('context' in flow) {\n        const flowPath = `/${pathPrefix}${flow.flow.__action.name}`;\n        logger.debug(` - ${flowPath}`);\n        server.post(\n          flowPath,\n          expressHandler(flow.flow, { contextProvider: flow.context })\n        );\n      } else {\n        const flowPath = `/${pathPrefix}${flow.__action.name}`;\n        logger.debug(` - ${flowPath}`);\n        server.post(flowPath, expressHandler(flow));\n      }\n    });\n    this.port =\n      this.options?.port ||\n      (process.env.PORT ? Number.parseInt(process.env.PORT) : 0) ||\n      3400;\n    this.server = server.listen(this.port, () => {\n      logger.debug(`Flow server running on http://localhost:${this.port}`);\n      FlowServer.RUNNING_SERVERS.push(this);\n    });\n  }\n\n  /**\n   * Stops the server and removes it from the list of running servers to clean up on exit.\n   */\n  async stop(): Promise<void> {\n    if (!this.server) {\n      return;\n    }\n    return new Promise<void>((resolve, reject) => {\n      this.server!.close((err) => {\n        if (err) {\n          logger.error(\n            `Error shutting down flow server on port ${this.port}: ${err}`\n          );\n          reject(err);\n        }\n        const index = FlowServer.RUNNING_SERVERS.indexOf(this);\n        if (index > -1) {\n          FlowServer.RUNNING_SERVERS.splice(index, 1);\n        }\n        logger.debug(\n          `Flow server on port ${this.port} has successfully shut down.`\n        );\n        this.port = null;\n        this.server = null;\n        resolve();\n      });\n    });\n  }\n\n  /**\n   * Stops all running servers.\n   */\n  static async stopAll() {\n    return Promise.all(\n      FlowServer.RUNNING_SERVERS.map((server) => server.stop())\n    );\n  }\n}\n"],"mappings":";;;;;;;;;;;;;;;;;;;;;;;;;;;;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAgBA,yBAAuB;AACvB,kBAAuC;AACvC,qBAAoB;AAEpB,qBAKO;AACP,qBAAuB;AAGvB,MAAM,kBAAkB;AAKjB,SAAS,eAMd,QACA,MAGwB;AACxB,SAAO,OACL,SACA,aACkB;AAClB,UAAM,EAAE,OAAO,IAAI,QAAQ;AAC3B,QAAI,CAAC,QAAQ,MAAM;AACjB,YAAM,SACJ;AAGF,4BAAO,MAAM,MAAM;AACnB,eACG,OAAO,GAAG,EACV,KAAK,EAAE,SAAS,QAAQ,QAAQ,mBAAmB,CAAC,EACpD,IAAI;AACP;AAAA,IACF;AAEA,UAAM,QAAQ,QAAQ,KAAK;AAC3B,QAAI;AAEJ,QAAI;AACF,gBACG,MAAM,MAAM,kBAAkB;AAAA,QAC7B,QAAQ,QAAQ;AAAA,QAChB,SAAS,OAAO;AAAA,UACd,OAAO,QAAQ,QAAQ,OAAO,EAAE,IAAI,CAAC,CAAC,KAAK,KAAK,MAAM;AAAA,YACpD,IAAI,YAAY;AAAA,YAChB,MAAM,QAAQ,KAAK,IAAI,MAAM,KAAK,GAAG,IAAI,OAAO,KAAK;AAAA,UACvD,CAAC;AAAA,QACH;AAAA,QACA;AAAA,MACF,CAAC,KAAM,CAAC;AAAA,IACZ,SAAS,GAAQ;AACf,4BAAO;AAAA,QACL,kCAAmC,EAAY,OAAO;AAAA,EAAM,EAAY,KAAK;AAAA,MAC/E;AACA,eAAS,WAAO,8BAAc,CAAC,CAAC,EAAE,SAAK,gCAAgB,CAAC,CAAC,EAAE,IAAI;AAC/D;AAAA,IACF;AAEA,UAAM,kBAAkB,IAAI,gBAAgB;AAC5C,YAAQ,GAAG,SAAS,MAAM;AACxB,sBAAgB,MAAM;AAAA,IACxB,CAAC;AAED,YAAQ,GAAG,WAAW,MAAM;AAC1B,sBAAgB,MAAM;AAAA,IACxB,CAAC;AAED,QAAI,QAAQ,IAAI,QAAQ,MAAM,uBAAuB,WAAW,QAAQ;AACtE,eAAS,UAAU,KAAK;AAAA,QACtB,gBAAgB;AAAA,QAChB,qBAAqB;AAAA,MACvB,CAAC;AACD,UAAI;AACF,cAAM,UAAU,CAAC,UAAsB;AACrC,mBAAS;AAAA,YACP,WAAW,KAAK,UAAU,EAAE,SAAS,MAAM,CAAC,IAAI;AAAA,UAClD;AAAA,QACF;AACA,cAAM,SAAS,MAAM,OAAO,IAAI,OAAO;AAAA,UACrC;AAAA,UACA;AAAA,UACA,aAAa,gBAAgB;AAAA,QAC/B,CAAC;AACD,iBAAS;AAAA,UACP,WAAW,KAAK,UAAU,EAAE,QAAQ,OAAO,OAAO,CAAC,IAAI;AAAA,QACzD;AACA,iBAAS,IAAI;AAAA,MACf,SAAS,GAAG;AACV,8BAAO;AAAA,UACL,wCAAyC,EAAY,OAAO;AAAA,EAAM,EAAY,KAAK;AAAA,QACrF;AACA,iBAAS;AAAA,UACP,UAAU,KAAK,UAAU,EAAE,WAAO,gCAAgB,CAAC,EAAE,CAAC,CAAC,GAAG,eAAe;AAAA,QAC3E;AACA,iBAAS,IAAI;AAAA,MACf;AAAA,IACF,OAAO;AACL,UAAI;AACF,cAAM,SAAS,MAAM,OAAO,IAAI,OAAO;AAAA,UACrC;AAAA,UACA,aAAa,gBAAgB;AAAA,QAC/B,CAAC;AACD,iBAAS,UAAU,qBAAqB,OAAO,UAAU,OAAO;AAChE,iBAAS,UAAU,oBAAoB,OAAO,UAAU,MAAM;AAE9D,iBACG,OAAO,GAAG,EACV,KAAK;AAAA,UACJ,QAAQ,OAAO;AAAA,QACjB,CAAC,EACA,IAAI;AAAA,MACT,SAAS,GAAG;AAEV,8BAAO;AAAA,UACL,4CAA6C,EAAY,OAAO;AAAA,EAAM,EAAY,KAAK;AAAA,QACzF;AACA,iBAAS,WAAO,8BAAc,CAAC,CAAC,EAAE,SAAK,gCAAgB,CAAC,CAAC,EAAE,IAAI;AAAA,MACjE;AAAA,IACF;AAAA,EACF;AACF;AAkBO,SAAS,oBAMd,MACA,SACqC;AACrC,SAAO;AAAA,IACL;AAAA,IACA;AAAA,EACF;AACF;AAqBO,SAAS,gBAAgB,SAAwC;AACtE,QAAM,SAAS,IAAI,WAAW,OAAO;AACrC,SAAO,MAAM;AACb,SAAO;AACT;AASO,MAAM,WAAW;AAAA;AAAA,EAEtB,OAAe,kBAAgC,CAAC;AAAA;AAAA,EAGxC;AAAA;AAAA,EAEA,OAAsB;AAAA;AAAA,EAEtB,SAAwB;AAAA,EAEhC,YAAY,SAA4B;AACtC,SAAK,UAAU;AAAA,MACb,GAAG;AAAA,IACL;AAAA,EACF;AAAA;AAAA;AAAA;AAAA,EAKA,MAAM,QAAQ;AACZ,UAAM,aAAS,eAAAA,SAAQ;AAEvB,WAAO,IAAI,mBAAAC,QAAW,KAAK,KAAK,QAAQ,iBAAiB,CAAC;AAC1D,WAAO,QAAI,YAAAC,SAAK,KAAK,QAAQ,IAAI,CAAC;AAElC,0BAAO,MAAM,sCAAsC;AACnD,UAAM,aAAa,KAAK,QAAQ,cAAc;AAC9C,SAAK,QAAQ,OAAO,QAAQ,CAAC,SAAS;AACpC,UAAI,aAAa,MAAM;AACrB,cAAM,WAAW,IAAI,UAAU,GAAG,KAAK,KAAK,SAAS,IAAI;AACzD,8BAAO,MAAM,MAAM,QAAQ,EAAE;AAC7B,eAAO;AAAA,UACL;AAAA,UACA,eAAe,KAAK,MAAM,EAAE,iBAAiB,KAAK,QAAQ,CAAC;AAAA,QAC7D;AAAA,MACF,OAAO;AACL,cAAM,WAAW,IAAI,UAAU,GAAG,KAAK,SAAS,IAAI;AACpD,8BAAO,MAAM,MAAM,QAAQ,EAAE;AAC7B,eAAO,KAAK,UAAU,eAAe,IAAI,CAAC;AAAA,MAC5C;AAAA,IACF,CAAC;AACD,SAAK,OACH,KAAK,SAAS,SACb,QAAQ,IAAI,OAAO,OAAO,SAAS,QAAQ,IAAI,IAAI,IAAI,MACxD;AACF,SAAK,SAAS,OAAO,OAAO,KAAK,MAAM,MAAM;AAC3C,4BAAO,MAAM,2CAA2C,KAAK,IAAI,EAAE;AACnE,iBAAW,gBAAgB,KAAK,IAAI;AAAA,IACtC,CAAC;AAAA,EACH;AAAA;AAAA;AAAA;AAAA,EAKA,MAAM,OAAsB;AAC1B,QAAI,CAAC,KAAK,QAAQ;AAChB;AAAA,IACF;AACA,WAAO,IAAI,QAAc,CAAC,SAAS,WAAW;AAC5C,WAAK,OAAQ,MAAM,CAAC,QAAQ;AAC1B,YAAI,KAAK;AACP,gCAAO;AAAA,YACL,2CAA2C,KAAK,IAAI,KAAK,GAAG;AAAA,UAC9D;AACA,iBAAO,GAAG;AAAA,QACZ;AACA,cAAM,QAAQ,WAAW,gBAAgB,QAAQ,IAAI;AACrD,YAAI,QAAQ,IAAI;AACd,qBAAW,gBAAgB,OAAO,OAAO,CAAC;AAAA,QAC5C;AACA,8BAAO;AAAA,UACL,uBAAuB,KAAK,IAAI;AAAA,QAClC;AACA,aAAK,OAAO;AACZ,aAAK,SAAS;AACd,gBAAQ;AAAA,MACV,CAAC;AAAA,IACH,CAAC;AAAA,EACH;AAAA;AAAA;AAAA;AAAA,EAKA,aAAa,UAAU;AACrB,WAAO,QAAQ;AAAA,MACb,WAAW,gBAAgB,IAAI,CAAC,WAAW,OAAO,KAAK,CAAC;AAAA,IAC1D;AAAA,EACF;AACF;","names":["express","bodyParser","cors"]}
+{"version":3,"sources":["../src/index.ts"],"sourcesContent":["/**\n * Copyright 2024 Google LLC\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\nimport bodyParser from 'body-parser';\nimport cors, { type CorsOptions } from 'cors';\nimport { randomUUID } from 'crypto';\nimport express from 'express';\nimport {\n  Action,\n  ActionStreamInput,\n  AsyncTaskQueue,\n  Flow,\n  StreamNotFoundError,\n  type ActionContext,\n  type StreamManager,\n  type z,\n} from 'genkit/beta';\nimport {\n  getCallableJSON,\n  getHttpStatus,\n  type ContextProvider,\n  type RequestData,\n} from 'genkit/context';\nimport { logger } from 'genkit/logging';\nimport type { Server } from 'http';\n\nconst streamDelimiter = '\\n\\n';\n\n/**\n * Exposes provided flow or an action as express handler.\n */\nexport function expressHandler<\n  C extends ActionContext = ActionContext,\n  I extends z.ZodTypeAny = z.ZodTypeAny,\n  O extends z.ZodTypeAny = z.ZodTypeAny,\n  S extends z.ZodTypeAny = z.ZodTypeAny,\n>(\n  action: Action<I, O, S>,\n  opts?: {\n    contextProvider?: ContextProvider<C, I>;\n    streamManager?: StreamManager;\n  }\n): express.RequestHandler {\n  return async (\n    request: express.Request,\n    response: express.Response\n  ): Promise<void> => {\n    const { stream } = request.query;\n    const streamIdHeader = request.headers['x-genkit-stream-id'];\n    const streamId = Array.isArray(streamIdHeader)\n      ? streamIdHeader[0]\n      : streamIdHeader;\n\n    if (!request.body) {\n      const errMsg =\n        `Error: request.body is undefined. ` +\n        `Possible reasons: missing 'content-type: application/json' in request ` +\n        `headers or misconfigured JSON middleware ('app.use(express.json()')? `;\n      logger.error(errMsg);\n      response\n        .status(400)\n        .json({ message: errMsg, status: 'INVALID ARGUMENT' })\n        .end();\n      return;\n    }\n\n    const input = request.body.data as z.infer<I>;\n    let context: Record<string, any>;\n\n    try {\n      context =\n        (await opts?.contextProvider?.({\n          method: request.method as RequestData['method'],\n          headers: Object.fromEntries(\n            Object.entries(request.headers).map(([key, value]) => [\n              key.toLowerCase(),\n              Array.isArray(value) ? value.join(' ') : String(value),\n            ])\n          ),\n          input,\n        })) || {};\n    } catch (e: any) {\n      logger.error(\n        `Auth policy failed with error: ${(e as Error).message}\\n${(e as Error).stack}`\n      );\n      response.status(getHttpStatus(e)).json(getCallableJSON(e)).end();\n      return;\n    }\n\n    const abortController = new AbortController();\n    request.on('close', () => {\n      abortController.abort();\n    });\n    // when/if using timeout middleware, it will emit 'timeout' event.\n    request.on('timeout', () => {\n      abortController.abort();\n    });\n\n    if (request.get('Accept') === 'text/event-stream' || stream === 'true') {\n      const streamManager = opts?.streamManager;\n      if (streamManager && streamId) {\n        await subscribeToStream(streamManager, streamId, response);\n        return;\n      }\n      const streamIdToUse = randomUUID();\n      const headers = {\n        'Content-Type': 'text/plain',\n        'Transfer-Encoding': 'chunked',\n      };\n      if (streamManager) {\n        headers['x-genkit-stream-id'] = streamIdToUse;\n      }\n      response.writeHead(200, headers);\n      runActionWithDurableStreaming(\n        action,\n        streamManager,\n        streamIdToUse,\n        input,\n        context,\n        response,\n        abortController.signal\n      );\n    } else {\n      try {\n        const result = await action.run(input, {\n          context,\n          abortSignal: abortController.signal,\n        });\n        response.setHeader('x-genkit-trace-id', result.telemetry.traceId);\n        response.setHeader('x-genkit-span-id', result.telemetry.spanId);\n        // Responses for non-streaming flows are passed back with the flow result stored in a field called \"result.\"\n        response\n          .status(200)\n          .json({\n            result: result.result,\n          })\n          .end();\n      } catch (e) {\n        // Errors for non-streaming flows are passed back as standard API errors.\n        logger.error(\n          `Non-streaming request failed with error: ${(e as Error).message}\\n${(e as Error).stack}`\n        );\n        response.status(getHttpStatus(e)).json(getCallableJSON(e)).end();\n      }\n    }\n  };\n}\n\nasync function runActionWithDurableStreaming<\n  I extends z.ZodTypeAny,\n  O extends z.ZodTypeAny,\n  S extends z.ZodTypeAny,\n>(\n  action: Action<I, O, S>,\n  streamManager: StreamManager | undefined,\n  streamId: string,\n  input: z.infer<I>,\n  context: ActionContext,\n  response: express.Response,\n  abortSignal: AbortSignal\n) {\n  let taskQueue: AsyncTaskQueue | undefined;\n  let durableStream: ActionStreamInput<any, any> | undefined;\n  if (streamManager) {\n    taskQueue = new AsyncTaskQueue();\n    durableStream = await streamManager.open(streamId);\n  }\n  try {\n    let onChunk = (chunk: z.infer<S>) => {\n      response.write(\n        'data: ' + JSON.stringify({ message: chunk }) + streamDelimiter\n      );\n    };\n    if (streamManager) {\n      const originalOnChunk = onChunk;\n      onChunk = (chunk: z.infer<S>) => {\n        originalOnChunk(chunk);\n        taskQueue!.enqueue(() => durableStream!.write(chunk));\n      };\n    }\n    const result = await action.run(input, {\n      onChunk,\n      context,\n      abortSignal,\n    });\n    if (streamManager) {\n      taskQueue!.enqueue(() => durableStream!.done(result.result));\n      await taskQueue!.merge();\n    }\n    response.write(\n      'data: ' + JSON.stringify({ result: result.result }) + streamDelimiter\n    );\n    response.end();\n  } catch (e) {\n    if (durableStream) {\n      taskQueue!.enqueue(() => durableStream!.error(e));\n      await taskQueue!.merge();\n    }\n    logger.error(\n      `Streaming request failed with error: ${(e as Error).message}\\n${\n        (e as Error).stack\n      }`\n    );\n    response.write(\n      `error: ${JSON.stringify({\n        error: getCallableJSON(e),\n      })}${streamDelimiter}`\n    );\n    response.end();\n  }\n}\n\nasync function subscribeToStream(\n  streamManager: StreamManager,\n  streamId: string,\n  response: express.Response\n): Promise<void> {\n  try {\n    await streamManager.subscribe(streamId, {\n      onChunk: (chunk) => {\n        response.write(\n          'data: ' + JSON.stringify({ message: chunk }) + streamDelimiter\n        );\n      },\n      onDone: (output) => {\n        response.write(\n          'data: ' + JSON.stringify({ result: output }) + streamDelimiter\n        );\n        response.end();\n      },\n      onError: (err) => {\n        logger.error(\n          `Streaming request failed with error: ${(err as Error).message}\\n${\n            (err as Error).stack\n          }`\n        );\n        response.write(\n          `error: ${JSON.stringify({\n            error: getCallableJSON(err),\n          })}${streamDelimiter}`\n        );\n        response.end();\n      },\n    });\n  } catch (e: any) {\n    if (e instanceof StreamNotFoundError) {\n      response.status(204).end();\n      return;\n    }\n    if (e.status === 'DEADLINE_EXCEEDED') {\n      response.write(\n        `error: ${JSON.stringify({\n          error: getCallableJSON(e),\n        })}${streamDelimiter}`\n      );\n      response.end();\n      return;\n    }\n    throw e;\n  }\n}\n\n/**\n * A wrapper object containing a flow with its associated auth policy.\n * @deprecated Use `withFlowOptions` instead.\n */\nexport type FlowWithContextProvider<\n  C extends ActionContext = ActionContext,\n  I extends z.ZodTypeAny = z.ZodTypeAny,\n  O extends z.ZodTypeAny = z.ZodTypeAny,\n  S extends z.ZodTypeAny = z.ZodTypeAny,\n> = {\n  flow: Flow<I, O, S>;\n  context: ContextProvider<C, I>;\n};\n\n/**\n * A wrapper object containing a flow with its associated options.\n */\nexport type FlowWithOptions<\n  I extends z.ZodTypeAny = z.ZodTypeAny,\n  O extends z.ZodTypeAny = z.ZodTypeAny,\n  S extends z.ZodTypeAny = z.ZodTypeAny,\n> = {\n  flow: Flow<I, O, S>;\n  options: {\n    contextProvider?: ContextProvider<any, I>;\n    streamManager?: StreamManager;\n    path?: string;\n  };\n};\n\n/**\n * Adds an auth policy to the flow.\n * @deprecated Use `withFlowOptions` instead.\n */\nexport function withContextProvider<\n  C extends ActionContext = ActionContext,\n  I extends z.ZodTypeAny = z.ZodTypeAny,\n  O extends z.ZodTypeAny = z.ZodTypeAny,\n  S extends z.ZodTypeAny = z.ZodTypeAny,\n>(\n  flow: Flow<I, O, S>,\n  context: ContextProvider<C, I>\n): FlowWithContextProvider<C, I, O, S> {\n  return {\n    flow,\n    context,\n  };\n}\n\n/**\n * Adds an auth policy to the flow.\n */\nexport function withFlowOptions<\n  I extends z.ZodTypeAny,\n  O extends z.ZodTypeAny,\n  S extends z.ZodTypeAny,\n>(\n  flow: Flow<I, O, S>,\n  options: {\n    contextProvider?: ContextProvider<any, I>;\n    streamManager?: StreamManager;\n    path?: string;\n  }\n): FlowWithOptions<I, O, S> {\n  return {\n    flow,\n    options,\n  };\n}\n\n/**\n * Options to configure the flow server.\n */\nexport interface FlowServerOptions {\n  /** List of flows to expose via the flow server. */\n  flows: (\n    | Flow<any, any, any>\n    | FlowWithContextProvider<any, any, any>\n    | FlowWithOptions<any, any, any>\n  )[];\n  /** Port to run the server on. Defaults to env.PORT or 3400. */\n  port?: number;\n  /** CORS options for the server. */\n  cors?: CorsOptions;\n  /** HTTP method path prefix for the exposed flows. */\n  pathPrefix?: string;\n  /** JSON body parser options. */\n  jsonParserOptions?: bodyParser.OptionsJson;\n}\n\n/**\n * Starts an express server with the provided flows and options.\n */\nexport function startFlowServer(options: FlowServerOptions): FlowServer {\n  const server = new FlowServer(options);\n  server.start();\n  return server;\n}\n\n/**\n * Flow server exposes registered flows as HTTP endpoints.\n *\n * This is for use in production environments.\n *\n * @hidden\n */\nexport class FlowServer {\n  /** List of all running servers needed to be cleaned up on process exit. */\n  private static RUNNING_SERVERS: FlowServer[] = [];\n\n  /** Options for the flow server configured by the developer. */\n  private options: FlowServerOptions;\n  /** Port the server is actually running on. This may differ from `options.port` if the original was occupied. Null is server is not running. */\n  private port: number | null = null;\n  /** Express server instance. Null if server is not running. */\n  private server: Server | null = null;\n\n  constructor(options: FlowServerOptions) {\n    this.options = {\n      ...options,\n    };\n  }\n\n  /**\n   * Starts the server and adds it to the list of running servers to clean up on exit.\n   */\n  async start() {\n    const server = express();\n\n    server.use(bodyParser.json(this.options.jsonParserOptions));\n    server.use(cors(this.options.cors));\n\n    logger.debug('Running flow server with flow paths:');\n    const pathPrefix = this.options.pathPrefix ?? '';\n    this.options.flows?.forEach((flow) => {\n      if ('flow' in flow) {\n        const flowPath = `/${pathPrefix}${\n          ('options' in flow && flow.options.path) || flow.flow.__action.name\n        }`;\n        logger.debug(` - ${flowPath}`);\n        const options =\n          'options' in flow ? flow.options : { contextProvider: flow.context };\n        server.post(flowPath, expressHandler(flow.flow, options));\n      } else {\n        const flowPath = `/${pathPrefix}${flow.__action.name}`;\n        logger.debug(` - ${flowPath}`);\n        server.post(flowPath, expressHandler(flow));\n      }\n    });\n    this.port =\n      this.options?.port ||\n      (process.env.PORT ? Number.parseInt(process.env.PORT) : 0) ||\n      3400;\n    this.server = server.listen(this.port, () => {\n      logger.debug(`Flow server running on http://localhost:${this.port}`);\n      FlowServer.RUNNING_SERVERS.push(this);\n    });\n  }\n\n  /**\n   * Stops the server and removes it from the list of running servers to clean up on exit.\n   */\n  async stop(): Promise<void> {\n    if (!this.server) {\n      return;\n    }\n    return new Promise<void>((resolve, reject) => {\n      this.server!.close((err) => {\n        if (err) {\n          logger.error(\n            `Error shutting down flow server on port ${this.port}: ${err}`\n          );\n          reject(err);\n        }\n        const index = FlowServer.RUNNING_SERVERS.indexOf(this);\n        if (index > -1) {\n          FlowServer.RUNNING_SERVERS.splice(index, 1);\n        }\n        logger.debug(\n          `Flow server on port ${this.port} has successfully shut down.`\n        );\n        this.port = null;\n        this.server = null;\n        resolve();\n      });\n    });\n  }\n\n  /**\n   * Stops all running servers.\n   */\n  static async stopAll() {\n    return Promise.all(\n      FlowServer.RUNNING_SERVERS.map((server) => server.stop())\n    );\n  }\n}\n"],"mappings":";;;;;;;;;;;;;;;;;;;;;;;;;;;;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAgBA,yBAAuB;AACvB,kBAAuC;AACvC,oBAA2B;AAC3B,qBAAoB;AACpB,kBASO;AACP,qBAKO;AACP,qBAAuB;AAGvB,MAAM,kBAAkB;AAKjB,SAAS,eAMd,QACA,MAIwB;AACxB,SAAO,OACL,SACA,aACkB;AAClB,UAAM,EAAE,OAAO,IAAI,QAAQ;AAC3B,UAAM,iBAAiB,QAAQ,QAAQ,oBAAoB;AAC3D,UAAM,WAAW,MAAM,QAAQ,cAAc,IACzC,eAAe,CAAC,IAChB;AAEJ,QAAI,CAAC,QAAQ,MAAM;AACjB,YAAM,SACJ;AAGF,4BAAO,MAAM,MAAM;AACnB,eACG,OAAO,GAAG,EACV,KAAK,EAAE,SAAS,QAAQ,QAAQ,mBAAmB,CAAC,EACpD,IAAI;AACP;AAAA,IACF;AAEA,UAAM,QAAQ,QAAQ,KAAK;AAC3B,QAAI;AAEJ,QAAI;AACF,gBACG,MAAM,MAAM,kBAAkB;AAAA,QAC7B,QAAQ,QAAQ;AAAA,QAChB,SAAS,OAAO;AAAA,UACd,OAAO,QAAQ,QAAQ,OAAO,EAAE,IAAI,CAAC,CAAC,KAAK,KAAK,MAAM;AAAA,YACpD,IAAI,YAAY;AAAA,YAChB,MAAM,QAAQ,KAAK,IAAI,MAAM,KAAK,GAAG,IAAI,OAAO,KAAK;AAAA,UACvD,CAAC;AAAA,QACH;AAAA,QACA;AAAA,MACF,CAAC,KAAM,CAAC;AAAA,IACZ,SAAS,GAAQ;AACf,4BAAO;AAAA,QACL,kCAAmC,EAAY,OAAO;AAAA,EAAM,EAAY,KAAK;AAAA,MAC/E;AACA,eAAS,WAAO,8BAAc,CAAC,CAAC,EAAE,SAAK,gCAAgB,CAAC,CAAC,EAAE,IAAI;AAC/D;AAAA,IACF;AAEA,UAAM,kBAAkB,IAAI,gBAAgB;AAC5C,YAAQ,GAAG,SAAS,MAAM;AACxB,sBAAgB,MAAM;AAAA,IACxB,CAAC;AAED,YAAQ,GAAG,WAAW,MAAM;AAC1B,sBAAgB,MAAM;AAAA,IACxB,CAAC;AAED,QAAI,QAAQ,IAAI,QAAQ,MAAM,uBAAuB,WAAW,QAAQ;AACtE,YAAM,gBAAgB,MAAM;AAC5B,UAAI,iBAAiB,UAAU;AAC7B,cAAM,kBAAkB,eAAe,UAAU,QAAQ;AACzD;AAAA,MACF;AACA,YAAM,oBAAgB,0BAAW;AACjC,YAAM,UAAU;AAAA,QACd,gBAAgB;AAAA,QAChB,qBAAqB;AAAA,MACvB;AACA,UAAI,eAAe;AACjB,gBAAQ,oBAAoB,IAAI;AAAA,MAClC;AACA,eAAS,UAAU,KAAK,OAAO;AAC/B;AAAA,QACE;AAAA,QACA;AAAA,QACA;AAAA,QACA;AAAA,QACA;AAAA,QACA;AAAA,QACA,gBAAgB;AAAA,MAClB;AAAA,IACF,OAAO;AACL,UAAI;AACF,cAAM,SAAS,MAAM,OAAO,IAAI,OAAO;AAAA,UACrC;AAAA,UACA,aAAa,gBAAgB;AAAA,QAC/B,CAAC;AACD,iBAAS,UAAU,qBAAqB,OAAO,UAAU,OAAO;AAChE,iBAAS,UAAU,oBAAoB,OAAO,UAAU,MAAM;AAE9D,iBACG,OAAO,GAAG,EACV,KAAK;AAAA,UACJ,QAAQ,OAAO;AAAA,QACjB,CAAC,EACA,IAAI;AAAA,MACT,SAAS,GAAG;AAEV,8BAAO;AAAA,UACL,4CAA6C,EAAY,OAAO;AAAA,EAAM,EAAY,KAAK;AAAA,QACzF;AACA,iBAAS,WAAO,8BAAc,CAAC,CAAC,EAAE,SAAK,gCAAgB,CAAC,CAAC,EAAE,IAAI;AAAA,MACjE;AAAA,IACF;AAAA,EACF;AACF;AAEA,eAAe,8BAKb,QACA,eACA,UACA,OACA,SACA,UACA,aACA;AACA,MAAI;AACJ,MAAI;AACJ,MAAI,eAAe;AACjB,gBAAY,IAAI,2BAAe;AAC/B,oBAAgB,MAAM,cAAc,KAAK,QAAQ;AAAA,EACnD;AACA,MAAI;AACF,QAAI,UAAU,CAAC,UAAsB;AACnC,eAAS;AAAA,QACP,WAAW,KAAK,UAAU,EAAE,SAAS,MAAM,CAAC,IAAI;AAAA,MAClD;AAAA,IACF;AACA,QAAI,eAAe;AACjB,YAAM,kBAAkB;AACxB,gBAAU,CAAC,UAAsB;AAC/B,wBAAgB,KAAK;AACrB,kBAAW,QAAQ,MAAM,cAAe,MAAM,KAAK,CAAC;AAAA,MACtD;AAAA,IACF;AACA,UAAM,SAAS,MAAM,OAAO,IAAI,OAAO;AAAA,MACrC;AAAA,MACA;AAAA,MACA;AAAA,IACF,CAAC;AACD,QAAI,eAAe;AACjB,gBAAW,QAAQ,MAAM,cAAe,KAAK,OAAO,MAAM,CAAC;AAC3D,YAAM,UAAW,MAAM;AAAA,IACzB;AACA,aAAS;AAAA,MACP,WAAW,KAAK,UAAU,EAAE,QAAQ,OAAO,OAAO,CAAC,IAAI;AAAA,IACzD;AACA,aAAS,IAAI;AAAA,EACf,SAAS,GAAG;AACV,QAAI,eAAe;AACjB,gBAAW,QAAQ,MAAM,cAAe,MAAM,CAAC,CAAC;AAChD,YAAM,UAAW,MAAM;AAAA,IACzB;AACA,0BAAO;AAAA,MACL,wCAAyC,EAAY,OAAO;AAAA,EACzD,EAAY,KACf;AAAA,IACF;AACA,aAAS;AAAA,MACP,UAAU,KAAK,UAAU;AAAA,QACvB,WAAO,gCAAgB,CAAC;AAAA,MAC1B,CAAC,CAAC,GAAG,eAAe;AAAA,IACtB;AACA,aAAS,IAAI;AAAA,EACf;AACF;AAEA,eAAe,kBACb,eACA,UACA,UACe;AACf,MAAI;AACF,UAAM,cAAc,UAAU,UAAU;AAAA,MACtC,SAAS,CAAC,UAAU;AAClB,iBAAS;AAAA,UACP,WAAW,KAAK,UAAU,EAAE,SAAS,MAAM,CAAC,IAAI;AAAA,QAClD;AAAA,MACF;AAAA,MACA,QAAQ,CAAC,WAAW;AAClB,iBAAS;AAAA,UACP,WAAW,KAAK,UAAU,EAAE,QAAQ,OAAO,CAAC,IAAI;AAAA,QAClD;AACA,iBAAS,IAAI;AAAA,MACf;AAAA,MACA,SAAS,CAAC,QAAQ;AAChB,8BAAO;AAAA,UACL,wCAAyC,IAAc,OAAO;AAAA,EAC3D,IAAc,KACjB;AAAA,QACF;AACA,iBAAS;AAAA,UACP,UAAU,KAAK,UAAU;AAAA,YACvB,WAAO,gCAAgB,GAAG;AAAA,UAC5B,CAAC,CAAC,GAAG,eAAe;AAAA,QACtB;AACA,iBAAS,IAAI;AAAA,MACf;AAAA,IACF,CAAC;AAAA,EACH,SAAS,GAAQ;AACf,QAAI,aAAa,iCAAqB;AACpC,eAAS,OAAO,GAAG,EAAE,IAAI;AACzB;AAAA,IACF;AACA,QAAI,EAAE,WAAW,qBAAqB;AACpC,eAAS;AAAA,QACP,UAAU,KAAK,UAAU;AAAA,UACvB,WAAO,gCAAgB,CAAC;AAAA,QAC1B,CAAC,CAAC,GAAG,eAAe;AAAA,MACtB;AACA,eAAS,IAAI;AACb;AAAA,IACF;AACA,UAAM;AAAA,EACR;AACF;AAoCO,SAAS,oBAMd,MACA,SACqC;AACrC,SAAO;AAAA,IACL;AAAA,IACA;AAAA,EACF;AACF;AAKO,SAAS,gBAKd,MACA,SAK0B;AAC1B,SAAO;AAAA,IACL;AAAA,IACA;AAAA,EACF;AACF;AAyBO,SAAS,gBAAgB,SAAwC;AACtE,QAAM,SAAS,IAAI,WAAW,OAAO;AACrC,SAAO,MAAM;AACb,SAAO;AACT;AASO,MAAM,WAAW;AAAA;AAAA,EAEtB,OAAe,kBAAgC,CAAC;AAAA;AAAA,EAGxC;AAAA;AAAA,EAEA,OAAsB;AAAA;AAAA,EAEtB,SAAwB;AAAA,EAEhC,YAAY,SAA4B;AACtC,SAAK,UAAU;AAAA,MACb,GAAG;AAAA,IACL;AAAA,EACF;AAAA;AAAA;AAAA;AAAA,EAKA,MAAM,QAAQ;AACZ,UAAM,aAAS,eAAAA,SAAQ;AAEvB,WAAO,IAAI,mBAAAC,QAAW,KAAK,KAAK,QAAQ,iBAAiB,CAAC;AAC1D,WAAO,QAAI,YAAAC,SAAK,KAAK,QAAQ,IAAI,CAAC;AAElC,0BAAO,MAAM,sCAAsC;AACnD,UAAM,aAAa,KAAK,QAAQ,cAAc;AAC9C,SAAK,QAAQ,OAAO,QAAQ,CAAC,SAAS;AACpC,UAAI,UAAU,MAAM;AAClB,cAAM,WAAW,IAAI,UAAU,GAC5B,aAAa,QAAQ,KAAK,QAAQ,QAAS,KAAK,KAAK,SAAS,IACjE;AACA,8BAAO,MAAM,MAAM,QAAQ,EAAE;AAC7B,cAAM,UACJ,aAAa,OAAO,KAAK,UAAU,EAAE,iBAAiB,KAAK,QAAQ;AACrE,eAAO,KAAK,UAAU,eAAe,KAAK,MAAM,OAAO,CAAC;AAAA,MAC1D,OAAO;AACL,cAAM,WAAW,IAAI,UAAU,GAAG,KAAK,SAAS,IAAI;AACpD,8BAAO,MAAM,MAAM,QAAQ,EAAE;AAC7B,eAAO,KAAK,UAAU,eAAe,IAAI,CAAC;AAAA,MAC5C;AAAA,IACF,CAAC;AACD,SAAK,OACH,KAAK,SAAS,SACb,QAAQ,IAAI,OAAO,OAAO,SAAS,QAAQ,IAAI,IAAI,IAAI,MACxD;AACF,SAAK,SAAS,OAAO,OAAO,KAAK,MAAM,MAAM;AAC3C,4BAAO,MAAM,2CAA2C,KAAK,IAAI,EAAE;AACnE,iBAAW,gBAAgB,KAAK,IAAI;AAAA,IACtC,CAAC;AAAA,EACH;AAAA;AAAA;AAAA;AAAA,EAKA,MAAM,OAAsB;AAC1B,QAAI,CAAC,KAAK,QAAQ;AAChB;AAAA,IACF;AACA,WAAO,IAAI,QAAc,CAAC,SAAS,WAAW;AAC5C,WAAK,OAAQ,MAAM,CAAC,QAAQ;AAC1B,YAAI,KAAK;AACP,gCAAO;AAAA,YACL,2CAA2C,KAAK,IAAI,KAAK,GAAG;AAAA,UAC9D;AACA,iBAAO,GAAG;AAAA,QACZ;AACA,cAAM,QAAQ,WAAW,gBAAgB,QAAQ,IAAI;AACrD,YAAI,QAAQ,IAAI;AACd,qBAAW,gBAAgB,OAAO,OAAO,CAAC;AAAA,QAC5C;AACA,8BAAO;AAAA,UACL,uBAAuB,KAAK,IAAI;AAAA,QAClC;AACA,aAAK,OAAO;AACZ,aAAK,SAAS;AACd,gBAAQ;AAAA,MACV,CAAC;AAAA,IACH,CAAC;AAAA,EACH;AAAA;AAAA;AAAA;AAAA,EAKA,aAAa,UAAU;AACrB,WAAO,QAAQ;AAAA,MACb,WAAW,gBAAgB,IAAI,CAAC,WAAW,OAAO,KAAK,CAAC;AAAA,IAC1D;AAAA,EACF;AACF;","names":["express","bodyParser","cors"]}

package/lib/index.mjs

@@ -1,6 +1,11 @@
 import bodyParser from "body-parser";
 import cors from "cors";
+import { randomUUID } from "crypto";
 import express from "express";
+import {
+  AsyncTaskQueue,
+  StreamNotFoundError
+} from "genkit/beta";
 import {
   getCallableJSON,
   getHttpStatus
@@ -10,6 +15,8 @@ const streamDelimiter = "\n\n";
 function expressHandler(action, opts) {
   return async (request, response) => {
     const { stream } = request.query;
+    const streamIdHeader = request.headers["x-genkit-stream-id"];
+    const streamId = Array.isArray(streamIdHeader) ? streamIdHeader[0] : streamIdHeader;
     if (!request.body) {
       const errMsg = `Error: request.body is undefined. Possible reasons: missing 'content-type: application/json' in request headers or misconfigured JSON middleware ('app.use(express.json()')? `;
       logger.error(errMsg);
@@ -45,35 +52,29 @@ ${e.stack}`
       abortController.abort();
     });
     if (request.get("Accept") === "text/event-stream" || stream === "true") {
-      response.writeHead(200, {
+      const streamManager = opts?.streamManager;
+      if (streamManager && streamId) {
+        await subscribeToStream(streamManager, streamId, response);
+        return;
+      }
+      const streamIdToUse = randomUUID();
+      const headers = {
         "Content-Type": "text/plain",
         "Transfer-Encoding": "chunked"
-      });
-      try {
-        const onChunk = (chunk) => {
-          response.write(
-            "data: " + JSON.stringify({ message: chunk }) + streamDelimiter
-          );
-        };
-        const result = await action.run(input, {
-          onChunk,
-          context,
-          abortSignal: abortController.signal
-        });
-        response.write(
-          "data: " + JSON.stringify({ result: result.result }) + streamDelimiter
-        );
-        response.end();
-      } catch (e) {
-        logger.error(
-          `Streaming request failed with error: ${e.message}
-${e.stack}`
-        );
-        response.write(
-          `error: ${JSON.stringify({ error: getCallableJSON(e) })}${streamDelimiter}`
-        );
-        response.end();
+      };
+      if (streamManager) {
+        headers["x-genkit-stream-id"] = streamIdToUse;
       }
+      response.writeHead(200, headers);
+      runActionWithDurableStreaming(
+        action,
+        streamManager,
+        streamIdToUse,
+        input,
+        context,
+        response,
+        abortController.signal
+      );
     } else {
       try {
         const result = await action.run(input, {
@@ -95,12 +96,112 @@ ${e.stack}`
     }
   };
 }
+async function runActionWithDurableStreaming(action, streamManager, streamId, input, context, response, abortSignal) {
+  let taskQueue;
+  let durableStream;
+  if (streamManager) {
+    taskQueue = new AsyncTaskQueue();
+    durableStream = await streamManager.open(streamId);
+  }
+  try {
+    let onChunk = (chunk) => {
+      response.write(
+        "data: " + JSON.stringify({ message: chunk }) + streamDelimiter
+      );
+    };
+    if (streamManager) {
+      const originalOnChunk = onChunk;
+      onChunk = (chunk) => {
+        originalOnChunk(chunk);
+        taskQueue.enqueue(() => durableStream.write(chunk));
+      };
+    }
+    const result = await action.run(input, {
+      onChunk,
+      context,
+      abortSignal
+    });
+    if (streamManager) {
+      taskQueue.enqueue(() => durableStream.done(result.result));
+      await taskQueue.merge();
+    }
+    response.write(
+      "data: " + JSON.stringify({ result: result.result }) + streamDelimiter
+    );
+    response.end();
+  } catch (e) {
+    if (durableStream) {
+      taskQueue.enqueue(() => durableStream.error(e));
+      await taskQueue.merge();
+    }
+    logger.error(
+      `Streaming request failed with error: ${e.message}
+${e.stack}`
+    );
+    response.write(
+      `error: ${JSON.stringify({
+        error: getCallableJSON(e)
+      })}${streamDelimiter}`
+    );
+    response.end();
+  }
+}
+async function subscribeToStream(streamManager, streamId, response) {
+  try {
+    await streamManager.subscribe(streamId, {
+      onChunk: (chunk) => {
+        response.write(
+          "data: " + JSON.stringify({ message: chunk }) + streamDelimiter
+        );
+      },
+      onDone: (output) => {
+        response.write(
+          "data: " + JSON.stringify({ result: output }) + streamDelimiter
+        );
+        response.end();
+      },
+      onError: (err) => {
+        logger.error(
+          `Streaming request failed with error: ${err.message}
+${err.stack}`
+        );
+        response.write(
+          `error: ${JSON.stringify({
+            error: getCallableJSON(err)
+          })}${streamDelimiter}`
+        );
+        response.end();
+      }
+    });
+  } catch (e) {
+    if (e instanceof StreamNotFoundError) {
+      response.status(204).end();
+      return;
+    }
+    if (e.status === "DEADLINE_EXCEEDED") {
+      response.write(
+        `error: ${JSON.stringify({
+          error: getCallableJSON(e)
+        })}${streamDelimiter}`
+      );
+      response.end();
+      return;
+    }
+    throw e;
+  }
+}
 function withContextProvider(flow, context) {
   return {
     flow,
     context
   };
 }
+function withFlowOptions(flow, options) {
+  return {
+    flow,
+    options
+  };
+}
 function startFlowServer(options) {
   const server = new FlowServer(options);
   server.start();
@@ -130,13 +231,11 @@ class FlowServer {
     logger.debug("Running flow server with flow paths:");
     const pathPrefix = this.options.pathPrefix ?? "";
     this.options.flows?.forEach((flow) => {
-      if ("context" in flow) {
-        const flowPath = `/${pathPrefix}${flow.flow.__action.name}`;
+      if ("flow" in flow) {
+        const flowPath = `/${pathPrefix}${"options" in flow && flow.options.path || flow.flow.__action.name}`;
         logger.debug(` - ${flowPath}`);
-        server.post(
-          flowPath,
-          expressHandler(flow.flow, { contextProvider: flow.context })
-        );
+        const options = "options" in flow ? flow.options : { contextProvider: flow.context };
+        server.post(flowPath, expressHandler(flow.flow, options));
       } else {
         const flowPath = `/${pathPrefix}${flow.__action.name}`;
         logger.debug(` - ${flowPath}`);
@@ -190,6 +289,7 @@ export {
   FlowServer,
   expressHandler,
   startFlowServer,
-  withContextProvider
+  withContextProvider,
+  withFlowOptions
 };
 //# sourceMappingURL=index.mjs.map

package/lib/index.mjs.map

@@ -1 +1 @@
-{"version":3,"sources":["../src/index.ts"],"sourcesContent":["/**\n * Copyright 2024 Google LLC\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\nimport bodyParser from 'body-parser';\nimport cors, { type CorsOptions } from 'cors';\nimport express from 'express';\nimport { type Action, type ActionContext, type Flow, type z } from 'genkit';\nimport {\n  getCallableJSON,\n  getHttpStatus,\n  type ContextProvider,\n  type RequestData,\n} from 'genkit/context';\nimport { logger } from 'genkit/logging';\nimport type { Server } from 'http';\n\nconst streamDelimiter = '\\n\\n';\n\n/**\n * Exposes provided flow or an action as express handler.\n */\nexport function expressHandler<\n  C extends ActionContext = ActionContext,\n  I extends z.ZodTypeAny = z.ZodTypeAny,\n  O extends z.ZodTypeAny = z.ZodTypeAny,\n  S extends z.ZodTypeAny = z.ZodTypeAny,\n>(\n  action: Action<I, O, S>,\n  opts?: {\n    contextProvider?: ContextProvider<C, I>;\n  }\n): express.RequestHandler {\n  return async (\n    request: express.Request,\n    response: express.Response\n  ): Promise<void> => {\n    const { stream } = request.query;\n    if (!request.body) {\n      const errMsg =\n        `Error: request.body is undefined. ` +\n        `Possible reasons: missing 'content-type: application/json' in request ` +\n        `headers or misconfigured JSON middleware ('app.use(express.json()')? `;\n      logger.error(errMsg);\n      response\n        .status(400)\n        .json({ message: errMsg, status: 'INVALID ARGUMENT' })\n        .end();\n      return;\n    }\n\n    const input = request.body.data as z.infer<I>;\n    let context: Record<string, any>;\n\n    try {\n      context =\n        (await opts?.contextProvider?.({\n          method: request.method as RequestData['method'],\n          headers: Object.fromEntries(\n            Object.entries(request.headers).map(([key, value]) => [\n              key.toLowerCase(),\n              Array.isArray(value) ? value.join(' ') : String(value),\n            ])\n          ),\n          input,\n        })) || {};\n    } catch (e: any) {\n      logger.error(\n        `Auth policy failed with error: ${(e as Error).message}\\n${(e as Error).stack}`\n      );\n      response.status(getHttpStatus(e)).json(getCallableJSON(e)).end();\n      return;\n    }\n\n    const abortController = new AbortController();\n    request.on('close', () => {\n      abortController.abort();\n    });\n    // when/if using timeout middleware, it will emit 'timeout' event.\n    request.on('timeout', () => {\n      abortController.abort();\n    });\n\n    if (request.get('Accept') === 'text/event-stream' || stream === 'true') {\n      response.writeHead(200, {\n        'Content-Type': 'text/plain',\n        'Transfer-Encoding': 'chunked',\n      });\n      try {\n        const onChunk = (chunk: z.infer<S>) => {\n          response.write(\n            'data: ' + JSON.stringify({ message: chunk }) + streamDelimiter\n          );\n        };\n        const result = await action.run(input, {\n          onChunk,\n          context,\n          abortSignal: abortController.signal,\n        });\n        response.write(\n          'data: ' + JSON.stringify({ result: result.result }) + streamDelimiter\n        );\n        response.end();\n      } catch (e) {\n        logger.error(\n          `Streaming request failed with error: ${(e as Error).message}\\n${(e as Error).stack}`\n        );\n        response.write(\n          `error: ${JSON.stringify({ error: getCallableJSON(e) })}${streamDelimiter}`\n        );\n        response.end();\n      }\n    } else {\n      try {\n        const result = await action.run(input, {\n          context,\n          abortSignal: abortController.signal,\n        });\n        response.setHeader('x-genkit-trace-id', result.telemetry.traceId);\n        response.setHeader('x-genkit-span-id', result.telemetry.spanId);\n        // Responses for non-streaming flows are passed back with the flow result stored in a field called \"result.\"\n        response\n          .status(200)\n          .json({\n            result: result.result,\n          })\n          .end();\n      } catch (e) {\n        // Errors for non-streaming flows are passed back as standard API errors.\n        logger.error(\n          `Non-streaming request failed with error: ${(e as Error).message}\\n${(e as Error).stack}`\n        );\n        response.status(getHttpStatus(e)).json(getCallableJSON(e)).end();\n      }\n    }\n  };\n}\n\n/**\n * A wrapper object containing a flow with its associated auth policy.\n */\nexport type FlowWithContextProvider<\n  C extends ActionContext = ActionContext,\n  I extends z.ZodTypeAny = z.ZodTypeAny,\n  O extends z.ZodTypeAny = z.ZodTypeAny,\n  S extends z.ZodTypeAny = z.ZodTypeAny,\n> = {\n  flow: Flow<I, O, S>;\n  context: ContextProvider<C, I>;\n};\n\n/**\n * Adds an auth policy to the flow.\n */\nexport function withContextProvider<\n  C extends ActionContext = ActionContext,\n  I extends z.ZodTypeAny = z.ZodTypeAny,\n  O extends z.ZodTypeAny = z.ZodTypeAny,\n  S extends z.ZodTypeAny = z.ZodTypeAny,\n>(\n  flow: Flow<I, O, S>,\n  context: ContextProvider<C, I>\n): FlowWithContextProvider<C, I, O, S> {\n  return {\n    flow,\n    context,\n  };\n}\n\n/**\n * Options to configure the flow server.\n */\nexport interface FlowServerOptions {\n  /** List of flows to expose via the flow server. */\n  flows: (Flow<any, any, any> | FlowWithContextProvider<any, any, any>)[];\n  /** Port to run the server on. Defaults to env.PORT or 3400. */\n  port?: number;\n  /** CORS options for the server. */\n  cors?: CorsOptions;\n  /** HTTP method path prefix for the exposed flows. */\n  pathPrefix?: string;\n  /** JSON body parser options. */\n  jsonParserOptions?: bodyParser.OptionsJson;\n}\n\n/**\n * Starts an express server with the provided flows and options.\n */\nexport function startFlowServer(options: FlowServerOptions): FlowServer {\n  const server = new FlowServer(options);\n  server.start();\n  return server;\n}\n\n/**\n * Flow server exposes registered flows as HTTP endpoints.\n *\n * This is for use in production environments.\n *\n * @hidden\n */\nexport class FlowServer {\n  /** List of all running servers needed to be cleaned up on process exit. */\n  private static RUNNING_SERVERS: FlowServer[] = [];\n\n  /** Options for the flow server configured by the developer. */\n  private options: FlowServerOptions;\n  /** Port the server is actually running on. This may differ from `options.port` if the original was occupied. Null is server is not running. */\n  private port: number | null = null;\n  /** Express server instance. Null if server is not running. */\n  private server: Server | null = null;\n\n  constructor(options: FlowServerOptions) {\n    this.options = {\n      ...options,\n    };\n  }\n\n  /**\n   * Starts the server and adds it to the list of running servers to clean up on exit.\n   */\n  async start() {\n    const server = express();\n\n    server.use(bodyParser.json(this.options.jsonParserOptions));\n    server.use(cors(this.options.cors));\n\n    logger.debug('Running flow server with flow paths:');\n    const pathPrefix = this.options.pathPrefix ?? '';\n    this.options.flows?.forEach((flow) => {\n      if ('context' in flow) {\n        const flowPath = `/${pathPrefix}${flow.flow.__action.name}`;\n        logger.debug(` - ${flowPath}`);\n        server.post(\n          flowPath,\n          expressHandler(flow.flow, { contextProvider: flow.context })\n        );\n      } else {\n        const flowPath = `/${pathPrefix}${flow.__action.name}`;\n        logger.debug(` - ${flowPath}`);\n        server.post(flowPath, expressHandler(flow));\n      }\n    });\n    this.port =\n      this.options?.port ||\n      (process.env.PORT ? Number.parseInt(process.env.PORT) : 0) ||\n      3400;\n    this.server = server.listen(this.port, () => {\n      logger.debug(`Flow server running on http://localhost:${this.port}`);\n      FlowServer.RUNNING_SERVERS.push(this);\n    });\n  }\n\n  /**\n   * Stops the server and removes it from the list of running servers to clean up on exit.\n   */\n  async stop(): Promise<void> {\n    if (!this.server) {\n      return;\n    }\n    return new Promise<void>((resolve, reject) => {\n      this.server!.close((err) => {\n        if (err) {\n          logger.error(\n            `Error shutting down flow server on port ${this.port}: ${err}`\n          );\n          reject(err);\n        }\n        const index = FlowServer.RUNNING_SERVERS.indexOf(this);\n        if (index > -1) {\n          FlowServer.RUNNING_SERVERS.splice(index, 1);\n        }\n        logger.debug(\n          `Flow server on port ${this.port} has successfully shut down.`\n        );\n        this.port = null;\n        this.server = null;\n        resolve();\n      });\n    });\n  }\n\n  /**\n   * Stops all running servers.\n   */\n  static async stopAll() {\n    return Promise.all(\n      FlowServer.RUNNING_SERVERS.map((server) => server.stop())\n    );\n  }\n}\n"],"mappings":"AAgBA,OAAO,gBAAgB;AACvB,OAAO,UAAgC;AACvC,OAAO,aAAa;AAEpB;AAAA,EACE;AAAA,EACA;AAAA,OAGK;AACP,SAAS,cAAc;AAGvB,MAAM,kBAAkB;AAKjB,SAAS,eAMd,QACA,MAGwB;AACxB,SAAO,OACL,SACA,aACkB;AAClB,UAAM,EAAE,OAAO,IAAI,QAAQ;AAC3B,QAAI,CAAC,QAAQ,MAAM;AACjB,YAAM,SACJ;AAGF,aAAO,MAAM,MAAM;AACnB,eACG,OAAO,GAAG,EACV,KAAK,EAAE,SAAS,QAAQ,QAAQ,mBAAmB,CAAC,EACpD,IAAI;AACP;AAAA,IACF;AAEA,UAAM,QAAQ,QAAQ,KAAK;AAC3B,QAAI;AAEJ,QAAI;AACF,gBACG,MAAM,MAAM,kBAAkB;AAAA,QAC7B,QAAQ,QAAQ;AAAA,QAChB,SAAS,OAAO;AAAA,UACd,OAAO,QAAQ,QAAQ,OAAO,EAAE,IAAI,CAAC,CAAC,KAAK,KAAK,MAAM;AAAA,YACpD,IAAI,YAAY;AAAA,YAChB,MAAM,QAAQ,KAAK,IAAI,MAAM,KAAK,GAAG,IAAI,OAAO,KAAK;AAAA,UACvD,CAAC;AAAA,QACH;AAAA,QACA;AAAA,MACF,CAAC,KAAM,CAAC;AAAA,IACZ,SAAS,GAAQ;AACf,aAAO;AAAA,QACL,kCAAmC,EAAY,OAAO;AAAA,EAAM,EAAY,KAAK;AAAA,MAC/E;AACA,eAAS,OAAO,cAAc,CAAC,CAAC,EAAE,KAAK,gBAAgB,CAAC,CAAC,EAAE,IAAI;AAC/D;AAAA,IACF;AAEA,UAAM,kBAAkB,IAAI,gBAAgB;AAC5C,YAAQ,GAAG,SAAS,MAAM;AACxB,sBAAgB,MAAM;AAAA,IACxB,CAAC;AAED,YAAQ,GAAG,WAAW,MAAM;AAC1B,sBAAgB,MAAM;AAAA,IACxB,CAAC;AAED,QAAI,QAAQ,IAAI,QAAQ,MAAM,uBAAuB,WAAW,QAAQ;AACtE,eAAS,UAAU,KAAK;AAAA,QACtB,gBAAgB;AAAA,QAChB,qBAAqB;AAAA,MACvB,CAAC;AACD,UAAI;AACF,cAAM,UAAU,CAAC,UAAsB;AACrC,mBAAS;AAAA,YACP,WAAW,KAAK,UAAU,EAAE,SAAS,MAAM,CAAC,IAAI;AAAA,UAClD;AAAA,QACF;AACA,cAAM,SAAS,MAAM,OAAO,IAAI,OAAO;AAAA,UACrC;AAAA,UACA;AAAA,UACA,aAAa,gBAAgB;AAAA,QAC/B,CAAC;AACD,iBAAS;AAAA,UACP,WAAW,KAAK,UAAU,EAAE,QAAQ,OAAO,OAAO,CAAC,IAAI;AAAA,QACzD;AACA,iBAAS,IAAI;AAAA,MACf,SAAS,GAAG;AACV,eAAO;AAAA,UACL,wCAAyC,EAAY,OAAO;AAAA,EAAM,EAAY,KAAK;AAAA,QACrF;AACA,iBAAS;AAAA,UACP,UAAU,KAAK,UAAU,EAAE,OAAO,gBAAgB,CAAC,EAAE,CAAC,CAAC,GAAG,eAAe;AAAA,QAC3E;AACA,iBAAS,IAAI;AAAA,MACf;AAAA,IACF,OAAO;AACL,UAAI;AACF,cAAM,SAAS,MAAM,OAAO,IAAI,OAAO;AAAA,UACrC;AAAA,UACA,aAAa,gBAAgB;AAAA,QAC/B,CAAC;AACD,iBAAS,UAAU,qBAAqB,OAAO,UAAU,OAAO;AAChE,iBAAS,UAAU,oBAAoB,OAAO,UAAU,MAAM;AAE9D,iBACG,OAAO,GAAG,EACV,KAAK;AAAA,UACJ,QAAQ,OAAO;AAAA,QACjB,CAAC,EACA,IAAI;AAAA,MACT,SAAS,GAAG;AAEV,eAAO;AAAA,UACL,4CAA6C,EAAY,OAAO;AAAA,EAAM,EAAY,KAAK;AAAA,QACzF;AACA,iBAAS,OAAO,cAAc,CAAC,CAAC,EAAE,KAAK,gBAAgB,CAAC,CAAC,EAAE,IAAI;AAAA,MACjE;AAAA,IACF;AAAA,EACF;AACF;AAkBO,SAAS,oBAMd,MACA,SACqC;AACrC,SAAO;AAAA,IACL;AAAA,IACA;AAAA,EACF;AACF;AAqBO,SAAS,gBAAgB,SAAwC;AACtE,QAAM,SAAS,IAAI,WAAW,OAAO;AACrC,SAAO,MAAM;AACb,SAAO;AACT;AASO,MAAM,WAAW;AAAA;AAAA,EAEtB,OAAe,kBAAgC,CAAC;AAAA;AAAA,EAGxC;AAAA;AAAA,EAEA,OAAsB;AAAA;AAAA,EAEtB,SAAwB;AAAA,EAEhC,YAAY,SAA4B;AACtC,SAAK,UAAU;AAAA,MACb,GAAG;AAAA,IACL;AAAA,EACF;AAAA;AAAA;AAAA;AAAA,EAKA,MAAM,QAAQ;AACZ,UAAM,SAAS,QAAQ;AAEvB,WAAO,IAAI,WAAW,KAAK,KAAK,QAAQ,iBAAiB,CAAC;AAC1D,WAAO,IAAI,KAAK,KAAK,QAAQ,IAAI,CAAC;AAElC,WAAO,MAAM,sCAAsC;AACnD,UAAM,aAAa,KAAK,QAAQ,cAAc;AAC9C,SAAK,QAAQ,OAAO,QAAQ,CAAC,SAAS;AACpC,UAAI,aAAa,MAAM;AACrB,cAAM,WAAW,IAAI,UAAU,GAAG,KAAK,KAAK,SAAS,IAAI;AACzD,eAAO,MAAM,MAAM,QAAQ,EAAE;AAC7B,eAAO;AAAA,UACL;AAAA,UACA,eAAe,KAAK,MAAM,EAAE,iBAAiB,KAAK,QAAQ,CAAC;AAAA,QAC7D;AAAA,MACF,OAAO;AACL,cAAM,WAAW,IAAI,UAAU,GAAG,KAAK,SAAS,IAAI;AACpD,eAAO,MAAM,MAAM,QAAQ,EAAE;AAC7B,eAAO,KAAK,UAAU,eAAe,IAAI,CAAC;AAAA,MAC5C;AAAA,IACF,CAAC;AACD,SAAK,OACH,KAAK,SAAS,SACb,QAAQ,IAAI,OAAO,OAAO,SAAS,QAAQ,IAAI,IAAI,IAAI,MACxD;AACF,SAAK,SAAS,OAAO,OAAO,KAAK,MAAM,MAAM;AAC3C,aAAO,MAAM,2CAA2C,KAAK,IAAI,EAAE;AACnE,iBAAW,gBAAgB,KAAK,IAAI;AAAA,IACtC,CAAC;AAAA,EACH;AAAA;AAAA;AAAA;AAAA,EAKA,MAAM,OAAsB;AAC1B,QAAI,CAAC,KAAK,QAAQ;AAChB;AAAA,IACF;AACA,WAAO,IAAI,QAAc,CAAC,SAAS,WAAW;AAC5C,WAAK,OAAQ,MAAM,CAAC,QAAQ;AAC1B,YAAI,KAAK;AACP,iBAAO;AAAA,YACL,2CAA2C,KAAK,IAAI,KAAK,GAAG;AAAA,UAC9D;AACA,iBAAO,GAAG;AAAA,QACZ;AACA,cAAM,QAAQ,WAAW,gBAAgB,QAAQ,IAAI;AACrD,YAAI,QAAQ,IAAI;AACd,qBAAW,gBAAgB,OAAO,OAAO,CAAC;AAAA,QAC5C;AACA,eAAO;AAAA,UACL,uBAAuB,KAAK,IAAI;AAAA,QAClC;AACA,aAAK,OAAO;AACZ,aAAK,SAAS;AACd,gBAAQ;AAAA,MACV,CAAC;AAAA,IACH,CAAC;AAAA,EACH;AAAA;AAAA;AAAA;AAAA,EAKA,aAAa,UAAU;AACrB,WAAO,QAAQ;AAAA,MACb,WAAW,gBAAgB,IAAI,CAAC,WAAW,OAAO,KAAK,CAAC;AAAA,IAC1D;AAAA,EACF;AACF;","names":[]}
+{"version":3,"sources":["../src/index.ts"],"sourcesContent":["/**\n * Copyright 2024 Google LLC\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\nimport bodyParser from 'body-parser';\nimport cors, { type CorsOptions } from 'cors';\nimport { randomUUID } from 'crypto';\nimport express from 'express';\nimport {\n  Action,\n  ActionStreamInput,\n  AsyncTaskQueue,\n  Flow,\n  StreamNotFoundError,\n  type ActionContext,\n  type StreamManager,\n  type z,\n} from 'genkit/beta';\nimport {\n  getCallableJSON,\n  getHttpStatus,\n  type ContextProvider,\n  type RequestData,\n} from 'genkit/context';\nimport { logger } from 'genkit/logging';\nimport type { Server } from 'http';\n\nconst streamDelimiter = '\\n\\n';\n\n/**\n * Exposes provided flow or an action as express handler.\n */\nexport function expressHandler<\n  C extends ActionContext = ActionContext,\n  I extends z.ZodTypeAny = z.ZodTypeAny,\n  O extends z.ZodTypeAny = z.ZodTypeAny,\n  S extends z.ZodTypeAny = z.ZodTypeAny,\n>(\n  action: Action<I, O, S>,\n  opts?: {\n    contextProvider?: ContextProvider<C, I>;\n    streamManager?: StreamManager;\n  }\n): express.RequestHandler {\n  return async (\n    request: express.Request,\n    response: express.Response\n  ): Promise<void> => {\n    const { stream } = request.query;\n    const streamIdHeader = request.headers['x-genkit-stream-id'];\n    const streamId = Array.isArray(streamIdHeader)\n      ? streamIdHeader[0]\n      : streamIdHeader;\n\n    if (!request.body) {\n      const errMsg =\n        `Error: request.body is undefined. ` +\n        `Possible reasons: missing 'content-type: application/json' in request ` +\n        `headers or misconfigured JSON middleware ('app.use(express.json()')? `;\n      logger.error(errMsg);\n      response\n        .status(400)\n        .json({ message: errMsg, status: 'INVALID ARGUMENT' })\n        .end();\n      return;\n    }\n\n    const input = request.body.data as z.infer<I>;\n    let context: Record<string, any>;\n\n    try {\n      context =\n        (await opts?.contextProvider?.({\n          method: request.method as RequestData['method'],\n          headers: Object.fromEntries(\n            Object.entries(request.headers).map(([key, value]) => [\n              key.toLowerCase(),\n              Array.isArray(value) ? value.join(' ') : String(value),\n            ])\n          ),\n          input,\n        })) || {};\n    } catch (e: any) {\n      logger.error(\n        `Auth policy failed with error: ${(e as Error).message}\\n${(e as Error).stack}`\n      );\n      response.status(getHttpStatus(e)).json(getCallableJSON(e)).end();\n      return;\n    }\n\n    const abortController = new AbortController();\n    request.on('close', () => {\n      abortController.abort();\n    });\n    // when/if using timeout middleware, it will emit 'timeout' event.\n    request.on('timeout', () => {\n      abortController.abort();\n    });\n\n    if (request.get('Accept') === 'text/event-stream' || stream === 'true') {\n      const streamManager = opts?.streamManager;\n      if (streamManager && streamId) {\n        await subscribeToStream(streamManager, streamId, response);\n        return;\n      }\n      const streamIdToUse = randomUUID();\n      const headers = {\n        'Content-Type': 'text/plain',\n        'Transfer-Encoding': 'chunked',\n      };\n      if (streamManager) {\n        headers['x-genkit-stream-id'] = streamIdToUse;\n      }\n      response.writeHead(200, headers);\n      runActionWithDurableStreaming(\n        action,\n        streamManager,\n        streamIdToUse,\n        input,\n        context,\n        response,\n        abortController.signal\n      );\n    } else {\n      try {\n        const result = await action.run(input, {\n          context,\n          abortSignal: abortController.signal,\n        });\n        response.setHeader('x-genkit-trace-id', result.telemetry.traceId);\n        response.setHeader('x-genkit-span-id', result.telemetry.spanId);\n        // Responses for non-streaming flows are passed back with the flow result stored in a field called \"result.\"\n        response\n          .status(200)\n          .json({\n            result: result.result,\n          })\n          .end();\n      } catch (e) {\n        // Errors for non-streaming flows are passed back as standard API errors.\n        logger.error(\n          `Non-streaming request failed with error: ${(e as Error).message}\\n${(e as Error).stack}`\n        );\n        response.status(getHttpStatus(e)).json(getCallableJSON(e)).end();\n      }\n    }\n  };\n}\n\nasync function runActionWithDurableStreaming<\n  I extends z.ZodTypeAny,\n  O extends z.ZodTypeAny,\n  S extends z.ZodTypeAny,\n>(\n  action: Action<I, O, S>,\n  streamManager: StreamManager | undefined,\n  streamId: string,\n  input: z.infer<I>,\n  context: ActionContext,\n  response: express.Response,\n  abortSignal: AbortSignal\n) {\n  let taskQueue: AsyncTaskQueue | undefined;\n  let durableStream: ActionStreamInput<any, any> | undefined;\n  if (streamManager) {\n    taskQueue = new AsyncTaskQueue();\n    durableStream = await streamManager.open(streamId);\n  }\n  try {\n    let onChunk = (chunk: z.infer<S>) => {\n      response.write(\n        'data: ' + JSON.stringify({ message: chunk }) + streamDelimiter\n      );\n    };\n    if (streamManager) {\n      const originalOnChunk = onChunk;\n      onChunk = (chunk: z.infer<S>) => {\n        originalOnChunk(chunk);\n        taskQueue!.enqueue(() => durableStream!.write(chunk));\n      };\n    }\n    const result = await action.run(input, {\n      onChunk,\n      context,\n      abortSignal,\n    });\n    if (streamManager) {\n      taskQueue!.enqueue(() => durableStream!.done(result.result));\n      await taskQueue!.merge();\n    }\n    response.write(\n      'data: ' + JSON.stringify({ result: result.result }) + streamDelimiter\n    );\n    response.end();\n  } catch (e) {\n    if (durableStream) {\n      taskQueue!.enqueue(() => durableStream!.error(e));\n      await taskQueue!.merge();\n    }\n    logger.error(\n      `Streaming request failed with error: ${(e as Error).message}\\n${\n        (e as Error).stack\n      }`\n    );\n    response.write(\n      `error: ${JSON.stringify({\n        error: getCallableJSON(e),\n      })}${streamDelimiter}`\n    );\n    response.end();\n  }\n}\n\nasync function subscribeToStream(\n  streamManager: StreamManager,\n  streamId: string,\n  response: express.Response\n): Promise<void> {\n  try {\n    await streamManager.subscribe(streamId, {\n      onChunk: (chunk) => {\n        response.write(\n          'data: ' + JSON.stringify({ message: chunk }) + streamDelimiter\n        );\n      },\n      onDone: (output) => {\n        response.write(\n          'data: ' + JSON.stringify({ result: output }) + streamDelimiter\n        );\n        response.end();\n      },\n      onError: (err) => {\n        logger.error(\n          `Streaming request failed with error: ${(err as Error).message}\\n${\n            (err as Error).stack\n          }`\n        );\n        response.write(\n          `error: ${JSON.stringify({\n            error: getCallableJSON(err),\n          })}${streamDelimiter}`\n        );\n        response.end();\n      },\n    });\n  } catch (e: any) {\n    if (e instanceof StreamNotFoundError) {\n      response.status(204).end();\n      return;\n    }\n    if (e.status === 'DEADLINE_EXCEEDED') {\n      response.write(\n        `error: ${JSON.stringify({\n          error: getCallableJSON(e),\n        })}${streamDelimiter}`\n      );\n      response.end();\n      return;\n    }\n    throw e;\n  }\n}\n\n/**\n * A wrapper object containing a flow with its associated auth policy.\n * @deprecated Use `withFlowOptions` instead.\n */\nexport type FlowWithContextProvider<\n  C extends ActionContext = ActionContext,\n  I extends z.ZodTypeAny = z.ZodTypeAny,\n  O extends z.ZodTypeAny = z.ZodTypeAny,\n  S extends z.ZodTypeAny = z.ZodTypeAny,\n> = {\n  flow: Flow<I, O, S>;\n  context: ContextProvider<C, I>;\n};\n\n/**\n * A wrapper object containing a flow with its associated options.\n */\nexport type FlowWithOptions<\n  I extends z.ZodTypeAny = z.ZodTypeAny,\n  O extends z.ZodTypeAny = z.ZodTypeAny,\n  S extends z.ZodTypeAny = z.ZodTypeAny,\n> = {\n  flow: Flow<I, O, S>;\n  options: {\n    contextProvider?: ContextProvider<any, I>;\n    streamManager?: StreamManager;\n    path?: string;\n  };\n};\n\n/**\n * Adds an auth policy to the flow.\n * @deprecated Use `withFlowOptions` instead.\n */\nexport function withContextProvider<\n  C extends ActionContext = ActionContext,\n  I extends z.ZodTypeAny = z.ZodTypeAny,\n  O extends z.ZodTypeAny = z.ZodTypeAny,\n  S extends z.ZodTypeAny = z.ZodTypeAny,\n>(\n  flow: Flow<I, O, S>,\n  context: ContextProvider<C, I>\n): FlowWithContextProvider<C, I, O, S> {\n  return {\n    flow,\n    context,\n  };\n}\n\n/**\n * Adds an auth policy to the flow.\n */\nexport function withFlowOptions<\n  I extends z.ZodTypeAny,\n  O extends z.ZodTypeAny,\n  S extends z.ZodTypeAny,\n>(\n  flow: Flow<I, O, S>,\n  options: {\n    contextProvider?: ContextProvider<any, I>;\n    streamManager?: StreamManager;\n    path?: string;\n  }\n): FlowWithOptions<I, O, S> {\n  return {\n    flow,\n    options,\n  };\n}\n\n/**\n * Options to configure the flow server.\n */\nexport interface FlowServerOptions {\n  /** List of flows to expose via the flow server. */\n  flows: (\n    | Flow<any, any, any>\n    | FlowWithContextProvider<any, any, any>\n    | FlowWithOptions<any, any, any>\n  )[];\n  /** Port to run the server on. Defaults to env.PORT or 3400. */\n  port?: number;\n  /** CORS options for the server. */\n  cors?: CorsOptions;\n  /** HTTP method path prefix for the exposed flows. */\n  pathPrefix?: string;\n  /** JSON body parser options. */\n  jsonParserOptions?: bodyParser.OptionsJson;\n}\n\n/**\n * Starts an express server with the provided flows and options.\n */\nexport function startFlowServer(options: FlowServerOptions): FlowServer {\n  const server = new FlowServer(options);\n  server.start();\n  return server;\n}\n\n/**\n * Flow server exposes registered flows as HTTP endpoints.\n *\n * This is for use in production environments.\n *\n * @hidden\n */\nexport class FlowServer {\n  /** List of all running servers needed to be cleaned up on process exit. */\n  private static RUNNING_SERVERS: FlowServer[] = [];\n\n  /** Options for the flow server configured by the developer. */\n  private options: FlowServerOptions;\n  /** Port the server is actually running on. This may differ from `options.port` if the original was occupied. Null is server is not running. */\n  private port: number | null = null;\n  /** Express server instance. Null if server is not running. */\n  private server: Server | null = null;\n\n  constructor(options: FlowServerOptions) {\n    this.options = {\n      ...options,\n    };\n  }\n\n  /**\n   * Starts the server and adds it to the list of running servers to clean up on exit.\n   */\n  async start() {\n    const server = express();\n\n    server.use(bodyParser.json(this.options.jsonParserOptions));\n    server.use(cors(this.options.cors));\n\n    logger.debug('Running flow server with flow paths:');\n    const pathPrefix = this.options.pathPrefix ?? '';\n    this.options.flows?.forEach((flow) => {\n      if ('flow' in flow) {\n        const flowPath = `/${pathPrefix}${\n          ('options' in flow && flow.options.path) || flow.flow.__action.name\n        }`;\n        logger.debug(` - ${flowPath}`);\n        const options =\n          'options' in flow ? flow.options : { contextProvider: flow.context };\n        server.post(flowPath, expressHandler(flow.flow, options));\n      } else {\n        const flowPath = `/${pathPrefix}${flow.__action.name}`;\n        logger.debug(` - ${flowPath}`);\n        server.post(flowPath, expressHandler(flow));\n      }\n    });\n    this.port =\n      this.options?.port ||\n      (process.env.PORT ? Number.parseInt(process.env.PORT) : 0) ||\n      3400;\n    this.server = server.listen(this.port, () => {\n      logger.debug(`Flow server running on http://localhost:${this.port}`);\n      FlowServer.RUNNING_SERVERS.push(this);\n    });\n  }\n\n  /**\n   * Stops the server and removes it from the list of running servers to clean up on exit.\n   */\n  async stop(): Promise<void> {\n    if (!this.server) {\n      return;\n    }\n    return new Promise<void>((resolve, reject) => {\n      this.server!.close((err) => {\n        if (err) {\n          logger.error(\n            `Error shutting down flow server on port ${this.port}: ${err}`\n          );\n          reject(err);\n        }\n        const index = FlowServer.RUNNING_SERVERS.indexOf(this);\n        if (index > -1) {\n          FlowServer.RUNNING_SERVERS.splice(index, 1);\n        }\n        logger.debug(\n          `Flow server on port ${this.port} has successfully shut down.`\n        );\n        this.port = null;\n        this.server = null;\n        resolve();\n      });\n    });\n  }\n\n  /**\n   * Stops all running servers.\n   */\n  static async stopAll() {\n    return Promise.all(\n      FlowServer.RUNNING_SERVERS.map((server) => server.stop())\n    );\n  }\n}\n"],"mappings":"AAgBA,OAAO,gBAAgB;AACvB,OAAO,UAAgC;AACvC,SAAS,kBAAkB;AAC3B,OAAO,aAAa;AACpB;AAAA,EAGE;AAAA,EAEA;AAAA,OAIK;AACP;AAAA,EACE;AAAA,EACA;AAAA,OAGK;AACP,SAAS,cAAc;AAGvB,MAAM,kBAAkB;AAKjB,SAAS,eAMd,QACA,MAIwB;AACxB,SAAO,OACL,SACA,aACkB;AAClB,UAAM,EAAE,OAAO,IAAI,QAAQ;AAC3B,UAAM,iBAAiB,QAAQ,QAAQ,oBAAoB;AAC3D,UAAM,WAAW,MAAM,QAAQ,cAAc,IACzC,eAAe,CAAC,IAChB;AAEJ,QAAI,CAAC,QAAQ,MAAM;AACjB,YAAM,SACJ;AAGF,aAAO,MAAM,MAAM;AACnB,eACG,OAAO,GAAG,EACV,KAAK,EAAE,SAAS,QAAQ,QAAQ,mBAAmB,CAAC,EACpD,IAAI;AACP;AAAA,IACF;AAEA,UAAM,QAAQ,QAAQ,KAAK;AAC3B,QAAI;AAEJ,QAAI;AACF,gBACG,MAAM,MAAM,kBAAkB;AAAA,QAC7B,QAAQ,QAAQ;AAAA,QAChB,SAAS,OAAO;AAAA,UACd,OAAO,QAAQ,QAAQ,OAAO,EAAE,IAAI,CAAC,CAAC,KAAK,KAAK,MAAM;AAAA,YACpD,IAAI,YAAY;AAAA,YAChB,MAAM,QAAQ,KAAK,IAAI,MAAM,KAAK,GAAG,IAAI,OAAO,KAAK;AAAA,UACvD,CAAC;AAAA,QACH;AAAA,QACA;AAAA,MACF,CAAC,KAAM,CAAC;AAAA,IACZ,SAAS,GAAQ;AACf,aAAO;AAAA,QACL,kCAAmC,EAAY,OAAO;AAAA,EAAM,EAAY,KAAK;AAAA,MAC/E;AACA,eAAS,OAAO,cAAc,CAAC,CAAC,EAAE,KAAK,gBAAgB,CAAC,CAAC,EAAE,IAAI;AAC/D;AAAA,IACF;AAEA,UAAM,kBAAkB,IAAI,gBAAgB;AAC5C,YAAQ,GAAG,SAAS,MAAM;AACxB,sBAAgB,MAAM;AAAA,IACxB,CAAC;AAED,YAAQ,GAAG,WAAW,MAAM;AAC1B,sBAAgB,MAAM;AAAA,IACxB,CAAC;AAED,QAAI,QAAQ,IAAI,QAAQ,MAAM,uBAAuB,WAAW,QAAQ;AACtE,YAAM,gBAAgB,MAAM;AAC5B,UAAI,iBAAiB,UAAU;AAC7B,cAAM,kBAAkB,eAAe,UAAU,QAAQ;AACzD;AAAA,MACF;AACA,YAAM,gBAAgB,WAAW;AACjC,YAAM,UAAU;AAAA,QACd,gBAAgB;AAAA,QAChB,qBAAqB;AAAA,MACvB;AACA,UAAI,eAAe;AACjB,gBAAQ,oBAAoB,IAAI;AAAA,MAClC;AACA,eAAS,UAAU,KAAK,OAAO;AAC/B;AAAA,QACE;AAAA,QACA;AAAA,QACA;AAAA,QACA;AAAA,QACA;AAAA,QACA;AAAA,QACA,gBAAgB;AAAA,MAClB;AAAA,IACF,OAAO;AACL,UAAI;AACF,cAAM,SAAS,MAAM,OAAO,IAAI,OAAO;AAAA,UACrC;AAAA,UACA,aAAa,gBAAgB;AAAA,QAC/B,CAAC;AACD,iBAAS,UAAU,qBAAqB,OAAO,UAAU,OAAO;AAChE,iBAAS,UAAU,oBAAoB,OAAO,UAAU,MAAM;AAE9D,iBACG,OAAO,GAAG,EACV,KAAK;AAAA,UACJ,QAAQ,OAAO;AAAA,QACjB,CAAC,EACA,IAAI;AAAA,MACT,SAAS,GAAG;AAEV,eAAO;AAAA,UACL,4CAA6C,EAAY,OAAO;AAAA,EAAM,EAAY,KAAK;AAAA,QACzF;AACA,iBAAS,OAAO,cAAc,CAAC,CAAC,EAAE,KAAK,gBAAgB,CAAC,CAAC,EAAE,IAAI;AAAA,MACjE;AAAA,IACF;AAAA,EACF;AACF;AAEA,eAAe,8BAKb,QACA,eACA,UACA,OACA,SACA,UACA,aACA;AACA,MAAI;AACJ,MAAI;AACJ,MAAI,eAAe;AACjB,gBAAY,IAAI,eAAe;AAC/B,oBAAgB,MAAM,cAAc,KAAK,QAAQ;AAAA,EACnD;AACA,MAAI;AACF,QAAI,UAAU,CAAC,UAAsB;AACnC,eAAS;AAAA,QACP,WAAW,KAAK,UAAU,EAAE,SAAS,MAAM,CAAC,IAAI;AAAA,MAClD;AAAA,IACF;AACA,QAAI,eAAe;AACjB,YAAM,kBAAkB;AACxB,gBAAU,CAAC,UAAsB;AAC/B,wBAAgB,KAAK;AACrB,kBAAW,QAAQ,MAAM,cAAe,MAAM,KAAK,CAAC;AAAA,MACtD;AAAA,IACF;AACA,UAAM,SAAS,MAAM,OAAO,IAAI,OAAO;AAAA,MACrC;AAAA,MACA;AAAA,MACA;AAAA,IACF,CAAC;AACD,QAAI,eAAe;AACjB,gBAAW,QAAQ,MAAM,cAAe,KAAK,OAAO,MAAM,CAAC;AAC3D,YAAM,UAAW,MAAM;AAAA,IACzB;AACA,aAAS;AAAA,MACP,WAAW,KAAK,UAAU,EAAE,QAAQ,OAAO,OAAO,CAAC,IAAI;AAAA,IACzD;AACA,aAAS,IAAI;AAAA,EACf,SAAS,GAAG;AACV,QAAI,eAAe;AACjB,gBAAW,QAAQ,MAAM,cAAe,MAAM,CAAC,CAAC;AAChD,YAAM,UAAW,MAAM;AAAA,IACzB;AACA,WAAO;AAAA,MACL,wCAAyC,EAAY,OAAO;AAAA,EACzD,EAAY,KACf;AAAA,IACF;AACA,aAAS;AAAA,MACP,UAAU,KAAK,UAAU;AAAA,QACvB,OAAO,gBAAgB,CAAC;AAAA,MAC1B,CAAC,CAAC,GAAG,eAAe;AAAA,IACtB;AACA,aAAS,IAAI;AAAA,EACf;AACF;AAEA,eAAe,kBACb,eACA,UACA,UACe;AACf,MAAI;AACF,UAAM,cAAc,UAAU,UAAU;AAAA,MACtC,SAAS,CAAC,UAAU;AAClB,iBAAS;AAAA,UACP,WAAW,KAAK,UAAU,EAAE,SAAS,MAAM,CAAC,IAAI;AAAA,QAClD;AAAA,MACF;AAAA,MACA,QAAQ,CAAC,WAAW;AAClB,iBAAS;AAAA,UACP,WAAW,KAAK,UAAU,EAAE,QAAQ,OAAO,CAAC,IAAI;AAAA,QAClD;AACA,iBAAS,IAAI;AAAA,MACf;AAAA,MACA,SAAS,CAAC,QAAQ;AAChB,eAAO;AAAA,UACL,wCAAyC,IAAc,OAAO;AAAA,EAC3D,IAAc,KACjB;AAAA,QACF;AACA,iBAAS;AAAA,UACP,UAAU,KAAK,UAAU;AAAA,YACvB,OAAO,gBAAgB,GAAG;AAAA,UAC5B,CAAC,CAAC,GAAG,eAAe;AAAA,QACtB;AACA,iBAAS,IAAI;AAAA,MACf;AAAA,IACF,CAAC;AAAA,EACH,SAAS,GAAQ;AACf,QAAI,aAAa,qBAAqB;AACpC,eAAS,OAAO,GAAG,EAAE,IAAI;AACzB;AAAA,IACF;AACA,QAAI,EAAE,WAAW,qBAAqB;AACpC,eAAS;AAAA,QACP,UAAU,KAAK,UAAU;AAAA,UACvB,OAAO,gBAAgB,CAAC;AAAA,QAC1B,CAAC,CAAC,GAAG,eAAe;AAAA,MACtB;AACA,eAAS,IAAI;AACb;AAAA,IACF;AACA,UAAM;AAAA,EACR;AACF;AAoCO,SAAS,oBAMd,MACA,SACqC;AACrC,SAAO;AAAA,IACL;AAAA,IACA;AAAA,EACF;AACF;AAKO,SAAS,gBAKd,MACA,SAK0B;AAC1B,SAAO;AAAA,IACL;AAAA,IACA;AAAA,EACF;AACF;AAyBO,SAAS,gBAAgB,SAAwC;AACtE,QAAM,SAAS,IAAI,WAAW,OAAO;AACrC,SAAO,MAAM;AACb,SAAO;AACT;AASO,MAAM,WAAW;AAAA;AAAA,EAEtB,OAAe,kBAAgC,CAAC;AAAA;AAAA,EAGxC;AAAA;AAAA,EAEA,OAAsB;AAAA;AAAA,EAEtB,SAAwB;AAAA,EAEhC,YAAY,SAA4B;AACtC,SAAK,UAAU;AAAA,MACb,GAAG;AAAA,IACL;AAAA,EACF;AAAA;AAAA;AAAA;AAAA,EAKA,MAAM,QAAQ;AACZ,UAAM,SAAS,QAAQ;AAEvB,WAAO,IAAI,WAAW,KAAK,KAAK,QAAQ,iBAAiB,CAAC;AAC1D,WAAO,IAAI,KAAK,KAAK,QAAQ,IAAI,CAAC;AAElC,WAAO,MAAM,sCAAsC;AACnD,UAAM,aAAa,KAAK,QAAQ,cAAc;AAC9C,SAAK,QAAQ,OAAO,QAAQ,CAAC,SAAS;AACpC,UAAI,UAAU,MAAM;AAClB,cAAM,WAAW,IAAI,UAAU,GAC5B,aAAa,QAAQ,KAAK,QAAQ,QAAS,KAAK,KAAK,SAAS,IACjE;AACA,eAAO,MAAM,MAAM,QAAQ,EAAE;AAC7B,cAAM,UACJ,aAAa,OAAO,KAAK,UAAU,EAAE,iBAAiB,KAAK,QAAQ;AACrE,eAAO,KAAK,UAAU,eAAe,KAAK,MAAM,OAAO,CAAC;AAAA,MAC1D,OAAO;AACL,cAAM,WAAW,IAAI,UAAU,GAAG,KAAK,SAAS,IAAI;AACpD,eAAO,MAAM,MAAM,QAAQ,EAAE;AAC7B,eAAO,KAAK,UAAU,eAAe,IAAI,CAAC;AAAA,MAC5C;AAAA,IACF,CAAC;AACD,SAAK,OACH,KAAK,SAAS,SACb,QAAQ,IAAI,OAAO,OAAO,SAAS,QAAQ,IAAI,IAAI,IAAI,MACxD;AACF,SAAK,SAAS,OAAO,OAAO,KAAK,MAAM,MAAM;AAC3C,aAAO,MAAM,2CAA2C,KAAK,IAAI,EAAE;AACnE,iBAAW,gBAAgB,KAAK,IAAI;AAAA,IACtC,CAAC;AAAA,EACH;AAAA;AAAA;AAAA;AAAA,EAKA,MAAM,OAAsB;AAC1B,QAAI,CAAC,KAAK,QAAQ;AAChB;AAAA,IACF;AACA,WAAO,IAAI,QAAc,CAAC,SAAS,WAAW;AAC5C,WAAK,OAAQ,MAAM,CAAC,QAAQ;AAC1B,YAAI,KAAK;AACP,iBAAO;AAAA,YACL,2CAA2C,KAAK,IAAI,KAAK,GAAG;AAAA,UAC9D;AACA,iBAAO,GAAG;AAAA,QACZ;AACA,cAAM,QAAQ,WAAW,gBAAgB,QAAQ,IAAI;AACrD,YAAI,QAAQ,IAAI;AACd,qBAAW,gBAAgB,OAAO,OAAO,CAAC;AAAA,QAC5C;AACA,eAAO;AAAA,UACL,uBAAuB,KAAK,IAAI;AAAA,QAClC;AACA,aAAK,OAAO;AACZ,aAAK,SAAS;AACd,gBAAQ;AAAA,MACV,CAAC;AAAA,IACH,CAAC;AAAA,EACH;AAAA;AAAA;AAAA;AAAA,EAKA,aAAa,UAAU;AACrB,WAAO,QAAQ;AAAA,MACb,WAAW,gBAAgB,IAAI,CAAC,WAAW,OAAO,KAAK,CAAC;AAAA,IAC1D;AAAA,EACF;AACF;","names":[]}

package/package.json

@@ -9,7 +9,7 @@
     "genai",
     "generative-ai"
   ],
-  "version": "1.24.0",
+  "version": "1.26.0-rc.0",
   "type": "commonjs",
   "repository": {
     "type": "git",
@@ -23,8 +23,8 @@
     "body-parser": "^1.20.3"
   },
   "peerDependencies": {
-    "express": "^4.21.1",
-    "genkit": "^1.24.0"
+    "express": "^4.21.1 || ^5",
+    "genkit": "^1.26.0-rc.0"
   },
   "devDependencies": {
     "get-port": "^5.1.0",

package/src/index.ts

@@ -16,8 +16,18 @@
 
 import bodyParser from 'body-parser';
 import cors, { type CorsOptions } from 'cors';
+import { randomUUID } from 'crypto';
 import express from 'express';
-import { type Action, type ActionContext, type Flow, type z } from 'genkit';
+import {
+  Action,
+  ActionStreamInput,
+  AsyncTaskQueue,
+  Flow,
+  StreamNotFoundError,
+  type ActionContext,
+  type StreamManager,
+  type z,
+} from 'genkit/beta';
 import {
   getCallableJSON,
   getHttpStatus,
@@ -41,6 +51,7 @@ export function expressHandler<
   action: Action<I, O, S>,
   opts?: {
     contextProvider?: ContextProvider<C, I>;
+    streamManager?: StreamManager;
   }
 ): express.RequestHandler {
   return async (
@@ -48,6 +59,11 @@ export function expressHandler<
     response: express.Response
   ): Promise<void> => {
     const { stream } = request.query;
+    const streamIdHeader = request.headers['x-genkit-stream-id'];
+    const streamId = Array.isArray(streamIdHeader)
+      ? streamIdHeader[0]
+      : streamIdHeader;
+
     if (!request.body) {
       const errMsg =
         `Error: request.body is undefined. ` +
@@ -94,34 +110,29 @@ export function expressHandler<
     });
 
     if (request.get('Accept') === 'text/event-stream' || stream === 'true') {
-      response.writeHead(200, {
+      const streamManager = opts?.streamManager;
+      if (streamManager && streamId) {
+        await subscribeToStream(streamManager, streamId, response);
+        return;
+      }
+      const streamIdToUse = randomUUID();
+      const headers = {
         'Content-Type': 'text/plain',
         'Transfer-Encoding': 'chunked',
-      });
-      try {
-        const onChunk = (chunk: z.infer<S>) => {
-          response.write(
-            'data: ' + JSON.stringify({ message: chunk }) + streamDelimiter
-          );
-        };
-        const result = await action.run(input, {
-          onChunk,
-          context,
-          abortSignal: abortController.signal,
-        });
-        response.write(
-          'data: ' + JSON.stringify({ result: result.result }) + streamDelimiter
-        );
-        response.end();
-      } catch (e) {
-        logger.error(
-          `Streaming request failed with error: ${(e as Error).message}\n${(e as Error).stack}`
-        );
-        response.write(
-          `error: ${JSON.stringify({ error: getCallableJSON(e) })}${streamDelimiter}`
-        );
-        response.end();
+      };
+      if (streamManager) {
+        headers['x-genkit-stream-id'] = streamIdToUse;
       }
+      response.writeHead(200, headers);
+      runActionWithDurableStreaming(
+        action,
+        streamManager,
+        streamIdToUse,
+        input,
+        context,
+        response,
+        abortController.signal
+      );
     } else {
       try {
         const result = await action.run(input, {
@@ -148,8 +159,123 @@ export function expressHandler<
   };
 }
 
+async function runActionWithDurableStreaming<
+  I extends z.ZodTypeAny,
+  O extends z.ZodTypeAny,
+  S extends z.ZodTypeAny,
+>(
+  action: Action<I, O, S>,
+  streamManager: StreamManager | undefined,
+  streamId: string,
+  input: z.infer<I>,
+  context: ActionContext,
+  response: express.Response,
+  abortSignal: AbortSignal
+) {
+  let taskQueue: AsyncTaskQueue | undefined;
+  let durableStream: ActionStreamInput<any, any> | undefined;
+  if (streamManager) {
+    taskQueue = new AsyncTaskQueue();
+    durableStream = await streamManager.open(streamId);
+  }
+  try {
+    let onChunk = (chunk: z.infer<S>) => {
+      response.write(
+        'data: ' + JSON.stringify({ message: chunk }) + streamDelimiter
+      );
+    };
+    if (streamManager) {
+      const originalOnChunk = onChunk;
+      onChunk = (chunk: z.infer<S>) => {
+        originalOnChunk(chunk);
+        taskQueue!.enqueue(() => durableStream!.write(chunk));
+      };
+    }
+    const result = await action.run(input, {
+      onChunk,
+      context,
+      abortSignal,
+    });
+    if (streamManager) {
+      taskQueue!.enqueue(() => durableStream!.done(result.result));
+      await taskQueue!.merge();
+    }
+    response.write(
+      'data: ' + JSON.stringify({ result: result.result }) + streamDelimiter
+    );
+    response.end();
+  } catch (e) {
+    if (durableStream) {
+      taskQueue!.enqueue(() => durableStream!.error(e));
+      await taskQueue!.merge();
+    }
+    logger.error(
+      `Streaming request failed with error: ${(e as Error).message}\n${
+        (e as Error).stack
+      }`
+    );
+    response.write(
+      `error: ${JSON.stringify({
+        error: getCallableJSON(e),
+      })}${streamDelimiter}`
+    );
+    response.end();
+  }
+}
+
+async function subscribeToStream(
+  streamManager: StreamManager,
+  streamId: string,
+  response: express.Response
+): Promise<void> {
+  try {
+    await streamManager.subscribe(streamId, {
+      onChunk: (chunk) => {
+        response.write(
+          'data: ' + JSON.stringify({ message: chunk }) + streamDelimiter
+        );
+      },
+      onDone: (output) => {
+        response.write(
+          'data: ' + JSON.stringify({ result: output }) + streamDelimiter
+        );
+        response.end();
+      },
+      onError: (err) => {
+        logger.error(
+          `Streaming request failed with error: ${(err as Error).message}\n${
+            (err as Error).stack
+          }`
+        );
+        response.write(
+          `error: ${JSON.stringify({
+            error: getCallableJSON(err),
+          })}${streamDelimiter}`
+        );
+        response.end();
+      },
+    });
+  } catch (e: any) {
+    if (e instanceof StreamNotFoundError) {
+      response.status(204).end();
+      return;
+    }
+    if (e.status === 'DEADLINE_EXCEEDED') {
+      response.write(
+        `error: ${JSON.stringify({
+          error: getCallableJSON(e),
+        })}${streamDelimiter}`
+      );
+      response.end();
+      return;
+    }
+    throw e;
+  }
+}
+
 /**
  * A wrapper object containing a flow with its associated auth policy.
+ * @deprecated Use `withFlowOptions` instead.
  */
 export type FlowWithContextProvider<
   C extends ActionContext = ActionContext,
@@ -161,8 +287,25 @@ export type FlowWithContextProvider<
   context: ContextProvider<C, I>;
 };
 
+/**
+ * A wrapper object containing a flow with its associated options.
+ */
+export type FlowWithOptions<
+  I extends z.ZodTypeAny = z.ZodTypeAny,
+  O extends z.ZodTypeAny = z.ZodTypeAny,
+  S extends z.ZodTypeAny = z.ZodTypeAny,
+> = {
+  flow: Flow<I, O, S>;
+  options: {
+    contextProvider?: ContextProvider<any, I>;
+    streamManager?: StreamManager;
+    path?: string;
+  };
+};
+
 /**
  * Adds an auth policy to the flow.
+ * @deprecated Use `withFlowOptions` instead.
  */
 export function withContextProvider<
   C extends ActionContext = ActionContext,
@@ -179,12 +322,37 @@ export function withContextProvider<
   };
 }
 
+/**
+ * Adds an auth policy to the flow.
+ */
+export function withFlowOptions<
+  I extends z.ZodTypeAny,
+  O extends z.ZodTypeAny,
+  S extends z.ZodTypeAny,
+>(
+  flow: Flow<I, O, S>,
+  options: {
+    contextProvider?: ContextProvider<any, I>;
+    streamManager?: StreamManager;
+    path?: string;
+  }
+): FlowWithOptions<I, O, S> {
+  return {
+    flow,
+    options,
+  };
+}
+
 /**
  * Options to configure the flow server.
  */
 export interface FlowServerOptions {
   /** List of flows to expose via the flow server. */
-  flows: (Flow<any, any, any> | FlowWithContextProvider<any, any, any>)[];
+  flows: (
+    | Flow<any, any, any>
+    | FlowWithContextProvider<any, any, any>
+    | FlowWithOptions<any, any, any>
+  )[];
   /** Port to run the server on. Defaults to env.PORT or 3400. */
   port?: number;
   /** CORS options for the server. */
@@ -240,13 +408,14 @@ export class FlowServer {
     logger.debug('Running flow server with flow paths:');
     const pathPrefix = this.options.pathPrefix ?? '';
     this.options.flows?.forEach((flow) => {
-      if ('context' in flow) {
-        const flowPath = `/${pathPrefix}${flow.flow.__action.name}`;
+      if ('flow' in flow) {
+        const flowPath = `/${pathPrefix}${
+          ('options' in flow && flow.options.path) || flow.flow.__action.name
+        }`;
         logger.debug(` - ${flowPath}`);
-        server.post(
-          flowPath,
-          expressHandler(flow.flow, { contextProvider: flow.context })
-        );
+        const options =
+          'options' in flow ? flow.options : { contextProvider: flow.context };
+        server.post(flowPath, expressHandler(flow.flow, options));
       } else {
         const flowPath = `/${pathPrefix}${flow.__action.name}`;
         logger.debug(` - ${flowPath}`);

package/tests/express_test.ts

@@ -23,6 +23,7 @@ import {
   type GenerateResponseData,
   type Genkit,
 } from 'genkit';
+import { InMemoryStreamManager } from 'genkit/beta';
 import { runFlow, streamFlow } from 'genkit/beta/client';
 import type { ContextProvider, RequestData } from 'genkit/context';
 import type { GenerateResponseChunkData, ModelAction } from 'genkit/model';
@@ -32,7 +33,7 @@ import { afterEach, beforeEach, describe, it } from 'node:test';
 import {
   expressHandler,
   startFlowServer,
-  withContextProvider,
+  withFlowOptions,
   type FlowServer,
 } from '../src/index.js';
 
@@ -144,6 +145,12 @@ describe('expressHandler', async () => {
     app.post('/stringInput', expressHandler(stringInput));
     app.post('/objectInput', expressHandler(objectInput));
     app.post('/streamingFlow', expressHandler(streamingFlow));
+    app.post(
+      '/streamingFlowDurable',
+      expressHandler(streamingFlow, {
+        streamManager: new InMemoryStreamManager(),
+      })
+    );
     app.post(
       '/flowWithAuth',
       expressHandler(flowWithAuth, { contextProvider })
@@ -333,6 +340,88 @@ describe('expressHandler', async () => {
       assert.strictEqual(await result.output, 'Echo: olleh');
     });
 
+    it('should create and subscribe to a durable stream', async () => {
+      const result = streamFlow({
+        url: `http://localhost:${port}/streamingFlowDurable`,
+        input: {
+          question: 'durable',
+        },
+      });
+
+      const streamId = await result.streamId;
+      assert.ok(streamId);
+
+      const subscription = streamFlow({
+        url: `http://localhost:${port}/streamingFlowDurable`,
+        input: {
+          question: 'durable',
+        },
+        streamId: streamId!,
+      });
+
+      const gotChunks: GenerateResponseChunkData[] = [];
+      for await (const chunk of subscription.stream) {
+        gotChunks.push(chunk);
+      }
+
+      // of note here is that we're consuming the original stream after re-subscription
+      // which should still work fine.
+      const originalChunks: GenerateResponseChunkData[] = [];
+      for await (const chunk of result.stream) {
+        originalChunks.push(chunk);
+      }
+
+      assert.deepStrictEqual(gotChunks, originalChunks);
+      assert.strictEqual(await subscription.output, 'Echo: durable');
+      assert.strictEqual(await result.output, 'Echo: durable');
+    });
+
+    it('should subscribe to a stream in progress', async () => {
+      const result = streamFlow({
+        url: `http://localhost:${port}/streamingFlowDurable`,
+        input: {
+          question: 'durable',
+        },
+      });
+
+      const streamId = await result.streamId;
+      assert.ok(streamId);
+
+      // Don't wait for the original stream to finish.
+      const subscription = streamFlow({
+        url: `http://localhost:${port}/streamingFlowDurable`,
+        input: {
+          question: 'durable',
+        },
+        streamId: streamId!,
+      });
+
+      const gotChunks: GenerateResponseChunkData[] = [];
+      for await (const chunk of subscription.stream) {
+        gotChunks.push(chunk);
+      }
+
+      assert.deepStrictEqual(gotChunks.length, 3);
+      assert.strictEqual(await subscription.output, 'Echo: durable');
+    });
+
+    it.only('should return 204 for a non-existent stream', async () => {
+      try {
+        const result = streamFlow({
+          url: `http://localhost:${port}/streamingFlowDurable`,
+          input: {
+            question: 'durable',
+          },
+          streamId: 'non-existent-stream-id',
+        });
+        for await (const _ of result.stream) {
+        }
+        assert.fail('should have thrown');
+      } catch (err: any) {
+        assert.strictEqual(err.message, 'NOT_FOUND: Stream not found.');
+      }
+    });
+
     it('stream a model', async () => {
       const result = streamFlow({
         url: `http://localhost:${port}/echoModel`,
@@ -447,8 +536,12 @@ describe('startFlowServer', async () => {
         stringInput,
         objectInput,
         streamingFlow,
+        withFlowOptions(streamingFlow, {
+          streamManager: new InMemoryStreamManager(),
+          path: 'streamingFlowDurable',
+        }),
         streamingFlowV2,
-        withContextProvider(flowWithAuth, contextProvider),
+        withFlowOptions(flowWithAuth, { contextProvider }),
       ],
       port,
     });
@@ -547,6 +640,88 @@ describe('startFlowServer', async () => {
 
       assert.strictEqual(await result.output, 'Echo: olleh');
     });
+
+    it('should create and subscribe to a durable stream', async () => {
+      const result = streamFlow({
+        url: `http://localhost:${port}/streamingFlowDurable`,
+        input: {
+          question: 'durable',
+        },
+      });
+
+      const streamId = await result.streamId;
+      assert.ok(streamId);
+
+      const subscription = streamFlow({
+        url: `http://localhost:${port}/streamingFlowDurable`,
+        input: {
+          question: 'durable',
+        },
+        streamId: streamId!,
+      });
+
+      const gotChunks: GenerateResponseChunkData[] = [];
+      for await (const chunk of subscription.stream) {
+        gotChunks.push(chunk);
+      }
+
+      // of note here is that we're consuming the original stream after re-subscription
+      // which should still work fine.
+      const originalChunks: GenerateResponseChunkData[] = [];
+      for await (const chunk of result.stream) {
+        originalChunks.push(chunk);
+      }
+
+      assert.deepStrictEqual(gotChunks, originalChunks);
+      assert.strictEqual(await subscription.output, 'Echo: durable');
+      assert.strictEqual(await result.output, 'Echo: durable');
+    });
+
+    it('should subscribe to a stream in progress', async () => {
+      const result = streamFlow({
+        url: `http://localhost:${port}/streamingFlowDurable`,
+        input: {
+          question: 'durable',
+        },
+      });
+
+      const streamId = await result.streamId;
+      assert.ok(streamId);
+
+      // Don't wait for the original stream to finish.
+      const subscription = streamFlow({
+        url: `http://localhost:${port}/streamingFlowDurable`,
+        input: {
+          question: 'durable',
+        },
+        streamId: streamId!,
+      });
+
+      const gotChunks: GenerateResponseChunkData[] = [];
+      for await (const chunk of subscription.stream) {
+        gotChunks.push(chunk);
+      }
+
+      assert.deepStrictEqual(gotChunks.length, 3);
+      assert.strictEqual(await subscription.output, 'Echo: durable');
+    });
+
+    it('should return 204 for a non-existent stream', async () => {
+      try {
+        const result = streamFlow({
+          url: `http://localhost:${port}/streamingFlowDurable`,
+          input: {
+            question: 'durable',
+          },
+          streamId: 'non-existent-stream-id',
+        });
+        for await (const _ of result.stream) {
+        }
+        assert.fail('should have thrown');
+      } catch (err: any) {
+        assert.strictEqual(err.message, 'NOT_FOUND: Stream not found.');
+      }
+    });
   });
 
   it('stream a flow (v2 model)', async () => {