gadget/gadget-drone/src/services/agent.ts

// src/services/agent.ts
// Copyright (C) 2026 Rob Colbert <rob.colbert@openplatform.us>
// Licensed under the Apache License, Version 2.0

import path from "node:path";
import fs from "node:fs/promises";

import env from "../config/env.ts";

import { Socket } from "socket.io-client";
import {
  type IAiChatOptions,
  type IAiChatResponse,
  type IAiResponseStreamFn,
  type IAiTool,
  type IContextChatMessage,
} from "@gadget/ai";
import {
  type IAiProvider,
  type IAiModelSettings,
  type IDroneModelConfig,
  type IChatSession,
  type IChatTurn,
  type IChatSubagentProcess,
  type IChatToolCall,
  type IUser,
  type GadgetId,
  type ServerToClientEvents,
  type ClientToServerEvents,
  ChatSessionMode,
  type IProject,
} from "@gadget/api";

import AiService from "./ai.ts";
import WorkspaceService from "./workspace.ts";

import { GadgetService } from "../lib/service.ts";
import {
  AiToolbox,
  FileEditTool,
  FileReadTool,
  FileWriteTool,
  FetchUrlTool,
  GlobTool,
  GoogleSearchTool,
  GrepTool,
  ListTool,
  PlanFileEditTool,
  PlanFileReadTool,
  PlanFileWriteTool,
  PlanListTool,
  ShellExecTool,
  SubagentTool,
  type DroneToolboxEnvironment,
} from "../tools/index.ts";

export interface IAgentWorkOrder {
  createdAt: Date;
  turn: IChatTurn;
  context: IChatTurn[];
}

type DroneSocket = Socket<ServerToClientEvents, ClientToServerEvents>;

const toolboxEnv: DroneToolboxEnvironment = {
  NODE_ENV: env.NODE_ENV || "develop",
  services: {
    google: {
      cse: {
        apiKey: env.google.cse.apiKey,
        engineId: env.google.cse.engineId,
      },
    },
  },
};

class AgentService extends GadgetService {
  private toolbox = new AiToolbox(toolboxEnv);
  private currentWorkOrder: IAgentWorkOrder | null = null;
  private currentSocket: DroneSocket | null = null;
  private currentToolCallId: string | null = null;
  private abortController: AbortController | null = null;

  get name(): string {
    return "AgentService";
  }
  get slug(): string {
    return "svc:agent";
  }

  async start(): Promise<void> {
    const readOnlyModes = [
      ChatSessionMode.Plan,
      ChatSessionMode.Build,
      ChatSessionMode.Test,
      ChatSessionMode.Ship,
      ChatSessionMode.Develop,
    ];
    const writeModes = [
      ChatSessionMode.Build,
      ChatSessionMode.Test,
      ChatSessionMode.Ship,
      ChatSessionMode.Develop,
    ];

    // Network tools — available in all modes for research
    this.toolbox.register(new GoogleSearchTool(this.toolbox), readOnlyModes);
    this.toolbox.register(new FetchUrlTool(this.toolbox), readOnlyModes);

    // System tools — read-only: available in all modes for exploration
    this.toolbox.register(new FileReadTool(this.toolbox), readOnlyModes);
    this.toolbox.register(new ListTool(this.toolbox), readOnlyModes);
    this.toolbox.register(new GrepTool(this.toolbox), readOnlyModes);
    this.toolbox.register(new GlobTool(this.toolbox), readOnlyModes);

    // System tools — write: restricted to build/test/ship/develop
    this.toolbox.register(new FileWriteTool(this.toolbox), writeModes);
    this.toolbox.register(new FileEditTool(this.toolbox), writeModes);
    this.toolbox.register(new ShellExecTool(this.toolbox), writeModes);

    // Plan tools — Gadget's own .gadget directory: only available in Plan mode
    this.toolbox.register(new PlanFileReadTool(this.toolbox), [ChatSessionMode.Plan]);
    this.toolbox.register(new PlanFileWriteTool(this.toolbox), [ChatSessionMode.Plan]);
    this.toolbox.register(new PlanFileEditTool(this.toolbox), [ChatSessionMode.Plan]);
    this.toolbox.register(new PlanListTool(this.toolbox), [ChatSessionMode.Plan]);

    // Chat tools — subagent spawning: available in all modes
    const subagentTool = new SubagentTool(this.toolbox);
    subagentTool.setSpawner((agentType, prompt) => this.spawnSubagent(agentType, prompt));
    this.toolbox.register(subagentTool, readOnlyModes);

    this.log.info("started");
  }

