feat: add connect zai provider (#523)

2026-01-12 14:21:12 -08:00
parent 4d76b849f9
commit 7140374879
4 changed files with 346 additions and 10 deletions
--- a/src/cli/App.tsx
+++ b/src/cli/App.tsx
@@ -3516,6 +3516,20 @@ export default function App({
          const provider = parts[1]?.toLowerCase();
          const hasCode = parts.length > 2;

+          // Handle /connect zai - create zai-coding-plan provider
+          if (provider === "zai") {
+            const { handleConnectZai } = await import("./commands/connect");
+            await handleConnectZai(
+              {
+                buffersRef,
+                refreshDerived,
+                setCommandRunning,
+              },
+              msg,
+            );
+            return { submitted: true };
+          }
+
          // If no code provided and provider is claude, show the OAuth dialog
          if (provider === "claude" && !hasCode) {
            setActiveOverlay("oauth");
--- a/src/cli/commands/connect.ts
+++ b/src/cli/commands/connect.ts
@@ -12,6 +12,12 @@ import {
  createOrUpdateAnthropicProvider,
  removeAnthropicProvider,
 } from "../../providers/anthropic-provider";
+import {
+  createOrUpdateZaiProvider,
+  getZaiProvider,
+  removeZaiProvider,
+  ZAI_PROVIDER_NAME,
+} from "../../providers/zai-provider";
 import { settingsManager } from "../../settings-manager";
 import { getErrorMessage } from "../../utils/error";
 import type { Buffers, Line } from "../helpers/accumulator";
@@ -102,23 +108,29 @@ export async function handleConnect(
      ctx.buffersRef,
      ctx.refreshDerived,
      msg,
-      "Usage: /connect claude\n\nConnect to Claude via OAuth to authenticate without an API key.",
+      "Usage: /connect <provider> [options]\n\nAvailable providers:\n  • claude          - Connect via OAuth to authenticate without an API key\n  • zai <api_key>   - Connect to Zai with your API key",
      false,
    );
    return;
  }

-  if (provider !== "claude") {
+  if (provider !== "claude" && provider !== "zai") {
    addCommandResult(
      ctx.buffersRef,
      ctx.refreshDerived,
      msg,
-      `Error: Unknown provider "${provider}"\n\nCurrently only 'claude' provider is supported.\nUsage: /connect claude`,
+      `Error: Unknown provider "${provider}"\n\nAvailable providers: claude, zai\nUsage: /connect <provider> [options]`,
      false,
    );
    return;
  }

+  // Zai is handled separately in App.tsx, but add a fallback just in case
+  if (provider === "zai") {
+    await handleConnectZai(ctx, msg);
+    return;
+  }
+
  // If authorization code is provided, complete the OAuth flow
  if (authCode && authCode.length > 0) {
    await completeOAuthFlow(ctx, msg, authCode);
@@ -382,7 +394,7 @@ async function completeOAuthFlow(

 /**
 * Handle /disconnect command
- * Usage: /disconnect [claude]
+ * Usage: /disconnect <provider>
 */
 export async function handleDisconnect(
  ctx: ConnectCommandContext,
@@ -391,18 +403,47 @@ export async function handleDisconnect(
  const parts = msg.trim().split(/\s+/);
  const provider = parts[1]?.toLowerCase();

-  // If no provider specified, show help or assume claude
-  if (provider && provider !== "claude") {
+  // If no provider specified, show usage
+  if (!provider) {
    addCommandResult(
      ctx.buffersRef,
      ctx.refreshDerived,
      msg,
-      `Error: Unknown provider "${provider}"\n\nCurrently only 'claude' provider is supported.\nUsage: /disconnect`,
+      "Usage: /disconnect <provider>\n\nAvailable providers: claude, zai",
      false,
    );
    return;
  }

+  // Handle /disconnect zai
+  if (provider === "zai") {
+    await handleDisconnectZai(ctx, msg);
+    return;
+  }
+
+  // Handle /disconnect claude
+  if (provider === "claude") {
+    await handleDisconnectClaude(ctx, msg);
+    return;
+  }
+
+  // Unknown provider
+  addCommandResult(
+    ctx.buffersRef,
+    ctx.refreshDerived,
+    msg,
+    `Error: Unknown provider "${provider}"\n\nAvailable providers: claude, zai\nUsage: /disconnect <provider>`,
+    false,
+  );
+}
+
+/**
+ * Handle /disconnect claude
+ */
+async function handleDisconnectClaude(
+  ctx: ConnectCommandContext,
+  msg: string,
+): Promise<void> {
  // Check if connected
  if (!settingsManager.hasAnthropicOAuth()) {
    addCommandResult(
@@ -463,3 +504,134 @@ export async function handleDisconnect(
    ctx.setCommandRunning(false);
  }
 }
+
+/**
+ * Handle /disconnect zai
+ */
+async function handleDisconnectZai(
+  ctx: ConnectCommandContext,
+  msg: string,
+): Promise<void> {
+  // Check if Zai provider exists
+  const existing = await getZaiProvider();
+  if (!existing) {
+    addCommandResult(
+      ctx.buffersRef,
+      ctx.refreshDerived,
+      msg,
+      "Not currently connected to Zai.\n\nUse /connect zai <api_key> to connect.",
+      false,
+    );
+    return;
+  }
+
+  // Show running status
+  const cmdId = addCommandResult(
+    ctx.buffersRef,
+    ctx.refreshDerived,
+    msg,
+    "Disconnecting from Zai...",
+    true,
+    "running",
+  );
+
+  ctx.setCommandRunning(true);
+
+  try {
+    // Remove provider from Letta
+    await removeZaiProvider();
+
+    updateCommandResult(
+      ctx.buffersRef,
+      ctx.refreshDerived,
+      cmdId,
+      msg,
+      `✓ Disconnected from Zai.\n\n` +
+        `Provider '${ZAI_PROVIDER_NAME}' removed from Letta.`,
+      true,
+      "finished",
+    );
+  } catch (error) {
+    updateCommandResult(
+      ctx.buffersRef,
+      ctx.refreshDerived,
+      cmdId,
+      msg,
+      `✗ Failed to disconnect from Zai: ${getErrorMessage(error)}`,
+      false,
+      "finished",
+    );
+  } finally {
+    ctx.setCommandRunning(false);
+  }
+}
+
+/**
+ * Handle /connect zai command
+ * Usage: /connect zai <api_key>
+ *
+ * Creates the zai-coding-plan provider with the provided API key
+ */
+export async function handleConnectZai(
+  ctx: ConnectCommandContext,
+  msg: string,
+): Promise<void> {
+  const parts = msg.trim().split(/\s+/);
+  // Join all remaining parts in case the API key got split
+  const apiKey = parts.slice(2).join("");
+
+  // If no API key provided, show usage
+  if (!apiKey || apiKey.length === 0) {
+    addCommandResult(
+      ctx.buffersRef,
+      ctx.refreshDerived,
+      msg,
+      "Usage: /connect zai <api_key>\n\n" +
+        "Connect to Zai by providing your API key.\n\n" +
+        "Example: /connect zai <api_key>...",
+      false,
+    );
+    return;
+  }
+
+  // Show running status
+  const cmdId = addCommandResult(
+    ctx.buffersRef,
+    ctx.refreshDerived,
+    msg,
+    "Creating Zai coding plan provider...",
+    true,
+    "running",
+  );
+
+  ctx.setCommandRunning(true);
+
+  try {
+    // Create or update the Zai provider with the API key
+    await createOrUpdateZaiProvider(apiKey);
+
+    updateCommandResult(
+      ctx.buffersRef,
+      ctx.refreshDerived,
+      cmdId,
+      msg,
+      `✓ Successfully connected to Zai!\n\n` +
+        `Provider '${ZAI_PROVIDER_NAME}' created in Letta.\n\n` +
+        `The models are populated in /model → "All Available Models"`,
+      true,
+      "finished",
+    );
+  } catch (error) {
+    updateCommandResult(
+      ctx.buffersRef,
+      ctx.refreshDerived,
+      cmdId,
+      msg,
+      `✗ Failed to create Zai provider: ${getErrorMessage(error)}`,
+      false,
+      "finished",
+    );
+  } finally {
+    ctx.setCommandRunning(false);
+  }
+}
--- a/src/cli/commands/registry.ts
+++ b/src/cli/commands/registry.ts
@@ -247,15 +247,15 @@ export const commands: Record<string, Command> = {

  // === Session management (order 40-49) ===
  "/connect": {
-    desc: "Connect an existing Claude account (/connect claude)",
+    desc: "Connect an existing account (/connect zai <api-key>)",
    order: 40,
    handler: () => {
      // Handled specially in App.tsx
-      return "Initiating OAuth connection...";
+      return "Initiating account connection...";
    },
  },
  "/disconnect": {
-    desc: "Disconnect from Claude OAuth",
+    desc: "Disconnect an existing account (/disconnect zai)",
    order: 41,
    handler: () => {
      // Handled specially in App.tsx
--- a/src/providers/zai-provider.ts
+++ b/src/providers/zai-provider.ts
@@ -0,0 +1,150 @@
+/**
+ * Direct API calls to Letta for managing Zai provider
+ */
+
+import { LETTA_CLOUD_API_URL } from "../auth/oauth";
+import { settingsManager } from "../settings-manager";
+
+// Provider name constant for Zai coding plan
+export const ZAI_PROVIDER_NAME = "zai-coding-plan";
+
+interface ProviderResponse {
+  id: string;
+  name: string;
+  provider_type: string;
+  api_key?: string;
+  base_url?: string;
+}
+
+/**
+ * Get the Letta API base URL and auth token
+ */
+async function getLettaConfig(): Promise<{ baseUrl: string; apiKey: string }> {
+  const settings = await settingsManager.getSettingsWithSecureTokens();
+  const baseUrl =
+    process.env.LETTA_BASE_URL ||
+    settings.env?.LETTA_BASE_URL ||
+    LETTA_CLOUD_API_URL;
+  const apiKey = process.env.LETTA_API_KEY || settings.env?.LETTA_API_KEY || "";
+  return { baseUrl, apiKey };
+}
+
+/**
+ * Make a request to the Letta providers API
+ */
+async function providersRequest<T>(
+  method: "GET" | "POST" | "PATCH" | "DELETE",
+  path: string,
+  body?: Record<string, unknown>,
+): Promise<T> {
+  const { baseUrl, apiKey } = await getLettaConfig();
+  const url = `${baseUrl}${path}`;
+
+  const response = await fetch(url, {
+    method,
+    headers: {
+      "Content-Type": "application/json",
+      Authorization: `Bearer ${apiKey}`,
+      "X-Letta-Source": "letta-code",
+    },
+    ...(body && { body: JSON.stringify(body) }),
+  });
+
+  if (!response.ok) {
+    const errorText = await response.text();
+    throw new Error(`Provider API error (${response.status}): ${errorText}`);
+  }
+
+  // Handle empty responses (e.g., DELETE)
+  const text = await response.text();
+  if (!text) {
+    return {} as T;
+  }
+  return JSON.parse(text) as T;
+}
+
+/**
+ * List all providers
+ */
+async function listProviders(): Promise<ProviderResponse[]> {
+  try {
+    const response = await providersRequest<ProviderResponse[]>(
+      "GET",
+      "/v1/providers",
+    );
+    return response;
+  } catch {
+    return [];
+  }
+}
+
+/**
+ * Get the zai-coding-plan provider if it exists
+ */
+export async function getZaiProvider(): Promise<ProviderResponse | null> {
+  const providers = await listProviders();
+  return providers.find((p) => p.name === ZAI_PROVIDER_NAME) || null;
+}
+
+/**
+ * Create the Zai coding plan provider with the given API key
+ */
+export async function createZaiProvider(
+  apiKey: string,
+): Promise<ProviderResponse> {
+  return providersRequest<ProviderResponse>("POST", "/v1/providers", {
+    name: ZAI_PROVIDER_NAME,
+    provider_type: "zai",
+    api_key: apiKey,
+  });
+}
+
+/**
+ * Update an existing Zai provider with a new API key
+ */
+export async function updateZaiProvider(
+  providerId: string,
+  apiKey: string,
+): Promise<ProviderResponse> {
+  return providersRequest<ProviderResponse>(
+    "PATCH",
+    `/v1/providers/${providerId}`,
+    {
+      api_key: apiKey,
+    },
+  );
+}
+
+/**
+ * Create or update the Zai coding plan provider
+ * If provider exists, updates it with the new API key
+ * If not, creates a new provider
+ */
+export async function createOrUpdateZaiProvider(
+  apiKey: string,
+): Promise<ProviderResponse> {
+  const existing = await getZaiProvider();
+
+  if (existing) {
+    return updateZaiProvider(existing.id, apiKey);
+  }
+
+  return createZaiProvider(apiKey);
+}
+
+/**
+ * Delete the Zai provider by ID
+ */
+async function deleteZaiProvider(providerId: string): Promise<void> {
+  await providersRequest<void>("DELETE", `/v1/providers/${providerId}`);
+}
+
+/**
+ * Remove the Zai provider (called on /disconnect zai)
+ */
+export async function removeZaiProvider(): Promise<void> {
+  const existing = await getZaiProvider();
+  if (existing) {
+    await deleteZaiProvider(existing.id);
+  }
+}