Skip to content

Commit 03393ab

Browse files
juliusmarmingeclaude
juliusmarminge
and
claude
committed
refactor(server): instrument WS RPCs in their own middleware
Authorization and instrumentation are separate RpcMiddleware services again. RpcScopeAuthorization is unchanged from main; RpcInstrumentation is added after it, so it wraps authorization and still records rejected calls. Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
1 parent 2422c9a commit 03393ab

9 files changed

Lines changed: 72 additions & 45 deletions

File tree

‎apps/server/src/auth/RpcAuthorization.test.ts‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,9 @@ import {
1515
RPC_REQUIRED_SCOPES,
1616
requiredScopeForRpcMethod,
1717
requiredScopeForDeviceList,
18-
wsRpcGuardLayer,
18+
rpcScopeAuthorizationLayer,
1919
} from "./RpcAuthorization.ts";
20+
import { rpcInstrumentationLayer } from "../observability/RpcInstrumentation.ts";
2021

2122
describe("RPC authorization scopes", () => {
2223
it("declares exactly one scope for every RPC in the server group", () => {
@@ -140,7 +141,8 @@ describe("RPC scope middleware", () => {
140141
group.toLayerHandler(WS_METHODS.serverRetryResourceTelemetry, () =>
141142
Effect.sync(() => handled.push("retry")).pipe(Effect.andThen(Effect.never)),
142143
),
143-
wsRpcGuardLayer([AuthOrchestrationReadScope]),
144+
rpcScopeAuthorizationLayer([AuthOrchestrationReadScope]),
145+
rpcInstrumentationLayer,
144146
),
145147
),
146148
);

‎apps/server/src/auth/RpcAuthorization.ts‎

Lines changed: 7 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -10,16 +10,14 @@ import {
1010
ORCHESTRATION_V2_WS_METHODS,
1111
type AuthEnvironmentScope,
1212
EnvironmentAuthorizationError,
13+
RpcScopeAuthorization,
1314
WS_METHODS,
1415
WsRpcGroup,
15-
WsRpcGuard,
1616
} from "@t3tools/contracts";
1717
import * as Effect from "effect/Effect";
1818
import * as Layer from "effect/Layer";
1919
import type * as RpcGroup from "effect/unstable/rpc/RpcGroup";
2020

21-
import { instrumentRpc } from "../observability/RpcInstrumentation.ts";
22-
2321
type WsRpcMethod = RpcGroup.Rpcs<typeof WsRpcGroup>["_tag"];
2422

2523
/**
@@ -225,17 +223,13 @@ export const rpcAuthorizationError = (requiredScope: AuthEnvironmentScope) =>
225223
requiredScope,
226224
});
227225

228-
/**
229-
* Guards every RPC on one connection: authorizes it against that connection's session scopes and
230-
* instruments it. A rejected call still gets its span and a failed request metric.
231-
*/
232-
export const wsRpcGuardLayer = (scopes: ReadonlyArray<AuthEnvironmentScope>) =>
233-
Layer.succeed(WsRpcGuard)((effect, { rpc }) => {
226+
/** Authorizes every RPC on one connection against that connection's session scopes. */
227+
export const rpcScopeAuthorizationLayer = (scopes: ReadonlyArray<AuthEnvironmentScope>) =>
228+
Layer.succeed(RpcScopeAuthorization)((effect, { rpc }) => {
234229
const requiredScope = requiredScopeForRpcMethod(rpc._tag);
235-
return instrumentRpc(
236-
rpc,
237-
scopes.includes(requiredScope) ? effect : Effect.fail(rpcAuthorizationError(requiredScope)),
238-
);
230+
return scopes.includes(requiredScope)
231+
? effect
232+
: Effect.fail(rpcAuthorizationError(requiredScope));
239233
});
240234

241235
/** Retrying can install or restart tools even though ordinary listing is readable. */

‎apps/server/src/mcp/PreviewAutomationBroker.test.ts‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,8 @@ import * as TestClock from "effect/testing/TestClock";
2929
import * as RpcGroup from "effect/unstable/rpc/RpcGroup";
3030
import * as RpcTest from "effect/unstable/rpc/RpcTest";
3131

32-
import { wsRpcGuardLayer } from "../auth/RpcAuthorization.ts";
32+
import { rpcScopeAuthorizationLayer } from "../auth/RpcAuthorization.ts";
33+
import { rpcInstrumentationLayer } from "../observability/RpcInstrumentation.ts";
3334
import * as PreviewAutomationBroker from "./PreviewAutomationBroker.ts";
3435

3536
const makeBroker = PreviewAutomationBroker.make.pipe(Effect.provide(NodeServices.layer));
@@ -1258,11 +1259,12 @@ it.effect("evicts an unanswered host and lets later calls use a healthy runtime"
12581259
);
12591260
const client = yield* RpcTest.makeClient(group).pipe(
12601261
Effect.provide(
1261-
Layer.merge(
1262+
Layer.mergeAll(
12621263
group.toLayer({
12631264
[WS_METHODS.previewAutomationConnect]: (host) => Stream.unwrap(broker.connect(host)),
12641265
}),
1265-
wsRpcGuardLayer([AuthOrchestrationOperateScope]),
1266+
rpcScopeAuthorizationLayer([AuthOrchestrationOperateScope]),
1267+
rpcInstrumentationLayer,
12661268
),
12671269
),
12681270
);

‎apps/server/src/observability/RpcInstrumentation.test.ts‎

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import { assert, describe, it } from "@effect/vitest";
22
import {
33
AuthOrchestrationOperateScope,
44
AuthOrchestrationReadScope,
5+
type AuthEnvironmentScope,
56
ScheduledTaskError,
67
ScheduledTaskId,
78
WS_METHODS,
@@ -20,7 +21,8 @@ import * as Tracer from "effect/Tracer";
2021
import * as RpcTest from "effect/unstable/rpc/RpcTest";
2122
import * as TestClock from "effect/testing/TestClock";
2223

23-
import { RPC_REQUIRED_SCOPES, wsRpcGuardLayer } from "../auth/RpcAuthorization.ts";
24+
import { RPC_REQUIRED_SCOPES, rpcScopeAuthorizationLayer } from "../auth/RpcAuthorization.ts";
25+
import { rpcInstrumentationLayer } from "./RpcInstrumentation.ts";
2426

2527
type WsRpcMethod = keyof typeof RPC_REQUIRED_SCOPES;
2628

@@ -33,7 +35,10 @@ const groupOf = <const Tags extends ReadonlyArray<WsRpcMethod>>(...tags: Tags) =
3335
),
3436
);
3537

36-
const readOnlyGuard = wsRpcGuardLayer([AuthOrchestrationReadScope]);
38+
/** The middleware ws.ts installs for a connection with `scopes`. */
39+
const connectionMiddleware = (scopes: ReadonlyArray<AuthEnvironmentScope>) =>
40+
Layer.merge(rpcScopeAuthorizationLayer(scopes), rpcInstrumentationLayer);
41+
const readOnlyConnection = connectionMiddleware([AuthOrchestrationReadScope]);
3742
const taskId = ScheduledTaskId.make("scheduled-task:instrumented");
3843
const rpcSpanDefaults = { "rpc.transport": "websocket", "rpc.system": "effect-rpc" };
3944

@@ -64,7 +69,7 @@ const rpcSpans = (ended: ReadonlyArray<Tracer.NativeSpan>) =>
6469
ended.filter((span) => span.name.startsWith("ws.rpc."));
6570

6671
// RpcTest keeps Effect's own RpcServer/RpcClient spans, which ws.ts turns off with
67-
// `disableTracing: true`. Everything else comes from the guard or the handlers.
72+
// `disableTracing: true`. Everything else comes from the middleware or the handlers.
6873
const appSpans = (ended: ReadonlyArray<Tracer.NativeSpan>) =>
6974
ended.filter((span) => !/^Rpc(Server|Client)\./.test(span.name));
7075

@@ -95,7 +100,7 @@ const requestDuration = (snapshots: ReadonlyArray<Metric.Metric.Snapshot>, metho
95100
snapshot.attributes?.["method"] === method,
96101
)?.state;
97102

98-
describe("WS RPC guard instrumentation", () => {
103+
describe("WS RPC instrumentation middleware", () => {
99104
it.effect("records one span and request metric per call, including rejected calls", () =>
100105
withTelemetry((ended) =>
101106
Effect.gen(function* () {
@@ -119,7 +124,7 @@ describe("WS RPC guard instrumentation", () => {
119124
),
120125
),
121126
group.toLayerHandler(WS_METHODS.serverRetryResourceTelemetry, () =>
122-
Effect.die("the guard let a rejected call through"),
127+
Effect.die("authorization let a rejected call through"),
123128
),
124129
group.toLayerHandler(WS_METHODS.pullRequestsSubscribeRefreshes, () =>
125130
Stream.make(1, 2),
@@ -131,7 +136,7 @@ describe("WS RPC guard instrumentation", () => {
131136
),
132137
),
133138
group.toLayerHandler(WS_METHODS.serverGetSettings, () => Effect.die("broken")),
134-
readOnlyGuard,
139+
readOnlyConnection,
135140
),
136141
),
137142
);
@@ -259,7 +264,7 @@ describe("WS RPC guard instrumentation", () => {
259264
),
260265
),
261266
),
262-
readOnlyGuard,
267+
readOnlyConnection,
263268
),
264269
),
265270
);
@@ -307,7 +312,7 @@ describe("WS RPC guard instrumentation", () => {
307312
message: Option.none(),
308313
}).pipe(Effect.withSpan("signalProcess.child")),
309314
),
310-
wsRpcGuardLayer([AuthOrchestrationReadScope, AuthOrchestrationOperateScope]),
315+
connectionMiddleware([AuthOrchestrationReadScope, AuthOrchestrationOperateScope]),
311316
),
312317
),
313318
);

‎apps/server/src/observability/RpcInstrumentation.ts‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,14 @@
1-
import { ORCHESTRATION_V2_WS_METHODS, WS_METHODS, type WsRpcGroup } from "@t3tools/contracts";
1+
import {
2+
ORCHESTRATION_V2_WS_METHODS,
3+
RpcInstrumentation,
4+
WS_METHODS,
5+
type WsRpcGroup,
6+
} from "@t3tools/contracts";
27
import * as Clock from "effect/Clock";
38
import * as Duration from "effect/Duration";
49
import * as Effect from "effect/Effect";
510
import type * as Exit from "effect/Exit";
11+
import * as Layer from "effect/Layer";
612
import * as Metric from "effect/Metric";
713
import * as References from "effect/References";
814
import type * as Rpc from "effect/unstable/rpc/Rpc";
@@ -269,3 +275,8 @@ export const instrumentRpc = <A, E, R>(
269275
}),
270276
);
271277
};
278+
279+
/** Instruments every WebSocket RPC. */
280+
export const rpcInstrumentationLayer = Layer.succeed(RpcInstrumentation)((effect, { rpc }) =>
281+
instrumentRpc(rpc, effect),
282+
);

