{"version":3,"sources":["../src/streaming.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 { GenkitError } from './error.mjs';\n\n/**\n * Error thrown when a stream cannot be found.\n */\nexport class StreamNotFoundError extends GenkitError {\n  constructor(message: string) {\n    super({ status: 'NOT_FOUND', message });\n    this.name = 'StreamNotFoundError';\n  }\n}\n\n/**\n * Interface for writing content to a stream.\n * @template S The type of the stream chunks.\n * @template O The type of the final output.\n */\nexport interface ActionStreamInput<S, O> {\n  /**\n   * Writes a chunk to the stream.\n   * @param chunk The chunk data to write.\n   */\n  write(chunk: S): Promise<void>;\n  /**\n   * Closes the stream with a final output.\n   * @param output The final output data.\n   */\n  done(output: O): Promise<void>;\n  /**\n   * Closes the stream with an error.\n   * @param err The error that occurred.\n   */\n  error(err: any): Promise<void>;\n}\n\n/**\n * Subscriber callbacks for receiving stream events.\n * @template S The type of the stream chunks.\n * @template O The type of the final output.\n */\nexport type ActionStreamSubscriber<S, O> = {\n  /** Called when a chunk is received. */\n  onChunk: (chunk: S) => void;\n  /** Called when the stream completes successfully. */\n  onDone: (output: O) => void;\n  /** Called when the stream encounters an error. */\n  onError: (error: any) => void;\n};\n\n/**\n * Interface for managing streaming actions, allowing creation and subscription to streams.\n * Implementations can provide different storage backends (e.g., in-memory, database, cache).\n */\nexport interface StreamManager {\n  /**\n   * Opens a stream for writing.\n   * @param streamId The unique identifier for the stream.\n   * @returns An object to write to the stream.\n   */\n  open<S, O>(streamId: string): Promise<ActionStreamInput<S, O>>;\n  /**\n   * Subscribes to a stream to receive its events.\n   * @param streamId The unique identifier for the stream.\n   * @param options The subscriber callbacks.\n   * @returns A promise resolving to an object containing an unsubscribe function.\n   */\n  subscribe<S, O>(\n    streamId: string,\n    options: ActionStreamSubscriber<S, O>\n  ): Promise<{ unsubscribe: () => void }>;\n}\n\ntype StreamState<S, O> =\n  | {\n      status: 'open';\n      chunks: S[];\n      subscribers: ActionStreamSubscriber<S, O>[];\n      lastTouched: number;\n    }\n  | { status: 'done'; chunks: S[]; output: O; lastTouched: number }\n  | { status: 'error'; chunks: S[]; error: any; lastTouched: number };\n\n/**\n * An in-memory implementation of StreamManager.\n * Useful for testing or single-instance deployments where persistence is not required.\n */\nexport class InMemoryStreamManager implements StreamManager {\n  private streams: Map<string, StreamState<any, any>> = new Map();\n\n  /**\n   * @param options Configuration options.\n   * @param options.ttlSeconds Time-to-live for streams in seconds. Defaults to 5 minutes.\n   */\n  constructor(private options: { ttlSeconds?: number } = {}) {}\n\n  private _cleanup() {\n    const ttl = (this.options.ttlSeconds ?? 5 * 60) * 1000;\n    const now = Date.now();\n    for (const [streamId, stream] of this.streams.entries()) {\n      if (stream.status !== 'open' && now - stream.lastTouched > ttl) {\n        this.streams.delete(streamId);\n      }\n    }\n  }\n\n  async open<S, O>(streamId: string): Promise<ActionStreamInput<S, O>> {\n    this._cleanup();\n    if (this.streams.has(streamId)) {\n      throw new Error(`Stream with id ${streamId} already exists.`);\n    }\n    this.streams.set(streamId, {\n      status: 'open',\n      chunks: [],\n      subscribers: [],\n      lastTouched: Date.now(),\n    });\n\n    return {\n      write: async (chunk: S) => {\n        const stream = this.streams.get(streamId);\n        if (stream?.status === 'open') {\n          stream.chunks.push(chunk);\n          stream.subscribers.forEach((s) => s.onChunk(chunk));\n          stream.lastTouched = Date.now();\n        }\n      },\n      done: async (output: O) => {\n        const stream = this.streams.get(streamId);\n        if (stream?.status === 'open') {\n          this.streams.set(streamId, {\n            status: 'done',\n            chunks: stream.chunks,\n            output,\n            lastTouched: Date.now(),\n          });\n          stream.subscribers.forEach((s) => s.onDone(output));\n        }\n      },\n      error: async (err: any) => {\n        const stream = this.streams.get(streamId);\n        if (stream?.status === 'open') {\n          stream.subscribers.forEach((s) => s.onError(err));\n          this.streams.set(streamId, {\n            status: 'error',\n            chunks: stream.chunks,\n            error: err,\n            lastTouched: Date.now(),\n          });\n        }\n      },\n    };\n  }\n\n  async subscribe<S, O>(\n    streamId: string,\n    subscriber: ActionStreamSubscriber<S, O>\n  ): Promise<{ unsubscribe: () => void }> {\n    const stream = this.streams.get(streamId);\n    if (!stream) {\n      throw new StreamNotFoundError(`Stream with id ${streamId} not found.`);\n    }\n\n    if (stream.status === 'done') {\n      for (const chunk of stream.chunks) {\n        subscriber.onChunk(chunk);\n      }\n      subscriber.onDone(stream.output);\n    } else if (stream.status === 'error') {\n      for (const chunk of stream.chunks) {\n        subscriber.onChunk(chunk);\n      }\n      subscriber.onError(stream.error);\n    } else {\n      stream.chunks.forEach((chunk) => subscriber.onChunk(chunk));\n      stream.subscribers.push(subscriber);\n    }\n\n    return {\n      unsubscribe: () => {\n        const currentStream = this.streams.get(streamId);\n        if (currentStream?.status === 'open') {\n          const index = currentStream.subscribers.indexOf(subscriber);\n          if (index > -1) {\n            currentStream.subscribers.splice(index, 1);\n          }\n        }\n      },\n    };\n  }\n}\n"],"mappings":"AAgBA,SAAS,mBAAmB;AAKrB,MAAM,4BAA4B,YAAY;AAAA,EACnD,YAAY,SAAiB;AAC3B,UAAM,EAAE,QAAQ,aAAa,QAAQ,CAAC;AACtC,SAAK,OAAO;AAAA,EACd;AACF;AA4EO,MAAM,sBAA+C;AAAA;AAAA;AAAA;AAAA;AAAA,EAO1D,YAAoB,UAAmC,CAAC,GAAG;AAAvC;AAAA,EAAwC;AAAA,EANpD,UAA8C,oBAAI,IAAI;AAAA,EAQtD,WAAW;AACjB,UAAM,OAAO,KAAK,QAAQ,cAAc,IAAI,MAAM;AAClD,UAAM,MAAM,KAAK,IAAI;AACrB,eAAW,CAAC,UAAU,MAAM,KAAK,KAAK,QAAQ,QAAQ,GAAG;AACvD,UAAI,OAAO,WAAW,UAAU,MAAM,OAAO,cAAc,KAAK;AAC9D,aAAK,QAAQ,OAAO,QAAQ;AAAA,MAC9B;AAAA,IACF;AAAA,EACF;AAAA,EAEA,MAAM,KAAW,UAAoD;AACnE,SAAK,SAAS;AACd,QAAI,KAAK,QAAQ,IAAI,QAAQ,GAAG;AAC9B,YAAM,IAAI,MAAM,kBAAkB,QAAQ,kBAAkB;AAAA,IAC9D;AACA,SAAK,QAAQ,IAAI,UAAU;AAAA,MACzB,QAAQ;AAAA,MACR,QAAQ,CAAC;AAAA,MACT,aAAa,CAAC;AAAA,MACd,aAAa,KAAK,IAAI;AAAA,IACxB,CAAC;AAED,WAAO;AAAA,MACL,OAAO,OAAO,UAAa;AACzB,cAAM,SAAS,KAAK,QAAQ,IAAI,QAAQ;AACxC,YAAI,QAAQ,WAAW,QAAQ;AAC7B,iBAAO,OAAO,KAAK,KAAK;AACxB,iBAAO,YAAY,QAAQ,CAAC,MAAM,EAAE,QAAQ,KAAK,CAAC;AAClD,iBAAO,cAAc,KAAK,IAAI;AAAA,QAChC;AAAA,MACF;AAAA,MACA,MAAM,OAAO,WAAc;AACzB,cAAM,SAAS,KAAK,QAAQ,IAAI,QAAQ;AACxC,YAAI,QAAQ,WAAW,QAAQ;AAC7B,eAAK,QAAQ,IAAI,UAAU;AAAA,YACzB,QAAQ;AAAA,YACR,QAAQ,OAAO;AAAA,YACf;AAAA,YACA,aAAa,KAAK,IAAI;AAAA,UACxB,CAAC;AACD,iBAAO,YAAY,QAAQ,CAAC,MAAM,EAAE,OAAO,MAAM,CAAC;AAAA,QACpD;AAAA,MACF;AAAA,MACA,OAAO,OAAO,QAAa;AACzB,cAAM,SAAS,KAAK,QAAQ,IAAI,QAAQ;AACxC,YAAI,QAAQ,WAAW,QAAQ;AAC7B,iBAAO,YAAY,QAAQ,CAAC,MAAM,EAAE,QAAQ,GAAG,CAAC;AAChD,eAAK,QAAQ,IAAI,UAAU;AAAA,YACzB,QAAQ;AAAA,YACR,QAAQ,OAAO;AAAA,YACf,OAAO;AAAA,YACP,aAAa,KAAK,IAAI;AAAA,UACxB,CAAC;AAAA,QACH;AAAA,MACF;AAAA,IACF;AAAA,EACF;AAAA,EAEA,MAAM,UACJ,UACA,YACsC;AACtC,UAAM,SAAS,KAAK,QAAQ,IAAI,QAAQ;AACxC,QAAI,CAAC,QAAQ;AACX,YAAM,IAAI,oBAAoB,kBAAkB,QAAQ,aAAa;AAAA,IACvE;AAEA,QAAI,OAAO,WAAW,QAAQ;AAC5B,iBAAW,SAAS,OAAO,QAAQ;AACjC,mBAAW,QAAQ,KAAK;AAAA,MAC1B;AACA,iBAAW,OAAO,OAAO,MAAM;AAAA,IACjC,WAAW,OAAO,WAAW,SAAS;AACpC,iBAAW,SAAS,OAAO,QAAQ;AACjC,mBAAW,QAAQ,KAAK;AAAA,MAC1B;AACA,iBAAW,QAAQ,OAAO,KAAK;AAAA,IACjC,OAAO;AACL,aAAO,OAAO,QAAQ,CAAC,UAAU,WAAW,QAAQ,KAAK,CAAC;AAC1D,aAAO,YAAY,KAAK,UAAU;AAAA,IACpC;AAEA,WAAO;AAAA,MACL,aAAa,MAAM;AACjB,cAAM,gBAAgB,KAAK,QAAQ,IAAI,QAAQ;AAC/C,YAAI,eAAe,WAAW,QAAQ;AACpC,gBAAM,QAAQ,cAAc,YAAY,QAAQ,UAAU;AAC1D,cAAI,QAAQ,IAAI;AACd,0BAAc,YAAY,OAAO,OAAO,CAAC;AAAA,UAC3C;AAAA,QACF;AAAA,MACF;AAAA,IACF;AAAA,EACF;AACF;","names":[]}