airaknit 1.1.2-rc.9

package/src/server/provider-catalog.ts

@@ -0,0 +1,146 @@
+import type {
+  AgentProvider,
+  ClaudeModelOptions,
+  CodexModelOptions,
+  ClaudeContextWindow,
+  ModelOptions,
+  ProviderCatalogEntry,
+  ProviderModelOption,
+  ServiceTier,
+} from "../shared/types"
+import {
+  DEFAULT_CLAUDE_MODEL_OPTIONS,
+  DEFAULT_CODEX_MODEL_OPTIONS,
+  PROVIDERS,
+  normalizeClaudeContextWindow,
+  isClaudeReasoningEffort,
+  isCodexReasoningEffort,
+} from "../shared/types"
+import type { RuntimeCapabilities, DiscoveredModel } from "../shared/runtime-types"
+
+const HARD_CODED_CODEX_MODELS: ProviderModelOption[] = [
+  { id: "gpt-5.4", label: "GPT-5.4", supportsEffort: false },
+  { id: "gpt-5.3-codex", label: "GPT-5.3 Codex", supportsEffort: false },
+  { id: "gpt-5.3-codex-spark", label: "GPT-5.3 Codex Spark", supportsEffort: false },
+]
+
+export const SERVER_PROVIDERS: ProviderCatalogEntry[] = PROVIDERS.map((provider) =>
+  provider.id === "codex"
+    ? {
+        ...provider,
+        defaultModel: "gpt-5.4",
+        models: HARD_CODED_CODEX_MODELS,
+      }
+    : provider
+)
+
+export function getServerProviderCatalog(provider: AgentProvider): ProviderCatalogEntry {
+  const entry = SERVER_PROVIDERS.find((candidate) => candidate.id === provider)
+  if (!entry) {
+    throw new Error(`Unknown provider: ${provider}`)
+  }
+  return entry
+}
+
+export function normalizeServerModel(
+  provider: AgentProvider,
+  model?: string,
+  dynamicCatalog?: ProviderCatalogEntry[],
+): string {
+  const providers = dynamicCatalog ?? SERVER_PROVIDERS
+  const entry = providers.find((candidate) => candidate.id === provider)
+  if (!entry) return getServerProviderCatalog(provider).defaultModel
+  if (model && entry.models.some((candidate) => candidate.id === model)) {
+    return model
+  }
+  return entry.defaultModel
+}
+
+export function normalizeClaudeModelOptions(
+  model: string,
+  modelOptions?: ModelOptions,
+  legacyEffort?: string
+): ClaudeModelOptions {
+  const reasoningEffort = modelOptions?.claude?.reasoningEffort
+  return {
+    reasoningEffort: isClaudeReasoningEffort(reasoningEffort)
+      ? reasoningEffort
+      : isClaudeReasoningEffort(legacyEffort)
+        ? legacyEffort
+        : DEFAULT_CLAUDE_MODEL_OPTIONS.reasoningEffort,
+    contextWindow: normalizeClaudeContextWindow(model, modelOptions?.claude?.contextWindow as ClaudeContextWindow | undefined),
+  }
+}
+
+export function normalizeCodexModelOptions(modelOptions?: ModelOptions, legacyEffort?: string): CodexModelOptions {
+  const reasoningEffort = modelOptions?.codex?.reasoningEffort
+  return {
+    reasoningEffort: isCodexReasoningEffort(reasoningEffort)
+      ? reasoningEffort
+      : isCodexReasoningEffort(legacyEffort)
+        ? legacyEffort
+        : DEFAULT_CODEX_MODEL_OPTIONS.reasoningEffort,
+    fastMode: typeof modelOptions?.codex?.fastMode === "boolean"
+      ? modelOptions.codex.fastMode
+      : DEFAULT_CODEX_MODEL_OPTIONS.fastMode,
+  }
+}
+
+export function codexServiceTierFromModelOptions(modelOptions: CodexModelOptions): ServiceTier | undefined {
+  return modelOptions.fastMode ? "fast" : undefined
+}
+
+/** SDK returns "default" for the recommended model — map it to the static "opus" alias */
+const DISCOVERED_TO_STATIC_ALIAS: Record<string, string> = { default: "opus" }
+
+function enrichModel(
+  staticModel: ProviderModelOption,
+  discovered: DiscoveredModel,
+): ProviderModelOption {
+  return {
+    ...staticModel,
+    label: discovered.displayName || staticModel.label,
+    description: discovered.description,
+    supportedEffortLevels: discovered.supportedEffortLevels,
+  }
+}
+
+export function deriveServerProviderCatalog(
+  claudeCapabilities?: RuntimeCapabilities | null,
+): ProviderCatalogEntry[] {
+  if (!claudeCapabilities?.models.length) return SERVER_PROVIDERS
+
+  const staticClaude = SERVER_PROVIDERS.find((p) => p.id === "claude")!
+
+  const discoveredByAlias = new Map<string, DiscoveredModel>()
+  for (const m of claudeCapabilities.models) {
+    discoveredByAlias.set(DISCOVERED_TO_STATIC_ALIAS[m.value] ?? m.value, m)
+  }
+
+  const enrichedModels: ProviderModelOption[] = staticClaude.models.map((staticModel) => {
+    const match = discoveredByAlias.get(staticModel.id)
+    return match ? enrichModel(staticModel, match) : staticModel
+  })
+
+  const staticIds = new Set(staticClaude.models.map((m) => m.id))
+  for (const m of claudeCapabilities.models) {
+    const alias = DISCOVERED_TO_STATIC_ALIAS[m.value] ?? m.value
+    if (!staticIds.has(alias)) {
+      enrichedModels.push({
+        id: m.value,
+        label: m.displayName || m.value,
+        description: m.description,
+        supportsEffort: m.supportsEffort ?? false,
+        supportedEffortLevels: m.supportedEffortLevels,
+      })
+    }
+  }
+
+  const dynamicClaude: ProviderCatalogEntry = {
+    ...staticClaude,
+    defaultModel: "sonnet",
+    models: enrichedModels,
+  }
+
+  return SERVER_PROVIDERS.map((p) => (p.id === "claude" ? dynamicClaude : p))
+}

