yzcode-cli 1.0.1 → 1.0.3

package/bridge/jwtUtils.ts

@@ -0,0 +1,256 @@
+import { logEvent } from '../services/analytics/index.js'
+import { logForDebugging } from '../utils/debug.js'
+import { logForDiagnosticsNoPII } from '../utils/diagLogs.js'
+import { errorMessage } from '../utils/errors.js'
+import { jsonParse } from '../utils/slowOperations.js'
+
+/** Format a millisecond duration as a human-readable string (e.g. "5m 30s"). */
+function formatDuration(ms: number): string {
+  if (ms < 60_000) return `${Math.round(ms / 1000)}s`
+  const m = Math.floor(ms / 60_000)
+  const s = Math.round((ms % 60_000) / 1000)
+  return s > 0 ? `${m}m ${s}s` : `${m}m`
+}
+
+/**
+ * Decode a JWT's payload segment without verifying the signature.
+ * Strips the `sk-ant-si-` session-ingress prefix if present.
+ * Returns the parsed JSON payload as `unknown`, or `null` if the
+ * token is malformed or the payload is not valid JSON.
+ */
+export function decodeJwtPayload(token: string): unknown | null {
+  const jwt = token.startsWith('sk-ant-si-')
+    ? token.slice('sk-ant-si-'.length)
+    : token
+  const parts = jwt.split('.')
+  if (parts.length !== 3 || !parts[1]) return null
+  try {
+    return jsonParse(Buffer.from(parts[1], 'base64url').toString('utf8'))
+  } catch {
+    return null
+  }
+}
+
+/**
+ * Decode the `exp` (expiry) claim from a JWT without verifying the signature.
+ * @returns The `exp` value in Unix seconds, or `null` if unparseable
+ */
+export function decodeJwtExpiry(token: string): number | null {
+  const payload = decodeJwtPayload(token)
+  if (
+    payload !== null &&
+    typeof payload === 'object' &&
+    'exp' in payload &&
+    typeof payload.exp === 'number'
+  ) {
+    return payload.exp
+  }
+  return null
+}
+
+/** Refresh buffer: request a new token before expiry. */
+const TOKEN_REFRESH_BUFFER_MS = 5 * 60 * 1000
+
+/** Fallback refresh interval when the new token's expiry is unknown. */
+const FALLBACK_REFRESH_INTERVAL_MS = 30 * 60 * 1000 // 30 minutes
+
+/** Max consecutive failures before giving up on the refresh chain. */
+const MAX_REFRESH_FAILURES = 3
+
+/** Retry delay when getAccessToken returns undefined. */
+const REFRESH_RETRY_DELAY_MS = 60_000
+
+/**
+ * Creates a token refresh scheduler that proactively refreshes session tokens
+ * before they expire. Used by both the standalone bridge and the REPL bridge.
+ *
+ * When a token is about to expire, the scheduler calls `onRefresh` with the
+ * session ID and the bridge's OAuth access token. The caller is responsible
+ * for delivering the token to the appropriate transport (child process stdin
+ * for standalone bridge, WebSocket reconnect for REPL bridge).
+ */
+export function createTokenRefreshScheduler({
+  getAccessToken,
+  onRefresh,
+  label,
+  refreshBufferMs = TOKEN_REFRESH_BUFFER_MS,
+}: {
+  getAccessToken: () => string | undefined | Promise<string | undefined>
+  onRefresh: (sessionId: string, oauthToken: string) => void
+  label: string
+  /** How long before expiry to fire refresh. Defaults to 5 min. */
+  refreshBufferMs?: number
+}): {
+  schedule: (sessionId: string, token: string) => void
+  scheduleFromExpiresIn: (sessionId: string, expiresInSeconds: number) => void
+  cancel: (sessionId: string) => void
+  cancelAll: () => void
+} {
+  const timers = new Map<string, ReturnType<typeof setTimeout>>()
+  const failureCounts = new Map<string, number>()
+  // Generation counter per session — incremented by schedule() and cancel()
+  // so that in-flight async doRefresh() calls can detect when they've been
+  // superseded and should skip setting follow-up timers.
+  const generations = new Map<string, number>()
+
+  function nextGeneration(sessionId: string): number {
+    const gen = (generations.get(sessionId) ?? 0) + 1
+    generations.set(sessionId, gen)
+    return gen
+  }
+
+  function schedule(sessionId: string, token: string): void {
+    const expiry = decodeJwtExpiry(token)
+    if (!expiry) {
+      // Token is not a decodable JWT (e.g. an OAuth token passed from the
+      // REPL bridge WebSocket open handler).  Preserve any existing timer
+      // (such as the follow-up refresh set by doRefresh) so the refresh
+      // chain is not broken.
+      logForDebugging(
+        `[${label}:token] Could not decode JWT expiry for sessionId=${sessionId}, token prefix=${token.slice(0, 15)}…, keeping existing timer`,
+      )
+      return
+    }
+
+    // Clear any existing refresh timer — we have a concrete expiry to replace it.
+    const existing = timers.get(sessionId)
+    if (existing) {
+      clearTimeout(existing)
+    }
+
+    // Bump generation to invalidate any in-flight async doRefresh.
+    const gen = nextGeneration(sessionId)
+
+    const expiryDate = new Date(expiry * 1000).toISOString()
+    const delayMs = expiry * 1000 - Date.now() - refreshBufferMs
+    if (delayMs <= 0) {
+      logForDebugging(
+        `[${label}:token] Token for sessionId=${sessionId} expires=${expiryDate} (past or within buffer), refreshing immediately`,
+      )
+      void doRefresh(sessionId, gen)
+      return
+    }
+
+    logForDebugging(
+      `[${label}:token] Scheduled token refresh for sessionId=${sessionId} in ${formatDuration(delayMs)} (expires=${expiryDate}, buffer=${refreshBufferMs / 1000}s)`,
+    )
+
+    const timer = setTimeout(doRefresh, delayMs, sessionId, gen)
+    timers.set(sessionId, timer)
+  }
+
+  /**
+   * Schedule refresh using an explicit TTL (seconds until expiry) rather
+   * than decoding a JWT's exp claim. Used by callers whose JWT is opaque
+   * (e.g. POST /v1/code/sessions/{id}/bridge returns expires_in directly).
+   */
+  function scheduleFromExpiresIn(
+    sessionId: string,
+    expiresInSeconds: number,
+  ): void {
+    const existing = timers.get(sessionId)
+    if (existing) clearTimeout(existing)
+    const gen = nextGeneration(sessionId)
+    // Clamp to 30s floor — if refreshBufferMs exceeds the server's expires_in
+    // (e.g. very large buffer for frequent-refresh testing, or server shortens
+    // expires_in unexpectedly), unclamped delayMs ≤ 0 would tight-loop.
+    const delayMs = Math.max(expiresInSeconds * 1000 - refreshBufferMs, 30_000)
+    logForDebugging(
+      `[${label}:token] Scheduled token refresh for sessionId=${sessionId} in ${formatDuration(delayMs)} (expires_in=${expiresInSeconds}s, buffer=${refreshBufferMs / 1000}s)`,
+    )
+    const timer = setTimeout(doRefresh, delayMs, sessionId, gen)
+    timers.set(sessionId, timer)
+  }
+
+  async function doRefresh(sessionId: string, gen: number): Promise<void> {
+    let oauthToken: string | undefined
+    try {
+      oauthToken = await getAccessToken()
+    } catch (err) {
+      logForDebugging(
+        `[${label}:token] getAccessToken threw for sessionId=${sessionId}: ${errorMessage(err)}`,
+        { level: 'error' },
+      )
+    }
+
+    // If the session was cancelled or rescheduled while we were awaiting,
+    // the generation will have changed — bail out to avoid orphaned timers.
+    if (generations.get(sessionId) !== gen) {
+      logForDebugging(
+        `[${label}:token] doRefresh for sessionId=${sessionId} stale (gen ${gen} vs ${generations.get(sessionId)}), skipping`,
+      )
+      return
+    }
+
+    if (!oauthToken) {
+      const failures = (failureCounts.get(sessionId) ?? 0) + 1
+      failureCounts.set(sessionId, failures)
+      logForDebugging(
+        `[${label}:token] No OAuth token available for refresh, sessionId=${sessionId} (failure ${failures}/${MAX_REFRESH_FAILURES})`,
+        { level: 'error' },
+      )
+      logForDiagnosticsNoPII('error', 'bridge_token_refresh_no_oauth')
+      // Schedule a retry so the refresh chain can recover if the token
+      // becomes available again (e.g. transient cache clear during refresh).
+      // Cap retries to avoid spamming on genuine failures.
+      if (failures < MAX_REFRESH_FAILURES) {
+        const retryTimer = setTimeout(
+          doRefresh,
+          REFRESH_RETRY_DELAY_MS,
+          sessionId,
+          gen,
+        )
+        timers.set(sessionId, retryTimer)
+      }
+      return
+    }
+
+    // Reset failure counter on successful token retrieval
+    failureCounts.delete(sessionId)
+
+    logForDebugging(
+      `[${label}:token] Refreshing token for sessionId=${sessionId}: new token prefix=${oauthToken.slice(0, 15)}…`,
+    )
+    logEvent('tengu_bridge_token_refreshed', {})
+    onRefresh(sessionId, oauthToken)
+
+    // Schedule a follow-up refresh so long-running sessions stay authenticated.
+    // Without this, the initial one-shot timer leaves the session vulnerable
+    // to token expiry if it runs past the first refresh window.
+    const timer = setTimeout(
+      doRefresh,
+      FALLBACK_REFRESH_INTERVAL_MS,
+      sessionId,
+      gen,
+    )
+    timers.set(sessionId, timer)
+    logForDebugging(
+      `[${label}:token] Scheduled follow-up refresh for sessionId=${sessionId} in ${formatDuration(FALLBACK_REFRESH_INTERVAL_MS)}`,
+    )
+  }
+
+  function cancel(sessionId: string): void {
+    // Bump generation to invalidate any in-flight async doRefresh.
+    nextGeneration(sessionId)
+    const timer = timers.get(sessionId)
+    if (timer) {
+      clearTimeout(timer)
+      timers.delete(sessionId)
+    }
+    failureCounts.delete(sessionId)
+  }
+
+  function cancelAll(): void {
+    // Bump all generations so in-flight doRefresh calls are invalidated.
+    for (const sessionId of generations.keys()) {
+      nextGeneration(sessionId)
+    }
+    for (const timer of timers.values()) {
+      clearTimeout(timer)
+    }
+    timers.clear()
+    failureCounts.clear()
+  }
+
+  return { schedule, scheduleFromExpiresIn, cancel, cancelAll }
+}