‎apps/server/src/ws.ts‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -206,8 +206,9 @@ import * as EnvironmentAuth from "./auth/EnvironmentAuth.ts";
206206
import {
207207
requiredScopeForDeviceList,
208208
rpcAuthorizationError,
209-
wsRpcGuardLayer,
209+
rpcScopeAuthorizationLayer,
210210
} from "./auth/RpcAuthorization.ts";
211+
import { rpcInstrumentationLayer } from "./observability/RpcInstrumentation.ts";
211212
import * as ProcessDiagnostics from "./diagnostics/ProcessDiagnostics.ts";
212213
import * as ProcessResourceMonitor from "./diagnostics/ProcessResourceMonitor.ts";
213214
import * as ResourceTelemetry from "./resourceTelemetry/ResourceTelemetry.ts";
@@ -1310,8 +1311,8 @@ const makeWsRpcLayer = (
13101311
const processResourceMonitor = yield* ProcessResourceMonitor.ProcessResourceMonitor;
13111312
const resourceTelemetry = yield* ResourceTelemetry.ResourceTelemetry;
13121313
const relayClient = yield* RelayClient.RelayClient;
1313-
// WsRpcGuard checks each RPC's declared scope before its handler runs.
1314-
// This covers the one RPC whose scope depends on its input.
1314+
// RpcScopeAuthorization checks each RPC's declared scope before its handler
1315+
// runs. This covers the one RPC whose scope depends on its input.
13151316
const authorizeEffect = <A, E, R>(
13161317
requiredScope: AuthEnvironmentScope,
13171318
effect: Effect.Effect<A, E, R>,
@@ -3128,7 +3129,9 @@ export const websocketRpcRouteLayer = Layer.unwrap(
31283129
const { protocol, httpEffect } = yield* RpcServer.makeProtocolWithHttpEffectWebsocket;
31293130
yield* RpcServer.make(ServerWsRpcGroup, { disableTracing: true }).pipe(
31303131
Effect.provideService(RpcServer.Protocol, withTerminalOutputWindow(protocol)),
3131-
Effect.provide(wsRpcGuardLayer(session.scopes)),
3132+
Effect.provide(
3133+
Layer.merge(rpcScopeAuthorizationLayer(session.scopes), rpcInstrumentationLayer),
3134+
),
31323135
Effect.forkScoped,
31333136
);
31343137
// @effect-diagnostics-next-line returnEffectInGen:off

‎docs/internals/effect-services.md‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -29,10 +29,10 @@ message) are fine next to the service; the capability itself is the method.
2929
.pipe(Effect.mapError((cause) => new ProjectCreateNewError({ cause }))),
3030
```
3131

32-
Handlers don't add their own spans or request metrics. The `WsRpcGroup` middleware
33-
([`RpcAuthorization.ts`](../../apps/server/src/auth/RpcAuthorization.ts)) authorizes and instruments
34-
every call. A handler with per-call context, such as a thread id, adds it with
35-
`Effect.annotateCurrentSpan`.
32+
Handlers don't add their own spans or request metrics. `WsRpcGroup` middleware authorizes every
33+
call ([`RpcAuthorization.ts`](../../apps/server/src/auth/RpcAuthorization.ts)) and instruments it
34+
([`RpcInstrumentation.ts`](../../apps/server/src/observability/RpcInstrumentation.ts)). A handler
35+
with per-call context, such as a thread id, adds it with `Effect.annotateCurrentSpan`.
3636

3737
## Shape of a service module
3838

‎docs/internals/environment-auth.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ HTTP so long-lived tokens stay out of socket URLs. Browser sessions can
2929
authenticate the upgrade with their cookie. A successful handshake grants no
3030
extra authority: [every RPC declares a required
3131
scope](../../apps/server/src/auth/RpcAuthorization.ts), and the WebSocket RPC
32-
group's `WsRpcGuard` middleware checks it before any handler runs.
32+
group's `RpcScopeAuthorization` middleware checks it before any handler runs.
3333

3434
Desktop restarts forget the previous local bearer token, so its reusable
3535
bootstrap grant replaces earlier sessions for the same subject and method.

‎packages/contracts/src/rpc.ts‎

Lines changed: 19 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1690,15 +1690,22 @@ const WsSubscribeResourceTelemetryRpc = Rpc.make(WS_METHODS.subscribeResourceTel
16901690
});
16911691

16921692
/**
1693-
* Wraps every WebSocket RPC on the server. It checks the connection's scopes
1694-
* against the scope the RPC declares before the handler runs, and records the
1695-
* call's span and metrics, including for calls it rejects. Every RPC in
1696-
* `WsRpcGroup` carries it, so a handler cannot be added without authorization
1697-
* or telemetry.
1693+
* Checks the connection's scopes against the scope each RPC declares, before
1694+
* the handler runs. Every RPC in `WsRpcGroup` carries it, so a handler cannot
1695+
* be added without authorization.
16981696
*/
1699-
export class WsRpcGuard extends RpcMiddleware.Service<WsRpcGuard>()("t3/contracts/WsRpcGuard", {
1700-
error: EnvironmentAuthorizationError,
1701-
}) {}
1697+
export class RpcScopeAuthorization extends RpcMiddleware.Service<RpcScopeAuthorization>()(
1698+
"t3/contracts/RpcScopeAuthorization",
1699+
{ error: EnvironmentAuthorizationError },
1700+
) {}
1701+
1702+
/**
1703+
* Records each RPC's span and request metrics on the server. It wraps
1704+
* `RpcScopeAuthorization`, so calls that authorization rejects are recorded too.
1705+
*/
1706+
export class RpcInstrumentation extends RpcMiddleware.Service<RpcInstrumentation>()(
1707+
"t3/contracts/RpcInstrumentation",
1708+
) {}
17021709

17031710
export const WsRpcGroup = RpcGroup.make(
17041711
WsServerProbeRpc,
@@ -1875,4 +1882,7 @@ export const WsRpcGroup = RpcGroup.make(
18751882
WsOrchestrationV2SubscribeArchivedShellRpc,
18761883
WsOrchestrationV2SubscribeShellRpc,
18771884
WsOrchestrationV2SubscribeThreadRpc,
1878-
).middleware(WsRpcGuard);
1885+
)
1886+
// Middleware added later wraps middleware added earlier.
1887+
.middleware(RpcScopeAuthorization)
1888+
.middleware(RpcInstrumentation);

0 commit comments

Comments
 (0)