  async stop(): Promise<void> {
    this.currentWorkOrder = null;
    this.currentSocket = null;
    this.log.info("stopped");
  }

  async process(
    workOrder: IAgentWorkOrder,
    socket: DroneSocket,
  ): Promise<void> {
    this.currentWorkOrder = workOrder;
    this.currentSocket = socket;
    this.abortController = new AbortController();

    const { turn } = workOrder;
    let toolCallCount = 0;
    let inputTokens = 0;
    let outputTokens = 0;

    // Build the full message array that grows with each iteration
    const messages: IContextChatMessage[] = [];

    if (turn.prompts.system) {
      messages.push({
        createdAt: turn.createdAt,
        role: "system",
        content: turn.prompts.system,
      });
    }

    messages.push(...this.buildSessionContext(workOrder));

    // Current turn's user prompt must be the last message before the AI call
    messages.push({
      createdAt: turn.createdAt,
      role: "user",
      content: turn.prompts.user,
    });

    const reasoningEffort = turn.reasoningEffort || "off";
    const reasoning: boolean | "low" | "medium" | "high" =
      reasoningEffort === "off" ? false : reasoningEffort;

    try {
      this.updateToolboxWorkspace(turn);
    } catch (cause) {
      socket.emit(
        "workOrderComplete",
        turn._id,
        false,
        `failed to update workspace: ${(cause as Error).message}`,
      );
      throw new Error("failed to update workspace", { cause });
    }

    this.log.info("agent loop starting", {
      turnId: turn._id,
      messageCount: messages.length,
      toolCount: this.toolbox.getToolNamesForMode(turn.mode).length,
    });

    try {
      let continueLoop = true;
      while (continueLoop) {
        continueLoop = false;

        this.log.info("agent loop iteration", {
          messagesCount: messages.length,
          toolsAvailable: this.toolbox.getToolNamesForMode(turn.mode),
        });

        const chatOptions: IAiChatOptions = {
          context: messages,
          tools: this.getToolsForMode(turn.mode),
          signal: this.abortController.signal,
        };

        let response: IAiChatResponse;
        let currentReasoning: boolean | "low" | "medium" | "high" = reasoning;

        for (let attempt = 0; ; attempt++) {
          try {
            response = await AiService.chat(
              turn.provider,
              this.buildDroneModelConfig(turn.provider, turn.llm, currentReasoning),
              chatOptions,
              this.makeStreamHandler(socket),
            );
            break;
          } catch (error) {
            if (
              attempt === 0 &&
              currentReasoning !== false &&
              error instanceof Error &&
              error.message.includes("empty chat response")
            ) {
              this.log.warn("model returned empty response with reasoning; retrying without reasoning", {
                originalError: error.message,
              });
              currentReasoning = false;
              continue;
            }
            throw error;
          }
        }

        if (response.doneReason === "load" && !response.response && !response.thinking) {
          throw new Error("Model failed to respond (still loading or error)");
        }

        // Process tool calls if present
        if (response.toolCalls && response.toolCalls.length > 0) {
          continueLoop = true;
          toolCallCount += response.toolCalls.length;

          messages.push({
            createdAt: turn.createdAt,
            role: "assistant",
            content: response.response,
            toolCalls: response.toolCalls.map((tc) => ({
              id: tc.callId,
              type: "function" as const,
              function: {
                name: tc.function.name,
                arguments: tc.function.arguments,
              },
            })),
          });

          for (const toolCall of response.toolCalls) {
            this.currentToolCallId = toolCall.callId;

            // Emit tool call start immediately so frontend shows it
            if (toolCall.function.name === "subagent") {
              socket.emit("toolCall", toolCall.callId, toolCall.function.name, toolCall.function.arguments, "");
            }

            const result = await this.executeTool(
              toolCall.function.name,
              toolCall.function.arguments,
            );

            if (toolCall.function.name === "subagent") {
              let responseText = result;
              let subagentData: Record<string, unknown> | undefined;
              try {
                const parsed = JSON.parse(result);
                if (parsed.success && parsed.data) {
                  subagentData = parsed.data;
                  responseText = parsed.data.response || result;
                }
              } catch {}
              socket.emit("agent:complete", {
                agentId: toolCall.callId,
                response: responseText,
                subagent: subagentData,
              });
            } else {
              socket.emit("toolCall", toolCall.callId, toolCall.function.name, toolCall.function.arguments, result);
            }

            this.currentToolCallId = null;

            messages.push({
              createdAt: turn.createdAt,
              role: "tool",
              content: result,
              toolCallId: toolCall.callId,
              toolName: toolCall.function.name,
            });

            inputTokens += Math.ceil(toolCall.function.arguments.length / 4);
            outputTokens += Math.ceil(result.length / 4);
          }
        }
      }

      socket.emit("workOrderComplete", turn._id, true);
    } catch (cause) {
      if (cause instanceof Error && cause.name === "AbortError") {
        this.log.info("work order aborted by user", { turnId: turn._id });
        socket.emit("workOrderComplete", turn._id, true, "aborted");
        return;
      }
      const msg = cause instanceof Error ? cause.message : String(cause);
      this.log.error("agent loop failed, sending workOrderComplete(false)", {
        turnId: turn._id,
        error: msg,
      });
      socket.emit("workOrderComplete", turn._id, false, msg);
      throw cause;
    } finally {
      this.abortController = null;
    }
  }

