import { tool, type LMStudioClient, type Tool } from "@lmstudio/sdk";
import { z } from "zod";

/**
 * Tools the sub-agent may use. Deliberately excludes `set_workspace` (workspace state is shared
 * process-wide; the main chat sets it) and `run` (no shell access for sub-agents).
 */
const SUBAGENT_TOOL_NAMES: ReadonlySet<string> = new Set([
  "list_dir",
  "read_file",
  "search",
  "write_file",
  "edit_file",
]);

/** The subset of core tools a sub-agent is allowed to call. */
export function subagentToolList(coreTools: Array<Tool>): Array<Tool> {
  return coreTools.filter((t) => SUBAGENT_TOOL_NAMES.has(t.name));
}

// Serialize sub-agent runs so concurrent delegations can't interleave file edits.
let chain: Promise<void> = Promise.resolve();
function serialized<T>(work: () => Promise<T>): Promise<T> {
  const run = chain.then(work);
  chain = run.then(
    () => undefined,
    () => undefined
  );
  return run;
}

const SUMMARY_CAP = 2000;

export function createSubagentTool(coreTools: Array<Tool>, client: LMStudioClient) {
  const subTools = subagentToolList(coreTools);
  return tool({
    name: "subagent",
    description:
      "Delegate a self-contained coding task to a background agent that can only list, read, search, create and edit files in the working folder (no shell, no deletes). Returns a short summary. Use for smaller tasks to keep this chat's context light.",
    parameters: { task: z.string().min(1) },
    implementation: async ({ task }) =>
      serialized(async () => {
        try {
          const model = await client.llm.model();
          let finalText = "";
          const result = await model.act(
            [
              "You are a focused sub-agent. Work only inside the current project working folder",
              "using your file tools (list_dir, read_file, search, write_file, edit_file).",
              "You cannot run commands or delete files; if a step requires either, skip it and say so.",
              "When done, reply with a concise final summary: what you did and the paths of all",
              "created/changed files. Keep the final message under 15 lines.",
              "",
              `Task: ${task}`,
            ].join("\n"),
            subTools,
            {
              onMessage: (m) => {
                if (m.isAssistantMessage()) finalText = m.getText().trim();
              },
            }
          );
          const secs = result.totalExecutionTimeSeconds.toFixed(1);
          if (!finalText) {
            return `Sub-agent finished (${result.rounds} rounds, ${secs}s) but produced no summary.`;
          }
          const text =
            finalText.length > SUMMARY_CAP ? "…[truncated]…" + finalText.slice(-SUMMARY_CAP) : finalText;
          return `[subagent: ${result.rounds} rounds, ${secs}s]\n${text}`;
        } catch (e) {
          const msg = e instanceof Error ? e.message : String(e);
          return `subagent error: ${msg}`;
        }
      }),
  });
}

import { tool, type LMStudioClient, type Tool } from "@lmstudio/sdk";
import { z } from "zod";

/**
 * Tools the sub-agent may use. Deliberately excludes `set_workspace` (workspace state is shared
 * process-wide; the main chat sets it) and `run` (no shell access for sub-agents).
 */
const SUBAGENT_TOOL_NAMES: ReadonlySet<string> = new Set([
  "list_dir",
  "read_file",
  "search",
  "write_file",
  "edit_file",
]);

/** The subset of core tools a sub-agent is allowed to call. */
export function subagentToolList(coreTools: Array<Tool>): Array<Tool> {
  return coreTools.filter((t) => SUBAGENT_TOOL_NAMES.has(t.name));
}

// Serialize sub-agent runs so concurrent delegations can't interleave file edits.
let chain: Promise<void> = Promise.resolve();
function serialized<T>(work: () => Promise<T>): Promise<T> {
  const run = chain.then(work);
  chain = run.then(
    () => undefined,
    () => undefined
  );
  return run;
}

const SUMMARY_CAP = 2000;

export function createSubagentTool(coreTools: Array<Tool>, client: LMStudioClient) {
  const subTools = subagentToolList(coreTools);
  return tool({
    name: "subagent",
    description:
      "Delegate a self-contained coding task to a background agent that can only list, read, search, create and edit files in the working folder (no shell, no deletes). Returns a short summary. Use for smaller tasks to keep this chat's context light.",
    parameters: { task: z.string().min(1) },
    implementation: async ({ task }) =>
      serialized(async () => {
        try {
          const model = await client.llm.model();
          let finalText = "";
          const result = await model.act(
            [
              "You are a focused sub-agent. Work only inside the current project working folder",
              "using your file tools (list_dir, read_file, search, write_file, edit_file).",
              "You cannot run commands or delete files; if a step requires either, skip it and say so.",
              "When done, reply with a concise final summary: what you did and the paths of all",
              "created/changed files. Keep the final message under 15 lines.",
              "",
              `Task: ${task}`,
            ].join("\n"),
            subTools,
            {
              onMessage: (m) => {
                if (m.isAssistantMessage()) finalText = m.getText().trim();
              },
            }
          );
          const secs = result.totalExecutionTimeSeconds.toFixed(1);
          if (!finalText) {
            return `Sub-agent finished (${result.rounds} rounds, ${secs}s) but produced no summary.`;
          }
          const text =
            finalText.length > SUMMARY_CAP ? "…[truncated]…" + finalText.slice(-SUMMARY_CAP) : finalText;
          return `[subagent: ${result.rounds} rounds, ${secs}s]\n${text}`;
        } catch (e) {
          const msg = e instanceof Error ? e.message : String(e);
          return `subagent error: ${msg}`;
        }
      }),
  });
}