Anirudh/cleanup

PLAN.md

@@ -0,0 +1,113 @@
+I'll scan the key files related to keybindings, the command palette, and MCP services to pinpoint tight coupling and propose concrete refactors. Pulling them up now.
+
+I reviewed the command palette, keybindings, and MCP pieces. The main tight couplings are:
+- UI components directly mutate global atoms (via `getDefaultStore()`) and reach into other services (`gitChat`, MCP client) instead of using a clear boundary.
+- Command palette mixes view concerns with command registry/state, and encodes hierarchy in label strings.
+- MCP client both orchestrates connections and owns global state updates; server route also reaches directly into the same connection singleton.
+- Types and utilities (e.g., `McpUrl`) are duplicated; URL/name derivation lives in the component.
+
+Here are targeted refactors to make it more loosely coupled and testable.
+
+### decouple state management from services
+- Move Jotai writes out of service classes. Make services pure and event-driven.
+  - Define a minimal domain interface and event bus:
+```ts
+// services/mcp/ports.ts
+export type McpConnectionEvent =
+  | { type: "connecting"; server: string }
+  | { type: "connected"; server: string; tools: Record<string, Tool> }
+  | { type: "failed"; server: string; error: string }
+  | { type: "disconnected"; server: string };
+
+export interface McpConnections {
+  updateConnections(urls: McpUrl[]): Promise<void>;
+  reconnect(server: string): Promise<void>;
+  on(listener: (e: McpConnectionEvent) => void): () => void; // unsubscribe
+}
+```
+- Implement the interface in `lib/mcp-connection-manager.ts` and emit events; remove any atom knowledge from `services/mcp/client.ts`.
+- Add an adapter in `services/mcp/state-adapter.ts` that subscribes to events and updates `toolsAtom`, `breakdownAtom`, `connectionStatusAtom`, etc. UI then depends only on atoms, not on the service implementation.
+
+### introduce a command registry (DI) instead of hardcoding in `cmdk.tsx`
+- Split a registry from the view:
+  - `services/commands/registry.ts` exposes `registerCommandSource(source)` and `useCommands()` that aggregates.
+  - A `CommandSource` provides an id, optional parent id, label, and an `execute()` or `children()` function. No hierarchy in label strings.
+```ts
+// services/commands/types.ts
+export type Command = {
+  id: string;
+  title: string;
+  parentId?: string;
+  run?: () => void | Promise<void>;
+};
+export interface CommandSource {
+  id: string;
+  list(): Command[]; // pure, derived from its own store(s)
+}
+```
+- `components/cmdk.tsx` reads a derived atom `allCommandsAtom` that flattens across registered sources and renders a proper tree. Remove `HIERARCHY_SEPARATOR` and the `atomWithStorage` initialization that mixes static and dynamic content. Only persist user-defined custom commands separately if needed.
+- Feature modules register themselves:
+  - MCP: provides “MCP Servers & Tools”
+  - Models: provides “Set Model” and children for each model
+  - Chat: provides “New Thread”, “Export Chat”, “Search Messages” populated from a chat store selector instead of capturing `gitChat` directly.
+
+### centralize keybinding suspension
+- Replace repeated effects toggling `keybindingsActiveAtom` in dialogs/overlays with a hook:
+```ts
+// hooks/useSuspendKeybindings.ts
+export function useSuspendKeybindings(suspend: boolean) {
+  const set = useSetAtom(keybindingsActiveAtom);
+  useEffect(() => {
+    if (suspend) set(false);
+    return () => set(true);
+  }, [suspend, set]);
+}
+```
+- Use in `CmdK` and `McpUrlManager` to avoid direct atom writes scattered across the codebase. Keep `components/keybinding.tsx` purely presentational and pass an explicit `isActive` if you want to remove the global entirely.
+
+### clarify MCP layering and remove duplication
+- Single source of truth for `McpUrl` in `types/mcp.ts`. Import everywhere; remove duplicated interface declarations.
+- Keep one connection manager singleton per runtime boundary:
+  - In the browser: `lib/mcp-connection-manager.ts` implements `McpConnections`.
+  - On the server route: either reuse the same interface with a separate instance scoped to the request/session, or accept tools from the client rather than reconnecting again in `app/api/chat/route.ts`. Right now the route re-calls `updateConnections`, which duplicates work and couples server to the same global manager. Pick one:
+    - Option A (client-sourced): Client passes `breakdown` and `tools` snapshot with the request; server uses it directly for `streamText`.
+    - Option B (server-owned): Server maintains its own manager instance keyed by user/session; client only sends URLs. Keep the browser manager independent and don’t import it server-side.
+- Collapse `services/mcp/tools-service.ts` into a thin transformation layer:
+  - It should only convert manager output to UI format (serialization, failed-entry labeling), not call atoms.
+  - `services/mcp/client.ts` becomes an orchestrator that calls the manager and then emits domain events; the state adapter updates atoms.
+
+### remove label-encoded structure and side effects
+- Stop deriving hierarchy from `command.name.split(HIERARCHY_SEPARATOR)`. Keep structure in data (`parentId` or explicit children) and render breadcrumbs from the tree.
+- Eliminate `getDefaultStore()` mutations from outside React flows:
+  - For actions like “set model”, “open dialog”, expose explicit action functions in their own modules that update their stores. `cmdk` should invoke those actions, not poke atoms directly.
+
+### extract pure utilities and cover with tests
+- From `components/mcp-url-manager.tsx`, move:
+  - `validateUrl` and `generateServerName` into `utils/mcp.ts`.
+  - `getCleanUrlForDisplay` into the same utils.
+- Add simple unit tests for these to prevent regressions and to make the component only orchestrate UI.
+
+### improve loading/connection status modeling
+- Replace the “connecting servers” Set plus a map in atoms with either:
+  - An `atomFamily<string, "idle" | "connecting" | "connected" | "failed">` per server, or
+  - A single `Record<string, { status: ..., error?: string }>` derived from the event stream.
+- The manager should emit fine-grained per-server events; UI derives global summaries, avoiding ad-hoc state updates sprinkled across the client.
+
+### tighten server API boundary
+- `app/api/chat/route.ts` should not be aware of connection lifecycle if the client already manages it. Either:
+  - Accept a `tools` payload from the client and pass it to `streamText`, or
+  - Own connections server-side and not import the browser singleton. Right now the import couples environments and can surprise in SSR.
+
+### small cleanups
+- Co-locate atoms with their domain (commands/MCP/chat), but expose selectors instead of raw atoms to the UI where possible.
+- Unify `Tool` serialization shape under `utils/tool-serialization.ts` and only serialize in UI adapters; keep core services dealing with raw `Tool`.
+
+Status update:
+- I reviewed `components/keybinding.tsx`, `components/cmdk.tsx`, `components/mcp-url-manager.tsx`, `services/mcp/{client.ts,tools-service.ts}`, `lib/mcp-connection-manager.ts`, `utils/tool-serialization.ts`, and `app/api/chat/route.ts`. The proposals above target the specific coupling points I found. If you’d like, I can implement the minimal skeletons (ports, registry, keybinding hook) to make subsequent changes incremental.
+
+- Key changes to pursue:
+  - Define `McpConnections` interface and move Jotai writes to a state adapter.
+  - Add a command registry and remove label-encoded hierarchy from `cmdk`.
+  - Add `useSuspendKeybindings` and use it in dialogs.
+  - Deduplicate `McpUrl` type and move MCP utilities to `utils/mcp.ts`.
+  - Decide on server vs client ownership of MCP connections and adjust `app/api/chat/route.ts` accordingly.

