@cuongtran001/kanna 0.43.2 → 0.45.0

package/src/server/agent.test.ts

@@ -7,6 +7,7 @@ import {
   normalizeClaudeStreamMessage,
   normalizeClaudeUsageSnapshot,
 } from "./agent"
+import { BackgroundTaskRegistry } from "./background-tasks"
 import type { HarnessTurn } from "./harness-types"
 import type { ChatAttachment, SlashCommand, TranscriptEntry } from "../shared/types"
 import type { AutoContinueEvent } from "./auto-continue/events"
@@ -1560,6 +1561,93 @@ describe("AgentCoordinator claude integration", () => {
     events.close()
   })
 
+  test("Claude steer + result without echoed cancel still clears running state", async () => {
+    // Repro: sometimes the underlying SDK closes its stream cleanly on cancel
+    // and never emits a `result.subtype=cancelled` (which would map to an
+    // `interrupted` entry). When the next prompt's `result` arrives the
+    // session.pendingPromptSeqs queue still holds the orphaned cancelled seq,
+    // so the FIFO shift returns the wrong seq and `activeTurns` is never
+    // cleared — leaving the UI stuck on "Running...".
+    const events = new AsyncEventQueue<any>()
+    const prompts: string[] = []
+
+    const store = createFakeStore()
+    await store.enqueueMessage("chat-1", {
+      id: "queued-1",
+      content: "follow up",
+      attachments: [],
+      provider: "claude",
+      model: "claude-opus-4-1",
+      planMode: false,
+    })
+
+    const coordinator = new AgentCoordinator({
+      store: store as never,
+      onStateChange: () => {},
+      startClaudeSession: async () => ({
+        provider: "claude",
+        stream: events,
+        getAccountInfo: async () => null,
+        interrupt: async () => {},
+        close: () => {},
+        setModel: async () => {},
+        setPermissionMode: async () => {},
+        getSupportedCommands: async () => [],
+        sendPrompt: async (content: string) => {
+          prompts.push(content)
+          if (prompts.length === 1) {
+            events.push({
+              type: "transcript" as const,
+              entry: timestamped({
+                kind: "system_init",
+                provider: "claude",
+                model: "claude-opus-4-1",
+                tools: [],
+                agents: [],
+                slashCommands: [],
+                mcpServers: [],
+              }),
+            })
+          }
+        },
+      }),
+    })
+
+    await coordinator.send({
+      type: "chat.send",
+      chatId: "chat-1",
+      provider: "claude",
+      content: "first prompt",
+      model: "claude-opus-4-1",
+    })
+
+    await coordinator.steer({
+      type: "message.steer",
+      chatId: "chat-1",
+      queuedMessageId: "queued-1",
+    })
+
+    expect(prompts).toHaveLength(2)
+
+    // SDK never echoes a cancelled result for the first prompt — only the
+    // result for the steered second prompt arrives.
+    events.push({
+      type: "transcript" as const,
+      entry: timestamped({
+        kind: "result",
+        subtype: "success",
+        isError: false,
+        durationMs: 0,
+        result: "done",
+      }),
+    })
+
+    await waitFor(() => !coordinator.getActiveStatuses().has("chat-1"))
+    expect(coordinator.getActiveStatuses().has("chat-1")).toBe(false)
+
+    events.close()
+  })
+
   test("uses Claude forkSession when starting a forked chat", async () => {
     const startSessionCalls: Array<{ sessionToken: string | null; forkSession: boolean }> = []
     const events = new AsyncEventQueue<any>()
@@ -2367,3 +2455,558 @@ describe("AgentCoordinator.listLiveSchedules", () => {
     expect(live.sort()).toEqual(["sched-proposed", "sched-scheduled"].sort())
   })
 })
