stagent 0.9.6 → 0.11.0

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,21 @@ 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 {
+  classifyTaskFailureReason,
+  RetryableRuntimeLaunchError,
+  toRetryableRuntimeLaunchError,
+  type RuntimeLaunchProgress,
+} from "./launch-failure";
 import type {
   AgentRuntimeAdapter,
   RuntimeConnectionResult,
@@ -227,6 +234,14 @@ async function finalizeTaskUsage(
     startedAt: state.startedAt,
     finishedAt: new Date(),
   });
+
+  await db
+    .update(tasks)
+    .set({
+      effectiveModelId: state.modelId ?? null,
+      updatedAt: new Date(),
+    })
+    .where(eq(tasks.id, state.taskId));
 }
 
 function buildTurnInput(prompt: string) {
@@ -372,6 +387,7 @@ async function markTaskFailed(taskId: string, title: string, message: string) {
     .set({
       status: "failed",
       result: message,
+      failureReason: classifyTaskFailureReason(new Error(message)),
       updatedAt: new Date(),
     })
     .where(eq(tasks.id, taskId));
@@ -585,43 +601,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 +623,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 +697,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;
@@ -718,6 +705,7 @@ async function executeOpenAICodexTask(
   let turnId: string | null = null;
   let agentOutput = "";
   let settled = false;
+  const launchProgress: RuntimeLaunchProgress = {};
   let resolveCompletion: (() => void) | null = null;
   let rejectCompletion: ((error: Error) => void) | null = null;
   const usageState = createTaskUsageState(task, Boolean(task.sessionId));
@@ -729,12 +717,21 @@ async function executeOpenAICodexTask(
   };
 
   try {
-    client = await CodexAppServerClient.connect({
-      cwd,
-      env: { OPENAI_API_KEY: apiKey },
-    });
+    client = await auth.connect(cwd);
 
     client.onProcessError = (error) => {
+      const retryableLaunchError =
+        !options.resume
+          ? toRetryableRuntimeLaunchError({
+              runtimeId: "openai-codex-app-server",
+              error,
+              progress: launchProgress,
+            })
+          : null;
+      if (retryableLaunchError) {
+        rejectCompletion?.(retryableLaunchError);
+        return;
+      }
       if (settled) return;
       void settle(async () => {
         await markTaskFailed(taskId, task.title, error.message);
@@ -775,6 +772,7 @@ async function executeOpenAICodexTask(
 
             case "turn/started": {
               turnId = extractTurnId(params);
+              launchProgress.hasTurnStarted = true;
               setExecution(taskId, {
                 abortController,
                 sessionId: threadId,
@@ -813,6 +811,7 @@ async function executeOpenAICodexTask(
             }
 
             case "item/commandExecution/outputDelta": {
+              launchProgress.hasToolUse = true;
               await insertLog(taskId, "command_output_delta", {
                 threadId,
                 turnId,
@@ -834,6 +833,7 @@ async function executeOpenAICodexTask(
 
             case "turn/completed": {
               const { status, errorMessage } = extractTurnStatus(params);
+              launchProgress.hasResult = true;
 
               if (status === "completed") {
                 const finalResult =
@@ -879,8 +879,7 @@ async function executeOpenAICodexTask(
       };
     });
 
-    await initializeOpenAIClient(client, apiKey);
-    await updateOpenAIAuthStatus(source);
+    await ensureOpenAICodexClientAuthenticated(client, auth);
 
     if (threadId) {
       await client.request("thread/resume", {
@@ -949,6 +948,10 @@ async function executeOpenAICodexTask(
       return;
     }
 
+    if (error instanceof RetryableRuntimeLaunchError) {
+      throw error;
+    }
+
     const message = error instanceof Error ? error.message : String(error);
     await settle(async () => {
       await markTaskFailed(taskId, task.title, message);
@@ -1051,29 +1054,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 {

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

@@ -396,6 +396,14 @@ async function executeOpenAIDirectTask(taskId: string, isResume = false): Promis
       startedAt: usageState.startedAt,
       finishedAt: new Date(),
     });
+
+    await db
+      .update(tasks)
+      .set({
+        effectiveModelId: result.totalUsage.modelId ?? modelId,
+        updatedAt: new Date(),
+      })
+      .where(eq(tasks.id, taskId));
   } catch (err) {
     if (!abortController.signal.aborted) {
       const errorMsg = err instanceof Error ? err.message : String(err);

package/src/lib/agents/runtime/types.ts

@@ -3,10 +3,18 @@ import type { ProfileTestReport } from "@/lib/agents/profiles/test-types";
 import type { RuntimeCapabilities, RuntimeCatalogEntry } from "./catalog";
 import type { TaskAssistResponse } from "./task-assist-types";
 import type { ProfileAssistRequest, ProfileAssistResponse } from "./profile-assist-types";
+import type {
+  OpenAIAccountInfo,
+  OpenAIAuthMode,
+  OpenAIRateLimitInfo,
+} from "@/lib/settings/openai-auth";
 
 export interface RuntimeConnectionResult {
   connected: boolean;
   apiKeySource?: ApiKeySource;
+  account?: OpenAIAccountInfo | null;
+  rateLimits?: OpenAIRateLimitInfo | null;
+  authMode?: OpenAIAuthMode;
   error?: string;
 }