From a4ea542f3e23096d3102414f1d249410b71de9f7 Mon Sep 17 00:00:00 2001 From: dtoro Date: Sat, 11 Jul 2026 19:15:32 +0200 Subject: [PATCH] =?UTF-8?q?fix(concurrency):=20per-session=20MCP=20client?= =?UTF-8?q?=20pool=20=E2=80=94=20removes=20cross-task=20tool-call=20blocki?= =?UTF-8?q?ng?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- cmd/nomos/agent.go | 32 +++++++--- cmd/nomos/main.go | 149 ++++++++++++++++++++++++++++++++++++++++++--- 2 files changed, 164 insertions(+), 17 deletions(-) diff --git a/cmd/nomos/agent.go b/cmd/nomos/agent.go index 56f0c1c..cf1529f 100644 --- 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 } diff --git a/cmd/nomos/main.go b/cmd/nomos/main.go index 1bc5f66..ac0c7b3 100644 --- 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) + } +}