fix: read-your-writes + global-scope idempotency correctness under the run-ops split

apps/webapp/app/models/runtimeEnvironment.server.ts

@@ -274,19 +274,17 @@ export async function findEnvironmentFromRun(
 ): Promise<EnvironmentFromRun | null> {
   // Run-ops scalars (runTags/batchId/runtimeEnvironmentId) from the run store; the env half is
   // resolved via the control-plane resolver so the run-ops DB can split without a cross-DB join.
-  const taskRun = await runStore.findRun(
-    {
-      id: runId,
-    },
-    {
-      select: {
-        runTags: true,
-        batchId: true,
-        runtimeEnvironmentId: true,
-      },
-    },
-    tx ?? $replica
-  );
+  const select = {
+    runTags: true,
+    batchId: true,
+    runtimeEnvironmentId: true,
+  } as const;
+  let taskRun = await runStore.findRun({ id: runId }, { select }, tx ?? $replica);
+  if (!taskRun) {
+    // Read-your-writes: a just-created run may not have replicated. Re-read the owning primary before
+    // treating it as absent, so runMetadataUpdated doesn't drop a live run's final metadata + publish.
+    taskRun = await runStore.findRun({ id: runId }, { select }, prisma);
+  }
   if (!taskRun) {
     return null;
   }

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

@@ -42,29 +42,36 @@ export class ApiWaitpointPresenter extends BasePresenter {
     return this.trace("call", async (span) => {
       // The store routes by the waitpointId's residency (id shape) and reads the owning
       // store's replica. waitpointId is pre-decoded from the friendlyId via WaitpointId.toId.
-      const waitpoint = await this.runStore.findWaitpoint({
-        where: {
-          id: waitpointId,
-          environmentId: environment.id,
-        },
-        select: {
-          id: true,
-          friendlyId: true,
-          type: true,
-          status: true,
-          idempotencyKey: true,
-          userProvidedIdempotencyKey: true,
-          idempotencyKeyExpiresAt: true,
-          inactiveIdempotencyKey: true,
-          output: true,
-          outputType: true,
-          outputIsError: true,
-          completedAfter: true,
-          completedAt: true,
-          createdAt: true,
-          tags: true,
-        },
-      });
+      const where = {
+        id: waitpointId,
+        environmentId: environment.id,
+      };
+      const select = {
+        id: true,
+        friendlyId: true,
+        type: true,
+        status: true,
+        idempotencyKey: true,
+        userProvidedIdempotencyKey: true,
+        idempotencyKeyExpiresAt: true,
+        inactiveIdempotencyKey: true,
+        output: true,
+        outputType: true,
+        outputIsError: true,
+        completedAfter: true,
+        completedAt: true,
+        createdAt: true,
+        tags: true,
+      } as const;
+
+      let waitpoint = await this.runStore.findWaitpoint({ where, select });
+
+      // Read-your-writes on a public GET: a just-minted token may not be on the owning store's
+      // replica yet, so a replica miss would 404 a live token. Re-read the owning primary before
+      // concluding it doesn't exist (mirrors the metadata GET loader + the complete/callback paths).
+      if (!waitpoint) {
+        waitpoint = await this.runStore.findWaitpointOnPrimary({ where, select });
+      }
 
       if (!waitpoint) {
         logger.error(`WaitpointPresenter: Waitpoint not found`, {

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

@@ -47,13 +47,25 @@ export class BatchPresenter extends BasePresenter {
     // The BatchTaskRun (run-ops) is read through the run store, which routes by residency. The
     // runtimeEnvironment (control-plane) is resolved separately because the cross-seam FK is
     // dropped, so the batch row cannot single-SQL join to control-plane RuntimeEnvironment.
-    const batch = await this.runStore.findBatchTaskRunByFriendlyId(
+    let batch = await this.runStore.findBatchTaskRunByFriendlyId(
       batchId,
       environmentId,
       { include: BATCH_INCLUDE },
       this._replica
     );
 
+    // Read-your-writes: findBatchTaskRunByFriendlyId defaults to (and here reads) the replica, so a
+    // batch created within the replica's apply window returns null under lag. Re-read from the owning
+    // primary on a miss so a live batch's detail page never spuriously 404s ("Batch not found").
+    if (!batch) {
+      batch = await this.runStore.findBatchTaskRunByFriendlyId(
+        batchId,
+        environmentId,
+        { include: BATCH_INCLUDE },
+        this._prisma
+      );
+    }
+
     if (!batch) {
       throw new Error("Batch not found");
     }

apps/webapp/app/routes/api.v1.batches.$batchId.ts

@@ -12,6 +12,11 @@ export const loader = createLoaderApiRoute(
     params: ParamsSchema,
     allowJWT: true,
     corsStrategy: "all",
+    // A just-created batch may not yet have replicated to the read replica this client-less
+    // findBatchTaskRunByFriendlyId lookup routes to; return a retryable 404 so the SDK retries through
+    // replica lag rather than stranding a live batch on a permanent 404 (mirrors the run-get routes,
+    // e.g. api.v3.runs.$runId).
+    shouldRetryNotFound: true,
     findResource: (params, auth) => {
       return runStore.findBatchTaskRunByFriendlyId(params.batchId, auth.environment.id, {
         include: { errors: true },

apps/webapp/app/routes/api.v1.runs.$runFriendlyId.input-streams.wait.ts

@@ -33,21 +33,23 @@ const { action, loader } = createActionApiRoute(
   },
   async ({ authentication, body, params }) => {
     try {
-      const run = await runStore.findRun(
-        {
-          friendlyId: params.runFriendlyId,
-          runtimeEnvironmentId: authentication.environment.id,
+      const where = {
+        friendlyId: params.runFriendlyId,
+        runtimeEnvironmentId: authentication.environment.id,
+      };
+      const args = {
+        select: {
+          id: true,
+          friendlyId: true,
+          realtimeStreamsVersion: true,
+          streamBasinName: true,
         },
-        {
-          select: {
-            id: true,
-            friendlyId: true,
-            realtimeStreamsVersion: true,
-            streamBasinName: true,
-          },
-        },
-        $replica
-      );
+      };
+      // Replica lag can null out a live run; a spurious 404 fails the .wait() registration on a run
+      // that exists. Re-read the owning primary on a replica miss.
+      const run =
+        (await runStore.findRun(where, args, $replica)) ??
+        (await runStore.findRunOnPrimary(where, args));
 
       if (!run) {
         return json({ error: "Run not found" }, { status: 404 });

apps/webapp/app/routes/api.v1.runs.$runFriendlyId.session-streams.wait.ts

@@ -39,20 +39,22 @@ const { action, loader } = createActionApiRoute(
   },
   async ({ authentication, body, params }) => {
     try {
-      const run = await runStore.findRun(
-        {
-          friendlyId: params.runFriendlyId,
-          runtimeEnvironmentId: authentication.environment.id,
+      const where = {
+        friendlyId: params.runFriendlyId,
+        runtimeEnvironmentId: authentication.environment.id,
+      };
+      const args = {
+        select: {
+          id: true,
+          friendlyId: true,
+          realtimeStreamsVersion: true,
         },
-        {
-          select: {
-            id: true,
-            friendlyId: true,
-            realtimeStreamsVersion: true,
-          },
-        },
-        $replica
-      );
+      };
+      // Replica lag can null out a live run; a spurious 404 fails the .wait() registration on a run
+      // that exists. Re-read the owning primary on a replica miss.
+      const run =
+        (await runStore.findRun(where, args, $replica)) ??
+        (await runStore.findRunOnPrimary(where, args));
 
       if (!run) {
         return json({ error: "Run not found" }, { status: 404 });

apps/webapp/app/routes/api.v1.runs.$runId.metadata.ts

@@ -64,6 +64,19 @@ export async function loader({ request, params }: LoaderFunctionArgs) {
     );
   }
 
+  // Read-your-writes: a run drained from the buffer to the primary but not yet replicated misses
+  // both the replica read and the buffer. Re-read the owning primary before 404ing.
+  const primaryRun = await runStore.findRunOnPrimary(
+    { friendlyId: parsed.data.runId, runtimeEnvironmentId: env.id },
+    { select: { metadata: true, metadataType: true } }
+  );
+  if (primaryRun) {
+    return json(
+      { metadata: primaryRun.metadata, metadataType: primaryRun.metadataType },
+      { status: 200 }
+    );
+  }
+
   return json({ error: "Run not found" }, { status: 404 });
 }
 

apps/webapp/app/routes/api.v1.sessions.$session.end-and-continue.ts

@@ -75,14 +75,27 @@ const { action, loader } = createActionApiRoute(
     // SDK exposes via `ctx.run.id`). Internally `Session.currentRunId`
     // stores the TaskRun.id cuid, so resolve before handing to the
     // optimistic-claim service.
-    const callingRun = await runStore.findRun(
+    let callingRun = await runStore.findRun(
       {
         friendlyId: body.callingRunId,
         runtimeEnvironmentId: authentication.environment.id,
       },
       { select: { id: true } },
       $replica
     );
+    if (!callingRun) {
+      // Replica lag: `callingRunId` is the agent's own live run (it is executing this request), so it
+      // exists on the owning primary even when the read replica has not caught up. Re-read the primary
+      // before 404ing — otherwise a lagging replica turns a legitimate handoff into a spurious
+      // "callingRunId not found in this environment".
+      callingRun = await runStore.findRunOnPrimary(
+        {
+          friendlyId: body.callingRunId,
+          runtimeEnvironmentId: authentication.environment.id,
+        },
+        { select: { id: true } }
+      );
+    }
     if (!callingRun) {
       return json({ error: "callingRunId not found in this environment" }, { status: 404 });
     }

apps/webapp/app/routes/api.v1.waitpoints.tokens.$waitpointFriendlyId.callback.$hash.ts

@@ -36,13 +36,22 @@ export async function action({ request, params }: ActionFunctionArgs) {
     // Resolve wherever the waitpoint resides. The store routes by the waitpoint id's residency
     // (id-shape) and probes both run-ops DBs, so a token on either store resolves; the env is
     // resolved below from the row via the control-plane resolver.
-    const waitpoint = await runStore.findWaitpoint({
+    let waitpoint = await runStore.findWaitpoint({
       where: {
         id: waitpointId,
       },
       select: { id: true, status: true, environmentId: true },
     });
 
+    if (!waitpoint) {
+      // Read-your-writes: a token whose callback fires right after mint may not have replicated
+      // yet. Re-read the owning primary before 404ing (mirrors complete.ts's primary fallback).
+      waitpoint = await runStore.findWaitpointOnPrimary({
+        where: { id: waitpointId },
+        select: { id: true, status: true, environmentId: true },
+      });
+    }
+
     if (!waitpoint) {
       return json({ error: "Waitpoint not found" }, { status: 404 });
     }

apps/webapp/app/routes/api.v2.batches.$batchId.ts

@@ -12,6 +12,11 @@ export const loader = createLoaderApiRoute(
     params: ParamsSchema,
     allowJWT: true,
     corsStrategy: "all",
+    // A just-created batch may not yet have replicated to the read replica this client-less
+    // findBatchTaskRunByFriendlyId lookup routes to; return a retryable 404 so the SDK retries through
+    // replica lag rather than stranding a live batch on a permanent 404 (mirrors the run-get routes,
+    // e.g. api.v3.runs.$runId).
+    shouldRetryNotFound: true,
     findResource: (params, auth) => {
       return runStore.findBatchTaskRunByFriendlyId(params.batchId, auth.environment.id, {
         include: { errors: true },

apps/webapp/app/routes/realtime.v1.batches.$batchId.ts

@@ -13,6 +13,11 @@ export const loader = createLoaderApiRoute(
     params: ParamsSchema,
     allowJWT: true,
     corsStrategy: "all",
+    // A just-created batch may not yet have replicated to the read replica this client-less
+    // findBatchTaskRunByFriendlyId lookup routes to; return a retryable 404 so the SDK retries through
+    // replica lag rather than stranding a live batch on a permanent 404 (mirrors the run-get routes,
+    // e.g. api.v3.runs.$runId).
+    shouldRetryNotFound: true,
     findResource: (params, auth) => {
       return runStore.findBatchTaskRunByFriendlyId(params.batchId, auth.environment.id);
     },

apps/webapp/app/routes/realtime.v1.runs.$runId.ts

@@ -15,22 +15,24 @@ export const loader = createLoaderApiRoute(
     allowJWT: true,
     corsStrategy: "all",
     findResource: async (params, authentication) => {
-      return runStore.findRun(
-        {
-          friendlyId: params.runId,
-          runtimeEnvironmentId: authentication.environment.id,
-        },
-        {
-          include: {
-            batch: {
-              select: {
-                friendlyId: true,
-              },
+      const where = {
+        friendlyId: params.runId,
+        runtimeEnvironmentId: authentication.environment.id,
+      };
+      const args = {
+        include: {
+          batch: {
+            select: {
+              friendlyId: true,
             },
           },
         },
-        $replica
-      );
+      };
+      // Replica lag can null out a run that already exists on the owning primary. A spurious 404
+      // here permanently fails the client's realtime subscription (the SSE client treats 404 as
+      // "stream gone" — nonRetryableStatuses). Re-read the primary on a replica miss.
+      const run = await runStore.findRun(where, args, $replica);
+      return run ?? runStore.findRunOnPrimary(where, args);
     },
     authorization: {
       action: "read",

apps/webapp/app/routes/realtime.v1.streams.$runId.$streamId.ts

@@ -16,29 +16,29 @@ export const loader = createLoaderApiRoute(
     allowJWT: true,
     corsStrategy: "all",
     findResource: async (params, auth) => {
-      const run = await runStore.findRun(
-        {
-          friendlyId: params.runId,
-          runtimeEnvironmentId: auth.environment.id,
-        },
-        {
-          select: {
-            id: true,
-            friendlyId: true,
-            taskIdentifier: true,
-            runTags: true,
-            realtimeStreamsVersion: true,
-            streamBasinName: true,
-            batch: {
-              select: {
-                friendlyId: true,
-              },
+      const where = {
+        friendlyId: params.runId,
+        runtimeEnvironmentId: auth.environment.id,
+      };
+      const args = {
+        select: {
+          id: true,
+          friendlyId: true,
+          taskIdentifier: true,
+          runTags: true,
+          realtimeStreamsVersion: true,
+          streamBasinName: true,
+          batch: {
+            select: {
+              friendlyId: true,
             },
           },
         },
-        $replica
-      );
-      return run;
+      };
+      // Replica lag can null out a live run; a spurious 404 permanently fails the SSE subscription
+      // (the client treats 404 as "stream gone"). Re-read the owning primary on a replica miss.
+      const run = await runStore.findRun(where, args, $replica);
+      return run ?? runStore.findRunOnPrimary(where, args);
     },
     authorization: {
       action: "read",