Skip to content

fix(accounts): avoid team auth snapshot collisions #104

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
tinywind wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(accounts): avoid team auth snapshot collisions #104

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
79 changes: 54 additions & 25 deletions scripts/import-working-accounts-from-db.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ def iso_now() -> str:
return datetime.utcnow().replace(microsecond=0).isoformat() + "Z"


def parse_account_id(access_token: str) -> Optional[str]:
def parse_token_identity(access_token: str) -> Optional[dict]:
parts = access_token.split(".")
if len(parts) < 2:
return None
Expand All @@ -59,8 +59,14 @@ def parse_account_id(access_token: str) -> Optional[str]:
except Exception:
return None
auth_claims = claims.get("https://api.openai.com/auth") or {}
value = auth_claims.get("chatgpt_account_id") or auth_claims.get("user_id")
return value.strip() if isinstance(value, str) and value.strip() else None
account_id = auth_claims.get("chatgpt_account_id") or auth_claims.get("user_id")
user_id = auth_claims.get("user_id")
if not isinstance(account_id, str) or not account_id.strip():
return None
return {
"account_id": account_id.strip(),
"user_id": user_id.strip() if isinstance(user_id, str) and user_id.strip() else None,
}


def exchange_refresh_token(refresh_token: str) -> Optional[dict]:
Expand Down Expand Up @@ -89,56 +95,71 @@ def exchange_refresh_token(refresh_token: str) -> Optional[dict]:
if not isinstance(access_token, str) or not access_token.strip():
return None

account_id = parse_account_id(access_token)
if not account_id:
identity = parse_token_identity(access_token)
if identity is None:
return None

return {
"access_token": access_token,
"refresh_token": payload.get("refresh_token") or refresh_token,
"id_token": payload.get("id_token"),
"account_id": account_id,
"account_id": identity["account_id"],
"user_id": identity["user_id"],
}


def storage_id(account_id: str) -> str:
return hashlib.sha256(account_id.encode("utf-8")).hexdigest()
def storage_id(account_id: str, user_id: Optional[str]) -> str:
key = f"{account_id}\0{user_id}" if user_id else account_id
return hashlib.sha256(key.encode("utf-8")).hexdigest()


def load_state() -> dict:
if not STATE_PATH.exists():
return {"activeAccountId": None, "accounts": []}
return {"activeAccountId": None, "activeStorageId": None, "accounts": []}
try:
parsed = json.loads(STATE_PATH.read_text(encoding="utf-8"))
except Exception:
return {"activeAccountId": None, "accounts": []}
return {"activeAccountId": None, "activeStorageId": None, "accounts": []}

accounts = parsed.get("accounts")
if not isinstance(accounts, list):
accounts = []
active = parsed.get("activeAccountId")
if not isinstance(active, str):
active = None
return {"activeAccountId": active, "accounts": accounts}
active_storage = parsed.get("activeStorageId")
if not isinstance(active_storage, str):
active_storage = None
return {"activeAccountId": active, "activeStorageId": active_storage, "accounts": accounts}


def save_state(state: dict) -> None:
STATE_PATH.write_text(json.dumps(state, indent=2), encoding="utf-8")
os.chmod(STATE_PATH, 0o600)


def upsert_account(state: dict, account_id: str, email: Optional[str], auth_mode: str = "chatgpt") -> None:
sid = storage_id(account_id)
def upsert_account(state: dict, account_id: str, user_id: Optional[str], email: Optional[str], auth_mode: str = "chatgpt") -> None:
sid = storage_id(account_id, user_id)
legacy_sid = storage_id(account_id, None)
now = iso_now()
existing = None
for item in state["accounts"]:
if isinstance(item, dict) and item.get("accountId") == account_id:
if isinstance(item, dict) and item.get("storageId") == sid:
existing = item
break
if (
sid != legacy_sid
and isinstance(item, dict)
and item.get("storageId") == legacy_sid
and item.get("accountId") == account_id
and item.get("userId") is None
):
existing = item