  /**
   * Signals the abort controller for the currently running work order.
   * Returns true if an abort was signaled, false if there was no active work order.
   */
  abortCurrentWorkOrder(): boolean {
    if (!this.abortController) return false;
    this.abortController.abort();
    return true;
  }

  buildSessionContext(workOrder: IAgentWorkOrder): IContextChatMessage[] {
    const session = workOrder.turn.session as IChatSession;
    if (!session.user) {
      throw new Error("ChatSession must be populated with user data");
    }
    const user: IUser = session.user as IUser;
    const messages: IContextChatMessage[] = [];

    for (const turn of workOrder.context) {
      /*
       * add the User message
       */
      messages.push({
        createdAt: turn.createdAt,
        role: "user",
        content: turn.prompts.user,
        user: {
          _id: user._id,
          username: user.email,
          displayName: user.displayName,
        },
      });

      /*
       * Add the assistant's output (if any), to include the thinking
       * (reasoning) output (if any).
       */
      let content = "";

      // Extract thinking and response from blocks
      for (const block of turn.blocks) {
        if (block.mode === "thinking" && typeof block.content === "string") {
          content += `<thinking>${block.content}</thinking>`;
        } else if (
          block.mode === "responding" &&
          typeof block.content === "string"
        ) {
          if (content && content.length) {
            content += "\n";
          }
          content += block.content;
        }
      }

      const assistantMsg: IContextChatMessage = {
        createdAt: turn.createdAt,
        role: "assistant",
        content:
          content && content.length
            ? content
            : "(you didn't say anything this turn)",
      };

      if (turn.toolCalls?.length > 0) {
        assistantMsg.toolCalls = turn.toolCalls.map((tc: IChatToolCall) => ({
          id: tc.callId,
          type: "function" as const,
          function: {
            name: tc.name,
            arguments: tc.parameters,
          },
        }));
      }

      messages.push(assistantMsg);

      if (turn.toolCalls?.length > 0) {
        for (const toolCall of turn.toolCalls) {
          messages.push({
            createdAt: turn.createdAt,
            role: "tool",
            content: toolCall.response,
            toolCallId: toolCall.callId,
            toolName: toolCall.name,
          });
        }
      }
    }

    return messages;
  }

  /**
   * To optimize context, reduce clutter, and help the agent focus, full outputs
   * of older file reads and edits are summarized once a newer version is available.
   */
  pruneSessionContext(messages: IContextChatMessage[]): void {
    // TODO
  }

