diff --git a/mobile/src/session/native-chat/transcript/MobileNativeChatUsageLimitStatus.tsx b/mobile/src/session/native-chat/transcript/MobileNativeChatUsageLimitStatus.tsx index bdca6cc312b..6d586d03903 100644 --- a/mobile/src/session/native-chat/transcript/MobileNativeChatUsageLimitStatus.tsx +++ b/mobile/src/session/native-chat/transcript/MobileNativeChatUsageLimitStatus.tsx @@ -43,14 +43,19 @@ export function MobileNativeChatUsageLimitStatus({ const { usageLimit, openAccounts } = useContext(MobileNativeChatNoticeContext) const failure = readWholeAgentSessionFailureFact(notice.failure) const resetsAt = failure?.kind === 'usageLimitReached' ? failure.usageLimit?.resetsAt : undefined - const resetPassed = useResetPassed(resetsAt) const actions = failure?.kind === 'usageLimitReached' && notice.usageLimitTurnItemId !== undefined && usageLimit?.turnItemId === notice.usageLimitTurnItemId ? usageLimit : null - // The line names the reset Orca waits for; Continue still keys off the cut's own reset. + // Continue applies at the cut's own reset, or earlier once the host says the chat moved accounts. + const resetPassed = useResetPassed( + resetsAt !== undefined && actions?.continueAt !== undefined + ? Math.min(resetsAt, actions.continueAt) + : resetsAt + ) + // The line names the reset Orca waits for. const shownAt = actions?.autoContinue ? (actions.continueAt ?? resetsAt) : resetsAt const shownPassed = useResetPassed(shownAt) if (failure?.kind !== 'usageLimitReached') { diff --git a/src/main/native-chat/agent-session-wire/host/testing/structured-agent-session-usage-limit-account-hold.test.ts b/src/main/native-chat/agent-session-wire/host/testing/structured-agent-session-usage-limit-account-hold.test.ts index 3bd4356bc7a..130bbfd2f55 100644 --- a/src/main/native-chat/agent-session-wire/host/testing/structured-agent-session-usage-limit-account-hold.test.ts +++ b/src/main/native-chat/agent-session-wire/host/testing/structured-agent-session-usage-limit-account-hold.test.ts @@ -75,6 +75,10 @@ async function heldChat() { launchUnder: (next: UsageLimitAccountScope) => { attributed = next }, + /** Switches the account a start of this chat would launch on. */ + switchAccount: (path: string) => { + current.host.deps.launchAccountHome = async () => ({ variable: 'CODEX_HOME', path }) + }, setEnabled: (next: boolean) => { enabled = next }, @@ -134,6 +138,15 @@ it('holds again on spending inside the grace, and leaves an unknown reset once r await settleOne(await refused.pass(room + 15_000, (now) => weekly(0, null, now))) }) +it('continues before the cut resets once the chat would start on another account with room', async () => { + const chat = await heldChat() + const start = chat.cut.resetsAt - 60_000 + chat.switchAccount('/homes/b') + expect(await chat.pass(start + 5_000, (now) => weekly(0, null, now))).toEqual([]) + expect(chat.state()).toMatchObject({ state: 'scheduled', continueAt: start + 5_000 }) + await settleOne(await chat.pass(start + 20_000, (now) => weekly(0, null, now))) +}) + it('retires a hold when the launch moves to billing the reading does not describe', async () => { const chat = await heldChat() const weeklyReset = chat.cut.resetsAt + 3 * HOUR diff --git a/src/main/native-chat/agent-session-wire/host/testing/structured-agent-session-usage-limit-continue.test.ts b/src/main/native-chat/agent-session-wire/host/testing/structured-agent-session-usage-limit-continue.test.ts index 34f1bbd5249..de2cab51473 100644 --- a/src/main/native-chat/agent-session-wire/host/testing/structured-agent-session-usage-limit-continue.test.ts +++ b/src/main/native-chat/agent-session-wire/host/testing/structured-agent-session-usage-limit-continue.test.ts @@ -119,6 +119,7 @@ it('continues every chat whose limit reset in the same pass', async () => { sessions: new Map(ids.map((id) => [id, {}])), readSavedCut: () => null, readLoadedCut: cut, + cutAccountHome: () => null, serialize: (_id, run) => run(), // Each send is answered as superseded, which spends no budget; only who was sent matters here. send: async (params) => { diff --git a/src/main/native-chat/agent-session-wire/recovery/structured-agent-session-host-background.ts b/src/main/native-chat/agent-session-wire/recovery/structured-agent-session-host-background.ts index 1a35b18449f..846498ed606 100644 --- a/src/main/native-chat/agent-session-wire/recovery/structured-agent-session-host-background.ts +++ b/src/main/native-chat/agent-session-wire/recovery/structured-agent-session-host-background.ts @@ -14,6 +14,15 @@ export function start( sessions: context.sessions, readLoadedCut: delivery.readUsageLimitCut, readSavedCut: (id) => readSavedUsageLimitCut(context.deps.journalDatabase.db, id), + cutAccountHome: (id, turnItemId) => { + // The child that ran the cut's turn knows its account; otherwise the chat's own. + const session = context.sessions.get(id) + const child = session?.child + if (child?.accountHome && session?.journal.itemFence(turnItemId) === child.fence) { + return child.accountHome + } + return context.deps.store.getRecord(id)?.accountHome ?? null + }, serialize: context.serialize, send: (params) => host.send({ callerKey: 'trusted-local:usage-limit-continue' }, params), publish: delivery.publishStatus, diff --git a/src/main/native-chat/agent-session-wire/recovery/structured-agent-session-usage-limit-account-hold.ts b/src/main/native-chat/agent-session-wire/recovery/structured-agent-session-usage-limit-account-hold.ts index 503e3471d8f..7cbad62333d 100644 --- a/src/main/native-chat/agent-session-wire/recovery/structured-agent-session-usage-limit-account-hold.ts +++ b/src/main/native-chat/agent-session-wire/recovery/structured-agent-session-usage-limit-account-hold.ts @@ -10,8 +10,11 @@ import type { AgentSessionRecord } from '../../../../shared/agent-session/lifecy import type { AgentSessionRecordStore } from '../../../runtime/agent-session/records/agent-session-record-store' import { agentSessionAccountHomePath, + agentSessionAccountHomesEqual, + isLegacyAgentSessionAccountHome, type AgentSessionAccountHome } from '../../../../shared/agent-session/lifecycle/agent-session-account-home' +import { normalizeRuntimePathForComparison } from '../../../../shared/cross-platform-path' import { usageLimitContinueIdentity as identity } from './structured-agent-session-usage-limit-chain' /** The polled usage of the account a chat continues on, as the host attributes it. */ @@ -75,6 +78,8 @@ export class UsageLimitAccountHolds { private readonly scopes = new Map() /** The source witness each chat's attribution was computed under. */ private readonly witnesses = new Map() + /** When each chat was first seen continuing on another account than the one its cut hit. */ + private readonly movedSince = new Map() /** Shared by every chat of a pass and keyed by source witness: work scales with accounts and * their changes, not chats. */ private launchHomes = new Map>() @@ -90,15 +95,22 @@ export class UsageLimitAccountHolds { this.readings = new Map() } - /** Re-attributed each pass: the account a chat continues on follows the host's selection. */ - async attribute(id: string, host: AttributionHost): Promise { + /** Re-attributed each pass: the account a chat continues on follows the host's selection. + * `cutHome` is the account the cut hit, when known. */ + async attribute( + id: string, + host: AttributionHost, + cutHome?: AgentSessionAccountHome | null, + now = Date.now() + ): Promise { const record = host.store.getRecord(id) const read = host.usageLimitAutoContinue?.accountUsage const witness = record ? liveWitness(host, record.provider) : undefined let scope: UsageLimitAccountScope | undefined + let moved = false try { // WSL readings never reach this host, so a WSL launch keeps the cut's own reset. - if (record && read && record.location.wslDistro === null) { + if (record && record.location.wslDistro === null) { // Off WSL a launch home follows only the provider's selection, never the chat. const target = JSON.stringify([witness, record.provider]) let launch = this.launchHomes.get(target) @@ -107,8 +119,9 @@ export class UsageLimitAccountHolds { this.launchHomes.set(target, launch) } const home = (await launch) ?? record.accountHome + moved = cutHome !== undefined && cutHome !== null && !sameAccountHome(home, cutHome) const path = agentSessionAccountHomePath(home) - if (home && path) { + if (read && home && path) { const key = JSON.stringify([witness, record.provider, path]) let reading = this.readings.get(key) if (!reading) { @@ -129,6 +142,11 @@ export class UsageLimitAccountHolds { } this.scopes.set(id, scope) this.witnesses.set(id, witness) + if (!moved) { + this.movedSince.delete(id) + } else if (!this.movedSince.has(id)) { + this.movedSince.set(id, now) + } } /** The witness the chat's attribution was computed under; a send binds to this one. */ @@ -155,20 +173,32 @@ export class UsageLimitAccountHolds { } } } + for (const id of this.movedSince.keys()) { + if (!shown.has(id)) { + this.movedSince.delete(id) + } + } } clear(): void { this.holds.clear() this.scopes.clear() this.witnesses.clear() + this.movedSince.clear() this.beginPass() } verdict( id: string, - cut: { turnItemId: string; resetsAt: number }, + limited: { turnItemId: string; resetsAt: number }, now: number ): UsageLimitAccountVerdict { + // The account the chat moved to never hit this cut: from the move on, only its reading holds. + const movedAt = this.movedSince.get(id) + const cut = + movedAt === undefined + ? limited + : { ...limited, resetsAt: Math.min(limited.resetsAt, movedAt) } const scope = this.scopes.get(id) let hold = this.holds.get(id) // Nothing attributed and nothing held: the cut decides, and nothing is kept for it. @@ -293,6 +323,14 @@ export class UsageLimitAccountHolds { } } +function sameAccountHome(left: AgentSessionAccountHome, right: AgentSessionAccountHome): boolean { + return isLegacyAgentSessionAccountHome(left) && isLegacyAgentSessionAccountHome(right) + ? left.variable === right.variable && + normalizeRuntimePathForComparison(left.path) === + normalizeRuntimePathForComparison(right.path) + : agentSessionAccountHomesEqual(left, right) +} + /** Changes whenever a source the provider's attribution reads changes. */ function liveWitness( host: AttributionHost, diff --git a/src/main/native-chat/agent-session-wire/recovery/structured-agent-session-usage-limit-continue.ts b/src/main/native-chat/agent-session-wire/recovery/structured-agent-session-usage-limit-continue.ts index 7160665287b..8fd72c3b9e1 100644 --- a/src/main/native-chat/agent-session-wire/recovery/structured-agent-session-usage-limit-continue.ts +++ b/src/main/native-chat/agent-session-wire/recovery/structured-agent-session-usage-limit-continue.ts @@ -16,6 +16,7 @@ import type { sendStructuredAgentSessionTurn } from '../operations/structured-ag import { RestartContinuationSupersededError } from './structured-agent-session-restart-continuation' import { UsageLimitAccountHolds } from './structured-agent-session-usage-limit-account-hold' import { UsageLimitColdCuts } from './structured-agent-session-usage-limit-cold-cuts' +import type { AgentSessionAccountHome } from '../../../../shared/agent-session/lifecycle/agent-session-account-home' type Cut = HostUsageLimitCut & { resetsAt: number; turnEndedAt: number } type Attempt = { identity: string; count: number; at: number; unconfirmed?: boolean } @@ -38,6 +39,8 @@ type Deps = { params: Parameters[2] ) => ReturnType readLoadedCut: (id: string) => HostUsageLimitCut | null + /** The account the cut's turn ran on, when known. */ + cutAccountHome: (id: string, turnItemId: string) => AgentSessionAccountHome | null publish: (id: string) => void now: () => number } @@ -132,7 +135,9 @@ export class StructuredAgentSessionUsageLimitContinue { turnItemId: cut.turnItemId, resetsAt: cut.resetsAt, // While the reset is unknown, `hold.resetsAt` is a reading time, not a reset. - ...(!hold.unknownReset && hold.resetsAt > cut.resetsAt ? { continueAt: hold.resetsAt } : {}), + ...(!hold.unknownReset && hold.resetsAt !== cut.resetsAt + ? { continueAt: hold.resetsAt } + : {}), state } } @@ -172,13 +177,7 @@ export class StructuredAgentSessionUsageLimitContinue { if (this.running || this.disposed) { return } - // Off, nothing can be scheduled; the settings listener publishes the switch itself. - if (!this.enabled()) { - if (this.holds.size + this.cold.size + this.attempts.size > 0) { - this.retainShown(this.deps.host.store.listVisibleSessionIds()) - } - return - } + // Off, nothing is sent, but an account switch still moves when Continue applies. this.running = true try { const visible = this.deps.host.store.listVisibleSessionIds() @@ -197,10 +196,10 @@ export class StructuredAgentSessionUsageLimitContinue { if (!cut || previous?.identity !== identity(cut)) { this.attempts.delete(id) } - // Before the overdue check: after a restart only the reading still knows a later reset. - // Only a reading taken after the cut's reset counts, so attribution waits for it. - if (cut?.resetsAt !== undefined && this.deps.now() >= cut.resetsAt) { - await this.holds.attribute(id, this.deps.host) + // Before the overdue check: after a restart only the reading still knows a later reset, + // and a switched account frees the chat before it. + if (cut?.resetsAt !== undefined) { + await this.attribute(id, cut) } const attempt = this.attempts.get(id) if (!cut || !this.eligible(id, cut) || attempt?.unconfirmed || (attempt?.count ?? 0) >= 5) { @@ -226,6 +225,15 @@ export class StructuredAgentSessionUsageLimitContinue { } } + private attribute(id: string, cut: HostUsageLimitCut): Promise { + return this.holds.attribute( + id, + this.deps.host, + this.deps.cutAccountHome(id, cut.turnItemId), + this.deps.now() + ) + } + private deferUnchangedCut(id: string, cut: Cut, at: number): void { const current = this.cold.read(id) // Admission also refuses temporary backlog; yield the next passes without spending a failure. @@ -243,7 +251,7 @@ export class StructuredAgentSessionUsageLimitContinue { private async send(id: string, cut: Cut): Promise { // Earlier sends take time; moved sources resolve again, once per new witness for every chat. if (!this.holds.attributionCurrent(id, this.deps.host)) { - await this.holds.attribute(id, this.deps.host) + await this.attribute(id, cut) } const record = this.deps.host.store.getRecord(id) if (!record) { diff --git a/src/renderer/src/components/native-chat/messages/notices/NativeChatNoticeRow.tsx b/src/renderer/src/components/native-chat/messages/notices/NativeChatNoticeRow.tsx index 9e569b40797..95ceca40858 100644 --- a/src/renderer/src/components/native-chat/messages/notices/NativeChatNoticeRow.tsx +++ b/src/renderer/src/components/native-chat/messages/notices/NativeChatNoticeRow.tsx @@ -321,7 +321,9 @@ export function NativeChatNoticeRow({ ) : usageLimitContinue ? ( ) : null} diff --git a/src/renderer/src/components/native-chat/messages/notices/native-chat-orca-stop-context.ts b/src/renderer/src/components/native-chat/messages/notices/native-chat-orca-stop-context.ts index bb995befe84..bb0395228f0 100644 --- a/src/renderer/src/components/native-chat/messages/notices/native-chat-orca-stop-context.ts +++ b/src/renderer/src/components/native-chat/messages/notices/native-chat-orca-stop-context.ts @@ -12,6 +12,8 @@ export type NativeChatOrcaStopView = { continueNow: () => void usageLimitTurnItemId: string | null usageLimitAutoContinue?: NativeChatUsageLimitAutoContinueState | null + /** When Continue past the usage limit applies; earlier than its reset after an account switch. */ + usageLimitContinueAt?: number onUsageLimitHandOffStarted?: () => void usageLimitContinueDisabled?: boolean continueUsageLimit: (hadFocus?: boolean) => void diff --git a/src/renderer/src/components/native-chat/usage/NativeChatInterruptedContinue.tsx b/src/renderer/src/components/native-chat/usage/NativeChatInterruptedContinue.tsx index 403a8dd907f..ba52745f3da 100644 --- a/src/renderer/src/components/native-chat/usage/NativeChatInterruptedContinue.tsx +++ b/src/renderer/src/components/native-chat/usage/NativeChatInterruptedContinue.tsx @@ -93,11 +93,8 @@ export function useNativeChatInterruptedContinuation(input: { () => latestNativeChatUsageLimitCut(input.journalItems, input.submissions), [input.journalItems, input.submissions] ) - const usageLimitAutoContinue = useNativeChatUsageLimitAutoContinue( - target, - sessionId, - usageLimitCut - ) + const { auto: usageLimitAutoContinue, continueAt: usageLimitContinueAt } = + useNativeChatUsageLimitAutoContinue(target, sessionId, usageLimitCut) const [attempt, setAttempt] = useState(NO_NATIVE_CHAT_INTERRUPTED_CONTINUE_ATTEMPT) const [usageFailure, setUsageFailure] = useState(null) if (usageFailure !== null && input.pending.some((entry) => entry.phase === 'sending')) { @@ -193,6 +190,7 @@ export function useNativeChatInterruptedContinuation(input: { continueNow, usageLimitTurnItemId, usageLimitAutoContinue, + usageLimitContinueAt, onUsageLimitHandOffStarted: usageLimitAutoContinue?.kind === 'scheduled' ? usageLimitAutoContinue.cancel : undefined, usageLimitContinueDisabled: usageCapability === 'unknown', @@ -207,6 +205,7 @@ export function useNativeChatInterruptedContinuation(input: { continueNow, usageLimitTurnItemId, usageLimitAutoContinue, + usageLimitContinueAt, usageCapability, continueUsageLimit ] diff --git a/src/renderer/src/components/native-chat/usage/use-native-chat-usage-limit-auto-continue.ts b/src/renderer/src/components/native-chat/usage/use-native-chat-usage-limit-auto-continue.ts index a6d8b8dee80..d2d9ae3514c 100644 --- a/src/renderer/src/components/native-chat/usage/use-native-chat-usage-limit-auto-continue.ts +++ b/src/renderer/src/components/native-chat/usage/use-native-chat-usage-limit-auto-continue.ts @@ -24,7 +24,11 @@ export function useNativeChatUsageLimitAutoContinue( target: RuntimeClientTarget, sessionId: string, cut: NativeChatUsageLimitCut | null -): NativeChatUsageLimitAutoContinueState | null { +): { + auto: NativeChatUsageLimitAutoContinueState | null + /** When the host lets Continue apply: earlier than the cut's reset once the chat moved accounts. */ + continueAt: number | undefined +} { const capability = useStructuredAgentSessionHostCapabilityState( target, AGENT_SESSION_USAGE_LIMIT_AUTO_CONTINUE_RUNTIME_CAPABILITY @@ -101,7 +105,7 @@ export function useNativeChatUsageLimitAutoContinue( }) .finally(() => setCancelling(false)) }, [cancelling, readCutState, target, sessionId, turnItemId, resetsAt]) - return useMemo(() => { + const auto = useMemo(() => { const kind = capability === 'supported' ? nativeChatUsageLimitAutoContinueKind(cutState, { @@ -116,4 +120,14 @@ export function useNativeChatUsageLimitAutoContinue( ? { kind, resetsAt: continueAt, cancelDisabled: cancelling, cancel } : { kind } }, [capability, continueAt, cutState, schedulable, target.kind, cancelling, cancel]) + return useMemo( + () => ({ + auto, + continueAt: + resetsAt === undefined || continueAt === undefined + ? resetsAt + : Math.min(resetsAt, continueAt) + }), + [auto, resetsAt, continueAt] + ) } diff --git a/src/shared/agent-session/agent-session-usage-limit-continue-wire.ts b/src/shared/agent-session/agent-session-usage-limit-continue-wire.ts index 95719a1c636..af5e02a1649 100644 --- a/src/shared/agent-session/agent-session-usage-limit-continue-wire.ts +++ b/src/shared/agent-session/agent-session-usage-limit-continue-wire.ts @@ -2,7 +2,8 @@ export type AgentSessionUsageLimitContinue = { turnItemId: string resetsAt: number - /** The later reset the account still shows spent, which Orca waits for; absent: `resetsAt`. */ + /** When Orca continues instead of `resetsAt`: later while the account still shows spent, + * earlier once the chat would continue on another account. Absent: `resetsAt`. */ continueAt?: number state: 'scheduled' | 'cancelled' | 'failed' | 'off' } diff --git a/src/shared/native-chat/native-chat-usage-limit-auto-continue.ts b/src/shared/native-chat/native-chat-usage-limit-auto-continue.ts index d09ef92c864..d9a62af102b 100644 --- a/src/shared/native-chat/native-chat-usage-limit-auto-continue.ts +++ b/src/shared/native-chat/native-chat-usage-limit-auto-continue.ts @@ -25,7 +25,7 @@ export function usageLimitContinueStateForCut( return publishedForCut(published, cut)?.state } -/** The reset Orca waits for before continuing this cut: a later one the account still shows spent. */ +/** When Orca continues this cut: a later reset the account still shows spent, or earlier after a switch. */ export function usageLimitContinueAtForCut( published: PublishedUsageLimitContinue, cut: CutKey