Skip to content

Commit bcbb186

Browse files
kalvenschrautkalvenschraut
committed
fix(usage): distinguish Codex workspaces by account ID
1 parent 08463e2 commit bcbb186

9 files changed

Lines changed: 535 additions & 21 deletions

File tree

‎apps/server/src/provider/Layers/CodexProvider.test.ts‎

Lines changed: 80 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,85 @@
11
import { assert, it } from "@effect/vitest";
2+
import * as NodeFSP from "node:fs/promises";
3+
import * as NodeOS from "node:os";
4+
import * as NodePath from "node:path";
5+
import * as Effect from "effect/Effect";
26

3-
import { applyPreferredCodexDefaultModel, mapCodexModelCapabilities } from "./CodexProvider.ts";
7+
import {
8+
applyPreferredCodexDefaultModel,
9+
mapCodexModelCapabilities,
10+
readCodexAccountId,
11+
resolveCodexAuthHome,
12+
} from "./CodexProvider.ts";
13+
14+
it("resolves the selected or shadow home ahead of inherited credentials", () => {
15+
const inherited = { CODEX_HOME: "/shared", HOME: "/user" };
16+
assert.strictEqual(
17+
resolveCodexAuthHome(
18+
{ homePath: "/shadow", environment: { CODEX_HOME: "/instance" } },
19+
inherited,
20+
),
21+
"/shadow",
22+
);
23+
assert.strictEqual(
24+
resolveCodexAuthHome({ environment: { CODEX_HOME: "/instance" } }, inherited),
25+
"/instance",
26+
);
27+
assert.strictEqual(resolveCodexAuthHome({}, inherited), "/shared");
28+
assert.strictEqual(
29+
resolveCodexAuthHome({ environment: { HOME: "/instance-user" } }, { HOME: "/user" }),
30+
"/instance-user/.codex",
31+
);
32+
assert.strictEqual(resolveCodexAuthHome({}, { HOME: "/user" }), "/user/.codex");
33+
});
34+
35+
it.effect(
36+
"reads workspace identity from the selected home without treating API keys as subscriptions",
37+
() =>
38+
Effect.acquireUseRelease(
39+
Effect.promise(() =>
40+
NodeFSP.mkdtemp(NodePath.join(NodeOS.tmpdir(), "codex-account-identity-")),
41+
),
42+
(root) =>
43+
Effect.gen(function* () {
44+
for (const accountId of ["personal", "work"]) {
45+
const home = NodePath.join(root, accountId);
46+
yield* Effect.promise(() => NodeFSP.mkdir(home));
47+
yield* Effect.promise(() =>
48+
NodeFSP.writeFile(
49+
NodePath.join(home, "auth.json"),
50+
JSON.stringify({ tokens: { account_id: accountId } }),
51+
),
52+
);
53+
assert.strictEqual(yield* readCodexAccountId(home), accountId);
54+
}
55+
yield* Effect.promise(() =>
56+
NodeFSP.writeFile(
57+
NodePath.join(root, "auth.json"),
58+
JSON.stringify({ OPENAI_API_KEY: "local-proxy" }),
59+
),
60+
);
61+
assert.strictEqual(yield* readCodexAccountId(root), undefined);
62+
assert.strictEqual(yield* readCodexAccountId(NodePath.join(root, "missing")), undefined);
63+
const payload = Buffer.from(
64+
JSON.stringify({
65+
"https://api.openai.com/auth": { chatgpt_account_id: "token-workspace" },
66+
}),
67+
).toString("base64url");
68+
for (const [tokens, expected] of [
69+
[{ id_token: `header.${payload}.signature` }, "token-workspace"],
70+
[{ account_id: null, id_token: `header.${payload}.signature` }, "token-workspace"],
71+
[{ account_id: "explicit", id_token: `header.${payload}.signature` }, "explicit"],
72+
[{ id_token: "malformed" }, undefined],
73+
] as const) {
74+
yield* Effect.promise(() =>
75+
NodeFSP.writeFile(NodePath.join(root, "auth.json"), JSON.stringify({ tokens })),
76+
);
77+
assert.strictEqual(yield* readCodexAccountId(root), expected);
78+
}
79+
}),
80+
(root) => Effect.promise(() => NodeFSP.rm(root, { recursive: true, force: true })),
81+
),
82+
);
483

