nanobot-ts/src/agent/subagent.ts

import type { MessageBus } from '../bus/queue.ts';
import type { ExecToolConfig } from '../config/types.ts';
import type { LLMProvider, ModelMessage } from '../provider/index.ts';
import { toolResultMessage } from '../provider/index.ts';
import { SessionManager } from '../session/manager.ts';
import { ToolRegistry } from './tools/base.ts';
import { EditFileTool, ListDirTool, ReadFileTool, WriteFileTool } from './tools/filesystem.ts';
import { ExecTool } from './tools/shell.ts';
import { WebFetchTool, WebSearchTool } from './tools/web.ts';

const MAX_SUBAGENT_ITERATIONS = 15;

interface SubagentTask {
  controller: AbortController;
  promise: Promise<void>;
}

export class SubagentManager {
  private _provider: LLMProvider;
  private _workspace: string;
  private _bus: MessageBus;
  private _model: string;
  private _braveApiKey: string | undefined;
  private _webProxy: string | undefined;
  private _execConfig: ExecToolConfig;
  private _restrictToWorkspace: boolean;
  private _tasks: Map<string, SubagentTask[]> = new Map();
  private _sessions: SessionManager;

  constructor(opts: {
    provider: LLMProvider;
    workspace: string;
    bus: MessageBus;
    model: string;
    braveApiKey?: string;
    webProxy?: string;
    execConfig: ExecToolConfig;
    restrictToWorkspace: boolean;
  }) {
    this._provider = opts.provider;
    this._workspace = opts.workspace;
    this._bus = opts.bus;
    this._model = opts.model;
    this._braveApiKey = opts.braveApiKey;
    this._webProxy = opts.webProxy;
    this._execConfig = opts.execConfig;
    this._restrictToWorkspace = opts.restrictToWorkspace;
    this._sessions = new SessionManager(opts.workspace);
  }

  spawn(sessionKey: string, task: string): string {
    const controller = new AbortController();
    const taskId = `subagent_${Date.now()}`;

    const promise = this._run(task, sessionKey, controller.signal).catch((err) => {
      console.error(`[subagent] Task failed for ${sessionKey}: ${err}`);
    });

    const entry: SubagentTask = { controller, promise };
    const list = this._tasks.get(sessionKey) ?? [];
    list.push(entry);
    this._tasks.set(sessionKey, list);

    // Clean up when done
    void promise.finally(() => {
      const current = this._tasks.get(sessionKey) ?? [];
      const idx = current.indexOf(entry);
      if (idx >= 0) current.splice(idx, 1);
    });

    return taskId;
  }

  async cancelBySession(sessionKey: string): Promise<number> {
    const tasks = this._tasks.get(sessionKey) ?? [];
    let count = 0;
    for (const t of tasks) {
      t.controller.abort();
      count++;
    }
    await Promise.allSettled(tasks.map((t) => t.promise));
    this._tasks.delete(sessionKey);
    return count;
  }

  private async _run(task: string, sessionKey: string, signal: AbortSignal): Promise<void> {
    const tools = this._buildTools();

    const systemPrompt = `You are a background subagent. Complete the following task autonomously using the available tools. When done, write a brief summary of what you accomplished. Do not ask for clarification — make your best effort.

Task: ${task}`;

    const messages: ModelMessage[] = [{ role: 'user', content: systemPrompt }];

    for (let i = 0; i < MAX_SUBAGENT_ITERATIONS; i++) {
      if (signal.aborted) break;

      const { response, responseMessages } = await this._provider.chatWithRetry({
        messages,
        tools: tools.getDefinitions(),
        model: this._model,
      });

      if (signal.aborted) break;

      messages.push(...responseMessages);

      if (response.finishReason !== 'tool-calls' || response.toolCalls.length === 0) {
        // Done — report result back to main agent via system channel
        const content = response.content ?? 'Subagent completed with no output.';
        this._bus.publishInbound({
          channel: 'system',
          senderId: 'subagent',
          chatId: sessionKey,
          content: `Subagent result:\n${content}`,
          metadata: {},
        });
        return;
      }

      // Execute tool calls
      for (const tc of response.toolCalls) {
        if (signal.aborted) break;
        const result = await tools.execute(tc.name, tc.arguments);
        messages.push(toolResultMessage(tc.id, tc.name, result));
      }
    }

    if (!signal.aborted) {
      this._bus.publishInbound({
        channel: 'system',
        senderId: 'subagent',
        chatId: sessionKey,
        content: 'Subagent reached max iterations without completing the task.',
        metadata: {},
      });
    }
  }

  private _buildTools(): ToolRegistry {
    const registry = new ToolRegistry();
    const allowed = this._restrictToWorkspace ? this._workspace : undefined;

    registry.register(new ReadFileTool({ workspace: this._workspace, allowedDir: allowed }));
    registry.register(new WriteFileTool({ workspace: this._workspace, allowedDir: allowed }));
    registry.register(new EditFileTool({ workspace: this._workspace, allowedDir: allowed }));
    registry.register(new ListDirTool({ workspace: this._workspace }));
    registry.register(
      new ExecTool({
        workspacePath: this._workspace,
        timeoutS: this._execConfig.timeout,
        restrictToWorkspace: this._restrictToWorkspace,
        pathAppend: this._execConfig.pathAppend,
      }),
    );
    registry.register(new WebSearchTool({ apiKey: this._braveApiKey, proxy: this._webProxy }));
    registry.register(new WebFetchTool({ proxy: this._webProxy }));

    return registry;
  }
}