Skip to content

Commit 2e6a18f

Browse files
ndycodendycode
authored
Merge pull request #86 from ndycode/fix/pr85-greptile-followups
fix workspace identity follow-ups in quota cleanup
2 parents 2175fe2 + 9a7cbec commit 2e6a18f

4 files changed

Lines changed: 211 additions & 68 deletions

File tree

‎index.ts‎

Lines changed: 1 addition & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,7 @@ import {
119119
saveFlaggedAccounts,
120120
withFlaggedAccountStorageTransaction,
121121
clearFlaggedAccounts,
122+
getWorkspaceIdentityKey,
122123
StorageError,
123124
formatStorageErrorHint,
124125
type AccountStorageV3,
@@ -185,23 +186,6 @@ import {
185186
getRecoveryToastContent,
186187
} from "./lib/recovery.js";
187188

188-
function getWorkspaceIdentityKey(account: {
189-
organizationId?: string;
190-
accountId?: string;
191-
refreshToken: string;
192-
}): string {
193-
const organizationId = account.organizationId?.trim();
194-
const accountId = account.accountId?.trim();
195-
const refreshToken = account.refreshToken.trim();
196-
if (organizationId) {
197-
return accountId
198-
? `organizationId:${organizationId}|accountId:${accountId}`
199-
: `organizationId:${organizationId}`;
200-
}
201-
if (accountId) return `accountId:${accountId}`;
202-
return `refreshToken:${refreshToken}`;
203-
}
204-
205189
function matchesWorkspaceIdentity(
206190
account: {
207191
organizationId?: string;

‎lib/storage.ts‎

Lines changed: 25 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,28 @@ export interface FlaggedAccountStorageV1 {
3434
accounts: FlaggedAccountMetadataV1[];
3535
}
3636

37+
const normalizeWorkspaceIdentityPart = (value: unknown): string | undefined =>
38+
typeof value === "string" && value.trim().length > 0 ? value.trim() : undefined;
39+
40+
export function getWorkspaceIdentityKey(account: {
41+
organizationId?: string;
42+
accountId?: string;
43+
refreshToken: string;
44+
}): string {
45+
const organizationId = normalizeWorkspaceIdentityPart(account.organizationId);
46+
const accountId = normalizeWorkspaceIdentityPart(account.accountId);
47+
const refreshToken = normalizeWorkspaceIdentityPart(account.refreshToken) ?? "";
48+
if (organizationId) {
49+
return accountId
50+
? `organizationId:${organizationId}|accountId:${accountId}`
51+
: `organizationId:${organizationId}`;
52+
}
53+
if (accountId) {
54+
return `accountId:${accountId}`;
55+
}
56+
return `refreshToken:${refreshToken}`;
57+
}
58+
3759
export type ImportBackupMode = "none" | "best-effort" | "required";
3860

3961
export interface ImportAccountsOptions {
@@ -921,30 +943,6 @@ function normalizeFlaggedStorage(data: unknown): FlaggedAccountStorageV1 {
921943
return { version: 1, accounts: [] };
922944
}
923945

924-
const normalizeFlaggedIdentityPart = (value: unknown): string | undefined =>
925-
typeof value === "string" && value.trim().length > 0 ? value.trim() : undefined;
926-
// Flagged storage must keep sibling workspace entries separate when they share an
927-
// organization but have different accountIds, so this key is more specific than
928-
// the normal account identity collapse used in active storage.
929-
const getFlaggedIdentityKey = (account: {
930-
organizationId?: string;
931-
accountId?: string;
932-
refreshToken: string;
933-
}): string => {
934-
const organizationId = normalizeFlaggedIdentityPart(account.organizationId);
935-
const accountId = normalizeFlaggedIdentityPart(account.accountId);
936-
const refreshToken = normalizeFlaggedIdentityPart(account.refreshToken) ?? "";
937-
if (organizationId) {
938-
return accountId
939-
? `organizationId:${organizationId}|accountId:${accountId}`
940-
: `organizationId:${organizationId}`;
941-
}
942-
if (accountId) {
943-
return `accountId:${accountId}`;
944-
}
945-
return `refreshToken:${refreshToken}`;
946-
};
947-
948946
const byIdentityKey = new Map<string, FlaggedAccountMetadataV1>();
949947
for (const rawAccount of data.accounts) {
950948
if (!isRecord(rawAccount)) continue;
@@ -1024,7 +1022,9 @@ function normalizeFlaggedStorage(data: unknown): FlaggedAccountStorageV1 {
10241022
flaggedReason: typeof rawAccount.flaggedReason === "string" ? rawAccount.flaggedReason : undefined,
10251023
lastError: typeof rawAccount.lastError === "string" ? rawAccount.lastError : undefined,
10261024
};
1027-
byIdentityKey.set(getFlaggedIdentityKey(normalized), normalized);
1025+
// Keep flagged dedup aligned with active cleanup so sibling workspaces only
1026+
// collapse when they resolve to the same shared workspace identity.
1027+
byIdentityKey.set(getWorkspaceIdentityKey(normalized), normalized);
10281028
}
10291029

10301030
return {

‎test/index.test.ts‎

Lines changed: 153 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -272,7 +272,10 @@ const mockSaveFlaggedAccounts = vi.fn(async (nextStorage: typeof mockFlaggedStor
272272
await persistMockFlaggedStorage(nextStorage);
273273
});
274274

275-
vi.mock("../lib/storage.js", () => ({
275+
vi.mock("../lib/storage.js", async () => {
276+
const actual = await vi.importActual<typeof import("../lib/storage.js")>("../lib/storage.js");
277+
278+
return {
276279
getStoragePath: () => "/mock/path/accounts.json",
277280
loadAccounts: vi.fn(async () => cloneMockStorage()),
278281
saveAccounts: vi.fn(async (nextStorage: typeof mockStorage) => {
@@ -311,6 +314,7 @@ vi.mock("../lib/storage.js", () => ({
311314
previewImportAccounts: vi.fn(async () => ({ imported: 2, skipped: 1, total: 5 })),
312315
createTimestampedBackupPath: vi.fn((prefix?: string) => `/tmp/${prefix ?? "codex-backup"}-20260101-000000.json`),
313316
loadFlaggedAccounts: vi.fn(async () => cloneMockFlaggedStorage()),
317+
getWorkspaceIdentityKey: vi.fn(actual.getWorkspaceIdentityKey),
314318
saveFlaggedAccounts: mockSaveFlaggedAccounts,
315319
withFlaggedAccountStorageTransaction: vi.fn(
316320
async <T>(
@@ -335,7 +339,8 @@ vi.mock("../lib/storage.js", () => ({
335339
}
336340
},
337341
formatStorageErrorHint: () => "Check file permissions",
338-
}));
342+
};
343+
});
339344

340345
vi.mock("../lib/accounts.js", () => {
341346
class MockAccountManager {
@@ -438,12 +443,15 @@ vi.mock("../lib/accounts.js", () => {
438443
),
439444
extractAccountEmail: vi.fn(() => "user@example.com"),
440445
extractAccountId: vi.fn(() => "account-1"),
441-
resolveRequestAccountId: (_storedId: string | undefined, _source: string | undefined, tokenId: string | undefined) => tokenId,
446+
resolveRequestAccountId: vi.fn(
447+
(_storedId: string | undefined, _source: string | undefined, tokenId: string | undefined) =>
448+
tokenId,
449+
),
442450
formatAccountLabel: (_account: unknown, index: number) => `Account ${index + 1}`,
443451
formatCooldown: () => null,
444452
formatWaitTime: (ms: number) => `${Math.round(ms / 1000)}s`,
445453
sanitizeEmail: (email: string) => email,
446-
shouldUpdateAccountIdFromToken: () => true,
454+
shouldUpdateAccountIdFromToken: vi.fn(() => true),
447455
parseRateLimitReason: () => "unknown",
448456
lookupCodexCliTokensByEmail: vi.fn(async () => null),
449457
};
@@ -2639,18 +2647,16 @@ describe("OpenAIOAuthPlugin fetch handler", () => {
26392647
expect(removeAccount).toHaveBeenCalledTimes(1);
26402648
expect(removeAccountsWithSameRefreshToken).not.toHaveBeenCalled();
26412649
expect(accounts.map((account) => account.accountId)).toEqual(["org-live"]);
2642-
expect(vi.mocked(storageModule.saveFlaggedAccounts)).toHaveBeenCalledTimes(1);
2643-
expect(vi.mocked(storageModule.saveFlaggedAccounts).mock.calls[0]?.[0]).toEqual(
2644-
expect.objectContaining({
2645-
accounts: expect.arrayContaining([
2646-
expect.objectContaining({
2647-
accountId: "org-dead",
2648-
organizationId: "org-dead",
2649-
flaggedReason: "workspace-deactivated",
2650-
lastError: "deactivated_workspace",
2651-
}),
2652-
]),
2653-
}),
2650+
expect(vi.mocked(storageModule.withFlaggedAccountStorageTransaction)).toHaveBeenCalledTimes(1);
2651+
expect(mockFlaggedStorage.accounts).toEqual(
2652+
expect.arrayContaining([
2653+
expect.objectContaining({
2654+
accountId: "org-dead",
2655+
organizationId: "org-dead",
2656+
flaggedReason: "workspace-deactivated",
2657+
lastError: "deactivated_workspace",
2658+
}),
2659+
]),
26542660
);
26552661
});
26562662

@@ -3618,16 +3624,137 @@ describe("OpenAIOAuthPlugin persistAccountPool", () => {
36183624
expect(mockStorage.accounts.some((account) => account.accountId === "workspace-dead")).toBe(false);
36193625
expect(mockStorage.accounts[0]?.organizationId).toBe("org-shared");
36203626
expect(mockStorage.accounts[0]?.refreshToken).toBe("shared-refresh");
3621-
expect(vi.mocked(storageModule.saveFlaggedAccounts)).toHaveBeenCalledWith(
3622-
expect.objectContaining({
3623-
accounts: expect.arrayContaining([
3624-
expect.objectContaining({
3625-
accountId: "workspace-dead",
3626-
organizationId: "org-shared",
3627-
flaggedReason: "token-invalid",
3628-
}),
3629-
]),
3630-
}),
3627+
expect(vi.mocked(storageModule.withFlaggedAccountStorageTransaction)).toHaveBeenCalled();
3628+
expect(mockFlaggedStorage.accounts).toEqual(
3629+
expect.arrayContaining([
3630+
expect.objectContaining({
3631+
accountId: "workspace-dead",
3632+
organizationId: "org-shared",
3633+
flaggedReason: "token-invalid",
3634+
}),
3635+
]),
3636+
);
3637+
});
3638+
3639+
it("removes only the deactivated org-scoped workspace during quota-check cleanup", async () => {
3640+
const accountsModule = await import("../lib/accounts.js");
3641+
const cliModule = await import("../lib/cli.js");
3642+
const refreshQueueModule = await import("../lib/refresh-queue.js");
3643+
const fetchHelpersModule = await import("../lib/request/fetch-helpers.js");
3644+
const storageModule = await import("../lib/storage.js");
3645+
3646+
mockStorage.accounts = [
3647+
{
3648+
refreshToken: "shared-refresh",
3649+
organizationId: "org-shared",
3650+
accountId: "workspace-dead",
3651+
accountIdSource: "manual",
3652+
email: "dead@example.com",
3653+
addedAt: 1,
3654+
lastUsed: 1,
3655+
},
3656+
{
3657+
refreshToken: "shared-refresh",
3658+
organizationId: "org-shared",
3659+
accountId: "workspace-live",
3660+
accountIdSource: "manual",
3661+
email: "live@example.com",
3662+
addedAt: 2,
3663+
lastUsed: 2,
3664+
},
3665+
];
3666+
3667+
vi.mocked(cliModule.promptLoginMode)
3668+
.mockResolvedValueOnce({ mode: "check" })
3669+
.mockResolvedValueOnce({ mode: "cancel" });
3670+
vi.mocked(accountsModule.shouldUpdateAccountIdFromToken).mockImplementation(
3671+
(source: string | undefined, currentAccountId?: string) => {
3672+
if (!currentAccountId) return true;
3673+
if (!source) return true;
3674+
return source === "token" || source === "id_token";
3675+
},
3676+
);
3677+
vi.mocked(accountsModule.resolveRequestAccountId).mockImplementation(
3678+
(storedId: string | undefined, source: string | undefined, tokenId: string | undefined) => {
3679+
if (!storedId) return tokenId;
3680+
if (!accountsModule.shouldUpdateAccountIdFromToken(source, storedId)) {
3681+
return storedId;
3682+
}
3683+
return tokenId ?? storedId;
3684+
},
3685+
);
3686+
vi.mocked(refreshQueueModule.queuedRefresh)
3687+
.mockResolvedValueOnce({
3688+
type: "success",
3689+
access: "access-dead",
3690+
refresh: "shared-refresh",
3691+
expires: Date.now() + 300_000,
3692+
})
3693+
.mockResolvedValueOnce({
3694+
type: "success",
3695+
access: "access-live",
3696+
refresh: "shared-refresh",
3697+
expires: Date.now() + 300_000,
3698+
});
3699+
vi.mocked(fetchHelpersModule.createCodexHeaders).mockImplementation(
3700+
(_requestId, _accountId, accessToken) => {
3701+
const headers = new Headers();
3702+
if (typeof accessToken === "string" && accessToken.length > 0) {
3703+
headers.set("x-test-access-token", accessToken);
3704+
}
3705+
return headers;
3706+
},
3707+
);
3708+
globalThis.fetch = vi.fn(async (_url, init) => {
3709+
const headers = new Headers(init?.headers);
3710+
const accessToken = headers.get("x-test-access-token");
3711+
if (accessToken === "access-dead") {
3712+
return new Response(JSON.stringify({ error: { code: "deactivated_workspace", message: "workspace dead" } }), {
3713+
status: 402,
3714+
headers: { "content-type": "application/json" },
3715+
});
3716+
}
3717+
return new Response("", {
3718+
status: 200,
3719+
headers: {
3720+
"x-codex-primary-used-percent": "20",
3721+
"x-codex-primary-window-minutes": "180",
3722+
"x-codex-primary-reset-after-seconds": "900",
3723+
"x-codex-secondary-used-percent": "10",
3724+
"x-codex-secondary-window-minutes": "10080",
3725+
"x-codex-secondary-reset-after-seconds": "86400",
3726+
"x-codex-plan-type": "plus",
3727+
"x-codex-active-limit": "40",
3728+
},
3729+
});
3730+
});
3731+
3732+
const mockClient = createMockClient();
3733+
const { OpenAIOAuthPlugin } = await import("../index.js");
3734+
const plugin = (await OpenAIOAuthPlugin({
3735+
client: mockClient,
3736+
} as never)) as unknown as PluginType;
3737+
const autoMethod = plugin.auth.methods[0] as unknown as {
3738+
authorize: (inputs?: Record<string, string>) => Promise<{ instructions: string }>;
3739+
};
3740+
3741+
const authResult = await autoMethod.authorize();
3742+
expect(authResult.instructions).toBe("Authentication cancelled");
3743+
3744+
expect(globalThis.fetch).toHaveBeenCalledTimes(2);
3745+
expect(mockStorage.accounts).toHaveLength(1);
3746+
expect(mockStorage.accounts.some((account) => account.accountId === "workspace-dead")).toBe(false);
3747+
expect(mockStorage.accounts[0]?.accountId).toBe("workspace-live");
3748+
expect(vi.mocked(storageModule.withFlaggedAccountStorageTransaction)).toHaveBeenCalled();
3749+
expect(mockFlaggedStorage.accounts).toEqual(
3750+
expect.arrayContaining([
3751+
expect.objectContaining({
3752+
accountId: "workspace-dead",
3753+
organizationId: "org-shared",
3754+
flaggedReason: "workspace-deactivated",
3755+
lastError: "deactivated_workspace",
3756+
}),
3757+
]),
36313758
);
36323759
});
36333760
});

‎test/storage.test.ts‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import {
2020
importAccounts,
2121
previewImportAccounts,
2222
createTimestampedBackupPath,
23+
getWorkspaceIdentityKey,
2324
withAccountStorageTransaction,
2425
withFlaggedAccountStorageTransaction,
2526
} from "../lib/storage.js";
@@ -30,6 +31,37 @@ import {
3031
// But Task 0 says: "Tests should fail initially (RED phase)"
3132

3233
describe("storage", () => {
34+
describe("getWorkspaceIdentityKey", () => {
35+
it("uses organizationId and accountId when both are present", () => {
36+
expect(
37+
getWorkspaceIdentityKey({
38+
organizationId: " org-shared ",
39+
accountId: " workspace-a ",
40+
refreshToken: " refresh-a ",
41+
}),
42+
).toBe("organizationId:org-shared|accountId:workspace-a");
43+
});
44+
45+
it("falls back to accountId when organizationId is missing", () => {
46+
expect(
47+
getWorkspaceIdentityKey({
48+
accountId: " workspace-only ",
49+
refreshToken: " refresh-b ",
50+
}),
51+
).toBe("accountId:workspace-only");
52+
});
53+
54+
it("falls back to refreshToken when workspace ids are missing", () => {
55+
expect(
56+
getWorkspaceIdentityKey({
57+
organizationId: " ",
58+
accountId: "",
59+
refreshToken: " refresh-c ",
60+
}),
61+
).toBe("refreshToken:refresh-c");
62+
});
63+
});
64+
3365
describe("deduplication", () => {
3466
it("remaps activeIndex after deduplication using active account key", () => {
3567
const now = Date.now();

0 commit comments

Comments
 (0)