584
it("maps current Codex model capability fields", () => {
685
const capabilities = mapCodexModelCapabilities({

‎apps/server/src/provider/Layers/CodexProvider.ts‎

Lines changed: 76 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,6 @@
1+
import * as NodeFSP from "node:fs/promises";
2+
import * as NodeOS from "node:os";
3+
import * as NodePath from "node:path";
14
import * as DateTime from "effect/DateTime";
25
import * as Duration from "effect/Duration";
36
import * as Effect from "effect/Effect";
@@ -67,6 +70,7 @@ const CODEX_PRESENTATION = {
6770

6871
export interface CodexAppServerProviderSnapshot {
6972
readonly account: CodexSchema.V2GetAccountResponse;
73+
readonly accountId?: string;
7074
readonly rateLimits?: CodexRateLimitsProbe;
7175
readonly version: string | undefined;
7276
readonly models: ReadonlyArray<ServerProviderModel>;
@@ -404,6 +408,65 @@ export const withCodexAppServerClient = Effect.fn("withCodexAppServerClient")(fu
404408
return { client, initialize };
405409
});
406410

411+
const decodeCodexAccountIdentity = Schema.decodeUnknownEffect(
412+
Schema.fromJsonString(
413+
Schema.Struct({
414+
tokens: Schema.Struct({
415+
account_id: Schema.optionalKey(Schema.NullOr(Schema.String)),
416+
id_token: Schema.optionalKey(Schema.NullOr(Schema.String)),
417+
}),
418+
}),
419+
),
420+
);
421+
422+
const decodeCodexIdTokenIdentity = Schema.decodeUnknownEffect(
423+
Schema.fromJsonString(
424+
Schema.Struct({
425+
"https://api.openai.com/auth": Schema.Struct({
426+
chatgpt_account_id: Schema.optionalKey(Schema.NullOr(Schema.String)),
427+
}),
428+
}),
429+
),
430+
);
431+
432+
export function resolveCodexAuthHome(
433+
input: {
434+
readonly homePath?: string;
435+
readonly environment?: NodeJS.ProcessEnv;
436+
},
437+
inheritedEnvironment: NodeJS.ProcessEnv = process.env,
438+
) {
439+
return expandHomePath(
440+
input.homePath?.trim() ||
441+
input.environment?.CODEX_HOME ||
442+
inheritedEnvironment.CODEX_HOME ||
443+
NodePath.join(
444+
input.environment?.HOME || inheritedEnvironment.HOME || NodeOS.homedir(),
445+
".codex",
446+
),
447+
);
448+
}
449+
450+
export const readCodexAccountId = Effect.fn("readCodexAccountId")(function* (homePath: string) {
451+
return yield* Effect.tryPromise(() =>
452+
NodeFSP.readFile(NodePath.join(homePath, "auth.json"), "utf8"),
453+
).pipe(
454+
Effect.flatMap(decodeCodexAccountIdentity),
455+
Effect.flatMap((auth) => {
456+
const accountId = auth.tokens.account_id?.trim();
457+
if (accountId) return Effect.succeed(accountId);
458+
const payload = auth.tokens.id_token?.split(".")[1];
459+
if (!payload) return Effect.succeed(undefined);
460+
return decodeCodexIdTokenIdentity(Buffer.from(payload, "base64url").toString("utf8")).pipe(
461+
Effect.map(
462+
(claims) => claims["https://api.openai.com/auth"].chatgpt_account_id?.trim() || undefined,
463+
),
464+
);
465+
}),
466+
Effect.orElseSucceed(() => undefined),
467+
);
468+
});
469+
407470
const probeCodexAppServerProvider = Effect.fn("probeCodexAppServerProvider")(function* (input: {
408471
readonly binaryPath: string;
409472
readonly homePath?: string;
@@ -419,6 +482,12 @@ const probeCodexAppServerProvider = Effect.fn("probeCodexAppServerProvider")(fun
419482
const version = versionMatch ? versionMatch[1] : undefined;
420483

421484
const accountResponse = yield* client.request("account/read", {});
485+
// account/read omits workspace identity. File-backed logins keep it beside
486+
// the credentials in the same home used by this app-server instance.
487+
const accountId =
488+
accountResponse.account?.type === "chatgpt"
489+
? yield* readCodexAccountId(resolveCodexAuthHome(input))
490+
: undefined;
422491
if (!accountResponse.account && accountResponse.requiresOpenaiAuth) {
423492
return {
424493
account: accountResponse,
@@ -460,6 +529,7 @@ const probeCodexAppServerProvider = Effect.fn("probeCodexAppServerProvider")(fun
460529

461530
return {
462531
account: accountResponse,
532+
...(accountId ? { accountId } : {}),
463533
rateLimits,
464534
version,
465535
models: applyPreferredCodexDefaultModel(
@@ -524,7 +594,10 @@ const makePendingCodexProvider = (
524594
});
525595
});
526596

527-
function accountProbeStatus(account: CodexAppServerProviderSnapshot["account"]): {
597+
function accountProbeStatus(
598+
account: CodexAppServerProviderSnapshot["account"],
599+
accountId?: string,
600+
): {
528601
readonly status: Exclude<ServerProviderState, "disabled">;
529602
readonly auth: ServerProvider["auth"];
530603
readonly message?: string;
@@ -536,6 +609,7 @@ function accountProbeStatus(account: CodexAppServerProviderSnapshot["account"]):
536609
...(account.account?.type ? { type: account.account?.type } : {}),
537610
...(authLabel ? { label: authLabel } : {}),
538611
...(authEmail ? { email: authEmail } : {}),
612+
...(accountId ? { accountId } : {}),
539613
} satisfies ServerProvider["auth"];
540614

541615
if (account.account) {
@@ -646,7 +720,7 @@ export const checkCodexProviderStatus = Effect.fn("checkCodexProviderStatus")(fu
646720
}
647721

648722
const snapshot = probeResult.success.value;
649-
const accountStatus = accountProbeStatus(snapshot.account);
723+
const accountStatus = accountProbeStatus(snapshot.account, snapshot.accountId);
650724
const usageLimits =
651725
snapshot.account.account?.type === "apiKey"
652726
? makeUnavailableUsageLimits({ checkedAt, reason: "unsupported" })

‎apps/server/src/provider/Layers/ProviderRegistry.test.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -370,6 +370,7 @@ it.layer(Layer.mergeAll(NodeServices.layer, ServerSettingsModule.layerTest(), Te
370370
const status = yield* checkCodexProviderStatus(defaultCodexSettings, () =>
371371
Effect.succeed(
372372
makeCodexProbeSnapshot({
373+
accountId: "workspace-a",
373374
skills: [
374375
{
375376
name: "github:gh-fix-ci",
@@ -389,6 +390,7 @@ it.layer(Layer.mergeAll(NodeServices.layer, ServerSettingsModule.layerTest(), Te
389390
assert.strictEqual(status.auth.type, "chatgpt");
390391
assert.strictEqual(status.auth.label, "ChatGPT Pro 20x Subscription");
391392
assert.strictEqual(status.auth.email, "[email protected]");
393+
assert.strictEqual(status.auth.accountId, "workspace-a");
392394
assert.deepStrictEqual(status.models, [
393395
{
394396
slug: "gpt-live-codex",

‎apps/server/src/usage/cliproxyApi.test.ts‎

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,12 @@ const encodeJson = Schema.encodeSync(Schema.fromJsonString(Schema.Unknown));
4747

4848
function fixture(
4949
options: {
50-
accounts?: Array<(typeof accounts)[number] & { disabled?: boolean }>;
50+
accounts?: Array<
51+
Omit<(typeof accounts)[number], "id_token"> & {
52+
disabled?: boolean;
53+
id_token?: { chatgpt_account_id?: string };
54+
}
55+
>;
5156
upstream?: (request: RequestBody) => { status: number; body: unknown };
5257
cooldownStatus?: number;
5358
} = {},
@@ -119,6 +124,7 @@ describe("CLIProxyAPI built-in management API", () => {
119124
const test = fixture();
120125
const api = yield* test.api;
121126
const result = yield* api.readAccounts(config);
127+
expect(result.map((account) => account.accountId)).toEqual(["account-a", "account-b"]);
122128
expect(result.map((account) => account.usageLimits.resetCredits)).toEqual([
123129
{ availableCount: 2, nextCreditId: "first", nextExpiresAt: "2099-01-01T00:00:00.000Z" },
124130
{ availableCount: 2, nextCreditId: "first", nextExpiresAt: "2099-01-01T00:00:00.000Z" },
@@ -141,6 +147,33 @@ describe("CLIProxyAPI built-in management API", () => {
141147
}),
142148
);
143149

150+
it.effect("uses trimmed account IDs for quota headers and omits empty IDs", () =>
151+
Effect.gen(function* () {
152+
const ids = [" account-a ", "", " ", undefined];
153+
const test = fixture({
154+
accounts: ids.map((accountId, index) => ({
155+
...accounts[0]!,
156+
id: `${index}.json`,
157+
auth_index: String(index),
158+
id_token: accountId === undefined ? {} : { chatgpt_account_id: accountId },
159+
})),
160+
});
161+
const api = yield* test.api;
162+
const result = yield* api.readAccounts(config);
163+
expect(result.map((account) => account.accountId)).toEqual([
164+
"account-a",
165+
undefined,
166+
undefined,
167+
undefined,
168+
]);
169+
for (const request of test.requests.filter((request) => request.body?.url)) {
170+
expect(request.body?.header?.["Chatgpt-Account-Id"]).toBe(
171+
request.body?.auth_index === "0" ? "account-a" : undefined,
172+
);
173+
}
174+
}),
175+
);
176+
144177
it.effect("keeps usage when the credits endpoint fails", () =>
145178
Effect.gen(function* () {
146179
const test = fixture({

‎apps/server/src/usage/cliproxyApi.ts‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -151,16 +151,15 @@ export const makeCliproxyApi = Effect.gen(function* () {
151151
url: string,
152152
data?: unknown,
153153
) {
154+
const accountId = account.id_token?.chatgpt_account_id?.trim();
154155
const header =
155156
account.provider === "codex"
156157
? {
157158
Authorization: "Bearer $TOKEN$",
158159
"Content-Type": "application/json",
159160
"OpenAI-Beta": "codex-1",
160161
Originator: "Codex Desktop",
161-
...(account.id_token?.chatgpt_account_id
162-
? { "Chatgpt-Account-Id": account.id_token.chatgpt_account_id }
163-
: {}),
162+
...(accountId ? { "Chatgpt-Account-Id": accountId } : {}),
164163
}
165164
: { Authorization: "Bearer $TOKEN$", "anthropic-beta": "oauth-2025-04-20" };
166165
const raw = yield* management(config, "api-call", {
@@ -205,6 +204,9 @@ export const makeCliproxyApi = Effect.gen(function* () {
205204
id: account.id,
206205
driver: ProviderDriverKind.make(account.provider === "codex" ? "codex" : "claudeAgent"),
207206
...(account.email ? { email: account.email } : {}),
207+
...(account.provider === "codex" && account.id_token?.chatgpt_account_id?.trim()
208+
? { accountId: account.id_token.chatgpt_account_id.trim() }
209+
: {}),
208210
};
209211
const read = Effect.gen(function* () {
210212
if (account.provider === "claude") {

‎packages/contracts/src/providerUsageLimits.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,8 @@ export const UsageLimitSourceAccount = Schema.Struct({
8181
driver: ProviderDriverKind,
8282
/** The signed-in address, when the source names one; clients blur it like provider auth. */
8383
email: Schema.optional(TrimmedNonEmptyString),
84+
/** Provider account or workspace ID; an email can belong to several quota buckets. */
85+
accountId: Schema.optional(TrimmedNonEmptyString),
8486
/** Plan as the matching provider would label it (`ChatGPT Pro 20x Subscription`). */
8587
plan: Schema.optional(TrimmedNonEmptyString),
8688
usageLimits: ServerProviderUsageLimits,

‎packages/contracts/src/server.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@ export const ServerProviderAuth = Schema.Struct({
6363
type: Schema.optional(TrimmedNonEmptyString),
6464
label: Schema.optional(TrimmedNonEmptyString),
6565
email: Schema.optional(TrimmedNonEmptyString),
66+
accountId: Schema.optional(TrimmedNonEmptyString),
6667
});
6768
export type ServerProviderAuth = typeof ServerProviderAuth.Type;
6869

0 commit comments

Comments
 (0)