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>
This commit is contained in:
2026-07-11 19:15:32 +02:00
parent 6a8fb435ad
commit a4ea542f3e
2 changed files with 164 additions and 17 deletions

View File

@@ -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
}