diff --git a/apps/server/src/orchestration-v2/Orchestrator.ts b/apps/server/src/orchestration-v2/Orchestrator.ts index 0c6ca268f1e0..aec6a05f8f85 100644 --- a/apps/server/src/orchestration-v2/Orchestrator.ts +++ b/apps/server/src/orchestration-v2/Orchestrator.ts @@ -30,6 +30,7 @@ import { type ModelSelection, OrchestrationV2Command, type OrchestrationV2InternalCommand, + type OrchestrationV2LimitRecovery, type OrchestrationV2LimitRecoveryUpdate, type OrchestrationV2ServerCommand, type ThreadPullRequestLink, @@ -759,6 +760,20 @@ function usageLimitHoldsQueue( ); } +// A run that hit the limit again before the previous reset outlasted what the +// usage reading showed, so its recovery waits for the reset instead of resuming +// early on that reading once more. +function limitRecoveryWaitsForReset( + previous: OrchestrationV2LimitRecovery | null | undefined, + runId: RunId, + now: DateTime.Utc, +): boolean { + if (previous == null) return false; + return previous.runId === runId + ? previous.waitForReset === true + : Date.parse(previous.resetAt) > DateTime.toEpochMillis(now); +} + function latestStableRun( projection: Pick, ): OrchestrationV2Run | null { @@ -3159,7 +3174,10 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio thread.settledOverride === "settled" || run?.id !== command.limitRecovery.runId || failure?.class !== "usage_limit" || - failure.resetAt !== command.limitRecovery.resetAt || + (failure.resetAt ?? + (thread.limitRecovery?.runId === run.id && thread.limitRecovery.waitForReset === true + ? thread.limitRecovery.resetAt + : null)) !== command.limitRecovery.resetAt || resetMs <= DateTime.toEpochMillis(run.completedAt ?? run.requestedAt) || projection.runtimeRequests.some((request) => request.status === "pending") ) { @@ -3363,6 +3381,13 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio command.limitRecovery.autoResume ?? previousRecovery?.autoResume ?? false, snooze: command.limitRecovery.snooze ?? previousRecovery?.snooze ?? false, requestId: command.commandId, + ...(limitRecoveryWaitsForReset( + thread.limitRecovery, + command.limitRecovery.runId, + now, + ) + ? { waitForReset: true } + : {}), }; return { ...thread, @@ -5018,8 +5043,13 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio !recovery?.autoResume || recovery.requestId !== command.usageLimitRecoveryRequestId || recovery.runId !== run.id || - recovery.resetAt !== failure.resetAt || - Date.parse(recovery.resetAt) > DateTime.toEpochMillis(now) || + recovery.resetAt !== + (failure.resetAt ?? (recovery.waitForReset === true ? recovery.resetAt : null)) || + (Date.parse(recovery.resetAt) > DateTime.toEpochMillis(now) && + (recovery.waitForReset === true || + command.usageLimitClearedAt === undefined || + Date.parse(command.usageLimitClearedAt) <= + DateTime.toEpochMillis(run.completedAt ?? run.requestedAt))) || projection.thread.archivedAt !== null || projection.thread.deletedAt !== null || projection.thread.settledOverride === "settled" || @@ -5040,14 +5070,38 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio ...(recovery?.autoResume && recovery.requestId === command.usageLimitRecoveryRequestId && recovery.runId === command.usageLimitContinuationOfRunId && - projection.thread.snoozedUntil != null && - DateTime.toEpochMillis(projection.thread.snoozedUntil) > DateTime.toEpochMillis(now) + (projection.runtimeRequests.some((request) => request.status === "pending") || + (projection.thread.snoozedUntil != null && + DateTime.toEpochMillis(projection.thread.snoozedUntil) > + DateTime.toEpochMillis(now))) ? { limitRecovery: { ...recovery, requestId: command.commandId } } : {}), }, }); return; } + if (Date.parse(recovery.resetAt) > DateTime.toEpochMillis(now)) { + const event = yield* emit( + events, + command, + )({ + type: "thread.metadata-updated", + threadId: command.threadId, + occurredAt: now, + payload: { + ...projection.thread, + limitRecovery: { + ...recovery, + runId: idAllocator.derive.run({ + threadId: command.threadId, + ordinal: nextRunOrdinal(projection), + }), + waitForReset: true, + }, + }, + }); + projection = applyToProjection(projection, event); + } } if (command.restartContinuationOfRunId !== undefined) { @@ -9613,6 +9667,9 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio resetAt: failure.resetAt, autoResume: false, snooze: false, + ...(limitRecoveryWaitsForReset(recovery, latest.id, input.now) + ? { waitForReset: true } + : {}), }; } } diff --git a/apps/server/src/orchestration-v2/ProjectionStore.test.ts b/apps/server/src/orchestration-v2/ProjectionStore.test.ts index 7f70a2ec25ee..c3b5fa1782fe 100644 --- a/apps/server/src/orchestration-v2/ProjectionStore.test.ts +++ b/apps/server/src/orchestration-v2/ProjectionStore.test.ts @@ -911,6 +911,101 @@ it.effect("SQL projection keeps background thread updates out of updatedAt", () backgroundThreadUpdatesKeepTimestamp.pipe(Effect.provide(layerTest)), ); +it.effect.each([ + ["memory", ProjectionStore.layerMemory], + ["SQL", layerTest], +] as const)("%s limit recovery loads armed threads before their reset", ([, layer]) => + Effect.gen(function* () { + const store = yield* ProjectionStore.ProjectionStoreV2; + const threadId = yield* addRolledBackRecoveryCandidate("armed-before-reset"); + const projection = yield* store.getThreadProjection(threadId); + const run = projection.runs[0]!; + const now = yield* DateTime.now; + const resetAt = "2099-01-01T00:00:00.000Z"; + yield* store.apply({ + id: EventId.make("event:armed-before-reset:error"), + type: "turn-item.updated", + threadId, + occurredAt: now, + payload: { + id: TurnItemId.make("armed-before-reset:error"), + threadId, + runId: run.id, + nodeId: run.rootNodeId, + providerThreadId: null, + providerTurnId: null, + nativeItemRef: null, + parentItemId: null, + ordinal: 2, + status: "failed", + title: "Usage limit reached", + startedAt: now, + completedAt: now, + updatedAt: now, + type: "error", + failure: { + class: "usage_limit", + message: "Plan limit reached.", + resetAt, + code: "usageLimitExceeded", + retryable: null, + }, + }, + }); + yield* store.apply({ + id: EventId.make("event:armed-before-reset:failed"), + type: "run.updated", + threadId, + occurredAt: now, + payload: { ...run, status: "failed" }, + }); + const recovery = { + runId: run.id, + resetAt, + autoResume: true, + requestId: CommandId.make("recovery:armed-before-reset"), + }; + const updateThread = (suffix: string, patch: Partial) => + store.apply({ + id: EventId.make(`event:armed-before-reset:${suffix}`), + type: "thread.metadata-updated", + threadId, + occurredAt: now, + payload: { ...projection.thread, limitRecovery: recovery, ...patch }, + }); + const candidate = Effect.map( + store.getLimitRecoveryCandidates({ now, autoResume: false, snooze: false }), + (rows) => rows.find((row) => row.id === threadId), + ); + + yield* updateThread("armed", {}); + assert.deepEqual((yield* candidate)?.limitRecovery, recovery); + const retry = { ...recovery, resetAt: "2099-01-02T00:00:00.000Z", waitForReset: true }; + yield* updateThread("retry", { limitRecovery: retry }); + assert.deepEqual((yield* candidate)?.limitRecovery, retry); + const failure = (yield* store.getThreadProjection(threadId)).turnItems.find( + (item) => item.id === TurnItemId.make("armed-before-reset:error"), + ); + assert(failure?.type === "error"); + yield* store.apply({ + id: EventId.make("event:armed-before-reset:no-reset"), + type: "turn-item.updated", + threadId, + occurredAt: now, + payload: { ...failure, failure: { ...failure.failure!, resetAt: null } }, + }); + assert.equal((yield* candidate)?.usageLimitResetAt, retry.resetAt); + yield* updateThread("no-reset-unarmed", {}); + assert.isUndefined(yield* candidate); + yield* updateThread("retry-cancelled", { limitRecovery: { ...retry, autoResume: false } }); + assert.isUndefined(yield* candidate); + yield* updateThread("snoozed", { snoozedUntil: DateTime.makeUnsafe(resetAt) }); + assert.isUndefined(yield* candidate); + yield* updateThread("cancelled", { limitRecovery: { ...recovery, autoResume: false } }); + assert.isUndefined(yield* candidate); + }).pipe(Effect.provide(layer)), +); + it.effect("memory recovery selection includes unfinished items from missing runs", () => Effect.gen(function* () { const projectionStore = yield* ProjectionStore.ProjectionStoreV2; @@ -3892,6 +3987,7 @@ it.layer(layerTest)("ProjectionStoreV2", (it) => { if (lastErrorClass === "usage_limit") { assert.deepEqual(candidate, { id: sqlShell.id, + providerInstanceId: sqlShell.providerInstanceId, status: sqlShell.status, lastErrorClass: sqlShell.lastErrorClass, usageLimitResetAt: sqlShell.usageLimitResetAt, @@ -4070,10 +4166,12 @@ it.layer(layerTest)("ProjectionStoreV2", (it) => { yield* sql`UPDATE orchestration_v2_projection_threads SET payload_json = json_set(payload_json, '$.limitRecovery', json(${encodeUnknownJsonString(recovery)})) WHERE thread_id = ${threadId}`; - assert.isUndefined( + // Armed rows load before the reset too, so the sweep can see an allowance restored early. + assert.deepEqual( (yield* store.getLimitRecoveryCandidates({ ...recoveryOptions, autoResume: true })).find( (row) => row.id === threadId, - ), + )?.limitRecovery, + recovery, ); const due = (yield* store.getLimitRecoveryCandidates({ ...recoveryOptions, diff --git a/apps/server/src/orchestration-v2/ProjectionStore.ts b/apps/server/src/orchestration-v2/ProjectionStore.ts index ed064325bfe1..0221194ca977 100644 --- a/apps/server/src/orchestration-v2/ProjectionStore.ts +++ b/apps/server/src/orchestration-v2/ProjectionStore.ts @@ -171,6 +171,7 @@ export type ProjectionRecoveryKind = export type ProjectionLimitRecoveryCandidate = Pick< OrchestrationV2ThreadShell, | "id" + | "providerInstanceId" | "status" | "lastErrorClass" | "latestRunId" @@ -3862,9 +3863,16 @@ export const layer: Layer.Layer = readonly completed_at: string | null; readonly failure_payload_json: string; readonly last_error: string | null; + readonly recovery_reset_at: string | null; }>` SELECT t.payload_json, r.run_id, r.completed_at, item.payload_json AS failure_payload_json, + COALESCE(json_extract(item.payload_json, '$.failure.resetAt'), CASE + WHEN json_extract(t.payload_json, '$.limitRecovery.runId') IS r.run_id + AND json_extract(t.payload_json, '$.limitRecovery.autoResume') = 1 + AND json_extract(t.payload_json, '$.limitRecovery.waitForReset') = 1 + THEN json_extract(t.payload_json, '$.limitRecovery.resetAt') + END) AS recovery_reset_at, ( SELECT json_extract(session.payload_json, '$.lastError') -- Keep bindings outermost: one provider instance can hold every session. @@ -3910,14 +3918,16 @@ export const layer: Layer.Layer = AND json_extract(t.payload_json, '$.archivedAt') IS NULL AND json_extract(t.payload_json, '$.settledOverride') IS NOT 'settled' AND json_extract(item.payload_json, '$.failure.class') = 'usage_limit' - AND json_extract(item.payload_json, '$.failure.resetAt') IS NOT NULL - AND julianday(json_extract(item.payload_json, '$.failure.resetAt')) > julianday(COALESCE(r.completed_at, json_extract(t.payload_json, '$.updatedAt'))) + AND recovery_reset_at IS NOT NULL + AND julianday(recovery_reset_at) > julianday(COALESCE(r.completed_at, json_extract(t.payload_json, '$.updatedAt'))) AND ( ( json_extract(t.payload_json, '$.limitRecovery.runId') IS r.run_id - AND json_extract(t.payload_json, '$.limitRecovery.resetAt') IS json_extract(item.payload_json, '$.failure.resetAt') + AND ( + json_extract(t.payload_json, '$.limitRecovery.resetAt') IS recovery_reset_at + OR json_extract(t.payload_json, '$.limitRecovery.waitForReset') = 1 + ) AND json_extract(t.payload_json, '$.limitRecovery.autoResume') = 1 - AND julianday(json_extract(item.payload_json, '$.failure.resetAt')) <= julianday(${DateTime.formatIso(options.now)}) AND ( json_extract(t.payload_json, '$.snoozedUntil') IS NULL OR julianday(json_extract(t.payload_json, '$.snoozedUntil')) <= julianday(${DateTime.formatIso(options.now)}) @@ -3926,11 +3936,11 @@ export const layer: Layer.Layer = OR ( ( json_extract(t.payload_json, '$.limitRecovery.runId') IS NOT r.run_id - OR json_extract(t.payload_json, '$.limitRecovery.resetAt') IS NOT json_extract(item.payload_json, '$.failure.resetAt') + OR json_extract(t.payload_json, '$.limitRecovery.resetAt') IS NOT recovery_reset_at ) AND ( ${booleanInt(options.autoResume)} - OR (${booleanInt(options.snooze)} AND julianday(json_extract(item.payload_json, '$.failure.resetAt')) > julianday(${DateTime.formatIso(options.now)})) + OR (${booleanInt(options.snooze)} AND julianday(recovery_reset_at) > julianday(${DateTime.formatIso(options.now)})) ) ) ) @@ -3951,9 +3961,10 @@ export const layer: Layer.Layer = if (summary.lastErrorClass !== "usage_limit") continue; candidates.push({ id: thread.id, + providerInstanceId: thread.providerInstanceId, status: "failed", lastErrorClass: summary.lastErrorClass, - usageLimitResetAt: summary.usageLimitResetAt, + usageLimitResetAt: row.recovery_reset_at, latestRunId: RunId.make(row.run_id), latestRunCompletedAt: row.completed_at === null ? null : DateTime.makeUnsafe(row.completed_at), @@ -7071,6 +7082,16 @@ export const layerMemory: Layer.Layer = Layer.effect( thread.settledOverride !== "settled", ) .map(threadShellFromProjection) + .map((thread) => ({ + ...thread, + usageLimitResetAt: + thread.usageLimitResetAt ?? + (thread.limitRecovery?.runId === thread.latestRunId && + thread.limitRecovery.autoResume && + thread.limitRecovery.waitForReset === true + ? thread.limitRecovery.resetAt + : null), + })) .filter( (thread) => thread.status === "failed" && @@ -7088,13 +7109,15 @@ export const layerMemory: Layer.Layer = Layer.effect( return false; if ( thread.limitRecovery?.runId !== thread.latestRunId || - thread.limitRecovery.resetAt !== thread.usageLimitResetAt + (thread.limitRecovery.resetAt !== thread.usageLimitResetAt && + !( + thread.limitRecovery.autoResume && thread.limitRecovery.waitForReset === true + )) ) { return options.autoResume || (options.snooze && resetMs > nowMs); } return ( thread.limitRecovery.autoResume && - resetMs <= nowMs && (thread.snoozedUntil == null || DateTime.toEpochMillis(thread.snoozedUntil) <= nowMs) ); diff --git a/apps/server/src/orchestration-v2/UsageLimitRecoveryWorker.ts b/apps/server/src/orchestration-v2/UsageLimitRecoveryWorker.ts index b6ebedfec4ca..59ac0c18afb7 100644 --- a/apps/server/src/orchestration-v2/UsageLimitRecoveryWorker.ts +++ b/apps/server/src/orchestration-v2/UsageLimitRecoveryWorker.ts @@ -1,7 +1,13 @@ -import { CommandId, MessageId, type OrchestrationV2Command } from "@t3tools/contracts"; +import { + CommandId, + MessageId, + type OrchestrationV2Command, + type ServerProviderUsageLimits, +} from "@t3tools/contracts"; import * as DateTime from "effect/DateTime"; import * as Effect from "effect/Effect"; import * as Layer from "effect/Layer"; +import * as ProviderRegistry from "../provider/ProviderRegistry.ts"; import * as Scheduler from "../scheduling/Scheduler.ts"; import * as ServerSettings from "../serverSettings.ts"; import * as ProjectionStore from "./ProjectionStore.ts"; @@ -13,6 +19,7 @@ export function limitRecoveryCommand( autoResume: boolean, nowMs: number, snooze = false, + usageLimits?: ServerProviderUsageLimits, ): OrchestrationV2Command | null { if ( thread.status !== "failed" || @@ -25,15 +32,16 @@ export function limitRecoveryCommand( ) return null; const resetMs = Date.parse(thread.usageLimitResetAt); + const failedAtMs = DateTime.toEpochMillis(thread.latestRunCompletedAt ?? thread.updatedAt); // An already-expired window reported with a fresh failure cannot start a retry loop. - if ( - !Number.isFinite(resetMs) || - resetMs <= DateTime.toEpochMillis(thread.latestRunCompletedAt ?? thread.updatedAt) - ) - return null; + if (!Number.isFinite(resetMs) || resetMs <= failedAtMs) return null; const identity = `${thread.id}:${thread.latestRunId}:${resetMs}`; const recovery = thread.limitRecovery; if (recovery?.runId !== thread.latestRunId || recovery.resetAt !== thread.usageLimitResetAt) { + autoResume ||= + recovery?.runId === thread.latestRunId && + recovery.autoResume && + recovery.waitForReset === true; if (!autoResume && (!snooze || resetMs <= nowMs)) return null; return { type: "thread.metadata.update", @@ -47,13 +55,29 @@ export function limitRecoveryCommand( }, }; } + // A reset credit or provider refresh can restore the allowance before the + // reported reset. Only a reading taken after the failure counts. A reading can + // omit the limit that stopped the run, so once an early run hits the limit + // again the server sets waitForReset and the thread waits for the reset. + const clearedAt = + resetMs > nowMs && + recovery.waitForReset !== true && + usageLimits !== undefined && + usageLimits.unavailable === undefined && + usageLimits.windows.length > 0 && + usageLimits.windows.every((window) => window.usedPercent < 100) && + Date.parse(usageLimits.checkedAt) > failedAtMs + ? usageLimits.checkedAt + : undefined; if ( !recovery.autoResume || - resetMs > nowMs || + (resetMs > nowMs && clearedAt === undefined) || (thread.snoozedUntil != null && DateTime.toEpochMillis(thread.snoozedUntil) > nowMs) ) return null; - const deliveryIdentity = `${identity}:${recovery.requestId ?? "legacy"}`; + // A separate id for early delivery keeps a rejected one from claiming the + // receipt of the delivery at the reset. + const deliveryIdentity = `${identity}:${recovery.requestId ?? "legacy"}${clearedAt === undefined ? "" : ":early"}`; return { type: "message.dispatch", commandId: CommandId.make(`limit-resume:${deliveryIdentity}`), @@ -63,6 +87,7 @@ export function limitRecoveryCommand( ...(recovery.requestId === undefined ? {} : { usageLimitRecoveryRequestId: recovery.requestId }), + ...(clearedAt === undefined ? {} : { usageLimitClearedAt: clearedAt }), text: "Continue where you left off.", attachments: [], dispatchMode: { type: "start_immediately" }, @@ -76,6 +101,7 @@ const makeSweep = Effect.gen(function* () { const threads = yield* ThreadManagement.ThreadManagementService; const orchestrator = yield* Orchestrator.OrchestratorV2; const settings = yield* ServerSettings.ServerSettingsService; + const providerRegistry = yield* ProviderRegistry.ProviderRegistry; return Effect.fn("UsageLimitRecoveryWorker.sweep")(function* () { const preferences = yield* settings.getSettings; const now = yield* DateTime.now; @@ -85,12 +111,15 @@ const makeSweep = Effect.gen(function* () { snooze: preferences.snoozeLimitedThreads, }); const nowMs = DateTime.toEpochMillis(now); + const providers = candidates.length === 0 ? [] : yield* providerRegistry.getProviders; for (const thread of candidates) { const command = limitRecoveryCommand( thread, preferences.autoResumeLimitedThreads, nowMs, preferences.snoozeLimitedThreads, + providers.find((provider) => provider.instanceId === thread.providerInstanceId) + ?.usageLimits, ); if (command === null) continue; yield* Effect.gen(function* () { diff --git a/apps/server/src/orchestration-v2/runtimeLayer.test.ts b/apps/server/src/orchestration-v2/runtimeLayer.test.ts index 025191f09882..1fa221812f83 100644 --- a/apps/server/src/orchestration-v2/runtimeLayer.test.ts +++ b/apps/server/src/orchestration-v2/runtimeLayer.test.ts @@ -1,5 +1,7 @@ import { SourceControlProviderKind } from "@t3tools/contracts"; import { limitRecoveryCommand } from "./UsageLimitRecoveryWorker.ts"; +import * as UsageLimitRecoveryWorker from "./UsageLimitRecoveryWorker.ts"; +import * as ProviderRegistry from "../provider/ProviderRegistry.ts"; import * as McpProviderSessions from "@t3tools/provider-core/server/McpProviderSessions"; import * as SourceControlProviderRegistry from "../sourceControl/SourceControlProviderRegistry.ts"; import * as NodeServices from "@effect/platform-node/NodeServices"; @@ -18,6 +20,7 @@ import { TurnItemId, type ModelSelection, type OrchestrationV2Run, + type ServerProviderUsageLimits, ProjectId, type PullRequestDetail, type PullRequestComment, @@ -6180,6 +6183,311 @@ it.layer(layerTest)("usage-limit recovery", (it) => { }), ); + it.effect.each([true, false])( + "preserves early recovery with global auto-resume off when retry reports reset: %s", + (retryReportsReset) => + Effect.gen(function* () { + const orchestrator = yield* Orchestrator.OrchestratorV2; + const events = yield* EventSink.EventSinkV2; + const projections = yield* ProjectionStore.ProjectionStoreV2; + const threadId = ThreadId.make(`recovery:early-reset:${retryReportsReset}`); + const projectId = ProjectId.make(`recovery:project:early-reset:${retryReportsReset}`); + const now = yield* DateTime.now; + yield* seedProject({ + projectId, + title: "Recovery project", + workspaceRoot: process.cwd(), + defaultModelSelection: modelSelection, + createdAt: DateTime.formatIso(now), + }); + yield* orchestrator.dispatch({ + type: "thread.create", + commandId: CommandId.make(`recovery:create:early-reset:${retryReportsReset}`), + threadId, + projectId, + title: "Limited thread", + modelSelection, + runtimeMode: "full-access", + interactionMode: "default", + branch: null, + worktreePath: null, + createdBy: "user", + creationSource: "web", + }); + yield* orchestrator.dispatch({ + type: "message.dispatch", + commandId: CommandId.make(`recovery:message:early-reset:${retryReportsReset}`), + threadId, + messageId: MessageId.make(`recovery:message:early-reset:${retryReportsReset}`), + text: "Work on this.", + attachments: [], + dispatchMode: { type: "defer_start" }, + createdBy: "user", + creationSource: "web", + }); + const run = (yield* orchestrator.getThreadProjection(threadId)).runs[0]!; + const resetAt = DateTime.formatIso(DateTime.add(now, { hours: 3 })); + yield* events.write({ + events: [ + { + id: EventId.make(`recovery:run:early-reset:${retryReportsReset}`), + type: "run.updated", + threadId, + occurredAt: now, + payload: { ...run, status: "failed", completedAt: now }, + }, + { + id: EventId.make(`recovery:error:early-reset:${retryReportsReset}`), + type: "turn-item.updated", + threadId, + occurredAt: now, + payload: { + id: TurnItemId.make(`recovery:error:early-reset:${retryReportsReset}`), + type: "error", + threadId, + runId: run.id, + nodeId: run.rootNodeId, + providerThreadId: null, + providerTurnId: null, + nativeItemRef: null, + parentItemId: null, + ordinal: 2, + status: "failed", + title: "Usage limit reached", + startedAt: now, + completedAt: now, + updatedAt: now, + failure: { + class: "usage_limit", + message: "Plan limit reached.", + code: "usageLimitExceeded", + retryable: null, + resetAt, + }, + }, + }, + ], + }); + const shell = (yield* orchestrator.getShellSnapshot()).threads.find( + (thread) => thread.id === threadId, + )!; + const arm = limitRecoveryCommand(shell, true, DateTime.toEpochMillis(now)); + assert(arm?.type === "thread.metadata.update" && arm.limitRecovery != null); + yield* orchestrator.autoArmLimitRecovery({ + commandId: arm.commandId, + threadId, + limitRecovery: arm.limitRecovery, + }); + yield* TestClock.adjust("1 minute"); + const checkedNow = yield* DateTime.now; + const reading = (checkedAt: DateTime.Utc, usedPercent: number) => ({ + checkedAt: DateTime.formatIso(checkedAt), + windows: [{ id: "five_hour", kind: "session" as const, label: "5h", usedPercent }], + }); + const sweep = ( + usageLimits: ServerProviderUsageLimits, + instanceId = shell.providerInstanceId, + ) => + Effect.scoped( + Layer.build( + UsageLimitRecoveryWorker.layer.pipe( + Layer.provide([ + Layer.succeed(ProjectionStore.ProjectionStoreV2, projections), + Layer.succeed(Orchestrator.OrchestratorV2, orchestrator), + Layer.mock(ThreadManagementService.ThreadManagementService)({ + dispatch: orchestrator.dispatch, + }), + ServerSettings.layerTest(), + Layer.mock(ProviderRegistry.ProviderRegistry)({ + getProviders: Effect.succeed([ + { + instanceId, + driver, + enabled: true, + installed: true, + version: null, + status: "ready", + auth: { status: "authenticated" }, + checkedAt: usageLimits.checkedAt, + models: [], + slashCommands: [], + skills: [], + usageLimits, + }, + ]), + }), + Layer.mock(Scheduler.Scheduler)({ + register: (_name, work) => Effect.orDie(work), + }), + ]), + ), + ), + ); + const armed = (yield* projections.getLimitRecoveryCandidates({ + now: checkedNow, + autoResume: true, + snooze: false, + })).find((row) => row.id === threadId); + assert.isDefined(armed); + const nowMs = DateTime.toEpochMillis(checkedNow); + // Readings from before the failure, or ones that still show the limit, keep waiting. + assert.isNull(limitRecoveryCommand(armed!, true, nowMs, false, reading(now, 0))); + assert.isNull(limitRecoveryCommand(armed!, true, nowMs, false, reading(checkedNow, 100))); + yield* sweep(reading(now, 0)); + yield* sweep(reading(checkedNow, 100)); + yield* sweep(reading(checkedNow, 0), alternateInstanceId); + assert.lengthOf((yield* orchestrator.getThreadProjection(threadId)).runs, 1); + const staleResume = limitRecoveryCommand( + armed!, + true, + nowMs, + false, + reading(checkedNow, 0), + ); + assert(staleResume?.type === "message.dispatch"); + yield* orchestrator.dispatch({ + type: "thread.metadata.update", + commandId: CommandId.make(`recovery:cancel:early-reset:${retryReportsReset}`), + threadId, + limitRecovery: { runId: run.id, resetAt, autoResume: false }, + }); + yield* orchestrator.dispatch(staleResume); + yield* sweep(reading(checkedNow, 0)); + assert.lengthOf((yield* orchestrator.getThreadProjection(threadId)).runs, 1); + yield* orchestrator.dispatch({ + type: "thread.metadata.update", + commandId: CommandId.make(`recovery:rearm:early-reset:${retryReportsReset}`), + threadId, + limitRecovery: { runId: run.id, resetAt, autoResume: true }, + }); + const ready = (yield* projections.getLimitRecoveryCandidates({ + now: checkedNow, + autoResume: true, + snooze: false, + })).find((row) => row.id === threadId)!; + const racedResume = limitRecoveryCommand(ready, true, nowMs, false, reading(checkedNow, 0)); + assert(racedResume?.type === "message.dispatch"); + const nodeId = run.rootNodeId; + assert(nodeId !== null); + const request = { + id: RuntimeRequestId.make(`recovery:pending-race:${retryReportsReset}`), + nodeId, + providerTurnId: null, + nativeRequestRef: null, + kind: "user_input" as const, + status: "pending" as const, + responseCapability: { type: "message" as const }, + createdAt: checkedNow, + resolvedAt: null, + }; + yield* events.write({ + events: [ + { + id: EventId.make(`recovery:pending-race:${retryReportsReset}`), + type: "runtime-request.updated", + threadId, + nodeId, + occurredAt: checkedNow, + payload: request, + }, + ], + }); + yield* orchestrator.dispatch(racedResume); + assert.lengthOf((yield* orchestrator.getThreadProjection(threadId)).runs, 1); + yield* events.write({ + events: [ + { + id: EventId.make(`recovery:pending-race:resolved:${retryReportsReset}`), + type: "runtime-request.updated", + threadId, + nodeId, + occurredAt: checkedNow, + payload: { ...request, status: "resolved", resolvedAt: checkedNow }, + }, + ], + }); + yield* sweep(reading(checkedNow, 0)); + const resumed = yield* orchestrator.getThreadProjection(threadId); + assert.lengthOf(resumed.runs, 2); + assert.equal(resumed.runs[1]?.status, "starting"); + + // The reading can omit the limit that stopped the run, such as a + // per-model weekly limit. The early run then fails on the same limit. + const earlyRun = resumed.runs[1]!; + const retryResetAt = retryReportsReset + ? DateTime.formatIso(DateTime.add(now, { hours: 4 })) + : resetAt; + yield* events.write({ + events: [ + { + id: EventId.make(`recovery:run:early-reset:again:${retryReportsReset}`), + type: "run.updated", + threadId, + occurredAt: checkedNow, + payload: { ...earlyRun, status: "failed", completedAt: checkedNow }, + }, + { + id: EventId.make(`recovery:error:early-reset:again:${retryReportsReset}`), + type: "turn-item.updated", + threadId, + occurredAt: checkedNow, + payload: { + id: TurnItemId.make(`recovery:error:early-reset:again:${retryReportsReset}`), + type: "error", + threadId, + runId: earlyRun.id, + nodeId: earlyRun.rootNodeId, + providerThreadId: null, + providerTurnId: null, + nativeItemRef: null, + parentItemId: null, + ordinal: 2, + status: "failed", + title: "Usage limit reached", + startedAt: checkedNow, + completedAt: checkedNow, + updatedAt: checkedNow, + failure: { + class: "usage_limit", + message: "Plan limit reached.", + code: "usageLimitExceeded", + retryable: null, + resetAt: retryReportsReset ? retryResetAt : null, + }, + }, + }, + ], + }); + yield* TestClock.adjust("1 minute"); + const laterNow = yield* DateTime.now; + yield* sweep(reading(laterNow, 0)); + const rearmed = (yield* projections.getLimitRecoveryCandidates({ + now: laterNow, + autoResume: false, + snooze: false, + })).find((row) => row.id === threadId)!; + assert.equal(rearmed.limitRecovery?.waitForReset, true); + assert.equal(rearmed.limitRecovery?.resetAt, retryResetAt); + yield* sweep(reading(laterNow, 0)); + assert.lengthOf((yield* orchestrator.getThreadProjection(threadId)).runs, 2); + assert.isNull( + limitRecoveryCommand( + rearmed, + true, + DateTime.toEpochMillis(laterNow), + false, + reading(laterNow, 0), + ), + ); + const onTime = limitRecoveryCommand(rearmed, false, Date.parse(retryResetAt)); + assert(onTime?.type === "message.dispatch"); + assert.isUndefined(onTime.usageLimitClearedAt); + yield* TestClock.setTime(Date.parse(retryResetAt)); + yield* orchestrator.dispatch(onTime); + assert.lengthOf((yield* orchestrator.getThreadProjection(threadId)).runs, 3); + }), + ); + it.effect.each([ "resume", "queued-resume", diff --git a/docs/user/thread-sidebar.md b/docs/user/thread-sidebar.md index 8bf55006bd7d..1809a0b31923 100644 --- a/docs/user/thread-sidebar.md +++ b/docs/user/thread-sidebar.md @@ -220,6 +220,9 @@ the continuation finishes. If the queue was held by a restart, resume it then. When the provider reports a reset time, choose **Resume at reset** to schedule a continuation. You can cancel it from the thread. Enable **Resume after usage limits** in **Settings → Threads** to schedule limit stops by default. +If the provider's usage shows the allowance back before the reset, for example +after you redeem a reset credit, the continuation starts then. A thread that hits +the limit again before the reset waits for the reset. The environment must be running when the reset arrives; it resumes overdue continuations after a restart. Stopping the thread, sending a new message, archiving, or settling it prevents a pending continuation from starting. diff --git a/packages/contracts/src/orchestrationV2.ts b/packages/contracts/src/orchestrationV2.ts index f4e830397226..da30c9c8b11f 100644 --- a/packages/contracts/src/orchestrationV2.ts +++ b/packages/contracts/src/orchestrationV2.ts @@ -360,6 +360,8 @@ export const OrchestrationV2LimitRecovery = Schema.Struct({ resetAt: IsoDateTime, autoResume: Schema.Boolean, snooze: Schema.optional(Schema.Boolean), + /** Set by the server when this run hit the limit again before the previous reset. */ + waitForReset: Schema.optional(Schema.Boolean), }); export type OrchestrationV2LimitRecovery = typeof OrchestrationV2LimitRecovery.Type; @@ -3010,6 +3012,8 @@ export const OrchestrationV2Command = Schema.Union([ usageLimitContinuationOfRunId: Schema.optional(RunId), manualContinuationOfRunId: Schema.optional(RunId), usageLimitRecoveryRequestId: Schema.optional(CommandId), + /** A usage reading after the failure showed the limits cleared before the reset. */ + usageLimitClearedAt: Schema.optional(IsoDateTime), deliveryIntent: Schema.optional(Schema.Literals(["auto", "steer", "restart"])), delegatedCompletion: Schema.optional( Schema.Struct({