From 7aa7f640639d986992d3911dd03b9d872d76aa86 Mon Sep 17 00:00:00 2001
From: Armaan Aggarwal <49625265+Armster15@users.noreply.github.com>
Date: Wed, 21 Oct 2026 23:35:37 +0000
Subject: [PATCH] Keep handler errors from claiming the reserved
 SERVER_DRAINING code

Only the runtime admission refusal may carry SERVER_DRAINING, so a
procedure that already ran is never resent. Remap a handler-returned
SERVER_DRAINING to UNCAUGHT_ERROR and log a warning.

Assisted-by: Replit
---
 __tests__/drain.test.ts | 30 ++++++++++++++++++++++++++++++
 router/server.ts        | 11 +++++++++++
 2 files changed, 41 insertions(+)

diff --git a/__tests__/drain.test.ts b/__tests__/drain.test.ts
index c04ca6be..31075049 100644
--- a/__tests__/drain.test.ts
+++ b/__tests__/drain.test.ts
@@ -5,9 +5,11 @@ import {
   createClient,
   createServer,
   createServiceSchema,
+  Err,
   Middleware,
   Ok,
   Procedure,
+  UNCAUGHT_ERROR_CODE,
 } from '../router';
 import { createMockTransportNetwork, readNextResult } from '../testUtil';
 import { waitFor } from '../testUtil/fixtures/cleanup';
@@ -62,6 +64,19 @@ function setup(
           resWritable.write(Ok(null));
         },
       }),
+      claim: Procedure.rpc({
+        requestInit: init,
+        responseData,
+        responseError: Type.Object({
+          code: Type.Literal('SERVER_DRAINING'),
+          message: Type.String(),
+        }),
+        async handler() {
+          effects.push('claim');
+
+          return Err({ code: 'SERVER_DRAINING', message: 'upstream' });
+        },
+      }),
     }),
   };
   function startServer() {
@@ -279,3 +294,18 @@ test('a refused upload is returned without replay', async () => {
   expect(await f.client.test.upload.upload({}).finalize()).toEqual(refusal);
   expect(f.effects).toEqual(['hold']);
 });
+
+test('a handler that returns the reserved drain code is not resent', async () => {
+  // Risk: a handler that already ran is replayed on the next session.
+  const f = setup(1);
+  const result = f.client.test.claim.rpc({});
+  await waitFor(() => expect(f.effects).toEqual(['claim']));
+  await vi.advanceTimersByTimeAsync(10);
+  f.reconnect();
+  await vi.advanceTimersByTimeAsync(20);
+  expect(f.effects).toEqual(['claim']);
+  expect(await result).toMatchObject({
+    ok: false,
+    payload: { code: UNCAUGHT_ERROR_CODE, message: 'upstream' },
+  });
+});
diff --git a/router/server.ts b/router/server.ts
index 827b5690..2ef33b96 100644
--- a/router/server.ts
+++ b/router/server.ts
@@ -604,6 +604,17 @@ class RiverServer<
       Result<Static<PayloadType>, Static<BaseErrorSchemaType>>
     >({
       writeCb: (response) => {
+        if (!response.ok && response.payload.code === SERVER_DRAINING_CODE) {
+          this.log?.warn(
+            `${serviceName}.${procedureName} handler returned the reserved ${SERVER_DRAINING_CODE} code`,
+            loggingMetadata,
+          );
+          response = Err({
+            code: UNCAUGHT_ERROR_CODE,
+            message: response.payload.message,
+          });
+        }
+
         if (!response.ok) {
           recordRiverError(span, response.payload);
         }