Simplify Claude agent session flow

packages/server/src/server/agent/providers/claude-agent.integration.test.ts

@@ -0,0 +1,414 @@
+import { describe, expect, test, beforeAll } from "vitest";
+import { mkdtempSync, readFileSync, rmSync } from "node:fs";
+import { tmpdir } from "node:os";
+import path from "node:path";
+import pino from "pino";
+
+import type {
+  AgentSession,
+  AgentStreamEvent,
+  ToolCallTimelineItem,
+} from "../agent-sdk-types.js";
+import { isCommandAvailable } from "../provider-launch-config.js";
+import { ClaudeAgentClient } from "./claude-agent.js";
+
+const logger = pino({ level: "silent" });
+const client = new ClaudeAgentClient({ logger });
+
+function tmpCwd(prefix: string): string {
+  return mkdtempSync(path.join(tmpdir(), prefix));
+}
+
+function compactText(value: string): string {
+  return value.replace(/\s+/g, "").toLowerCase();
+}
+
+function isTerminalEvent(event: AgentStreamEvent): boolean {
+  return (
+    event.type === "turn_completed" ||
+    event.type === "turn_failed" ||
+    event.type === "turn_canceled"
+  );
+}
+
+async function nextStreamEvent(
+  stream: AsyncGenerator<AgentStreamEvent>,
+  timeoutMs: number,
+  label: string
+): Promise<IteratorResult<AgentStreamEvent>> {
+  let timer: ReturnType<typeof setTimeout> | null = null;
+  try {
+    return await Promise.race([
+      stream.next(),
+      new Promise<never>((_, reject) => {
+        timer = setTimeout(() => {
+          reject(new Error(`Timed out waiting for ${label}`));
+        }, timeoutMs);
+      }),
+    ]);
+  } finally {
+    if (timer) {
+      clearTimeout(timer);
+    }
+  }
+}
+
+async function collectUntilTerminal(
+  stream: AsyncGenerator<AgentStreamEvent>,
+  options?: {
+    timeoutMs?: number;
+    onEvent?: (event: AgentStreamEvent) => Promise<void> | void;
+  }
+): Promise<AgentStreamEvent[]> {
+  const events: AgentStreamEvent[] = [];
+  while (true) {
+    const next = await nextStreamEvent(
+      stream,
+      options?.timeoutMs ?? 45_000,
+      "stream event"
+    );
+    if (next.done || !next.value) {
+      return events;
+    }
+    const event = next.value;
+    events.push(event);
+    await options?.onEvent?.(event);
+    if (isTerminalEvent(event)) {
+      return events;
+    }
+  }
+}
+
+async function collectUntil(
+  stream: AsyncGenerator<AgentStreamEvent>,
+  predicate: (event: AgentStreamEvent) => boolean,
+  timeoutMs = 45_000
+): Promise<AgentStreamEvent[]> {
+  const events: AgentStreamEvent[] = [];
+  while (true) {
+    const next = await nextStreamEvent(stream, timeoutMs, "matching stream event");
+    if (next.done || !next.value) {
+      throw new Error("Stream ended before the expected event arrived");
+    }
+    const event = next.value;
+    events.push(event);
+    if (predicate(event) || isTerminalEvent(event)) {
+      return events;
+    }
+  }
+}
+
+function getAssistantText(events: AgentStreamEvent[]): string {
+  return events
+    .flatMap((event) => {
+      if (event.type !== "timeline" || event.item.type !== "assistant_message") {
+        return [];
+      }
+      return [event.item.text];
+    })
+    .join("\n");
+}
+
+function getToolCalls(events: AgentStreamEvent[]): ToolCallTimelineItem[] {
+  return events.flatMap((event) => {
+    if (event.type !== "timeline" || event.item.type !== "tool_call") {
+      return [];
+    }
+    return [event.item];
+  });
+}
+
+function getLatestCompletedBashCall(
+  events: AgentStreamEvent[]
+): ToolCallTimelineItem | undefined {
+  return [...getToolCalls(events)]
+    .reverse()
+    .find(
+      (item) =>
+        item.status === "completed" &&
+        item.name.toLowerCase() === "bash"
+    );
+}
+
+function getInternalQuery(session: AgentSession): unknown {
+  return (session as AgentSession & { query?: unknown | null }).query ?? null;
+}
+
+async function createSession(params?: {
+  cwdPrefix?: string;
+  modeId?: string;
+  title?: string;
+}): Promise<{ cwd: string; session: AgentSession }> {
+  const cwd = tmpCwd(params?.cwdPrefix ?? "claude-agent-integration-");
+  const session = await client.createSession({
+    provider: "claude",
+    cwd,
+    title: params?.title ?? "ClaudeAgentSession integration",
+    modeId: params?.modeId ?? "acceptEdits",
+    model: "haiku",
+  });
+  return { cwd, session };
+}
+
+async function cleanupSession(handle: {
+  cwd: string;
+  session: AgentSession;
+}): Promise<void> {
+  await handle.session.close().catch(() => undefined);
+  rmSync(handle.cwd, { recursive: true, force: true });
+}
+
+describe("ClaudeAgentSession integration", () => {
+  beforeAll(() => {
+    expect(isCommandAvailable("claude")).toBe(true);
+  });
+
+  test(
+    "streams a basic response turn end-to-end",
+    async () => {
+      const handle = await createSession({
+        cwdPrefix: "claude-agent-basic-response-",
+      });
+
+      try {
+        const events = await collectUntilTerminal(
+          handle.session.stream("Respond with exactly: HELLO_WORLD")
+        );
+
+        expect(events[0]).toMatchObject({
+          type: "turn_started",
+          provider: "claude",
+        });
+        expect(
+          events.some(
+            (event) =>
+              event.type === "timeline" &&
+              event.item.type === "assistant_message" &&
+              compactText(event.item.text).includes("hello_world")
+          )
+        ).toBe(true);
+        expect(events.at(-1)).toMatchObject({
+          type: "turn_completed",
+          provider: "claude",
+        });
+      } finally {
+        await cleanupSession(handle);
+      }
+    },
+    60_000
+  );
+
+  test(
+    "runs a real Bash tool call and completes it",
+    async () => {
+      const handle = await createSession({
+        cwdPrefix: "claude-agent-basic-tool-",
+      });
+
+      try {
+        const events = await collectUntilTerminal(
+          handle.session.stream(
+            [
+              "Use the Bash tool.",
+              "Run exactly: echo TOOL_TEST_OUTPUT",
+              "After the command completes, reply with exactly: TOOL_DONE",
+            ].join(" ")
+          )
+        );
+
+        const bashCalls = getToolCalls(events).filter(
+          (item) => item.name.toLowerCase() === "bash"
+        );
+        const completedBashCall = getLatestCompletedBashCall(events);
+
+        expect(bashCalls.length).toBeGreaterThan(0);
+        expect(completedBashCall).toBeDefined();
+        expect(completedBashCall?.detail.type).toBe("shell");
+        expect(
+          completedBashCall?.detail.type === "shell" &&
+            completedBashCall.detail.output?.includes("TOOL_TEST_OUTPUT")
+        ).toBe(true);
+        expect(compactText(getAssistantText(events))).toContain("tool_done");
+        expect(events.at(-1)).toMatchObject({
+          type: "turn_completed",
+          provider: "claude",
+        });
+      } finally {
+        await cleanupSession(handle);
+      }
+    },
+    60_000
+  );
+
+  test(
+    "interrupts a running Bash turn and continues on the same query",
+    async () => {
+      const handle = await createSession({
+        cwdPrefix: "claude-agent-interrupt-continue-",
+      });
+
+      try {
+        const firstStream = handle.session.stream(
+          [
+            "Use the Bash tool.",
+            "Run exactly: sleep 10",
+            "Do not use a background task.",
+            "Do not do anything after starting the command.",
+          ].join(" ")
+        );
+
+        const initialEvents = await collectUntil(
+          firstStream,
+          (event) =>
+            event.type === "timeline" &&
+            event.item.type === "tool_call" &&
+            event.item.name.toLowerCase() === "bash",
+          45_000
+        );
+        const firstQuery = getInternalQuery(handle.session);
+
+        expect(firstQuery).toBeTruthy();
+
+        await handle.session.interrupt();
+
+        const canceledEvents = await collectUntilTerminal(firstStream, {
+          timeoutMs: 20_000,
+        });
+        const allFirstTurnEvents = [...initialEvents, ...canceledEvents];
+
+        expect(
+          allFirstTurnEvents.some(
+            (event) =>
+              event.type === "turn_canceled" && event.provider === "claude"
+          )
+        ).toBe(true);
+
+        const followUpEvents = await collectUntilTerminal(
+          handle.session.stream("Respond with exactly: AFTER_INTERRUPT_OK")
+        );
+        const secondQuery = getInternalQuery(handle.session);
+
+        expect(secondQuery).toBe(firstQuery);
+        expect(compactText(getAssistantText(followUpEvents))).toContain(
+          "after_interrupt_ok"
+        );
+        expect(followUpEvents.at(-1)).toMatchObject({
+          type: "turn_completed",
+          provider: "claude",
+        });
+      } finally {
+        await cleanupSession(handle);
+      }
+    },
+    60_000
+  );
+
+  test(
+    "creates an autonomous live turn when a background task completes",
+    async () => {
+      const handle = await createSession({
+        cwdPrefix: "claude-agent-autonomous-",
+      });
+      const autonomousWakeToken = `AUTONOMOUS_WAKE_${Date.now().toString(36)}`;
+
+      try {
+        const liveEventsStream = handle.session.streamLiveEvents();
+        const foregroundEvents = await collectUntilTerminal(
+          handle.session.stream(
+            [
+              "Use the Task tool to start a background sub-agent.",
+              "In that task, run the Bash command exactly: sleep 3 && echo BACKGROUND_DONE",
+              "Do not wait for task completion.",
+              "Reply immediately with exactly: SPAWNED",
+              `When the background task completes later, reply with exactly: ${autonomousWakeToken}`,
+            ].join(" ")
+          ),
+          { timeoutMs: 45_000 }
+        );
+
+        expect(compactText(getAssistantText(foregroundEvents))).toContain("spawned");
+
+        const liveEvents = await collectUntilTerminal(liveEventsStream, {
+          timeoutMs: 45_000,
+        });
+
+        expect(
+          liveEvents.some(
+            (event) => event.type === "turn_started" && event.provider === "claude"
+          )
+        ).toBe(true);
+        expect(
+          compactText(getAssistantText(liveEvents))
+        ).toContain(autonomousWakeToken.toLowerCase());
+        expect(liveEvents.at(-1)).toMatchObject({
+          type: "turn_completed",
+          provider: "claude",
+        });
+      } finally {
+        await cleanupSession(handle);
+      }
+    },
+    60_000
+  );
+
+  test(
+    "surfaces permission requests and resumes after approval",
+    async () => {
+      const handle = await createSession({
+        cwdPrefix: "claude-agent-permission-",
+        modeId: "default",
+      });
+      const permissionFile = path.join(handle.cwd, "permission.txt");
+
+      try {
+        const events = await collectUntilTerminal(
+          handle.session.stream(
+            [
+              "Use the Bash tool to run exactly: printf 'PERM_TEST' > permission.txt",
+              "If approval is required, wait for approval.",
+              "After the command succeeds, reply with exactly: PERM_DONE",
+            ].join(" ")
+          ),
+          {
+            timeoutMs: 45_000,
+            onEvent: async (event) => {
+              if (event.type !== "permission_requested") {
+                return;
+              }
+              await handle.session.respondToPermission(event.request.id, {
+                behavior: "allow",
+              });
+            },
+          }
+        );
+
+        const permissionRequest = events.find(
+          (event): event is Extract<AgentStreamEvent, { type: "permission_requested" }> =>
+            event.type === "permission_requested"
+        );
+        const permissionResolved = events.find(
+          (event): event is Extract<AgentStreamEvent, { type: "permission_resolved" }> =>
+            event.type === "permission_resolved"
+        );
+        const completedBashCall = getLatestCompletedBashCall(events);
+
+        expect(permissionRequest?.request.kind).toBe("tool");
+        expect(permissionResolved).toMatchObject({
+          type: "permission_resolved",
+          provider: "claude",
+          resolution: { behavior: "allow" },
+        });
+        expect(completedBashCall).toBeDefined();
+        expect(readFileSync(permissionFile, "utf8")).toBe("PERM_TEST");
+        expect(compactText(getAssistantText(events))).toContain("perm_done");
+        expect(events.at(-1)).toMatchObject({
+          type: "turn_completed",
+          provider: "claude",
+        });
+      } finally {
+        await cleanupSession(handle);
+      }
+    },
+    60_000
+  );
+});

