@teamkeel/functions-runtime 0.421.2 → 0.421.3

package/dist/index.cjs

@@ -247,6 +247,8 @@ var import_pg = require("pg");
 // src/tracing.js
 var tracing_exports = {};
 __export(tracing_exports, {
+  KEEL_INTERNAL_ATTR: () => KEEL_INTERNAL_ATTR,
+  KEEL_INTERNAL_CHILDREN: () => KEEL_INTERNAL_CHILDREN,
   forceFlush: () => forceFlush,
   getTracer: () => getTracer,
   init: () => init,
@@ -263,11 +265,13 @@ async function withSpan(name, fn) {
     try {
       return await fn(span);
     } catch (err) {
-      span.recordException(err);
-      span.setStatus({
-        code: opentelemetry.SpanStatusCode.ERROR,
-        message: err.message
-      });
+      if (err instanceof Error) {
+        span.recordException(err);
+        span.setStatus({
+          code: opentelemetry.SpanStatusCode.ERROR,
+          message: err.message
+        });
+      }
       throw err;
     } finally {
       span.end();
@@ -405,6 +409,8 @@ function spanNameForModelAPI(modelName, action) {
   return `Database ${modelName}.${action}`;
 }
 __name(spanNameForModelAPI, "spanNameForModelAPI");
+var KEEL_INTERNAL_ATTR = "keel_internal";
+var KEEL_INTERNAL_CHILDREN = "includeChildrenSpans";
 
 // src/database.ts
 var import_ws = __toESM(require("ws"), 1);
@@ -473,6 +479,7 @@ var InstrumentedPool = class extends import_pg.Pool {
     const _super = super.connect.bind(this);
     return withSpan("Database Connect", function(span) {
       span.setAttribute("dialect", process.env["KEEL_DB_CONN_TYPE"]);
+      span.setAttribute(KEEL_INTERNAL_ATTR, true);
       return _super.apply(null, args);
     });
   }
@@ -485,6 +492,7 @@ var InstrumentedNeonServerlessPool = class extends neon.Pool {
     const _super = super.connect.bind(this);
     return withSpan("Database Connect", function(span) {
       span.setAttribute("dialect", process.env["KEEL_DB_CONN_TYPE"]);
+      span.setAttribute(KEEL_INTERNAL_ATTR, true);
       return _super.apply(null, args);
     });
   }
@@ -3036,112 +3044,19 @@ function createFlowContext(runId, data, action, callback, element, spanId, ctx)
       };
     }, "complete"),
     step: /* @__PURE__ */ __name(async (name, optionsOrFn, fn) => {
-      const options = typeof optionsOrFn === "function" ? {} : optionsOrFn;
-      const actualFn = typeof optionsOrFn === "function" ? optionsOrFn : fn;
-      options.retries = options.retries ?? defaultOpts.retries;
-      options.timeout = options.timeout ?? defaultOpts.timeout;
-      const db = useDatabase();
-      if (usedNames.has(name)) {
-        await db.insertInto("keel.flow_step").values({
-          run_id: runId,
-          name,
-          stage: options.stage,
-          status: "FAILED" /* FAILED */,
-          type: "FUNCTION" /* FUNCTION */,
-          error: `Duplicate step name: ${name}`,
-          startTime: /* @__PURE__ */ new Date(),
-          endTime: /* @__PURE__ */ new Date()
-        }).returningAll().executeTakeFirst();
-        throw new Error(`Duplicate step name: ${name}`);
-      }
-      usedNames.add(name);
-      const past = await db.selectFrom("keel.flow_step").where("run_id", "=", runId).where("name", "=", name).selectAll().execute();
-      const newSteps = past.filter((step) => step.status === "NEW" /* NEW */);
-      const completedSteps = past.filter(
-        (step) => step.status === "COMPLETED" /* COMPLETED */
-      );
-      const failedSteps = past.filter(
-        (step) => step.status === "FAILED" /* FAILED */
-      );
-      if (newSteps.length > 1) {
-        throw new Error("Multiple NEW steps found for the same step");
-      }
-      if (completedSteps.length > 1) {
-        throw new Error("Multiple completed steps found for the same step");
-      }
-      if (completedSteps.length > 1 && newSteps.length > 1) {
-        throw new Error(
-          "Multiple completed and new steps found for the same step"
-        );
-      }
-      if (completedSteps.length === 1) {
-        return completedSteps[0].value;
-      }
-      if (newSteps.length === 1) {
-        let result = null;
-        await db.updateTable("keel.flow_step").set({
-          startTime: /* @__PURE__ */ new Date()
-        }).where("id", "=", newSteps[0].id).returningAll().executeTakeFirst();
-        try {
-          const stepArgs = {
-            attempt: failedSteps.length + 1,
-            stepOptions: options
-          };
-          result = await withTimeout(actualFn(stepArgs), options.timeout);
-        } catch (e) {
-          await db.updateTable("keel.flow_step").set({
-            status: "FAILED" /* FAILED */,
-            spanId,
-            endTime: /* @__PURE__ */ new Date(),
-            error: e instanceof Error ? e.message : "An error occurred"
-          }).where("id", "=", newSteps[0].id).returningAll().executeTakeFirst();
-          if (failedSteps.length >= options.retries || e instanceof NonRetriableError) {
-            if (options.onFailure) {
-              await options.onFailure();
-            }
-            throw new ExhuastedRetriesDisrupt();
-          }
-          await db.insertInto("keel.flow_step").values({
-            run_id: runId,
-            name,
-            stage: options.stage,
-            status: "NEW" /* NEW */,
-            type: "FUNCTION" /* FUNCTION */
-          }).returningAll().executeTakeFirst();
-          throw new StepCreatedDisrupt(
-            options.retryPolicy ? new Date(
-              Date.now() + options.retryPolicy(failedSteps.length + 1)
-            ) : void 0
-          );
-        }
-        await db.updateTable("keel.flow_step").set({
-          status: "COMPLETED" /* COMPLETED */,
-          value: JSON.stringify(result),
-          spanId,
-          endTime: /* @__PURE__ */ new Date()
-        }).where("id", "=", newSteps[0].id).returningAll().executeTakeFirst();
-        return result;
-      }
-      await db.insertInto("keel.flow_step").values({
-        run_id: runId,
-        name,
-        stage: options.stage,
-        status: "NEW" /* NEW */,
-        type: "FUNCTION" /* FUNCTION */
-      }).returningAll().executeTakeFirst();
-      throw new StepCreatedDisrupt();
-    }, "step"),
-    ui: {
-      page: /* @__PURE__ */ __name(async (name, options) => {
+      return withSpan(`Step - ${name}`, async (span) => {
+        const options = typeof optionsOrFn === "function" ? {} : optionsOrFn;
+        const actualFn = typeof optionsOrFn === "function" ? optionsOrFn : fn;
+        options.retries = options.retries ?? defaultOpts.retries;
+        options.timeout = options.timeout ?? defaultOpts.timeout;
         const db = useDatabase();
-        const isCallback = element && callback;
         if (usedNames.has(name)) {
           await db.insertInto("keel.flow_step").values({
             run_id: runId,
             name,
             stage: options.stage,
             status: "FAILED" /* FAILED */,
-            type: "UI" /* UI */,
+            type: "FUNCTION" /* FUNCTION */,
             error: `Duplicate step name: ${name}`,
             startTime: /* @__PURE__ */ new Date(),
             endTime: /* @__PURE__ */ new Date()
@@ -3149,77 +3064,179 @@ function createFlowContext(runId, data, action, callback, element, spanId, ctx)
           throw new Error(`Duplicate step name: ${name}`);
         }
         usedNames.add(name);
-        let step = await db.selectFrom("keel.flow_step").where("run_id", "=", runId).where("name", "=", name).selectAll().executeTakeFirst();
-        if (step && step.status === "COMPLETED" /* COMPLETED */) {
-          if (step.action) {
-            return { data: step.value, action: step.action };
-          }
-          return step.value;
+        const past = await db.selectFrom("keel.flow_step").where("run_id", "=", runId).where("name", "=", name).selectAll().execute();
+        const newSteps = past.filter((step) => step.status === "NEW" /* NEW */);
+        const completedSteps = past.filter(
+          (step) => step.status === "COMPLETED" /* COMPLETED */
+        );
+        const failedSteps = past.filter(
+          (step) => step.status === "FAILED" /* FAILED */
+        );
+        if (newSteps.length > 1) {
+          throw new Error("Multiple NEW steps found for the same step");
         }
-        if (!step) {
-          step = await db.insertInto("keel.flow_step").values({
-            run_id: runId,
-            name,
-            stage: options.stage,
-            status: "PENDING" /* PENDING */,
-            type: "UI" /* UI */,
-            startTime: /* @__PURE__ */ new Date()
-          }).returningAll().executeTakeFirst();
-          throw new UIRenderDisrupt(
-            step?.id,
-            (await page(options, null, null)).page
+        if (completedSteps.length > 1) {
+          throw new Error("Multiple completed steps found for the same step");
+        }
+        if (completedSteps.length > 1 && newSteps.length > 1) {
+          throw new Error(
+            "Multiple completed and new steps found for the same step"
           );
         }
-        if (isCallback) {
+        if (completedSteps.length === 1) {
+          span.setAttribute(KEEL_INTERNAL_ATTR, KEEL_INTERNAL_CHILDREN);
+          return completedSteps[0].value;
+        }
+        if (newSteps.length === 1) {
+          let result = null;
+          await db.updateTable("keel.flow_step").set({
+            startTime: /* @__PURE__ */ new Date()
+          }).where("id", "=", newSteps[0].id).returningAll().executeTakeFirst();
           try {
-            const response = await callbackFn(
-              options.content,
-              element,
-              callback,
-              data
-            );
-            throw new CallbackDisrupt(response, false);
+            const stepArgs = {
+              attempt: failedSteps.length + 1,
+              stepOptions: options
+            };
+            result = await withTimeout(actualFn(stepArgs), options.timeout);
           } catch (e) {
-            if (e instanceof CallbackDisrupt) {
-              throw e;
+            await db.updateTable("keel.flow_step").set({
+              status: "FAILED" /* FAILED */,
+              spanId,
+              endTime: /* @__PURE__ */ new Date(),
+              error: e instanceof Error ? e.message : "An error occurred"
+            }).where("id", "=", newSteps[0].id).returningAll().executeTakeFirst();
+            if (failedSteps.length >= options.retries || e instanceof NonRetriableError) {
+              if (options.onFailure) {
+                await options.onFailure();
+              }
+              throw new ExhuastedRetriesDisrupt();
             }
-            throw new CallbackDisrupt(
-              e instanceof Error ? e.message : `An error occurred`,
-              true
+            await db.insertInto("keel.flow_step").values({
+              run_id: runId,
+              name,
+              stage: options.stage,
+              status: "NEW" /* NEW */,
+              type: "FUNCTION" /* FUNCTION */
+            }).returningAll().executeTakeFirst();
+            throw new StepCreatedDisrupt(
+              options.retryPolicy ? new Date(
+                Date.now() + options.retryPolicy(failedSteps.length + 1)
+              ) : void 0
             );
           }
+          await db.updateTable("keel.flow_step").set({
+            status: "COMPLETED" /* COMPLETED */,
+            value: JSON.stringify(result),
+            spanId,
+            endTime: /* @__PURE__ */ new Date()
+          }).where("id", "=", newSteps[0].id).returningAll().executeTakeFirst();
+          return result;
         }
-        if (!data) {
-          throw new UIRenderDisrupt(
-            step?.id,
-            (await page(options, null, null)).page
-          );
-        }
-        try {
-          const p = await page(options, data, action);
-          if (p.hasValidationErrors) {
-            throw new UIRenderDisrupt(step?.id, p.page);
+        await db.insertInto("keel.flow_step").values({
+          run_id: runId,
+          name,
+          stage: options.stage,
+          status: "NEW" /* NEW */,
+          type: "FUNCTION" /* FUNCTION */
+        }).returningAll().executeTakeFirst();
+        span.setAttribute(KEEL_INTERNAL_ATTR, KEEL_INTERNAL_CHILDREN);
+        throw new StepCreatedDisrupt();
+      });
+    }, "step"),
+    ui: {
+      page: /* @__PURE__ */ __name(async (name, options) => {
+        return withSpan(`Page - ${name}`, async (span) => {
+          const db = useDatabase();
+          const isCallback = element && callback;
+          if (usedNames.has(name)) {
+            await db.insertInto("keel.flow_step").values({
+              run_id: runId,
+              name,
+              stage: options.stage,
+              status: "FAILED" /* FAILED */,
+              type: "UI" /* UI */,
+              error: `Duplicate step name: ${name}`,
+              startTime: /* @__PURE__ */ new Date(),
+              endTime: /* @__PURE__ */ new Date()
+            }).returningAll().executeTakeFirst();
+            throw new Error(`Duplicate step name: ${name}`);
           }
-        } catch (e) {
-          if (e instanceof UIRenderDisrupt) {
+          usedNames.add(name);
+          let step = await db.selectFrom("keel.flow_step").where("run_id", "=", runId).where("name", "=", name).selectAll().executeTakeFirst();
+          if (step && step.status === "COMPLETED" /* COMPLETED */) {
+            span.setAttribute(KEEL_INTERNAL_ATTR, KEEL_INTERNAL_CHILDREN);
+            if (step.action) {
+              return { data: step.value, action: step.action };
+            }
+            return step.value;
+          }
+          if (!step) {
+            step = await db.insertInto("keel.flow_step").values({
+              run_id: runId,
+              name,
+              stage: options.stage,
+              status: "PENDING" /* PENDING */,
+              type: "UI" /* UI */,
+              startTime: /* @__PURE__ */ new Date()
+            }).returningAll().executeTakeFirst();
+            span.setAttribute("rendered", true);
+            throw new UIRenderDisrupt(
+              step?.id,
+              (await page(options, null, null)).page
+            );
+          }
+          if (isCallback) {
+            span.setAttribute("callback", callback);
+            try {
+              const response = await callbackFn(
+                options.content,
+                element,
+                callback,
+                data
+              );
+              throw new CallbackDisrupt(response, false);
+            } catch (e) {
+              if (e instanceof CallbackDisrupt) {
+                throw e;
+              }
+              throw new CallbackDisrupt(
+                e instanceof Error ? e.message : `An error occurred`,
+                true
+              );
+            }
+          }
+          if (!data) {
+            throw new UIRenderDisrupt(
+              step?.id,
+              (await page(options, null, null)).page
+            );
+          }
+          try {
+            const p = await page(options, data, action);
+            if (p.hasValidationErrors) {
+              throw new UIRenderDisrupt(step?.id, p.page);
+            }
+          } catch (e) {
+            if (e instanceof UIRenderDisrupt) {
+              throw e;
+            }
+            await db.updateTable("keel.flow_step").set({
+              status: "FAILED" /* FAILED */,
+              spanId,
+              endTime: /* @__PURE__ */ new Date(),
+              error: e instanceof Error ? e.message : "An error occurred"
+            }).where("id", "=", step?.id).returningAll().executeTakeFirst();
             throw e;
           }
           await db.updateTable("keel.flow_step").set({
-            status: "FAILED" /* FAILED */,
+            status: "COMPLETED" /* COMPLETED */,
+            value: JSON.stringify(data),
+            action,
             spanId,
-            endTime: /* @__PURE__ */ new Date(),
-            error: e instanceof Error ? e.message : "An error occurred"
-          }).where("id", "=", step?.id).returningAll().executeTakeFirst();
-          throw e;
-        }
-        await db.updateTable("keel.flow_step").set({
-          status: "COMPLETED" /* COMPLETED */,
-          value: JSON.stringify(data),
-          action,
-          spanId,
-          endTime: /* @__PURE__ */ new Date()
-        }).where("id", "=", step.id).returningAll().executeTakeFirst();
-        return { data, action };
+            endTime: /* @__PURE__ */ new Date()
+          }).where("id", "=", step.id).returningAll().executeTakeFirst();
+          return { data, action };
+        });
       }, "page"),
       inputs: {
         text: textInput,
@@ -3298,6 +3315,7 @@ async function handleFlow(request, config) {
   );
   return opentelemetry6.context.with(activeContext, () => {
     return withSpan(request.method, async (span) => {
+      span.setAttribute(KEEL_INTERNAL_ATTR, true);
       let db = null;
       let flowConfig = null;
       const runId = request.meta?.runId;
@@ -3355,6 +3373,7 @@ async function handleFlow(request, config) {
           });
         } catch (e) {
           if (e instanceof StepCreatedDisrupt) {
+            span.setAttribute(KEEL_INTERNAL_ATTR, KEEL_INTERNAL_CHILDREN);
             return (0, import_json_rpc_26.createJSONRPCSuccessResponse)(request.id, {
               runId,
               runCompleted: false,
@@ -3376,11 +3395,13 @@ async function handleFlow(request, config) {
               ui: e.contents
             });
           }
-          span.recordException(e);
-          span.setStatus({
-            code: opentelemetry6.SpanStatusCode.ERROR,
-            message: e instanceof Error ? e.message : "unknown error"
-          });
+          if (e instanceof Error) {
+            span.recordException(e);
+            span.setStatus({
+              code: opentelemetry6.SpanStatusCode.ERROR,
+              message: e instanceof Error ? e.message : "unknown error"
+            });
+          }
           if (e instanceof ExhuastedRetriesDisrupt) {
             return (0, import_json_rpc_26.createJSONRPCSuccessResponse)(request.id, {
               runId,