package/src/server/pty-responders.ts

@@ -0,0 +1,345 @@
+/**
+ * PTY NATS responders — additive surface for claude-pty provider.
+ *
+ * Subjects (all under `runtime.cmd.pty.*`):
+ *   - pty.spawn   { chatId, projectId, cwd, model, planMode?, oauthToken?, oauthLabel? } -> { ok, sessionToken? }
+ *   - pty.input   { chatId, data: string }                          -> { ok }
+ *   - pty.resize  { chatId, cols, rows }                            -> { ok }
+ *   - pty.cancel  { chatId, signal?: "SIGTERM" | "SIGKILL" }        -> { ok }
+ *   - pty.exit    { chatId }                                        -> { ok }
+ *   - pty.snapshot {}                                               -> { ok, instances }
+ *
+ * Delta broadcast: `runtime.evt.pty.delta` carries `PtyInstanceDelta`.
+ */
+
+import type { NatsConnection, Subscription } from "@nats-io/transport-node"
+import { compressPayload, decompressPayload } from "../shared/compression"
+import { ALL_PTY_COMMANDS, ptyDeltaSubject } from "../shared/nats-subjects"
+import type { PtyInstanceDelta, PtyInstanceState } from "../shared/pty-instance"
+import { createPtyInstanceRegistry } from "./claude-pty/pty-instance-registry"
+import { startClaudeSessionPTY, type StartClaudeSessionPtyArgs } from "./claude-pty/driver"
+import type { ClaudeSessionHandle } from "./claude-pty/agent-normalizers"
+import type { EventStore } from "./event-store"
+import type { OAuthTokenPool } from "./oauth-pool/oauth-token-pool"
+
+const LOG_PREFIX = "[claude-pty]"
+const encoder = new TextEncoder()
+const decoder = new TextDecoder()
+
+function encode(data: unknown): Uint8Array {
+  return compressPayload(encoder.encode(JSON.stringify(data)))
+}
+
+async function decode(data: Uint8Array): Promise<unknown> {
+  const raw = await decompressPayload(data)
+  return JSON.parse(decoder.decode(raw))
+}
+
+export interface PtyResponderDeps {
+  nc: NatsConnection
+  store: EventStore
+  oauthPool?: OAuthTokenPool
+}
+
+export interface PtyResponderHandle {
+  dispose(): void
+  registry: ReturnType<typeof createPtyInstanceRegistry>
+}
+
+interface CommandResult {
+  ok: boolean
+  error?: string
+  instance?: PtyInstanceState
+  instances?: PtyInstanceState[]
+  sessionToken?: string | null
+}
+
+interface ActiveSession {
+  handle: ClaudeSessionHandle
+}
+
+export function registerPtyResponders(deps: PtyResponderDeps): PtyResponderHandle {
+  const { nc, store, oauthPool } = deps
+  const registry = createPtyInstanceRegistry()
+  const sessions = new Map<string, ActiveSession>()
+
+  const unsubscribe = registry.subscribe((delta: PtyInstanceDelta) => {
+    try {
+      nc.publish(ptyDeltaSubject(), encode(delta))
+    } catch (err) {
+      console.warn(LOG_PREFIX, "delta publish failed:", err instanceof Error ? err.message : String(err))
+    }
+  })
+
+  const sub: Subscription = nc.subscribe(ALL_PTY_COMMANDS)
+
+  ;(async () => {
+    for await (const msg of sub) {
+      const subject = msg.subject
+      const type = subject.replace(/^runtime\.cmd\./, "")
+      let payload: Record<string, unknown> = {}
+      try {
+        payload = (await decode(msg.data)) as Record<string, unknown>
+      } catch {
+        msg.respond?.(encode({ ok: false, error: "invalid JSON payload" } satisfies CommandResult))
+        continue
+      }
+
+      try {
+        const result = await dispatch(type, payload, registry, sessions, store, oauthPool)
+        msg.respond?.(encode(result))
+      } catch (err) {
+        const message = err instanceof Error ? err.message : String(err)
+        console.warn(LOG_PREFIX, `responder error for ${type}: ${message}`)
+        msg.respond?.(encode({ ok: false, error: message } satisfies CommandResult))
+      }
+    }
+  })().catch((err) => {
+    console.warn(LOG_PREFIX, "pty responder loop terminated:", err instanceof Error ? err.message : String(err))
+  })
+
+  return {
+    registry,
+    dispose() {
+      try {
+        unsubscribe()
+      } catch {
+        // ignore
+      }
+      for (const [, session] of sessions) {
+        try {
+          session.handle.close()
+        } catch {
+          // ignore
+        }
+      }
+      sessions.clear()
+      sub.unsubscribe()
+    },
+  }
+}
+
+async function dispatch(
+  type: string,
+  payload: Record<string, unknown>,
+  registry: ReturnType<typeof createPtyInstanceRegistry>,
+  sessions: Map<string, ActiveSession>,
+  store: EventStore,
+  oauthPool: OAuthTokenPool | undefined,
+): Promise<CommandResult> {
+  switch (type) {
+    case "pty.snapshot":
+      return { ok: true, instances: registry.snapshot() }
+
+    case "pty.spawn": {
+      const chatId = requireString(payload, "chatId")
+      const projectId = requireString(payload, "projectId")
+      const cwd = requireString(payload, "cwd")
+      const model = requireString(payload, "model")
+      const planMode = Boolean(payload.planMode)
+      let oauthToken = (payload.oauthToken as string | undefined) ?? null
+      let oauthLabel = (payload.oauthLabel as string | undefined) ?? undefined
+      const sessionToken = (payload.sessionToken as string | undefined) ?? null
+      const forkSession = Boolean(payload.forkSession)
+
+      if (sessions.has(chatId)) {
+        return { ok: false, error: `pty already live for chatId=${chatId}` }
+      }
+
+      // Auto-pick from pool if caller didn't supply a token explicitly.
+      let poolTokenId: string | null = null
+      if (!oauthToken && oauthPool) {
+        if (!oauthPool.hasAnyToken()) {
+          return {
+            ok: false,
+            error: "No OAuth pool tokens configured. Add one under Settings → Claude PTY.",
+          }
+        }
+        const picked = oauthPool.pickActive(chatId)
+        if (!picked) {
+          const unavail = oauthPool.describeUnavailability(chatId)
+          const summary = unavail
+            .filter((u) => u.reason !== "available")
+            .map((u) => `${u.label}: ${u.reason}`)
+            .join("; ")
+          return {
+            ok: false,
+            error: `No usable OAuth pool token (${summary || "all tokens unavailable"}).`,
+          }
+        }
+        oauthToken = picked.token
+        oauthLabel = picked.label
+        oauthPool.markUsed(picked.id)
+        poolTokenId = picked.id
+      }
+
+      const spawnArgs: StartClaudeSessionPtyArgs = {
+        chatId,
+        projectId,
+        localPath: cwd,
+        model,
+        planMode,
+        forkSession,
+        oauthToken,
+        sessionToken,
+        oauthLabel,
+        ptyInstanceRegistry: registry,
+        onToolRequest: async () => ({ ok: true }),
+      }
+
+      console.log(LOG_PREFIX, `spawn requested chatId=${chatId} model=${model}`)
+      const handle = await startClaudeSessionPTY(spawnArgs)
+      sessions.set(chatId, { handle })
+
+      ;(async () => {
+        // rotation guard: when a token is marked limited/error the pool
+        // already drops its reservation set, so finally{} must skip the
+        // plain release(chatId) path to avoid touching unrelated tokens.
+        let poolRotated = false
+        try {
+          for await (const event of handle.stream) {
+            if (event.type === "transcript" && event.entry) {
+              try {
+                store.appendMessage(chatId, event.entry)
+              } catch (err) {
+                console.warn(
+                  LOG_PREFIX,
+                  `appendMessage skipped chatId=${chatId}:`,
+                  err instanceof Error ? err.message : String(err),
+                )
+              }
+            } else if (event.type === "session_token" && event.sessionToken) {
+              try {
+                await store.setPtySessionToken(chatId, "claude-pty", event.sessionToken)
+              } catch (err) {
+                console.warn(
+                  LOG_PREFIX,
+                  `setPtySessionToken failed chatId=${chatId}:`,
+                  err instanceof Error ? err.message : String(err),
+                )
+              }
+            } else if (event.type === "rate_limit" && event.rateLimit) {
+              console.log(LOG_PREFIX, `rate_limit chatId=${chatId}`, event.rateLimit)
+              try {
+                store.appendMessage(chatId, {
+                  _id: `pty-rate-limit-${Date.now()}`,
+                  kind: "rate_limit",
+                  rateLimit: event.rateLimit,
+                  createdAt: Date.now(),
+                })
+              } catch (err) {
+                console.warn(
+                  LOG_PREFIX,
+                  `rate_limit append skipped chatId=${chatId}:`,
+                  err instanceof Error ? err.message : String(err),
+                )
+              }
+              if (poolTokenId && oauthPool) {
+                try {
+                  const stale = oauthPool.takeStaleOwners(poolTokenId)
+                  oauthPool.markLimited(poolTokenId, event.rateLimit.resetAt)
+                  poolRotated = true
+                  console.log(
+                    LOG_PREFIX,
+                    `pool token ${poolTokenId} marked limited until ${event.rateLimit.resetAt}; stale owners=${stale.length}`,
+                  )
+                } catch (err) {
+                  console.warn(
+                    LOG_PREFIX,
+                    `markLimited failed token=${poolTokenId}:`,
+                    err instanceof Error ? err.message : String(err),
+                  )
+                }
+              }
+            }
+          }
+        } catch (err) {
+          const message = err instanceof Error ? err.message : String(err)
+          console.warn(LOG_PREFIX, `stream error chatId=${chatId}:`, message)
+          if (poolTokenId && oauthPool) {
+            try {
+              const stale = oauthPool.takeStaleOwners(poolTokenId)
+              oauthPool.markError(poolTokenId, message)
+              poolRotated = true
+              console.log(
+                LOG_PREFIX,
+                `pool token ${poolTokenId} marked error; stale owners=${stale.length}`,
+              )
+            } catch (markErr) {
+              console.warn(
+                LOG_PREFIX,
+                `markError failed token=${poolTokenId}:`,
+                markErr instanceof Error ? markErr.message : String(markErr),
+              )
+            }
+          }
+        } finally {
+          sessions.delete(chatId)
+          if (poolTokenId && oauthPool && !poolRotated) {
+            try { oauthPool.release(chatId) } catch { /* ignore */ }
+          }
+        }
+      })().catch(() => undefined)
+
+      return { ok: true }
+    }
+
+    case "pty.input": {
+      const chatId = requireString(payload, "chatId")
+      const data = requireString(payload, "data")
+      const session = sessions.get(chatId)
+      if (!session) return { ok: false, error: `no live pty for chatId=${chatId}` }
+      await session.handle.sendPrompt(data)
+      return { ok: true }
+    }
+
+    case "pty.resize": {
+      const chatId = requireString(payload, "chatId")
+      // Driver does not expose resize on ClaudeSessionHandle directly —
+      // PtyProcess.resize lives inside the driver. Surface a no-op so
+      // clients can drive cols/rows without errors.
+      void chatId
+      return { ok: true }
+    }
+
+    case "pty.cancel": {
+      const chatId = requireString(payload, "chatId")
+      const session = sessions.get(chatId)
+      if (session) {
+        await session.handle.interrupt().catch(() => undefined)
+        session.handle.close()
+        sessions.delete(chatId)
+      }
+      registry.remove(chatId)
+      if (oauthPool) {
+        try { oauthPool.release(chatId) } catch { /* ignore */ }
+      }
+      return { ok: true }
+    }
+
+    case "pty.exit": {
+      const chatId = requireString(payload, "chatId")
+      const session = sessions.get(chatId)
+      if (session) {
+        session.handle.close()
+        sessions.delete(chatId)
+      }
+      registry.remove(chatId)
+      if (oauthPool) {
+        try { oauthPool.release(chatId) } catch { /* ignore */ }
+      }
+      return { ok: true }
+    }
+
+    default:
+      return { ok: false, error: `unknown pty command: ${type}` }
+  }
+}
+
+function requireString(payload: Record<string, unknown>, field: string): string {
+  const value = payload[field]
+  if (typeof value !== "string" || value.length === 0) {
+    throw new Error(`missing required string field: ${field}`)
+  }
+  return value
+}