Claude Agent SDK Migration Implementation Plan

// packages/protocol/src/toolSchemas.ts
//
// Shared Zod tool schemas for frontend UI tools.
// Exported as ZodRawShape (plain objects of Zod types) for:
//   - Agent SDK MCP server registration (electron/src/agent.ts)
//   - Frontend tool validation (web/src/lib/tools/builtin/*.ts wraps with z.object())
//
// If you update a schema here, update the corresponding builtin tool file too.

import { z } from "zod";

// --- Shared sub-schemas ---

export const xyPositionSchema = z.object({ x: z.number(), y: z.number() });
export const positionInputSchema = z.union([xyPositionSchema, z.string()]);
export const nodePropertySchema = z.record(z.string(), z.any());
export const optionalWorkflowIdSchema = z
  .string()
  .nullable()
  .optional()
  .describe("Optional workflow id; when omitted/null, the current workflow is used.");

// --- Tool parameter shapes (ZodRawShape for SDK tool()) ---

export const uiSearchNodesParams = {
  query: z.string(),
  input_type: z.string().optional(),
  output_type: z.string().optional(),
  strict_match: z.boolean().optional(),
  include_properties: z.boolean().optional(),
  include_outputs: z.boolean().optional(),
  limit: z.number().min(1).max(100).optional(),
};

export const uiAddNodeParams = {
  id: z.string().optional(),
  type: z.string().optional(),
  node_type: z.string().optional(),
  position: positionInputSchema.optional(),
  properties: nodePropertySchema.optional(),
  workflow_id: optionalWorkflowIdSchema,
};

export const uiConnectNodesParams = {
  source_id: z.string(),
  source_handle: z.string(),
  target_id: z.string(),
  target_handle: z.string(),
  workflow_id: optionalWorkflowIdSchema,
};

export const uiGetGraphParams = {
  workflow_id: optionalWorkflowIdSchema,
};

export const uiUpdateNodeDataParams = {
  node_id: z.string(),
  data: z.record(z.string(), z.any()),
  workflow_id: optionalWorkflowIdSchema,
};

export const uiDeleteNodeParams = {
  node_id: z.string(),
  workflow_id: optionalWorkflowIdSchema,
};

export const uiDeleteEdgeParams = {
  edge_id: z.string(),
  workflow_id: optionalWorkflowIdSchema,
};

export const uiMoveNodeParams = {
  node_id: z.string(),
  position: xyPositionSchema,
  workflow_id: optionalWorkflowIdSchema,
};

export const uiSetNodeTitleParams = {
  node_id: z.string(),
  title: z.string(),
  workflow_id: optionalWorkflowIdSchema,
};

export const uiSetNodeSyncModeParams = {
  node_id: z.string(),
  mode: z.enum(["on_any", "zip_all"]),
  workflow_id: optionalWorkflowIdSchema,
};

export const uiOpenWorkflowParams = {
  workflow_id: z.string().optional().describe("Workflow id to target."),
  id: z.string().optional().describe("Alias for workflow_id."),
};

export const uiRunWorkflowParams = {
  workflow_id: z.string().optional().describe("Workflow id to target."),
  id: z.string().optional().describe("Alias for workflow_id."),
  params: z.record(z.string(), z.unknown()).optional().describe("Optional workflow run parameters."),
};

export const uiSwitchTabParams = {
  tab_index: z.number().int().min(0).describe("Zero-based tab index (0 is the first tab)."),
};

export const uiCopyParams = {
  text: z.string().describe("The text to copy to clipboard."),
};

export const uiPasteParams = {};

// --- Registry for MCP server registration ---
// Only non-hidden tools. ui_graph (hidden) is excluded.

export interface UiToolSchema {
  description: string;
  parameters: Record<string, z.ZodTypeAny>;
}

