1import { CommandId, MessageId, type OrchestrationV2Command } from "@t3tools/contracts";
2import * as DateTime from "effect/DateTime";
3import * as Effect from "effect/Effect";
4import * as Layer from "effect/Layer";
5import * as Stream from "effect/Stream";
6import type { OrchestrationV2DomainEvent } from "@t3tools/contracts";
7import * as DueHint from "../scheduling/DueHint.ts";
8import * as Scheduler from "../scheduling/Scheduler.ts";
9import { forkParked } from "../serverActivation.ts";
10import * as ServerSettings from "../serverSettings.ts";
11import * as ProjectionStore from "./ProjectionStore.ts";
12import * as Orchestrator from "./Orchestrator.ts";
13import * as ThreadManagement from "./ThreadManagementService.ts";
14
15export const MAX_OVERDUE_ARM_MS = 24 * 60 * 60 * 1000;
16
17export function limitRecoveryCommand(
18 thread: ProjectionStore.ProjectionLimitRecoveryCandidate,
19 autoResume: boolean,
20 nowMs: number,
21 snooze = false,
22): OrchestrationV2Command | null {
23 if (
24 thread.status !== "failed" ||
25 thread.lastErrorClass !== "usage_limit" ||
26 !thread.latestRunId ||
27 !thread.usageLimitResetAt ||
28 thread.archivedAt !== null ||
29 thread.settledOverride === "settled" ||
30 thread.pendingRuntimeRequest !== null
31 )
32 return null;
33 const resetMs = Date.parse(thread.usageLimitResetAt);
34 // An already-expired window reported with a fresh failure cannot start a retry loop.
35 if (
36 !Number.isFinite(resetMs) ||
37 resetMs <= DateTime.toEpochMillis(thread.latestRunCompletedAt ?? thread.updatedAt)
38 )
39 return null;
40 const identity = `${thread.id}:${thread.latestRunId}:${resetMs}`;
41 const recovery = thread.limitRecovery;
42 if (recovery?.runId !== thread.latestRunId || recovery.resetAt !== thread.usageLimitResetAt) {
43 if (nowMs - resetMs > MAX_OVERDUE_ARM_MS) return null;
44 if (!autoResume && (!snooze || resetMs <= nowMs)) return null;
45 return {
46 type: "thread.metadata.update",
47 commandId: CommandId.make(`limit-arm:${identity}`),
48 threadId: thread.id,
49 limitRecovery: {
50 runId: thread.latestRunId,
51 resetAt: thread.usageLimitResetAt,
52 autoResume,
53 snooze: snooze && resetMs > nowMs,
54 },
55 };
56 }
57 if (
58 !recovery.autoResume ||
59 resetMs > nowMs ||
60 (thread.snoozedUntil != null && DateTime.toEpochMillis(thread.snoozedUntil) > nowMs)
61 )
62 return null;
63 const deliveryIdentity = `${identity}:${recovery.requestId ?? "legacy"}`;
64 return {
65 type: "message.dispatch",
66 commandId: CommandId.make(`limit-resume:${deliveryIdentity}`),
67 messageId: MessageId.make(`limit-resume:${deliveryIdentity}`),
68 threadId: thread.id,
69 usageLimitContinuationOfRunId: thread.latestRunId,
70 ...(recovery.requestId === undefined
71 ? {}
72 : { usageLimitRecoveryRequestId: recovery.requestId }),
73 text: "Continue where you left off.",
74 attachments: [],
75 dispatchMode: { type: "start_immediately" },
76 createdBy: "agent",
77 creationSource: "server",
78 };
79}
80
81// Events that can make a thread a candidate or change its command; streaming items cannot.
82const affectsLimitRecovery = (event: OrchestrationV2DomainEvent) =>
83 event.type.startsWith("thread.") ||
84 event.type === "run.created" ||
85 event.type === "run.updated" ||
86 event.type === "runtime-request.updated" ||
87 event.type.startsWith("provider-session.") ||
88 (event.type === "turn-item.updated" && event.payload.type === "error");
89
90const makeSweep = Effect.gen(function* () {
91 const projections = yield* ProjectionStore.ProjectionStoreV2;
92 const threads = yield* ThreadManagement.ThreadManagementService;
93 const orchestrator = yield* Orchestrator.OrchestratorV2;
94 const settings = yield* ServerSettings.ServerSettingsService;
95 const due = DueHint.make();
96 yield* forkParked(
97 Stream.runForEach(orchestrator.streamDomainEvents, (event) =>
98 affectsLimitRecovery(event) ? due.invalidate : Effect.void,
99 ),
100 );
101 yield* forkParked(Stream.runForEach(settings.streamChanges, () => due.invalidate));
102 return Effect.fn("UsageLimitRecoveryWorker.sweep")(function* () {
103 const now = yield* DateTime.now;
104 const nowMs = DateTime.toEpochMillis(now);
105 if (yield* due.isIdle(nowMs)) return;
106 const readVersion = yield* due.version;
107 const preferences = yield* settings.getSettings;
108 const candidates = yield* projections.getLimitRecoveryCandidates({
109 now,
110 autoResume: preferences.autoResumeLimitedThreads,
111 snooze: preferences.snoozeLimitedThreads,
112 });
113 let acted = false;
114 for (const thread of candidates) {
115 const command = limitRecoveryCommand(
116 thread,
117 preferences.autoResumeLimitedThreads,
118 nowMs,
119 preferences.snoozeLimitedThreads,
120 );
121 if (command === null) continue;
122 acted = true;
123 yield* Effect.gen(function* () {
124 if (command.type === "thread.metadata.update" && command.limitRecovery != null) {
125 yield* orchestrator.autoArmLimitRecovery({
126 commandId: command.commandId,
127 threadId: command.threadId,
128 limitRecovery: command.limitRecovery,
129 });
130 } else {
131 yield* threads.dispatch(command);
132 }
133 }).pipe(
134 Effect.catchCause((cause) =>
135 Effect.logWarning("orchestration-v2.limit-recovery.dispatch-failed", {
136 threadId: thread.id,
137 cause,
138 }),
139 ),
140 );
141 }
142 // A sweep that acted, or failed to, re-checks on the next tick.
143 if (acted) return;
144 const wakeAt = yield* projections.getLimitRecoveryWakeAt(now);
145 yield* due.record(readVersion, wakeAt ?? Infinity, nowMs);
146 });
147});
148
149export const layer = Layer.effectDiscard(
150 Effect.gen(function* () {
151 const sweep = yield* makeSweep;
152 const scheduler = yield* Scheduler.Scheduler;
153 yield* scheduler.register("usage-limit-recovery", sweep());
154 }),
155);