stagent 0.9.6 → 0.11.0

package/src/lib/settings/__tests__/runtime-setup.test.ts

@@ -0,0 +1,33 @@
+import { describe, expect, it, vi } from "vitest";
+
+vi.mock("../auth", () => ({
+  getAuthSettings: vi.fn(async () => ({
+    method: "oauth",
+    hasKey: false,
+    apiKeySource: "oauth",
+  })),
+}));
+
+vi.mock("../openai-auth", () => ({
+  getOpenAIAuthSettings: vi.fn(async () => ({
+    method: "oauth",
+    hasKey: true,
+    apiKeySource: "db",
+    oauthConnected: true,
+    account: { type: "chatgpt", email: "dev@example.com", planType: "pro" },
+    rateLimits: null,
+  })),
+}));
+
+describe("runtime setup states", () => {
+  it("marks Codex App Server subscription-backed when ChatGPT auth is connected", async () => {
+    const { getRuntimeSetupStates } = await import("../runtime-setup");
+    const states = await getRuntimeSetupStates();
+
+    expect(states["openai-codex-app-server"].configured).toBe(true);
+    expect(states["openai-codex-app-server"].authMethod).toBe("oauth");
+    expect(states["openai-codex-app-server"].billingMode).toBe("subscription");
+    expect(states["openai-direct"].configured).toBe(true);
+    expect(states["openai-direct"].billingMode).toBe("usage");
+  });
+});

package/src/lib/settings/openai-auth.ts

@@ -1,26 +1,77 @@
-import { SETTINGS_KEYS, type ApiKeySource } from "@/lib/constants/settings";
+import { existsSync } from "node:fs";
+import { SETTINGS_KEYS, type ApiKeySource, type AuthMethod } from "@/lib/constants/settings";
 import { decrypt, encrypt } from "@/lib/utils/crypto";
+import { getStagentCodexAuthPath } from "@/lib/utils/stagent-paths";
 import { getSetting, setSetting } from "./helpers";
 
+export type OpenAIAccountType = "apiKey" | "chatgpt" | "chatgptAuthTokens";
+export type OpenAIAuthMode = "apikey" | "chatgpt" | "chatgptAuthTokens" | null;
+
+export interface OpenAIAccountInfo {
+  type: OpenAIAccountType;
+  email?: string | null;
+  planType?: string | null;
+}
+
+export interface OpenAIRateLimitWindow {
+  usedPercent: number | null;
+  windowDurationMins: number | null;
+  resetsAt: number | null;
+}
+
+export interface OpenAIRateLimitInfo {
+  limitId: string | null;
+  limitName: string | null;
+  primary: OpenAIRateLimitWindow | null;
+  secondary: OpenAIRateLimitWindow | null;
+}
+
 export interface OpenAIAuthSettings {
+  method: AuthMethod;
   hasKey: boolean;
-  apiKeySource: ApiKeySource;
+  apiKeySource: Exclude<ApiKeySource, "oauth">;
+  oauthConnected: boolean;
+  account: OpenAIAccountInfo | null;
+  rateLimits: OpenAIRateLimitInfo | null;
 }
 
 export interface OpenAIAuthConfigInput {
-  apiKey: string;
+  method: AuthMethod;
+  apiKey?: string;
+}
+
+interface PersistedOpenAIAccountPayload {
+  account: OpenAIAccountInfo | null;
+  authMode?: OpenAIAuthMode;
+}
+
+function parseJson<T>(raw: string | null): T | null {
+  if (!raw) return null;
+  try {
+    return JSON.parse(raw) as T;
+  } catch {
+    return null;
+  }
 }
 
 export async function getOpenAIAuthSettings(): Promise<OpenAIAuthSettings> {
+  const method = ((await getSetting(SETTINGS_KEYS.OPENAI_AUTH_METHOD)) as AuthMethod) ?? "api_key";
   const encryptedKey = await getSetting(SETTINGS_KEYS.OPENAI_AUTH_API_KEY);
   const storedSource = (await getSetting(
     SETTINGS_KEYS.OPENAI_AUTH_API_KEY_SOURCE
-  )) as ApiKeySource | null;
+  )) as Exclude<ApiKeySource, "oauth"> | null;
+  const storedOauthConnected = await getSetting(SETTINGS_KEYS.OPENAI_AUTH_OAUTH_CONNECTED);
+  const storedAccount = parseJson<PersistedOpenAIAccountPayload>(
+    await getSetting(SETTINGS_KEYS.OPENAI_AUTH_ACCOUNT)
+  );
+  const storedRateLimits = parseJson<OpenAIRateLimitInfo>(
+    await getSetting(SETTINGS_KEYS.OPENAI_AUTH_RATE_LIMITS)
+  );
 
   const hasDbKey = encryptedKey !== null;
   const hasEnvKey = !!process.env.OPENAI_API_KEY;
 