  private getToolsForMode(mode: ChatSessionMode): any[] {
    return Array.from(this.toolbox.getModeSet(mode) || []);
  }

  getToolNamesForMode(mode: ChatSessionMode): string[] {
    return this.toolbox.getToolNamesForMode(mode);
  }

  private makeStreamHandler(socket: DroneSocket): IAiResponseStreamFn {
    return async (chunk) => {
      switch (chunk.type) {
        case "thinking":
          socket.emit("thinking", chunk.data);
          break;
        case "response":
          socket.emit("response", chunk.data);
          break;
      }
    };
  }

  private makeAgentStreamHandler(socket: DroneSocket, agentId: string): IAiResponseStreamFn {
    return async (chunk) => {
      switch (chunk.type) {
        case "thinking":
          socket.emit("agent:thinking", { agentId, thinking: chunk.data });
          break;
        case "response":
          socket.emit("agent:response", { agentId, chunk: chunk.data });
          break;
      }
    };
  }

  /**
   * Builds an IDroneModelConfig by looking up the model's stored settings
   * from the provider's cached models array. Falls back to safe defaults
   * when stored settings aren't available (e.g., model not yet probed).
   */
  private buildDroneModelConfig(
    provider: IAiProvider | GadgetId,
    modelId: string,
    reasoning: boolean | "low" | "medium" | "high",
  ): Omit<IDroneModelConfig, "provider"> {
    const settings = this.getModelSettings(provider, modelId);
    return {
      modelId,
      params: {
        reasoning,
        temperature: settings?.temperature ?? 0.1,
        topP: settings?.topP ?? 0.9,
        topK: settings?.topK ?? 40,
        numPredict: settings?.numPredict ?? -1,         // -1 = unlimited (Ollama)
        numCtx: settings?.numCtx ?? 131072,             // 128k fallback
        maxCompletionTokens: settings?.maxCompletionTokens ?? 16384, // 16k fallback
      },
    };
  }

  /**
   * Looks up the stored IAiModelSettings for a model from the provider's
   * cached models array. Returns undefined if the provider or model isn't found.
   */
  private getModelSettings(
    provider: IAiProvider | GadgetId,
    modelId: string,
  ): IAiModelSettings | undefined {
    if (typeof provider === "string") return undefined;
    const modelRecord = (provider as IAiProvider).models?.find(
      (m) => m.id === modelId,
    );
    return modelRecord?.settings;
  }

  private async executeTool(name: string, argsJson: string): Promise<string> {
    const tool = this.toolbox.getTool(name);
    if (!tool) {
      const msg = `Unknown tool: ${name}`;
      this.log.error("tool not found", { toolName: name });
      return JSON.stringify({ success: false, error: msg });
    }

    try {
      const args = JSON.parse(argsJson);
      this.log.info("executing tool", { name, params: argsJson });
      const result = await tool.execute(args, this.log);
      this.log.info("tool result", {
        name,
        resultLength: result.length,
        preview: result.length > 100 ? `${result.slice(0, 100)}...` : result,
      });
      return result;
    } catch (error) {
      const msg = error instanceof Error ? error.message : String(error);
      this.log.error("tool execution failed", { toolName: name, args: argsJson, error: msg });
      return JSON.stringify({ success: false, error: msg });
    }
  }

  private updateToolboxWorkspace(turn: IChatTurn): void {
    const project = turn.project as IProject;
    if (!project || typeof project === "string") {
      throw new Error("ChatTurn must be populated with project data");
    }

    const workspaceDir = WorkspaceService.workspaceDir;
    const cacheDir = WorkspaceService.workspaceCacheDir;
    if (!workspaceDir || !cacheDir) {
      throw new Error("Workspace must be initialized before agent tools are used");
    }

    this.toolbox.updateWorkspace({
      workspaceDir,
      projectDir: WorkspaceService.getProjectDirectory(project.slug),
      cacheDir,
    });
  }