package/bridge/pollConfig.ts

@@ -0,0 +1,110 @@
+import { z } from 'zod/v4'
+import { getFeatureValue_CACHED_WITH_REFRESH } from '../services/analytics/growthbook.js'
+import { lazySchema } from '../utils/lazySchema.js'
+import {
+  DEFAULT_POLL_CONFIG,
+  type PollIntervalConfig,
+} from './pollConfigDefaults.js'
+
+// .min(100) on the seek-work intervals restores the old Math.max(..., 100)
+// defense-in-depth floor against fat-fingered GrowthBook values. Unlike a
+// clamp, Zod rejects the whole object on violation — a config with one bad
+// field falls back to DEFAULT_POLL_CONFIG entirely rather than being
+// partially trusted.
+//
+// The at_capacity intervals use a 0-or-≥100 refinement: 0 means "disabled"
+// (heartbeat-only mode), ≥100 is the fat-finger floor. Values 1–99 are
+// rejected so unit confusion (ops thinks seconds, enters 10) doesn't poll
+// every 10ms against the VerifyEnvironmentSecretAuth DB path.
+//
+// The object-level refines require at least one at-capacity liveness
+// mechanism enabled: heartbeat OR the relevant poll interval. Without this,
+// the hb=0, atCapMs=0 drift config (ops disables heartbeat without
+// restoring at_capacity) falls through every throttle site with no sleep —
+// tight-looping /poll at HTTP-round-trip speed.
+const zeroOrAtLeast100 = {
+  message: 'must be 0 (disabled) or ≥100ms',
+}
+const pollIntervalConfigSchema = lazySchema(() =>
+  z
+    .object({
+      poll_interval_ms_not_at_capacity: z.number().int().min(100),
+      // 0 = no at-capacity polling. Independent of heartbeat — both can be
+      // enabled (heartbeat runs, periodically breaks out to poll).
+      poll_interval_ms_at_capacity: z
+        .number()
+        .int()
+        .refine(v => v === 0 || v >= 100, zeroOrAtLeast100),
+      // 0 = disabled; positive value = heartbeat at this interval while at
+      // capacity. Runs alongside at-capacity polling, not instead of it.
+      // Named non_exclusive to distinguish from the old heartbeat_interval_ms
+      // (either-or semantics in pre-#22145 clients). .default(0) so existing
+      // GrowthBook configs without this field parse successfully.
+      non_exclusive_heartbeat_interval_ms: z.number().int().min(0).default(0),
+      // Multisession (bridgeMain.ts) intervals. Defaults match the
+      // single-session values so existing configs without these fields
+      // preserve current behavior.
+      multisession_poll_interval_ms_not_at_capacity: z
+        .number()
+        .int()
+        .min(100)
+        .default(
+          DEFAULT_POLL_CONFIG.multisession_poll_interval_ms_not_at_capacity,
+        ),
+      multisession_poll_interval_ms_partial_capacity: z
+        .number()
+        .int()
+        .min(100)
+        .default(
+          DEFAULT_POLL_CONFIG.multisession_poll_interval_ms_partial_capacity,
+        ),
+      multisession_poll_interval_ms_at_capacity: z
+        .number()
+        .int()
+        .refine(v => v === 0 || v >= 100, zeroOrAtLeast100)
+        .default(DEFAULT_POLL_CONFIG.multisession_poll_interval_ms_at_capacity),
+      // .min(1) matches the server's ge=1 constraint (work_v1.py:230).
+      reclaim_older_than_ms: z.number().int().min(1).default(5000),
+      session_keepalive_interval_v2_ms: z
+        .number()
+        .int()
+        .min(0)
+        .default(120_000),
+    })
+    .refine(
+      cfg =>
+        cfg.non_exclusive_heartbeat_interval_ms > 0 ||
+        cfg.poll_interval_ms_at_capacity > 0,
+      {
+        message:
+          'at-capacity liveness requires non_exclusive_heartbeat_interval_ms > 0 or poll_interval_ms_at_capacity > 0',
+      },
+    )
+    .refine(
+      cfg =>
+        cfg.non_exclusive_heartbeat_interval_ms > 0 ||
+        cfg.multisession_poll_interval_ms_at_capacity > 0,
+      {
+        message:
+          'at-capacity liveness requires non_exclusive_heartbeat_interval_ms > 0 or multisession_poll_interval_ms_at_capacity > 0',
+      },
+    ),
+)
+
+/**
+ * Fetch the bridge poll interval config from GrowthBook with a 5-minute
+ * refresh window. Validates the served JSON against the schema; falls back
+ * to defaults if the flag is absent, malformed, or partially-specified.
+ *
+ * Shared by bridgeMain.ts (standalone) and replBridge.ts (REPL) so ops
+ * can tune both poll rates fleet-wide with a single config push.
+ */
+export function getPollIntervalConfig(): PollIntervalConfig {
+  const raw = getFeatureValue_CACHED_WITH_REFRESH<unknown>(
+    'tengu_bridge_poll_interval_config',
+    DEFAULT_POLL_CONFIG,
+    5 * 60 * 1000,
+  )
+  const parsed = pollIntervalConfigSchema().safeParse(raw)
+  return parsed.success ? parsed.data : DEFAULT_POLL_CONFIG
+}