export const uiToolSchemas: Record<string, UiToolSchema> = {
  ui_search_nodes: {
    description: "Search available node types from metadata store by query/type filters.",
    parameters: uiSearchNodesParams,
  },
  ui_add_node: {
    description: "Add a node to the current workflow graph.",
    parameters: uiAddNodeParams,
  },
  ui_connect_nodes: {
    description: "Connect two nodes via their handles.",
    parameters: uiConnectNodesParams,
  },
  ui_get_graph: {
    description: "Get the current workflow graph (nodes and edges).",
    parameters: uiGetGraphParams,
  },
  ui_update_node_data: {
    description: "Update an existing node's data/properties.",
    parameters: uiUpdateNodeDataParams,
  },
  ui_delete_node: {
    description: "Delete a node from the workflow graph.",
    parameters: uiDeleteNodeParams,
  },
  ui_delete_edge: {
    description: "Delete an edge from the workflow graph.",
    parameters: uiDeleteEdgeParams,
  },
  ui_move_node: {
    description: "Move a node to a new position.",
    parameters: uiMoveNodeParams,
  },
  ui_set_node_title: {
    description: "Set a node's display title.",
    parameters: uiSetNodeTitleParams,
  },
  ui_set_node_sync_mode: {
    description: "Set a node's sync mode (on_any or zip_all).",
    parameters: uiSetNodeSyncModeParams,
  },
  ui_open_workflow: {
    description: "Open a workflow by ID in a new tab.",
    parameters: uiOpenWorkflowParams,
  },
  ui_run_workflow: {
    description: "Run a workflow, optionally passing parameters.",
    parameters: uiRunWorkflowParams,
  },
  ui_switch_tab: {
    description: "Switch to a workflow tab by index.",
    parameters: uiSwitchTabParams,
  },
  ui_copy: {
    description: "Copy text to the clipboard.",
    parameters: uiCopyParams,
  },
  ui_paste: {
    description: "Paste text from the clipboard.",
    parameters: uiPasteParams,
  },
};

export * from "./toolSchemas.js";

// Before:
import { z } from "zod";
import { optionalWorkflowIdSchema } from "./workflow";

FrontendToolRegistry.register({
  name: "ui_connect_nodes",
  parameters: z.object({
    source_id: z.string(),
    source_handle: z.string(),
    target_id: z.string(),
    target_handle: z.string(),
    workflow_id: optionalWorkflowIdSchema
  }),
  ...
});

// After:
import { z } from "zod";
import { uiConnectNodesParams } from "@nodetool/protocol";

FrontendToolRegistry.register({
  name: "ui_connect_nodes",
  parameters: z.object(uiConnectNodesParams),
  ...
});

// addNode.ts — wrap shared params and add refinement
import { uiAddNodeParams, positionInputSchema, nodePropertySchema } from "@nodetool/protocol";

const nodeInputSchema = z.object({
  id: z.string(),
  type: z.string().optional(),
  node_type: z.string().optional(),
  position: positionInputSchema,
  properties: nodePropertySchema.optional(),
});

const addNodeParametersSchema = z
  .object({
    node: nodeInputSchema.optional(),
    ...uiAddNodeParams,
  })
  .refine(
    (data) => data.node || (data.id !== undefined && data.position !== undefined),
    { message: "Either 'node' object or both 'id' and 'position' must be provided" }
  );

// uiActions.ts — ui_open_workflow and ui_run_workflow re-apply refinements
import { uiOpenWorkflowParams, uiRunWorkflowParams } from "@nodetool/protocol";

const openWorkflowSchema = z.object(uiOpenWorkflowParams).refine(
  (value) => Boolean(value.workflow_id ?? value.id),
  { message: "workflow_id (or id) is required" }
);

const runWorkflowSchema = z.object(uiRunWorkflowParams).refine(
  (value) => Boolean(value.workflow_id ?? value.id),
  { message: "workflow_id (or id) is required" }
);

// searchNodes.ts — wrap shared params, add preprocess for string→boolean coercion
import { uiSearchNodesParams } from "@nodetool/protocol";

const booleanLikeOptional = z.preprocess((value) => {
  if (typeof value === "boolean") return value;
  if (typeof value === "string") {
    const n = value.trim().toLowerCase();
    if (n === "true") return true;
    if (n === "false") return false;
  }
  return value;
}, z.boolean().optional());

// Override boolean fields with preprocess version for frontend resilience
const searchNodesSchema = z.object({
  ...uiSearchNodesParams,
  strict_match: booleanLikeOptional,
  include_properties: booleanLikeOptional,
  include_outputs: booleanLikeOptional,
});

// Before (lines 1-10):
import { z } from "zod";
import type { FrontendToolState } from "../frontendTools";

export const optionalWorkflowIdSchema = z
  .string()
  .nullable()
  .optional()
  .describe("Optional workflow id; when omitted/null, the current workflow is used.");

export const optionalWorkflowIdSchemaCompact = z.string().nullable().optional();

// After:
import { z } from "zod";
import type { FrontendToolState } from "../frontendTools";
export { optionalWorkflowIdSchema } from "@nodetool/protocol";

