1import {
2 WS_METHODS,
3 PullRequestRef,
4 PullRequestInvalidateInput,
5 PullRequestState,
6 type EnvironmentId,
7 type PullRequestRoutingResult,
8 type PullRequestRoutingIdentityResult,
9} from "@t3tools/contracts";
10import * as Effect from "effect/Effect";
11import * as Cause from "effect/Cause";
12import * as Deferred from "effect/Deferred";
13import * as Exit from "effect/Exit";
14import * as Option from "effect/Option";
15import * as Schema from "effect/Schema";
16import * as SubscriptionRef from "effect/SubscriptionRef";
17
18import * as EnvironmentRegistry from "../connection/registry.ts";
19import * as EnvironmentSupervisor from "../connection/supervisor.ts";
20import {
21 PullRequestRoutingPermissions,
22 pullRequestRoutingConnectionKey,
23} from "../connection/pullRequestRoutingPermissions.ts";
24import type { ConnectionCatalogEntry } from "../connection/catalog.ts";
25import * as ConnectionProfileStore from "../connection/profileStore.ts";
26import type { RpcSession } from "../rpc/session.ts";
27import {
28 requestGuarded as request,
29 EnvironmentRpcUnavailableError,
30 type EnvironmentRpcInput,
31 type EnvironmentUnaryRpcTag,
32} from "../rpc/client.ts";
33
34const reads = new Set<string>([
35 WS_METHODS.pullRequestsSummary,
36 WS_METHODS.pullRequestsStack,
37 WS_METHODS.pullRequestsDetail,
38 WS_METHODS.pullRequestsPreview,
39 WS_METHODS.pullRequestsChecks,
40 WS_METHODS.pullRequestsActivity,
41 WS_METHODS.pullRequestsThreadComments,
42 WS_METHODS.pullRequestsDiffFileContents,
43 WS_METHODS.pullRequestsFilesViewed,
44 WS_METHODS.pullRequestsChangeBreakdown,
45 WS_METHODS.pullRequestsReviewerCandidates,
46 WS_METHODS.pullRequestsLabelCandidates,
47]);
48// Review guides are not routed: each server keeps its own, so another server answering would
49// report none, or start a run nobody on the pull request's own server can see.
50const writes = new Set<string>([
51 WS_METHODS.pullRequestsSetFilesViewed,
52 WS_METHODS.pullRequestsRunAction,
53 WS_METHODS.pullRequestsUpdate,
54 WS_METHODS.pullRequestsComment,
55 WS_METHODS.pullRequestsUpdateComment,
56 WS_METHODS.pullRequestsSubmitReview,
57 WS_METHODS.pullRequestsReplyToThread,
58 WS_METHODS.pullRequestsSetReaction,
59 WS_METHODS.pullRequestsSetThreadResolution,
60 WS_METHODS.pullRequestsRequestReviewers,
61 WS_METHODS.pullRequestsSetLabels,
62]);
63const isRef = Schema.is(PullRequestRef);
64const hasState = Schema.is(
65 Schema.Struct({ state: PullRequestState, isDraft: Schema.optional(Schema.Boolean) }),
66);
67const isInvalidation = Schema.is(PullRequestInvalidateInput);
68const readTimeout = (environmentId: EnvironmentId) =>
69 Effect.timeoutOrElse({
70 duration: "30 seconds",
71 orElse: () =>
72 Effect.fail(
73 new EnvironmentRpcUnavailableError({
74 environmentId,
75 message: "The environment did not respond to the PR request.",
76 }),
77 ),
78 });
79interface RoutedRead {
80 origin: EnvironmentId;
81 reference: PullRequestRef;
82 targets: Set<EnvironmentId>;
83}
84const routedReads = new WeakMap<
85 EnvironmentRegistry.EnvironmentRegistry["Service"],
86 Map<string, RoutedRead>
87>();
88const isUnregistered = Schema.is(EnvironmentRegistry.EnvironmentNotRegisteredError);
89const encodeKey = Schema.encodeSync(
90 Schema.fromJsonString(Schema.Array(Schema.NullOr(Schema.String))),
91);
92type PendingMetadata<A> = WeakMap<RpcSession, Map<string, Deferred.Deferred<Exit.Exit<A>>>>;
93const pendingRouting: PendingMetadata<PullRequestRoutingResult | null> = new WeakMap();
94const pendingIdentities: PendingMetadata<PullRequestRoutingIdentityResult | null> = new WeakMap();
95
96function shareMetadataRead<A, R>(
97 pending: PendingMetadata<A>,
98 key: string,
99 probe: Effect.Effect<A, never, R>,
100): Effect.Effect<A, never, R | EnvironmentSupervisor.EnvironmentSupervisor> {
101 return Effect.uninterruptibleMask((restore) =>
102 Effect.gen(function* () {
103 const supervisor = yield* EnvironmentSupervisor.EnvironmentSupervisor;
104 const session = yield* SubscriptionRef.get(supervisor.session);
105 if (Option.isNone(session)) return yield* restore(probe);
106 const entries =
107 pending.get(session.value) ?? new Map<string, Deferred.Deferred<Exit.Exit<A>>>();
108 pending.set(session.value, entries);
109 const existing = entries.get(key);
110 if (existing !== undefined) {
111 const result = yield* restore(Deferred.await(existing));
112 // A cancelled first reader must not cancel the other readers waiting on its probe.
113 return yield* restore(
114 Exit.hasInterrupts(result) ? shareMetadataRead(pending, key, probe) : result,
115 );
116 }
117 const result = Deferred.makeUnsafe<Exit.Exit<A>>();
118 entries.set(key, result);
119 return yield* restore(probe).pipe(
120 Effect.onExit((exit) =>
121 Effect.sync(() => entries.delete(key)).pipe(
122 Effect.andThen(Deferred.succeed(result, exit)),
123 ),
124 ),
125 );
126 }),
127 );
128}
129
130function isLocal(entry: ConnectionCatalogEntry): boolean {
131 const url =
132 entry.target._tag === "PrimaryConnectionTarget"
133 ? entry.target.httpBaseUrl
134 : Option.isSome(entry.profile) && entry.profile.value._tag === "BearerConnectionProfile"
135 ? entry.profile.value.httpBaseUrl
136 : undefined;
137 if (url === undefined) return false;
138 try {
139 return ["localhost", "127.0.0.1", "[::1]"].includes(new URL(url).hostname);
140 } catch {
141 return false;
142 }
143}
144
145function rejectedBeforeDispatch(error: unknown): boolean {
146 return (
147 typeof error === "object" &&
148 error !== null &&
149 "_tag" in error &&
150 (error._tag === "EnvironmentRpcUnavailableError" ||
151 error._tag === "EnvironmentNotRegisteredError" ||
152 (error._tag === "PullRequestOperationError" &&
153 "operation" in error &&
154 error.operation === "routeIdentity"))
155 );
156}
157
158const routingAllowed = Effect.fn("PullRequestRouting.allowed")(function* (
159 registry: EnvironmentRegistry.EnvironmentRegistry["Service"],
160 originId: EnvironmentId,
161 destinationId: EnvironmentId,
162 write: boolean,
163) {
164 const entries = yield* SubscriptionRef.get(registry.entries);
165 const origin = entries.get(originId);
166 const destination = entries.get(destinationId);
167 if (
168 !origin?.enabled ||
169 !destination?.enabled ||
170 origin.unsupportedReason !== undefined ||
171 destination.unsupportedReason !== undefined
172 )
173 return false;
174 const permissions = yield* PullRequestRoutingPermissions;
175 const source = yield* permissions.get(origin);
176 const target = yield* permissions.get(destination);
177 const allowed = write
178 ? source === "read-write" && target === "read-write"
179 : source !== "off" && target !== "off";
180 if (!allowed) return false;
181 for (const entry of [origin, destination]) {
182 if (entry.target._tag !== "SshConnectionTarget") continue;
183 const profiles = yield* Effect.serviceOption(ConnectionProfileStore.ConnectionProfileStore);
184 if (Option.isNone(profiles)) return false;
185 const profile = yield* profiles.value
186 .get(entry.target.connectionId)
187 .pipe(Effect.orElseSucceed(() => Option.none()));
188 const key = pullRequestRoutingConnectionKey(entry);
189 if (key === null || key !== pullRequestRoutingConnectionKey({ ...entry, profile }))
190 return false;
191 }
192 return true;
193});
194
195function matchesReference(reference: PullRequestRef, filter: PullRequestRef): boolean {
196 return (
197 reference.projectId === filter.projectId &&
198 reference.repository.toLowerCase() === filter.repository.toLowerCase() &&
199 reference.number === filter.number &&
200 (filter.host === undefined || reference.host?.toLowerCase() === filter.host.toLowerCase())
201 );
202}
203
204const invalidateTarget = Effect.fn("PullRequestRouting.invalidateTarget")(function* (
205 registry: EnvironmentRegistry.EnvironmentRegistry["Service"],
206 origin: EnvironmentId,
207 target: EnvironmentId,
208 inputs: ReadonlyArray<PullRequestInvalidateInput>,
209) {
210 if (target !== origin && !(yield* routingAllowed(registry, origin, target, false))) return;
211 yield* Effect.forEach(
212 inputs,
213 (input) =>
214 registry.run(target, request(WS_METHODS.pullRequestsInvalidate, input)).pipe(
215 // A stalled reader must not delay the write's answer.
216 Effect.timeoutOption("1 second"),
217 Effect.orElseSucceed(() => undefined),
218 ),
219 { concurrency: 3, discard: true },
220 );
221});
222
223/**
224 * The origin owns the threads linked to this pull request, but a routed read never reaches it.
225 * Tell it which state the read saw so it checks again.
226 */
227const reportState = (
228 { allowStale: _allowStale, expectedAccountId: _expectedAccountId, ...reference }: PullRequestRef,
229 status: { readonly state: PullRequestState; readonly isDraft?: boolean | undefined },
230) =>
231 request(WS_METHODS.pullRequestsReportState, {
232 reference,
233 state: status.state,
234 ...(status.isDraft === undefined ? {} : { isDraft: status.isDraft }),
235 }).pipe(Effect.timeoutOption("5 seconds"), Effect.ignore, Effect.forkDetach, Effect.asVoid);
236
237/** Credentials stay on their environments. Only a verified host and account cross the wire. */
238export function createPullRequestRouter() {
239 const routedRequest = Effect.fn("PullRequestRouting.request")(function* <
240 T extends EnvironmentUnaryRpcTag,
241 >(tag: T, input: EnvironmentRpcInput<T>, sourceRead?: ReturnType<typeof request<T>>) {
242 const source = sourceRead ?? request(tag, input);
243 if (tag === WS_METHODS.pullRequestsInvalidate && isInvalidation(input)) {
244 const result = yield* request(tag, input);
245 const origin = yield* EnvironmentSupervisor.EnvironmentSupervisor;
246 const registry = yield* EnvironmentRegistry.EnvironmentRegistry;
247 const used = routedReads.get(registry);
248 const targets = new Map<EnvironmentId, PullRequestRef>();
249 for (const entry of used?.values() ?? []) {
250 if (entry.origin !== origin.target.environmentId) continue;
251 const ref = input.reference;
252 if (ref !== undefined && !matchesReference(entry.reference, ref)) continue;
253 for (const target of entry.targets) targets.set(target, entry.reference);
254 if (input.reference !== undefined)
255 targets.set(origin.target.environmentId, entry.reference);
256 }
257 yield* Effect.forEach(
258 targets,
259 ([target, reference]) =>
260 invalidateTarget(registry, origin.target.environmentId, target, [
261 { ...input, ...(input.reference === undefined ? {} : { reference }) },
262 ]),
263 { concurrency: 4, discard: true },
264 );
265 return result;
266 }
267 if ((!reads.has(tag) && !writes.has(tag)) || !isRef(input)) {
268 return yield* request(tag, input);
269 }
270 const ref = input;
271 const origin = yield* EnvironmentSupervisor.EnvironmentSupervisor;
272 const registry = yield* EnvironmentRegistry.EnvironmentRegistry;
273 const entries = yield* SubscriptionRef.get(registry.entries);
274 const sourceEntry = entries.get(origin.target.environmentId);
275 const used = routedReads.get(registry) ?? new Map<string, RoutedRead>();
276 routedReads.set(registry, used);
277 const refKey = encodeKey([
278 origin.target.environmentId,
279 ref.projectId,
280 ref.host?.toLowerCase() ?? null,
281 ref.repository.toLowerCase(),
282 String(ref.number),
283 ]);
284 const invalidateReaders = Effect.suspend(() => {
285 if (!writes.has(tag) || used.size === 0) return Effect.void;
286 const targets = new Map<EnvironmentId, PullRequestRef[]>([
287 [origin.target.environmentId, [ref]],
288 ]);
289 for (const entry of used.values()) {
290 if (entry.origin !== origin.target.environmentId || !matchesReference(entry.reference, ref))
291 continue;
292 for (const target of [...entry.targets, origin.target.environmentId]) {
293 const refs = targets.get(target) ?? [];
294 if (!refs.some((existing) => existing.host === entry.reference.host))
295 refs.push(entry.reference);
296 targets.set(target, refs);
297 }
298 }
299 return Effect.forEach(
300 targets,
301 ([target, refs]) =>
302 invalidateTarget(registry, origin.target.environmentId, target, [
303 ...refs.map((reference) => ({
304 reference,
305 ...(tag === WS_METHODS.pullRequestsSetFilesViewed ? { filesViewedOnly: true } : {}),
306 })),
307 ...(tag === WS_METHODS.pullRequestsSetFilesViewed ? [] : [{}]),
308 ]),
309 { concurrency: 4, discard: true },
310 );
311 });
312 // A failed viewed-files batch may have applied some paths. Other writes invalidate on success.
313 const finish = (operation: ReturnType<typeof request<T>>) =>
314 tag === WS_METHODS.pullRequestsSetFilesViewed
315 ? operation.pipe(Effect.ensuring(invalidateReaders))
316 : operation.pipe(Effect.tap(() => invalidateReaders));
317 const alternatives = [];
318 for (const [id, entry] of entries) {
319 if (id === origin.target.environmentId) continue;
320 if (!(yield* routingAllowed(registry, origin.target.environmentId, id, writes.has(tag))))
321 continue;
322 const connected = yield* registry
323 .run(
324 id,
325 EnvironmentSupervisor.EnvironmentSupervisor.pipe(
326 Effect.flatMap((s) => SubscriptionRef.get(s.session)),
327 ),
328 )
329 .pipe(Effect.orElseSucceed(() => Option.none()));
330 if (Option.isNone(connected)) continue;
331 if (tag === WS_METHODS.pullRequestsChecks) {
332 const supported = yield* connected.value.initialConfig.pipe(
333 Effect.map((config) => config.environment.capabilities.pullRequestChecks === true),
334 Effect.timeout("2 seconds"),
335 Effect.catchCause((cause) =>
336 Cause.hasInterrupts(cause) ? Effect.interrupt : Effect.succeed(false),
337 ),
338 );
339 if (!supported) continue;
340 }
341 alternatives.push({ id, local: isLocal(entry) });
342 }
343 if (alternatives.length === 0) return yield* finish(source);
344
345 const identityProbe = request(WS_METHODS.pullRequestsRouting, ref).pipe(
346 Effect.timeout("2 seconds"),
347 Effect.catchCause((cause) =>
348 Cause.hasInterrupts(cause) ? Effect.interrupt : Effect.succeed(null),
349 ),
350 );
351 const identity = yield* reads.has(tag)
352 ? shareMetadataRead(
353 pendingRouting,
354 encodeKey([ref.projectId, ref.host?.toLowerCase() ?? null, ref.repository.toLowerCase()]),
355 identityProbe,
356 )
357 : identityProbe;
358 // The source server answers only for a host whose account it can verify.
359 if (identity === null) return yield* finish(source);
360
361 const routedInput = {
362 ...input,
363 host: identity.host,
364 expectedAccountId: identity.accountId,
365 };
366 const guardedSource = sourceRead ?? (yield* Effect.cached(request(tag, routedInput)));
367 const local = alternatives.filter((entry) => entry.local);
368 const remote = alternatives.filter((entry) => !entry.local);
369 const sourceLocal = sourceEntry !== undefined && isLocal(sourceEntry);
370 const candidates = [
371 ...(sourceLocal && writes.has(tag) ? [origin.target.environmentId] : []),
372 ...local.map((entry) => entry.id),
373 ...(!sourceLocal && writes.has(tag) ? [origin.target.environmentId] : []),
374 ...remote.map((entry) => entry.id),
375 ...(reads.has(tag) ? [origin.target.environmentId] : []),
376 ];
377 const run = (id: EnvironmentId): ReturnType<typeof request<T>> =>
378 (id === origin.target.environmentId
379 ? guardedSource
380 : registry.run(id, request(tag, routedInput))
381 ).pipe(
382 Effect.tap((result) =>
383 id !== origin.target.environmentId && hasState(result)
384 ? reportState(ref, result)
385 : Effect.void,
386 ),
387 Effect.tap(() =>
388 Effect.sync(() => {
389 const entry = used.get(refKey) ?? {
390 origin: origin.target.environmentId,
391 reference: { ...ref, host: identity.host },
392 targets: new Set<EnvironmentId>(),
393 };
394 entry.targets.add(id);
395 if (used.size >= 256 && !used.has(refKey)) {
396 const oldest = used.keys().next().value;
397 if (oldest !== undefined) used.delete(oldest);
398 }
399 used.set(refKey, entry);
400 }),
401 ),
402 Effect.mapError((error) =>
403 isUnregistered(error)
404 ? new EnvironmentRpcUnavailableError({
405 environmentId: id,
406 message: "The environment was removed.",
407 })
408 : error,
409 ),
410 );
411 const visit = (index: number): ReturnType<typeof request<T>> => {
412 const id = candidates[index];
413 if (id === undefined) return guardedSource;
414 return Effect.gen(function* () {
415 if (id !== origin.target.environmentId) {
416 if (!(yield* routingAllowed(registry, origin.target.environmentId, id, writes.has(tag))))
417 return yield* visit(index + 1);
418 // An older server would discard expectedAccountId. Verify it implements the guard first.
419 const alternateProbe = request(WS_METHODS.pullRequestsRoutingIdentity, {
420 host: identity.host,
421 }).pipe(
422 Effect.timeout("2 seconds"),
423 Effect.catchCause((cause) =>
424 Cause.hasInterrupts(cause) ? Effect.interrupt : Effect.succeed(null),
425 ),
426 );
427 const alternate = yield* registry
428 .run(
429 id,
430 reads.has(tag)
431 ? shareMetadataRead(pendingIdentities, identity.host.toLowerCase(), alternateProbe)
432 : alternateProbe,
433 )
434 .pipe(
435 Effect.timeout("2 seconds"),
436 Effect.catchCause((cause) =>
437 Cause.hasInterrupts(cause) ? Effect.interrupt : Effect.succeed(null),
438 ),
439 );
440 if (
441 alternate === null ||
442 alternate.provider !== identity.provider ||
443 alternate.host.toLowerCase() !== identity.host.toLowerCase() ||
444 alternate.accountId !== identity.accountId
445 ) {
446 return yield* visit(index + 1);
447 }
448 if (!(yield* routingAllowed(registry, origin.target.environmentId, id, writes.has(tag))))
449 return yield* visit(index + 1);
450 }
451 const operation = run(id);
452 return yield* (reads.has(tag) ? operation.pipe(readTimeout(id)) : operation).pipe(
453 Effect.catchIf(
454 (error) =>
455 (reads.has(tag) || rejectedBeforeDispatch(error)) && index + 1 < candidates.length,
456 () => visit(index + 1),
457 ),
458 Effect.map((result) =>
459 typeof result === "object" && result !== null && "projectId" in result
460 ? {
461 ...result,
462 projectId: ref.projectId,
463 ...(tag === WS_METHODS.pullRequestsDetail
464 ? { projectTitle: identity.projectTitle, workspaceRoot: identity.workspaceRoot }
465 : {}),
466 }
467 : result,
468 ),
469 );
470 });
471 };
472 return yield* finish(visit(0));
473 });
474
475 return Effect.fn("PullRequestRouting.readOrWrite")(function* <T extends EnvironmentUnaryRpcTag>(
476 tag: T,
477 input: EnvironmentRpcInput<T>,
478 ) {
479 if (!reads.has(tag) || !isRef(input)) return yield* routedRequest(tag, input);
480 const registry = yield* EnvironmentRegistry.EnvironmentRegistry;
481 const entries = yield* SubscriptionRef.get(registry.entries);
482 if (entries.size < 2) return yield* request(tag, input);
483 const origin = yield* EnvironmentSupervisor.EnvironmentSupervisor;
484 let allowed = false;
485 for (const id of entries.keys()) {
486 if (
487 id !== origin.target.environmentId &&
488 (yield* routingAllowed(registry, origin.target.environmentId, id, false))
489 ) {
490 allowed = true;
491 break;
492 }
493 }
494 if (!allowed) return yield* request(tag, input);
495 const strictInput = { ...input, allowStale: false };
496 const source = yield* Effect.cached(
497 request(tag, strictInput).pipe(readTimeout(origin.target.environmentId)),
498 );
499 const sourceEntry = entries.get(origin.target.environmentId);
500 const routed = routedRequest(tag, strictInput, source);
501 return yield* (
502 sourceEntry !== undefined && isLocal(sourceEntry)
503 ? source.pipe(Effect.catch(() => routed))
504 : routed
505 ).pipe(
506 Effect.catch((error) =>
507 input.allowStale !== false &&
508 (tag === WS_METHODS.pullRequestsSummary || tag === WS_METHODS.pullRequestsDetail)
509 ? request(tag, input)
510 : Effect.fail(error),
511 ),
512 );
513 });
514}