stagent 0.9.5 → 0.10.0

package/src/lib/agents/runtime/__tests__/openai-codex-auth.test.ts

@@ -0,0 +1,118 @@
+import { describe, expect, it, vi } from "vitest";
+
+import {
+  extractPlanTypeFromIdToken,
+  readCodexAuthStateFromClient,
+} from "@/lib/agents/runtime/openai-codex-auth";
+
+describe("openai codex auth", () => {
+  it("keeps a ChatGPT session connected when rate limit decoding fails", async () => {
+    const client = {
+      request: vi.fn(async (method: string) => {
+        if (method === "account/read") {
+          return {
+            account: {
+              type: "chatgpt",
+              email: "sehgal.manav@gmail.com",
+              planType: "prolite",
+            },
+            requiresOpenaiAuth: false,
+          };
+        }
+
+        if (method === "account/rateLimits/read") {
+          throw new Error("unknown variant `prolite`");
+        }
+
+        throw new Error(`Unexpected method: ${method}`);
+      }),
+    };
+
+    const state = await readCodexAuthStateFromClient(client as never, {
+      refreshToken: true,
+    });
+
+    expect(state.connected).toBe(true);
+    expect(state.account).toEqual({
+      type: "chatgpt",
+      email: "sehgal.manav@gmail.com",
+      planType: "prolite",
+    });
+    expect(state.rateLimits).toBeNull();
+    expect(state.authMode).toBe("chatgpt");
+  });
+
+  it("recovers the plan type from the rate limit error payload when account/read omits it", async () => {
+    const client = {
+      request: vi.fn(async (method: string) => {
+        if (method === "account/read") {
+          return {
+            account: {
+              type: "chatgpt",
+              email: "sehgal.manav@gmail.com",
+              planType: null,
+            },
+            requiresOpenaiAuth: false,
+          };
+        }
+
+        if (method === "account/rateLimits/read") {
+          throw new Error('Decode error: body={ "plan_type": "prolite" }');
+        }
+
+        throw new Error(`Unexpected method: ${method}`);
+      }),
+    };
+
+    const state = await readCodexAuthStateFromClient(client as never);
+
+    expect(state.connected).toBe(true);
+    expect(state.account?.planType).toBe("prolite");
+    expect(state.rateLimits).toBeNull();
+  });
+
+  it("treats account/read planType=unknown as missing and recovers the upstream plan", async () => {
+    const client = {
+      request: vi.fn(async (method: string) => {
+        if (method === "account/read") {
+          return {
+            account: {
+              type: "chatgpt",
+              email: "sehgal.manav@gmail.com",
+              planType: "unknown",
+            },
+            requiresOpenaiAuth: false,
+          };
+        }
+
+        if (method === "account/rateLimits/read") {
+          throw new Error('Decode error: body={ "plan_type": "prolite" }');
+        }
+
+        throw new Error(`Unexpected method: ${method}`);
+      }),
+    };
+
+    const state = await readCodexAuthStateFromClient(client as never);
+
+    expect(state.connected).toBe(true);
+    expect(state.account?.planType).toBe("prolite");
+    expect(state.rateLimits).toBeNull();
+  });
+
+  it("extracts the plan type from the stored id token payload", async () => {
+    const payload = Buffer.from(
+      JSON.stringify({
+        "https://api.openai.com/auth": {
+          chatgpt_plan_type: "prolite",
+        },
+      })
+    )
+      .toString("base64")
+      .replace(/\+/g, "-")
+      .replace(/\//g, "_")
+      .replace(/=+$/g, "");
+
+    expect(extractPlanTypeFromIdToken(`header.${payload}.signature`)).toBe("prolite");
+  });
+});

package/src/lib/agents/runtime/codex-app-server-client.ts

@@ -24,7 +24,7 @@ interface JsonRpcNotification {
 type JsonRpcMessage = JsonRpcRequest | JsonRpcResponse | JsonRpcNotification;
 
 export interface CodexAppServerClientOptions {
-  env?: Record<string, string>;
+  env?: Record<string, string | undefined>;
   cwd?: string;
 }
 
@@ -59,15 +59,21 @@ export class CodexAppServerClient {
   ): Promise<CodexAppServerClient> {
     const port = await reservePort();
     const listenUrl = `ws://127.0.0.1:${port}`;
+    const env: NodeJS.ProcessEnv = { ...process.env };
+    for (const [key, value] of Object.entries(options.env ?? {})) {
+      if (value === undefined) {
+        delete env[key];
+      } else {
+        env[key] = value;
+      }
+    }
+
     const child = spawn(
       "codex",
       ["app-server", "--listen", listenUrl],
       {
         cwd: options.cwd,
-        env: {
-          ...process.env,
-          ...(options.env ?? {}),
-        },
+        env,
         stdio: ["ignore", "pipe", "pipe"],
       }
     );

package/src/lib/agents/runtime/openai-codex-auth.ts

@@ -0,0 +1,389 @@
+import { mkdir, readFile, rm, writeFile } from "fs/promises";
+import { dirname } from "path";
+import { getLaunchCwd } from "@/lib/environment/workspace-context";
+import {
+  getOpenAIApiKey,
+  getOpenAIAuthSettings,
+  updateOpenAIAuthStatus,
+  clearOpenAIOAuthStatus,
+  updateOpenAIOAuthStatus,
+  type OpenAIAccountInfo,
+  type OpenAIAuthMode,
+  type OpenAIRateLimitInfo,
+  type OpenAIRateLimitWindow,
+} from "@/lib/settings/openai-auth";
+import {
+  getStagentCodexAuthPath,
+  getStagentCodexConfigPath,
+  getStagentCodexDir,
+} from "@/lib/utils/stagent-paths";
+import { CodexAppServerClient } from "./codex-app-server-client";
+
+const STAGENT_CODEX_CONFIG = `cli_auth_credentials_store = "file"
+`;
+
+interface AccountReadResult {
+  account?: {
+    type?: string;
+    email?: string | null;
+    planType?: string | null;
+  } | null;
+  requiresOpenaiAuth?: boolean;
+}
+
+interface RateLimitsReadResult {
+  rateLimits?: {
+    limitId?: string | null;
+    limitName?: string | null;
+    primary?: RateLimitWindowLike | null;
+    secondary?: RateLimitWindowLike | null;
+  } | null;
+}
+
+interface RateLimitWindowLike {
+  usedPercent?: unknown;
+  windowDurationMins?: unknown;
+  resetsAt?: unknown;
+}
+
+function parseRateLimitWindow(
+  value: RateLimitWindowLike | null | undefined
+): OpenAIRateLimitWindow | null {
+  if (!value || typeof value !== "object") return null;
+  return {
+    usedPercent:
+      typeof value.usedPercent === "number" ? value.usedPercent : null,
+    windowDurationMins:
+      typeof value.windowDurationMins === "number" ? value.windowDurationMins : null,
+    resetsAt: typeof value.resetsAt === "number" ? value.resetsAt : null,
+  };
+}
+
+function parseAccountInfo(
+  value: AccountReadResult["account"]
+): OpenAIAccountInfo | null {
+  if (!value?.type) return null;
+  if (
+    value.type !== "apiKey" &&
+    value.type !== "chatgpt" &&
+    value.type !== "chatgptAuthTokens"
+  ) {
+    return null;
+  }
+
+  return {
+    type: value.type,
+    email: value.email ?? null,
+    planType:
+      value.planType && value.planType.toLowerCase() !== "unknown"
+        ? value.planType
+        : null,
+  };
+}
+
+function extractPlanTypeFromError(error: unknown): string | null {
+  const message = error instanceof Error ? error.message : String(error);
+  const match = message.match(/"plan_type"\s*:\s*"([^"]+)"/);
+  return match?.[1] ?? null;
+}
+
+function decodeBase64Url(value: string): string | null {
+  try {
+    const normalized = value.replace(/-/g, "+").replace(/_/g, "/");
+    const padded = normalized.padEnd(Math.ceil(normalized.length / 4) * 4, "=");
+    return Buffer.from(padded, "base64").toString("utf8");
+  } catch {
+    return null;
+  }
+}
+
+export function extractPlanTypeFromIdToken(idToken: string): string | null {
+  const [, payload] = idToken.split(".");
+  if (!payload) return null;
+
+  const decoded = decodeBase64Url(payload);
+  if (!decoded) return null;
+
+  try {
+    const parsed = JSON.parse(decoded) as {
+      "https://api.openai.com/auth"?: {
+        chatgpt_plan_type?: string | null;
+      };
+    };
+    return parsed["https://api.openai.com/auth"]?.chatgpt_plan_type ?? null;
+  } catch {
+    return null;
+  }
+}
+
+async function readStagentCodexPlanTypeFromAuthFile(): Promise<string | null> {
+  try {
+    const raw = await readFile(getStagentCodexAuthPath(), "utf8");
+    const parsed = JSON.parse(raw) as {
+      tokens?: {
+        id_token?: string | null;
+      } | null;
+    };
+    const idToken = parsed.tokens?.id_token;
+    return idToken ? extractPlanTypeFromIdToken(idToken) : null;
+  } catch {
+    return null;
+  }
+}
+
+async function ensureCodexHomeConfig() {
+  const codexDir = getStagentCodexDir();
+  const configPath = getStagentCodexConfigPath();
+
+  await mkdir(codexDir, { recursive: true });
+  await mkdir(dirname(configPath), { recursive: true });
+
+  let current = "";
+  try {
+    current = await readFile(configPath, "utf8");
+  } catch {
+    // File will be created below.
+  }
+
+  if (current.includes('cli_auth_credentials_store = "file"')) {
+    return;
+  }
+
+  const next = current.trim().length > 0
+    ? `${current.trimEnd()}\n\n${STAGENT_CODEX_CONFIG}`
+    : STAGENT_CODEX_CONFIG;
+  await writeFile(configPath, next, "utf8");
+}
+
+export async function buildCodexAuthEnv(
+  env?: Record<string, string | undefined>
+): Promise<Record<string, string | undefined>> {
+  await ensureCodexHomeConfig();
+
+  return {
+    ...env,
+    CODEX_HOME: getStagentCodexDir(),
+    OPENAI_API_KEY: env?.OPENAI_API_KEY,
+  };
+}
+
+export async function connectStagentCodexClient(options: {
+  cwd?: string;
+  env?: Record<string, string | undefined>;
+} = {}) {
+  const env = await buildCodexAuthEnv(options.env);
+  return CodexAppServerClient.connect({
+    cwd: options.cwd ?? getLaunchCwd(),
+    env,
+  });
+}
+
+export async function initializeCodexClient(client: CodexAppServerClient) {
+  await client.request("initialize", {
+    clientInfo: {
+      name: "Stagent",
+      version: "0.1.1",
+    },
+    capabilities: null,
+  });
+}
+
+export async function readCodexAuthStateFromClient(
+  client: CodexAppServerClient,
+  options: { refreshToken?: boolean } = {}
+) {
+  const accountResult = (await client.request("account/read", {
+    refreshToken: options.refreshToken ?? false,
+  })) as AccountReadResult;
+
+  const account = parseAccountInfo(accountResult.account ?? null);
+  if (account?.type === "chatgpt" && !account.planType) {
+    account.planType = await readStagentCodexPlanTypeFromAuthFile();
+  }
+  let rateLimits: OpenAIRateLimitInfo | null = null;
+  if (account?.type === "chatgpt") {
+    try {
+      const rateLimitResult = (await client.request(
+        "account/rateLimits/read"
+      )) as RateLimitsReadResult;
+      const payload = rateLimitResult.rateLimits ?? null;
+      if (payload) {
+        rateLimits = {
+          limitId: payload.limitId ?? null,
+          limitName: payload.limitName ?? null,
+          primary: parseRateLimitWindow(payload.primary),
+          secondary: parseRateLimitWindow(payload.secondary),
+        };
+      }
+    } catch (error) {
+      if (!account.planType) {
+        account.planType = extractPlanTypeFromError(error);
+      }
+      rateLimits = null;
+    }
+  }
+
+  const authMode: OpenAIAuthMode =
+    account?.type === "apiKey"
+      ? "apikey"
+      : account?.type === "chatgpt"
+        ? "chatgpt"
+        : account?.type === "chatgptAuthTokens"
+          ? "chatgptAuthTokens"
+          : null;
+
+  return {
+    connected: account?.type === "chatgpt",
+    account,
+    rateLimits,
+    requiresOpenaiAuth: Boolean(accountResult.requiresOpenaiAuth),
+    authMode,
+  };
+}
+
+export type ResolvedOpenAICodexAuthContext =
+  | {
+      method: "api_key";
+      apiKeySource: "db" | "env" | "unknown";
+      connect: (cwd?: string) => Promise<CodexAppServerClient>;
+    }
+  | {
+      method: "oauth";
+      apiKeySource: "oauth";
+      connect: (cwd?: string) => Promise<CodexAppServerClient>;
+    };
+
+export async function resolveOpenAICodexAuthContext(): Promise<ResolvedOpenAICodexAuthContext> {
+  const settings = await getOpenAIAuthSettings();
+
+  if (settings.method === "oauth") {
+    if (!settings.oauthConnected) {
+      try {
+        const state = await readStagentCodexAuthState({ refreshToken: false });
+        if (!state.connected) {
+          throw new Error("OpenAI ChatGPT sign-in is not configured.");
+        }
+      } catch {
+        throw new Error(
+          "OpenAI ChatGPT sign-in is not configured. Sign in from Settings > Providers & Runtimes."
+        );
+      }
+    }
+
+    return {
+      method: "oauth",
+      apiKeySource: "oauth",
+      connect: (cwd?: string) =>
+        connectStagentCodexClient({
+          cwd,
+          env: { OPENAI_API_KEY: undefined },
+        }),
+    };
+  }
+
+  const { apiKey, source } = await getOpenAIApiKey();
+  if (!apiKey) {
+    throw new Error("OpenAI API key is not configured");
+  }
+
+  return {
+    method: "api_key",
+    apiKeySource: source,
+    connect: (cwd?: string) =>
+      connectStagentCodexClient({
+        cwd,
+        env: { OPENAI_API_KEY: apiKey },
+      }),
+  };
+}
+
+export async function ensureOpenAICodexClientAuthenticated(
+  client: CodexAppServerClient,
+  auth: ResolvedOpenAICodexAuthContext
+) {
+  await initializeCodexClient(client);
+
+  if (auth.method === "api_key") {
+    const { apiKey } = await getOpenAIApiKey();
+    if (!apiKey) {
+      throw new Error("OpenAI API key is not configured");
+    }
+    await client.request("account/login/start", {
+      type: "apiKey",
+      apiKey,
+    });
+    await updateOpenAIAuthStatus(auth.apiKeySource);
+    return;
+  }
+
+  const state = await readCodexAuthStateFromClient(client, {
+    refreshToken: true,
+  });
+  if (!state.connected || state.account?.type !== "chatgpt") {
+    throw new Error(
+      "OpenAI ChatGPT sign-in is not configured. Sign in from Settings > Providers & Runtimes."
+    );
+  }
+  await updateOpenAIOAuthStatus({
+    connected: true,
+    account: state.account,
+    authMode: state.authMode,
+    rateLimits: state.rateLimits,
+  });
+}
+
+export async function readStagentCodexAuthState(options: {
+  refreshToken?: boolean;
+  cwd?: string;
+} = {}) {
+  let client: CodexAppServerClient | null = null;
+
+  try {
+    client = await connectStagentCodexClient({ cwd: options.cwd });
+    await initializeCodexClient(client);
+
+    const state = await readCodexAuthStateFromClient(client, {
+      refreshToken: options.refreshToken,
+    });
+
+    await updateOpenAIOAuthStatus({
+      connected: state.connected,
+      account: state.account,
+      authMode: state.authMode,
+      rateLimits: state.rateLimits,
+    });
+
+    return state;
+  } catch (error) {
+    await clearOpenAIOAuthStatus();
+    throw error;
+  } finally {
+    if (client) {
+      await client.close();
+    }
+  }
+}
+
+export async function logoutStagentCodexAuth() {
+  let client: CodexAppServerClient | null = null;
+
+  try {
+    client = await connectStagentCodexClient();
+    await initializeCodexClient(client);
+    await client.request("account/logout");
+  } catch {
+    // Even if app-server logout fails, clear the isolated credential cache below.
+  } finally {
+    if (client) {
+      await client.close();
+    }
+  }
+
+  try {
+    await rm(getStagentCodexAuthPath(), { force: true });
+  } catch {
+    // Ignore cleanup failures.
+  }
+
+  await clearOpenAIOAuthStatus();
+}

package/src/lib/agents/runtime/openai-codex.ts

@@ -14,14 +14,15 @@ import {
   prepareTaskOutputDirectory,
   scanTaskOutputDocuments,
 } from "@/lib/documents/output-scanner";
-import {
-  getOpenAIApiKey,
-  updateOpenAIAuthStatus,
-} from "@/lib/settings/openai-auth";
 import { isToolAllowed } from "@/lib/settings/permissions";
 import { getRuntimeCatalogEntry } from "./catalog";
 import { getLaunchCwd } from "@/lib/environment/workspace-context";
 import { CodexAppServerClient } from "./codex-app-server-client";
+import {
+  ensureOpenAICodexClientAuthenticated,
+  readCodexAuthStateFromClient,
+  resolveOpenAICodexAuthContext,
+} from "./openai-codex-auth";
 import type {
   AgentRuntimeAdapter,
   RuntimeConnectionResult,
@@ -585,43 +586,19 @@ async function handleServerRequest(
   }
 }
 
-async function initializeOpenAIClient(
-  client: CodexAppServerClient,
-  apiKey: string
-) {
-  await client.request("initialize", {
-    clientInfo: {
-      name: "Stagent",
-      version: "0.1.1",
-    },
-    capabilities: null,
-  });
-
-  await client.request("account/login/start", {
-    type: "apiKey",
-    apiKey,
-  });
-}
-
 async function runAssistTurn({
   prompt,
   developerInstructions,
   cwd,
 }: AssistTurnOptions): Promise<{ text: string; usage: UsageSnapshot }> {
-  const { apiKey, source } = await getOpenAIApiKey();
-  if (!apiKey) {
-    throw new Error("OpenAI API key is not configured");
-  }
+  const auth = await resolveOpenAICodexAuthContext();
 
   let client: CodexAppServerClient | null = null;
   let text = "";
   let usage: UsageSnapshot = {};
 
   try {
-    client = await CodexAppServerClient.connect({
-      cwd,
-      env: { OPENAI_API_KEY: apiKey },
-    });
+    client = await auth.connect(cwd);
 
     client.onNotification = (notification: JsonRpcLikeNotification) => {
       if (notification.method !== "item/agentMessage/delta") return;
@@ -631,8 +608,7 @@ async function runAssistTurn({
       }
     };
 
-    await initializeOpenAIClient(client, apiKey);
-    await updateOpenAIAuthStatus(source);
+    await ensureOpenAICodexClientAuthenticated(client, auth);
 
     const threadResponse = (await client.request("thread/start", {
       cwd,
@@ -706,11 +682,7 @@ async function executeOpenAICodexTask(
 ): Promise<void> {
   const { task, profileId, instructions, prompt, cwd } =
     await resolveTaskExecutionContext(taskId, options);
-  const { apiKey, source } = await getOpenAIApiKey();
-
-  if (!apiKey) {
-    throw new Error("OpenAI API key is not configured");
-  }
+  const auth = await resolveOpenAICodexAuthContext();
 
   const abortController = new AbortController();
   let client: CodexAppServerClient | null = null;
@@ -729,10 +701,7 @@ async function executeOpenAICodexTask(
   };
 
   try {
-    client = await CodexAppServerClient.connect({
-      cwd,
-      env: { OPENAI_API_KEY: apiKey },
-    });
+    client = await auth.connect(cwd);
 
     client.onProcessError = (error) => {
       if (settled) return;
@@ -879,8 +848,7 @@ async function executeOpenAICodexTask(
       };
     });
 
-    await initializeOpenAIClient(client, apiKey);
-    await updateOpenAIAuthStatus(source);
+    await ensureOpenAICodexClientAuthenticated(client, auth);
 
     if (threadId) {
       await client.request("thread/resume", {
@@ -1051,29 +1019,30 @@ async function runOpenAITaskAssist(
 }
 
 async function testOpenAIConnection(): Promise<RuntimeConnectionResult> {
-  const { apiKey, source } = await getOpenAIApiKey();
-  if (!apiKey) {
-    return {
-      connected: false,
-      apiKeySource: "unknown",
-      error: "OpenAI API key is not configured",
-    };
-  }
-
   let client: CodexAppServerClient | null = null;
   try {
-    client = await CodexAppServerClient.connect({
-      cwd: getLaunchCwd(),
-      env: { OPENAI_API_KEY: apiKey },
+    const auth = await resolveOpenAICodexAuthContext();
+    client = await auth.connect(getLaunchCwd());
+    await ensureOpenAICodexClientAuthenticated(client, auth);
+    const accountState = await readCodexAuthStateFromClient(client, {
+      refreshToken: auth.apiKeySource === "oauth",
     });
-    await initializeOpenAIClient(client, apiKey);
-    await client.request("account/read", { refreshToken: false });
-    await updateOpenAIAuthStatus(source);
-    return { connected: true, apiKeySource: source };
+
+    return {
+      connected: auth.apiKeySource === "oauth" ? accountState.connected : true,
+      apiKeySource: auth.apiKeySource,
+      account: accountState.account,
+      rateLimits: accountState.rateLimits,
+      authMode: accountState.authMode,
+    };
   } catch (error) {
     return {
       connected: false,
-      apiKeySource: source,
+      apiKeySource:
+        error instanceof Error &&
+        error.message.includes("ChatGPT sign-in is not configured")
+          ? "oauth"
+          : "unknown",
       error: error instanceof Error ? error.message : String(error),
     };
   } finally {