From 722abeca9a8812aa19d97fb90fbb0049bef7783f Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 23 Oct 2026 09:34:38 +0000 Subject: [PATCH] fix(effect): make middlewareGen next() structured and propagate Effect context `next()` in `middlewareGen` wrapped `opts.next` in `Effect.tryPromise` without a signal, so Effect abandoned the downstream promise on interruption: the middleware's finalizers ran while downstream middleware and the handler were still running, `call()` settled early, and `Effect.timeout`/`Effect.race` around `next()` never stopped the handler. Services, spans and log annotations applied around `next()` also never reached downstream effects. - server: middleware `next({ signal })` replaces the signal for the rest of the chain (downstream middleware and the handler), including through `decorateMiddleware(...).use(...)`. Documented on the middleware page. - effect: `next()` runs via `Effect.callback`. Interrupting it aborts the downstream signal (linked with the middleware's own signal) and waits for the chain to settle before finalizers run, like `try/finally` around `await next()`. - effect: `next()` passes the fiber's Effect context downstream as `effect/context`. An `effect/context` passed to `next()` is the incoming one plus changes, so only the changed services are layered over the fiber context, keeping the fiber's spans and log annotations. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01C5AXgMUccu8A2egzJusu8m --- apps/content/docs/middleware.mdx | 14 + packages/effect/src/middleware.test.ts | 320 +++++++++++++++++- packages/effect/src/middleware.ts | 46 ++- .../server/src/middleware-decorated.test.ts | 22 ++ packages/server/src/middleware-decorated.ts | 5 +- packages/server/src/middleware.ts | 9 +- packages/server/src/procedure-client.test.ts | 34 ++ packages/server/src/procedure-client.ts | 7 +- 8 files changed, 446 insertions(+), 11 deletions(-) diff --git a/apps/content/docs/middleware.mdx b/apps/content/docs/middleware.mdx index 6b0441ed58..a687f08a5a 100644 --- a/apps/content/docs/middleware.mdx +++ b/apps/content/docs/middleware.mdx @@ -148,6 +148,20 @@ const cache = os.middleware(async ({ context, next, path }, input, done) => { }) ``` +## Middleware Signal + +Pass `signal` to `next` to replace the signal for downstream middleware and the handler, such as aborting them after a timeout. + +```ts +const timeout = os.middleware(async ({ signal, next }) => { + const timeoutSignal = AbortSignal.timeout(5000) + + return next({ + signal: signal ? AbortSignal.any([signal, timeoutSignal]) : timeoutSignal, // <- received by downstream middleware and the handler + }) +}) +``` + ## Inline Middleware Middleware is simply a function that can be defined inline with `.use`, which is useful for simple middleware cases. diff --git a/packages/effect/src/middleware.test.ts b/packages/effect/src/middleware.test.ts index 786c5932d0..3366ac0c60 100644 --- a/packages/effect/src/middleware.test.ts +++ b/packages/effect/src/middleware.test.ts @@ -1,6 +1,7 @@ import type { EffectWrapOptions, WithEffectContext } from './context' import { call, ORPCError, os } from '@orpc/server' -import { Context, Effect } from 'effect' +import { sleep } from '@orpc/shared' +import { Cause, Context, Effect, Option, References } from 'effect' import { z } from 'zod' import { handlerGen } from './handler' import { middlewareGen } from './middleware' @@ -19,6 +20,13 @@ class Service1 extends Context.Service< } >()('Service1') {} +class Service2 extends Context.Service< + Service2, + { + readonly id: string + } +>()('Service2') {} + describe('middlewareGen', () => { it('continues the chain and can extend the context', async () => { const procedure = os @@ -206,6 +214,316 @@ describe('middlewareGen', () => { ) }) + describe('interruption', () => { + let log: string[] = [] + + beforeEach(() => { + log = [] + }) + + const tracked = (name: string) => middlewareGen(function* ({ next }) { + return yield* Effect.acquireUseRelease( + Effect.sync(() => log.push(`${name}-acquire`)), + () => next(), + () => Effect.sync(() => log.push(`${name}-release`)), + ) + }) + + const timeoutFallback = middlewareGen(function* ({ next }, _, done) { + return yield* next().pipe( + Effect.timeout('20 millis'), + Effect.catchIf(Cause.isTimeoutError, () => Effect.sync(() => { + log.push('mw-timeout') + return done({ output: 'fallback' }) + })), + ) + }) + + it('waits for the downstream chain before running finalizers when the request aborts', async () => { + const procedure = os + .use(tracked('mw')) + .handler(async ({ signal }) => { + await sleep(50) + log.push(`handler-done:aborted=${signal?.aborted}`) + return 1 + }) + + const signal = AbortSignal.timeout(10) + await expect(call(procedure, undefined, { signal })).rejects.toSatisfy(error => error === signal.reason) + expect(log).toEqual(['mw-acquire', 'handler-done:aborted=true', 'mw-release']) + }) + + it('unwinds nested middleware in order when the request aborts', async () => { + const procedure = os + .use(tracked('outer')) + .use(async ({ next }) => { + try { + return await next() + } + finally { + log.push('plain-finally') + } + }) + .use(tracked('inner')) + .handler(async ({ signal }) => { + await sleep(50) + log.push(`handler-done:aborted=${signal?.aborted}`) + return 1 + }) + + await expect(call(procedure, undefined, { signal: AbortSignal.timeout(10) })).rejects.toThrow() + expect(log).toEqual([ + 'outer-acquire', + 'inner-acquire', + 'handler-done:aborted=true', + 'inner-release', + 'plain-finally', + 'outer-release', + ]) + }) + + it('interrupts a downstream handlerGen when the request aborts', async () => { + const procedure = os + .use(middlewareGen(function* ({ next }) { + return yield* next().pipe( + Effect.ensuring(Effect.sync(() => log.push('mw-release'))), + ) + })) + .handler(handlerGen(function* () { + return yield* Effect.never.pipe( + Effect.onInterrupt(() => Effect.sleep('10 millis').pipe( + Effect.andThen(Effect.sync(() => log.push('handler-interrupted'))), + )), + ) + })) + + const signal = AbortSignal.timeout(10) + await expect(call(procedure, undefined, { signal })).rejects.toSatisfy(error => error === signal.reason) + expect(log).toEqual(['handler-interrupted', 'mw-release']) + }) + + it('interrupts a downstream handlerGen when Effect.timeout interrupts next', async () => { + const procedure = os + .use(timeoutFallback) + .handler(handlerGen(function* () { + yield* Effect.sleep('1 second').pipe( + Effect.onInterrupt(() => Effect.sync(() => log.push('handler-interrupted'))), + ) + return 'handler' + })) + + await expect(call(procedure)).resolves.toEqual('fallback') + expect(log).toEqual(['handler-interrupted', 'mw-timeout']) + }) + + it('aborts the signal of a plain handler and waits for it when Effect.timeout interrupts next', async () => { + const procedure = os + .use(timeoutFallback) + .handler(async ({ signal }) => { + await new Promise(resolve => signal?.addEventListener('abort', resolve)) + await sleep(10) + log.push('handler-aborted') + throw new Error('dropped after interruption') + }) + + await expect(call(procedure)).resolves.toEqual('fallback') + expect(log).toEqual(['handler-aborted', 'mw-timeout']) + }) + + it('links a signal passed to next with interruption', async () => { + const controller = new AbortController() + let handlerSignal: AbortSignal | undefined + + const procedure = os + .use(middlewareGen(function* ({ next }) { + return yield* next({ signal: controller.signal }) + })) + .handler(({ signal }) => { + handlerSignal = signal + return 'output' + }) + + await expect(call(procedure)).resolves.toEqual('output') + expect(handlerSignal?.aborted).toBe(false) + + controller.abort() + expect(handlerSignal?.aborted).toBe(true) + }) + }) + + describe('effect context propagation', () => { + it('propagates services, spans and log annotations around next to downstream effects', async () => { + const procedure = os + .use(middlewareGen(function* ({ next }) { + return yield* next().pipe( + Effect.provideService(Service1, { id: 'Service1' }), + Effect.withSpan('mw-span'), + Effect.annotateLogs({ requestId: 'req-1' }), + ) + })) + .handler(handlerGen(function* () { + const service = yield* Effect.serviceOption(Service1) + const span = yield* Effect.currentSpan + const annotations = yield* References.CurrentLogAnnotations + + return { service: Option.getOrUndefined(service)?.id, span: span.name, annotations } + })) + + await expect(call(procedure)).resolves.toEqual({ + service: 'Service1', + span: 'mw-span', + annotations: { requestId: 'req-1' }, + }) + }) + + it('passes the propagated context through non-effect middleware', async () => { + const procedure = os + .$context>() + .use(middlewareGen(function* ({ next }) { + return yield* next().pipe(Effect.provideService(Service1, { id: 'Service1' })) + })) + .use(({ context, next }) => next({ + context: { 'effect/context': context['effect/context'].pipe(Context.add(Service2, { id: 'Service2' })) }, + })) + .handler(handlerGen(function* () { + const service1 = yield* Effect.serviceOption(Service1) + const service2 = yield* Effect.serviceOption(Service2) + + return [Option.getOrUndefined(service1)?.id, Option.getOrUndefined(service2)?.id] + })) + + await expect(call(procedure, undefined, { context: { 'effect/context': Context.empty() } })) + .resolves + .toEqual(['Service1', 'Service2']) + }) + + it('layers services passed to next over the fiber context', async () => { + const procedure = os + .$context>() + .use(middlewareGen(function* ({ context, next }) { + return yield* next({ + context: { 'effect/context': context['effect/context'].pipe(Context.add(Service2, { id: 'from-next' })) }, + }).pipe( + Effect.provideService(Service2, { id: 'from-fiber' }), + Effect.annotateLogs({ requestId: 'req-1' }), + ) + })) + .handler(handlerGen(function* () { + const service1 = yield* Service1 + const service2 = yield* Service2 + const annotations = yield* References.CurrentLogAnnotations + + return { service1: service1.id, service2: service2.id, annotations } + })) + + await expect(call( + procedure, + undefined, + { context: { 'effect/context': Context.make(Service1, { id: 'Service1' }) } }, + )).resolves.toEqual({ + service1: 'Service1', + service2: 'from-next', + annotations: { requestId: 'req-1' }, + }) + }) + + it('keeps the inner span and annotations when a nested middleware passes effect/context to next', async () => { + const procedure = os + .$context>() + .use(middlewareGen(function* ({ next }) { + return yield* next().pipe( + Effect.withSpan('outer'), + Effect.annotateLogs({ who: 'outer' }), + ) + })) + .use(middlewareGen(function* ({ context, next }) { + return yield* next({ + context: { 'effect/context': context['effect/context'].pipe(Context.add(Service2, { id: 'Service2' })) }, + }).pipe( + Effect.withSpan('inner'), + Effect.annotateLogs({ who: 'inner' }), + ) + })) + .handler(handlerGen(function* () { + const span = yield* Effect.currentSpan + const annotations = yield* References.CurrentLogAnnotations + + return { service2: (yield* Service2).id, span: span.name, annotations } + })) + + await expect(call( + procedure, + undefined, + { context: { 'effect/context': Context.make(Service1, { id: 'Service1' }) } }, + )).resolves.toEqual({ + service2: 'Service2', + span: 'inner', + annotations: { who: 'inner' }, + }) + }) + + it('layers a service whose value is undefined', async () => { + class Marker extends Context.Service()('Marker') {} + + const procedure = os + .$context>() + .use(middlewareGen(function* ({ context, next }) { + return yield* next({ + context: { 'effect/context': context['effect/context'].pipe(Context.add(Marker, undefined)) }, + }) + })) + .handler(handlerGen(function* () { + return Option.isSome(yield* Effect.serviceOption(Marker)) + })) + + await expect(call(procedure, undefined, { context: { 'effect/context': Context.empty() } })).resolves.toBe(true) + }) + + it('lets effect/context passed to next override an incoming service', async () => { + const procedure = os + .$context>() + .use(middlewareGen(function* ({ context, next }) { + return yield* next({ + context: { 'effect/context': context['effect/context'].pipe(Context.add(Service2, { id: 'overridden' })) }, + }) + })) + .handler(handlerGen(function* () { + return (yield* Service2).id + })) + + await expect(call( + procedure, + undefined, + { context: { 'effect/context': Context.make(Service2, { id: 'incoming' }) } }, + )).resolves.toEqual('overridden') + }) + + it.each([false, true])('nests effect/wrap spans of downstream effects under the middleware span (adds a service: %s)', async (addsService) => { + let index = 0 + + const procedure = os + .$context>() + .use(middlewareGen(function* ({ context, next }) { + return yield* (addsService + ? next({ context: { 'effect/context': context['effect/context'].pipe(Context.add(Service2, { id: 'Service2' })) } }) + : next()) + })) + .handler(handlerGen(function* () { + const span = yield* Effect.currentSpan + const parent = Option.getOrUndefined(span.parent) + + return [span.name, parent?._tag === 'Span' ? parent.name : undefined] + })) + + await expect(call(procedure, undefined, { + context: { + 'effect/context': Context.empty(), + 'effect/wrap': effect => effect.pipe(Effect.withSpan(`wrap-${index++}`)), + }, + })).resolves.toEqual(['wrap-1', 'wrap-0']) + }) + }) + it('forwards effect/context and the provided signal to runPromise', async () => { const controller = new AbortController() const effectContext = Context.empty().pipe(Context.add(Service1, { id: 'Service1' })) diff --git a/packages/effect/src/middleware.ts b/packages/effect/src/middleware.ts index 6fd0a8b3b0..bd5f456ccc 100644 --- a/packages/effect/src/middleware.ts +++ b/packages/effect/src/middleware.ts @@ -2,6 +2,7 @@ import type { ErrorMap } from '@orpc/contract' import type { Context, Middleware, MiddlewareDone, MiddlewareNextOptions, MiddlewareOptions, MiddlewareResult, ORPCErrorConstructorMap } from '@orpc/server' import type { MaybeOptionalOptions } from '@orpc/shared' import type { InferEffectServices, WithEffectContext } from './context' +import { anyAbortSignal, resolveMaybeOptionalOptions } from '@orpc/shared' import { Effect, Context as EffectContext } from 'effect' import { runPromise } from './runtime' @@ -59,10 +60,27 @@ export function middlewareGen< middleware: MiddlewareGen>, ): Middleware { const mid: Middleware = (opts, input, done) => { - const next: MiddlewareGenNext = (...rest: MaybeOptionalOptions>) => Effect.tryPromise({ - try: async () => opts.next(...rest), - catch: error => error, - }) + const next: MiddlewareGenNext = (...rest: MaybeOptionalOptions>) => Effect.contextWith( + fiberContext => Effect.callback>, unknown>((resume, interruptSignal) => { + const nextOptions = resolveMaybeOptionalOptions(rest) + + const promise = Promise.resolve(opts.next({ + context: { + ...nextOptions.context, + 'effect/context': layerChangedServices(fiberContext, nextOptions.context?.['effect/context'], opts.context['effect/context']), + }, + signal: anyAbortSignal([nextOptions.signal ?? opts.signal, interruptSignal]), + } as any)) + + promise.then( + result => resume(Effect.succeed(result)), + error => resume(Effect.fail(error)), + ) + + // On interruption, wait for the chain to settle. + return Effect.promise(() => promise.catch(() => {})) + }), + ) const effectContext = opts.context['effect/context'] const wrap = opts.context['effect/wrap'] as WithEffectContext['effect/wrap'] @@ -84,3 +102,23 @@ export function middlewareGen< return mid } + +/** + * An `effect/context` passed to `next` is the incoming one plus changes, so only the changed services + * are layered over the fiber context, keeping the fiber's newer services, spans and log annotations. + */ +function layerChangedServices(fiberContext: EffectContext.Context, nextServices: unknown, incomingServices: unknown): EffectContext.Context { + if (!EffectContext.isContext(nextServices) || nextServices === incomingServices) { + return fiberContext + } + + const services = new Map(fiberContext.mapUnsafe) + + for (const [key, service] of nextServices.mapUnsafe) { + if (!EffectContext.isContext(incomingServices) || !incomingServices.mapUnsafe.has(key) || incomingServices.mapUnsafe.get(key) !== service) { + services.set(key, service) + } + } + + return EffectContext.makeUnsafe(services) +} diff --git a/packages/server/src/middleware-decorated.test.ts b/packages/server/src/middleware-decorated.test.ts index 9694baa4fb..6ca1e5a4a5 100644 --- a/packages/server/src/middleware-decorated.test.ts +++ b/packages/server/src/middleware-decorated.test.ts @@ -150,6 +150,28 @@ describe('decorateMiddleware', () => { expect(done).not.toHaveBeenCalled() }) + const outerSignal = new AbortController().signal + const signal1 = new AbortController().signal + const signal2 = new AbortController().signal + + it.each([ + [undefined, undefined, outerSignal, outerSignal], + [signal1, undefined, signal1, signal1], + [signal1, signal2, signal1, signal2], + ])('forwards signal overrides through the combined chain (%#)', async (mid1Signal, mid2Signal, expectedMid2Signal, expectedNextSignal) => { + const mid1 = vi.fn().mockImplementation(async ({ next }) => next({ signal: mid1Signal })) + const mid2 = vi.fn().mockImplementation(async ({ next }) => next({ signal: mid2Signal })) + + const combined = decorateMiddleware(mid1 as any).use(mid2 as any) + const next = vi.fn().mockImplementation(async ({ context }) => ({ output: 'final', context })) + const done = vi.fn() + + await (combined as any)({ context: {}, signal: outerSignal, next }, 'input', done) + + expect(mid2).toHaveBeenCalledWith(expect.objectContaining({ signal: expectedMid2Signal }), 'input', done) + expect(next).toHaveBeenCalledWith(expect.objectContaining({ signal: expectedNextSignal })) + }) + it('can stop early with done', async () => { const mid1 = vi.fn().mockImplementation(async ({ next }) => { return await next({ context: { mid1: true } }) diff --git a/packages/server/src/middleware-decorated.ts b/packages/server/src/middleware-decorated.ts index e72f1b15c2..64cf7ce76f 100644 --- a/packages/server/src/middleware-decorated.ts +++ b/packages/server/src/middleware-decorated.ts @@ -90,16 +90,19 @@ export function decorateMiddleware< { ...opts, async next(nextOpts1: undefined | MiddlewareNextOptions) { + const signal = nextOpts1?.signal ?? opts.signal + const result: MiddlewareResult = await usedMiddleware({ ...opts, context: { ...opts.context, ...nextOpts1?.context }, + signal, next(nextOpts2: undefined | MiddlewareNextOptions, ...args: []) { return opts.next({ - ...opts, context: { ...nextOpts1?.context, ...nextOpts2?.context, }, + signal: nextOpts2?.signal ?? signal, }, ...args) }, } as any, ...args) diff --git a/packages/server/src/middleware.ts b/packages/server/src/middleware.ts index 517eb8441e..3445ef7fa7 100644 --- a/packages/server/src/middleware.ts +++ b/packages/server/src/middleware.ts @@ -8,9 +8,12 @@ export type MiddlewareResult = Promisable< context: TOutContext }> -export type MiddlewareNextOptions = object extends TOutContext - ? { context?: TOutContext | undefined } - : { context: TOutContext } +export type MiddlewareNextOptions + = & (object extends TOutContext ? { context?: TOutContext | undefined } : { context: TOutContext }) + & { + /** Overrides the signal for the rest of the chain. */ + signal?: AbortSignal | undefined + } export interface MiddlewareNext { ( diff --git a/packages/server/src/procedure-client.test.ts b/packages/server/src/procedure-client.test.ts index 16019ddb03..26fe0e77ec 100644 --- a/packages/server/src/procedure-client.test.ts +++ b/packages/server/src/procedure-client.test.ts @@ -298,6 +298,40 @@ describe('createProcedureClient', () => { ) }) + it('middleware can replace the signal for the rest of the chain', async () => { + const signal = new AbortController().signal + const overrideSignal = new AbortController().signal + + mid1.mockImplementationOnce(async ({ next }) => next({ signal: overrideSignal })) + + await expect(client(undefined, { signal })).resolves.toBe('__OUTPUT__') + + expect(interceptor).toHaveBeenCalledWith(expect.objectContaining({ signal })) + expect(mid1).toHaveBeenCalledWith(expect.objectContaining({ signal }), undefined, expect.any(Function)) + expect(mid2).toHaveBeenCalledWith(expect.objectContaining({ signal: overrideSignal }), undefined, expect.any(Function)) + expect(handler).toHaveBeenCalledWith(expect.objectContaining({ signal: overrideSignal }), undefined) + }) + + it('middleware signal override applies to its own next call only', async () => { + const signal = new AbortController().signal + const overrideSignal = new AbortController().signal + + mid1.mockImplementationOnce(async ({ next }) => { + await next({ signal: overrideSignal }) + return next({ context: { second: true } }) + }) + + await expect(client(undefined, { signal })).resolves.toBe('__OUTPUT__') + + expect(mid2).toHaveBeenCalledTimes(2) + expect(mid2).toHaveBeenNthCalledWith(1, expect.objectContaining({ signal: overrideSignal }), undefined, expect.any(Function)) + expect(mid2).toHaveBeenNthCalledWith(2, expect.objectContaining({ signal, context: { second: true } }), undefined, expect.any(Function)) + + expect(handler).toHaveBeenCalledTimes(2) + expect(handler).toHaveBeenNthCalledWith(1, expect.objectContaining({ signal: overrideSignal }), undefined) + expect(handler).toHaveBeenNthCalledWith(2, expect.objectContaining({ signal }), undefined) + }) + it('middleware can early return with done helper', async () => { mid1.mockImplementationOnce((options, input, done) => done({ output: 'MID1__EARLY_RETURN__' })) diff --git a/packages/server/src/procedure-client.ts b/packages/server/src/procedure-client.ts index 40a9f3b0b6..cca947d5c7 100644 --- a/packages/server/src/procedure-client.ts +++ b/packages/server/src/procedure-client.ts @@ -200,6 +200,7 @@ async function executeProcedureInternal(procedure: AnyProcedure, options: Proced midIndex: number, context: Context, input: unknown, + signal: AbortSignal | undefined, ): Promise<{ output: unknown, context: Record }> => { let currentInput = input let middlewareInput = input @@ -245,6 +246,7 @@ async function executeProcedureInternal(procedure: AnyProcedure, options: Proced { ...options, context, + signal, next: (...rest) => { const nextOptions = resolveMaybeOptionalOptions(rest) // context can be undefined when all field is optional @@ -254,6 +256,7 @@ async function executeProcedureInternal(procedure: AnyProcedure, options: Proced midIndex + 1, { ...context, ...nextContext }, currentInput, + nextOptions.signal ?? signal, ) }, lastEventId: options.lastEventId, @@ -269,7 +272,7 @@ async function executeProcedureInternal(procedure: AnyProcedure, options: Proced else { currentOutput = await runWithSpan( 'handler', - () => procedure['~orpc'].handler({ ...options, context, input: currentInput }, currentInput), + () => procedure['~orpc'].handler({ ...options, context, input: currentInput, signal }, currentInput), ) } @@ -289,6 +292,6 @@ async function executeProcedureInternal(procedure: AnyProcedure, options: Proced return { output: currentOutput, context: currentContext } } - const { output } = await next(0, options.context, options.input) + const { output } = await next(0, options.context, options.input, options.signal) return output }