entry = {
"accountId": account_id,
"storageId": sid,
"userId": user_id,
"authMode": auth_mode,
"email": email,
"planType": existing.get("planType") if isinstance(existing, dict) else None,
Expand All @@ -151,14 +172,20 @@ def upsert_account(state: dict, account_id: str, email: Optional[str], auth_mode
"unavailableReason": existing.get("unavailableReason") if isinstance(existing, dict) else None,
}

removed_storage_ids = {sid}
if isinstance(existing, dict) and isinstance(existing.get("storageId"), str):
removed_storage_ids.add(existing["storageId"])
if state.get("activeStorageId") == existing["storageId"]:
state["activeStorageId"] = sid

state["accounts"] = [
entry,
*[x for x in state["accounts"] if not (isinstance(x, dict) and x.get("accountId") == account_id)],
*[x for x in state["accounts"] if not (isinstance(x, dict) and x.get("storageId") in removed_storage_ids)],
]


def write_snapshot(account_id: str, token_payload: dict) -> None:
sid = storage_id(account_id)
def write_snapshot(account_id: str, user_id: Optional[str], token_payload: dict) -> None:
sid = storage_id(account_id, user_id)
dest_dir = SNAP_ROOT / sid
dest_dir.mkdir(parents=True, exist_ok=True)
os.chmod(dest_dir, 0o700)
Expand Down Expand Up @@ -190,10 +217,10 @@ rows = conn.execute(
conn.close()

state = load_state()
existing_ids = {
item.get("accountId")
existing_storage_ids = {
item.get("storageId")
for item in state["accounts"]
if isinstance(item, dict) and isinstance(item.get("accountId"), str)
if isinstance(item, dict) and isinstance(item.get("storageId"), str)
}

imported = 0
Expand All @@ -207,13 +234,15 @@ for row_id, email, refresh_token in rows:
if token_payload is None:
continue
account_id = token_payload["account_id"]
if account_id in existing_ids:
user_id = token_payload.get("user_id")
sid = storage_id(account_id, user_id)
if sid in existing_storage_ids:
continue
write_snapshot(account_id, token_payload)
upsert_account(state, account_id=account_id, email=email)
existing_ids.add(account_id)
write_snapshot(account_id, user_id, token_payload)
upsert_account(state, account_id=account_id, user_id=user_id, email=email)
existing_storage_ids.add(sid)
imported += 1
print(f"imported row_id={row_id} email={email} account_id={account_id}")
print(f"imported row_id={row_id} email={email} account_id={account_id} user_id={user_id}")

save_state(state)
print(f"done imported={imported} attempted={attempted} limit={LIMIT}")
Expand Down
40 changes: 20 additions & 20 deletions src/App.vue
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -122,7 +122,7 @@
<div v-else class="sidebar-settings-account-list">
<article
v-for="account in accounts"
:key="account.accountId"
:key="account.storageId"
class="sidebar-settings-account-item"
:class="{
'is-active': account.isActive,
Expand All @@ -131,8 +131,8 @@
'is-remove-visible': isRemoveVisible(account),
}"
:title="buildAccountTitle(account)"
@mouseenter="onAccountCardPointerEnter(account.accountId)"
@mouseleave="onAccountCardPointerLeave(account.accountId)"
@mouseenter="onAccountCardPointerEnter(account.storageId)"
@mouseleave="onAccountCardPointerLeave(account.storageId)"
>
<div class="sidebar-settings-account-main">
<p class="sidebar-settings-account-email">{{ account.email || t('Account') }}</p>
Expand All @@ -151,7 +151,7 @@
class="sidebar-settings-account-switch"
type="button"
:disabled="isAccountActionDisabled(account) || account.isActive || isAccountUnavailable(account)"
@click="onSwitchAccount(account.accountId)"
@click="onSwitchAccount(account.storageId)"
>
{{ getAccountSwitchLabel(account) }}
</button>
Expand All @@ -163,7 +163,7 @@
}"
type="button"
:disabled="isAccountActionDisabled(account)"
@click="onRemoveAccount(account.accountId)"
@click="onRemoveAccount(account.storageId)"
>
{{ getAccountRemoveLabel(account) }}
</button>
Expand Down Expand Up @@ -1777,15 +1777,15 @@ function isAccountUnavailable(account: UiAccountEntry): boolean {

function isAccountActionDisabled(account: UiAccountEntry): boolean {
return isRefreshingAccounts.value || isSwitchingAccounts.value || removingAccountId.value.length > 0
|| (account.isActive && removingAccountId.value !== account.accountId && isAccountSwitchBlocked.value)
|| (account.isActive && removingAccountId.value !== account.storageId && isAccountSwitchBlocked.value)
}

function isRemoveConfirmationActive(account: UiAccountEntry): boolean {
return confirmingRemoveAccountId.value === account.accountId
return confirmingRemoveAccountId.value === account.storageId
}

function isRemoveVisible(account: UiAccountEntry): boolean {
return hoveredAccountId.value === account.accountId || isRemoveConfirmationActive(account)
return hoveredAccountId.value === account.storageId || isRemoveConfirmationActive(account)
}

function getAccountSwitchLabel(account: UiAccountEntry): string {
Expand All @@ -1796,7 +1796,7 @@ function getAccountSwitchLabel(account: UiAccountEntry): string {
}

function getAccountRemoveLabel(account: UiAccountEntry): string {
if (removingAccountId.value === account.accountId) return t('Removing…')
if (removingAccountId.value === account.storageId) return t('Removing…')
if (isRemoveConfirmationActive(account)) return t('Click again to remove')
return t('Remove')
}
Expand Down Expand Up @@ -1885,10 +1885,10 @@ async function loadAccountsState(options: { silent?: boolean } = {}): Promise<vo
try {
const result = await getAccounts()
accounts.value = result.accounts
if (!result.accounts.some((account) => account.accountId === hoveredAccountId.value)) {
if (!result.accounts.some((account) => account.storageId === hoveredAccountId.value)) {
hoveredAccountId.value = ''
}
if (!result.accounts.some((account) => account.accountId === confirmingRemoveAccountId.value)) {
if (!result.accounts.some((account) => account.storageId === confirmingRemoveAccountId.value)) {
confirmingRemoveAccountId.value = ''
}
} catch (error) {
Expand Down Expand Up @@ -1918,7 +1918,7 @@ async function onRefreshAccounts(): Promise<void> {
}
}

async function onSwitchAccount(accountId: string): Promise<void> {
async function onSwitchAccount(storageId: string): Promise<void> {
if (isSwitchingAccounts.value || isRefreshingAccounts.value) return
if (isAccountSwitchBlocked.value) {
accountActionError.value = t('Finish the current turn and pending requests before switching accounts.')
Expand All @@ -1929,9 +1929,9 @@ async function onSwitchAccount(accountId: string): Promise<void> {
confirmingRemoveAccountId.value = ''
isSwitchingAccounts.value = true
try {
const nextActiveAccount = await switchAccount(accountId)
const nextActiveAccount = await switchAccount(storageId)
accounts.value = accounts.value.map((account) => (
account.accountId === accountId
account.storageId === storageId
? nextActiveAccount
: { ...account, isActive: false }
))
Expand All @@ -1948,12 +1948,12 @@ async function onSwitchAccount(accountId: string): Promise<void> {
}
}

async function onRemoveAccount(accountId: string): Promise<void> {
async function onRemoveAccount(storageId: string): Promise<void> {
if (isRefreshingAccounts.value || isSwitchingAccounts.value || removingAccountId.value.length > 0) return
const targetAccount = accounts.value.find((account) => account.accountId === accountId) ?? null
const targetAccount = accounts.value.find((account) => account.storageId === storageId) ?? null
if (!targetAccount) return
if (confirmingRemoveAccountId.value !== accountId) {
confirmingRemoveAccountId.value = accountId
if (confirmingRemoveAccountId.value !== storageId) {
confirmingRemoveAccountId.value = storageId
return
}
if (targetAccount.isActive && isAccountSwitchBlocked.value) {
Expand All @@ -1964,9 +1964,9 @@ async function onRemoveAccount(accountId: string): Promise<void> {
const removedWasActive = targetAccount.isActive
accountActionError.value = ''
confirmingRemoveAccountId.value = ''
removingAccountId.value = accountId
removingAccountId.value = storageId
try {
const result = await removeAccount(accountId)
const result = await removeAccount(storageId)
accounts.value = result.accounts
stopPolling()
startPolling()
Expand Down
28 changes: 20 additions & 8 deletions src/api/codexGateway.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -364,8 +364,10 @@ export type ThreadTerminalQuickCommand = {

export type AccountsListResult = {
activeAccountId: string | null
activeStorageId: string | null
accounts: UiAccountEntry[]
importedAccountId?: string
importedStorageId?: string
}

type ThreadFileChangeFallbackEntry = {
Expand Down Expand Up @@ -459,16 +461,23 @@ function normalizeRateLimitSnapshot(value: unknown): UiRateLimitSnapshot | null
}
}

function normalizeAccountEntry(value: unknown, activeAccountId: string | null = null): UiAccountEntry | null {
function normalizeAccountEntry(
value: unknown,
activeAccountId: string | null = null,
activeStorageId: string | null = null,
): UiAccountEntry | null {
const record = asRecord(value)
if (!record) return null
const accountId = readString(record.accountId)
const storageId = readString(record.storageId) ?? accountId
const quotaStatusRaw = readString(record.quotaStatus)
const quotaStatus: UiAccountQuotaStatus =
quotaStatusRaw === 'loading' || quotaStatusRaw === 'ready' || quotaStatusRaw === 'error' ? quotaStatusRaw : 'idle'
if (!accountId) return null
return {
accountId,
storageId: storageId ?? accountId,
userId: readString(record.userId),
authMode: readString(record.authMode),
email: readString(record.email),
planType: readString(record.planType),
Expand All @@ -480,7 +489,7 @@ function normalizeAccountEntry(value: unknown, activeAccountId: string | null =
quotaError: readString(record.quotaError),
unavailableReason: normalizeAccountUnavailableReason(record.unavailableReason)
?? (isPaymentRequiredErrorMessage(readString(record.quotaError)) ? 'payment_required' : null),
isActive: readBoolean(record.isActive) ?? accountId === activeAccountId,
isActive: readBoolean(record.isActive) ?? (storageId === activeStorageId || accountId === activeAccountId),
}
}

Expand Down Expand Up @@ -1175,12 +1184,15 @@ export async function getAccountRateLimits(): Promise<UiRateLimitSnapshot | null
function normalizeAccountsListResult(payload: unknown): AccountsListResult {
const record = asRecord(payload)
const activeAccountId = readString(record?.activeAccountId)
const activeStorageId = readString(record?.activeStorageId)
const data = Array.isArray(record?.accounts) ? record?.accounts : []
return {
activeAccountId,
activeStorageId,
importedAccountId: readString(record?.importedAccountId) ?? undefined,
importedStorageId: readString(record?.importedStorageId) ?? undefined,
accounts: data
.map((entry) => normalizeAccountEntry(entry, activeAccountId))
.map((entry) => normalizeAccountEntry(entry, activeAccountId, activeStorageId))
.filter((entry): entry is UiAccountEntry => entry !== null),
}
}
Expand All @@ -1207,30 +1219,30 @@ export async function refreshAccountsFromAuth(): Promise<AccountsListResult> {
return normalizeAccountsListResult(envelope?.data)
}

export async function switchAccount(accountId: string): Promise<UiAccountEntry> {
export async function switchAccount(storageId: string): Promise<UiAccountEntry> {
const response = await fetch('/codex-api/accounts/switch', {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ accountId }),
body: JSON.stringify({ storageId }),
})
const payload = (await response.json()) as unknown
if (!response.ok) {
throw new Error(getErrorMessageFromPayload(payload, 'Failed to switch account'))
}
const envelope = asRecord(payload)
const data = asRecord(envelope?.data)
const account = normalizeAccountEntry(data?.account, readString(data?.activeAccountId))
const account = normalizeAccountEntry(data?.account, readString(data?.activeAccountId), readString(data?.activeStorageId))
if (!account) {
throw new Error('Failed to switch account')
}
return account
}

export async function removeAccount(accountId: string): Promise<AccountsListResult> {
export async function removeAccount(storageId: string): Promise<AccountsListResult> {
const response = await fetch('/codex-api/accounts/remove', {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ accountId }),
body: JSON.stringify({ storageId }),
})
const payload = (await response.json()) as unknown
if (!response.ok) {
Expand Down
Loading