stagent 0.5.0 → 0.6.1

package/src/lib/agents/agentic-loop.ts

@@ -0,0 +1,235 @@
+/**
+ * Provider-agnostic agentic loop for direct API runtimes.
+ *
+ * The loop handles turn counting, budget tracking, abort signaling,
+ * and HITL tool permission checks. Provider-specific logic (API calls,
+ * event mapping, tool result formatting) is injected via callbacks.
+ */
+
+import type { ToolResult } from "@/lib/chat/tool-registry";
+import type { ToolPermissionResponse } from "./tool-permissions";
+
+// ── Types ────────────────────────────────────────────────────────────
+
+/** A single tool call extracted from the model response. */
+export interface ToolCall {
+  id: string;
+  name: string;
+  arguments: Record<string, unknown>;
+}
+
+/** Usage snapshot from a single model turn. */
+export interface TurnUsage {
+  inputTokens?: number;
+  outputTokens?: number;
+  totalTokens?: number;
+  modelId?: string;
+  costUsd?: number;
+}
+
+/** Events emitted during the loop for SSE streaming. */
+export type AgentStreamEvent =
+  | { type: "status"; phase: "running" | "tool_use" | "thinking"; message?: string }
+  | { type: "delta"; content: string }
+  | { type: "done"; finalText: string }
+  | { type: "error"; message: string };
+
+/** Result of a single model API call (accumulated from stream). */
+export interface ModelTurnResult {
+  /** Concatenated text output from the model. */
+  text: string;
+  /** Tool calls requested by the model. */
+  toolCalls: ToolCall[];
+  /** Whether the model indicated it is done (end_turn / stop). */
+  isComplete: boolean;
+  /** Whether output was truncated by max_tokens. */
+  needsContinuation: boolean;
+  /** Usage for this turn. */
+  usage: TurnUsage;
+}
+
+/** Message in the conversation history (provider-agnostic shape). */
+export type LoopMessage = Record<string, unknown>;
+
+/** Configuration for the agentic loop — provider injects callbacks. */
+export interface AgenticLoopConfig {
+  /**
+   * Call the model API with the current messages. Must stream events
+   * via `emitEvent` and return the accumulated turn result.
+   */
+  callModel: (
+    messages: LoopMessage[],
+    signal: AbortSignal,
+  ) => Promise<ModelTurnResult>;
+
+  /** Format a tool result for appending to the message history. */
+  formatToolResult: (
+    toolCallId: string,
+    toolName: string,
+    result: ToolResult,
+  ) => LoopMessage;
+
+  /** Format a continuation message (e.g. after max_tokens truncation). */
+  formatContinuation: () => LoopMessage;
+
+  /** Execute a Stagent tool by name. */
+  executeTool: (
+    name: string,
+    args: Record<string, unknown>,
+  ) => Promise<ToolResult>;
+
+  /** HITL permission check. Return allow/deny. */
+  checkPermission: (
+    toolName: string,
+    args: Record<string, unknown>,
+  ) => Promise<ToolPermissionResponse>;
+
+  /** Emit SSE event for real-time UI streaming. */
+  emitEvent: (event: AgentStreamEvent) => void;
+
+  /** Maximum model turns before stopping. */
+  maxTurns: number;
+
+  /** Maximum budget in USD before stopping. */
+  maxBudgetUsd?: number;
+
+  /** Abort signal for cancellation. */
+  signal: AbortSignal;
+}
+
+/** Result of the agentic loop. */
+export interface AgenticLoopResult {
+  finalText: string;
+  turnCount: number;
+  totalUsage: TurnUsage;
+  stopReason: "complete" | "max_turns" | "budget_exceeded" | "cancelled" | "error";
+}
+
+// ── Loop implementation ──────────────────────────────────────────────
+
+function mergeTurnUsage(total: TurnUsage, turn: TurnUsage): TurnUsage {
+  return {
+    inputTokens: (total.inputTokens ?? 0) + (turn.inputTokens ?? 0),
+    outputTokens: (total.outputTokens ?? 0) + (turn.outputTokens ?? 0),
+    totalTokens: (total.totalTokens ?? 0) + (turn.totalTokens ?? 0),
+    modelId: turn.modelId ?? total.modelId,
+    costUsd: (total.costUsd ?? 0) + (turn.costUsd ?? 0),
+  };
+}
+
+/**
+ * Run a provider-agnostic agentic loop.
+ *
+ * Repeatedly calls the model, handles tool execution with HITL checks,
+ * and enforces turn/budget limits until the model completes or a limit
+ * is reached.
+ */
+export async function runAgenticLoop(
+  initialMessages: LoopMessage[],
+  config: AgenticLoopConfig,
+): Promise<AgenticLoopResult> {
+  const messages = [...initialMessages];
+  let turnCount = 0;
+  let totalUsage: TurnUsage = {};
+  let lastText = "";
+
+  while (turnCount < config.maxTurns) {
+    // Check cancellation
+    if (config.signal.aborted) {
+      return { finalText: lastText, turnCount, totalUsage, stopReason: "cancelled" };
+    }
+
+    // Check budget
+    if (config.maxBudgetUsd && (totalUsage.costUsd ?? 0) >= config.maxBudgetUsd) {
+      config.emitEvent({ type: "error", message: "Budget limit exceeded" });
+      return { finalText: lastText, turnCount, totalUsage, stopReason: "budget_exceeded" };
+    }
+
+    // Call model
+    turnCount++;
+    let turnResult: ModelTurnResult;
+
+    try {
+      turnResult = await config.callModel(messages, config.signal);
+    } catch (err) {
+      if (config.signal.aborted) {
+        return { finalText: lastText, turnCount, totalUsage, stopReason: "cancelled" };
+      }
+      const message = err instanceof Error ? err.message : "Model API call failed";
+      config.emitEvent({ type: "error", message });
+      return { finalText: lastText, turnCount, totalUsage, stopReason: "error" };
+    }
+
+    totalUsage = mergeTurnUsage(totalUsage, turnResult.usage);
+    if (turnResult.text) lastText = turnResult.text;
+
+    // Handle completion
+    if (turnResult.isComplete && turnResult.toolCalls.length === 0) {
+      config.emitEvent({ type: "done", finalText: lastText });
+      return { finalText: lastText, turnCount, totalUsage, stopReason: "complete" };
+    }
+
+    // Handle tool calls
+    if (turnResult.toolCalls.length > 0) {
+      for (const toolCall of turnResult.toolCalls) {
+        if (config.signal.aborted) {
+          return { finalText: lastText, turnCount, totalUsage, stopReason: "cancelled" };
+        }
+
+        config.emitEvent({
+          type: "status",
+          phase: "tool_use",
+          message: toolCall.name,
+        });
+
+        // HITL permission check
+        const permission = await config.checkPermission(
+          toolCall.name,
+          toolCall.arguments,
+        );
+
+        let result: ToolResult;
+        if (permission.behavior === "deny") {
+          result = {
+            content: [{ type: "text", text: JSON.stringify({ error: permission.message ?? "Tool denied by user" }) }],
+            isError: true,
+          };
+        } else {
+          try {
+            result = await config.executeTool(
+              toolCall.name,
+              (permission.updatedInput as Record<string, unknown>) ?? toolCall.arguments,
+            );
+          } catch (err) {
+            result = {
+              content: [{ type: "text", text: JSON.stringify({ error: err instanceof Error ? err.message : "Tool execution failed" }) }],
+              isError: true,
+            };
+          }
+        }
+
+        // Append tool result to messages
+        messages.push(
+          config.formatToolResult(toolCall.id, toolCall.name, result),
+        );
+      }
+
+      // Continue loop — model needs to process tool results
+      continue;
+    }
+
+    // Handle max_tokens continuation
+    if (turnResult.needsContinuation) {
+      messages.push(config.formatContinuation());
+      continue;
+    }
+
+    // Shouldn't reach here — safeguard
+    config.emitEvent({ type: "done", finalText: lastText });
+    return { finalText: lastText, turnCount, totalUsage, stopReason: "complete" };
+  }
+
+  // Max turns exceeded
+  config.emitEvent({ type: "error", message: `Max turns (${config.maxTurns}) reached` });
+  return { finalText: lastText, turnCount, totalUsage, stopReason: "max_turns" };
+}

