Add production source location plugins

.gitignore

@@ -7,3 +7,5 @@ mcp/dist
 .env.*
 !.env.example
 .claude/settings.local.json
+
+examples

mcp/package.json

@@ -27,6 +27,7 @@
   ],
   "scripts": {
     "build": "tsup",
+    "test": "pnpm build && node --test test/*.test.cjs",
     "watch": "tsup --watch",
     "dev": "pnpm build && pnpm watch",
     "start": "pnpm build && node dist/cli.js server",

mcp/src/server/http.ts

@@ -4,13 +4,9 @@
  */
 
 import { createServer, type IncomingMessage, type ServerResponse } from "http";
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
-import { Server } from "@modelcontextprotocol/sdk/server/index.js";
-import {
-  CallToolRequestSchema,
-  ListToolsRequestSchema,
-} from "@modelcontextprotocol/sdk/types.js";
-import { TOOLS, handleTool, error as toolError } from "./mcp.js";
+import { createAgentationMcpServer } from "./mcp.js";
 import {
   createSession,
   getSession,
@@ -66,33 +62,25 @@ const agentConnections = new Set<ServerResponse>();
 // MCP HTTP Transport
 // -----------------------------------------------------------------------------
 
-// Store transports by session ID for stateful sessions
-const mcpTransports = new Map<string, StreamableHTTPServerTransport>();
+type McpSession = {
+  server: McpServer;
+  transport: StreamableHTTPServerTransport;
+};
+
+// Keep both objects alive for the lifetime of each stateful MCP session.
+const mcpSessions = new Map<string, McpSession>();
 
 /**
  * Initialize a new MCP server with HTTP transport for a session.
  */
-function createMcpSession(): { server: Server; transport: StreamableHTTPServerTransport } {
+async function createMcpSession(): Promise<McpSession> {
   const transport = new StreamableHTTPServerTransport({
     sessionIdGenerator: () => crypto.randomUUID(),
   });
 
-  const server = new Server(
-    { name: "agentation", version: "0.0.1" },
-    { capabilities: { tools: {} } }
-  );
-
-  server.setRequestHandler(ListToolsRequestSchema, async () => ({ tools: TOOLS }));
-  server.setRequestHandler(CallToolRequestSchema, async (req) => {
-    try {
-      return await handleTool(req.params.name, req.params.arguments);
-    } catch (err) {
-      const message = err instanceof Error ? err.message : "Unknown error";
-      return toolError(message);
-    }
-  });
+  const server = createAgentationMcpServer();
 
-  server.connect(transport);
+  await server.connect(transport);
   return { server, transport };
 }
 
@@ -710,11 +698,11 @@ async function handleMcp(req: IncomingMessage, res: ServerResponse): Promise<voi
 
   // POST: Handle JSON-RPC requests
   if (method === "POST") {
-    let transport: StreamableHTTPServerTransport;
+    let mcpSession: McpSession;
 
     if (sessionId) {
       // Session ID provided - must exist in our map
-      if (!mcpTransports.has(sessionId)) {
+      if (!mcpSessions.has(sessionId)) {
         res.writeHead(404, { "Content-Type": "application/json" });
         res.end(JSON.stringify({
           jsonrpc: "2.0",
@@ -723,14 +711,14 @@ async function handleMcp(req: IncomingMessage, res: ServerResponse): Promise<voi
         }));
         return;
       }
-      transport = mcpTransports.get(sessionId)!;
+      mcpSession = mcpSessions.get(sessionId)!;
     } else {
       // No session ID - this should be an initialize request, create new session
-      const { transport: newTransport } = createMcpSession();
-      transport = newTransport;
+      mcpSession = await createMcpSession();
     }
 
     try {
+      const { transport } = mcpSession;
       // Read the request body
       const body = await new Promise<string>((resolve, reject) => {
         let data = "";
@@ -744,10 +732,10 @@ async function handleMcp(req: IncomingMessage, res: ServerResponse): Promise<voi
       // Handle the request through the transport (it writes directly to res)
       await transport.handleRequest(req, res, parsedBody);
 
-      // Store the transport with its session ID after the request is handled (for new sessions)
+      // Store the server and transport after initialization assigns a session ID.
       const newSessionId = transport.sessionId;
-      if (newSessionId && !mcpTransports.has(newSessionId)) {
-        mcpTransports.set(newSessionId, transport);
+      if (newSessionId && !mcpSessions.has(newSessionId)) {
+        mcpSessions.set(newSessionId, mcpSession);
         log(`[MCP HTTP] New session created: ${newSessionId}`);
       }
     } catch (err) {
@@ -762,13 +750,13 @@ async function handleMcp(req: IncomingMessage, res: ServerResponse): Promise<voi
 
   // GET: SSE stream for notifications
   if (method === "GET") {
-    if (!sessionId || !mcpTransports.has(sessionId)) {
+    if (!sessionId || !mcpSessions.has(sessionId)) {
       res.writeHead(400, { "Content-Type": "application/json" });
       res.end(JSON.stringify({ error: "Missing or invalid Mcp-Session-Id" }));
       return;
     }
 
-    const transport = mcpTransports.get(sessionId)!;
+    const { transport } = mcpSessions.get(sessionId)!;
 
     try {
       // Handle the SSE request (transport writes directly to res)
@@ -785,10 +773,10 @@ async function handleMcp(req: IncomingMessage, res: ServerResponse): Promise<voi
 
   // DELETE: Session cleanup
   if (method === "DELETE") {
-    if (sessionId && mcpTransports.has(sessionId)) {
-      const transport = mcpTransports.get(sessionId)!;
-      await transport.close();
-      mcpTransports.delete(sessionId);
+    if (sessionId && mcpSessions.has(sessionId)) {
+      const { server } = mcpSessions.get(sessionId)!;
+      await server.close();
+      mcpSessions.delete(sessionId);
       res.writeHead(204);
       res.end();
     } else {

mcp/src/server/mcp.ts

@@ -6,12 +6,8 @@
  * rather than maintaining its own store.
  */
 
-import { Server } from "@modelcontextprotocol/sdk/server/index.js";
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
-import {
-  CallToolRequestSchema,
-  ListToolsRequestSchema,
-} from "@modelcontextprotocol/sdk/types.js";
 import { z } from "zod";
 import type { ActionRequest } from "../types.js";
 
@@ -155,7 +151,7 @@ export const TOOLS = [
   {
     name: "agentation_get_pending",
     description:
-      "Get all pending (unacknowledged) annotations for a session. Annotations have a `kind` field: \"feedback\" (default), \"placement\" (design component placements), or \"rearrange\" (section reorder/resize). Placement and rearrange annotations include structured data.",
+      "Get all pending (unacknowledged) annotations for a session. When available, `sourceFile` is the first code lookup target for the selected element. Annotations have a `kind` field: \"feedback\" (default), \"placement\" (design component placements), or \"rearrange\" (section reorder/resize). Placement and rearrange annotations include structured data.",
     inputSchema: {
       type: "object" as const,
       properties: {
@@ -170,7 +166,7 @@ export const TOOLS = [
   {
     name: "agentation_get_all_pending",
     description:
-      "Get all pending annotations across ALL sessions. Includes feedback, design placements, and rearrange annotations. Each annotation has a `kind` field.",
+      "Get all pending annotations across ALL sessions. When available, `sourceFile` is the first code lookup target for the selected element. Includes feedback, design placements, and rearrange annotations. Each annotation has a `kind` field.",
     inputSchema: {
       type: "object" as const,
       properties: {},
@@ -254,7 +250,8 @@ export const TOOLS = [
     description:
       "Block until new annotations appear, then collect a batch and return them. " +
       "Triggers automatically when annotations are created — the user just annotates in the browser " +
-      "and the agent picks them up. Includes all annotation kinds: feedback, placement (design components), " +
+      "and the agent picks them up. When available, `sourceFile` is the first code lookup target. " +
+      "Includes all annotation kinds: feedback, placement (design components), " +
       "and rearrange (section reorder/resize). After detecting the first new annotation, waits for a batch window " +
       "to collect more before returning. Use in a loop for hands-free processing. " +
       "After addressing each annotation, call agentation_resolve with the annotation ID and a summary " +
@@ -304,6 +301,7 @@ type Annotation = {
   timestamp?: number;
   nearbyText?: string;
   reactComponents?: string;
+  sourceFile?: string;
   status: string;
   kind?: "feedback" | "placement" | "rearrange";
   placement?: {
@@ -345,6 +343,7 @@ function mapAnnotationForMcp(a: Annotation) {
     timestamp: a.timestamp,
     nearbyText: a.nearbyText,
     reactComponents: a.reactComponents,
+    sourceFile: a.sourceFile,
     ...(a.kind === "placement" && a.placement ? { placement: a.placement } : {}),
     ...(a.kind === "rearrange" && a.rearrange ? { rearrange: a.rearrange } : {}),
   };
@@ -695,6 +694,98 @@ export async function handleTool(name: string, args: unknown): Promise<ToolResul
 // Server
 // -----------------------------------------------------------------------------
 
+async function executeTool(name: string, args: unknown): Promise<ToolResult> {
+  try {
+    return await handleTool(name, args);
+  } catch (err) {
+    const message = err instanceof Error ? err.message : "Unknown error";
+    return error(message);
+  }
+}
+
+function getToolDescription(name: string): string {
+  const tool = TOOLS.find((candidate) => candidate.name === name);
+  if (!tool) throw new Error(`Missing MCP tool definition: ${name}`);
+  return tool.description;
+}
+
+/** Create a high-level MCP server with all Agentation tools registered. */
+export function createAgentationMcpServer(): McpServer {
+  const server = new McpServer({
+    name: "agentation",
+    version: "0.0.1",
+  });
+
+  server.registerTool(
+    "agentation_list_sessions",
+    { description: getToolDescription("agentation_list_sessions") },
+    () => executeTool("agentation_list_sessions", undefined),
+  );
+  server.registerTool(
+    "agentation_get_session",
+    {
+      description: getToolDescription("agentation_get_session"),
+      inputSchema: GetSessionSchema.shape,
+    },
+    (args) => executeTool("agentation_get_session", args),
+  );
+  server.registerTool(
+    "agentation_get_pending",
+    {
+      description: getToolDescription("agentation_get_pending"),
+      inputSchema: GetPendingSchema.shape,
+    },
+    (args) => executeTool("agentation_get_pending", args),
+  );
+  server.registerTool(
+    "agentation_get_all_pending",
+    { description: getToolDescription("agentation_get_all_pending") },
+    () => executeTool("agentation_get_all_pending", undefined),
+  );
+  server.registerTool(
+    "agentation_acknowledge",
+    {
+      description: getToolDescription("agentation_acknowledge"),
+      inputSchema: AcknowledgeSchema.shape,
+    },
+    (args) => executeTool("agentation_acknowledge", args),
+  );
+  server.registerTool(
+    "agentation_resolve",
+    {
+      description: getToolDescription("agentation_resolve"),
+      inputSchema: ResolveSchema.shape,
+    },
+    (args) => executeTool("agentation_resolve", args),
+  );
+  server.registerTool(
+    "agentation_dismiss",
+    {
+      description: getToolDescription("agentation_dismiss"),
+      inputSchema: DismissSchema.shape,
+    },
+    (args) => executeTool("agentation_dismiss", args),
+  );
+  server.registerTool(
+    "agentation_reply",
+    {
+      description: getToolDescription("agentation_reply"),
+      inputSchema: ReplySchema.shape,
+    },
+    (args) => executeTool("agentation_reply", args),
+  );
+  server.registerTool(
+    "agentation_watch_annotations",
+    {
+      description: getToolDescription("agentation_watch_annotations"),
+      inputSchema: WatchAnnotationsSchema.shape,
+    },
+    (args) => executeTool("agentation_watch_annotations", args),
+  );
+
+  return server;
+}
+
 /**
  * Create and start the MCP server on stdio.
  * @param baseUrl - Optional HTTP server URL to fetch from (default: http://localhost:4747)
@@ -704,33 +795,7 @@ export async function startMcpServer(baseUrl?: string): Promise<void> {
     setHttpBaseUrl(baseUrl);
   }
 
-  const server = new Server(
-    {
-      name: "agentation",
-      version: "0.0.1",
-    },
-    {
-      capabilities: {
-        tools: {},
-      },
-    }
-  );
-
-  // List available tools
-  server.setRequestHandler(ListToolsRequestSchema, async () => {
-    return { tools: TOOLS };
-  });
-
-  // Handle tool calls
-  server.setRequestHandler(CallToolRequestSchema, async (request) => {
-    const { name, arguments: args } = request.params;
-    try {
-      return await handleTool(name, args);
-    } catch (err) {
-      const message = err instanceof Error ? err.message : "Unknown error";
-      return error(message);
-    }
-  });
+  const server = createAgentationMcpServer();
 
   // Connect via stdio
   const transport = new StdioServerTransport();

mcp/src/server/sqlite.ts

@@ -102,6 +102,7 @@ function initDatabase(db: Database.Database): void {
       is_multi_select INTEGER DEFAULT 0,
       is_fixed INTEGER DEFAULT 0,
       react_components TEXT,
+      source_file TEXT,
       url TEXT,
       intent TEXT,
       severity TEXT,
@@ -215,6 +216,7 @@ function rowToAnnotation(row: Record<string, unknown>): Annotation {
     isMultiSelect: Boolean(row.is_multi_select),
     isFixed: Boolean(row.is_fixed),
     reactComponents: row.react_components as string | undefined,
+    sourceFile: row.source_file as string | undefined,
     kind,
     ...(kind === "placement" && extra?.placement ? { placement: extra.placement } : {}),
     ...(kind === "rearrange" && extra?.rearrange ? { rearrange: extra.rearrange } : {}),
@@ -243,6 +245,7 @@ export function createSQLiteStore(dbPath?: string): AFSStore {
   // Safe migrations for new columns (no-ops if already exist)
   try { db.exec("ALTER TABLE annotations ADD COLUMN kind TEXT DEFAULT 'feedback'"); } catch {}
   try { db.exec("ALTER TABLE annotations ADD COLUMN extra TEXT"); } catch {}
+  try { db.exec("ALTER TABLE annotations ADD COLUMN source_file TEXT"); } catch {}
 
   // Restore event sequence from last event
   const lastEvent = db.prepare("SELECT MAX(sequence) as seq FROM events").get() as { seq: number | null };
@@ -269,13 +272,13 @@ export function createSQLiteStore(dbPath?: string): AFSStore {
         id, session_id, x, y, comment, element, element_path, timestamp,
         selected_text, bounding_box, nearby_text, css_classes, nearby_elements,
         computed_styles, full_path, accessibility, is_multi_select, is_fixed,
-        react_components, url, intent, severity, status, thread, created_at,
+        react_components, source_file, url, intent, severity, status, thread, created_at,
         updated_at, resolved_at, resolved_by, author_id, kind, extra
       ) VALUES (
         @id, @sessionId, @x, @y, @comment, @element, @elementPath, @timestamp,
         @selectedText, @boundingBox, @nearbyText, @cssClasses, @nearbyElements,
         @computedStyles, @fullPath, @accessibility, @isMultiSelect, @isFixed,
-        @reactComponents, @url, @intent, @severity, @status, @thread, @createdAt,
+        @reactComponents, @sourceFile, @url, @intent, @severity, @status, @thread, @createdAt,
         @updatedAt, @resolvedAt, @resolvedBy, @authorId, @kind, @extra
       )
     `),
@@ -430,6 +433,7 @@ export function createSQLiteStore(dbPath?: string): AFSStore {
         isMultiSelect: annotation.isMultiSelect ? 1 : 0,
         isFixed: annotation.isFixed ? 1 : 0,
         reactComponents: annotation.reactComponents ?? null,
+        sourceFile: annotation.sourceFile ?? null,
         url: annotation.url ?? null,
         intent: annotation.intent ?? null,
         severity: annotation.severity ?? null,
@@ -612,6 +616,7 @@ export function createTenantStore(dbPath?: string): TenantStore {
   // Safe migrations for new columns (no-ops if already exist)
   try { db.exec("ALTER TABLE annotations ADD COLUMN kind TEXT DEFAULT 'feedback'"); } catch {}
   try { db.exec("ALTER TABLE annotations ADD COLUMN extra TEXT"); } catch {}
+  try { db.exec("ALTER TABLE annotations ADD COLUMN source_file TEXT"); } catch {}
 
   // Restore event sequence from last event
   const lastEvent = db.prepare("SELECT MAX(sequence) as seq FROM events").get() as { seq: number | null };