fix(concurrency): per-session MCP client pool — removes cross-task tool-call blocking

Fix 3 of plans/2026-07-11-concurrent-task-execution.md, the throughput one. nomos held exactly one *mcpClient for the whole process, shared by every /chat goroutine. Its mutex was held for the full duration of each tool round-trip, and `run` executes its SSH command SYNCHRONOUSLY inside that round-trip (capped at up to 10 minutes) — so while Task A was mid-`run`, every other task's tool calls, even a trivial get_entity, queued behind that single lock. Tasks could think (LLM calls) in parallel but never act in parallel. The MCP server has no per-connection state to protect (newServer returns one shared *mcp.Server instance whose handlers close only over the DB connection pool, already safe for concurrent use) — the mutex existed purely because the client reused one stateful transport session. So the fix doesn't touch the server at all: - New mcpClientPool (cmd/nomos/main.go): one *mcpClient per session id, created lazily (a real MCP initialize handshake) on first use and cached; session-less traffic (the ephemeral no-DB-store path, the structured /query endpoint) gets its own fixed, reused key instead of a fresh connection per request. Idle clients (20 min past last use — long enough to outlive a single slow `run`) are evicted on a 5-minute sweep ticker. - agent.go: `client *mcpClient` → `clients *mcpClientPool`; every call site (buildTools, fleetSnapshot, the tool-dispatch loop) now resolves its own session's client via clients.get(sessionID) instead of reaching for one shared field. A task's own tool calls stay sequential (already true — the agent loop calls tools one at a time within a turn) but no longer block anyone else's. - main.go: handleQuery takes the pool instead of a client (keyed "query", a fixed non-session slot); shutdown calls pool.closeAll(). Verified live against the deployed stack: fired a slow-but-ungated command (`ping -c 15 127.0.0.1`, read-only per policy's allowlist, no approval needed) as Task A, then — 2s into A's run — a trivial hostname lookup as Task B, both through the real /chat endpoint. Task A's ping genuinely ran ~14.3s (confirmed via its own execution record and the agent's reported output). Task B returned in 6s total, well before A finished — proving it was never queued behind A's connection. Before this fix, B would have been forced to wait out A's entire ~14.3s hold on the single shared client. This completes plans/2026-07-11-concurrent-task-execution.md's required scope — only the explicitly optional/deferred Fix 4 (a concurrency/cost cap, pending real usage data) remains. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-11 19:15:32 +02:00
parent 6a8fb435ad
commit a4ea542f3e
2 changed files with 164 additions and 17 deletions
--- a/cmd/nomos/agent.go
+++ b/cmd/nomos/agent.go
@@ -33,7 +33,7 @@ var refusalDenylist = []string{
 }

 type agent struct {
-	client     *mcpClient
+	clients    *mcpClientPool // one MCP client PER SESSION, not shared — see mcpClientPool's doc comment
 	system     string
 	provider   *openai.Client
 	model      string
@@ -44,7 +44,7 @@ type agent struct {
 	httpClient *http.Client
 }

-func newAgent(ctx context.Context, mcpClient *mcpClient, st *store, agentSlug string) (*agent, error) {
+func newAgent(ctx context.Context, clients *mcpClientPool, st *store, agentSlug string) (*agent, error) {
 	system := loadSoul()
 	apiKey := os.Getenv("OPENROUTER_API_KEY")
 	model := os.Getenv("NOMOS_MODEL")
@@ -92,7 +92,7 @@ func newAgent(ctx context.Context, mcpClient *mcpClient, st *store, agentSlug st
 	}

 	return &agent{
-		client:     mcpClient,
+		clients:    clients,
 		system:     system,
 		provider:   &provider,
 		model:      model,
@@ -171,14 +171,14 @@ func (a *agent) chat(ctx context.Context, sessionID, message string, emit func(a
 func (a *agent) chatWith(ctx context.Context, sessionID, message, systemInject string, emit func(agentEvent)) {
 	correlationID := uuid.New().String()

-	tools, err := a.buildTools()
+	tools, err := a.buildTools(sessionID)
 	if err != nil {
 		emit(agentEvent{Type: "error", Data: fmt.Sprintf("build tools: %v", err), SessionID: sessionID})
 		return
 	}

 	system := a.system
-	if snapshot := a.fleetSnapshot(); snapshot != "" {
+	if snapshot := a.fleetSnapshot(sessionID); snapshot != "" {
 		system += "\n\n" + snapshot
 	}
 	messages := []openai.ChatCompletionMessageParamUnion{openai.SystemMessage(system)}
@@ -382,7 +382,11 @@ func (a *agent) chatWith(ctx context.Context, sessionID, message, systemInject s
 					wireArgs[k] = v
 				}
 				wireArgs["_session_id"] = sessionID
-				result, callErr = a.client.callTool(tc.Function.Name, wireArgs)
+				var client *mcpClient
+				client, callErr = a.clients.get(sessionID)
+				if callErr == nil {
+					result, callErr = client.callTool(tc.Function.Name, wireArgs)
+				}
 			}
 			elapsed := int(time.Since(start).Milliseconds())

@@ -595,8 +599,12 @@ func assistantToolCallMessage(calls []persistedCall) openai.ChatCompletionMessag
 // of spending its first iteration rediscovering topology it already has
 // tools to query. Best-effort: an empty string on any failure just means no
 // snapshot, not an error for the turn.
-func (a *agent) fleetSnapshot() string {
-	result, err := a.client.callTool("get_health_summary", map[string]any{})
+func (a *agent) fleetSnapshot(sessionID string) string {
+	client, err := a.clients.get(sessionID)
+	if err != nil {
+		return ""
+	}
+	result, err := client.callTool("get_health_summary", map[string]any{})
 	if err != nil {
 		return ""
 	}
@@ -659,8 +667,12 @@ func isRefusalOrEmpty(text string) bool {
 	return false
 }

-func (a *agent) buildTools() ([]openai.ChatCompletionToolParam, error) {
-	defs, err := a.client.listToolsFull()
+func (a *agent) buildTools(sessionID string) ([]openai.ChatCompletionToolParam, error) {
+	client, err := a.clients.get(sessionID)
+	if err != nil {
+		return nil, err
+	}
+	defs, err := client.listToolsFull()
 	if err != nil {
 		return nil, err
 	}
--- a/cmd/nomos/main.go
+++ b/cmd/nomos/main.go
@@ -42,10 +42,19 @@ func main() {
 		ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGTERM, syscall.SIGINT)
 		defer cancel()

-		client, err := newMCPClient(mcpURL)
-		if err != nil {
+		// One MCP client PER SESSION, not one shared client for the whole
+		// process — see mcpClientPool's doc comment. A dedicated client is
+		// created lazily on each session's first tool call.
+		clientPool := newMCPClientPool(mcpURL)
+		// Prove connectivity at startup the same way the old single-client
+		// constructor did, so a misconfigured/unreachable MCP endpoint still
+		// fails fast on boot instead of only on the first real chat. Doesn't
+		// reuse the pool (nothing to key it by yet) — just a throwaway probe.
+		if probe, err := newMCPClient(mcpURL); err != nil {
 			slog.Error("nomos: mcp connect", "url", mcpURL, "error", err)
 			os.Exit(1)
+		} else {
+			probe.close()
 		}

 		st, err := newStore(ctx, databaseURL)
@@ -57,7 +66,7 @@ func main() {
 			defer st.close()
 		}

-		nAgent, err := newAgent(ctx, client, st, agentSlug)
+		nAgent, err := newAgent(ctx, clientPool, st, agentSlug)
 		if err != nil {
 			slog.Error("nomos: agent init", "error", err)
 			os.Exit(1)
@@ -68,13 +77,26 @@ func main() {
 		// from failures) without the operator ticking it forward each step.
 		go nAgent.runContinuationWorker(ctx)

+		go func() {
+			ticker := time.NewTicker(5 * time.Minute)
+			defer ticker.Stop()
+			for {
+				select {
+				case <-ctx.Done():
+					return
+				case <-ticker.C:
+					clientPool.sweep()
+				}
+			}
+		}()
+
 		mux := http.NewServeMux()
 		mux.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {
 			w.WriteHeader(200)
 			w.Write([]byte("ok"))
 		})
 		mux.HandleFunc("/query", func(w http.ResponseWriter, r *http.Request) {
-			handleQuery(w, r, client, agentSlug, mcpURL)
+			handleQuery(w, r, clientPool, agentSlug, mcpURL)
 		})
 		mux.HandleFunc("/chat", func(w http.ResponseWriter, r *http.Request) {
 			handleChat(w, r, nAgent, st)
@@ -102,7 +124,7 @@ func main() {
 		<-ctx.Done()
 		slog.Info("nomos: shutting down")
 		srv.Shutdown(context.Background())
-		client.close()
+		clientPool.closeAll()

 	default:
 		fmt.Fprintf(os.Stderr, "unknown command: %s\n", os.Args[1])
@@ -327,7 +349,7 @@ func handleAnswerQuestion(w http.ResponseWriter, r *http.Request, st *store, a *
 	w.WriteHeader(202)
 }

-func handleQuery(w http.ResponseWriter, r *http.Request, client *mcpClient, agentSlug, mcpURL string) {
+func handleQuery(w http.ResponseWriter, r *http.Request, pool *mcpClientPool, agentSlug, mcpURL string) {
 	if r.Method != http.MethodPost {
 		http.Error(w, "method not allowed", 405)
 		return
@@ -343,6 +365,16 @@ func handleQuery(w http.ResponseWriter, r *http.Request, client *mcpClient, agen
 		return
 	}

+	// The structured /query endpoint is stateless/session-less — "query" is a
+	// fixed pool key (not a real session id) so repeated calls reuse one
+	// dedicated connection instead of paying a fresh MCP handshake every time,
+	// while still never sharing a connection with an actual chat task.
+	client, err := pool.get("query")
+	if err != nil {
+		http.Error(w, "mcp unavailable: "+err.Error(), 502)
+		return
+	}
+
 	start := time.Now()

 	if req.Tool != "" {
@@ -416,7 +448,7 @@ type mcpClient struct {
 	sessionID string
 	http      *http.Client
 	nextID    int
-	mu        sync.Mutex // MCP is one stateful session; serialize concurrent calls
+	mu        sync.Mutex // this client is one stateful MCP session; serialize ITS OWN calls
 }

 func newMCPClient(baseURL string) (*mcpClient, error) {
@@ -616,3 +648,106 @@ func (c *mcpClient) listTools() ([]string, error) {

 func (c *mcpClient) close() {
 }
+
+// ─── Per-session MCP client pool ────────────────────────────────────────
+//
+// A single shared mcpClient serializes EVERY tool call across EVERY
+// concurrently-running task through one mutex (see mcpClient.mu) — `run`
+// executes its SSH command synchronously inside that lock and is capped at
+// up to 10 minutes, so one task mid-`run` stalled every other task's tool
+// calls, even trivial reads, behind it. The MCP *server* has no per-
+// connection state to protect (newServer in internal/mcp/server.go returns
+// one shared *mcp.Server instance whose tool handlers close only over the DB
+// pool, which is already safe for concurrent use) — the mutex existed purely
+// because the *client* reused one stateful transport session, not because
+// the server needed it. Giving each task's own session its own client
+// removes the cross-task serialization entirely: a task's own tool calls
+// stay sequential (which they already are — the agent loop calls tools one
+// at a time within a turn), but no longer block anyone else's.
+type mcpClientPool struct {
+	baseURL string
+	mu      sync.Mutex
+	clients map[string]*pooledMCPClient
+}
+
+type pooledMCPClient struct {
+	client   *mcpClient
+	lastUsed time.Time
+}
+
+func newMCPClientPool(baseURL string) *mcpClientPool {
+	return &mcpClientPool{baseURL: baseURL, clients: make(map[string]*pooledMCPClient)}
+}
+
+// get returns the client for sessionID, creating and initializing one (a
+// real MCP handshake) on first use. Session ids that don't identify a real
+// persisted conversation ("" / "ephemeral", the no-DB-store path; "query",
+// the structured /query endpoint) still get exactly one dedicated,
+// reused client each via the same map — just keyed on a fixed string instead
+// of a real session id — so that traffic doesn't pay a fresh handshake per
+// request while still never sharing a connection with an actual task.
+func (p *mcpClientPool) get(sessionID string) (*mcpClient, error) {
+	key := sessionID
+	if key == "" {
+		key = "ephemeral"
+	}
+
+	p.mu.Lock()
+	if pc, ok := p.clients[key]; ok {
+		pc.lastUsed = time.Now()
+		p.mu.Unlock()
+		return pc.client, nil
+	}
+	p.mu.Unlock()
+
+	// Initialize outside the lock — it's a network round-trip, and holding
+	// the pool mutex for it would serialize unrelated sessions' first calls
+	// behind each other, undermining the whole point of this pool.
+	c, err := newMCPClient(p.baseURL)
+	if err != nil {
+		return nil, err
+	}
+
+	p.mu.Lock()
+	// Another goroutine may have created one for the same key while we were
+	// initializing (two of this session's tool calls racing on a cold
+	// start); keep whichever won, close out the loser's connection (a no-op
+	// today, but future-proof if mcpClient.close ever does real teardown).
+	if existing, ok := p.clients[key]; ok {
+		p.mu.Unlock()
+		c.close()
+		return existing.client, nil
+	}
+	p.clients[key] = &pooledMCPClient{client: c, lastUsed: time.Now()}
+	p.mu.Unlock()
+	return c, nil
+}
+
+// mcpClientIdleTimeout is how long an idle session's MCP client is kept
+// before eviction — long enough to outlive a single slow `run` (capped at 10
+// minutes server-side) plus normal think-time between a task's tool calls,
+// short enough not to accumulate one abandoned connection per finished task
+// forever.
+const mcpClientIdleTimeout = 20 * time.Minute
+
+// sweep evicts clients idle past mcpClientIdleTimeout. Call on a ticker.
+func (p *mcpClientPool) sweep() {
+	cutoff := time.Now().Add(-mcpClientIdleTimeout)
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	for key, pc := range p.clients {
+		if pc.lastUsed.Before(cutoff) {
+			pc.client.close()
+			delete(p.clients, key)
+		}
+	}
+}
+
+func (p *mcpClientPool) closeAll() {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	for key, pc := range p.clients {
+		pc.client.close()
+		delete(p.clients, key)
+	}
+}