package/src/lib/agents/browser-mcp.ts

@@ -1,13 +1,22 @@
 import { getSetting } from "@/lib/settings/helpers";
 import { SETTINGS_KEYS } from "@/lib/constants/settings";
 
-// ── MCP server config type (matches Claude Agent SDK shape) ──────────
+// ── MCP server config types (matches Claude Agent SDK shape) ─────────
 
-interface McpServerConfig {
+interface McpStdioConfig {
+  type?: "stdio";
   command: string;
   args: string[];
 }
 
+interface McpHttpConfig {
+  type: "http";
+  url: string;
+  headers?: Record<string, string>;
+}
+
+type AnyMcpServerConfig = McpStdioConfig | McpHttpConfig;
+
 // ── Read-only browser tools — auto-approved in chat & task permission callbacks
 
 export const BROWSER_READ_ONLY_TOOLS = new Set([
@@ -72,7 +81,7 @@ function parseExtraArgs(config: string | null): string[] {
  *
  * Returns `{}` when neither server is enabled — zero overhead.
  */
-export async function getBrowserMcpServers(): Promise<Record<string, McpServerConfig>> {
+export async function getBrowserMcpServers(): Promise<Record<string, McpStdioConfig>> {
   const [chromeEnabled, playwrightEnabled, chromeConfig, playwrightConfig] =
     await Promise.all([
       getSetting(SETTINGS_KEYS.BROWSER_MCP_CHROME_DEVTOOLS_ENABLED),
@@ -81,7 +90,7 @@ export async function getBrowserMcpServers(): Promise<Record<string, McpServerCo
       getSetting(SETTINGS_KEYS.BROWSER_MCP_PLAYWRIGHT_CONFIG),
     ]);
 
-  const servers: Record<string, McpServerConfig> = {};
+  const servers: Record<string, McpStdioConfig> = {};
 
   if (chromeEnabled === "true") {
     const extraArgs = parseExtraArgs(chromeConfig);
@@ -117,3 +126,49 @@ export async function getBrowserAllowedToolPatterns(): Promise<string[]> {
   if (playwrightEnabled === "true") patterns.push("mcp__playwright__*");
   return patterns;
 }
+
+// ── Exa Search MCP — semantic web search ────────────────────────────
+
+/** All Exa tools are read-only (search, similarity, content fetch) */
+export const EXA_READ_ONLY_TOOLS = new Set([
+  "mcp__exa__web_search_exa",
+  "mcp__exa__find_similar",
+  "mcp__exa__get_contents",
+]);
+
+export function isExaTool(toolName: string): boolean {
+  return toolName.startsWith("mcp__exa__");
+}
+
+export function isExaReadOnly(toolName: string): boolean {
+  return EXA_READ_ONLY_TOOLS.has(toolName);
+}
+
+/**
+ * Read external MCP server settings from DB and return configs
+ * for any enabled servers. Currently supports Exa Search.
+ *
+ * Returns `{}` when nothing is enabled — zero overhead.
+ */
+export async function getExternalMcpServers(): Promise<Record<string, AnyMcpServerConfig>> {
+  const exaEnabled = await getSetting(SETTINGS_KEYS.EXA_SEARCH_MCP_ENABLED);
+
+  const servers: Record<string, AnyMcpServerConfig> = {};
+
+  if (exaEnabled === "true") {
+    servers.exa = { type: "http", url: "https://mcp.exa.ai/mcp" };
+  }
+
+  return servers;
+}
+
+/**
+ * Build the allowedTools glob patterns for enabled external MCP servers.
+ */
+export async function getExternalAllowedToolPatterns(): Promise<string[]> {
+  const exaEnabled = await getSetting(SETTINGS_KEYS.EXA_SEARCH_MCP_ENABLED);
+
+  const patterns: string[] = [];
+  if (exaEnabled === "true") patterns.push("mcp__exa__*");
+  return patterns;
+}

package/src/lib/agents/claude-agent.ts

@@ -1,5 +1,4 @@
 import { query } from "@anthropic-ai/claude-agent-sdk";
-import { z } from "zod";
 import { db } from "@/lib/db";
 import { tasks, projects, agentLogs, notifications } from "@/lib/db/schema";
 import { eq } from "drizzle-orm";
@@ -20,7 +19,7 @@ import { getActiveLearnedContext } from "./learned-context";
 import { getLaunchCwd, getWorkspaceContext } from "@/lib/environment/workspace-context";
 import { analyzeForLearnedPatterns } from "./pattern-extractor";
 import { processSweepResult } from "./sweep";
-import { getBrowserMcpServers } from "./browser-mcp";
+import { getBrowserMcpServers, getExternalMcpServers } from "./browser-mcp";
 import { persistScreenshot, SCREENSHOT_TOOL_NAMES } from "@/lib/screenshots/persist";
 import {
   extractUsageSnapshot,
@@ -30,6 +29,10 @@ import {
   type UsageActivityType,
   type UsageSnapshot,
 } from "@/lib/usage/ledger";
+import {
+  handleToolPermission,
+  clearPermissionCache,
+} from "./tool-permissions";
 
 /** Typed representation of messages from the Agent SDK stream */
 interface AgentStreamMessage {
@@ -44,7 +47,7 @@ interface AgentStreamMessage {
   result?: unknown;
 }
 
-interface TaskUsageState extends UsageSnapshot {
+export interface TaskUsageState extends UsageSnapshot {
   activityType: UsageActivityType;
   startedAt: Date;
   taskId: string;
@@ -53,44 +56,7 @@ interface TaskUsageState extends UsageSnapshot {
   scheduleId?: string | null;
 }
 
-const toolPermissionResponseSchema = z.object({
-  behavior: z.enum(["allow", "deny"]),
-  updatedInput: z.unknown().optional(),
-  message: z.string().optional(),
-});
-
-type ToolPermissionResponse = z.infer<typeof toolPermissionResponseSchema>;
-
-const inFlightPermissionRequests = new Map<
-  string,
-  Promise<ToolPermissionResponse>
->();
-const settledPermissionRequests = new Map<string, ToolPermissionResponse>();
-
-function buildAllowedToolPermissionResponse(
-  input: Record<string, unknown>
-): ToolPermissionResponse {
-  return {
-    behavior: "allow",
-    updatedInput: input,
-  };
-}
-
-function normalizeToolPermissionResponse(
-  response: ToolPermissionResponse,
-  input: Record<string, unknown>
-): ToolPermissionResponse {
-  if (response.behavior !== "allow" || response.updatedInput !== undefined) {
-    return response;
-  }
-
-  return {
-    ...response,
-    updatedInput: input,
-  };
-}
-
-function createTaskUsageState(
+export function createTaskUsageState(
   task: {
     id: string;
     projectId?: string | null;
@@ -117,64 +83,7 @@ function applyUsageSnapshot(state: TaskUsageState, source: unknown) {
   Object.assign(state, mergeUsageSnapshot(state, extractUsageSnapshot(source)));
 }
 
-function buildPermissionCacheKey(
-  taskId: string,
-  toolName: string,
-  input: Record<string, unknown>
-): string {
-  return `${taskId}::${toolName}::${JSON.stringify(input)}`;
-}
-
-function clearPermissionCache(taskId: string) {
-  const prefix = `${taskId}::`;
-
-  for (const key of inFlightPermissionRequests.keys()) {
-    if (key.startsWith(prefix)) {
-      inFlightPermissionRequests.delete(key);
-    }
-  }
-
-  for (const key of settledPermissionRequests.keys()) {
-    if (key.startsWith(prefix)) {
-      settledPermissionRequests.delete(key);
-    }
-  }
-}
-
-async function waitForToolPermissionResponse(
-  notificationId: string
-): Promise<ToolPermissionResponse> {
-  const deadline = Date.now() + 55_000;
-  const pollInterval = 1500;
-
-  while (Date.now() < deadline) {
-    const [notification] = await db
-      .select()
-      .from(notifications)
-      .where(eq(notifications.id, notificationId));
-
-    if (notification?.response) {
-      try {
-        const parsed = JSON.parse(notification.response);
-        const validated = toolPermissionResponseSchema.safeParse(parsed);
-        if (validated.success) {
-          return validated.data;
-        }
-        console.error("[claude-agent] Invalid permission response shape:", validated.error.message);
-        return { behavior: "deny", message: "Invalid response format" };
-      } catch (err) {
-        console.error("[claude-agent] Failed to parse permission response:", err);
-        return { behavior: "deny", message: "Invalid response format" };
-      }
-    }
-
-    await new Promise((resolve) => setTimeout(resolve, pollInterval));
-  }
-
-  return { behavior: "deny", message: "Permission request timed out" };
-}
-
-async function finalizeTaskUsage(
+export async function finalizeTaskUsage(
   state: TaskUsageState,
   status: "completed" | "failed" | "cancelled"
 ) {
@@ -272,8 +181,9 @@ async function processAgentStream(
       for (const block of message.message.content) {
         if (block.type === "tool_use") {
           // Track screenshot tool_use IDs for result interception
-          if (typeof block.name === "string" && SCREENSHOT_TOOL_NAMES.has(block.name) && typeof block.id === "string") {
-            pendingScreenshotTools.add(block.id);
+          const toolBlock = block as { type: string; id?: string; name?: string; input?: unknown };
+          if (typeof toolBlock.name === "string" && SCREENSHOT_TOOL_NAMES.has(toolBlock.name) && typeof toolBlock.id === "string") {
+            pendingScreenshotTools.add(toolBlock.id);
           }
           await db.insert(agentLogs).values({
             id: crypto.randomUUID(),
@@ -429,7 +339,7 @@ async function processAgentStream(
 // Shared prompt & query context builder (F12: eliminate duplication)
 // ---------------------------------------------------------------------------
 
-interface TaskQueryContext {
+export interface TaskQueryContext {
   /** User task content — goes into `prompt` */
   userPrompt: string;
   /** System instructions — goes into `options.systemPrompt` */
@@ -444,7 +354,7 @@ interface TaskQueryContext {
   canUseToolPolicy?: CanUseToolPolicy;
 }
 
-async function buildTaskQueryContext(
+export async function buildTaskQueryContext(
   task: { id: string; title: string; description?: string | null; projectId?: string | null },
   profileId: string
 ): Promise<TaskQueryContext> {
@@ -462,7 +372,7 @@ async function buildTaskQueryContext(
   const outputInstructions = buildTaskOutputInstructions(task.id);
   const learnedCtx = getActiveLearnedContext(profileId);
   const learnedCtxBlock = learnedCtx
-    ? `## Learned Context\nPatterns and insights learned from previous tasks:\n\n${learnedCtx}`
+    ? `## Learned Context\n<learned-context>\n${learnedCtx}\n</learned-context>`
     : "";
 
   // Resolve working directory: project's workingDirectory > launch cwd
@@ -520,10 +430,13 @@ export async function executeClaudeTask(taskId: string): Promise<void> {
     await prepareTaskOutputDirectory(taskId, { clearExisting: true });
     const ctx = await buildTaskQueryContext(task, agentProfileId);
 
-    // Merge browser MCP servers when enabled globally
-    const browserServers = await getBrowserMcpServers();
+    // Merge browser + external MCP servers when enabled globally
+    const [browserServers, externalServers] = await Promise.all([
+      getBrowserMcpServers(),
+      getExternalMcpServers(),
+    ]);
     const profileMcpServers = ctx.payload?.mcpServers ?? {};
-    const mergedMcpServers = { ...profileMcpServers, ...browserServers };
+    const mergedMcpServers = { ...profileMcpServers, ...browserServers, ...externalServers };
 
     const authEnv = await getAuthEnv();
     const response = query({
@@ -630,10 +543,13 @@ export async function resumeClaudeTask(taskId: string): Promise<void> {
     await prepareTaskOutputDirectory(taskId);
     const ctx = await buildTaskQueryContext(task, profileId);
 
-    // Merge browser MCP servers when enabled globally
-    const browserServers = await getBrowserMcpServers();
+    // Merge browser + external MCP servers when enabled globally
+    const [browserServers, externalServers] = await Promise.all([
+      getBrowserMcpServers(),
+      getExternalMcpServers(),
+    ]);
     const profileMcpServers = ctx.payload?.mcpServers ?? {};
-    const mergedMcpServers = { ...profileMcpServers, ...browserServers };
+    const mergedMcpServers = { ...profileMcpServers, ...browserServers, ...externalServers };
 
     const authEnv = await getAuthEnv();
     const response = query({
@@ -781,93 +697,4 @@ async function handleExecutionError(
   }
 }
 
-/**
- * Handle tool permission by inserting a notification and polling for response.
- * Uses database polling pattern — the Inbox UI writes the response.
- */
-async function handleToolPermission(
-  taskId: string,
-  toolName: string,
-  input: Record<string, unknown>,
-  canUseToolPolicy?: CanUseToolPolicy
-): Promise<ToolPermissionResponse> {
-  const isQuestion = toolName === "AskUserQuestion";
-
-  // Layer 1: Profile-level canUseToolPolicy — fastest check, no I/O
-  if (!isQuestion && canUseToolPolicy) {
-    if (canUseToolPolicy.autoApprove?.includes(toolName)) {
-      return buildAllowedToolPermissionResponse(input);
-    }
-    if (canUseToolPolicy.autoDeny?.includes(toolName)) {
-      return { behavior: "deny", message: `Profile policy denies ${toolName}` };
-    }
-  }
-
-  // Layer 2: Saved user permissions — skip notification for pre-approved tools
-  if (!isQuestion) {
-    const { isToolAllowed } = await import("@/lib/settings/permissions");
-    if (await isToolAllowed(toolName, input)) {
-      return buildAllowedToolPermissionResponse(input);
-    }
-  }
-
-  if (!isQuestion) {
-    const cacheKey = buildPermissionCacheKey(taskId, toolName, input);
-    const settledResponse = settledPermissionRequests.get(cacheKey);
-    if (settledResponse) {
-      return normalizeToolPermissionResponse(settledResponse, input);
-    }
-
-    const pendingRequest = inFlightPermissionRequests.get(cacheKey);
-    if (pendingRequest) {
-      return pendingRequest;
-    }
-
-    const requestPromise = (async () => {
-      const notificationId = crypto.randomUUID();
-
-      await db.insert(notifications).values({
-        id: notificationId,
-        taskId,
-        type: "permission_required",
-        title: `Permission required: ${toolName}`,
-        body: JSON.stringify(input).slice(0, 1000),
-        toolName,
-        toolInput: JSON.stringify(input),
-        createdAt: new Date(),
-      });
-
-      const response = normalizeToolPermissionResponse(
-        await waitForToolPermissionResponse(notificationId),
-        input
-      );
-      settledPermissionRequests.set(cacheKey, response);
-      return response;
-    })();
-
-    inFlightPermissionRequests.set(cacheKey, requestPromise);
-
-    try {
-      return await requestPromise;
-    } finally {
-      inFlightPermissionRequests.delete(cacheKey);
-    }
-  }
-
-  const notificationId = crypto.randomUUID();
-
-  await db.insert(notifications).values({
-    id: notificationId,
-    taskId,
-    type: isQuestion ? "agent_message" : "permission_required",
-    title: isQuestion
-      ? "Agent has a question"
-      : `Permission required: ${toolName}`,
-    body: JSON.stringify(input).slice(0, 1000),
-    toolName,
-    toolInput: JSON.stringify(input),
-    createdAt: new Date(),
-  });
-
-  return waitForToolPermissionResponse(notificationId);
-}
+// handleToolPermission and clearPermissionCache imported from ./tool-permissions