airaknit 1.1.2-rc.9

package/src/server/claude-pty/tool-callback.ts

@@ -0,0 +1,274 @@
+import { createHmac } from "node:crypto"
+import type {
+  ChatPermissionPolicy,
+  ToolRequest,
+  ToolRequestDecision,
+  ToolRequestStatus,
+} from "../../shared/permission-policy"
+import { POLICY_TERMINAL_STATUSES } from "../../shared/permission-policy"
+import { policy } from "../claude-pty-mcp/permission-gate"
+import { canonicalArgsHash } from "../claude-pty-mcp/canonical-args"
+import type { EventStore } from "../event-store"
+
+export interface ToolCallbackServiceArgs {
+  store: EventStore
+  serverSecret: string
+  now: () => number
+  timeoutMs: number
+}
+
+export interface ToolCallbackSubmitArgs {
+  chatId: string
+  sessionId: string
+  toolUseId: string
+  toolName: string
+  args: Record<string, unknown>
+  chatPolicy: ChatPermissionPolicy
+  cwd: string
+}
+
+export interface ToolCallbackResult {
+  status: ToolRequestStatus
+  decision: ToolRequestDecision
+  mismatchReason?: string
+}
+
+export interface ToolCallbackService {
+  submit(args: ToolCallbackSubmitArgs): Promise<ToolCallbackResult>
+  answer(id: string, decision: ToolRequestDecision): Promise<void>
+  cancel(id: string, reason: string): Promise<void>
+  cancelAllForChat(chatId: string, reason: string): Promise<void>
+  cancelAllForSession(sessionId: string, reason: string): Promise<void>
+  recoverOnStartup(): Promise<void>
+  tickTimeouts(): Promise<void>
+}
+
+export function createToolCallbackService(opts: ToolCallbackServiceArgs): ToolCallbackService {
+  interface PendingWaiter {
+    resolve: (r: ToolCallbackResult) => void
+    expiresAt: number
+  }
+  const waiters = new Map<string, PendingWaiter[]>()
+  // Tracks the canonical (toolName, canonicalArgsHash) per (chatId, sessionId,
+  // toolUseId) so we detect mismatched retries. toolUseId is generated by the
+  // claude CLI per-session and restarts at "1" for each new spawn, so keying
+  // by toolUseId alone collides across chats / restarts.
+  const seenToolUseIds = new Map<string, { id: string; toolName: string; canonicalArgsHash: string }>()
+  function seenKey(s: { chatId: string; sessionId: string; toolUseId: string }): string {
+    return `${s.chatId}|${s.sessionId}|${s.toolUseId}`
+  }
+  // In-memory mirror of persisted records keyed by id — lets submit() check
+  // existing state synchronously (before the first await) so that concurrent
+  // calls within the same event-loop turn see the correct state.
+  const inMemory = new Map<string, ToolRequest>()
+
+  function hmacId(s: ToolCallbackSubmitArgs, hash: string): string {
+    const h = createHmac("sha256", opts.serverSecret)
+    h.update(`${s.chatId}|${s.sessionId}|${s.toolUseId}|${s.toolName}|${hash}`)
+    return h.digest("hex")
+  }
+
+  function resolveWaiters(id: string, result: ToolCallbackResult) {
+    const ws = waiters.get(id) ?? []
+    waiters.delete(id)
+    for (const w of ws) w.resolve(result)
+  }
+
+  async function persistPut(req: ToolRequest): Promise<void> {
+    inMemory.set(req.id, { ...req })
+    await opts.store.putToolRequest(req)
+  }
+
+  async function persistResolve(
+    id: string,
+    update: { status: ToolRequestStatus; decision: ToolRequestDecision; resolvedAt: number; mismatchReason?: string },
+  ): Promise<void> {
+    const existing = inMemory.get(id)
+    if (existing) {
+      inMemory.set(id, { ...existing, ...update })
+    }
+    await opts.store.resolveToolRequest(id, update)
+  }
+
+  const svc: ToolCallbackService = {
+    submit(args) {
+      const hash = canonicalArgsHash(args.args)
+      const id = hmacId(args, hash)
+
+      // ── Arg-mismatch check (synchronous, no I/O) ──────────────────────────
+      const seen = seenToolUseIds.get(seenKey(args))
+      if (seen && (seen.toolName !== args.toolName || seen.canonicalArgsHash !== hash)) {
+        const reason = `argument_mismatch: canonicalArgsHash differs from prior submission for toolUseId=${args.toolUseId}`
+        const decision: ToolRequestDecision = { kind: "deny", reason }
+        const now = opts.now()
+        const mismatchReq: ToolRequest = {
+          id,
+          chatId: args.chatId,
+          sessionId: args.sessionId,
+          toolUseId: args.toolUseId,
+          toolName: args.toolName,
+          arguments: args.args,
+          canonicalArgsHash: hash,
+          policyVerdict: "auto-deny",
+          status: "arg_mismatch",
+          decision,
+          mismatchReason: reason,
+          createdAt: now,
+          resolvedAt: now,
+          expiresAt: now,
+        }
+        // Await persistence so a caller scanning the store after submit() sees the record.
+        return persistPut(mismatchReq).then(() => ({ status: "arg_mismatch" as const, decision, mismatchReason: reason }))
+      }
+
+      // ── Idempotency: check in-memory mirror (synchronous) ─────────────────
+      const existing = inMemory.get(id)
+      if (existing && POLICY_TERMINAL_STATUSES.has(existing.status)) {
+        return Promise.resolve({
+          status: existing.status,
+          decision: existing.decision ?? { kind: "deny", reason: "unknown" },
+          mismatchReason: existing.mismatchReason,
+        })
+      }
+      if (existing) {
+        // Already pending — attach a new waiter.
+        return new Promise<ToolCallbackResult>((resolve) => {
+          const list = waiters.get(id) ?? []
+          list.push({ resolve, expiresAt: existing.expiresAt })
+          waiters.set(id, list)
+        })
+      }
+
+      // ── New request ───────────────────────────────────────────────────────
+      const verdict = policy.evaluate({
+        toolName: args.toolName,
+        args: args.args,
+        chatPolicy: args.chatPolicy,
+        cwd: args.cwd,
+      })
+      const now = opts.now()
+      const expiresAt = now + opts.timeoutMs
+      const req: ToolRequest = {
+        id,
+        chatId: args.chatId,
+        sessionId: args.sessionId,
+        toolUseId: args.toolUseId,
+        toolName: args.toolName,
+        arguments: args.args,
+        canonicalArgsHash: hash,
+        policyVerdict: verdict.verdict,
+        status: "pending",
+        createdAt: now,
+        expiresAt,
+      }
+
+      // Register synchronously so subsequent calls within the same tick see it.
+      inMemory.set(id, { ...req })
+      seenToolUseIds.set(seenKey(args), { id, toolName: args.toolName, canonicalArgsHash: hash })
+
+      if (verdict.verdict === "auto-allow" || verdict.verdict === "auto-deny") {
+        const decision: ToolRequestDecision = verdict.verdict === "auto-allow"
+          ? { kind: "allow", reason: verdict.reason }
+          : { kind: "deny", reason: verdict.reason }
+        const resolvedReq: ToolRequest = { ...req, status: "answered", decision, resolvedAt: now }
+        inMemory.set(id, resolvedReq)
+        // Persist in background; caller gets immediate result.
+        void (async () => {
+          await opts.store.putToolRequest(req)
+          await opts.store.resolveToolRequest(id, { status: "answered", decision, resolvedAt: now })
+        })()
+        return Promise.resolve({ status: "answered", decision })
+      }
+
+      // "ask" verdict → persist then wait for external answer.
+      const pendingPromise = new Promise<ToolCallbackResult>((resolve) => {
+        const list = waiters.get(id) ?? []
+        list.push({ resolve, expiresAt })
+        waiters.set(id, list)
+      })
+      void opts.store.putToolRequest(req)
+      return pendingPromise
+    },
+
+    async answer(id, decision) {
+      const existing = inMemory.get(id) ?? opts.store.getToolRequest(id)
+      if (!existing || POLICY_TERMINAL_STATUSES.has(existing.status)) return
+      await persistResolve(id, { status: "answered", decision, resolvedAt: opts.now() })
+      resolveWaiters(id, { status: "answered", decision })
+    },
+
+    async cancel(id, reason) {
+      const existing = inMemory.get(id) ?? opts.store.getToolRequest(id)
+      if (!existing || POLICY_TERMINAL_STATUSES.has(existing.status)) return
+      const decision: ToolRequestDecision = { kind: "deny", reason: `canceled: ${reason}` }
+      await persistResolve(id, { status: "canceled", decision, resolvedAt: opts.now() })
+      resolveWaiters(id, { status: "canceled", decision })
+    },
+
+    async cancelAllForChat(chatId, reason) {
+      // Collect pending ids from in-memory mirror first (synchronous), then
+      // also check store for any records loaded before this service started.
+      const pendingIds = new Set<string>()
+      for (const [id, req] of inMemory.entries()) {
+        if (req.chatId === chatId && req.status === "pending") pendingIds.add(id)
+      }
+      const storeList = opts.store.listPendingToolRequests(chatId)
+      for (const req of storeList) pendingIds.add(req.id)
+      for (const id of pendingIds) await svc.cancel(id, reason)
+    },
+
+    async cancelAllForSession(sessionId, reason) {
+      const ids = Array.from(waiters.keys())
+      for (const id of ids) {
+        const req = inMemory.get(id) ?? opts.store.getToolRequest(id)
+        if (req && req.sessionId === sessionId) await svc.cancel(id, reason)
+      }
+    },
+
+    async recoverOnStartup() {
+      const all = opts.store.scanAllToolRequests()
+      for (const req of all) {
+        if (req.status !== "pending") continue
+        const decision: ToolRequestDecision = { kind: "deny", reason: "server_restarted" }
+        await persistResolve(req.id, { status: "session_closed", decision, resolvedAt: opts.now() })
+      }
+    },
+
+    async tickTimeouts() {
+      const now = opts.now()
+      for (const [id, list] of waiters.entries()) {
+        if (list.length === 0) continue
+        if (list[0].expiresAt > now) continue
+        const decision: ToolRequestDecision = { kind: "deny", reason: "timeout" }
+        await persistResolve(id, { status: "timeout", decision, resolvedAt: now })
+        resolveWaiters(id, { status: "timeout", decision })
+      }
+    },
+  }
+
+  return svc
+}
+
+/**
+ * Creates a ToolCallbackService and immediately calls recoverOnStartup()
+ * to fail-close any pending tool requests left over from a previous server
+ * run.  KANNA_SERVER_SECRET should be set in the environment for stable
+ * HMAC ids within a process lifetime; if unset, a fresh random UUID is used
+ * (cross-restart idempotency is not required because recoverOnStartup()
+ * already closes all pending records).
+ */
+export async function initToolCallbackOnBoot(args: {
+  store: EventStore
+  serverSecret: string
+  now?: () => number
+  timeoutMs?: number
+}): Promise<ToolCallbackService> {
+  const svc = createToolCallbackService({
+    store: args.store,
+    serverSecret: args.serverSecret,
+    now: args.now ?? (() => Date.now()),
+    timeoutMs: args.timeoutMs ?? 600_000,
+  })
+  await svc.recoverOnStartup()
+  return svc
+}

