From 54a4e244933b315aa4d6402fc857e095ce3d7020 Mon Sep 17 00:00:00 2001
From: swyxio <shawnthe1@gmail.com>
Date: Thu, 7 Aug 2025 01:35:54 -0700
Subject: [PATCH] elderberry

---
 .../run-engine/src/engine/systems/runAttemptSystem.ts         | 4 ++++
 packages/core/src/v3/schemas/common.ts                        | 4 ++++
 2 files changed, 8 insertions(+)

diff --git a/internal-packages/run-engine/src/engine/systems/runAttemptSystem.ts b/internal-packages/run-engine/src/engine/systems/runAttemptSystem.ts
index ce0f8abe4d5..76897dc0bff 100644
--- a/internal-packages/run-engine/src/engine/systems/runAttemptSystem.ts
+++ b/internal-packages/run-engine/src/engine/systems/runAttemptSystem.ts
@@ -181,6 +181,7 @@ export class RunAttemptSystem {
         lockedById: true,
         lockedQueueId: true,
         queue: true,
+        concurrencyKey: true,
         attemptNumber: true,
         status: true,
         ttl: true,
@@ -248,6 +249,7 @@ export class RunAttemptSystem {
     ]);
 
     return {
+      concurrencyKey: run.concurrencyKey ?? undefined,
       run: {
         id: run.friendlyId,
         tags: run.runTags,
@@ -396,6 +398,7 @@ export class RunAttemptSystem {
                   lockedById: true,
                   lockedQueueId: true,
                   queue: true,
+                  concurrencyKey: true,
                   attemptNumber: true,
                   status: true,
                   ttl: true,
@@ -533,6 +536,7 @@ export class RunAttemptSystem {
             ]);
 
           const execution: BackwardsCompatibleTaskRunExecution = {
+            concurrencyKey: updatedRun.concurrencyKey ?? undefined,
             attempt: {
               number: nextAttemptNumber,
               startedAt: latestSnapshot.updatedAt,
diff --git a/packages/core/src/v3/schemas/common.ts b/packages/core/src/v3/schemas/common.ts
index 2928995606b..487849df195 100644
--- a/packages/core/src/v3/schemas/common.ts
+++ b/packages/core/src/v3/schemas/common.ts
@@ -334,6 +334,8 @@ export const TaskRunExecution = z.object({
       traceContext: z.record(z.unknown()).optional(),
     })
   ),
+  // The concurrencyKey creates a copy of the queue per unique value
+  concurrencyKey: z.string().optional(),
   ...StaticTaskRunExecutionShape,
 });
 
@@ -406,6 +408,8 @@ export const TaskRunContext = z.object({
     durationMs: true,
     costInCents: true,
   }),
+  // The concurrencyKey used for this run, if any
+  concurrencyKey: z.string().optional(),
   ...StaticTaskRunExecutionShape,
 });