+
+// ── AgentCoordinator: BackgroundTaskRegistry integration ──
+
+describe("AgentCoordinator background task registry", () => {
+  function makeCoordinatorWithRegistry(registry: BackgroundTaskRegistry) {
+    const store = createFakeStore()
+    const coordinator = new AgentCoordinator({
+      store: store as never,
+      onStateChange: () => {},
+      backgroundTasks: registry,
+    })
+    return { store, coordinator }
+  }
+
+  test("registers a bash_shell task on tool_result when run_in_background is true", () => {
+    const registry = new BackgroundTaskRegistry()
+    const { coordinator } = makeCoordinatorWithRegistry(registry)
+    const chatId = "chat-1"
+    const toolId = "tool-bg-1"
+
+    // Simulate tool_call with run_in_background: true
+    const toolCallEntry = timestamped({
+      kind: "tool_call",
+      tool: {
+        kind: "tool",
+        toolKind: "bash",
+        toolName: "Bash",
+        toolId,
+        input: {
+          command: "bun run dev",
+          runInBackground: true,
+        },
+        rawInput: { command: "bun run dev", run_in_background: true },
+      },
+    })
+
+    // Simulate tool_result with text containing "pid 12345"
+    const toolResultEntry = timestamped({
+      kind: "tool_result",
+      toolId,
+      content: "Started process pid: 12345\nServer running on port 3000",
+    })
+
+    // Invoke trackBashToolEntry via the public appendMessage path is not
+    // available directly; call the private method via casting.
+    const coord = coordinator as unknown as {
+      trackBashToolEntry(chatId: string, entry: TranscriptEntry): void
+    }
+    coord.trackBashToolEntry(chatId, toolCallEntry)
+    coord.trackBashToolEntry(chatId, toolResultEntry)
+
+    expect(registry.list()).toHaveLength(1)
+    const task = registry.list()[0]
+    expect(task?.kind).toBe("bash_shell")
+    if (task?.kind !== "bash_shell") throw new Error("unexpected task kind")
+    expect(task.chatId).toBe(chatId)
+    expect(task.command).toBe("bun run dev")
+    expect(task.pid).toBe(12345)
+    expect(task.status).toBe("running")
+  })
+
+  test("does not register when run_in_background is false", () => {
+    const registry = new BackgroundTaskRegistry()
+    const { coordinator } = makeCoordinatorWithRegistry(registry)
+    const chatId = "chat-1"
+    const toolId = "tool-fg-1"
+
+    const coord = coordinator as unknown as {
+      trackBashToolEntry(chatId: string, entry: TranscriptEntry): void
+    }
+
+    coord.trackBashToolEntry(chatId, timestamped({
+      kind: "tool_call",
+      tool: {
+        kind: "tool",
+        toolKind: "bash",
+        toolName: "Bash",
+        toolId,
+        input: {
+          command: "ls",
+          runInBackground: false,
+        },
+        rawInput: { command: "ls" },
+      },
+    }))
+    coord.trackBashToolEntry(chatId, timestamped({
+      kind: "tool_result",
+      toolId,
+      content: "file.txt",
+    }))
+
+    expect(registry.list()).toHaveLength(0)
+  })
+
+  test("does not register when backgroundTasks is not provided", () => {
+    const store = createFakeStore()
+    const coordinator = new AgentCoordinator({
+      store: store as never,
+      onStateChange: () => {},
+      // no backgroundTasks
+    })
+    const toolId = "tool-no-reg-1"
+
+    const coord = coordinator as unknown as {
+      trackBashToolEntry(chatId: string, entry: TranscriptEntry): void
+    }
+
+    // Should not throw
+    coord.trackBashToolEntry("chat-1", timestamped({
+      kind: "tool_call",
+      tool: {
+        kind: "tool",
+        toolKind: "bash",
+        toolName: "Bash",
+        toolId,
+        input: { command: "sleep 999", runInBackground: true },
+        rawInput: { command: "sleep 999", run_in_background: true },
+      },
+    }))
+    coord.trackBashToolEntry("chat-1", timestamped({
+      kind: "tool_result",
+      toolId,
+      content: "pid: 99999",
+    }))
+    // No registry to check — just verify no exception was thrown
+  })
+
+  test("uses toolId as shellId fallback when output has no shell_id text", () => {
+    const registry = new BackgroundTaskRegistry()
+    const { coordinator } = makeCoordinatorWithRegistry(registry)
+    const toolId = "tool-no-shell-id"
+
+    const coord = coordinator as unknown as {
+      trackBashToolEntry(chatId: string, entry: TranscriptEntry): void
+    }
+
+    coord.trackBashToolEntry("chat-1", timestamped({
+      kind: "tool_call",
+      tool: {
+        kind: "tool",
+        toolKind: "bash",
+        toolName: "Bash",
+        toolId,
+        input: { command: "tail -f /tmp/log.txt", runInBackground: true },
+        rawInput: { command: "tail -f /tmp/log.txt", run_in_background: true },
+      },
+    }))
+    coord.trackBashToolEntry("chat-1", timestamped({
+      kind: "tool_result",
+      toolId,
+      content: "Tailing log file…",
+    }))
+
+    const tasks = registry.list()
+    expect(tasks).toHaveLength(1)
+    const task = tasks[0]
+    if (task?.kind !== "bash_shell") throw new Error("unexpected task kind")
+    expect(task.shellId).toBe(toolId)
+    expect(task.pid).toBeNull()
+  })
+
+  test("uses SDK backgroundTaskId from structured content array when present", () => {
+    const registry = new BackgroundTaskRegistry()
+    const { coordinator } = makeCoordinatorWithRegistry(registry)
+    const toolId = "tool-struct-bg-1"
+
+    const coord = coordinator as unknown as {
+      trackBashToolEntry(chatId: string, entry: TranscriptEntry): void
+    }
+
+    coord.trackBashToolEntry("chat-1", timestamped({
+      kind: "tool_call",
+      tool: {
+        kind: "tool",
+        toolKind: "bash",
+        toolName: "Bash",
+        toolId,
+        input: { command: "bun run dev", runInBackground: true },
+        rawInput: { command: "bun run dev", run_in_background: true },
+      },
+    }))
+
+    // Simulate a tool_result where content is an array of content blocks
+    // and one block carries the SDK's canonical backgroundTaskId field.
+    const structuredContent: Array<{ type: string; text?: string; backgroundTaskId?: string }> = [
+      { type: "text", text: "Started pid: 99" },
+      { type: "tool_result", backgroundTaskId: "shell_abc123" },
+    ]
+    coord.trackBashToolEntry("chat-1", timestamped({
+      kind: "tool_result",
+      toolId,
+      content: structuredContent,
+    }))
+
+    const tasks = registry.list()
+    expect(tasks).toHaveLength(1)
+    const task = tasks[0]
+    if (task?.kind !== "bash_shell") throw new Error("unexpected task kind")
+    // Structured backgroundTaskId must win over regex-parsed shell_id or toolId fallback
+    expect(task.shellId).toBe("shell_abc123")
+  })
+
+  test("uses SDK backgroundTaskId from direct BashOutput object when present", () => {
+    const registry = new BackgroundTaskRegistry()
+    const { coordinator } = makeCoordinatorWithRegistry(registry)
+    const toolId = "tool-direct-bg-1"
+
+    const coord = coordinator as unknown as {
+      trackBashToolEntry(chatId: string, entry: TranscriptEntry): void
+    }
+
+    coord.trackBashToolEntry("chat-1", timestamped({
+      kind: "tool_call",
+      tool: {
+        kind: "tool",
+        toolKind: "bash",
+        toolName: "Bash",
+        toolId,
+        input: { command: "sleep 300", runInBackground: true },
+        rawInput: { command: "sleep 300", run_in_background: true },
+      },
+    }))
+
+    // Simulate a BashOutput object as content (direct-object form)
+    const bashOutputContent = {
+      stdout: "Started pid: 55555",
+      stderr: "",
+      interrupted: false,
+      backgroundTaskId: "shell_direct789",
+    }
+    coord.trackBashToolEntry("chat-1", timestamped({
+      kind: "tool_result",
+      toolId,
+      content: bashOutputContent,
+    }))
+
+    const tasks = registry.list()
+    expect(tasks).toHaveLength(1)
+    const task = tasks[0]
+    if (task?.kind !== "bash_shell") throw new Error("unexpected task kind")
+    expect(task.shellId).toBe("shell_direct789")
+  })
+
+  // ── draining_stream registry tracking ──
+
+  function makeDrainingStreamSetup() {
+    let resolveStream!: () => void
+    const fakeCodexManager = {
+      async startSession() {},
+      async startTurn(): Promise<HarnessTurn> {
+        async function* stream() {
+          yield {
+            type: "transcript" as const,
+            entry: timestamped({
+              kind: "system_init",
+              provider: "codex" as const,
+              model: "gpt-5.4",
+              tools: [],
+              agents: [],
+              slashCommands: [],
+              mcpServers: [],
+            }),
+          }
+          yield {
+            type: "transcript" as const,
+            entry: timestamped({
+              kind: "result",
+              subtype: "success" as const,
+              isError: false,
+              durationMs: 0,
+              result: "done",
+            }),
+          }
+          await new Promise<void>((resolve) => {
+            resolveStream = resolve
+          })
+        }
+        return {
+          provider: "codex" as const,
+          stream: stream(),
+          interrupt: async () => {},
+          close: () => {
+            resolveStream?.()
+          },
+        }
+      },
+    }
+    return { fakeCodexManager, resolveStream: () => resolveStream() }
+  }
+
+  test("registers a draining_stream entry in the registry when a turn result arrives with stream still open", async () => {
+    const registry = new BackgroundTaskRegistry()
+    const { fakeCodexManager } = makeDrainingStreamSetup()
+    const store = createFakeStore()
+    const coordinator = new AgentCoordinator({
+      store: store as never,
+      onStateChange: () => {},
+      backgroundTasks: registry,
+      codexManager: fakeCodexManager as never,
+    })
+
+    await coordinator.send({
+      type: "chat.send",
+      chatId: "chat-1",
+      provider: "codex",
+      content: "run with bg task",
+    })
+
+    await waitFor(() => coordinator.getDrainingChatIds().has("chat-1"))
+
+    const tasks = registry.list()
+    expect(tasks).toHaveLength(1)
+    const task = tasks[0]
+    if (task?.kind !== "draining_stream") throw new Error("unexpected task kind")
+    expect(task.id).toBe("drain:chat-1")
+    expect(task.chatId).toBe("chat-1")
+
+    // Clean up the hanging stream
+    await coordinator.stopDraining("chat-1")
+  })
+
+  test("unregisters the draining_stream entry when stopDraining is called directly", async () => {
+    const registry = new BackgroundTaskRegistry()
+    const { fakeCodexManager } = makeDrainingStreamSetup()
+    const store = createFakeStore()
+    const coordinator = new AgentCoordinator({
+      store: store as never,
+      onStateChange: () => {},
+      backgroundTasks: registry,
+      codexManager: fakeCodexManager as never,
+    })
+
+    await coordinator.send({
+      type: "chat.send",
+      chatId: "chat-1",
+      provider: "codex",
+      content: "run with bg task",
+    })
+
+    await waitFor(() => coordinator.getDrainingChatIds().has("chat-1"))
+    expect(registry.list()).toHaveLength(1)
+
+    await coordinator.stopDraining("chat-1")
+
+    expect(coordinator.getDrainingChatIds().has("chat-1")).toBe(false)
+    expect(registry.list()).toHaveLength(0)
+  })
+
+  test("registry.stop('drain:CHATID') invokes closeStream strategy and clears both drainingStreams and registry", async () => {
+    const registry = new BackgroundTaskRegistry()
+    const { fakeCodexManager } = makeDrainingStreamSetup()
+    const store = createFakeStore()
+    const coordinator = new AgentCoordinator({
+      store: store as never,
+      onStateChange: () => {},
+      backgroundTasks: registry,
+      codexManager: fakeCodexManager as never,
+    })
+
+    await coordinator.send({
+      type: "chat.send",
+      chatId: "chat-1",
+      provider: "codex",
+      content: "run with bg task",
+    })
+
+    await waitFor(() => coordinator.getDrainingChatIds().has("chat-1"))
+    expect(registry.list()).toHaveLength(1)
+
+    const result = await registry.stop("drain:chat-1")
+
+    expect(result).toEqual({ ok: true, method: "close" })
+    expect(coordinator.getDrainingChatIds().has("chat-1")).toBe(false)
+    expect(registry.list()).toHaveLength(0)
+  })
+
+  test("clearDrainingStream: natural completion clears both drainingStreams and registry", async () => {
+    // Drives a turn through to its stream finally block (natural completion —
+    // the most common path). We use the private-method-cast pattern to exercise
+    // clearDrainingStream directly because driving the full async stream through
+    // its finally block in a reliable, race-free way would require
+    // substantially more scaffolding than the existing helpers provide.
+    const registry = new BackgroundTaskRegistry()
+    const store = createFakeStore()
+    const coordinator = new AgentCoordinator({
+      store: store as never,
+      onStateChange: () => {},
+      backgroundTasks: registry,
+    })
+
+    // Seed the registry with a draining_stream entry for chat-1 to simulate
+    // the state that exists just before the finally block runs.
+    const coord = coordinator as unknown as {
+      drainingStreams: Map<string, { turn: { close(): void } }>
+      clearDrainingStream(chatId: string): void
+    }
+    coord.drainingStreams.set("chat-1", { turn: { close: () => {} } })
+    registry.register({
+      kind: "draining_stream",
+      id: "drain:chat-1",
+      chatId: "chat-1",
+      startedAt: Date.now(),
+      lastOutput: "",
+    })
+
+    expect(coord.drainingStreams.has("chat-1")).toBe(true)
+    expect(registry.list()).toHaveLength(1)
+
+    // Call clearDrainingStream — this is what the stream finally block now calls.
+    coord.clearDrainingStream("chat-1")
+
+    expect(coord.drainingStreams.has("chat-1")).toBe(false)
+    expect(registry.list()).toHaveLength(0)
+  })
+
+  test("clearDrainingStream: startTurnForChat clears stale draining_stream for same chat", async () => {
+    // Verifies clearDrainingStream correctly removes both the Map entry and the
+    // registry entry when a new turn starts on a chat that already has a
+    // draining stream. Uses the private-method cast for the same reason as the
+    // natural-completion test above.
+    const registry = new BackgroundTaskRegistry()
+    const store = createFakeStore()
+    const coordinator = new AgentCoordinator({
+      store: store as never,
+      onStateChange: () => {},
+      backgroundTasks: registry,
+    })
+
+    const coord = coordinator as unknown as {
+      drainingStreams: Map<string, { turn: { close(): void } }>
+      clearDrainingStream(chatId: string): void
+    }
+
+    // Seed state: a previous turn left a draining stream for chat-1.
+    coord.drainingStreams.set("chat-1", { turn: { close: () => {} } })
+    registry.register({
+      kind: "draining_stream",
+      id: "drain:chat-1",
+      chatId: "chat-1",
+      startedAt: Date.now(),
+      lastOutput: "",
+    })
+
+    expect(registry.list()).toHaveLength(1)
+
+    // A new turn starting for chat-1 calls clearDrainingStream before proceeding.
+    coord.clearDrainingStream("chat-1")
+
+    expect(coord.drainingStreams.has("chat-1")).toBe(false)
+    expect(registry.list()).toHaveLength(0)
+  })
+
+  test("clearDrainingStream: cancel() clears draining_stream from registry", async () => {
+    // Verifies cancel()'s cleanup path (clearDrainingStream) removes the
+    // registry entry. Uses the private-method cast because wiring cancel()
+    // through to a live draining stream requires the full harness scaffolding.
+    const registry = new BackgroundTaskRegistry()
+    const store = createFakeStore()
+    const coordinator = new AgentCoordinator({
+      store: store as never,
+      onStateChange: () => {},
+      backgroundTasks: registry,
+    })
+
+    const coord = coordinator as unknown as {
+      drainingStreams: Map<string, { turn: { close(): void } }>
+      clearDrainingStream(chatId: string): void
+    }
+
+    // Seed state: a result event left a draining_stream registered for chat-1.
+    coord.drainingStreams.set("chat-1", { turn: { close: () => {} } })
+    registry.register({
+      kind: "draining_stream",
+      id: "drain:chat-1",
+      chatId: "chat-1",
+      startedAt: Date.now(),
+      lastOutput: "",
+    })
+
+    expect(registry.list()).toHaveLength(1)
+
+    // cancel() calls clearDrainingStream when a draining entry exists.
+    coord.clearDrainingStream("chat-1")
+
+    expect(coord.drainingStreams.has("chat-1")).toBe(false)
+    expect(registry.list()).toHaveLength(0)
+  })
+})
+
+describe("parseBackgroundPid regex variants", () => {
+  // These tests exercise parseBackgroundPid indirectly via trackBashToolEntry,
+  // verifying the regex accepts all supported PID output formats.
+  function makeRegistryAndCoord() {
+    const registry = new BackgroundTaskRegistry()
+    const store = createFakeStore()
+    const coordinator = new AgentCoordinator({
+      store: store as never,
+      onStateChange: () => {},
+      backgroundTasks: registry,
+    })
+    const coord = coordinator as unknown as {
+      trackBashToolEntry(chatId: string, entry: TranscriptEntry): void
+    }
+    return { registry, coord }
+  }
+
+  function registerWithOutput(output: string, toolId: string) {
+    const { registry, coord } = makeRegistryAndCoord()
+    coord.trackBashToolEntry("chat-1", timestamped({
+      kind: "tool_call",
+      tool: {
+        kind: "tool",
+        toolKind: "bash",
+        toolName: "Bash",
+        toolId,
+        input: { command: "sleep 60", runInBackground: true },
+        rawInput: { command: "sleep 60", run_in_background: true },
+      },
+    }))
+    coord.trackBashToolEntry("chat-1", timestamped({
+      kind: "tool_result",
+      toolId,
+      content: output,
+    }))
+    return registry
+  }
+
+  test("parses 'pid: N' format", () => {
+    const registry = registerWithOutput("Started pid: 11111", "t1")
+    const task = registry.list()[0]
+    if (task?.kind !== "bash_shell") throw new Error("unexpected task kind")
+    expect(task.pid).toBe(11111)
+  })
+
+  test("parses 'pid N' format (space only)", () => {
+    const registry = registerWithOutput("Server pid 22222 running", "t2")
+    const task = registry.list()[0]
+    if (task?.kind !== "bash_shell") throw new Error("unexpected task kind")
+    expect(task.pid).toBe(22222)
+  })
+
+  test("parses 'PID=N' format", () => {
+    const registry = registerWithOutput("PID=33333", "t3")
+    const task = registry.list()[0]
+    if (task?.kind !== "bash_shell") throw new Error("unexpected task kind")
+    expect(task.pid).toBe(33333)
+  })
+
+  test("returns null when no pid in output", () => {
+    const registry = registerWithOutput("No process info here", "t4")
+    const task = registry.list()[0]
+    if (task?.kind !== "bash_shell") throw new Error("unexpected task kind")
+    expect(task.pid).toBeNull()
+  })
+})