package/bridge/pollConfigDefaults.ts

@@ -0,0 +1,82 @@
+/**
+ * Bridge poll interval defaults. Extracted from pollConfig.ts so callers
+ * that don't need live GrowthBook tuning (daemon via Agent SDK) can avoid
+ * the growthbook.ts → config.ts → file.ts → sessionStorage.ts → commands.ts
+ * transitive dependency chain.
+ */
+
+/**
+ * Poll interval when actively seeking work (no transport / below maxSessions).
+ * Governs user-visible "connecting…" latency on initial work pickup and
+ * recovery speed after the server re-dispatches a work item.
+ */
+const POLL_INTERVAL_MS_NOT_AT_CAPACITY = 2000
+
+/**
+ * Poll interval when the transport is connected. Runs independently of
+ * heartbeat — when both are enabled, the heartbeat loop breaks out to poll
+ * at this interval. Set to 0 to disable at-capacity polling entirely.
+ *
+ * Server-side constraints that bound this value:
+ * - BRIDGE_LAST_POLL_TTL = 4h (Redis key expiry → environment auto-archived)
+ * - max_poll_stale_seconds = 24h (session-creation health gate, currently disabled)
+ *
+ * 10 minutes gives 24× headroom on the Redis TTL while still picking up
+ * server-initiated token-rotation redispatches within one poll cycle.
+ * The transport auto-reconnects internally for 10 minutes on transient WS
+ * failures, so poll is not the recovery path — it's strictly a liveness
+ * signal plus a backstop for permanent close.
+ */
+const POLL_INTERVAL_MS_AT_CAPACITY = 600_000
+
+/**
+ * Multisession bridge (bridgeMain.ts) poll intervals. Defaults match the
+ * single-session values so existing GrowthBook configs without these fields
+ * preserve current behavior. Ops can tune these independently via the
+ * tengu_bridge_poll_interval_config GB flag.
+ */
+const MULTISESSION_POLL_INTERVAL_MS_NOT_AT_CAPACITY =
+  POLL_INTERVAL_MS_NOT_AT_CAPACITY
+const MULTISESSION_POLL_INTERVAL_MS_PARTIAL_CAPACITY =
+  POLL_INTERVAL_MS_NOT_AT_CAPACITY
+const MULTISESSION_POLL_INTERVAL_MS_AT_CAPACITY = POLL_INTERVAL_MS_AT_CAPACITY
+
+export type PollIntervalConfig = {
+  poll_interval_ms_not_at_capacity: number
+  poll_interval_ms_at_capacity: number
+  non_exclusive_heartbeat_interval_ms: number
+  multisession_poll_interval_ms_not_at_capacity: number
+  multisession_poll_interval_ms_partial_capacity: number
+  multisession_poll_interval_ms_at_capacity: number
+  reclaim_older_than_ms: number
+  session_keepalive_interval_v2_ms: number
+}
+
+export const DEFAULT_POLL_CONFIG: PollIntervalConfig = {
+  poll_interval_ms_not_at_capacity: POLL_INTERVAL_MS_NOT_AT_CAPACITY,
+  poll_interval_ms_at_capacity: POLL_INTERVAL_MS_AT_CAPACITY,
+  // 0 = disabled. When > 0, at-capacity loops send per-work-item heartbeats
+  // at this interval. Independent of poll_interval_ms_at_capacity — both may
+  // run (heartbeat periodically yields to poll). 60s gives 5× headroom under
+  // the server's 300s heartbeat TTL. Named non_exclusive to distinguish from
+  // the old heartbeat_interval_ms field (either-or semantics in pre-#22145
+  // clients — heartbeat suppressed poll). Old clients ignore this key; ops
+  // can set both fields during rollout.
+  non_exclusive_heartbeat_interval_ms: 0,
+  multisession_poll_interval_ms_not_at_capacity:
+    MULTISESSION_POLL_INTERVAL_MS_NOT_AT_CAPACITY,
+  multisession_poll_interval_ms_partial_capacity:
+    MULTISESSION_POLL_INTERVAL_MS_PARTIAL_CAPACITY,
+  multisession_poll_interval_ms_at_capacity:
+    MULTISESSION_POLL_INTERVAL_MS_AT_CAPACITY,
+  // Poll query param: reclaim unacknowledged work items older than this.
+  // Matches the server's DEFAULT_RECLAIM_OLDER_THAN_MS (work_service.py:24).
+  // Enables picking up stale-pending work after JWT expiry, when the prior
+  // ack failed because the session_ingress_token was already stale.
+  reclaim_older_than_ms: 5000,
+  // 0 = disabled. When > 0, push a silent {type:'keep_alive'} frame to
+  // session-ingress at this interval so upstream proxies don't GC an idle
+  // remote-control session. 2 min is the default. _v2: bridge-only gate
+  // (pre-v2 clients read the old key, new clients ignore it).
+  session_keepalive_interval_v2_ms: 120_000,
+}