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:
@@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user