export const optionalWorkflowIdSchemaCompact = z.string().nullable().optional();
// rest of file (resolveWorkflowId) unchanged

import { query, tool, createSdkMcpServer } from "@anthropic-ai/claude-agent-sdk";
import type { Query } from "@anthropic-ai/claude-agent-sdk";
import { uiToolSchemas } from "@nodetool/protocol";

function buildMcpServer(
  webContents: WebContents,
  sessionId: string,
): ReturnType<typeof createSdkMcpServer> {
  const tools = Object.entries(uiToolSchemas).map(([name, schema]) =>
    tool(name, schema.description, schema.parameters, async (args) => {
      const result = await executeFrontendTool(
        webContents,
        sessionId,
        randomUUID(),
        name,
        args,
      );
      const text =
        typeof result === "string" ? result : JSON.stringify(result ?? null);
      return { content: [{ type: "text" as const, text }] };
    }),
  );

  return createSdkMcpServer({ name: "nodetool-ui", version: "1.0.0", tools });
}

function convertSdkMessage(
  msg: { type: string; [key: string]: unknown },
  sessionId: string,
): AgentMessage | null {
  const uuid = randomUUID();

  if (msg.type === "assistant") {
    const message = msg.message as Record<string, unknown> | undefined;
    const content = message?.content;
    if (!Array.isArray(content)) return null;

    const textBlocks: Array<{ type: string; text: string }> = [];
    const toolCalls: AgentMessage["tool_calls"] = [];

    for (const block of content) {
      if (!block || typeof block !== "object") continue;
      const b = block as Record<string, unknown>;
      if (b.type === "text" && typeof b.text === "string") {
        textBlocks.push({ type: "text", text: b.text });
      } else if (
        b.type === "tool_use" &&
        typeof b.id === "string" &&
        typeof b.name === "string"
      ) {
        toolCalls.push({
          id: b.id,
          type: "function",
          function: {
            name: b.name,
            arguments: JSON.stringify(b.input ?? {}),
          },
        });
      }
    }

    if (textBlocks.length === 0 && toolCalls.length === 0) return null;

    return {
      type: "assistant",
      uuid,
      session_id: sessionId,
      content: textBlocks.length > 0 ? textBlocks : undefined,
      tool_calls: toolCalls.length > 0 ? toolCalls : undefined,
    };
  }

  if (msg.type === "result") {
    const subtype = msg.subtype as string | undefined;
    if (subtype === "success") {
      return {
        type: "result",
        uuid,
        session_id: sessionId,
        subtype: "success",
        text: typeof msg.result === "string" ? msg.result : undefined,
      };
    }
    return {
      type: "result",
      uuid,
      session_id: sessionId,
      subtype: subtype ?? "error",
      is_error: true,
      errors: Array.isArray(msg.errors)
        ? (msg.errors as string[])
        : [String(msg.result ?? "Unknown error")],
    };
  }

  // system init — handled by caller for session ID capture
  if (msg.type === "system") return null;

  // stream_event — partial assistant message for real-time streaming
  if (msg.type === "stream_event") {
    const event = msg as Record<string, unknown>;
    // Extract partial text if available
    const partial = event.message as Record<string, unknown> | undefined;
    const content = partial?.content;
    if (Array.isArray(content)) {
      const textBlocks: Array<{ type: string; text: string }> = [];
      for (const block of content) {
        if (block && typeof block === "object") {
          const b = block as Record<string, unknown>;
          if (b.type === "text" && typeof b.text === "string") {
            textBlocks.push({ type: "text", text: b.text });
          }
        }
      }
      if (textBlocks.length > 0) {
        return {
          type: "assistant",
          uuid,
          session_id: sessionId,
          content: textBlocks,
        };
      }
    }
    return null;
  }

  return null;
}

class ClaudeAgentSession implements AgentQuerySession {
  private closed = false;
  private readonly model: string;
  private readonly workspacePath: string;
  private readonly systemPrompt: string;
  private resolvedSessionId: string | null;
  private inFlight = false;
  private activeQuery: Query | null = null;

  constructor(options: {
    model: string;
    workspacePath: string;
    resumeSessionId?: string;
    systemPrompt?: string;
  }) {
    this.model = options.model;
    this.workspacePath = options.workspacePath;
    this.systemPrompt = options.systemPrompt ?? FAST_WORKFLOW_SYSTEM_PROMPT;
    this.resolvedSessionId = options.resumeSessionId ?? null;
  }

