1import {
2 DEFAULT_PROVIDER_HEALTH_REFRESH_INTERVAL,
3 type ServerProvider,
4 ServerSettingsError,
5} from "@t3tools/contracts";
6import { resolveServerBackgroundActivitySettings } from "@t3tools/shared/backgroundActivitySettings";
7import * as Duration from "effect/Duration";
8import * as Effect from "effect/Effect";
9import * as Equal from "effect/Equal";
10import * as Fiber from "effect/Fiber";
11import * as PubSub from "effect/PubSub";
12import * as Queue from "effect/Queue";
13import * as Ref from "effect/Ref";
14import * as Scope from "effect/Scope";
15import * as Stream from "effect/Stream";
16import * as Semaphore from "effect/Semaphore";
17
18import * as ProviderHost from "./ProviderHost.ts";
19import {
20 applyUsageLimitsUpdate,
21 makeUsageLimits,
22 resolveUsageLimitsAfterProbe,
23} from "./usageLimits.ts";
24import type { ManagedServerProvider } from "./snapshot.ts";
25
26interface ProviderSnapshotState {
27 readonly snapshot: ServerProvider;
28 readonly enrichmentGeneration: number;
29 // null preserves an explicitly cleared plan across probes and enrichment.
30 readonly usageAuthLabel: string | null | undefined;
31 readonly checked: boolean;
32}
33
34function withUsageLimits(
35 snapshot: ServerProvider,
36 usageLimits: ServerProvider["usageLimits"],
37): ServerProvider {
38 if (snapshot.usageLimits === usageLimits) {
39 return snapshot;
40 }
41 const { usageLimits: _previous, ...rest } = snapshot;
42 return usageLimits ? { ...rest, usageLimits } : rest;
43}
44
45const FAILED_CHECK_RECHECK_DELAY = "30 seconds";
46
47export const makeManagedServerProvider = Effect.fn("makeManagedServerProvider")(function* <
48 Settings,
49>(input: {
50 readonly resolveMaintenance: ManagedServerProvider["resolveMaintenance"];
51 readonly getSettings: Effect.Effect<Settings, ServerSettingsError>;
52 readonly streamSettings: Stream.Stream<Settings>;
53 readonly haveSettingsChanged: (previous: Settings, next: Settings) => boolean;
54 readonly initialSnapshot: (settings: Settings) => Effect.Effect<ServerProvider>;
55 readonly checkProvider: Effect.Effect<ServerProvider, ServerSettingsError>;
56 readonly checkUsageLimits?: Effect.Effect<ServerProvider["usageLimits"]>;
57 readonly invalidateUsageCache?: Effect.Effect<void>;
58 readonly enrichSnapshot?: (input: {
59 readonly settings: Settings;
60 readonly snapshot: ServerProvider;
61 readonly getSnapshot: Effect.Effect<ServerProvider>;
62 readonly publishSnapshot: (snapshot: ServerProvider) => Effect.Effect<void>;
63 }) => Effect.Effect<void>;
64 readonly refreshInterval?: Duration.Input;
65 readonly refreshOnInterval?: boolean | ((snapshot: ServerProvider) => boolean);
66 readonly checkProviderOnSettingsChange?: (previous: Settings, next: Settings) => boolean;
67}): Effect.fn.Return<
68 ManagedServerProvider,
69 ServerSettingsError,
70 Scope.Scope | ProviderHost.ProviderHost
71> {
72 const host = yield* ProviderHost.ProviderHost;
73 const refreshSemaphore = yield* Semaphore.make(1);
74 const usageRefreshQueued = yield* Ref.make<symbol | null>(null);
75 const changesPubSub = yield* Effect.acquireRelease(
76 PubSub.unbounded<ServerProvider>(),
77 PubSub.shutdown,
78 );
79 const initialSettings = yield* input.getSettings;
80 const initialSnapshot = yield* input.initialSnapshot(initialSettings);
81 const snapshotStateRef = yield* Ref.make<ProviderSnapshotState>({
82 snapshot: initialSnapshot,
83 enrichmentGeneration: 0,
84 usageAuthLabel: undefined,
85 checked: false,
86 });
87 const settingsRef = yield* Ref.make(initialSettings);
88 const recheckPendingRef = yield* Ref.make(false);
89 const enrichmentFiberRef = yield* Ref.make<Fiber.Fiber<void, unknown> | null>(null);
90 const scope = yield* Effect.scope;
91
92 const publishEnrichedSnapshot = Effect.fn("publishEnrichedSnapshot")(function* (
93 generation: number,
94 nextSnapshot: ServerProvider,
95 ) {
96 const snapshotToPublish = yield* Ref.modify(snapshotStateRef, (state) => {
97 if (state.enrichmentGeneration !== generation) {
98 return [null, state] as const;
99 }
100 // Enrichment must not revert usage updates newer than its input snapshot.
101 const merged = withUsageLimits(
102 state.usageAuthLabel === undefined
103 ? nextSnapshot
104 : {
105 ...nextSnapshot,
106 auth: { ...nextSnapshot.auth, label: state.usageAuthLabel ?? undefined },
107 },
108 state.snapshot.usageLimits,
109 );
110 if (Equal.equals(state.snapshot, merged)) {
111 return [null, state] as const;
112 }
113 return [merged, { ...state, snapshot: merged }] as const;
114 });
115 if (snapshotToPublish === null) {
116 return;
117 }
118 yield* PubSub.publish(changesPubSub, snapshotToPublish);
119 });
120
121 const restartSnapshotEnrichment = Effect.fn("restartSnapshotEnrichment")(function* (
122 settings: Settings,
123 snapshot: ServerProvider,
124 generation: number,
125 ) {
126 const previousFiber = yield* Ref.getAndSet(enrichmentFiberRef, null);
127 if (previousFiber) {
128 yield* Fiber.interrupt(previousFiber).pipe(Effect.ignore);
129 }
130
131 if (!input.enrichSnapshot || !snapshot.enabled) {
132 return;
133 }
134
135 const fiber = yield* input
136 .enrichSnapshot({
137 settings,
138 snapshot,
139 getSnapshot: Ref.get(snapshotStateRef).pipe(Effect.map((state) => state.snapshot)),
140 publishSnapshot: (nextSnapshot) => publishEnrichedSnapshot(generation, nextSnapshot),
141 })
142 .pipe(Effect.ignoreCause({ log: true }), Effect.forkIn(scope));
143
144 yield* Ref.set(enrichmentFiberRef, fiber);
145 });
146
147 const applySnapshotBase = Effect.fn("applySnapshot")(function* (
148 nextSettings: Settings,
149 options?: { readonly forceRefresh?: boolean; readonly deferFailure?: boolean },
150 ) {
151 const previousSettings = yield* Ref.get(settingsRef);
152 const settingsChanged =
153 options?.forceRefresh === true || input.haveSettingsChanged(previousSettings, nextSettings);
154 // Nothing watches credential files, so any settings apply rechecks a failed sign-in.
155 const { snapshot: currentSnapshot } = yield* Ref.get(snapshotStateRef);
156 const forceRefresh =
157 options?.forceRefresh === true ||
158 (currentSnapshot.enabled && currentSnapshot.auth.status === "unauthenticated");
159 if (!forceRefresh && !settingsChanged) {
160 yield* Ref.set(settingsRef, nextSettings);
161 return yield* Ref.get(snapshotStateRef).pipe(Effect.map((state) => state.snapshot));
162 }
163
164 if (
165 !forceRefresh &&
166 input.checkProviderOnSettingsChange?.(previousSettings, nextSettings) === false
167 ) {
168 const state = yield* Ref.get(snapshotStateRef);
169 const nextGeneration = state.enrichmentGeneration + 1;
170 yield* Ref.set(snapshotStateRef, {
171 ...state,
172 enrichmentGeneration: nextGeneration,
173 });
174 yield* Ref.set(settingsRef, nextSettings);
175 yield* restartSnapshotEnrichment(nextSettings, state.snapshot, nextGeneration);
176 return state.snapshot;
177 }
178
179 // A probe without usage must not erase a report published while its I/O was in flight.
180 const usageBeforeProbe = (yield* Ref.get(snapshotStateRef)).snapshot.usageLimits;
181 const probedSnapshot = yield* input.checkProvider;
182 // One failed periodic check over a ready provider is often transient; recheck before showing it.
183 if (
184 options?.deferFailure === true &&
185 probedSnapshot.status === "error" &&
186 probedSnapshot.auth.status !== "unauthenticated"
187 ) {
188 const { snapshot } = yield* Ref.get(snapshotStateRef);
189 if (snapshot.enabled && snapshot.status === "ready") {
190 yield* Ref.set(recheckPendingRef, true);
191 return snapshot;
192 }
193 }
194 const { snapshot: nextSnapshot, generation: nextGeneration } = yield* Ref.modify(
195 snapshotStateRef,
196 (state) => {
197 const generation = input.enrichSnapshot
198 ? state.enrichmentGeneration + 1
199 : state.enrichmentGeneration;
200 // Failed reads retain limits only for the same account and original workspace.
201 const previous = state.snapshot.auth;
202 const differs = (a: string | undefined, b: string | undefined) =>
203 a !== undefined && b !== undefined && a !== b;
204 const switchedAccount =
205 differs(probedSnapshot.auth.email?.toLowerCase(), previous.email?.toLowerCase()) ||
206 (probedSnapshot.auth.workspaceId !== undefined &&
207 probedSnapshot.auth.workspaceId !== previous.workspaceId);
208 const usageLimits = switchedAccount
209 ? probedSnapshot.usageLimits
210 : probedSnapshot.usageLimits === undefined &&
211 state.snapshot.usageLimits !== usageBeforeProbe
212 ? state.snapshot.usageLimits
213 : resolveUsageLimitsAfterProbe({
214 published: state.snapshot.usageLimits,
215 probed: probedSnapshot.usageLimits,
216 });
217 const workspaceId =
218 usageLimits === probedSnapshot.usageLimits ? undefined : state.snapshot.auth.workspaceId;
219 const usageAuthLabel =
220 !switchedAccount && usageLimits?.checkedAt === state.snapshot.usageLimits?.checkedAt
221 ? state.usageAuthLabel
222 : undefined;
223 const labeledProbe =
224 usageAuthLabel === undefined
225 ? probedSnapshot
226 : {
227 ...probedSnapshot,
228 auth: { ...probedSnapshot.auth, label: usageAuthLabel ?? undefined },
229 };
230 const snapshot = withUsageLimits(
231 workspaceId && !probedSnapshot.auth.workspaceId
232 ? { ...labeledProbe, auth: { ...labeledProbe.auth, workspaceId } }
233 : labeledProbe,
234 usageLimits,
235 );
236 return [
237 { snapshot, generation },
238 { snapshot, enrichmentGeneration: generation, usageAuthLabel, checked: true },
239 ] as const;
240 },
241 );
242 yield* Ref.set(settingsRef, nextSettings);
243 yield* PubSub.publish(changesPubSub, nextSnapshot);
244 yield* restartSnapshotEnrichment(nextSettings, nextSnapshot, nextGeneration);
245 return nextSnapshot;
246 });
247 const applySnapshot = (nextSettings: Settings, options?: { readonly forceRefresh?: boolean }) =>
248 refreshSemaphore.withPermits(1)(applySnapshotBase(nextSettings, options));
249
250 // Providers call this on every token-usage tick, usually without changing the snapshot.
251 const applyUsageLimits: ManagedServerProvider["applyUsageLimits"] = Effect.fnUntraced(
252 function* (update) {
253 const snapshotToPublish = yield* Ref.modify(snapshotStateRef, (state) => {
254 const complete = update.complete && update.windows.length > 0;
255 const usageLimits = complete
256 ? makeUsageLimits(update)
257 : applyUsageLimitsUpdate({
258 previous: state.snapshot.usageLimits,
259 update,
260 checkedAt: update.checkedAt,
261 });
262 const { auth } = state.snapshot;
263 const label = update.windows.length > 0 ? update.authLabel : undefined;
264 const relabel =
265 (label !== undefined || (complete && state.usageAuthLabel !== undefined)) &&
266 label !== auth.label;
267 if (usageLimits === state.snapshot.usageLimits && !relabel) {
268 return [null, state] as const;
269 }
270 const snapshot = withUsageLimits(
271 relabel ? { ...state.snapshot, auth: { ...auth, label } } : state.snapshot,
272 usageLimits,
273 );
274 return [
275 snapshot,
276 {
277 ...state,
278 snapshot,
279 usageAuthLabel: label ?? (complete ? null : state.usageAuthLabel),
280 },
281 ] as const;
282 });
283 if (snapshotToPublish !== null) {
284 yield* PubSub.publish(changesPubSub, snapshotToPublish);
285 }
286 },
287 );
288
289 const refreshSnapshot = Effect.fn("refreshSnapshot")(function* () {
290 const nextSettings = yield* input.getSettings;
291 return yield* applySnapshot(nextSettings, { forceRefresh: true });
292 });
293
294 const refreshUsage = Effect.fn("ManagedServerProvider.refreshUsage")(function* () {
295 const { snapshot } = yield* Ref.get(snapshotStateRef);
296 if (
297 !snapshot.enabled ||
298 !snapshot.usageLimits ||
299 snapshot.usageLimits.unavailable?.reason === "unsupported"
300 )
301 return;
302 // An in-flight refresh may have read cached usage, so one refresh waits behind it.
303 const claim = Symbol();
304 const release = Ref.update(usageRefreshQueued, (queued) => (queued === claim ? null : queued));
305 yield* Effect.uninterruptibleMask((restore) =>
306 Ref.modify(usageRefreshQueued, (queued) => [queued !== null, queued ?? claim]).pipe(
307 Effect.flatMap((held) =>
308 held
309 ? Effect.void
310 : restore(
311 refreshSemaphore.withPermits(1)(
312 release.pipe(
313 Effect.andThen(input.invalidateUsageCache ?? Effect.void),
314 Effect.andThen(input.getSettings),
315 Effect.flatMap((settings) =>
316 applySnapshotBase(settings, { forceRefresh: true }),
317 ),
318 ),
319 ),
320 ).pipe(Effect.onInterrupt(() => release)),
321 ),
322 ),
323 );
324 });
325
326 const needsUsageRecovery = (snapshot: ServerProvider) =>
327 input.checkUsageLimits !== undefined &&
328 snapshot.installed &&
329 snapshot.status === "ready" &&
330 snapshot.auth.status === "authenticated" &&
331 snapshot.usageLimits?.unavailable?.reason === "probeFailed";
332
333 const needsProviderRefresh = (snapshot: ServerProvider) =>
334 input.refreshOnInterval !== false &&
335 (typeof input.refreshOnInterval !== "function" || input.refreshOnInterval(snapshot));
336
337 const refreshOnInterval = Effect.fn("ManagedServerProvider.refreshOnInterval")(function* (
338 deferFailure: boolean,
339 ) {
340 const { snapshot } = yield* Ref.get(snapshotStateRef);
341 if (!snapshot.enabled) return;
342 if (!needsUsageRecovery(snapshot) || !input.checkUsageLimits) {
343 if (needsProviderRefresh(snapshot)) {
344 yield* applySnapshotBase(yield* input.getSettings, { forceRefresh: true, deferFailure });
345 }
346 return;
347 }
348 const probed = yield* input.checkUsageLimits;
349 const next = yield* Ref.modify(snapshotStateRef, (state) => {
350 const usageLimits = resolveUsageLimitsAfterProbe({
351 published: state.snapshot.usageLimits,
352 probed,
353 });
354 const snapshot = withUsageLimits(state.snapshot, usageLimits);
355 return [snapshot, { ...state, snapshot }] as const;
356 });
357 yield* PubSub.publish(changesPubSub, next);
358 });
359
360 const shouldRefreshOnInterval = Effect.gen(function* () {
361 const state = yield* Ref.get(snapshotStateRef);
362 // A disabled instance has nothing to probe; a periodic check would only
363 // republish a snapshot that differs in checkedAt.
364 if (
365 !state.snapshot.enabled ||
366 (!needsUsageRecovery(state.snapshot) && !needsProviderRefresh(state.snapshot))
367 ) {
368 return false;
369 }
370 const instanceId = state.snapshot.instanceId;
371 const [genericDemand, instanceDemand] = yield* Effect.all([
372 host.shouldRunBackgroundWork({ type: "provider-status" }),
373 host.shouldRunBackgroundWork({ type: "provider-status", instanceId }),
374 ]);
375 return genericDemand || instanceDemand;
376 });
377
378 const getRefreshInterval =
379 input.refreshInterval !== undefined
380 ? Effect.succeed(input.refreshInterval)
381 : host.settings.get.pipe(
382 Effect.map(
383 (settings) =>
384 resolveServerBackgroundActivitySettings(settings).providerHealthRefreshInterval,
385 ),
386 Effect.orElseSucceed(() => DEFAULT_PROVIDER_HEALTH_REFRESH_INTERVAL),
387 );
388
389 const refreshIntervalChanges = yield* Queue.sliding<void>(1);
390 if (input.refreshInterval === undefined) {
391 const serverSettingsChanges = yield* host.settings.subscribe;
392 yield* serverSettingsChanges.pipe(
393 Stream.map((settings) =>
394 Duration.toMillis(
395 resolveServerBackgroundActivitySettings(settings).providerHealthRefreshInterval,
396 ),
397 ),
398 Stream.changes,
399 Stream.runForEach(() => Queue.offer(refreshIntervalChanges, undefined).pipe(Effect.asVoid)),
400 Effect.forkScoped,
401 );
402 }
403
404 // A failed check must not end this subscription: the next settings change
405 // retries, like the periodic and initial checks below.
406 yield* Stream.runForEach(input.streamSettings, (nextSettings) =>
407 Effect.asVoid(applySnapshot(nextSettings)).pipe(Effect.ignoreCause({ log: true })),
408 ).pipe(Effect.forkScoped);
409
410 yield* Effect.forever(
411 Effect.all([getRefreshInterval, Ref.getAndSet(recheckPendingRef, false)]).pipe(
412 Effect.flatMap(([configuredInterval, recheckPending]) => {
413 const refreshInterval = recheckPending ? FAILED_CHECK_RECHECK_DELAY : configuredInterval;
414 // @effect-diagnostics-next-line raceFirstWithSleepToTimeout:off - races the interval against a settings-change signal, not a timeout
415 return Effect.raceFirst(
416 Effect.sleep(
417 Duration.toMillis(Duration.fromInputUnsafe(refreshInterval)) <= 0
418 ? "60 seconds"
419 : refreshInterval,
420 ).pipe(Effect.as(true)),
421 Queue.take(refreshIntervalChanges).pipe(Effect.as(false)),
422 ).pipe(
423 Effect.flatMap((intervalElapsed) =>
424 intervalElapsed && Duration.toMillis(Duration.fromInputUnsafe(refreshInterval)) > 0
425 ? shouldRefreshOnInterval.pipe(
426 Effect.flatMap((shouldRefresh) =>
427 shouldRefresh
428 ? refreshSemaphore.withPermits(1)(refreshOnInterval(!recheckPending))
429 : Effect.void,
430 ),
431 )
432 : Effect.void,
433 ),
434 );
435 }),
436 Effect.ignoreCause({ log: true }),
437 ),
438 ).pipe(Effect.forkScoped);
439
440 yield* applySnapshot(initialSettings, { forceRefresh: true }).pipe(
441 Effect.ignoreCause({ log: true }),
442 Effect.forkScoped,
443 );
444
445 return {
446 resolveMaintenance: input.resolveMaintenance,
447 getSnapshot: Ref.get(snapshotStateRef).pipe(Effect.map((state) => state.snapshot)),
448 isAwaitingFirstCheck: Ref.get(snapshotStateRef).pipe(Effect.map((state) => !state.checked)),
449 refresh: input
450 .resolveMaintenance({ fresh: true })
451 .pipe(Effect.andThen(refreshSnapshot), Effect.tapError(Effect.logError), Effect.orDie),
452 applyUsageLimits,
453 refreshUsage: refreshUsage(),
454 get streamChanges() {
455 return Stream.fromPubSub(changesPubSub);
456 },
457 } satisfies ManagedServerProvider;
458});