-  let apiKeySource: ApiKeySource;
+  let apiKeySource: Exclude<ApiKeySource, "oauth">;
   if (storedSource) {
     apiKeySource = storedSource;
   } else if (hasDbKey) {
@@ -31,20 +82,32 @@ export async function getOpenAIAuthSettings(): Promise<OpenAIAuthSettings> {
     apiKeySource = "unknown";
   }
 
+  const oauthConnected =
+    storedOauthConnected === "true" ||
+    (storedOauthConnected == null && existsSync(getStagentCodexAuthPath()));
+
   return {
+    method,
     hasKey: hasDbKey || hasEnvKey,
     apiKeySource,
+    oauthConnected,
+    account: storedAccount?.account ?? null,
+    rateLimits: storedRateLimits,
   };
 }
 
 export async function setOpenAIAuthSettings(
   input: OpenAIAuthConfigInput
 ): Promise<void> {
-  await setSetting(
-    SETTINGS_KEYS.OPENAI_AUTH_API_KEY,
-    encrypt(input.apiKey)
-  );
-  await setSetting(SETTINGS_KEYS.OPENAI_AUTH_API_KEY_SOURCE, "db");
+  await setSetting(SETTINGS_KEYS.OPENAI_AUTH_METHOD, input.method);
+
+  if (input.apiKey) {
+    await setSetting(
+      SETTINGS_KEYS.OPENAI_AUTH_API_KEY,
+      encrypt(input.apiKey)
+    );
+    await setSetting(SETTINGS_KEYS.OPENAI_AUTH_API_KEY_SOURCE, "db");
+  }
 }
 
 export async function getOpenAIApiKey(): Promise<{
@@ -78,3 +141,35 @@ export async function updateOpenAIAuthStatus(
 ): Promise<void> {
   await setSetting(SETTINGS_KEYS.OPENAI_AUTH_API_KEY_SOURCE, source);
 }
+
+export async function updateOpenAIOAuthStatus(input: {
+  connected: boolean;
+  account?: OpenAIAccountInfo | null;
+  authMode?: OpenAIAuthMode;
+  rateLimits?: OpenAIRateLimitInfo | null;
+}): Promise<void> {
+  await setSetting(
+    SETTINGS_KEYS.OPENAI_AUTH_OAUTH_CONNECTED,
+    input.connected ? "true" : "false"
+  );
+  await setSetting(
+    SETTINGS_KEYS.OPENAI_AUTH_ACCOUNT,
+    JSON.stringify({
+      account: input.account ?? null,
+      authMode: input.authMode ?? null,
+    } satisfies PersistedOpenAIAccountPayload)
+  );
+  await setSetting(
+    SETTINGS_KEYS.OPENAI_AUTH_RATE_LIMITS,
+    JSON.stringify(input.rateLimits ?? null)
+  );
+}
+
+export async function clearOpenAIOAuthStatus(): Promise<void> {
+  await updateOpenAIOAuthStatus({
+    connected: false,
+    account: null,
+    authMode: null,
+    rateLimits: null,
+  });
+}

package/src/lib/settings/openai-login-manager.ts

@@ -0,0 +1,260 @@
+import {
+  connectStagentCodexClient,
+  initializeCodexClient,
+  readStagentCodexAuthState,
+} from "@/lib/agents/runtime/openai-codex-auth";
+import { clearOpenAIOAuthStatus } from "./openai-auth";
+import type { CodexAppServerClient } from "@/lib/agents/runtime/codex-app-server-client";
+import type {
+  OpenAIAccountInfo,
+  OpenAIRateLimitInfo,
+} from "./openai-auth";
+
+type LoginPhase =
+  | "idle"
+  | "pending"
+  | "connected"
+  | "cancelled"
+  | "failed";
+
+export interface OpenAILoginState {
+  phase: LoginPhase;
+  loginId: string | null;
+  authUrl: string | null;
+  account: OpenAIAccountInfo | null;
+  rateLimits: OpenAIRateLimitInfo | null;
+  error: string | null;
+  startedAt: string | null;
+  updatedAt: string;
+}
+
+interface LoginAttempt {
+  client: CodexAppServerClient;
+  state: OpenAILoginState;
+  cancelRequested: boolean;
+}
+
+const idleState = (): OpenAILoginState => ({
+  phase: "idle",
+  loginId: null,
+  authUrl: null,
+  account: null,
+  rateLimits: null,
+  error: null,
+  startedAt: null,
+  updatedAt: new Date().toISOString(),
+});
+
+let activeAttempt: LoginAttempt | null = null;
+let lastKnownState: OpenAILoginState = idleState();
+
+function updateState(
+  next: Partial<OpenAILoginState> & Pick<OpenAILoginState, "phase">
+): OpenAILoginState {
+  lastKnownState = {
+    ...lastKnownState,
+    ...next,
+    updatedAt: new Date().toISOString(),
+  };
+  if (next.phase === "idle") {
+    lastKnownState = idleState();
+  }
+  return lastKnownState;
+}
+
+async function closeAttempt() {
+  const attempt = activeAttempt;
+  activeAttempt = null;
+  if (attempt) {
+    await attempt.client.close();
+  }
+}
+
+export function getOpenAILoginState(): OpenAILoginState {
+  return activeAttempt?.state ?? lastKnownState;
+}
+
+export async function startOpenAIChatGPTLogin(): Promise<OpenAILoginState> {
+  if (activeAttempt) {
+    return activeAttempt.state;
+  }
+
+  try {
+    const current = await readStagentCodexAuthState({ refreshToken: false });
+    if (current.connected) {
+      return updateState({
+        phase: "connected",
+        loginId: null,
+        authUrl: null,
+        account: current.account,
+        rateLimits: current.rateLimits,
+        error: null,
+        startedAt: null,
+      });
+    }
+  } catch {
+    await clearOpenAIOAuthStatus();
+  }
+
+  const client = await connectStagentCodexClient();
+  await initializeCodexClient(client);
+
+  const startedAt = new Date().toISOString();
+  const state = updateState({
+    phase: "pending",
+    loginId: null,
+    authUrl: null,
+    account: null,
+    rateLimits: null,
+    error: null,
+    startedAt,
+  });
+  const attempt: LoginAttempt = { client, state, cancelRequested: false };
+  activeAttempt = attempt;
+
+  client.onProcessError = (error) => {
+    void closeAttempt();
+    updateState({
+      phase: "failed",
+      error: error.message,
+      loginId: state.loginId,
+      authUrl: state.authUrl,
+      account: null,
+      rateLimits: null,
+      startedAt,
+    });
+  };
+
+  client.onNotification = (notification) => {
+    if (notification.method !== "account/login/completed") return;
+    const params = notification.params as
+      | { loginId?: string | null; success?: boolean; error?: string | null }
+      | undefined;
+
+    const completedLoginId = params?.loginId ?? state.loginId;
+    if (params?.success) {
+      void (async () => {
+        try {
+          const current = await readStagentCodexAuthState({ refreshToken: true });
+          updateState({
+            phase: "connected",
+            loginId: completedLoginId ?? null,
+            authUrl: state.authUrl,
+            account: current.account,
+            rateLimits: current.rateLimits,
+            error: null,
+            startedAt,
+          });
+        } catch (error) {
+          updateState({
+            phase: "failed",
+            loginId: completedLoginId ?? null,
+            authUrl: state.authUrl,
+            account: null,
+            rateLimits: null,
+            error: error instanceof Error ? error.message : String(error),
+            startedAt,
+          });
+        } finally {
+          await closeAttempt();
+        }
+      })();
+      return;
+    }
+
+    void (async () => {
+      const cancelled =
+        attempt.cancelRequested ||
+        params?.error?.toLowerCase().includes("cancel") ||
+        params?.error === "Login was not completed";
+      updateState({
+        phase: cancelled ? "cancelled" : "failed",
+        loginId: completedLoginId ?? null,
+        authUrl: state.authUrl,
+        account: null,
+        rateLimits: null,
+        error: cancelled ? null : params?.error ?? "ChatGPT login failed",
+        startedAt,
+      });
+      await clearOpenAIOAuthStatus();
+      await closeAttempt();
+    })();
+  };
+
+  try {
+    const result = (await client.request("account/login/start", {
+      type: "chatgpt",
+    })) as {
+      type: "chatgpt";
+      loginId: string;
+      authUrl: string;
+    };
+
+    const next = updateState({
+      phase: "pending",
+      loginId: result.loginId,
+      authUrl: result.authUrl,
+      account: null,
+      rateLimits: null,
+      error: null,
+      startedAt,
+    });
+    if (activeAttempt) {
+      activeAttempt.state = next;
+    }
+    return next;
+  } catch (error) {
+    await clearOpenAIOAuthStatus();
+    await closeAttempt();
+    return updateState({
+      phase: "failed",
+      loginId: null,
+      authUrl: null,
+      account: null,
+      rateLimits: null,
+      error: error instanceof Error ? error.message : String(error),
+      startedAt,
+    });
+  }
+}
+
+export async function cancelOpenAIChatGPTLogin(): Promise<OpenAILoginState> {
+  const attempt = activeAttempt;
+  if (!attempt?.state.loginId) {
+    return getOpenAILoginState();
+  }
+
+  attempt.cancelRequested = true;
+
+  try {
+    await attempt.client.request("account/login/cancel", {
+      loginId: attempt.state.loginId,
+    });
+    const next = updateState({
+      phase: "cancelled",
+      loginId: attempt.state.loginId,
+      authUrl: attempt.state.authUrl,
+      account: null,
+      rateLimits: null,
+      error: null,
+      startedAt: attempt.state.startedAt,
+    });
+    attempt.state = next;
+    await clearOpenAIOAuthStatus();
+    await closeAttempt();
+    return next;
+  } catch (error) {
+    updateState({
+      phase: "failed",
+      loginId: attempt.state.loginId,
+      authUrl: attempt.state.authUrl,
+      account: null,
+      rateLimits: null,
+      error: error instanceof Error ? error.message : String(error),
+      startedAt: attempt.state.startedAt,
+    });
+    await closeAttempt();
+  }
+
+  return getOpenAILoginState();
+}

package/src/lib/settings/runtime-setup.ts

@@ -54,10 +54,20 @@ export async function getRuntimeSetupStates(): Promise<
       runtimeId: "openai-codex-app-server",
       label: openAIRuntime.label,
       providerId: openAIRuntime.providerId,
-      configured: openAIAuth.hasKey,
-      authMethod: openAIAuth.hasKey ? "api_key" : "none",
-      apiKeySource: openAIAuth.apiKeySource,
-      billingMode: "usage",
+      configured:
+        openAIAuth.method === "oauth" ? openAIAuth.oauthConnected : openAIAuth.hasKey,
+      authMethod:
+        openAIAuth.method === "oauth"
+          ? "oauth"
+          : openAIAuth.hasKey
+            ? "api_key"
+            : "none",
+      apiKeySource:
+        openAIAuth.method === "oauth" ? "oauth" : openAIAuth.apiKeySource,
+      billingMode:
+        openAIAuth.method === "oauth" && openAIAuth.oauthConnected
+          ? "subscription"
+          : "usage",
     },
     "anthropic-direct": {
       runtimeId: "anthropic-direct",

package/src/lib/tables/__tests__/enrichment-planner.test.ts

@@ -0,0 +1,124 @@
+import { describe, expect, it } from "vitest";
+import {
+  buildEnrichmentPlan,
+  buildTargetContract,
+  normalizeEnrichmentOutput,
+} from "../enrichment-planner";
+import type { ColumnDef } from "../types";
+
+const baseColumn: ColumnDef = {
+  name: "linkedin_url",
+  displayName: "LinkedIn URL",
+  dataType: "url",
+  position: 0,
+};
+
+describe("buildTargetContract", () => {
+  it("preserves select options for categorical columns", () => {
+    const contract = buildTargetContract({
+      name: "status",
+      displayName: "Status",
+      dataType: "select",
+      position: 1,
+      config: { options: ["Lead", "Qualified", "Closed Won"] },
+    });
+
+    expect(contract.allowedOptions).toEqual(["Lead", "Qualified", "Closed Won"]);
+  });
+});
+
+describe("buildEnrichmentPlan", () => {
+  it("chooses single-pass-classify for boolean columns", () => {
+    const plan = buildEnrichmentPlan({
+      targetColumn: {
+        name: "qualified",
+        displayName: "Qualified",
+        dataType: "boolean",
+        position: 0,
+      },
+      sampleRows: [{ id: "row_1", data: { company: "Acme" } }],
+      eligibleRowCount: 1,
+      promptMode: "auto",
+    });
+
+    expect(plan.strategy).toBe("single-pass-classify");
+    expect(plan.steps).toHaveLength(1);
+  });
+
+  it("chooses research-and-synthesize when text guidance asks for synthesis", () => {
+    const plan = buildEnrichmentPlan({
+      targetColumn: {
+        name: "summary",
+        displayName: "Summary",
+        dataType: "text",
+        position: 0,
+      },
+      sampleRows: [{ id: "row_1", data: { company: "Acme" } }],
+      eligibleRowCount: 1,
+      promptMode: "auto",
+      prompt: "Research the company and synthesize a short account summary.",
+    });
+
+    expect(plan.strategy).toBe("research-and-synthesize");
+    expect(plan.steps).toHaveLength(2);
+    expect(plan.steps[1].prompt).toContain("{{previous}}");
+  });
+
+  it("keeps custom mode single-step and wraps the typed contract", () => {
+    const plan = buildEnrichmentPlan({
+      targetColumn: baseColumn,
+      sampleRows: [{ id: "row_1", data: { company: "Acme" } }],
+      eligibleRowCount: 1,
+      promptMode: "custom",
+      prompt: "Find the profile URL.",
+    });
+
+    expect(plan.promptMode).toBe("custom");
+    expect(plan.steps).toHaveLength(1);
+    expect(plan.steps[0].prompt).toContain("RESPONSE FORMAT");
+  });
+});
+
+describe("normalizeEnrichmentOutput", () => {
+  it("canonicalizes select output to the configured option casing", () => {
+    const result = normalizeEnrichmentOutput("qualified", {
+      columnName: "status",
+      columnLabel: "Status",
+      dataType: "select",
+      allowedOptions: ["Lead", "Qualified", "Closed Won"],
+    });
+
+    expect(result).toEqual({ kind: "valid", value: "Qualified" });
+  });
+
+  it("parses booleans and numbers into typed values", () => {
+    expect(
+      normalizeEnrichmentOutput("true", {
+        columnName: "qualified",
+        columnLabel: "Qualified",
+        dataType: "boolean",
+      })
+    ).toEqual({ kind: "valid", value: true });
+
+    expect(
+      normalizeEnrichmentOutput("42", {
+        columnName: "score",
+        columnLabel: "Score",
+        dataType: "number",
+      })
+    ).toEqual({ kind: "valid", value: 42 });
+  });
+
+  it("rejects invalid urls", () => {
+    expect(
+      normalizeEnrichmentOutput("not-a-url", {
+        columnName: "linkedin_url",
+        columnLabel: "LinkedIn URL",
+        dataType: "url",
+      })
+    ).toEqual({
+      kind: "invalid",
+      reason: "Expected a valid URL or NOT_FOUND",
+    });
+  });
+});