Keep handler errors from claiming the reserved SERVER_DRAINING code

__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' },
+  });
+});

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);
         }