  async spawnSubagent(
    agentType: string,
    prompt: string,
  ): Promise<string> {
    const workOrder = this.currentWorkOrder;
    if (!workOrder) {
      return JSON.stringify({
        success: false,
        error: "OPERATION_FAILED",
        message: "No active work order. Subagent cannot be spawned outside of a work order.",
      });
    }

    const turn = workOrder.turn;
    const session = turn.session as IChatSession;

    if (!session.user) {
      return JSON.stringify({
        success: false,
        error: "OPERATION_FAILED",
        message: "ChatSession must be populated with user data.",
      });
    }

    const user = session.user as IUser;
    const provider = turn.provider as IAiProvider;

    if (!provider || typeof provider === "string") {
      return JSON.stringify({
        success: false,
        error: "OPERATION_FAILED",
        message: "Turn provider must be populated.",
      });
    }

    // Build the subagent loop
    const startTime = Date.now();
    const systemPrompt = await this.buildSubagentSystemPrompt(agentType);
    const messages: IContextChatMessage[] = [
      { createdAt: new Date(), role: "system", content: systemPrompt },
      { createdAt: new Date(), role: "user", content: prompt },
    ];

    const agentId = this.currentToolCallId || `sa_${Date.now()}_${Math.random().toString(36).slice(2, 8)}`;
    const socket = this.currentSocket;
    const tools = this.getToolsForSubagent(agentType);
    const streamHandler = socket
      ? this.makeAgentStreamHandler(socket, agentId)
      : undefined;

    let fullResponse = "";
    let fullThinking = "";
    const subagentToolCalls: IChatToolCall[] = [];
    let toolCallCount = 0;
    let inputTokens = 0;
    let outputTokens = 0;
    let thinkingTokenCount = 0;

    try {
      let continueLoop = true;
      while (continueLoop) {
        continueLoop = false;

        const chatOptions: IAiChatOptions = {
          context: messages,
          tools,
          signal: this.abortController?.signal,
        };

        this.log.info("subagent loop iteration", {
          agentType,
          messagesCount: messages.length,
          toolsAvailable: tools.map((t) => t.name),
        });

        const response = await AiService.chat(
          provider,
          this.buildDroneModelConfig(provider, turn.llm, false),
          chatOptions,
          streamHandler,
        );

        if (response.doneReason === "load" && !response.response && !response.thinking) {
          throw new Error("Subagent model failed to respond (still loading or error)");
        }

        inputTokens += response.stats?.tokenCounts?.input ?? 0;
        if (response.thinking) {
          thinkingTokenCount += Math.ceil(response.thinking.length / 4);
        }
        if (response.response) {
          fullResponse += response.response;
          outputTokens += Math.ceil(response.response.length / 4);
        }
        if (response.thinking) {
          fullThinking = (fullThinking + response.thinking).trim();
        }

        // Process tool calls
        if (response.toolCalls && response.toolCalls.length > 0) {
          continueLoop = true;
          toolCallCount += response.toolCalls.length;

          messages.push({
            createdAt: new Date(),
            role: "assistant",
            content: response.response,
            toolCalls: response.toolCalls.map((tc) => ({
              id: tc.callId,
              type: "function" as const,
              function: {
                name: tc.function.name,
                arguments: tc.function.arguments,
              },
            })),
          });

          for (const toolCall of response.toolCalls) {
            const toolArgsRaw = toolCall.function.arguments;

            if (socket) {
              socket.emit("agent:tool-call", { agentId, tool: toolCall.function.name, args: toolArgsRaw });
            }

            const result = await this.executeTool(
              toolCall.function.name,
              toolArgsRaw,
            );

            if (socket) {
              socket.emit("agent:tool-result", { agentId, tool: toolCall.function.name, result });
            }

            subagentToolCalls.push({
              callId: toolCall.callId,
              name: toolCall.function.name,
              parameters: toolArgsRaw,
              response: result,
            });

            messages.push({
              createdAt: new Date(),
              role: "tool",
              content: result,
              toolCallId: toolCall.callId,
              toolName: toolCall.function.name,
            });

            inputTokens += Math.ceil(toolArgsRaw.length / 4);
            outputTokens += Math.ceil(result.length / 4);
          }
        }
      }

      // Ensure we have a response if the subagent only made tool calls
      if (!fullResponse.trim()) {
        if (toolCallCount > 0) {
          fullResponse = `Subagent task completed. Executed ${toolCallCount} tool calls to perform the requested actions.`;
        } else {
          fullResponse = "Subagent task completed without additional actions.";
        }
      }

      const endTime = Date.now();
      const durationMs = endTime - startTime;
      const durationSeconds = Math.floor(durationMs / 1000);
      const minutes = Math.floor(durationSeconds / 60);
      const seconds = durationSeconds % 60;
      const durationLabel = `${String(minutes).padStart(2, "0")}:${String(seconds).padStart(2, "00")}`;

      const result: IChatSubagentProcess = {
        prompt,
        thinking: fullThinking || undefined,
        response: fullResponse,
        toolCalls: subagentToolCalls,
        stats: {
          toolCallCount,
          inputTokens,
          thinkingTokenCount,
          responseTokens: outputTokens,
          durationMs,
          durationLabel,
        },
      };

      return JSON.stringify({ success: true, data: result });
    } catch (error) {
      const errorMessage = error instanceof Error ? error.message : String(error);
      this.log.error("subagent execution failed", { agentType, error: errorMessage });

      return JSON.stringify({
        success: false,
        error: "SUBAGENT_FAILED",
        message: errorMessage,
      });
    }
  }

