{"version":3,"sources":["../src/background-action.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 type { JSONSchema7 } from 'json-schema';\nimport * as z from 'zod';\nimport { Action, ActionMetadata, Middleware, action } from './action.mjs';\nimport { ActionContext } from './context.mjs';\nimport { GenkitError } from './error.mjs';\nimport { ActionType, Registry } from './registry.mjs';\nimport { toJsonSchema } from './schema.mjs';\n\n/**\n * Zod schema of an opration representing a background task.\n */\nexport const OperationSchema = z.object({\n  action: z.string().optional(),\n  id: z.string(),\n  done: z.boolean().optional(),\n  output: z.any().optional(),\n  error: z.object({ message: z.string() }).passthrough().optional(),\n  metadata: z.record(z.string(), z.any()).optional(),\n});\n\n/**\n * Background operation.\n */\nexport interface Operation<O = any> {\n  action?: string;\n  id: string;\n  done?: boolean;\n  output?: O;\n  error?: { message: string; [key: string]: unknown };\n  metadata?: Record<string, any>;\n}\n\n/**\n * Background action. Unlike regular action, background action can run for a long time in the background.\n * The returned operation can used to check the status of the background operation and retrieve the response.\n */\nexport interface BackgroundAction<\n  I extends z.ZodTypeAny = z.ZodTypeAny,\n  O extends z.ZodTypeAny = z.ZodTypeAny,\n  RunOptions extends BackgroundActionRunOptions = BackgroundActionRunOptions,\n> {\n  __action: ActionMetadata<I, O>;\n\n  readonly startAction: Action<I, typeof OperationSchema>;\n  readonly checkAction: Action<typeof OperationSchema, typeof OperationSchema>;\n  readonly cancelAction?: Action<\n    typeof OperationSchema,\n    typeof OperationSchema\n  >;\n  readonly supportsCancel: boolean;\n\n  start(\n    input?: z.infer<I>,\n    options?: RunOptions\n  ): Promise<Operation<z.infer<O>>>;\n\n  check(operation: Operation<z.infer<O>>): Promise<Operation<z.infer<O>>>;\n\n  cancel(operation: Operation<z.infer<O>>): Promise<Operation<z.infer<O>>>;\n}\n\nexport async function lookupBackgroundAction<\n  I extends z.ZodTypeAny = z.ZodTypeAny,\n  O extends z.ZodTypeAny = z.ZodTypeAny,\n>(\n  registry: Registry,\n  key: string\n): Promise<BackgroundAction<I, O> | undefined> {\n  const root: Action<I, typeof OperationSchema> = await registry.lookupAction<\n    I,\n    typeof OperationSchema,\n    Action<I, typeof OperationSchema>\n  >(key);\n  if (!root) return undefined;\n  const actionName = key.substring(key.indexOf('/', 1) + 1);\n  return new BackgroundActionImpl(\n    root,\n    await registry.lookupAction<\n      typeof OperationSchema,\n      typeof OperationSchema,\n      Action<typeof OperationSchema, typeof OperationSchema>\n    >(`/check-operation/${actionName}/check`),\n    await registry.lookupAction<\n      typeof OperationSchema,\n      typeof OperationSchema,\n      Action<typeof OperationSchema, typeof OperationSchema>\n    >(`/cancel-operation/${actionName}/cancel`)\n  );\n}\n\nclass BackgroundActionImpl<\n  I extends z.ZodTypeAny = z.ZodTypeAny,\n  O extends z.ZodTypeAny = z.ZodTypeAny,\n  RunOptions extends BackgroundActionRunOptions = BackgroundActionRunOptions,\n> implements BackgroundAction<I, O, RunOptions>\n{\n  __action: ActionMetadata<I, O>;\n\n  readonly startAction: Action<I, typeof OperationSchema>;\n  readonly checkAction: Action<typeof OperationSchema, typeof OperationSchema>;\n  readonly cancelAction?: Action<\n    typeof OperationSchema,\n    typeof OperationSchema\n  >;\n\n  constructor(\n    startAction: Action<I, typeof OperationSchema>,\n    checkAction: Action<typeof OperationSchema, typeof OperationSchema>,\n    cancelAction:\n      | Action<typeof OperationSchema, typeof OperationSchema>\n      | undefined\n  ) {\n    this.__action = {\n      name: startAction.__action.name,\n      description: startAction.__action.description,\n      inputSchema: startAction.__action.inputSchema,\n      inputJsonSchema: startAction.__action.inputJsonSchema,\n      metadata: startAction.__action.metadata,\n      actionType: startAction.__action.actionType,\n    };\n    this.startAction = startAction;\n    this.checkAction = checkAction;\n    this.cancelAction = cancelAction;\n  }\n\n  async start(\n    input?: z.infer<I>,\n    options?: RunOptions\n  ): Promise<Operation<z.infer<O>>> {\n    return await this.startAction(input, options);\n  }\n\n  async check(\n    operation: Operation<z.infer<O>>\n  ): Promise<Operation<z.infer<O>>> {\n    return await this.checkAction(operation);\n  }\n\n  get supportsCancel(): boolean {\n    return !!this.cancelAction;\n  }\n\n  async cancel(\n    operation: Operation<z.infer<O>>\n  ): Promise<Operation<z.infer<O>>> {\n    if (!this.cancelAction) {\n      return operation;\n    }\n    return await this.cancelAction(operation);\n  }\n}\n\n/**\n * Options (side channel) data to pass to the model.\n */\nexport interface BackgroundActionRunOptions {\n  /**\n   * Additional runtime context data (ex. auth context data).\n   */\n  context?: ActionContext;\n\n  /**\n   * Additional span attributes to apply to OT spans.\n   */\n  telemetryLabels?: Record<string, string>;\n}\n\n/**\n * Options (side channel) data to pass to the model.\n */\nexport interface BackgroundActionFnArg<S> {\n  /**\n   * Additional runtime context data (ex. auth context data).\n   */\n  context?: ActionContext;\n\n  /**\n   * Trace context containing trace and span IDs.\n   */\n  trace: {\n    traceId: string;\n    spanId: string;\n  };\n}\n\n/**\n * Action factory params.\n */\nexport type BackgroundActionParams<\n  I extends z.ZodTypeAny,\n  O extends z.ZodTypeAny,\n  S extends z.ZodTypeAny = z.ZodTypeAny,\n> = {\n  name: string;\n  start: (\n    input: z.infer<I>,\n    options: BackgroundActionFnArg<z.infer<S>>\n  ) => Promise<Operation<z.infer<O>>>;\n  check: (input: Operation<z.infer<O>>) => Promise<Operation<z.infer<O>>>;\n  cancel?: (input: Operation<z.infer<O>>) => Promise<Operation<z.infer<O>>>;\n  actionType: ActionType;\n\n  description?: string;\n  inputSchema?: I;\n  inputJsonSchema?: JSONSchema7;\n  outputSchema?: O;\n  outputJsonSchema?: JSONSchema7;\n  metadata?: Record<string, any>;\n  use?: Middleware<z.infer<I>, z.infer<O>>[];\n  streamSchema?: S;\n};\n\n/**\n * Defines an action with the given config and registers it in the registry.\n */\nexport function defineBackgroundAction<\n  I extends z.ZodTypeAny,\n  O extends z.ZodTypeAny,\n  S extends z.ZodTypeAny = z.ZodTypeAny,\n>(\n  registry: Registry,\n  config: BackgroundActionParams<I, O, S>\n): BackgroundAction<I, O> {\n  const act = backgroundAction(config);\n\n  registerBackgroundAction(registry, act);\n\n  return act;\n}\n\nexport function registerBackgroundAction(\n  registry: Registry,\n  act: BackgroundAction<any, any>,\n  opts?: { namespace?: string }\n) {\n  registry.registerAction(\n    act.startAction.__action.actionType!,\n    act.startAction,\n    opts\n  );\n  registry.registerAction(\n    act.checkAction.__action.actionType!,\n    act.checkAction,\n    opts\n  );\n  if (act.cancelAction) {\n    registry.registerAction(\n      act.cancelAction.__action.actionType!,\n      act.cancelAction,\n      opts\n    );\n  }\n}\n\n/**\n * Creates a background action with the given config.\n */\nexport function backgroundAction<\n  I extends z.ZodTypeAny,\n  O extends z.ZodTypeAny,\n  S extends z.ZodTypeAny = z.ZodTypeAny,\n>(config: BackgroundActionParams<I, O, S>): BackgroundAction<I, O> {\n  const startAction = action(\n    {\n      actionType: config.actionType,\n      name: config.name,\n      description: config.description,\n      inputSchema: config.inputSchema,\n      inputJsonSchema: config.inputJsonSchema,\n      outputSchema: OperationSchema,\n      metadata: {\n        ...config.metadata,\n        outputSchema: toJsonSchema({\n          schema: config.outputSchema,\n          jsonSchema: config.outputJsonSchema,\n        }),\n      },\n      use: config.use,\n    },\n    async (input, options) => {\n      const operation = await config.start(input, options);\n      operation.action = `/${config.actionType}/${config.name}`;\n      return operation;\n    }\n  );\n  const checkAction = action(\n    {\n      actionType: 'check-operation',\n      name: `${config.name}/check`,\n      description: config.description,\n      inputSchema: OperationSchema,\n      inputJsonSchema: config.inputJsonSchema,\n      outputSchema: OperationSchema,\n      metadata: {\n        ...config.metadata,\n        outputSchema: toJsonSchema({\n          schema: config.outputSchema,\n          jsonSchema: config.outputJsonSchema,\n        }),\n      },\n    },\n    async (input) => {\n      const operation = await config.check(input);\n      operation.action = `/${config.actionType}/${config.name}`;\n      return operation;\n    }\n  );\n  let cancelAction:\n    | Action<typeof OperationSchema, typeof OperationSchema>\n    | undefined = undefined;\n  if (config.cancel) {\n    cancelAction = action(\n      {\n        actionType: 'cancel-operation',\n        name: `${config.name}/cancel`,\n        description: config.description,\n        inputSchema: OperationSchema,\n        inputJsonSchema: config.inputJsonSchema,\n        outputSchema: OperationSchema,\n        metadata: {\n          ...config.metadata,\n          outputSchema: toJsonSchema({\n            schema: config.outputSchema,\n            jsonSchema: config.outputJsonSchema,\n          }),\n        },\n      },\n      async (input) => {\n        if (!config.cancel) {\n          throw new GenkitError({\n            status: 'UNAVAILABLE',\n            message: `${config.name} does not support cancellation.`,\n          });\n        }\n        const operation = await config.cancel(input);\n        operation.action = `/${config.actionType}/${config.name}`;\n        return operation;\n      }\n    );\n  }\n\n  return new BackgroundActionImpl(startAction, checkAction, cancelAction);\n}\n\nexport function isBackgroundAction(a: unknown): a is BackgroundAction {\n  return a instanceof BackgroundActionImpl;\n}\n"],"mappings":"AAiBA,YAAY,OAAO;AACnB,SAA6C,cAAc;AAE3D,SAAS,mBAAmB;AAE5B,SAAS,oBAAoB;AAKtB,MAAM,kBAAkB,EAAE,OAAO;AAAA,EACtC,QAAQ,EAAE,OAAO,EAAE,SAAS;AAAA,EAC5B,IAAI,EAAE,OAAO;AAAA,EACb,MAAM,EAAE,QAAQ,EAAE,SAAS;AAAA,EAC3B,QAAQ,EAAE,IAAI,EAAE,SAAS;AAAA,EACzB,OAAO,EAAE,OAAO,EAAE,SAAS,EAAE,OAAO,EAAE,CAAC,EAAE,YAAY,EAAE,SAAS;AAAA,EAChE,UAAU,EAAE,OAAO,EAAE,OAAO,GAAG,EAAE,IAAI,CAAC,EAAE,SAAS;AACnD,CAAC;AA2CD,eAAsB,uBAIpB,UACA,KAC6C;AAC7C,QAAM,OAA0C,MAAM,SAAS,aAI7D,GAAG;AACL,MAAI,CAAC,KAAM,QAAO;AAClB,QAAM,aAAa,IAAI,UAAU,IAAI,QAAQ,KAAK,CAAC,IAAI,CAAC;AACxD,SAAO,IAAI;AAAA,IACT;AAAA,IACA,MAAM,SAAS,aAIb,oBAAoB,UAAU,QAAQ;AAAA,IACxC,MAAM,SAAS,aAIb,qBAAqB,UAAU,SAAS;AAAA,EAC5C;AACF;AAEA,MAAM,qBAKN;AAAA,EACE;AAAA,EAES;AAAA,EACA;AAAA,EACA;AAAA,EAKT,YACE,aACA,aACA,cAGA;AACA,SAAK,WAAW;AAAA,MACd,MAAM,YAAY,SAAS;AAAA,MAC3B,aAAa,YAAY,SAAS;AAAA,MAClC,aAAa,YAAY,SAAS;AAAA,MAClC,iBAAiB,YAAY,SAAS;AAAA,MACtC,UAAU,YAAY,SAAS;AAAA,MAC/B,YAAY,YAAY,SAAS;AAAA,IACnC;AACA,SAAK,cAAc;AACnB,SAAK,cAAc;AACnB,SAAK,eAAe;AAAA,EACtB;AAAA,EAEA,MAAM,MACJ,OACA,SACgC;AAChC,WAAO,MAAM,KAAK,YAAY,OAAO,OAAO;AAAA,EAC9C;AAAA,EAEA,MAAM,MACJ,WACgC;AAChC,WAAO,MAAM,KAAK,YAAY,SAAS;AAAA,EACzC;AAAA,EAEA,IAAI,iBAA0B;AAC5B,WAAO,CAAC,CAAC,KAAK;AAAA,EAChB;AAAA,EAEA,MAAM,OACJ,WACgC;AAChC,QAAI,CAAC,KAAK,cAAc;AACtB,aAAO;AAAA,IACT;AACA,WAAO,MAAM,KAAK,aAAa,SAAS;AAAA,EAC1C;AACF;AAiEO,SAAS,uBAKd,UACA,QACwB;AACxB,QAAM,MAAM,iBAAiB,MAAM;AAEnC,2BAAyB,UAAU,GAAG;AAEtC,SAAO;AACT;AAEO,SAAS,yBACd,UACA,KACA,MACA;AACA,WAAS;AAAA,IACP,IAAI,YAAY,SAAS;AAAA,IACzB,IAAI;AAAA,IACJ;AAAA,EACF;AACA,WAAS;AAAA,IACP,IAAI,YAAY,SAAS;AAAA,IACzB,IAAI;AAAA,IACJ;AAAA,EACF;AACA,MAAI,IAAI,cAAc;AACpB,aAAS;AAAA,MACP,IAAI,aAAa,SAAS;AAAA,MAC1B,IAAI;AAAA,MACJ;AAAA,IACF;AAAA,EACF;AACF;AAKO,SAAS,iBAId,QAAiE;AACjE,QAAM,cAAc;AAAA,IAClB;AAAA,MACE,YAAY,OAAO;AAAA,MACnB,MAAM,OAAO;AAAA,MACb,aAAa,OAAO;AAAA,MACpB,aAAa,OAAO;AAAA,MACpB,iBAAiB,OAAO;AAAA,MACxB,cAAc;AAAA,MACd,UAAU;AAAA,QACR,GAAG,OAAO;AAAA,QACV,cAAc,aAAa;AAAA,UACzB,QAAQ,OAAO;AAAA,UACf,YAAY,OAAO;AAAA,QACrB,CAAC;AAAA,MACH;AAAA,MACA,KAAK,OAAO;AAAA,IACd;AAAA,IACA,OAAO,OAAO,YAAY;AACxB,YAAM,YAAY,MAAM,OAAO,MAAM,OAAO,OAAO;AACnD,gBAAU,SAAS,IAAI,OAAO,UAAU,IAAI,OAAO,IAAI;AACvD,aAAO;AAAA,IACT;AAAA,EACF;AACA,QAAM,cAAc;AAAA,IAClB;AAAA,MACE,YAAY;AAAA,MACZ,MAAM,GAAG,OAAO,IAAI;AAAA,MACpB,aAAa,OAAO;AAAA,MACpB,aAAa;AAAA,MACb,iBAAiB,OAAO;AAAA,MACxB,cAAc;AAAA,MACd,UAAU;AAAA,QACR,GAAG,OAAO;AAAA,QACV,cAAc,aAAa;AAAA,UACzB,QAAQ,OAAO;AAAA,UACf,YAAY,OAAO;AAAA,QACrB,CAAC;AAAA,MACH;AAAA,IACF;AAAA,IACA,OAAO,UAAU;AACf,YAAM,YAAY,MAAM,OAAO,MAAM,KAAK;AAC1C,gBAAU,SAAS,IAAI,OAAO,UAAU,IAAI,OAAO,IAAI;AACvD,aAAO;AAAA,IACT;AAAA,EACF;AACA,MAAI,eAEY;AAChB,MAAI,OAAO,QAAQ;AACjB,mBAAe;AAAA,MACb;AAAA,QACE,YAAY;AAAA,QACZ,MAAM,GAAG,OAAO,IAAI;AAAA,QACpB,aAAa,OAAO;AAAA,QACpB,aAAa;AAAA,QACb,iBAAiB,OAAO;AAAA,QACxB,cAAc;AAAA,QACd,UAAU;AAAA,UACR,GAAG,OAAO;AAAA,UACV,cAAc,aAAa;AAAA,YACzB,QAAQ,OAAO;AAAA,YACf,YAAY,OAAO;AAAA,UACrB,CAAC;AAAA,QACH;AAAA,MACF;AAAA,MACA,OAAO,UAAU;AACf,YAAI,CAAC,OAAO,QAAQ;AAClB,gBAAM,IAAI,YAAY;AAAA,YACpB,QAAQ;AAAA,YACR,SAAS,GAAG,OAAO,IAAI;AAAA,UACzB,CAAC;AAAA,QACH;AACA,cAAM,YAAY,MAAM,OAAO,OAAO,KAAK;AAC3C,kBAAU,SAAS,IAAI,OAAO,UAAU,IAAI,OAAO,IAAI;AACvD,eAAO;AAAA,MACT;AAAA,IACF;AAAA,EACF;AAEA,SAAO,IAAI,qBAAqB,aAAa,aAAa,YAAY;AACxE;AAEO,SAAS,mBAAmB,GAAmC;AACpE,SAAO,aAAa;AACtB;","names":[]}