diff --git a/apps/server/src/provider/Layers/CodexAdapter.test.ts b/apps/server/src/provider/Layers/CodexAdapter.test.ts index f96ca33fd958..c16a9cdb6880 100644 --- a/apps/server/src/provider/Layers/CodexAdapter.test.ts +++ b/apps/server/src/provider/Layers/CodexAdapter.test.ts @@ -34,6 +34,7 @@ import * as Schema from "effect/Schema"; import * as Scope from "effect/Scope"; import * as Stream from "effect/Stream"; import * as CodexErrors from "effect-codex-app-server/errors"; +import type * as EffectCodexSchema from "effect-codex-app-server/schema"; import { ServerConfig } from "../../config.ts"; import { ServerSettingsService } from "../../serverSettings.ts"; @@ -63,6 +64,10 @@ const asItemId = (value: string): ProviderItemId => ProviderItemId.make(value); class FakeCodexRuntime implements CodexSessionRuntimeShape { private readonly eventQueue = Effect.runSync(Queue.unbounded()); private readonly now = "2026-01-01T00:00:00.000Z"; + public rateLimitsShouldFail = false; + public accountShouldFail = false; + public rateLimitsGate: Deferred.Deferred | null = null; + public rateLimitsStarted: Deferred.Deferred | null = null; public readonly startImpl = vi.fn(() => Promise.resolve({ @@ -116,6 +121,28 @@ class FakeCodexRuntime implements CodexSessionRuntimeShape { ); public readonly closeImpl = vi.fn(() => Promise.resolve(undefined)); + public readonly readAccountRateLimitsImpl = vi.fn( + (): Promise => + Promise.resolve({ + rateLimits: { + limitId: "codex", + primary: { usedPercent: 25, windowDurationMins: 300 }, + }, + rateLimitsByLimitId: { + "gpt-5.3-codex": { + limitId: "gpt-5.3-codex", + primary: { usedPercent: 40, windowDurationMins: 300, resetsAt: 1_800_000_000 }, + }, + }, + }), + ); + public readonly readAccountImpl = vi.fn( + (): Promise => + Promise.resolve({ + account: { type: "chatgpt", email: "test@example.com", planType: "plus" }, + requiresOpenaiAuth: true, + }), + ); readonly options: CodexSessionRuntimeOptions; @@ -128,6 +155,33 @@ class FakeCodexRuntime implements CodexSessionRuntimeShape { } getSession = Effect.promise(() => this.startImpl()); + get readAccountRateLimits() { + const started = this.rateLimitsStarted; + const gate = this.rateLimitsGate; + const read = () => this.readAccountRateLimitsImpl(); + return this.rateLimitsShouldFail + ? Effect.fail( + new CodexErrors.CodexAppServerTransportError({ + operation: "read-input-stream", + cause: new Error("temporary usage failure"), + }), + ) + : Effect.gen(function* () { + if (started) yield* Deferred.succeed(started, undefined); + if (gate) yield* Deferred.await(gate); + return yield* Effect.promise(read); + }); + } + readAccount = Effect.suspend(() => + this.accountShouldFail + ? Effect.fail( + new CodexErrors.CodexAppServerTransportError({ + operation: "read-input-stream", + cause: new Error("temporary account failure"), + }), + ) + : Effect.promise(() => this.readAccountImpl()), + ); sendTurn(input: CodexSessionRuntimeSendTurnInput) { return Effect.promise(() => this.sendTurnImpl(input)); @@ -164,8 +218,10 @@ class FakeCodexRuntime implements CodexSessionRuntimeShape { function makeRuntimeFactory() { const runtimes: Array = []; + let accountResponse: EffectCodexSchema.V2GetAccountResponse | null = null; const factory = vi.fn((options: CodexSessionRuntimeOptions) => { const runtime = new FakeCodexRuntime(options); + if (accountResponse) runtime.readAccountImpl.mockResolvedValue(accountResponse); runtimes.push(runtime); return Effect.succeed(runtime); }); @@ -175,6 +231,9 @@ function makeRuntimeFactory() { get lastRuntime(): FakeCodexRuntime | undefined { return runtimes.at(-1); }, + set accountResponse(value: EffectCodexSchema.V2GetAccountResponse | null) { + accountResponse = value; + }, }; } @@ -241,6 +300,38 @@ const validationLayer = it.layer( ); validationLayer("CodexAdapterLive validation", (it) => { + it.effect("reads model-specific usage without starting a thread", () => + Effect.gen(function* () { + validationRuntimeFactory.factory.mockClear(); + const adapter = yield* CodexAdapter; + const usage = yield* adapter.readCodexUsage!("gpt-5.3-codex"); + + NodeAssert.equal(validationRuntimeFactory.factory.mock.calls.length, 1); + NodeAssert.equal(validationRuntimeFactory.lastRuntime?.startImpl.mock.calls.length, 0); + NodeAssert.deepStrictEqual( + usage?.windows.map((window) => window.remainingPercent), + [60], + ); + NodeAssert.equal(usage?.model, "gpt-5.3-codex"); + validationRuntimeFactory.factory.mockClear(); + }), + ); + it.effect("suppresses API-key usage without reading rate limits", () => + Effect.gen(function* () { + const adapter = yield* CodexAdapter; + validationRuntimeFactory.factory.mockClear(); + validationRuntimeFactory.accountResponse = { + account: { type: "apiKey" }, + requiresOpenaiAuth: false, + }; + const usage = yield* adapter.readCodexUsage!("gpt-5.3-codex"); + const runtime = validationRuntimeFactory.lastRuntime; + NodeAssert.equal(usage, null); + NodeAssert.equal(runtime?.readAccountRateLimitsImpl.mock.calls.length ?? 0, 0); + validationRuntimeFactory.accountResponse = null; + validationRuntimeFactory.factory.mockClear(); + }), + ); it.effect("returns validation error for non-codex provider on startSession", () => Effect.gen(function* () { const adapter = yield* CodexAdapter; @@ -482,6 +573,7 @@ sessionErrorLayer("CodexAdapterLive session errors", (it) => { }); const lifecycleRuntimeFactory = makeRuntimeFactory(); +let lifecycleUsageNow = new Date("2026-01-01T00:00:00.000Z"); const lifecycleLayer = it.layer( Layer.effect( CodexAdapter, @@ -489,6 +581,7 @@ const lifecycleLayer = it.layer( const codexConfig = decodeCodexSettings({}); return yield* makeCodexAdapter(codexConfig, { makeRuntime: lifecycleRuntimeFactory.factory, + now: () => lifecycleUsageNow, }); }), ).pipe( @@ -514,6 +607,59 @@ function startLifecycleRuntime() { } lifecycleLayer("CodexAdapterLive lifecycle", (it) => { + it.effect( + "merges sparse notifications, preserves observation time, and reconciles full reads", + () => + Effect.gen(function* () { + const { adapter, runtime } = yield* startLifecycleRuntime(); + lifecycleUsageNow = new Date("2026-01-02T00:00:00.000Z"); + const initial = yield* adapter.readCodexUsage!("gpt-5.3-codex"); + NodeAssert.equal(initial?.windows[0]?.remainingPercent, 60); + NodeAssert.equal(initial?.checkedAt, "2026-01-02T00:00:00.000Z"); + + const eventFiber = yield* Stream.runHead(adapter.streamEvents).pipe(Effect.forkChild); + yield* runtime.emit({ + id: asEventId("evt-rate-limits"), + kind: "notification", + provider: ProviderDriverKind.make("codex"), + createdAt: "2026-02-03T04:05:06.000Z", + method: "account/rateLimits/updated", + threadId: asThreadId("thread-1"), + payload: { + rateLimits: { + limitId: "gpt-5.3-codex", + primary: { usedPercent: 65 }, + }, + }, + }); + yield* Fiber.join(eventFiber); + + runtime.rateLimitsShouldFail = true; + const retained = yield* adapter.readCodexUsage!("gpt-5.3-codex"); + NodeAssert.equal(retained?.source, "notification"); + NodeAssert.equal(retained?.checkedAt, "2026-02-03T04:05:06.000Z"); + NodeAssert.equal(retained?.windows[0]?.remainingPercent, 35); + NodeAssert.equal(retained?.windows[0]?.windowDurationMins, 300); + NodeAssert.equal(retained?.windows[0]?.resetsAt, "2027-01-15T08:00:00.000Z"); + + runtime.rateLimitsShouldFail = false; + lifecycleUsageNow = new Date("2026-02-04T00:00:00.000Z"); + runtime.readAccountRateLimitsImpl.mockResolvedValue({ + rateLimits: { limitId: "codex", primary: { usedPercent: 5 } }, + rateLimitsByLimitId: { + "gpt-5.3-codex": { + limitId: "gpt-5.3-codex", + primary: { usedPercent: 10, windowDurationMins: 300 }, + }, + }, + }); + const reconciled = yield* adapter.readCodexUsage!("gpt-5.3-codex"); + NodeAssert.equal(reconciled?.source, "read"); + NodeAssert.equal(reconciled?.windows[0]?.remainingPercent, 90); + NodeAssert.equal(reconciled?.checkedAt, "2026-02-04T00:00:00.000Z"); + }), + ); + it.effect("holds a liveness marker behind a paused lifecycle enqueue", () => Effect.gen(function* () { const mutationObserved = yield* Deferred.make(); @@ -587,6 +733,116 @@ lifecycleLayer("CodexAdapterLive lifecycle", (it) => { }), ); + it.effect("invalidates retained usage across account transitions", () => + Effect.gen(function* () { + const { adapter, runtime } = yield* startLifecycleRuntime(); + NodeAssert.ok(yield* adapter.readCodexUsage!("gpt-5.3-codex")); + const rateLimitReads = runtime.readAccountRateLimitsImpl.mock.calls.length; + + runtime.readAccountImpl.mockResolvedValue({ + account: { type: "apiKey" }, + requiresOpenaiAuth: false, + }); + NodeAssert.equal(yield* adapter.readCodexUsage!("gpt-5.3-codex"), null); + NodeAssert.equal(runtime.readAccountRateLimitsImpl.mock.calls.length, rateLimitReads); + + runtime.accountShouldFail = true; + NodeAssert.equal(yield* adapter.readCodexUsage!("gpt-5.3-codex"), null); + }), + ); + + it.effect("invalidates retained usage on account update notifications", () => + Effect.gen(function* () { + const { adapter, runtime } = yield* startLifecycleRuntime(); + NodeAssert.ok(yield* adapter.readCodexUsage!("gpt-5.3-codex")); + const eventFiber = yield* Stream.runHead(adapter.streamEvents).pipe(Effect.forkChild); + yield* runtime.emit({ + id: asEventId("evt-account-updated"), + kind: "notification", + provider: ProviderDriverKind.make("codex"), + createdAt: "2026-02-03T04:05:06.000Z", + method: "account/updated", + threadId: asThreadId("thread-1"), + payload: { authMode: "apikey" }, + }); + yield* Fiber.join(eventFiber); + runtime.accountShouldFail = true; + NodeAssert.equal(yield* adapter.readCodexUsage!("gpt-5.3-codex"), null); + }), + ); + + it.effect("does not let an older in-flight read overwrite a newer notification", () => + Effect.gen(function* () { + const { adapter, runtime } = yield* startLifecycleRuntime(); + NodeAssert.ok(yield* adapter.readCodexUsage!("gpt-5.3-codex")); + runtime.rateLimitsStarted = yield* Deferred.make(); + runtime.rateLimitsGate = yield* Deferred.make(); + runtime.readAccountRateLimitsImpl.mockResolvedValue({ + rateLimits: { limitId: "codex", primary: { usedPercent: 5 } }, + rateLimitsByLimitId: { + "gpt-5.3-codex": { + limitId: "gpt-5.3-codex", + primary: { usedPercent: 10, windowDurationMins: 300 }, + }, + }, + }); + const readFiber = yield* adapter.readCodexUsage!("gpt-5.3-codex").pipe(Effect.forkChild); + yield* Deferred.await(runtime.rateLimitsStarted); + + const eventFiber = yield* Stream.runHead(adapter.streamEvents).pipe(Effect.forkChild); + yield* runtime.emit({ + id: asEventId("evt-race-rate-limits"), + kind: "notification", + provider: ProviderDriverKind.make("codex"), + createdAt: "2026-03-01T00:00:00.000Z", + method: "account/rateLimits/updated", + threadId: asThreadId("thread-1"), + payload: { + rateLimits: { + limitId: "gpt-5.3-codex", + primary: { usedPercent: 70 }, + }, + }, + }); + yield* Fiber.join(eventFiber); + yield* Deferred.succeed(runtime.rateLimitsGate, undefined); + const raced = yield* Fiber.join(readFiber); + NodeAssert.equal(raced?.source, "notification"); + NodeAssert.equal(raced?.checkedAt, "2026-03-01T00:00:00.000Z"); + NodeAssert.equal(raced?.windows[0]?.remainingPercent, 30); + runtime.rateLimitsGate = null; + runtime.rateLimitsStarted = null; + }), + ); + + it.effect("applies a notification that arrives after a completed read", () => + Effect.gen(function* () { + const { adapter, runtime } = yield* startLifecycleRuntime(); + const read = yield* adapter.readCodexUsage!("gpt-5.3-codex"); + NodeAssert.equal(read?.source, "read"); + const eventFiber = yield* Stream.runHead(adapter.streamEvents).pipe(Effect.forkChild); + yield* runtime.emit({ + id: asEventId("evt-after-read-rate-limits"), + kind: "notification", + provider: ProviderDriverKind.make("codex"), + createdAt: "2026-03-02T00:00:00.000Z", + method: "account/rateLimits/updated", + threadId: asThreadId("thread-1"), + payload: { + rateLimits: { + limitId: "gpt-5.3-codex", + primary: { usedPercent: 75 }, + }, + }, + }); + yield* Fiber.join(eventFiber); + runtime.rateLimitsShouldFail = true; + const retained = yield* adapter.readCodexUsage!("gpt-5.3-codex"); + NodeAssert.equal(retained?.source, "notification"); + NodeAssert.equal(retained?.windows[0]?.remainingPercent, 25); + }), + ); + it.effect("maps completed agent message items to canonical item.completed events", () => Effect.gen(function* () { const { adapter, runtime } = yield* startLifecycleRuntime(); diff --git a/apps/server/src/provider/Layers/CodexAdapter.ts b/apps/server/src/provider/Layers/CodexAdapter.ts index f6b2d8896da2..71b1238ebd0d 100644 --- a/apps/server/src/provider/Layers/CodexAdapter.ts +++ b/apps/server/src/provider/Layers/CodexAdapter.ts @@ -11,6 +11,7 @@ import { type CanonicalItemType, type CanonicalRequestType, type CodexSettings, + type CodexUsageSnapshot, EventId, ProviderDriverKind, type ProviderEvent, @@ -42,6 +43,11 @@ import * as CodexErrors from "effect-codex-app-server/errors"; import * as EffectCodexSchema from "effect-codex-app-server/schema"; import { getModelSelectionStringOptionValue } from "@t3tools/shared/model"; +import { + resolveCodexUsageSnapshot, + type CodexUsageRawPayload, + type CodexUsageRawWindow, +} from "@t3tools/shared/codexUsage"; import { getCodexServiceTierOptionValue } from "../../codexModelOptions.ts"; import * as McpProviderSession from "../../mcp/McpProviderSession.ts"; @@ -90,6 +96,7 @@ export interface CodexAdapterLiveOptions { readonly nativeEventLogPath?: string; readonly nativeEventLogger?: EventNdjsonLogger; readonly beforeRuntimeEventEnqueue?: (event: ProviderRuntimeEvent) => Effect.Effect; + readonly now?: () => Date; } interface CodexAdapterSessionContext { @@ -164,6 +171,49 @@ function readPayload( return isPayload(payload) ? payload : undefined; } +function readRateLimitsUpdate(payload: ProviderEvent["payload"]) { + const wrapped = readPayload(EffectCodexSchema.V2AccountRateLimitsUpdatedNotification, payload); + return ( + wrapped?.rateLimits ?? + readPayload( + EffectCodexSchema.V2AccountRateLimitsUpdatedNotification__RateLimitSnapshot, + payload, + ) + ); +} + +function mergeCodexUsageWindow( + previous: CodexUsageRawWindow | null | undefined, + update: CodexUsageRawWindow | null | undefined, +) { + if (update === undefined) return previous; + if (update === null) return null; + return { + ...(previous && typeof previous === "object" ? previous : {}), + ...(update.usedPercent === undefined ? {} : { usedPercent: update.usedPercent }), + ...(update.resetsAt === undefined ? {} : { resetsAt: update.resetsAt }), + ...(update.windowDurationMins === undefined + ? {} + : { windowDurationMins: update.windowDurationMins }), + }; +} + +function mergeCodexUsageBucket( + previous: CodexUsageRawPayload["rateLimits"], + update: NonNullable, +): NonNullable { + return { + ...previous, + ...(update.limitId === undefined ? {} : { limitId: update.limitId }), + ...(update.limitName === undefined ? {} : { limitName: update.limitName }), + ...(update.rateLimitReachedType === undefined + ? {} + : { rateLimitReachedType: update.rateLimitReachedType }), + primary: mergeCodexUsageWindow(previous?.primary, update.primary), + secondary: mergeCodexUsageWindow(previous?.secondary, update.secondary), + }; +} + function trimText(value: string | undefined | null): string | undefined { const trimmed = value?.trim(); return trimmed && trimmed.length > 0 ? trimmed : undefined; @@ -1395,6 +1445,12 @@ export const makeCodexAdapter = Effect.fn("makeCodexAdapter")(function* ( options?.nativeEventLogger === undefined ? nativeEventLogger : undefined; const runtimeEventQueue = yield* Queue.unbounded(); const sessions = new Map(); + let cachedCodexUsage: { + readonly payload: CodexUsageRawPayload; + readonly checkedAt: string; + readonly source: "read" | "notification"; + } | null = null; + let codexUsageGeneration = 0; const startEpochs = new Map(); const transitions = yield* makeTargetTransitionLock(); const invalidateStart = (threadId: ThreadId) => @@ -1492,6 +1548,37 @@ export const makeCodexAdapter = Effect.fn("makeCodexAdapter")(function* ( event.threadId, Effect.gen(function* () { yield* writeNativeEvent(event); + if (event.method === "account/updated") { + codexUsageGeneration += 1; + cachedCodexUsage = null; + } + if (event.method === "account/rateLimits/updated") { + const update = readRateLimitsUpdate(event.payload); + const limitId = update?.limitId?.trim(); + if (update && limitId) codexUsageGeneration += 1; + if (cachedCodexUsage && update && limitId) { + cachedCodexUsage = { + checkedAt: event.createdAt, + source: "notification", + payload: { + ...cachedCodexUsage.payload, + rateLimits: + cachedCodexUsage.payload.rateLimits?.limitId === limitId + ? mergeCodexUsageBucket(cachedCodexUsage.payload.rateLimits, update) + : cachedCodexUsage.payload.rateLimits, + rateLimitsByLimitId: { + ...cachedCodexUsage.payload.rateLimitsByLimitId, + [limitId]: { + ...mergeCodexUsageBucket( + cachedCodexUsage.payload.rateLimitsByLimitId?.[limitId], + update, + ), + }, + }, + }, + }; + } + } const runtimeEvents = mapToRuntimeEvents(event, event.threadId); if (runtimeEvents.length === 0) { yield* Effect.logDebug("ignoring unhandled Codex provider event", { @@ -1809,6 +1896,121 @@ export const makeCodexAdapter = Effect.fn("makeCodexAdapter")(function* ( const hasSession: CodexAdapterShape["hasSession"] = (threadId) => Effect.succeed(Boolean(sessions.get(threadId) && !sessions.get(threadId)?.stopped)); + const readCodexUsageWithoutSession = Effect.fn("readCodexUsageWithoutSession")(function* () { + const usageThreadId = ThreadId.make("codex-usage"); + const createRuntime = options?.makeRuntime ?? makeCodexSessionRuntime; + return yield* Effect.acquireUseRelease( + Scope.make("sequential"), + (usageScope) => + Effect.gen(function* () { + const runtime = yield* createRuntime({ + threadId: usageThreadId, + providerInstanceId: boundInstanceId, + cwd: process.cwd(), + binaryPath: codexConfig.binaryPath, + launchArgs: resolveCodexLaunchArgs(codexConfig.launchArgs, options?.environment), + ...(options?.environment ? { environment: options.environment } : {}), + ...(codexConfig.homePath ? { homePath: codexConfig.homePath } : {}), + runtimeMode: "full-access", + }).pipe( + Effect.provideService(Scope.Scope, usageScope), + Effect.provideService(ChildProcessSpawner.ChildProcessSpawner, childProcessSpawner), + Effect.mapError( + (cause) => + new ProviderAdapterProcessError({ + provider: PROVIDER, + threadId: usageThreadId, + detail: cause.message, + cause, + }), + ), + ); + return yield* Effect.gen(function* () { + const account = yield* runtime.readAccount.pipe( + Effect.mapError((cause) => + mapCodexRuntimeError(usageThreadId, "account/read", cause), + ), + ); + if (account.account?.type !== "chatgpt") { + codexUsageGeneration += 1; + cachedCodexUsage = null; + return null; + } + return yield* runtime.readAccountRateLimits.pipe( + Effect.mapError((cause) => + mapCodexRuntimeError(usageThreadId, "account/rateLimits/read", cause), + ), + ); + }).pipe(Effect.ensuring(runtime.close)); + }), + (usageScope) => Scope.close(usageScope, Exit.void), + ); + }); + + const readCodexUsage: NonNullable = Effect.fn( + "CodexAdapter.readCodexUsage", + )(function* (model): Effect.fn.Return { + const readGeneration = codexUsageGeneration; + return yield* Effect.gen(function* () { + const session = Array.from(sessions.values()).findLast((candidate) => !candidate.stopped); + const payload = yield* session + ? Effect.gen(function* () { + const account = yield* session.runtime.readAccount.pipe( + Effect.mapError((cause) => + mapCodexRuntimeError(session.threadId, "account/read", cause), + ), + ); + if (account.account?.type !== "chatgpt") { + codexUsageGeneration += 1; + cachedCodexUsage = null; + return null; + } + return yield* session.runtime.readAccountRateLimits.pipe( + Effect.mapError((cause) => + mapCodexRuntimeError(session.threadId, "account/rateLimits/read", cause), + ), + ); + }) + : readCodexUsageWithoutSession(); + if (payload === null) return null; + if (codexUsageGeneration !== readGeneration) { + return cachedCodexUsage + ? resolveCodexUsageSnapshot({ + providerInstanceId: boundInstanceId, + model, + payload: cachedCodexUsage.payload, + source: cachedCodexUsage.source, + checkedAt: cachedCodexUsage.checkedAt, + }) + : null; + } + const checkedAt = (options?.now?.() ?? new Date()).toISOString(); + codexUsageGeneration += 1; + cachedCodexUsage = { payload, checkedAt, source: "read" }; + return resolveCodexUsageSnapshot({ + providerInstanceId: boundInstanceId, + model, + payload, + source: "read", + checkedAt, + }); + }).pipe( + Effect.catch((_error) => + cachedCodexUsage + ? Effect.succeed( + resolveCodexUsageSnapshot({ + providerInstanceId: boundInstanceId, + model, + payload: cachedCodexUsage.payload, + source: cachedCodexUsage.source === "notification" ? "notification" : "cache", + checkedAt: cachedCodexUsage.checkedAt, + }), + ) + : Effect.succeed(null), + ), + ); + }); + const requestLivenessSample: CodexAdapterShape["requestLivenessSample"] = ( threadId, markerId, @@ -1896,6 +2098,7 @@ export const makeCodexAdapter = Effect.fn("makeCodexAdapter")(function* ( listSessions, hasSession, requestLivenessSample, + readCodexUsage, stopAll, get streamEvents() { return Stream.fromQueue(runtimeEventQueue).pipe(Stream.filter(isProviderRuntimeEvent)); diff --git a/apps/server/src/provider/Layers/CodexSessionRuntime.ts b/apps/server/src/provider/Layers/CodexSessionRuntime.ts index 67108dd4dbba..a5433d0e69c5 100644 --- a/apps/server/src/provider/Layers/CodexSessionRuntime.ts +++ b/apps/server/src/provider/Layers/CodexSessionRuntime.ts @@ -27,6 +27,7 @@ import * as Layer from "effect/Layer"; import * as Queue from "effect/Queue"; import * as Ref from "effect/Ref"; import * as Schema from "effect/Schema"; +import * as Semaphore from "effect/Semaphore"; import * as Scope from "effect/Scope"; import * as Stream from "effect/Stream"; import { ChildProcess, ChildProcessSpawner } from "effect/unstable/process"; @@ -141,6 +142,14 @@ export interface CodexSessionRuntimeShape { readonly rollbackThread: ( numTurns: number, ) => Effect.Effect; + readonly readAccountRateLimits: Effect.Effect< + EffectCodexSchema.V2GetAccountRateLimitsResponse, + CodexSessionRuntimeError + >; + readonly readAccount: Effect.Effect< + EffectCodexSchema.V2GetAccountResponse, + CodexSessionRuntimeError + >; readonly respondToRequest: ( requestId: ApprovalRequestId, decision: ProviderApprovalDecision, @@ -724,6 +733,8 @@ export const makeCodexSessionRuntime = ( const pendingUserInputsRef = yield* Ref.make(new Map()); const collabReceiverTurnsRef = yield* Ref.make(new Map()); const closedRef = yield* Ref.make(false); + const initializedRef = yield* Ref.make(false); + const initializeSemaphore = yield* Semaphore.make(1); // `~` is not shell-expanded when env vars are set via // `child_process.spawn`; `expandHomePath` lets a configured @@ -767,6 +778,14 @@ export const makeCodexSessionRuntime = ( const client = yield* Effect.service(CodexClient.CodexAppServerClient).pipe( Effect.provide(clientContext), ); + const ensureInitialized = initializeSemaphore.withPermits(1)( + Effect.gen(function* () { + if (yield* Ref.get(initializedRef)) return; + yield* client.request("initialize", buildCodexInitializeParams()); + yield* client.notify("initialized", undefined); + yield* Ref.set(initializedRef, true); + }), + ); const serverNotifications = yield* Queue.unbounded(); const nowIso = Effect.map(DateTime.now, DateTime.formatIso); const randomUUIDv4 = (purpose: CodexErrors.CodexAppServerIdentifierPurpose) => @@ -1214,8 +1233,7 @@ export const makeCodexSessionRuntime = ( const start = Effect.fn("CodexSessionRuntime.start")(function* () { yield* emitSessionEvent("session/connecting", "Starting Codex App Server session."); - yield* client.request("initialize", buildCodexInitializeParams()); - yield* client.notify("initialized", undefined); + yield* ensureInitialized; const requestedModel = normalizeCodexModelSlug(options.model); @@ -1277,6 +1295,10 @@ export const makeCodexSessionRuntime = ( return { start, getSession: Ref.get(sessionRef), + readAccountRateLimits: ensureInitialized.pipe( + Effect.andThen(client.request("account/rateLimits/read", undefined)), + ), + readAccount: ensureInitialized.pipe(Effect.andThen(client.request("account/read", {}))), sendTurn: (input) => Effect.gen(function* () { const providerThreadId = yield* readProviderThreadId; diff --git a/apps/server/src/provider/Layers/ProviderService.test.ts b/apps/server/src/provider/Layers/ProviderService.test.ts index a178e3380622..262b49f3eca2 100644 --- a/apps/server/src/provider/Layers/ProviderService.test.ts +++ b/apps/server/src/provider/Layers/ProviderService.test.ts @@ -223,6 +223,7 @@ function makeFakeCodexAdapter(provider: ProviderDriverKind = CODEX_DRIVER) { sessions.clear(); }), ); + const readCodexUsage = vi.fn((_model: string) => Effect.succeed(null)); const adapter: ProviderAdapterShape = { provider, @@ -240,6 +241,7 @@ function makeFakeCodexAdapter(provider: ProviderDriverKind = CODEX_DRIVER) { requestLivenessSample, readThread, rollbackThread, + readCodexUsage, stopAll, get streamEvents() { return Stream.fromPubSub(runtimeEventPubSub).pipe(Stream.filter(isProviderRuntimeEvent)); @@ -268,6 +270,7 @@ function makeFakeCodexAdapter(provider: ProviderDriverKind = CODEX_DRIVER) { adapter, emit, updateSession, + readCodexUsage, startSession, sendTurn, interruptTurn, @@ -1935,6 +1938,41 @@ fanout.layer("ProviderServiceLive fanout", (it) => { const validation = makeProviderServiceLayer(); validation.layer("ProviderServiceLive validation", (it) => { + it.effect("contains Codex usage adapter failures", () => + Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + validation.codex.readCodexUsage.mockImplementation(() => + Effect.fail( + new ProviderAdapterRequestError({ + provider: "codex", + method: "account/rateLimits/read", + detail: "temporary failure", + }), + ), + ); + + const usage = yield* provider.getCodexUsage({ + providerInstanceId: codexInstanceId, + model: "gpt-5.3-codex", + }); + assert.equal(usage, null); + validation.codex.readCodexUsage.mockImplementation(() => Effect.succeed(null)); + }), + ); + + it.effect("does not ask non-Codex adapters for Codex usage", () => + Effect.gen(function* () { + const provider = yield* ProviderService.ProviderService; + validation.claude.readCodexUsage.mockClear(); + const usage = yield* provider.getCodexUsage({ + providerInstanceId: claudeAgentInstanceId, + model: "claude-opus", + }); + assert.equal(usage, null); + assert.equal(validation.claude.readCodexUsage.mock.calls.length, 0); + }), + ); + it.effect("rejects session starts without an explicit provider instance id", () => Effect.gen(function* () { const provider = yield* ProviderService.ProviderService; diff --git a/apps/server/src/provider/Layers/ProviderService.ts b/apps/server/src/provider/Layers/ProviderService.ts index abb0b68c7193..63baa9391fbb 100644 --- a/apps/server/src/provider/Layers/ProviderService.ts +++ b/apps/server/src/provider/Layers/ProviderService.ts @@ -1045,6 +1045,22 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( const getInstanceInfo: ProviderServiceMethod<"getInstanceInfo"> = (instanceId) => registry.getInstanceInfo(instanceId); + const getCodexUsage: ProviderServiceMethod<"getCodexUsage"> = Effect.fn( + "ProviderService.getCodexUsage", + )(function* (input) { + const info = yield* registry + .getInstanceInfo(input.providerInstanceId) + .pipe(Effect.catch(() => Effect.succeed(null))); + if (!info || info.driverKind !== "codex") return null; + const adapter = yield* registry + .getByInstance(input.providerInstanceId) + .pipe(Effect.catch(() => Effect.succeed(null))); + if (!adapter?.readCodexUsage) return null; + return yield* adapter + .readCodexUsage(input.model) + .pipe(Effect.catch(() => Effect.succeed(null))); + }); + const inspectTarget: NonNullable> = Effect.fn( "ProviderService.inspectTarget", )(function* (input) { @@ -1191,6 +1207,7 @@ const makeProviderService = Effect.fn("makeProviderService")(function* ( listSessions, getCapabilities, getInstanceInfo, + getCodexUsage, inspectTarget, rollbackConversation, // Each access creates a fresh PubSub subscription so that multiple diff --git a/apps/server/src/provider/Services/ProviderAdapter.ts b/apps/server/src/provider/Services/ProviderAdapter.ts index 72117354301d..64b18c87b054 100644 --- a/apps/server/src/provider/Services/ProviderAdapter.ts +++ b/apps/server/src/provider/Services/ProviderAdapter.ts @@ -9,6 +9,7 @@ */ import type { ApprovalRequestId, + CodexUsageSnapshot, ProviderApprovalDecision, ProviderDriverKind, ProviderUserInputAnswers, @@ -166,6 +167,8 @@ export interface ProviderAdapterShape { numTurns: number, ) => Effect.Effect; + readonly readCodexUsage?: (model: string) => Effect.Effect; + /** * Stop all sessions owned by this adapter. */ diff --git a/apps/server/src/provider/Services/ProviderService.ts b/apps/server/src/provider/Services/ProviderService.ts index 6b39054f9b5f..8644db24a14d 100644 --- a/apps/server/src/provider/Services/ProviderService.ts +++ b/apps/server/src/provider/Services/ProviderService.ts @@ -12,6 +12,7 @@ * @module ProviderService */ import type { + CodexUsageSnapshot, ProviderInterruptTurnInput, ProviderInstanceId, ProviderRespondToRequestInput, @@ -98,6 +99,11 @@ export interface ProviderServiceShape { instanceId: ProviderInstanceId, ) => Effect.Effect; + readonly getCodexUsage: (input: { + readonly providerInstanceId: ProviderInstanceId; + readonly model: string; + }) => Effect.Effect; + readonly inspectTarget?: (input: { readonly providerInstanceId: ProviderInstanceId; readonly threadId: ThreadId; diff --git a/apps/server/src/ws.ts b/apps/server/src/ws.ts index bc78053f9ad0..a6722e887775 100644 --- a/apps/server/src/ws.ts +++ b/apps/server/src/ws.ts @@ -83,6 +83,7 @@ import { observeRpcStreamEffect as instrumentRpcStreamEffect, } from "./observability/RpcInstrumentation.ts"; import * as ProviderRegistry from "./provider/Services/ProviderRegistry.ts"; +import * as ProviderService from "./provider/Services/ProviderService.ts"; import * as ProviderMaintenanceRunner from "./provider/providerMaintenanceRunner.ts"; import * as ServerSelfUpdate from "./cloud/selfUpdate.ts"; import * as ForkUpdate from "./cloud/forkUpdate.ts"; @@ -308,6 +309,7 @@ const RPC_REQUIRED_SCOPE = new Map([ [ORCHESTRATION_WS_METHODS.subscribeThread, AuthOrchestrationReadScope], [WS_METHODS.serverProbe, AuthOrchestrationReadScope], [WS_METHODS.serverGetConfig, AuthOrchestrationReadScope], + [WS_METHODS.serverGetCodexUsage, AuthOrchestrationReadScope], [WS_METHODS.serverRefreshProviders, AuthOrchestrationOperateScope], [WS_METHODS.serverUpdateProvider, AuthOrchestrationOperateScope], [WS_METHODS.serverUpdateServer, AuthOrchestrationOperateScope], @@ -373,6 +375,14 @@ const RPC_REQUIRED_SCOPE = new Map([ [WS_METHODS.subscribeAuthAccess, AuthAccessReadScope], ]); +export function requiredScopeForRpcMethod(method: string): AuthEnvironmentScope { + const requiredScope = RPC_REQUIRED_SCOPE.get(method); + if (requiredScope === undefined) { + throw new Error(`RPC method ${method} has no declared authorization scope.`); + } + return requiredScope; +} + function toAuthAccessStreamEvent( change: PairingGrantStore.BootstrapCredentialChange | SessionStore.SessionCredentialChange, revision: number, @@ -434,6 +444,7 @@ const makeWsRpcLayer = ( const previewManager = yield* PreviewManager.PreviewManager; const portDiscovery = yield* PortScanner.PortDiscovery; const providerRegistry = yield* ProviderRegistry.ProviderRegistry; + const providerService = yield* ProviderService.ProviderService; const providerMaintenanceRunner = yield* ProviderMaintenanceRunner.ProviderMaintenanceRunner; const serverSelfUpdate = yield* ServerSelfUpdate.ServerSelfUpdate; const forkUpdate = yield* ForkUpdate.ForkUpdate; @@ -483,13 +494,6 @@ const makeWsRpcLayer = ( currentSession.scopes.includes(requiredScope) ? stream : Stream.fail(authorizationError(requiredScope)); - const requiredScopeForMethod = (method: string): AuthEnvironmentScope => { - const requiredScope = RPC_REQUIRED_SCOPE.get(method); - if (requiredScope === undefined) { - throw new Error(`RPC method ${method} has no declared authorization scope.`); - } - return requiredScope; - }; const observeRpcEffect = ( method: string, effect: Effect.Effect, @@ -497,7 +501,7 @@ const makeWsRpcLayer = ( ) => instrumentRpcEffect( method, - authorizeEffect(requiredScopeForMethod(method), effect), + authorizeEffect(requiredScopeForRpcMethod(method), effect), traceAttributes, ); const observeRpcStream = ( @@ -507,7 +511,7 @@ const makeWsRpcLayer = ( ) => instrumentRpcStream( method, - authorizeStream(requiredScopeForMethod(method), stream), + authorizeStream(requiredScopeForRpcMethod(method), stream), traceAttributes, ); const observeRpcStreamEffect = ( @@ -521,7 +525,7 @@ const makeWsRpcLayer = ( ) => instrumentRpcStreamEffect( method, - authorizeEffect(requiredScopeForMethod(method), effect), + authorizeEffect(requiredScopeForRpcMethod(method), effect), traceAttributes, ); const toDispatchCommandError = (cause: unknown, fallbackMessage: string) => @@ -1463,6 +1467,12 @@ const makeWsRpcLayer = ( observeRpcEffect(WS_METHODS.serverGetConfig, loadServerConfig, { "rpc.aggregate": "server", }), + [WS_METHODS.serverGetCodexUsage]: (input) => + observeRpcEffect( + WS_METHODS.serverGetCodexUsage, + providerService.getCodexUsage(input).pipe(Effect.catch(() => Effect.succeed(null))), + { "rpc.aggregate": "server" }, + ), [WS_METHODS.serverRefreshProviders]: (input) => observeRpcEffect( WS_METHODS.serverRefreshProviders, diff --git a/apps/server/src/wsAuthorization.test.ts b/apps/server/src/wsAuthorization.test.ts new file mode 100644 index 000000000000..a917eac11785 --- /dev/null +++ b/apps/server/src/wsAuthorization.test.ts @@ -0,0 +1,11 @@ +import { AuthOrchestrationReadScope, WS_METHODS } from "@t3tools/contracts"; +import { describe, expect, it } from "vite-plus/test"; +import { requiredScopeForRpcMethod } from "./ws.ts"; + +describe("WebSocket RPC authorization", () => { + it("allows read-authorized sessions to query Codex usage", () => { + expect(requiredScopeForRpcMethod(WS_METHODS.serverGetCodexUsage)).toBe( + AuthOrchestrationReadScope, + ); + }); +}); diff --git a/apps/web/src/components/chat/ChatComposer.tsx b/apps/web/src/components/chat/ChatComposer.tsx index b427037bdf7a..492f801b02fa 100644 --- a/apps/web/src/components/chat/ChatComposer.tsx +++ b/apps/web/src/components/chat/ChatComposer.tsx @@ -88,6 +88,8 @@ import { renderProviderTraitsPicker, } from "./composerProviderState"; import { ContextWindowMeter } from "./ContextWindowMeter"; +import { CodexUsageIndicator } from "./CodexUsageIndicator"; +import { canShowCodexUsage } from "./codexUsagePresentation"; import { buildExpandedImagePreview, type ExpandedImagePreview } from "./ExpandedImagePreview"; import { basenameOfPath } from "../../pierre-icons"; import { cn, randomUUID } from "~/lib/utils"; @@ -2191,6 +2193,34 @@ export const ChatComposer = memo(function ChatComposer(props: ChatComposerProps) ], ); + const showCodexUsage = canShowCodexUsage(selectedProviderStatus) && !noProviderAvailable; + const providerModelPicker = ( + { + setIsComposerModelPickerOpen(open); + }} + getModelDisabledReason={getModelDisabledReason} + onInstanceModelChange={onProviderModelSelect} + /> + ); + // Render // ------------------------------------------------------------------ return ( @@ -2645,30 +2675,18 @@ export const ChatComposer = memo(function ChatComposer(props: ChatComposerProps) No provider available + ) : showCodexUsage ? ( + + {providerModelPicker} + ) : ( - { - setIsComposerModelPickerOpen(open); - }} - getModelDisabledReason={getModelDisabledReason} - onInstanceModelChange={onProviderModelSelect} - /> + providerModelPicker )} {isComposerFooterCompact ? ( diff --git a/apps/web/src/components/chat/CodexUsageIndicator.tsx b/apps/web/src/components/chat/CodexUsageIndicator.tsx new file mode 100644 index 000000000000..2448e4231da7 --- /dev/null +++ b/apps/web/src/components/chat/CodexUsageIndicator.tsx @@ -0,0 +1,79 @@ +import type { EnvironmentId, ProviderInstanceId } from "@t3tools/contracts"; +import { GaugeIcon } from "lucide-react"; +import { Fragment, memo, type ReactNode, useMemo, useState } from "react"; +import { useEnvironmentQuery } from "../../state/query"; +import { providerUsageEnvironment } from "../../state/providerUsage"; +import { Tooltip, TooltipPopup, TooltipTrigger } from "../ui/tooltip"; +import { cn } from "../../lib/utils"; +import { codexUsagePresentation } from "./codexUsagePresentation"; + +export const CodexUsageIndicator = memo(function CodexUsageIndicator(props: { + readonly environmentId: EnvironmentId; + readonly providerInstanceId: ProviderInstanceId; + readonly model: string; + readonly modelPickerOpen: boolean; + readonly children: ReactNode; +}) { + const [isHoveredOrFocused, setIsHoveredOrFocused] = useState(false); + const atom = useMemo( + () => + providerUsageEnvironment.codex({ + environmentId: props.environmentId, + input: { + providerInstanceId: props.providerInstanceId, + model: props.model, + }, + }), + [props.environmentId, props.model, props.providerInstanceId], + ); + const usage = useEnvironmentQuery(atom).data; + if ( + !usage || + usage.model !== props.model || + usage.providerInstanceId !== props.providerInstanceId + ) { + return {props.children}; + } + const presentation = codexUsagePresentation(usage); + + return ( + + + } + > + {props.children} + + + + + + Codex usage + + {presentation.details} + + + + + + ); +}); diff --git a/apps/web/src/components/chat/ProviderModelPicker.tsx b/apps/web/src/components/chat/ProviderModelPicker.tsx index 74c92a044b3d..015a711a49c2 100644 --- a/apps/web/src/components/chat/ProviderModelPicker.tsx +++ b/apps/web/src/components/chat/ProviderModelPicker.tsx @@ -39,6 +39,8 @@ export const ProviderModelPicker = memo(function ProviderModelPicker(props: { open?: boolean; triggerVariant?: VariantProps["variant"]; triggerClassName?: string; + hideTriggerTooltip?: boolean; + popoverSide?: "top" | "bottom"; onOpenChange?: (open: boolean) => void; getModelDisabledReason?: (instanceId: ProviderInstanceId, model: string) => string | null; onInstanceModelChange: (instanceId: ProviderInstanceId, model: string) => void; @@ -174,12 +176,16 @@ export const ProviderModelPicker = memo(function ProviderModelPicker(props: { )} /> ) : null} - - }> - {triggerTitle} - - {triggerLabel} - + {props.hideTriggerTooltip ? ( + {triggerTitle} + ) : ( + + }> + {triggerTitle} + + {triggerLabel} + + )}