feat(sdk): expose user-provided idempotency key and scope in task context

Author: ericallam

apps/webapp/app/presenters/v3/ApiRetrieveRunPresenter.server.ts

@@ -18,6 +18,26 @@ import { generatePresignedUrl } from "~/v3/r2.server";
 import { tracer } from "~/v3/tracer.server";
 import { startSpanWithEnv } from "~/v3/tracing.server";
 
+/**
+ * Returns the user-provided idempotency key if available (from idempotencyKeyOptions),
+ * otherwise falls back to the stored idempotency key (which is the hash).
+ */
+function getUserProvidedIdempotencyKey(run: {
+  idempotencyKey: string | null;
+  idempotencyKeyOptions: unknown;
+}): string | undefined {
+  // If we have the user-provided key options, return the original key
+  const options = run.idempotencyKeyOptions as {
+    key?: string;
+    scope?: string;
+  } | null;
+  if (options?.key) {
+    return options.key;
+  }
+  // Fallback to the hash (for runs created before this feature)
+  return run.idempotencyKey ?? undefined;
+}
+
 // Build 'select' object
 const commonRunSelect = {
   id: true,
@@ -38,6 +58,7 @@ const commonRunSelect = {
   baseCostInCents: true,
   usageDurationMs: true,
   idempotencyKey: true,
+  idempotencyKeyOptions: true,
   isTest: true,
   depth: true,
   scheduleId: true,
@@ -442,7 +463,7 @@ async function createCommonRunStructure(run: CommonRelatedRun, apiVersion: API_V
   return {
     id: run.friendlyId,
     taskIdentifier: run.taskIdentifier,
-    idempotencyKey: run.idempotencyKey ?? undefined,
+    idempotencyKey: getUserProvidedIdempotencyKey(run),
     version: run.lockedToVersion?.version,
     status: ApiRetrieveRunPresenter.apiStatusFromRunStatus(run.status, apiVersion),
     createdAt: run.createdAt,

apps/webapp/app/presenters/v3/SpanPresenter.server.ts

@@ -229,7 +229,7 @@ export class SpanPresenter extends BasePresenter {
       isTest: run.isTest,
       replayedFromTaskRunFriendlyId: run.replayedFromTaskRunFriendlyId,
       environmentId: run.runtimeEnvironment.id,
-      idempotencyKey: run.idempotencyKey,
+      idempotencyKey: this.getUserProvidedIdempotencyKey(run),
       idempotencyKeyExpiresAt: run.idempotencyKeyExpiresAt,
       debounce: run.debounce as { key: string; delay: string; createdAt: Date } | null,
       schedule: await this.resolveSchedule(run.scheduleId ?? undefined),
@@ -355,6 +355,7 @@ export class SpanPresenter extends BasePresenter {
         //idempotency
         idempotencyKey: true,
         idempotencyKeyExpiresAt: true,
+        idempotencyKeyOptions: true,
         //debounce
         debounce: true,
         //delayed
@@ -644,7 +645,7 @@ export class SpanPresenter extends BasePresenter {
         createdAt: run.createdAt,
         tags: run.runTags,
         isTest: run.isTest,
-        idempotencyKey: run.idempotencyKey ?? undefined,
+        idempotencyKey: this.getUserProvidedIdempotencyKey(run) ?? undefined,
         startedAt: run.startedAt ?? run.createdAt,
         durationMs: run.usageDurationMs,
         costInCents: run.costInCents,
@@ -704,4 +705,23 @@ export class SpanPresenter extends BasePresenter {
 
     return parsedTraceparent?.traceId;
   }
+
+  /**
+   * Returns the user-provided idempotency key if available (from idempotencyKeyOptions),
+   * otherwise falls back to the stored idempotency key (which is the hash).
+   */
+  getUserProvidedIdempotencyKey(
+    run: Pick<FindRunResult, "idempotencyKey" | "idempotencyKeyOptions">
+  ): string | null {
+    // If we have the user-provided key options, return the original key
+    const options = run.idempotencyKeyOptions as {
+      key?: string;
+      scope?: string;
+    } | null;
+    if (options?.key) {
+      return options.key;
+    }
+    // Fallback to the hash (for runs created before this feature)
+    return run.idempotencyKey;
+  }
 }

apps/webapp/app/runEngine/services/triggerTask.server.ts

@@ -304,6 +304,7 @@ export class RunEngineTriggerTaskService {
                 environment: environment,
                 idempotencyKey,
                 idempotencyKeyExpiresAt: idempotencyKey ? idempotencyKeyExpiresAt : undefined,
+                idempotencyKeyOptions: body.options?.idempotencyKeyOptions,
                 taskIdentifier: taskId,
                 payload: payloadPacket.data ?? "",
                 payloadType: payloadPacket.dataType,

apps/webapp/app/services/runsReplicationService.server.ts

@@ -891,6 +891,8 @@ export class RunsReplicationService {
       run.spanId, // span_id
       run.traceId, // trace_id
       run.idempotencyKey ?? "", // idempotency_key
+      this.#extractIdempotencyKeyUser(run), // idempotency_key_user
+      this.#extractIdempotencyKeyScope(run), // idempotency_key_scope
       run.ttl ?? "", // expiration_ttl
       run.isTest ?? false, // is_test
       _version.toString(), // _version
@@ -951,6 +953,16 @@ export class RunsReplicationService {
 
     return { data: parsedData };
   }
+
+  #extractIdempotencyKeyUser(run: TaskRun): string {
+    const options = run.idempotencyKeyOptions as { key?: string; scope?: string } | null;
+    return options?.key ?? "";
+  }
+
+  #extractIdempotencyKeyScope(run: TaskRun): string {
+    const options = run.idempotencyKeyOptions as { key?: string; scope?: string } | null;
+    return options?.scope ?? "";
+  }
 }
 
 export type ConcurrentFlushSchedulerConfig<T> = {

internal-packages/clickhouse/schema/013_add_task_runs_v2_idempotency_key_options.sql

@@ -0,0 +1,16 @@
+-- +goose Up
+
+-- Add columns for storing user-provided idempotency key and scope for searching
+ALTER TABLE trigger_dev.task_runs_v2
+ADD COLUMN idempotency_key_user String DEFAULT '';
+
+ALTER TABLE trigger_dev.task_runs_v2
+ADD COLUMN idempotency_key_scope String DEFAULT '';
+
+-- +goose Down
+
+ALTER TABLE trigger_dev.task_runs_v2
+DROP COLUMN idempotency_key_user;
+
+ALTER TABLE trigger_dev.task_runs_v2
+DROP COLUMN idempotency_key_scope;

internal-packages/clickhouse/src/taskRuns.ts

@@ -40,6 +40,8 @@ export const TaskRunV2 = z.object({
   span_id: z.string(),
   trace_id: z.string(),
   idempotency_key: z.string(),
+  idempotency_key_user: z.string().default(""),
+  idempotency_key_scope: z.string().default(""),
   expiration_ttl: z.string(),
   is_test: z.boolean().default(false),
   concurrency_key: z.string().default(""),
@@ -91,6 +93,8 @@ export const TASK_RUN_COLUMNS = [
   "span_id",
   "trace_id",
   "idempotency_key",
+  "idempotency_key_user",
+  "idempotency_key_scope",
   "expiration_ttl",
   "is_test",
   "_version",
@@ -151,6 +155,8 @@ export type TaskRunFieldTypes = {
   span_id: string;
   trace_id: string;
   idempotency_key: string;
+  idempotency_key_user: string;
+  idempotency_key_scope: string;
   expiration_ttl: string;
   is_test: boolean;
   _version: string;
@@ -282,6 +288,8 @@ export type TaskRunInsertArray = [
   span_id: string,
   trace_id: string,
   idempotency_key: string,
+  idempotency_key_user: string,
+  idempotency_key_scope: string,
   expiration_ttl: string,
   is_test: boolean,
   _version: string,

internal-packages/database/prisma/migrations/20260116154810_add_idempotency_key_options_to_task_run/migration.sql

@@ -0,0 +1,2 @@
+-- AlterTable
+ALTER TABLE "TaskRun" ADD COLUMN "idempotencyKeyOptions" JSONB;

internal-packages/database/prisma/schema.prisma

@@ -589,6 +589,8 @@ model TaskRun {
 
   idempotencyKey          String?
   idempotencyKeyExpiresAt DateTime?
+  /// Stores the user-provided key and scope: { key: string, scope: "run" | "attempt" | "global" }
+  idempotencyKeyOptions   Json?
 
   /// Debounce options: { key: string, delay: string, createdAt: Date }
   debounce Json?

internal-packages/run-engine/src/engine/index.ts

@@ -395,6 +395,7 @@ export class RunEngine {
       environment,
       idempotencyKey,
       idempotencyKeyExpiresAt,
+      idempotencyKeyOptions,
       taskIdentifier,
       payload,
       payloadType,
@@ -544,6 +545,7 @@ export class RunEngine {
               projectId: environment.project.id,
               idempotencyKey,
               idempotencyKeyExpiresAt,
+              idempotencyKeyOptions,
               taskIdentifier,
               payload,
               payloadType,

internal-packages/run-engine/src/engine/systems/runAttemptSystem.ts

@@ -194,6 +194,7 @@ export class RunAttemptSystem {
         runTags: true,
         isTest: true,
         idempotencyKey: true,
+        idempotencyKeyOptions: true,
         startedAt: true,
         maxAttempts: true,
         taskVersion: true,
@@ -261,7 +262,8 @@ export class RunAttemptSystem {
         isTest: run.isTest,
         createdAt: run.createdAt,
         startedAt: run.startedAt ?? run.createdAt,
-        idempotencyKey: run.idempotencyKey ?? undefined,
+        idempotencyKey: this.#getUserProvidedIdempotencyKey(run) ?? undefined,
+        idempotencyKeyScope: this.#getIdempotencyKeyScope(run),
         maxAttempts: run.maxAttempts ?? undefined,
         version: run.taskVersion ?? "unknown",
         maxDuration: run.maxDurationInSeconds ?? undefined,
@@ -422,6 +424,7 @@ export class RunAttemptSystem {
                   runTags: true,
                   isTest: true,
                   idempotencyKey: true,
+                  idempotencyKeyOptions: true,
                   startedAt: true,
                   maxAttempts: true,
                   taskVersion: true,
@@ -570,7 +573,8 @@ export class RunAttemptSystem {
               createdAt: updatedRun.createdAt,
               tags: updatedRun.runTags,
               isTest: updatedRun.isTest,
-              idempotencyKey: updatedRun.idempotencyKey ?? undefined,
+              idempotencyKey: this.#getUserProvidedIdempotencyKey(updatedRun) ?? undefined,
+              idempotencyKeyScope: this.#getIdempotencyKeyScope(updatedRun),
               startedAt: updatedRun.startedAt ?? updatedRun.createdAt,
               maxAttempts: updatedRun.maxAttempts ?? undefined,
               version: updatedRun.taskVersion ?? "unknown",
@@ -1914,6 +1918,25 @@ export class RunAttemptSystem {
       stackTrace: truncateString(error.stackTrace, 1024 * 16), // 16kb
     };
   }
+
+  #getUserProvidedIdempotencyKey(
+    run: { idempotencyKey: string | null; idempotencyKeyOptions: unknown }
+  ): string | null {
+    const options = run.idempotencyKeyOptions as { key?: string; scope?: string } | null;
+    // Return user-provided key if available, otherwise fall back to the hash
+    return options?.key ?? run.idempotencyKey;
+  }
+
+  #getIdempotencyKeyScope(
+    run: { idempotencyKeyOptions: unknown }
+  ): "run" | "attempt" | "global" | undefined {
+    const options = run.idempotencyKeyOptions as { key?: string; scope?: string } | null;
+    const scope = options?.scope;
+    if (scope === "run" || scope === "attempt" || scope === "global") {
+      return scope;
+    }
+    return undefined;
+  }
 }
 
 export function safeParseGitMeta(git: unknown): GitMeta | undefined {