packages/server/src/server/agent/providers/claude/sidechain-tracker.ts

@@ -0,0 +1,324 @@
+import type { SDKMessage } from "@anthropic-ai/claude-agent-sdk";
+
+import { mapClaudeRunningToolCall } from "./tool-call-mapper.js";
+import { buildToolCallDisplayModel } from "../../../../shared/tool-call-display.js";
+
+import type {
+  AgentMetadata,
+  AgentStreamEvent,
+  AgentTimelineItem,
+} from "../../agent-sdk-types.js";
+
+type ClaudeContentChunk = { type: string; [key: string]: unknown };
+
+type SubAgentActionEntry = {
+  index: number;
+  toolName: string;
+  summary?: string;
+};
+
+type SubAgentActivityState = {
+  subAgentType?: string;
+  description?: string;
+  actions: SubAgentActionEntry[];
+  actionKeys: string[];
+  nextActionIndex: number;
+  actionIndexByKey: Map<string, number>;
+};
+
+type SubAgentActionCandidate = {
+  key: string;
+  toolName: string;
+  input: unknown;
+};
+
+const MAX_SUB_AGENT_LOG_ENTRIES = 200;
+const MAX_SUB_AGENT_SUMMARY_CHARS = 160;
+
+function readTrimmedString(value: unknown): string | undefined {
+  if (typeof value !== "string") {
+    return undefined;
+  }
+  const trimmed = value.trim();
+  return trimmed.length > 0 ? trimmed : undefined;
+}
+
+function isClaudeContentChunk(value: unknown): value is ClaudeContentChunk {
+  return Boolean(
+    value &&
+      typeof value === "object" &&
+      typeof (value as { type?: unknown }).type === "string"
+  );
+}
+
+export class ClaudeSidechainTracker {
+  private readonly activeSidechains = new Map<string, SubAgentActivityState>();
+  private readonly getToolInput: (toolUseId: string) => AgentMetadata | null | undefined;
+
+  constructor(input: {
+    getToolInput: (toolUseId: string) => AgentMetadata | null | undefined;
+  }) {
+    this.getToolInput = input.getToolInput;
+  }
+
+  handleMessage(message: SDKMessage, parentToolUseId: string): AgentStreamEvent[] {
+    const state =
+      this.activeSidechains.get(parentToolUseId) ??
+      ({
+        actions: [],
+        actionKeys: [],
+        nextActionIndex: 1,
+        actionIndexByKey: new Map<string, number>(),
+      } satisfies SubAgentActivityState);
+    this.activeSidechains.set(parentToolUseId, state);
+
+    const contextUpdated = this.updateSubAgentContextFromTaskInput(
+      state,
+      parentToolUseId
+    );
+    const actionCandidates = this.extractSubAgentActionCandidates(message);
+    let actionUpdated = false;
+    for (const action of actionCandidates) {
+      if (this.appendSubAgentAction(state, action)) {
+        actionUpdated = true;
+      }
+    }
+
+    if (!contextUpdated && !actionUpdated) {
+      return [];
+    }
+
+    const toolCall = mapClaudeRunningToolCall({
+      name: "Task",
+      callId: parentToolUseId,
+      input: null,
+      output: null,
+    });
+    if (!toolCall) {
+      return [];
+    }
+
+    const detail: Extract<AgentTimelineItem, { type: "tool_call" }>["detail"] = {
+      type: "sub_agent",
+      ...(state.subAgentType ? { subAgentType: state.subAgentType } : {}),
+      ...(state.description ? { description: state.description } : {}),
+      log: state.actions
+        .map((action) =>
+          action.summary ? `[${action.toolName}] ${action.summary}` : `[${action.toolName}]`
+        )
+        .join("\n"),
+      actions: state.actions.map((action) => ({
+        index: action.index,
+        toolName: action.toolName,
+        ...(action.summary ? { summary: action.summary } : {}),
+      })),
+    };
+
+    return [
+      {
+        type: "timeline",
+        item: {
+          ...toolCall,
+          detail,
+        },
+        provider: "claude",
+      },
+    ];
+  }
+
+  delete(toolUseId: string): void {
+    this.activeSidechains.delete(toolUseId);
+  }
+
+  clear(): void {
+    this.activeSidechains.clear();
+  }
+
+  private updateSubAgentContextFromTaskInput(
+    state: SubAgentActivityState,
+    parentToolUseId: string
+  ): boolean {
+    const taskInput = this.getToolInput(parentToolUseId);
+    const nextSubAgentType = this.normalizeSubAgentText(taskInput?.subagent_type);
+    const nextDescription = this.normalizeSubAgentText(taskInput?.description);
+
+    let changed = false;
+    if (nextSubAgentType && nextSubAgentType !== state.subAgentType) {
+      state.subAgentType = nextSubAgentType;
+      changed = true;
+    }
+    if (nextDescription && nextDescription !== state.description) {
+      state.description = nextDescription;
+      changed = true;
+    }
+    return changed;
+  }
+
+  private normalizeSubAgentText(value: unknown): string | undefined {
+    const normalized = readTrimmedString(value)?.replace(/\s+/g, " ");
+    if (!normalized) {
+      return undefined;
+    }
+    if (normalized.length <= MAX_SUB_AGENT_SUMMARY_CHARS) {
+      return normalized;
+    }
+    return `${normalized.slice(0, MAX_SUB_AGENT_SUMMARY_CHARS)}...`;
+  }
+
+  private extractSubAgentActionCandidates(message: SDKMessage): SubAgentActionCandidate[] {
+    if (message.type === "assistant") {
+      const content = message.message?.content;
+      if (!Array.isArray(content)) {
+        return [];
+      }
+      const actions: SubAgentActionCandidate[] = [];
+      for (const block of content) {
+        if (
+          !isClaudeContentChunk(block) ||
+          !(
+            block.type === "tool_use" ||
+            block.type === "mcp_tool_use" ||
+            block.type === "server_tool_use"
+          ) ||
+          typeof block.name !== "string"
+        ) {
+          continue;
+        }
+        const key = readTrimmedString(block.id) ?? `assistant:${block.name}:${actions.length}`;
+        actions.push({
+          key,
+          toolName: block.name,
+          input: block.input ?? null,
+        });
+      }
+      return actions;
+    }
+
+    if (message.type === "stream_event") {
+      const event = message.event;
+      if (event.type !== "content_block_start") {
+        return [];
+      }
+      const block = isClaudeContentChunk(event.content_block) ? event.content_block : null;
+      if (
+        !block ||
+        !(
+          block.type === "tool_use" ||
+          block.type === "mcp_tool_use" ||
+          block.type === "server_tool_use"
+        ) ||
+        typeof block.name !== "string"
+      ) {
+        return [];
+      }
+      const key =
+        readTrimmedString(block.id) ??
+        `stream:${block.name}:${typeof event.index === "number" ? event.index : 0}`;
+      return [
+        {
+          key,
+          toolName: block.name,
+          input: block.input ?? null,
+        },
+      ];
+    }
+
+    if (message.type === "tool_progress") {
+      const toolName = readTrimmedString(message.tool_name);
+      if (!toolName) {
+        return [];
+      }
+      const key = readTrimmedString(message.tool_use_id) ?? `progress:${toolName}`;
+      return [{ key, toolName, input: null }];
+    }
+
+    return [];
+  }
+
+  private appendSubAgentAction(
+    state: SubAgentActivityState,
+    candidate: SubAgentActionCandidate
+  ): boolean {
+    const normalizedToolName = readTrimmedString(candidate.toolName);
+    if (!normalizedToolName) {
+      return false;
+    }
+
+    const summary = this.deriveSubAgentActionSummary(
+      normalizedToolName,
+      candidate.input
+    );
+    const existingIndex = state.actionIndexByKey.get(candidate.key);
+
+    if (existingIndex !== undefined) {
+      const existing = state.actions[existingIndex];
+      if (!existing) {
+        return false;
+      }
+      const nextSummary = existing.summary ?? summary;
+      if (
+        existing.toolName === normalizedToolName &&
+        existing.summary === nextSummary
+      ) {
+        return false;
+      }
+      state.actions[existingIndex] = {
+        ...existing,
+        toolName: normalizedToolName,
+        ...(nextSummary ? { summary: nextSummary } : {}),
+      };
+      return true;
+    }
+
+    state.actions.push({
+      index: state.nextActionIndex,
+      toolName: normalizedToolName,
+      ...(summary ? { summary } : {}),
+    });
+    state.nextActionIndex += 1;
+    state.actionKeys.push(candidate.key);
+    this.trimSubAgentTail(state);
+    this.rebuildSubAgentActionIndex(state);
+    return true;
+  }
+
+  private trimSubAgentTail(state: SubAgentActivityState): void {
+    while (state.actions.length > MAX_SUB_AGENT_LOG_ENTRIES) {
+      state.actions.shift();
+      state.actionKeys.shift();
+    }
+  }
+
+  private rebuildSubAgentActionIndex(state: SubAgentActivityState): void {
+    state.actionIndexByKey.clear();
+    for (let index = 0; index < state.actionKeys.length; index += 1) {
+      const key = state.actionKeys[index];
+      if (key) {
+        state.actionIndexByKey.set(key, index);
+      }
+    }
+  }
+
+  private deriveSubAgentActionSummary(
+    toolName: string,
+    input: unknown
+  ): string | undefined {
+    const runningToolCall = mapClaudeRunningToolCall({
+      name: toolName,
+      callId: `sub-agent-summary-${toolName}`,
+      input,
+      output: null,
+    });
+    if (!runningToolCall) {
+      return undefined;
+    }
+    const display = buildToolCallDisplayModel({
+      name: runningToolCall.name,
+      status: runningToolCall.status,
+      error: runningToolCall.error,
+      detail: runningToolCall.detail,
+      metadata: runningToolCall.metadata,
+    });
+    return this.normalizeSubAgentText(display.summary);
+  }
+}