@teamkeel/functions-runtime 0.449.3 → 0.451.0

package/dist/index.cjs

@@ -50,6 +50,7 @@ __export(index_exports, {
   TaskAPI: () => TaskAPI,
   checkBuiltInPermissions: () => checkBuiltInPermissions,
   createFlowContext: () => createFlowContext,
+  createTraceAPI: () => createTraceAPI2,
   handleFlow: () => handleFlow,
   handleJob: () => handleJob,
   handleRequest: () => handleRequest,
@@ -251,17 +252,20 @@ var tracing_exports = {};
 __export(tracing_exports, {
   KEEL_INTERNAL_ATTR: () => KEEL_INTERNAL_ATTR,
   KEEL_INTERNAL_CHILDREN: () => KEEL_INTERNAL_CHILDREN,
+  createTraceAPI: () => createTraceAPI,
   forceFlush: () => forceFlush,
   getTracer: () => getTracer,
   init: () => init,
   spanNameForModelAPI: () => spanNameForModelAPI,
-  withSpan: () => withSpan
+  withSpan: () => withSpan,
+  withUserSpan: () => withUserSpan
 });
 var opentelemetry = __toESM(require("@opentelemetry/api"), 1);
 var import_sdk_trace_base = require("@opentelemetry/sdk-trace-base");
 var import_exporter_trace_otlp_proto = require("@opentelemetry/exporter-trace-otlp-proto");
 var import_sdk_trace_node = require("@opentelemetry/sdk-trace-node");
 var import_resources = require("@opentelemetry/resources");
+var userSpanContextKey = opentelemetry.createContextKey("keel.userSpan");
 async function withSpan(name, fn) {
   return getTracer().startActiveSpan(name, async (span) => {
     try {
@@ -281,6 +285,73 @@ async function withSpan(name, fn) {
   });
 }
 __name(withSpan, "withSpan");
+function withUserSpan(span, fn) {
+  const ctx = opentelemetry.context.active().setValue(userSpanContextKey, span);
+  return opentelemetry.context.with(ctx, fn);
+}
+__name(withUserSpan, "withUserSpan");
+function getUserSpan() {
+  return opentelemetry.context.active().getValue(userSpanContextKey);
+}
+__name(getUserSpan, "getUserSpan");
+function resolveTraceSpan(defaultSpan) {
+  const span = getUserSpan() ?? defaultSpan;
+  if (!span) {
+    throw new Error("no user span available for ctx.trace");
+  }
+  return span;
+}
+__name(resolveTraceSpan, "resolveTraceSpan");
+function normalizeTraceAttributeValue(value) {
+  if (value === void 0) {
+    return void 0;
+  }
+  if (typeof value === "string" || typeof value === "number" || typeof value === "boolean") {
+    return value;
+  }
+  if (Array.isArray(value)) {
+    if (value.length === 0) {
+      return value;
+    }
+    const firstType = typeof value[0];
+    const validArrayType = firstType === "string" || firstType === "number" || firstType === "boolean";
+    if (validArrayType && value.every((item) => typeof item === firstType)) {
+      return value;
+    }
+  }
+  throw new TypeError(
+    "trace attribute values must be strings, numbers, booleans, or arrays of a single primitive type"
+  );
+}
+__name(normalizeTraceAttributeValue, "normalizeTraceAttributeValue");
+function createTraceAPI(defaultSpan) {
+  return {
+    get traceId() {
+      return resolveTraceSpan(defaultSpan).spanContext().traceId;
+    },
+    get spanId() {
+      return resolveTraceSpan(defaultSpan).spanContext().spanId;
+    },
+    setAttribute(key, value) {
+      const normalized = normalizeTraceAttributeValue(value);
+      if (normalized === void 0) {
+        return;
+      }
+      resolveTraceSpan(defaultSpan).setAttribute(key, normalized);
+    },
+    setAttributes(attributes) {
+      const normalized = {};
+      for (const [key, value] of Object.entries(attributes)) {
+        const normalizedValue = normalizeTraceAttributeValue(value);
+        if (normalizedValue !== void 0) {
+          normalized[key] = normalizedValue;
+        }
+      }
+      resolveTraceSpan(defaultSpan).setAttributes(normalized);
+    }
+  };
+}
+__name(createTraceAPI, "createTraceAPI");
 function patchFetch() {
   if (!globalThis.fetch.patched) {
     const originalFetch = globalThis.fetch;
@@ -1056,7 +1127,11 @@ var TimePeriod = class _TimePeriod {
     let value = rawValue ? parseInt(rawValue, 10) : 1;
     let complete2 = Boolean(isComplete);
     let offset = 0;
-    switch (direction?.toLowerCase()) {
+    const dir = direction?.toLowerCase();
+    if (!complete2 && !rawValue && (dir === "last" || dir === "next") && ["day", "week", "month", "year"].includes(period)) {
+      complete2 = true;
+    }
+    switch (dir) {
       case "this":
         offset = 0;
         complete2 = true;
@@ -1114,10 +1189,14 @@ var opMapping = {
   },
   notEquals: { op: import_kysely4.sql`is distinct from` },
   equalsRelative: {
-    op: import_kysely4.sql`BETWEEN`,
-    value: /* @__PURE__ */ __name((v) => import_kysely4.sql`${import_kysely4.sql.raw(
-      TimePeriod.fromExpression(v).periodStartSQL()
-    )} AND ${import_kysely4.sql.raw(TimePeriod.fromExpression(v).periodEndSQL())}`, "value")
+    custom: /* @__PURE__ */ __name((qb, fieldName, v) => {
+      const tp = TimePeriod.fromExpression(v);
+      return qb.where(
+        import_kysely4.sql`${import_kysely4.sql.raw(fieldName)} >= ${import_kysely4.sql.raw(
+          tp.periodStartSQL()
+        )} AND ${import_kysely4.sql.raw(fieldName)} < ${import_kysely4.sql.raw(tp.periodEndSQL())}`
+      );
+    }, "custom")
   },
   beforeRelative: {
     op: "<",
@@ -1154,18 +1233,18 @@ var opMapping = {
     notEquals: { op: "=", value: /* @__PURE__ */ __name((v) => import_kysely4.sql`NOT ${v}`, "value") }
   }
 };
-function applyWhereConditions(context6, qb, where = {}) {
-  const conf = context6.tableConfig();
+function applyWhereConditions(context7, qb, where = {}) {
+  const conf = context7.tableConfig();
   for (const key of Object.keys(where)) {
     const v = where[key];
     if (conf && conf[(0, import_change_case.snakeCase)(key)]) {
       const rel = conf[(0, import_change_case.snakeCase)(key)];
-      context6.withJoin(rel.referencesTable, () => {
-        qb = applyWhereConditions(context6, qb, v);
+      context7.withJoin(rel.referencesTable, () => {
+        qb = applyWhereConditions(context7, qb, v);
       });
       continue;
     }
-    const fieldName = `${context6.tableAlias()}.${(0, import_change_case.snakeCase)(key)}`;
+    const fieldName = `${context7.tableAlias()}.${(0, import_change_case.snakeCase)(key)}`;
     if (Object.prototype.toString.call(v) !== "[object Object]") {
       const operator = v === null || Array.isArray(v) ? import_kysely4.sql`is not distinct from` : "=";
       qb = qb.where(fieldName, operator, import_kysely4.sql`${v}`);
@@ -1176,7 +1255,9 @@ function applyWhereConditions(context6, qb, where = {}) {
       if (!mapping) {
         throw new Error(`invalid where condition: ${op}`);
       }
-      if (mapping.isArrayQuery) {
+      if (mapping.custom) {
+        qb = mapping.custom(qb, fieldName, v[op]);
+      } else if (mapping.isArrayQuery) {
         for (const arrayOp of Object.keys(v[op])) {
           qb = qb.where(
             mapping[arrayOp].value ? mapping[arrayOp].value(v[op][arrayOp]) : import_kysely4.sql`${v[op][arrayOp]}`,
@@ -1199,15 +1280,15 @@ function applyWhereConditions(context6, qb, where = {}) {
 __name(applyWhereConditions, "applyWhereConditions");
 
 // src/applyAdditionalQueryConstraints.js
-function applyLimit(context6, qb, limit) {
+function applyLimit(context7, qb, limit) {
   return qb.limit(limit);
 }
 __name(applyLimit, "applyLimit");
-function applyOffset(context6, qb, offset) {
+function applyOffset(context7, qb, offset) {
   return qb.offset(offset);
 }
 __name(applyOffset, "applyOffset");
-function applyOrderBy(context6, qb, tableName, orderBy = {}) {
+function applyOrderBy(context7, qb, tableName, orderBy = {}) {
   Object.entries(orderBy).forEach(([key, sortOrder]) => {
     qb = qb.orderBy(`${tableName}.${(0, import_change_case.snakeCase)(key)}`, sortOrder.toLowerCase());
   });
@@ -1216,41 +1297,41 @@ function applyOrderBy(context6, qb, tableName, orderBy = {}) {
 __name(applyOrderBy, "applyOrderBy");
 
 // src/applyJoins.js
-function applyJoins(context6, qb, where) {
-  const conf = context6.tableConfig();
+function applyJoins(context7, qb, where) {
+  const conf = context7.tableConfig();
   if (!conf) {
     return qb;
   }
-  const srcTable = context6.tableAlias();
+  const srcTable = context7.tableAlias();
   for (const key of Object.keys(where)) {
     const rel = conf[(0, import_change_case.snakeCase)(key)];
     if (!rel) {
       continue;
     }
     const targetTable = rel.referencesTable;
-    if (context6.hasJoin(targetTable)) {
+    if (context7.hasJoin(targetTable)) {
       continue;
     }
-    context6.withJoin(targetTable, () => {
+    context7.withJoin(targetTable, () => {
       switch (rel.relationshipType) {
         case "hasMany":
           qb = qb.innerJoin(
-            `${targetTable} as ${context6.tableAlias()}`,
+            `${targetTable} as ${context7.tableAlias()}`,
             `${srcTable}.id`,
-            `${context6.tableAlias()}.${rel.foreignKey}`
+            `${context7.tableAlias()}.${rel.foreignKey}`
           );
           break;
         case "belongsTo":
           qb = qb.innerJoin(
-            `${targetTable} as ${context6.tableAlias()}`,
+            `${targetTable} as ${context7.tableAlias()}`,
             `${srcTable}.${rel.foreignKey}`,
-            `${context6.tableAlias()}.id`
+            `${context7.tableAlias()}.id`
           );
           break;
         default:
           throw new Error(`unknown relationshipType: ${rel.relationshipType}`);
       }
-      qb = applyJoins(context6, qb, where[key]);
+      qb = applyJoins(context7, qb, where[key]);
     });
   }
   return qb;
@@ -1480,17 +1561,17 @@ var QueryBuilder = class _QueryBuilder {
    * @param {import("./QueryContext").QueryContext} context
    * @param {import("kysely").Kysely} db
    */
-  constructor(tableName, context6, db) {
+  constructor(tableName, context7, db) {
     this._tableName = tableName;
-    this._context = context6;
+    this._context = context7;
     this._db = db;
     this._modelName = upperCamelCase(this._tableName);
   }
   where(where) {
-    const context6 = this._context.clone();
-    let builder = applyJoins(context6, this._db, where);
-    builder = applyWhereConditions(context6, builder, where);
-    return new _QueryBuilder(this._tableName, context6, builder);
+    const context7 = this._context.clone();
+    let builder = applyJoins(context7, this._db, where);
+    builder = applyWhereConditions(context7, builder, where);
+    return new _QueryBuilder(this._tableName, context7, builder);
   }
   sql() {
     return this._db.compile().sql;
@@ -1555,19 +1636,19 @@ var QueryBuilder = class _QueryBuilder {
     const name = spanNameForModelAPI(this._modelName, "findMany");
     const db = useDatabase();
     return withSpan(name, async (span) => {
-      const context6 = new QueryContext([this._tableName], this._tableConfigMap);
+      const context7 = new QueryContext([this._tableName], this._tableConfigMap);
       let builder = db.selectFrom((qb) => {
         return this._db.as(this._tableName);
       }).selectAll();
       if (params?.limit) {
-        builder = applyLimit(context6, builder, params.limit);
+        builder = applyLimit(context7, builder, params.limit);
       }
       if (params?.offset) {
-        builder = applyOffset(context6, builder, params.offset);
+        builder = applyOffset(context7, builder, params.offset);
       }
       if (params?.orderBy !== void 0 && Object.keys(params?.orderBy).length > 0) {
         builder = applyOrderBy(
-          context6,
+          context7,
           builder,
           this._tableName,
           params.orderBy
@@ -1601,7 +1682,7 @@ var QueryBuilder = class _QueryBuilder {
     );
     const db = useDatabase();
     return withSpan(name, async (span) => {
-      const context6 = new QueryContext([this._tableName], this._tableConfigMap);
+      const context7 = new QueryContext([this._tableName], this._tableConfigMap);
       const isOffsetPagination = params.limit != null && params.limit > 0;
       const isBackward = params.last != null && params.last > 0;
       const DEFAULT_PAGE_SIZE = 50;
@@ -1649,7 +1730,7 @@ var QueryBuilder = class _QueryBuilder {
           ])
         ) : normalizedOrderBy;
         builder = applyOrderBy(
-          context6,
+          context7,
           builder,
           this._tableName,
           effectiveOrderBy
@@ -1722,10 +1803,10 @@ var QueryBuilder = class _QueryBuilder {
         }
       }
       if (limit != null) {
-        builder = applyLimit(context6, builder, limit);
+        builder = applyLimit(context7, builder, limit);
       }
       if (isOffsetPagination && params.offset) {
-        builder = applyOffset(context6, builder, params.offset);
+        builder = applyOffset(context7, builder, params.offset);
       }
       span.setAttribute("sql", builder.compile().sql);
       let rows = await builder.execute();
@@ -1830,9 +1911,9 @@ var ModelAPI = class {
     const db = useDatabase();
     return withSpan(name, async (span) => {
       let builder = db.selectFrom(this._tableName).distinctOn(`${this._tableName}.id`).selectAll(this._tableName);
-      const context6 = new QueryContext([this._tableName], this._tableConfigMap);
-      builder = applyJoins(context6, builder, where);
-      builder = applyWhereConditions(context6, builder, where);
+      const context7 = new QueryContext([this._tableName], this._tableConfigMap);
+      builder = applyJoins(context7, builder, where);
+      builder = applyWhereConditions(context7, builder, where);
       span.setAttribute("sql", builder.compile().sql);
       const row = await builder.executeTakeFirst();
       if (!row) {
@@ -1846,23 +1927,23 @@ var ModelAPI = class {
     const db = useDatabase();
     const where = params?.where || {};
     return withSpan(name, async (span) => {
-      const context6 = new QueryContext([this._tableName], this._tableConfigMap);
+      const context7 = new QueryContext([this._tableName], this._tableConfigMap);
       let builder = db.selectFrom((qb) => {
         let builder2 = qb.selectFrom(this._tableName).distinctOn(`${this._tableName}.id`).selectAll(this._tableName);
-        builder2 = applyJoins(context6, builder2, where);
-        builder2 = applyWhereConditions(context6, builder2, where);
+        builder2 = applyJoins(context7, builder2, where);
+        builder2 = applyWhereConditions(context7, builder2, where);
         builder2 = builder2.as(this._tableName);
         return builder2;
       }).selectAll();
       if (params?.limit) {
-        builder = applyLimit(context6, builder, params.limit);
+        builder = applyLimit(context7, builder, params.limit);
       }
       if (params?.offset) {
-        builder = applyOffset(context6, builder, params.offset);
+        builder = applyOffset(context7, builder, params.offset);
       }
       if (params?.orderBy !== void 0 && Object.keys(params?.orderBy).length > 0) {
         builder = applyOrderBy(
-          context6,
+          context7,
           builder,
           this._tableName,
           params.orderBy
@@ -1913,8 +1994,8 @@ var ModelAPI = class {
         }
       }
       builder = builder.set(snakeCaseObject(row));
-      const context6 = new QueryContext([this._tableName], this._tableConfigMap);
-      builder = applyWhereConditions(context6, builder, where);
+      const context7 = new QueryContext([this._tableName], this._tableConfigMap);
+      builder = applyWhereConditions(context7, builder, where);
       span.setAttribute("sql", builder.compile().sql);
       try {
         const row2 = await builder.executeTakeFirstOrThrow();
@@ -1929,8 +2010,8 @@ var ModelAPI = class {
     const db = useDatabase();
     return withSpan(name, async (span) => {
       let builder = db.deleteFrom(this._tableName).returning(["id"]);
-      const context6 = new QueryContext([this._tableName], this._tableConfigMap);
-      builder = applyWhereConditions(context6, builder, where);
+      const context7 = new QueryContext([this._tableName], this._tableConfigMap);
+      builder = applyWhereConditions(context7, builder, where);
       span.setAttribute("sql", builder.compile().sql);
       try {
         const row = await builder.executeTakeFirstOrThrow();
@@ -1943,10 +2024,10 @@ var ModelAPI = class {
   where(where) {
     const db = useDatabase();
     let builder = db.selectFrom(this._tableName).distinctOn(`${this._tableName}.id`).selectAll(this._tableName);
-    const context6 = new QueryContext([this._tableName], this._tableConfigMap);
-    builder = applyJoins(context6, builder, where);
-    builder = applyWhereConditions(context6, builder, where);
-    return new QueryBuilder(this._tableName, context6, builder);
+    const context7 = new QueryContext([this._tableName], this._tableConfigMap);
+    builder = applyJoins(context7, builder, where);
+    builder = applyWhereConditions(context7, builder, where);
+    return new QueryBuilder(this._tableName, context7, builder);
   }
 };
 async function create(conn, tableName, tableConfigs, values) {
@@ -2866,6 +2947,11 @@ async function handleRequest(request, config) {
   }
   return opentelemetry2.context.with(activeContext, () => {
     return withSpan(request.method, async (span) => {
+      span.setAttribute("action.name", request.method);
+      span.setAttribute("action.type", request.type);
+      if (request.meta?.httpPath) {
+        span.setAttribute("http.path", request.meta.httpPath);
+      }
       let db = null;
       try {
         const { createContextAPI, functions, permissionFns, actionTypes } = config;
@@ -2884,7 +2970,8 @@ async function handleRequest(request, config) {
         const headers = new Headers();
         const ctx = createContextAPI({
           responseHeaders: headers,
-          meta: request.meta
+          meta: request.meta,
+          span
         });
         const permitted = request.meta && request.meta.permissionState.status === "granted" ? true : null;
         db = createDatabaseClient({
@@ -3003,7 +3090,8 @@ async function handleJob(request, config) {
           );
         }
         const ctx = createJobContextAPI({
-          meta: request.meta
+          meta: request.meta,
+          span
         });
         const permitted = request.meta && request.meta.permissionState.status === "granted" ? true : null;
         db = createDatabaseClient({
@@ -3091,7 +3179,8 @@ async function handleSubscriber(request, config) {
           );
         }
         const ctx = createSubscriberContextAPI({
-          meta: request.meta
+          meta: request.meta,
+          span
         });
         db = createDatabaseClient({
           connString: request.meta?.secrets?.KEEL_DB_CONN
@@ -3167,7 +3256,8 @@ async function handleRoute(request, config) {
           ...ctx
         } = createContextAPI({
           responseHeaders: new Headers(),
-          meta: request.meta
+          meta: request.meta,
+          span
         });
         request.params.headers = headers;
         db = createDatabaseClient({
@@ -4010,6 +4100,7 @@ function createFlowContext(runId, data, action, callback, element, spanId, ctx)
     env: ctx.env,
     now: ctx.now,
     secrets: ctx.secrets,
+    trace: ctx.trace,
     complete: /* @__PURE__ */ __name((options) => {
       return {
         __type: "ui.complete",
@@ -4018,116 +4109,19 @@ function createFlowContext(runId, data, action, callback, element, spanId, ctx)
     }, "complete"),
     step: /* @__PURE__ */ __name(async (name, optionsOrFn, fn) => {
       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();
-        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) {
-          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 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();
-        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) => {
+        return withUserSpan(span, async () => {
+          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()
@@ -4135,85 +4129,188 @@ 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 */) {
-            span.setAttribute(KEEL_INTERNAL_ATTR, KEEL_INTERNAL_CHILDREN);
-            const parsedData2 = transformRichDataTypes(step.value);
-            if (step.action) {
-              return { data: parsedData2, action: step.action };
-            }
-            return parsedData2;
+          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();
-            span.setAttribute("rendered", true);
-            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) {
-            span.setAttribute("callback", callback);
+          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) => {
+          return withUserSpan(span, async () => {
+            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}`);
+            }
+            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);
+              const parsedData2 = transformRichDataTypes(step.value);
+              if (step.action) {
+                return { data: parsedData2, action: step.action };
+              }
+              return parsedData2;
+            }
+            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
+                );
+              }
             }
-          } catch (e) {
-            if (e instanceof UIRenderDisrupt) {
+            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();
-          const parsedData = transformRichDataTypes(data);
-          if (action) {
-            return { data: parsedData, action };
-          }
-          return parsedData;
+              endTime: /* @__PURE__ */ new Date()
+            }).where("id", "=", step.id).returningAll().executeTakeFirst();
+            const parsedData = transformRichDataTypes(data);
+            if (action) {
+              return { data: parsedData, action };
+            }
+            return parsedData;
+          });
         });
       }), "page"),
       inputs: {
@@ -4327,7 +4424,8 @@ async function handleFlow(request, config) {
           request.meta.element,
           span.spanContext().spanId,
           createFlowContextAPI({
-            meta: request.meta
+            meta: request.meta,
+            span
           })
         );
         const flowFunction = flows[request.method].fn;
@@ -4455,6 +4553,7 @@ __name(handleFlow, "handleFlow");
 
 // src/index.ts
 var import_ksuid2 = __toESM(require("ksuid"), 1);
+var createTraceAPI2 = createTraceAPI;
 function ksuid() {
   return import_ksuid2.default.randomSync().string;
 }
@@ -4480,6 +4579,7 @@ __name(ksuid, "ksuid");
   TaskAPI,
   checkBuiltInPermissions,
   createFlowContext,
+  createTraceAPI,
   handleFlow,
   handleJob,
   handleRequest,