package/src/server/claude-pty/tui-control.test.ts

@@ -0,0 +1,272 @@
+import { describe, expect, test } from "bun:test"
+import {
+  sendUserPrompt,
+  sendExitCommand,
+  dismissTrustDialogIfPresent,
+  waitForTuiReady,
+  waitForTuiReadyWithTrustDismiss,
+  TRUST_DIALOG_MARKER,
+  TUI_READY_MARKER,
+} from "./tui-control"
+import { OutputRing } from "./output-ring"
+import type { PtyProcess } from "./pty-process.adapter"
+
+function fakePty(): PtyProcess & { sent: string[] } {
+  const sent: string[] = []
+  return {
+    sent,
+    pid: 99997,
+    async sendInput(data: string) { sent.push(data) },
+    resize() { /* noop */ },
+    exited: new Promise<number>(() => { /* never */ }),
+    close() { /* noop */ },
+    kill() { /* noop */ },
+  } as PtyProcess & { sent: string[] }
+}
+
+describe("sendUserPrompt", () => {
+  test("writes bracketed-paste wrapped text then separate carriage return once ring grows", async () => {
+    const pty = fakePty()
+    const ring = new OutputRing()
+    // Simulate the TUI rendering the paste preview shortly after the paste write.
+    setTimeout(() => ring.append("[Pasted text #1 +56 lines]"), 5)
+    await sendUserPrompt(pty, ring, "say hi", { commitTimeoutMs: 500, pollMs: 1 })
+    expect(pty.sent).toEqual(["\x1b[200~say hi\x1b[201~", "\r"])
+  })
+
+  test("empty string still emits paste markers + carriage return after commit", async () => {
+    const pty = fakePty()
+    const ring = new OutputRing()
+    setTimeout(() => ring.append("x"), 5)
+    await sendUserPrompt(pty, ring, "", { commitTimeoutMs: 500, pollMs: 1 })
+    expect(pty.sent).toEqual(["\x1b[200~\x1b[201~", "\r"])
+  })
+
+  test("sends Enter after commitTimeoutMs even if the ring never grows (degraded fallback)", async () => {
+    const pty = fakePty()
+    const ring = new OutputRing() // never grows
+    const start = Date.now()
+    await sendUserPrompt(pty, ring, "no echo", { commitTimeoutMs: 50, pollMs: 5 })
+    const elapsed = Date.now() - start
+    expect(pty.sent).toEqual(["\x1b[200~no echo\x1b[201~", "\r"])
+    expect(elapsed).toBeGreaterThanOrEqual(45)
+  })
+})
+
+describe("sendExitCommand", () => {
+  test("writes /exit + carriage return", async () => {
+    const pty = fakePty()
+    await sendExitCommand(pty)
+    expect(pty.sent).toEqual(["/exit\r"])
+  })
+})
+
+describe("dismissTrustDialogIfPresent", () => {
+  test("sends carriage return when ringbuf contains trust marker", async () => {
+    const pty = fakePty()
+    const ring = new OutputRing()
+    ring.append("Quick safety check: Do you trust this folder? trust this folder")
+    const dismissed = await dismissTrustDialogIfPresent(pty, ring)
+    expect(dismissed).toBe(true)
+    expect(pty.sent).toEqual(["\r"])
+  })
+
+  test("does nothing when ringbuf lacks trust marker", async () => {
+    const pty = fakePty()
+    const ring = new OutputRing()
+    ring.append("Welcome back c!")
+    const dismissed = await dismissTrustDialogIfPresent(pty, ring)
+    expect(dismissed).toBe(false)
+    expect(pty.sent).toEqual([])
+  })
+
+  test("exported TRUST_DIALOG_MARKER is the substring matched", () => {
+    expect(TRUST_DIALOG_MARKER).toBe("trust this folder")
+  })
+})
+
+describe("waitForTuiReady", () => {
+  test("returns 'marker' when ringbuf already contains the input-box marker", async () => {
+    const ring = new OutputRing()
+    ring.append("❯ ")
+    const result = await waitForTuiReady(ring, { hardCapMs: 1000, pollMs: 10, quietPeriodMs: 0 })
+    expect(result).toBe("marker")
+  })
+
+  test("returns 'timeout' when no marker appears within hardCapMs", async () => {
+    const ring = new OutputRing()
+    const result = await waitForTuiReady(ring, { hardCapMs: 200, pollMs: 10, quietPeriodMs: 0 })
+    expect(result).toBe("timeout")
+  })
+
+  test("polls until marker appears", async () => {
+    const ring = new OutputRing()
+    setTimeout(() => ring.append("❯ "), 50)
+    const start = Date.now()
+    const result = await waitForTuiReady(ring, { hardCapMs: 1000, pollMs: 10, quietPeriodMs: 0 })
+    const elapsed = Date.now() - start
+    expect(result).toBe("marker")
+    expect(elapsed).toBeLessThan(300)
+  })
+
+  // Regression: claude TUI v2.1.146 leaks `❯ ` during splash/trust/MCP render
+  // before Ink mounts the keyboard handler. Returning "marker" on first hit
+  // makes the driver send the first prompt into a TUI that discards stdin.
+  // The quiet-period gate waits for ring growth to settle as a proxy for
+  // "input handler attached" — drops the race that causes PTY chats to hang
+  // with no transcript file ever created.
+  test("waits for ring to stay quiet for quietPeriodMs after marker hit", async () => {
+    const ring = new OutputRing()
+    ring.append("splash banner ❯ ")
+    // Keep appending bytes for 200 ms after the marker first appears — the
+    // gate must NOT resolve while the TUI is still rendering.
+    const interval = setInterval(() => ring.append("."), 30)
+    setTimeout(() => clearInterval(interval), 200)
+    const start = Date.now()
+    const result = await waitForTuiReady(ring, { hardCapMs: 2000, pollMs: 10, quietPeriodMs: 150 })
+    const elapsed = Date.now() - start
+    expect(result).toBe("marker")
+    // Render bursts (200 ms) + quiet period (150 ms) = at least 350 ms.
+    expect(elapsed).toBeGreaterThanOrEqual(300)
+  })
+
+  test("resolves immediately when ring stays quiet from the start", async () => {
+    const ring = new OutputRing()
+    ring.append("❯ ")
+    const start = Date.now()
+    const result = await waitForTuiReady(ring, { hardCapMs: 2000, pollMs: 10, quietPeriodMs: 80 })
+    const elapsed = Date.now() - start
+    expect(result).toBe("marker")
+    expect(elapsed).toBeGreaterThanOrEqual(70)
+    expect(elapsed).toBeLessThan(250)
+  })
+
+  test("exported TUI_READY_MARKER is the input-box prompt", () => {
+    expect(TUI_READY_MARKER).toBe("❯ ")
+  })
+
+  test("returns 'marker' when TUI renders ❯ with \\x1b[1C instead of space", async () => {
+    const ring = new OutputRing()
+    // Real TUI output: space after ❯ is cursor-forward-1, not a literal space
+    ring.append("❯\x1b[1C")
+    const result = await waitForTuiReady(ring, { hardCapMs: 1000, pollMs: 10 })
+    expect(result).toBe("marker")
+  })
+
+  test("returns 'marker' when TUI renders ❯ followed by U+00A0 (non-breaking space)", async () => {
+    const ring = new OutputRing()
+    // Real TUI output: ❯ is followed by NBSP (U+00A0), not a regular space
+    ring.append("❯ ")
+    const result = await waitForTuiReady(ring, { hardCapMs: 1000, pollMs: 10 })
+    expect(result).toBe("marker")
+  })
+})
+
+describe("dismissTrustDialogIfPresent (ANSI-encoded ring)", () => {
+  test("detects trust dialog when words separated by \\x1b[1C (TUI rendering)", async () => {
+    const pty = fakePty()
+    const ring = new OutputRing()
+    // Real TUI output: spaces rendered as cursor-forward-1 escape sequences
+    ring.append("\x1b[1C❯\x1b[1C1.\x1b[1CYes,\x1b[1CI\x1b[1Ctrust\x1b[1Cthis\x1b[1Cfolder\r\n")
+    const dismissed = await dismissTrustDialogIfPresent(pty, ring)
+    expect(dismissed).toBe(true)
+    expect(pty.sent).toEqual(["\r"])
+  })
+
+  test("does not false-trigger on plain text without trust marker", async () => {
+    const pty = fakePty()
+    const ring = new OutputRing()
+    ring.append("\x1b[1CWelcome\x1b[1Cback!\x1b[0m❯ ")
+    const dismissed = await dismissTrustDialogIfPresent(pty, ring)
+    expect(dismissed).toBe(false)
+    expect(pty.sent).toEqual([])
+  })
+})
+
+describe("waitForTuiReadyWithTrustDismiss", () => {
+  test("returns 'ready' immediately when input box already present", async () => {
+    const pty = fakePty()
+    const ring = new OutputRing()
+    ring.append("❯ ")
+    const result = await waitForTuiReadyWithTrustDismiss(pty, ring, { hardCapMs: 500, pollMs: 10 })
+    expect(result).toBe("ready")
+    expect(pty.sent).toEqual([])
+  })
+
+  test("dismisses ANSI trust dialog then resolves when input box appears", async () => {
+    const pty = fakePty()
+    const ring = new OutputRing()
+    // Trust dialog with ANSI-encoded text appears first
+    ring.append("\x1b[1Ctrust\x1b[1Cthis\x1b[1Cfolder")
+    // After 50ms simulate TUI loading: input box appears after dismiss sends \r
+    setTimeout(() => ring.append("❯ "), 80)
+    const result = await waitForTuiReadyWithTrustDismiss(pty, ring, { hardCapMs: 1000, pollMs: 10 })
+    expect(result).toBe("ready")
+    expect(pty.sent).toContain("\r")
+  })
+
+  test("returns 'timeout' when neither marker appears", async () => {
+    const pty = fakePty()
+    const ring = new OutputRing()
+    const result = await waitForTuiReadyWithTrustDismiss(pty, ring, { hardCapMs: 150, pollMs: 10 })
+    expect(result).toBe("timeout")
+    expect(pty.sent).toEqual([])
+  })
+
+  test("returns 'ready' when TUI renders ❯ with \\x1b[1C (ANSI cursor-forward)", async () => {
+    const pty = fakePty()
+    const ring = new OutputRing()
+    ring.append("❯\x1b[1C")
+    const result = await waitForTuiReadyWithTrustDismiss(pty, ring, { hardCapMs: 500, pollMs: 10 })
+    expect(result).toBe("ready")
+    expect(pty.sent).toEqual([])
+  })
+
+  test("does not false-trigger on trust dialog's own ❯ selection cursor", async () => {
+    // Trust dialog renders "❯ 1. Yes, I trust this folder" as ANSI:
+    // "❯\x1b[1C1.\x1b[1CYes,...trust\x1b[1Cthis\x1b[1Cfolder"
+    // stripAnsi gives "❯ 1. Yes, I trust this folder" which contains "❯ "
+    // → must NOT trigger "ready" while trust dialog is present
+    const pty = fakePty()
+    const ring = new OutputRing()
+    ring.append("\x1b[1C❯\x1b[1C1.\x1b[1CYes,\x1b[1CI\x1b[1Ctrust\x1b[1Cthis\x1b[1Cfolder\r\n")
+    setTimeout(() => ring.append("❯\x1b[1C"), 80)
+    const result = await waitForTuiReadyWithTrustDismiss(pty, ring, { hardCapMs: 1000, pollMs: 10 })
+    expect(result).toBe("ready")
+    expect(pty.sent).toContain("\r")
+  })
+
+  test("dismisses trust dialog only once even if marker persists in ring", async () => {
+    const pty = fakePty()
+    const ring = new OutputRing()
+    ring.append("\x1b[1Ctrust\x1b[1Cthis\x1b[1Cfolder")
+    setTimeout(() => ring.append("❯ "), 80)
+    await waitForTuiReadyWithTrustDismiss(pty, ring, { hardCapMs: 1000, pollMs: 10 })
+    // \r should appear exactly once (dismiss sent once, not repeated each poll)
+    expect(pty.sent.filter((s) => s === "\r")).toHaveLength(1)
+  })
+
+  // Regression: same race as waitForTuiReady — after trust dismiss the input
+  // box marker can render before Ink's keyboard handler mounts. Quiet-period
+  // gate must apply on this path too so prompts don't land into a discarding
+  // TUI under load.
+  test("waits for ring to stay quiet for quietPeriodMs after post-dismiss marker hit", async () => {
+    const pty = fakePty()
+    const ring = new OutputRing()
+    ring.append("\x1b[1Ctrust\x1b[1Cthis\x1b[1Cfolder")
+    // After dismiss, marker arrives at t=50, then more bytes for 200 ms.
+    setTimeout(() => ring.append("❯ "), 50)
+    const interval = setInterval(() => ring.append("."), 30)
+    setTimeout(() => clearInterval(interval), 250)
+    const start = Date.now()
+    const result = await waitForTuiReadyWithTrustDismiss(pty, ring, {
+      hardCapMs: 2000,
+      pollMs: 10,
+      quietPeriodMs: 150,
+    })
+    const elapsed = Date.now() - start
+    expect(result).toBe("ready")
+    // Bursts end ~250 ms + quiet 150 ms.
+    expect(elapsed).toBeGreaterThanOrEqual(350)
+  })
+})