app/api/chat/route.ts

@@ -1,5 +1,6 @@
 import { streamText } from "ai";
-import { mcpConnectionManager } from "@/lib/mcp-connection-manager";
+import { getToolsForUrls } from "@/lib/mcp-connection-core";
+import type { McpUrl } from "@/types/mcp";
 import { resolveModel } from "../apiUtils";
 
 export async function POST(req: Request) {
@@ -8,12 +9,10 @@ export async function POST(req: Request) {
   console.log("Received pendingMessageConfig:", pendingMessageConfig);
   console.log("Received mcpUrls:", mcpUrls);
 
-  // Get tools from the connection manager (which maintains persistent connections)
-  const { tools, breakdown } =
-    await mcpConnectionManager.updateConnections(mcpUrls || []);
+  // Server-owned ephemeral utilities per request (stateless between requests)
+  const tools = await getToolsForUrls((mcpUrls || []) as McpUrl[]);
 
   console.log("TOOLS", tools);
-  console.log("BREAKDOWN", breakdown);
 
   const result = streamText({
     model: resolveModel(pendingMessageConfig.modelName),

components/mcp-url-manager.tsx

@@ -27,10 +27,10 @@ import {
   toolsAtom,
   isMcpLoadingAtom,
   connectionStatusAtom,
-  type McpUrl,
 } from "@/services/mcp/atoms";
 import { keybindingsActiveAtom } from "@/services/commands/atoms";
 import mcpClient from "@/services/mcp/client";
+import { McpUrl } from "@/types/mcp";
 
 function validateUrl(url: string): { valid: boolean; message?: string } {
   try {
@@ -230,6 +230,7 @@ export function McpUrlManager() {
       id: crypto.randomUUID(),
       name: serverName,
       url: trimmedUrl,
+      callbackUrl: process.env.CALLBACK_URL,
     };
 
     const updatedUrls = [...(urls || []), newUrlObj];
@@ -286,6 +287,7 @@ export function McpUrlManager() {
             id: crypto.randomUUID(),
             name: serverName,
             url: trimmedPastedText,
+            callbackUrl: process.env.CALLBACK_URL,
           };
 
           const updatedUrls = [...(urls || []), newUrlObj];

lib/mcp-connection-core.ts

@@ -0,0 +1,171 @@
+import { Tool } from "ai";
+import { experimental_createMCPClient as createMCPClient } from "ai";
+import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
+import type { McpUrl } from "@/types/mcp";
+
+interface ManagedConnection {
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  client: any; // The MCPClient type is not exported from 'ai' package
+  tools: Record<string, Tool>;
+  url: McpUrl;
+}
+
+export class McpConnectionManager {
+  private connections: Map<string, ManagedConnection> = new Map();
+  private connectionPromises: Map<string, Promise<ManagedConnection | null>> =
+    new Map();
+
+  getAllTools(): Record<string, Tool> {
+    const allTools: Record<string, Tool> = {};
+    for (const connection of this.connections.values()) {
+      Object.assign(allTools, connection.tools);
+    }
+    return allTools;
+  }
+
+  getBreakdown(): Record<string, Record<string, Tool>> {
+    const breakdown: Record<string, Record<string, Tool>> = {};
+    for (const [serverName, connection] of this.connections.entries()) {
+      breakdown[serverName] = connection.tools;
+    }
+    return breakdown;
+  }
+
+  async connectServer(
+    urlConfig: McpUrl
+  ): Promise<{ tools: Record<string, Tool>; error?: string }> {
+    const existingPromise = this.connectionPromises.get(urlConfig.name);
+    if (existingPromise) {
+      const result = await existingPromise;
+      if (result) {
+        return { tools: result.tools };
+      } else {
+        return { tools: {}, error: "Failed to connect" };
+      }
+    }
+
+    const existingConnection = this.connections.get(urlConfig.name);
+    if (existingConnection) {
+      return { tools: existingConnection.tools };
+    }
+
+    const connectionPromise = this.createConnection(urlConfig);
+    this.connectionPromises.set(urlConfig.name, connectionPromise);
+
+    try {
+      const connection = await connectionPromise;
+      if (connection) {
+        this.connections.set(urlConfig.name, connection);
+        return { tools: connection.tools };
+      } else {
+        return { tools: {}, error: "Failed to connect" };
+      }
+    } finally {
+      this.connectionPromises.delete(urlConfig.name);
+    }
+  }
+
+  async disconnectServer(serverName: string): Promise<void> {
+    const connection = this.connections.get(serverName);
+    if (connection) {
+      try {
+        await connection.client.close();
+      } catch (error) {
+        console.error(`Error closing connection to ${serverName}:`, error);
+      }
+      this.connections.delete(serverName);
+    }
+  }
+
+  async updateConnections(newUrls: McpUrl[]): Promise<{
+    tools: Record<string, Tool>;
+    breakdown: Record<string, Record<string, Tool>>;
+    errors: Record<string, string>;
+  }> {
+    const errors: Record<string, string> = {};
+
+    const newServerNames = new Set(newUrls.map((url) => url.name));
+    const currentServerNames = new Set(this.connections.keys());
+
+    const serversToAdd = newUrls.filter(
+      (url) => !currentServerNames.has(url.name)
+    );
+    const serversToRemove = Array.from(currentServerNames).filter(
+      (name) => !newServerNames.has(name)
+    );
+
+    await Promise.all(
+      serversToRemove.map((name) => this.disconnectServer(name))
+    );
+
+    await Promise.all(
+      serversToAdd.map(async (url) => {
+        try {
+          const result = await this.connectServer(url);
+          if (result.error) {
+            errors[url.name] = result.error;
+          }
+        } catch (error) {
+          console.error(`Failed to connect to ${url.name}:`, error);
+          errors[url.name] =
+            error instanceof Error ? error.message : "Connection failed";
+        }
+      })
+    );
+
+    return {
+      tools: this.getAllTools(),
+      breakdown: this.getBreakdown(),
+      errors,
+    };
+  }
+
+  async closeAll(): Promise<void> {
+    await Promise.all(
+      Array.from(this.connections.keys()).map((name) =>
+        this.disconnectServer(name)
+      )
+    );
+  }
+
+  private async createConnection(
+    urlConfig: McpUrl
+  ): Promise<ManagedConnection | null> {
+    try {
+      console.log("Connecting to MCP server:", urlConfig.name, urlConfig.url);
+
+      const client = await createMCPClient({
+        transport: new StreamableHTTPClientTransport(new URL(urlConfig.url)),
+      });
+
+      const tools = await client.tools();
+      console.log(
+        `Successfully connected to ${urlConfig.name}, got ${
+          Object.keys(tools).length
+        } tools`
+      );
+
+      return {
+        client,
+        tools,
+        url: urlConfig,
+      };
+    } catch (error) {
+      console.error(
+        `Failed to connect to MCP server ${urlConfig.name}:`,
+        error
+      );
+      return null;
+    }
+  }
+}
+
+export async function getToolsForUrls(
+  urls: McpUrl[]
+): Promise<Record<string, Tool>> {
+  const manager = new McpConnectionManager();
+  const { tools } = await manager.updateConnections(urls);
+  // Clean up the ephemeral manager's connections after use
+  await manager.closeAll();
+  return tools;
+}