  async send(
    message: string,
    webContents: WebContents | null,
    sessionId: string,
    _manifest: FrontendToolManifest[],
    onMessage?: (message: AgentMessage) => void,
  ): Promise<AgentMessage[]> {
    if (this.closed) {
      throw new Error("Cannot send to a closed session");
    }
    if (this.inFlight) {
      throw new Error("A request is already in progress for this session");
    }

    this.inFlight = true;

    try {
      const mcpServer = webContents
        ? buildMcpServer(webContents, sessionId)
        : undefined;

      const allowedTools = mcpServer
        ? Object.keys(uiToolSchemas).map((n) => `mcp__nodetool-ui__${n}`)
        : [];

      const queryHandle = query({
        prompt: message,
        options: {
          model: this.model,
          cwd: this.workspacePath,
          systemPrompt: {
            type: "preset",
            preset: "claude_code",
            append: this.systemPrompt,
          },
          permissionMode: "bypassPermissions",
          allowDangerouslySkipPermissions: true,
          disallowedTools: CLAUDE_DISALLOWED_TOOLS,
          allowedTools,
          maxTurns: 50,
          ...(mcpServer && { mcpServers: { "nodetool-ui": mcpServer } }),
          ...(this.resolvedSessionId && { resume: this.resolvedSessionId }),
        },
      });
      this.activeQuery = queryHandle;

      const outputMessages: AgentMessage[] = [];
      for await (const msg of queryHandle) {
        if (
          msg.type === "system" &&
          (msg as Record<string, unknown>).subtype === "init"
        ) {
          this.resolvedSessionId =
            (msg as Record<string, unknown>).session_id as string;
        }

        const agentMsg = convertSdkMessage(
          msg as Record<string, unknown>,
          this.resolvedSessionId ?? sessionId,
        );
        if (!agentMsg) continue;

        outputMessages.push(agentMsg);
        if (onMessage) {
          onMessage(agentMsg);
        }
      }

      return outputMessages;
    } catch (error) {
      const errorMsg: AgentMessage = {
        type: "result",
        uuid: randomUUID(),
        session_id: this.resolvedSessionId ?? sessionId,
        subtype: "error",
        is_error: true,
        errors: [error instanceof Error ? error.message : String(error)],
      };
      if (onMessage) {
        onMessage(errorMsg);
      }
      return [errorMsg];
    } finally {
      this.activeQuery = null;
      this.inFlight = false;
    }
  }

  async interrupt(): Promise<void> {
    if (this.activeQuery) {
      this.activeQuery.interrupt();
    }
  }

  close(): void {
    if (this.closed) {
      return;
    }
    this.closed = true;
    if (this.activeQuery) {
      this.activeQuery.close();
      this.activeQuery = null;
    }
  }
}

// In createAgentSession, replace the existing claude session creation:
const useCliAgent = process.env.NODETOOL_AGENT_USE_CLI === "1";

const session: AgentQuerySession =
  provider === "codex"
    ? new CodexQuerySession({
        model: options.model,
        workspacePath: options.workspacePath,
        resumeSessionId: options.resumeSessionId,
        systemPrompt: CODEX_SYSTEM_PROMPT,
      })
    : useCliAgent
      ? new ClaudeCliSession({
          model: options.model,
          workspacePath: options.workspacePath,
          resumeSessionId: options.resumeSessionId,
          systemPrompt:
            process.env.NODETOOL_AGENT_VERBOSE_PROMPT === "1"
              ? HELP_SYSTEM_PROMPT
              : FAST_WORKFLOW_SYSTEM_PROMPT,
        })
      : new ClaudeAgentSession({
          model: options.model,
          workspacePath: options.workspacePath,
          resumeSessionId: options.resumeSessionId,
          systemPrompt:
            process.env.NODETOOL_AGENT_VERBOSE_PROMPT === "1"
              ? HELP_SYSTEM_PROMPT
              : FAST_WORKFLOW_SYSTEM_PROMPT,
        });

const DEFAULT_CLAUDE_MODELS: AgentModelDescriptor[] = [
  {
    id: "claude-sonnet-4-6",
    label: "Claude Sonnet 4.6",
    isDefault: true,
  },
  {
    id: "claude-opus-4-6",
    label: "Claude Opus 4.6",
  },
  {
    id: "claude-haiku-4-5",
    label: "Claude Haiku 4.5",
  },
];

/**
 * Legacy Claude session using CLI pipe approach.
 * Kept behind NODETOOL_AGENT_USE_CLI=1 for rollback.
 * Will be removed after SDK migration is validated.
 */