Skip to content
Source
Sign inSign up

RpcInstrumentation.ts

271 lines13.2 KB

1import { ORCHESTRATION_V2_WS_METHODS, WS_METHODS, type WsRpcGroup } from "@t3tools/contracts";
2import * as Effect from "effect/Effect";
3import * as Layer from "effect/Layer";
4import * as References from "effect/References";
5import type * as RpcGroup from "effect/rpc/RpcGroup";
6import * as RpcMiddleware from "effect/rpc/RpcMiddleware";
7​
8import { rpcRequestDuration, rpcRequestsTotal, withMetrics } from "./Metrics.ts";
9​
10type WsRpcMethod = RpcGroup.Rpcs<typeof WsRpcGroup>["_tag"];
11​
12// Trace queries and dashboards depend on these labels, including deviations from method prefixes.
13const RPC_AGGREGATES = {
14 [ORCHESTRATION_V2_WS_METHODS.dispatchCommand]: "orchestrationV2",
15 [ORCHESTRATION_V2_WS_METHODS.getWorkflowScript]: "orchestration",
16 [ORCHESTRATION_V2_WS_METHODS.getNativeChildWork]: "orchestration",
17 [ORCHESTRATION_V2_WS_METHODS.stopWorkflow]: "orchestration",
18 [ORCHESTRATION_V2_WS_METHODS.forceStopThread]: "orchestration",
19 [ORCHESTRATION_V2_WS_METHODS.getTurnItem]: "orchestration",
20 [ORCHESTRATION_V2_WS_METHODS.getTurnDiff]: "orchestration",
21 [ORCHESTRATION_V2_WS_METHODS.getFullThreadDiff]: "orchestration",
22 [ORCHESTRATION_V2_WS_METHODS.searchThreads]: "orchestration",
23 [ORCHESTRATION_V2_WS_METHODS.searchThread]: "orchestration",
24 [ORCHESTRATION_V2_WS_METHODS.searchThreadStream]: "orchestration",
25 [ORCHESTRATION_V2_WS_METHODS.getArchivedShellSnapshot]: "orchestration",
26 [ORCHESTRATION_V2_WS_METHODS.getThreadFamily]: "orchestration",
27 [ORCHESTRATION_V2_WS_METHODS.getThreadTitles]: "orchestration",
28 [ORCHESTRATION_V2_WS_METHODS.getThreadLinkMetadata]: "orchestration",
29 [ORCHESTRATION_V2_WS_METHODS.getWorktreeThreadIds]: "orchestration",
30 [ORCHESTRATION_V2_WS_METHODS.getOlderSettledThreads]: "orchestration",
31 [ORCHESTRATION_V2_WS_METHODS.getThreadProjection]: "orchestrationV2",
32 [ORCHESTRATION_V2_WS_METHODS.launchThread]: "orchestration",
33 [ORCHESTRATION_V2_WS_METHODS.subscribeShell]: "orchestrationV2",
34 [ORCHESTRATION_V2_WS_METHODS.subscribeThread]: "orchestrationV2",
35 [ORCHESTRATION_V2_WS_METHODS.subscribeCommandOutput]: "orchestrationV2",
36 [ORCHESTRATION_V2_WS_METHODS.subscribePromptSuggestion]: "orchestrationV2",
37 [WS_METHODS.projectsMutate]: "orchestration",
38 [WS_METHODS.serverProbe]: "server",
39 [WS_METHODS.serverGetConfig]: "server",
40 [WS_METHODS.serverGetNpmRelease]: "server",
41 [WS_METHODS.serverRefreshProviders]: "server",
42 [WS_METHODS.serverUpdateProvider]: "server",
43 [WS_METHODS.providerAuthStart]: "provider",
44 [WS_METHODS.providerConsumeResetCredit]: "provider",
45 [WS_METHODS.providerAuthComplete]: "provider",
46 [WS_METHODS.chatGptReconnectProfile]: "provider",
47 [WS_METHODS.chatGptImportProfile]: "provider",
48 [WS_METHODS.chatGptHandoffSubscribe]: "provider",
49 [WS_METHODS.providerAuthRespond]: "provider",
50 [WS_METHODS.providerAuthCancel]: "provider",
51 [WS_METHODS.providerAuthRemoveMethod]: "provider",
52 [WS_METHODS.providerAuthLogout]: "provider",
53 [WS_METHODS.providerAuthSubscribe]: "provider",
54 [WS_METHODS.providerInstallStart]: "provider",
55 [WS_METHODS.providerInstallCancel]: "provider",
56 [WS_METHODS.providerInstallSubscribe]: "provider",
57 [WS_METHODS.providerInstallRemove]: "provider",
58 [WS_METHODS.serverUpdateServer]: "server",
59 [WS_METHODS.serverUpdateServerWithProgress]: "server",
60 [WS_METHODS.serverCommitDesktopUpdate]: "server",
61 [WS_METHODS.serverUpsertKeybinding]: "server",
62 [WS_METHODS.serverRemoveKeybinding]: "server",
63 [WS_METHODS.serverRunStorageCleanup]: "server",
64 [WS_METHODS.serverGetStorageCleanupReport]: "server",
65 [WS_METHODS.serverGetDatabaseStorage]: "server",
66 [WS_METHODS.serverReclaimDatabaseSpace]: "server",
67 [WS_METHODS.peerLinksList]: "peerLinks",
68 [WS_METHODS.peerLinksLink]: "peerLinks",
69 [WS_METHODS.peerLinksUnlink]: "peerLinks",
70 [WS_METHODS.peerLinksAnswerRequest]: "peerLinks",
71 [WS_METHODS.serverUpdateSettings]: "server",
72 [WS_METHODS.serverSearchAcpRegistry]: "server",
73 [WS_METHODS.serverPrepareAcpRegistryAgent]: "server",
74 [WS_METHODS.serverUninstallAcpRegistryManagedBinary]: "server",
75 [WS_METHODS.serverAcceptAcpRegistryUrlAuth]: "server",
76 [WS_METHODS.serverListAcpRegistrySessions]: "server",
77 [WS_METHODS.serverImportAcpRegistrySession]: "server",
78 [WS_METHODS.serverDeleteAcpRegistrySession]: "server",
79 [WS_METHODS.serverListAcpRegistryProviders]: "server",
80 [WS_METHODS.serverSetAcpRegistryProvider]: "server",
81 [WS_METHODS.serverDisableAcpRegistryProvider]: "server",
82 [WS_METHODS.serverLogoutAcpRegistry]: "server",
83 [WS_METHODS.serverDiscoverSourceControl]: "server",
84 [WS_METHODS.serverCheckSecretSource]: "server",
85 [WS_METHODS.serverGetTraceDiagnostics]: "server",
86 [WS_METHODS.serverGetProcessDiagnostics]: "server",
87 [WS_METHODS.serverGetHostResources]: "server",
88 [WS_METHODS.serverGetProcessResourceHistory]: "server",
89 [WS_METHODS.serverGetResourceTelemetryHistory]: "server",
90 [WS_METHODS.serverRetryResourceTelemetry]: "server",
91 [WS_METHODS.serverCheckOtlpEndpoint]: "server",
92 [WS_METHODS.serverListTextGeneratorModels]: "server",
93 [WS_METHODS.serverPrepareTriage]: "server",
94 [WS_METHODS.serverGetUsageSummary]: "server",
95 [WS_METHODS.serverRefreshUsageRates]: "server",
96 [WS_METHODS.serverSignalProcess]: "server",
97 [WS_METHODS.serverReportClientActivity]: "server",
98 [WS_METHODS.scheduledTasksSubscribe]: "scheduledTasks",
99 [WS_METHODS.scheduledTasksUpsert]: "scheduledTasks",
100 [WS_METHODS.scheduledTasksSetEnabled]: "scheduledTasks",
101 [WS_METHODS.scheduledTasksDelete]: "scheduledTasks",
102 [WS_METHODS.scheduledTasksRunNow]: "scheduledTasks",
103 [WS_METHODS.scheduledTasksRotateWebhookToken]: "scheduledTasks",
104 [WS_METHODS.scheduledTasksListWebhookDeliveries]: "scheduledTasks",
105 [WS_METHODS.scheduledTasksGetWebhookDelivery]: "scheduledTasks",
106 [WS_METHODS.secretsAnswerRequest]: "secrets",
107 [WS_METHODS.guestShareCreate]: "guest-share",
108 [WS_METHODS.cloudGetRelayClientStatus]: "cloud",
109 [WS_METHODS.cloudInstallRelayClient]: "cloud",
110 [WS_METHODS.pullRequestsList]: "pull-requests",
111 [WS_METHODS.pullRequestsListStats]: "pull-requests",
112 [WS_METHODS.pullRequestsFilterOptions]: "pull-requests",
113 [WS_METHODS.pullRequestsSummary]: "pull-requests",
114 [WS_METHODS.pullRequestsRouting]: "pull-requests",
115 [WS_METHODS.pullRequestsRoutingIdentity]: "pull-requests",
116 [WS_METHODS.pullRequestsStack]: "pull-requests",
117 [WS_METHODS.pullRequestsLinkedThreads]: "pull-requests",
118 [WS_METHODS.pullRequestsDetail]: "pull-requests",
119 [WS_METHODS.pullRequestsPreview]: "pull-requests",
120 [WS_METHODS.pullRequestsChecks]: "pull-requests",
121 [WS_METHODS.pullRequestsActivity]: "pull-requests",
122 [WS_METHODS.pullRequestsReviewGuide]: "pull-requests",
123 [WS_METHODS.pullRequestsGenerateReviewGuide]: "pull-requests",
124 [WS_METHODS.pullRequestsThreadComments]: "pull-requests",
125 [WS_METHODS.pullRequestsDiffFileContents]: "pull-requests",
126 [WS_METHODS.pullRequestsFilesViewed]: "pull-requests",
127 [WS_METHODS.pullRequestsChangeBreakdown]: "pull-requests",
128 [WS_METHODS.pullRequestsRunAction]: "pull-requests",
129 [WS_METHODS.pullRequestsUpdate]: "pull-requests",
130 [WS_METHODS.pullRequestsComment]: "pull-requests",
131 [WS_METHODS.pullRequestsUpdateComment]: "pull-requests",
132 [WS_METHODS.pullRequestsSubmitReview]: "pull-requests",
133 [WS_METHODS.pullRequestsReplyToThread]: "pull-requests",
134 [WS_METHODS.pullRequestsSetThreadResolution]: "pull-requests",
135 [WS_METHODS.pullRequestsSetReaction]: "pull-requests",
136 [WS_METHODS.pullRequestsSetFilesViewed]: "pull-requests",
137 [WS_METHODS.pullRequestsInvalidate]: "pull-requests",
138 [WS_METHODS.pullRequestsReportState]: "pull-requests",
139 [WS_METHODS.pullRequestsSubscribeRefreshes]: "pull-requests",
140 [WS_METHODS.pullRequestsSubscribeLinkedThreadRefreshes]: "pull-requests",
141 [WS_METHODS.pullRequestsReviewerCandidates]: "pull-requests",
142 [WS_METHODS.pullRequestsRequestReviewers]: "pull-requests",
143 [WS_METHODS.pullRequestsLabelCandidates]: "pull-requests",
144 [WS_METHODS.pullRequestsSetLabels]: "pull-requests",
145 [WS_METHODS.sourceControlLookupRepository]: "source-control",
146 [WS_METHODS.sourceControlCloneRepository]: "source-control",
147 [WS_METHODS.sourceControlPublishRepository]: "source-control",
148 [WS_METHODS.sourceControlSignInStart]: "source-control",
149 [WS_METHODS.sourceControlSignInWait]: "source-control",
150 [WS_METHODS.sourceControlSignInCancel]: "source-control",
151 [WS_METHODS.sourceControlAppAccess]: "source-control",
152 [WS_METHODS.sourceControlStarred]: "source-control",
153 [WS_METHODS.sourceControlStar]: "source-control",
154 [WS_METHODS.projectCloneStart]: "source-control",
155 [WS_METHODS.projectCloneCancel]: "source-control",
156 [WS_METHODS.projectCloneRetry]: "source-control",
157 [WS_METHODS.subscribeProjectClones]: "source-control",
158 [WS_METHODS.projectsListEntries]: "workspace",
159 [WS_METHODS.projectsListIgnoredFiles]: "workspace",
160 [WS_METHODS.projectsReadFile]: "workspace",
161 [WS_METHODS.projectsSearchContents]: "workspace",
162 [WS_METHODS.projectsSearchEntries]: "workspace",
163 [WS_METHODS.projectsWriteFile]: "workspace",
164 [WS_METHODS.projectsEnsureScratch]: "orchestration",
165 [WS_METHODS.projectsCreateNew]: "orchestration",
166 [WS_METHODS.projectsCreateFromThread]: "orchestration",
167 [WS_METHODS.shellOpenInEditor]: "workspace",
168 [WS_METHODS.filesystemBrowse]: "workspace",
169 [WS_METHODS.filesystemGetMetadata]: "workspace",
170 [WS_METHODS.agentSessionsScan]: "workspace",
171 [WS_METHODS.agentSessionsImport]: "workspace",
172 [WS_METHODS.assetsCreateUrl]: "workspace",
173 [WS_METHODS.assetsPersistChatAttachments]: "orchestration",
174 [WS_METHODS.attachmentsCreateUploadUrl]: "workspace",
175 [WS_METHODS.attachmentsDelete]: "workspace",
176 [WS_METHODS.providerUploadFeedback]: "provider",
177 [WS_METHODS.mcpAppsCallTool]: "provider",
178 [WS_METHODS.mcpAppsToolInfo]: "provider",
179 [WS_METHODS.mcpAppsUpdateModelContext]: "provider",
180 [WS_METHODS.mcpAppsReadResource]: "provider",
181 [WS_METHODS.subscribeVcsStatus]: "vcs",
182 [WS_METHODS.subscribeVcsBranch]: "vcs",
183 [WS_METHODS.vcsRefreshBranch]: "vcs",
184 [WS_METHODS.subscribeWorktreeSetup]: "vcs",
185 [WS_METHODS.worktreeSetupCancel]: "vcs",
186 [WS_METHODS.subscribeResourceTelemetry]: "server",
187 [WS_METHODS.vcsRefreshStatus]: "vcs",
188 [WS_METHODS.vcsPull]: "git",
189 [WS_METHODS.gitRunStackedAction]: "vcs",
190 [WS_METHODS.gitResolvePullRequest]: "git",
191 [WS_METHODS.gitPreparePullRequestThread]: "git",
192 [WS_METHODS.vcsListRefs]: "vcs",
193 [WS_METHODS.vcsCreateWorktree]: "vcs",
194 [WS_METHODS.vcsRemoveWorktree]: "vcs",
195 [WS_METHODS.vcsCreateRef]: "vcs",
196 [WS_METHODS.vcsSwitchRef]: "vcs",
197 [WS_METHODS.vcsInit]: "vcs",
198 [WS_METHODS.reviewGetBranchCommits]: "review",
199 [WS_METHODS.reviewGetDiffPreview]: "review",
200 [WS_METHODS.reviewGetDiffFileContents]: "review",
201 [WS_METHODS.terminalOpen]: "terminal",
202 [WS_METHODS.terminalAttach]: "terminal",
203 [WS_METHODS.terminalObserve]: "terminal",
204 [WS_METHODS.terminalWrite]: "terminal",
205 [WS_METHODS.terminalResize]: "terminal",
206 [WS_METHODS.terminalClear]: "terminal",
207 [WS_METHODS.terminalClose]: "terminal",
208 [WS_METHODS.subscribeTerminalMetadata]: "terminal",
209 [WS_METHODS.previewOpen]: "preview",
210 [WS_METHODS.previewResize]: "preview",
211 [WS_METHODS.previewAdjust]: "preview",
212 [WS_METHODS.previewClose]: "preview",
213 [WS_METHODS.previewList]: "preview",
214 [WS_METHODS.previewClearProfile]: "preview",
215 [WS_METHODS.previewReportProfiles]: "preview",
216 [WS_METHODS.previewReportStatus]: "preview",
217 [WS_METHODS.subscribePreviewEvents]: "preview",
218 [WS_METHODS.subscribeDiscoveredLocalServers]: "preview",
219 [WS_METHODS.previewStopLocalServer]: "preview",
220 [WS_METHODS.deviceConfigure]: "device",
221 [WS_METHODS.deviceTestHost]: "device",
222 [WS_METHODS.deviceEasAccount]: "device",
223 [WS_METHODS.deviceSimlockIdentity]: "device",
224 [WS_METHODS.deviceList]: "device",
225 [WS_METHODS.deviceOpen]: "device",
226 [WS_METHODS.deviceClose]: "device",
227 [WS_METHODS.deviceShutdown]: "device",
228 [WS_METHODS.deviceDetail]: "device",
229 [WS_METHODS.deviceAction]: "device",
230 [WS_METHODS.subscribeDeviceState]: "device",
231 [WS_METHODS.subscribeServerConfig]: "server",
232 [WS_METHODS.subscribeServerLifecycle]: "server",
233 [WS_METHODS.subscribeAuthAccess]: "auth",
234} as const satisfies Readonly<Record<WsRpcMethod, string>>;
235​
236const RPC_SPAN_PREFIX = "ws.rpc";
237const DEFAULT_RPC_SPAN_ATTRIBUTES = {
238 "rpc.transport": "websocket",
239 "rpc.system": "effect-rpc",
240} as const;
241const RPC_METHODS_WITH_TRACING_DISABLED: ReadonlySet<string> = new Set([
242 WS_METHODS.serverGetTraceDiagnostics,
243 WS_METHODS.serverGetProcessDiagnostics,
244 WS_METHODS.serverGetProcessResourceHistory,
245 WS_METHODS.serverSignalProcess,
246]);
247​
248/** Add after RpcScopeAuthorization to wrap authorization and record rejected calls. */
249export class RpcInstrumentation extends RpcMiddleware.Service<RpcInstrumentation>()(
250 "t3/server/RpcInstrumentation",
251) {}
252​
253export const rpcInstrumentationLayer = Layer.succeed(RpcInstrumentation)((effect, { rpc }) => {
254 const method = rpc._tag;
255 const measured = effect.pipe(
256 withMetrics({ counter: rpcRequestsTotal, timer: rpcRequestDuration, attributes: { method } }),
257 );
258​
259 if (RPC_METHODS_WITH_TRACING_DISABLED.has(method)) {
260 return measured.pipe(Effect.provideService(References.TracerEnabled, false));
261 }
262 return measured.pipe(
263 Effect.withSpan(`${RPC_SPAN_PREFIX}.${method}`, {
264 attributes: {
265 ...DEFAULT_RPC_SPAN_ATTRIBUTES,
266 "rpc.method": method,
267 "rpc.aggregate": RPC_AGGREGATES[method as WsRpcMethod],
268 },
269 }),
270 );
271});