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