1import * as NodeCrypto from "@effect/platform-node/NodeCrypto";
2import { assert, expect, it } from "@effect/vitest";
3import {
4 CommandId,
5 DEFAULT_SERVER_SETTINGS,
6 ProjectId,
7 ProviderInstanceId,
8 RunId,
9 ThreadId,
10 type OrchestrationV2ServerCommand,
11 type OrchestrationV2ThreadShell,
12} from "@t3tools/contracts";
13import * as DateTime from "effect/DateTime";
14import * as Effect from "effect/Effect";
15import * as Layer from "effect/Layer";
16import * as Queue from "effect/Queue";
17import * as Ref from "effect/Ref";
18import * as SqlClient from "effect/sql/SqlClient";
19import * as TestClock from "effect/testing/TestClock";
20
21import * as ProjectionStore from "../orchestration-v2/ProjectionStore.ts";
22import * as ThreadLaunchService from "../orchestration-v2/ThreadLaunchService.ts";
23import * as ThreadManagementService from "../orchestration-v2/ThreadManagementService.ts";
24import * as UsageLimitRecoveryWorker from "../orchestration-v2/UsageLimitRecoveryWorker.ts";
25import * as SqlitePersistence from "../persistence/Sqlite.ts";
26import * as ScheduledTasks from "../scheduledTasks/ScheduledTaskService.ts";
27import * as SecretRequests from "../secrets/SecretRequests.ts";
28import * as ServerSettings from "../serverSettings.ts";
29import * as Scheduler from "./Scheduler.ts";
30
31it.effect.each(["on time", "after restart"])(
32 "runs Scheduled Tasks and a persisted limit retry through the same scheduler %s",
33 (scenario) =>
34 Effect.gen(function* () {
35 const now = yield* DateTime.now;
36 const resetAt = DateTime.formatIso(DateTime.add(now, { seconds: 60 }));
37 const thread: OrchestrationV2ThreadShell = {
38 id: ThreadId.make("thread:limited"),
39 projectId: ProjectId.make("project:test"),
40 title: "Limited",
41 providerInstanceId: ProviderInstanceId.make("codex"),
42 modelSelection: { instanceId: ProviderInstanceId.make("codex"), model: "gpt-5.4" },
43 runtimeMode: "full-access",
44 interactionMode: "default",
45 createdBy: "user",
46 creationSource: "web",
47 branch: null,
48 worktreePath: null,
49 lineage: {
50 rootThreadId: ThreadId.make("thread:limited"),
51 parentThreadId: null,
52 relationshipToParent: null,
53 },
54 forkedFrom: null,
55 activeProviderThreadId: null,
56 latestRunId: RunId.make("run:limited"),
57 latestRunCompletedAt: now,
58 activeRunId: null,
59 status: "failed",
60 lastErrorClass: "usage_limit",
61 usageLimitResetAt: resetAt,
62 pendingRuntimeRequest: null,
63 latestVisibleMessage: null,
64 latestUserMessageAt: now,
65 hasActionableProposedPlan: false,
66 itemCount: 1,
67 visibleItemCount: 1,
68 createdAt: now,
69 updatedAt: now,
70 archivedAt: null,
71 settledOverride: null,
72 settledAt: null,
73 lastVisitedAt: null,
74 pinnedAt: null,
75 deletedAt: null,
76 limitRecovery: {
77 runId: RunId.make("run:limited"),
78 resetAt,
79 autoResume: true,
80 requestId: CommandId.make("recovery:choice"),
81 },
82 };
83 if (scenario === "after restart") {
84 yield* TestClock.adjust("65 seconds");
85 }
86 const current = yield* Ref.make(thread);
87 const commands = yield* Ref.make<ReadonlyArray<OrchestrationV2ServerCommand>>([]);
88 const receipts = yield* Queue.unbounded<"task" | "retry">();
89 const layerDependencies = Layer.mergeAll(
90 NodeCrypto.layer,
91 Layer.mock(ThreadLaunchService.ThreadLaunchService)({
92 launch: () =>
93 Queue.offer(receipts, "task").pipe(
94 Effect.andThen(Effect.die("fixture dispatch failure")),
95 ),
96 }),
97 Layer.mock(ThreadManagementService.ThreadManagementService)({
98 dispatch: (command) =>
99 Ref.update(commands, (all) => [...all, command]).pipe(
100 Effect.andThen(
101 Ref.update(current, (shell) => ({ ...shell, status: "running" as const })),
102 ),
103 Effect.andThen(Queue.offer(receipts, "retry")),
104 Effect.as({ sequence: 1, storedEvents: [] }),
105 ),
106 }),
107 Layer.mock(ProjectionStore.ProjectionStoreV2)({
108 getLimitRecoveryCandidates: () =>
109 Ref.get(current).pipe(
110 Effect.map((shell) => (shell.status === "failed" ? [shell] : [])),
111 ),
112 }),
113 Layer.mock(ServerSettings.ServerSettingsService)({
114 getSettings: Effect.succeed(DEFAULT_SERVER_SETTINGS),
115 }),
116 Layer.mock(SecretRequests.SecretRequests)({}),
117 );
118 const layerWorkers = Layer.mergeAll(
119 ScheduledTasks.layer,
120 UsageLimitRecoveryWorker.layer,
121 ).pipe(Layer.provide(layerDependencies), Layer.provide(Scheduler.layer));
122 yield* Effect.gen(function* () {
123 const tasks = yield* ScheduledTasks.ScheduledTaskService;
124 const { task } = yield* tasks.upsert({
125 title: "Scheduled work",
126 prompt: "Run scheduled work",
127 enabled: true,
128 schedule: { type: "interval", everyMs: 60_000 },
129 projectId: thread.projectId,
130 workspaceStrategy: { type: "root" },
131 modelSelection: thread.modelSelection,
132 runtimeMode: thread.runtimeMode,
133 interactionMode: thread.interactionMode,
134 });
135 const sql = yield* SqlClient.SqlClient;
136 yield* sql`UPDATE scheduled_tasks SET next_run_at = ${resetAt} WHERE task_id = ${task.id}`;
137 if (scenario === "on time") {
138 yield* TestClock.adjust("55 seconds");
139 assert.deepEqual(yield* Ref.get(commands), []);
140 }
141 yield* TestClock.adjust("5 seconds");
142 assert.deepEqual([yield* Queue.take(receipts), yield* Queue.take(receipts)].sort(), [
143 "retry",
144 "task",
145 ]);
146 const delivered = yield* Ref.get(commands);
147 assert.equal(delivered.length, 1);
148 expect(delivered[0]).toMatchObject({
149 type: "message.dispatch",
150 usageLimitContinuationOfRunId: thread.latestRunId,
151 usageLimitRecoveryRequestId: thread.limitRecovery!.requestId,
152 text: "Continue where you left off.",
153 });
154 }).pipe(Effect.provide(layerWorkers));
155 }).pipe(Effect.provide(SqlitePersistence.layerMemory)),
156);