diff --git a/apps/server/src/orchestration-v2/ProviderTurnStartService.ts b/apps/server/src/orchestration-v2/ProviderTurnStartService.ts index a506cc6e122c6..cf52173f45c6a 100644 --- a/apps/server/src/orchestration-v2/ProviderTurnStartService.ts +++ b/apps/server/src/orchestration-v2/ProviderTurnStartService.ts @@ -45,6 +45,7 @@ import { handoffTokenCapConfig, handoffBudget, attachmentTokenAllowance, + estimateTextTokens, contextUsageForHandoff, historicalMessage, latestNativeContextUsage, @@ -998,7 +999,7 @@ export const layer: Layer.Layer< : 0; return ( sum + - (historical === null ? 0 : Buffer.byteLength(historical.text)) + + (historical === null ? 0 : estimateTextTokens(historical.text)) + nativeAttachments ); }, 0) diff --git a/apps/server/src/orchestration-v2/testkit/ProviderSwitch.integration.test.ts b/apps/server/src/orchestration-v2/testkit/ProviderSwitch.integration.test.ts index f252380bf7f70..306935cdd0975 100644 --- a/apps/server/src/orchestration-v2/testkit/ProviderSwitch.integration.test.ts +++ b/apps/server/src/orchestration-v2/testkit/ProviderSwitch.integration.test.ts @@ -510,7 +510,7 @@ describe("orchestration v2 provider switching", () => { capabilities: ClaudeProviderCapabilitiesV2, modelSelection: CLAUDE_MODEL_SELECTION, responseByRunOrdinal: scenario.includes("target-estimate") - ? { 2: "x".repeat(10_000) } + ? { 2: "x".repeat(64_000) } : {}, capturedTurns, ...(scenario.includes("retry") || scenario.includes("unsent") ? { failStartOnce } : {}), @@ -1230,6 +1230,106 @@ describe("orchestration v2 provider switching", () => { ), ); + it.live("keeps a native session whose stored prose is far below an unreported window", () => + Effect.scoped( + Effect.gen(function* () { + const cwd = yield* checkpointWorkspace("handoff-estimate-headroom"); + const capturedTurns = yield* Ref.make>([]); + const generation = yield* Ref.make(0); + // About 112 KB of ordinary English is about 25k tokens: a quarter of a 128k window. + const prose = "The quick brown fox jumps over the lazy dog. ".repeat(2_500); + const layerRegistry = ProviderAdapterRegistry.layerFromAdapters([ + makeTestAdapter({ + instanceId: CODEX_MODEL_SELECTION.instanceId, + driver: CODEX_DRIVER, + capabilities: CodexProviderCapabilitiesV2, + modelSelection: CODEX_MODEL_SELECTION, + responseByRunOrdinal: { 2: "Source work" }, + capturedTurns, + }), + // No window and no usage, like an ACP registry agent that sends no usage_update. + makeTestAdapter({ + instanceId: CLAUDE_MODEL_SELECTION.instanceId, + driver: CLAUDE_DRIVER, + capabilities: ClaudeProviderCapabilitiesV2, + modelSelection: CLAUDE_MODEL_SELECTION, + responseByRunOrdinal: { 1: prose }, + capturedTurns, + nativeThreadGeneration: generation, + }), + ]); + yield* Effect.gen(function* () { + const orchestrator = yield* Orchestrator.OrchestratorV2; + const worker = yield* EffectWorker.OrchestrationEffectWorkerV2; + const dispatch = (ordinal: number, text: string, selection: ModelSelection) => + orchestrator.dispatch({ + type: "message.dispatch", + commandId: CommandId.make(`headroom:${ordinal}`), + threadId, + messageId: MessageId.make(`headroom:${ordinal}`), + createdBy: "user", + creationSource: "web", + text, + attachments: [], + modelSelection: selection, + dispatchMode: { type: "start_immediately" }, + }); + const wait = (ordinal: number) => + orchestrator.streamStoredEvents.pipe( + Stream.filter( + ({ event }) => + event.type === "run.updated" && + event.payload.ordinal === ordinal && + (event.payload.status === "completed" || event.payload.status === "failed"), + ), + Stream.runHead, + Effect.andThen(worker.drain()), + ); + yield* orchestrator.dispatch({ + type: "thread.create", + commandId: CommandId.make("headroom:create"), + threadId, + projectId, + createdBy: "user", + creationSource: "web", + title: "Handoff estimate headroom", + modelSelection: CLAUDE_MODEL_SELECTION, + runtimeMode: "full-access", + interactionMode: "default", + branch: null, + worktreePath: null, + }); + yield* dispatch(1, "Write the long report", CLAUDE_MODEL_SELECTION); + yield* wait(1); + yield* dispatch(2, "Continue elsewhere", CODEX_MODEL_SELECTION); + yield* wait(2); + yield* dispatch(3, "Back to the first provider", CLAUDE_MODEL_SELECTION); + yield* wait(3); + assert.equal( + (yield* orchestrator.getThreadProjection(threadId)).runs.at(-1)?.status, + "completed", + ); + // A second generation means T3 discarded the native session as full. + assert.equal(yield* Ref.get(generation), 1); + }).pipe( + Effect.provide( + ProviderReplayHarness.layerWithRegistry( + { + name: "handoff-estimate-headroom", + runtimePolicyOverride: { + cwd, + approvalPolicy: "never", + sandboxPolicy: { type: "readOnly" }, + }, + }, + layerRegistry, + ), + ), + ); + }), + ), + ); + it.live.each(["native", "legacy"] as const)( "starts a replacement native thread as new after the resume fallback with %s attempts", (attempts) => diff --git a/packages/provider-core/src/server/handoffBudget.ts b/packages/provider-core/src/server/handoffBudget.ts index 5f9f2f3ae33f6..735647693d06d 100644 --- a/packages/provider-core/src/server/handoffBudget.ts +++ b/packages/provider-core/src/server/handoffBudget.ts @@ -95,7 +95,12 @@ export function attachmentTokenAllowance(attachments: ReadonlyArray