socaichat/langgraph/src/agent/coder/nodes/tool-executor.ts

/**
 * Coder tool executor: executes code in sandbox, pushes Gen-UI cards.
 */
import { typedUi } from "@langchain/langgraph-sdk/react-ui/server";
import type ComponentMap from "../../../agent-uis/index.js";
import { LangGraphRunnableConfig } from "@langchain/langgraph";
import { AIMessage } from "@langchain/core/messages";
import { CoderState, CoderUpdate } from "../types.js";
import { sandboxRun } from "../../enterprise/tools/soc-client.js";
import { codeExecuteSchema, codeInstallSchema } from "./agent.js";
import { executeWithRetry, formatToolError } from "@/agent/utils/retry";

export async function toolExecutorNode(
  state: CoderState,
  config: LangGraphRunnableConfig,
): Promise<CoderUpdate> {
  const ui = typedUi<typeof ComponentMap>(config);

  const lastAiMessage = [...state.messages]
    .reverse()
    .find(
      (m): m is AIMessage =>
        (m as AIMessage).tool_calls !== undefined &&
        ((m as AIMessage).tool_calls?.length ?? 0) > 0,
    ) as AIMessage | undefined;

  if (!lastAiMessage?.tool_calls?.length) {
    return { ui: ui.items, timestamp: Date.now() };
  }

  const toolCalls = lastAiMessage.tool_calls!;
  const toolMessages: Array<{
    role: "tool";
    tool_call_id: string;
    content: string;
  }> = [];

  const executions = toolCalls.map(async (tc) => {
    const name = tc.name;
    const args = tc.args;
    const id = tc.id ?? "";

    try {
      switch (name) {
        case "code_execute": {
          const parsed = codeExecuteSchema.parse(args);
          const result = await executeWithRetry(() =>
            sandboxRun(parsed.code, parsed.language ?? "python"),
          );
          ui.push(
            {
              name: "sandbox-result",
              props: {
                language: parsed.language ?? "python",
                exit_code: result.exit_code,
                stdout: result.stdout,
                has_more: result.stdout.length >= 2000,
                duration_ms: result.duration_ms,
              },
            },
            { message: lastAiMessage },
          );
          return {
            role: "tool" as const,
            tool_call_id: id,
            content: JSON.stringify(result),
          };
        }

        case "code_install": {
          const parsed = codeInstallSchema.parse(args);
          // Install packages by running pip/npm in sandbox
          const installCmd = `pip install ${parsed.packages.join(" ")}`;
          const result = await sandboxRun(installCmd, "bash");
          return {
            role: "tool" as const,
            tool_call_id: id,
            content: JSON.stringify({
              success: result.exit_code === 0,
              packages: parsed.packages,
              output: result.stdout.slice(0, 500),
            }),
          };
        }

        default:
          return {
            role: "tool" as const,
            tool_call_id: id,
            content: `Unknown tool: ${name}`,
          };
      }
    } catch (e) {
      return {
        role: "tool" as const,
        tool_call_id: id,
        content: formatToolError(name, e),
      };
    }
  });

  const results = await Promise.all(executions);
  toolMessages.push(...results);

  return {
    messages: toolMessages,
    ui: ui.items,
    timestamp: Date.now(),
  };
}