  private getToolsForSubagent(agentType: string): IAiTool[] {
    const mode = agentType === "explore"
      ? ChatSessionMode.Plan
      : ChatSessionMode.Build;

    const tools = Array.from(this.toolbox.getModeSet(mode) || []);

    if (agentType === "explore") {
      return tools.filter((t) =>
        !t.name.startsWith("plan_") &&
        t.name !== "subagent"
      );
    }

    return tools.filter((t) => t.name !== "subagent");
  }

  private async buildSubagentSystemPrompt(agentType: string): Promise<string> {
    const promptDir = path.resolve(env.installDir, "data", "prompts", "subagent", agentType);
    const templatePath = path.join(promptDir, "system.md");

    let template: string;
    try {
      template = await fs.readFile(templatePath, "utf-8");
    } catch {
      this.log.error("subagent prompt template not found", { agentType, templatePath });
      throw new Error(`Subagent prompt template not found: ${templatePath}`);
    }

    const workOrder = this.currentWorkOrder;
    if (!workOrder) return template;

    const session = workOrder.turn.session as IChatSession;
    const user = session?.user as IUser | undefined;
    const project = session?.project as IProject | undefined;
    const provider = workOrder.turn.provider as IAiProvider | undefined;

    // Build session block matching the format used in gadget-code's buildSystemPrompt
    const sessionCreated = session?.createdAt
      ? (session.createdAt instanceof Date
        ? session.createdAt.toISOString()
        : String(session.createdAt))
      : "unknown";

    const sessionBlock = [
      `Session ID: ${session?._id ?? "unknown"}`,
      `Session Name: ${session?.name ?? "unknown"}`,
      `Session Created: ${sessionCreated}`,
      `AI Provider: ${provider?.name ?? "unknown"}`,
      `AI API: ${provider?.apiType ?? "unknown"}`,
      `Model: ${workOrder.turn.llm}`,
      `Project: ${project?.name ?? "unknown"}`,
    ].join("\n");

    const personaBlock = [
      `User ID: ${user?._id ?? "unknown"}`,
      `Name: ${user?.displayName ?? "unknown"}`,
      `Is Admin: ${user?.flags?.isAdmin ? "Yes" : "No"}`,
    ].join("\n");

    template = template.replace("{{session_block}}", sessionBlock);
    template = template.replace("{{persona_block}}", personaBlock);

    // Append AGENTS.md if available in the project root
    try {
      const projectRoot = this.toolbox.env.workspace?.projectDir;
      if (projectRoot) {
        const agentsMdPath = path.join(projectRoot, "AGENTS.md");
        const agentsMd = await fs.readFile(agentsMdPath, "utf-8");
        template += "\n\n## AGENT CONFIGURATION\n\n" + agentsMd;
      }
    } catch {
      // AGENTS.md not found or inaccessible - that's fine
    }

    return template;
  }
}

export { AgentService };
export default new AgentService();