From 2b9fed741a07e3e3eaa7cdaa37183f76cd0c0e9d Mon Sep 17 00:00:00 2001 From: Kyrill <1942093+poolski@users.noreply.github.com> Date: Sun, 1 Nov 2026 00:37:33 +0000 Subject: [PATCH 1/3] perf(runtime): make the capture lock per server captureSnapMu makes a server's "approval records persisted + pruned + capture stamp published" one unit for GetServerReview. It was a single runtime-wide mutex, held for the whole of a discovery pass's apply, including the TPA scans, approval record writes and index writes. A server with a large or changing tool set therefore blocked discovery commits of every other server and the review API of every server for as long as its pass ran. Every holder and reader works on one server, so the mutex is replaced by one RWMutex per server (captureSnapLock), stored like the per-server approval locks. No path held captureSnapMu for one server while taking it for another (that would have self-deadlocked), so splitting it adds no lock-order edges. --- .../runtime/capture_lock_per_server_test.go | 57 +++++++++++++++++++ internal/runtime/commit_testhooks.go | 2 +- internal/runtime/inventory_order.go | 2 +- internal/runtime/lifecycle.go | 22 +++---- internal/runtime/review.go | 8 +-- internal/runtime/runtime.go | 10 ++-- internal/runtime/tool_approval_lock.go | 7 +++ internal/runtime/tool_quarantine.go | 8 +-- 8 files changed, 91 insertions(+), 25 deletions(-) create mode 100644 internal/runtime/capture_lock_per_server_test.go diff --git a/internal/runtime/capture_lock_per_server_test.go b/internal/runtime/capture_lock_per_server_test.go new file mode 100644 index 000000000..478245ae9 --- /dev/null +++ b/internal/runtime/capture_lock_per_server_test.go @@ -0,0 +1,57 @@ +package runtime + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/smart-mcp-proxy/mcpproxy-go/internal/config" +) + +// A capture of one server must not block the review or the capture of +// another: a server with a large tool set otherwise stalls every other +// server's discovery and the review API for as long as its pass runs. +func TestCaptureOfOneServerDoesNotBlockAnother(t *testing.T) { + rt := setupQuarantineRuntime(t, nil, []*config.ServerConfig{ + {Name: "srv", Enabled: true, Quarantined: true}, + {Name: "other", Enabled: true, Quarantined: true}, + }) + otherTool := func(name string) *config.ToolMetadata { + return &config.ToolMetadata{ServerName: "other", Name: "other:" + name, Description: name} + } + require.NoError(t, rt.persistQuarantinedToolDefinitions("other", []*config.ToolMetadata{otherTool("a")})) + + paused, release := make(chan struct{}), make(chan struct{}) + rt.captureAfterPrune = func() { + rt.captureAfterPrune = nil + close(paused) + <-release + } + captureDone := make(chan error, 1) + go func() { + captureDone <- rt.persistQuarantinedToolDefinitions("srv", []*config.ToolMetadata{snapTool("a")}) + }() + <-paused + defer func() { + close(release) + require.NoError(t, <-captureDone) + }() + + otherDone := make(chan error, 1) + go func() { + if err := rt.persistQuarantinedToolDefinitions("other", []*config.ToolMetadata{otherTool("b")}); err != nil { + otherDone <- err + return + } + _, err := rt.GetServerReview(context.Background(), "other") + otherDone <- err + }() + select { + case err := <-otherDone: + require.NoError(t, err) + case <-time.After(2 * time.Second): + t.Fatal("a capture of srv blocked the capture and review of other") + } +} diff --git a/internal/runtime/commit_testhooks.go b/internal/runtime/commit_testhooks.go index 19cf88765..32dd860b5 100644 --- a/internal/runtime/commit_testhooks.go +++ b/internal/runtime/commit_testhooks.go @@ -51,7 +51,7 @@ type commitTestHooks struct { // inventoryBeforeApplyLock runs in every inventory write path (sweep, // single-server discovery, quarantined review capture, lenient empty // observation) after the tools/list returned and before the apply takes - // captureSnapMu and the server's approval lock, with the capture's + // captureSnapLock and the server's approval lock, with the capture's // sequence (Spec 117 PR-3). inventoryBeforeApplyLock func(serverName string, seq uint64) diff --git a/internal/runtime/inventory_order.go b/internal/runtime/inventory_order.go index 2d89e4a7f..2bd6d4716 100644 --- a/internal/runtime/inventory_order.go +++ b/internal/runtime/inventory_order.go @@ -36,7 +36,7 @@ import ( // (FR-037); otherwise the approval reindex re-lists and the sweep fallback // skips the server. // -// Lock order is unchanged: captureSnapMu → approval lock → storage / index / +// Lock order is unchanged: captureSnapLock → approval lock → storage / index / // supervisor state. The currency check reads two atomics of the client // (CaptureIsCurrent) and takes no manager or client lock. diff --git a/internal/runtime/lifecycle.go b/internal/runtime/lifecycle.go index 9fe03d436..4edee9e33 100644 --- a/internal/runtime/lifecycle.go +++ b/internal/runtime/lifecycle.go @@ -423,7 +423,7 @@ func (r *Runtime) discoverAndIndexTools(ctx context.Context, dueOnly bool) error // Prune entries for servers removed from config. Fresh snapshots are NOT // published here: each server's snapshot, capture stamp and approval - // records are committed together under captureSnapMu by + // records are committed together under captureSnapLock by // commitDiscoveredTools, so a concurrent review never reads one // generation's records with another's stamp. r.lastGoodToolsMu.Lock() @@ -1039,13 +1039,13 @@ func (r *Runtime) persistQuarantinedToolDefinitionsAt(serverName string, tools [ // applied for the server, or from a retired or reconnected client, persists // nothing. func (r *Runtime) persistInventoryDefinitions(serverName string, tools []*config.ToolMetadata, instanceGen int64, inv managed.InventoryCapture) (persisted bool, err error) { - r.captureSnapMu.Lock() - defer r.captureSnapMu.Unlock() + r.captureSnapLock(serverName).Lock() + defer r.captureSnapLock(serverName).Unlock() return r.persistToolDefinitionsLocked(serverName, tools, instanceGen, inv) } // persistToolDefinitionsLocked records, prunes and publishes one review -// capture. Caller holds captureSnapMu for writing. +// capture. Caller holds captureSnapLock for writing. // // Spec 116 FR-008a (review R2.2): the capture's final generation check and // all of its writes (review records, prune, last-good snapshot, capture @@ -1121,7 +1121,7 @@ func (r *Runtime) noteCaptureLocked(serverName string) { } // commitDiscoveredTools publishes a discovery result and updates the stored -// approval records under captureSnapMu, so GetServerReview never observes the +// approval records under captureSnapLock, so GetServerReview never observes the // new capture stamp with the previous generation's records (or vice versa). // For an eligible server the last-good snapshot and capture stamp are written, // together with the differential index update, in ONE hold of the server's @@ -1156,8 +1156,8 @@ func (r *Runtime) commitDiscoveredToolsAt(ctx context.Context, serverName string // re-lists the current client. func (r *Runtime) commitInventoryAt(ctx context.Context, serverName string, tools []*config.ToolMetadata, instanceGen int64, capture *supervisor.DiscoveryCapture, inv inventoryRef) (eligible, published bool, err error) { r.inventoryBeforeApply(serverName, inv.capture) - r.captureSnapMu.Lock() - defer r.captureSnapMu.Unlock() + r.captureSnapLock(serverName).Lock() + defer r.captureSnapLock(serverName).Unlock() if !r.serverEligibleForIndexing(serverName) { if _, perr := r.persistToolDefinitionsLocked(serverName, tools, instanceGen, inv.capture); perr != nil { r.logger.Warn("Failed to capture review definitions for ineligible server", @@ -1395,7 +1395,7 @@ func (r *Runtime) discoverAndIndexToolsForServerOnce(ctx context.Context, server } // Publish the last-good snapshot, capture stamp and approval records as - // one generation (captureSnapMu), then apply the differential index + // one generation (captureSnapLock), then apply the differential index // update. The snapshot gives the approval-driven reindex (issue #873) a // source without a network round-trip; on an authoritative empty refresh // it stores an empty slice so no stale set lingers. @@ -2941,13 +2941,13 @@ func configsEquivalent(a, b *config.Config) bool { // reapplyLastGoodSnapshot re-applies a server's last-good snapshot to the index // (the sweep fallback for a connected server missing from discovery). The -// snapshot is selected and applied under one captureSnapMu hold: a snapshot +// snapshot is selected and applied under one captureSnapLock hold: a snapshot // read outside the lock could be superseded by a concurrent capture that marks // a tool changed, and applying the stale one afterwards would look like a // genuine revert and clear that hold. func (r *Runtime) reapplyLastGoodSnapshot(ctx context.Context, serverName string) { - r.captureSnapMu.Lock() - defer r.captureSnapMu.Unlock() + r.captureSnapLock(serverName).Lock() + defer r.captureSnapLock(serverName).Unlock() snapshot, inv, seq, snapshotGen := r.lastGoodSnapshotWithCaptureGen(serverName) if len(snapshot) == 0 { diff --git a/internal/runtime/review.go b/internal/runtime/review.go index 97cfbac26..aa3601526 100644 --- a/internal/runtime/review.go +++ b/internal/runtime/review.go @@ -300,12 +300,12 @@ func (r *Runtime) GetServerReview(ctx context.Context, serverName string) (*Serv } } // Stored records and the live capture state are published by different - // writers. A capture holds captureSnapMu for persist + prune + stamp, so + // writers. A capture holds captureSnapLock for persist + prune + stamp, so // reading both under the read lock yields one generation, never a mix. - r.captureSnapMu.RLock() + r.captureSnapLock(serverName).RLock() records, err := r.storageManager.ListToolApprovals(serverName) if err != nil { - r.captureSnapMu.RUnlock() + r.captureSnapLock(serverName).RUnlock() return nil, fmt.Errorf("list tool reviews for %q: %w", serverName, err) } if hook := r.reviewAfterRecords; hook != nil { @@ -321,7 +321,7 @@ func (r *Runtime) GetServerReview(ctx context.Context, serverName string) (*Serv liveNames[config.RawToolName(t)] = struct{}{} } r.lastGoodToolsMu.RUnlock() - r.captureSnapMu.RUnlock() + r.captureSnapLock(serverName).RUnlock() for _, rec := range records { if _, ok := liveNames[rec.ToolName]; ok { reviewServer.LiveToolCount++ diff --git a/internal/runtime/runtime.go b/internal/runtime/runtime.go index 3d6d0c9f2..48322f5f3 100644 --- a/internal/runtime/runtime.go +++ b/internal/runtime/runtime.go @@ -309,10 +309,12 @@ type Runtime struct { lastGoodCapture map[string]managed.InventoryCapture // pruneAfterSnapshot is a test hook run after the prune snapshot. pruneAfterSnapshot func(toolName string) - // captureSnapMu makes "records persisted + pruned + stamp published" one - // unit for readers: a capture holds the write lock for all of it, a review - // holds the read lock while it reads records and live state. - captureSnapMu sync.RWMutex + // captureSnapLocks holds one *sync.RWMutex per server (captureSnapLock). + // It makes "records persisted + pruned + stamp published" one unit for + // readers: a capture holds the server's write lock for all of it, a + // review holds its read lock while it reads records and live state. The + // lock is per server so one server's long pass never blocks another's. + captureSnapLocks sync.Map // reviewAfterRecords is a test hook run after GetServerReview lists records. reviewAfterRecords func() // captureAfterPrune is a test hook run after a capture persisted and diff --git a/internal/runtime/tool_approval_lock.go b/internal/runtime/tool_approval_lock.go index e5fe4a283..f2656e80b 100644 --- a/internal/runtime/tool_approval_lock.go +++ b/internal/runtime/tool_approval_lock.go @@ -88,6 +88,13 @@ func (r *Runtime) toolApprovalServerState(serverName string) *toolApprovalServer return v.(*toolApprovalServerLock) } +// captureSnapLock returns serverName's capture/review lock, creating it. Like +// the approval locks it is never removed: a re-added server reuses it. +func (r *Runtime) captureSnapLock(serverName string) *sync.RWMutex { + v, _ := r.captureSnapLocks.LoadOrStore(serverName, &sync.RWMutex{}) + return v.(*sync.RWMutex) +} + // lockToolApprovals acquires the per-server tool-approval lock and returns its // unlock function. func (r *Runtime) lockToolApprovals(serverName string) func() { diff --git a/internal/runtime/tool_quarantine.go b/internal/runtime/tool_quarantine.go index 173d225c1..dc146f793 100644 --- a/internal/runtime/tool_quarantine.go +++ b/internal/runtime/tool_quarantine.go @@ -2690,7 +2690,7 @@ func (r *Runtime) reindexServerToolsAfterApprovalChange(serverName string) { // Prefer the last-good snapshot so the reindex is a pure local diff with no // network round-trip. It is populated by the full sweep and by // DiscoverAndIndexToolsForServer. Selection and application happen under one - // captureSnapMu hold so a concurrent capture cannot be overwritten by a + // captureSnapLock hold so a concurrent capture cannot be overwritten by a // stale snapshot (which would look like a revert and clear a changed hold). if r.reindexFromLastGoodSnapshot(serverName) { return @@ -2707,14 +2707,14 @@ func (r *Runtime) reindexServerToolsAfterApprovalChange(serverName string) { } // reindexFromLastGoodSnapshot re-runs the differential update from the -// last-good snapshot under captureSnapMu. It reports false when no snapshot +// last-good snapshot under captureSnapLock. It reports false when no snapshot // exists yet, or (Spec 117 PR-3, FR-037) when the snapshot no longer // describes the live connection — its client was replaced, retired or // reconnected, or a newer capture with definitions was applied — so the caller // re-lists the server instead of re-applying older definitions. func (r *Runtime) reindexFromLastGoodSnapshot(serverName string) bool { - r.captureSnapMu.Lock() - defer r.captureSnapMu.Unlock() + r.captureSnapLock(serverName).Lock() + defer r.captureSnapLock(serverName).Unlock() snapshot, inv, seq, snapshotGen := r.lastGoodSnapshotWithCaptureGen(serverName) if len(snapshot) == 0 { From c2dce9a1e0a7bfd4c3ab60ee0cd15861620bd4aa Mon Sep 17 00:00:00 2001 From: Algis Dumbris Date: Mon, 2 Nov 2026 19:37:16 +0200 Subject: [PATCH 2/3] fix(runtime): keep peer scans and signature reconcile cross-server safe under per-server capture locks The runtime-wide capture lock serialized every server's discovery pass, so the TPA shadowing/clone check of a later pass always saw the earlier pass's tools in the StateView. With per-server locks two passes overlap and each publishes only at its end, so a same-name near-duplicate added on two servers at once scanned clean on both. A pass now registers its inventory as peer context before its approval scan and unregisters it when the apply returns (after its publish); collectPeerToolMetadata overlays other servers' in-flight inventories on the StateView snapshot, so the later-scanning pass holds the clone again. reconcileSignatureCache enumerates the index and then evicts every hash it did not see. A concurrent pass of another server could index and warm a hash in between and lose it until its next pass. sigReconcileMu now covers reconcile from enumeration through RetainHashes and every warm, so a warm lands either before the enumeration (its tool already indexed) or after the eviction. Co-Authored-By: Claude Fable 5.1 --- .../capture_lock_per_server_peer_test.go | 230 ++++++++++++++++++ internal/runtime/lifecycle.go | 22 ++ internal/runtime/runtime.go | 22 ++ internal/runtime/tool_quarantine.go | 80 ++++++ 4 files changed, 354 insertions(+) create mode 100644 internal/runtime/capture_lock_per_server_peer_test.go diff --git a/internal/runtime/capture_lock_per_server_peer_test.go b/internal/runtime/capture_lock_per_server_peer_test.go new file mode 100644 index 000000000..77e55886b --- /dev/null +++ b/internal/runtime/capture_lock_per_server_peer_test.go @@ -0,0 +1,230 @@ +package runtime + +import ( + "context" + "fmt" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/smart-mcp-proxy/mcpproxy-go/internal/config" + "github.com/smart-mcp-proxy/mcpproxy-go/internal/storage" +) + +// peerToolNames returns the raw tool names collectPeerToolMetadata reports +// for server. +func peerToolNames(peers map[string][]*config.ToolMetadata, server string) []string { + var names []string + for _, t := range peers[server] { + names = append(names, t.RawName) + } + return names +} + +// With per-server capture locks, two servers' discovery passes overlap. A pass +// publishes its tools to the StateView only at its end, so the peer context of +// another server's scan must also include in-flight inventories, or two +// near-duplicate tools added at once both scan clean. +func TestCollectPeerToolMetadata_IncludesOtherServersInflightInventory(t *testing.T) { + rt := setupQuarantineRuntime(t, nil, []*config.ServerConfig{scanServer("a"), scanServer("b")}) + + require.Nil(t, rt.collectPeerToolMetadata("a"), "no peer has published or in-flight tools yet") + + unregisterB := rt.registerInflightInventory("b", []*config.ToolMetadata{ + {ServerName: "b", Name: "sync_database", RawName: "sync_database", Description: "b desc", ParamsJSON: `{"type":"object"}`}, + }) + unregisterA := rt.registerInflightInventory("a", []*config.ToolMetadata{ + {ServerName: "a", Name: "own_tool", RawName: "own_tool", Description: "a desc"}, + }) + + peers := rt.collectPeerToolMetadata("a") + assert.Equal(t, []string{"sync_database"}, peerToolNames(peers, "b"), "b's in-flight inventory is peer context for a") + assert.Empty(t, peers["a"], "a's own in-flight inventory is never its own peer") + require.Len(t, peers["b"], 1) + got := peers["b"][0] + assert.Equal(t, "b", got.ServerName) + assert.Equal(t, "sync_database", got.Name) + assert.Equal(t, "b desc", got.Description) + assert.Equal(t, `{"type":"object"}`, got.ParamsJSON) + + peersOfB := rt.collectPeerToolMetadata("b") + assert.Equal(t, []string{"own_tool"}, peerToolNames(peersOfB, "a")) + assert.Empty(t, peersOfB["b"]) + + unregisterB() + unregisterB() // idempotent + peers = rt.collectPeerToolMetadata("a") + assert.Empty(t, peers["b"], "a finished pass is no longer in flight") + + unregisterA() + assert.Nil(t, rt.collectPeerToolMetadata("b")) +} + +// A stale unregister (from a pass that already finished) must not remove the +// registration of the server's next pass. +func TestRegisterInflightInventory_StaleUnregisterKeepsNewerPass(t *testing.T) { + rt := setupQuarantineRuntime(t, nil, []*config.ServerConfig{scanServer("a"), scanServer("b")}) + + first := rt.registerInflightInventory("b", []*config.ToolMetadata{{ServerName: "b", Name: "old"}}) + second := rt.registerInflightInventory("b", []*config.ToolMetadata{{ServerName: "b", Name: "new"}}) + first() + assert.Equal(t, []string{"new"}, peerToolNames(rt.collectPeerToolMetadata("a"), "b")) + second() + assert.Nil(t, rt.collectPeerToolMetadata("a")) +} + +// Two scan-mode servers gain a same-name, near-duplicate tool in overlapping +// discovery passes. Before captureSnapMu became per server the passes were +// serialized, so the second one's scan saw the first's published tool and +// held its clone. With per-server locks, server a's pass is paused after its +// approval check (before it publishes); server b's whole pass runs in that +// window and must still see a's new tool as a peer and hold its own clone. +func TestOverlappingDiscoveryPasses_SecondScanSeesInflightPeer(t *testing.T) { + rt := setupQuarantineRuntime(t, nil, []*config.ServerConfig{scanServer("a"), scanServer("b")}) + const shared = "sync_database" + aDescription := "Synchronize production database records to nightly encrypted cold storage backup bucket automatically without manual intervention" + bDescription := "Synchronize database records to production nightly encrypted cold storage backup bucket automatically without manual intervention" + schema := `{"type":"object"}` + + seedApprovedBaseline(t, rt, "a", "a_baseline", "Lists the a things", schema) + seedApprovedBaseline(t, rt, "b", "b_baseline", "Lists the b things", schema) + + paused, release := make(chan struct{}), make(chan struct{}) + var once sync.Once + rt.testHooks.discoveryBeforeIndex = func(name string) { + if name != "a" { + return + } + once.Do(func() { + close(paused) + <-release + }) + } + + ctx := context.Background() + aDone := make(chan error, 1) + go func() { + aDone <- rt.applyDifferentialToolUpdate(ctx, "a", []*config.ToolMetadata{ + {ServerName: "a", Name: "a_baseline", RawName: "a_baseline", Description: "Lists the a things", ParamsJSON: schema}, + {ServerName: "a", Name: shared, RawName: shared, Description: aDescription, ParamsJSON: schema}, + }) + }() + select { + case <-paused: + case <-time.After(10 * time.Second): + t.Fatal("a's pass never reached its index step") + } + released := false + defer func() { + if !released { + close(release) + <-aDone + } + }() + + bErr := make(chan error, 1) + go func() { + bErr <- rt.applyDifferentialToolUpdate(ctx, "b", []*config.ToolMetadata{ + {ServerName: "b", Name: "b_baseline", RawName: "b_baseline", Description: "Lists the b things", ParamsJSON: schema}, + {ServerName: "b", Name: shared, RawName: shared, Description: bDescription, ParamsJSON: schema}, + }) + }() + select { + case err := <-bErr: + require.NoError(t, err) + case <-time.After(10 * time.Second): + t.Fatal("b's pass blocked behind a's paused pass") + } + + rec, err := rt.storageManager.GetToolApproval("b", shared) + require.NoError(t, err) + assert.Equal(t, storage.ToolApprovalStatusPending, rec.Status, + "b's near-duplicate of a's in-flight tool must be held, as it was when the passes were serialized") + + released = true + close(release) + require.NoError(t, <-aDone) +} + +// reconcileSignatureCache enumerates the index and then evicts every cached +// hash it did not see. A concurrent pass that indexes and warms a new hash +// between those two steps must not lose it. +func TestReconcileSignatureCache_DoesNotEvictConcurrentlyWarmedHash(t *testing.T) { + rt := setupQuarantineRuntime(t, nil, nil) + require.NotNil(t, rt.sigCache) + require.NotNil(t, rt.indexManager) + + stop := make(chan struct{}) + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + for { + select { + case <-stop: + return + default: + rt.reconcileSignatureCache() + } + } + }() + + var warmed []*config.ToolMetadata + deadline := time.Now().Add(2 * time.Second) + for i := 0; i < 400 && time.Now().Before(deadline); i++ { + server := fmt.Sprintf("s%d", i%4) + tool := &config.ToolMetadata{ + ServerName: server, + Name: fmt.Sprintf("tool_%d", i), + RawName: fmt.Sprintf("tool_%d", i), + Description: fmt.Sprintf("tool number %d", i), + ParamsJSON: fmt.Sprintf(`{"type":"object","properties":{"p%d":{"type":"string"}}}`, i), + Hash: fmt.Sprintf("hash-%d", i), + } + require.NoError(t, rt.indexManager.BatchIndexTools([]*config.ToolMetadata{tool})) + rt.warmSignatureCache([]*config.ToolMetadata{tool}) + warmed = append(warmed, tool) + } + close(stop) + wg.Wait() + + var lost []string + for _, tool := range warmed { + if _, ok := rt.sigCache.Peek(tool.Hash); !ok { + lost = append(lost, tool.Hash) + } + } + assert.Empty(t, lost, "indexed and warmed hashes were evicted by a concurrent reconcile") +} + +// An in-flight tool replaces the peer's published tool of the same raw name; +// published tools the in-flight inventory does not carry stay (the union is +// what the scan would see with the peer's pass ordered before or after it). +func TestOverlayInflightPeers_ReplacesSameNameKeepsOthers(t *testing.T) { + rt := setupQuarantineRuntime(t, nil, []*config.ServerConfig{scanServer("a"), scanServer("b")}) + unregister := rt.registerInflightInventory("b", []*config.ToolMetadata{ + {ServerName: "b", Name: "changed", RawName: "changed", Description: "new"}, + {ServerName: "b", Name: "added", RawName: "added", Description: "added"}, + {ServerName: "other", Name: "foreign", RawName: "foreign"}, // not b's: dropped like the publish drops it + }) + defer unregister() + + peers := map[string][]*config.ToolMetadata{ + "b": { + {ServerName: "b", Name: "changed", RawName: "changed", Description: "old"}, + {ServerName: "b", Name: "kept", RawName: "kept", Description: "kept"}, + }, + "c": {{ServerName: "c", Name: "c_tool", RawName: "c_tool"}}, + } + rt.overlayInflightPeers("a", peers) + + byName := map[string]string{} + for _, tool := range peers["b"] { + byName[tool.RawName] = tool.Description + } + assert.Equal(t, map[string]string{"changed": "new", "added": "added", "kept": "kept"}, byName) + assert.Len(t, peers["c"], 1, "servers without an in-flight inventory are untouched") +} diff --git a/internal/runtime/lifecycle.go b/internal/runtime/lifecycle.go index 4edee9e33..b3d119e83 100644 --- a/internal/runtime/lifecycle.go +++ b/internal/runtime/lifecycle.go @@ -1748,6 +1748,14 @@ func (r *Runtime) applyInventoryGuarded(ctx context.Context, serverName string, return incomplete(err) } + // Make this inventory peer context for other servers' scans from before + // its own scan until this function returns, which is after complete() + // published it to the StateView (or after the pass failed). With + // per-server capture locks another server's pass can scan while this one + // runs; without the registration neither pass would see the other's new + // tools until both published (see Runtime.inflightInventories). + defer r.registerInflightInventory(serverName, newTools)() + // Check tool-level quarantine approvals before indexing. Any error // (errCaptureIncomplete: a record write failed; errBaselineMarker) leaves // the capture incomplete: nothing below runs and nothing is published. @@ -2076,10 +2084,15 @@ func (r *Runtime) applyInventoryGuarded(ctx context.Context, serverName string, // later compact retrieve_tools is a pure cache read — never a per-request // compile. Hashless tools are skipped: warming them would memoize distinct // schemas under one "" key (indexed tools always carry a hash). +// +// It holds sigReconcileMu, so a warm never lands inside a concurrent +// reconcile's enumerate-then-retain window (see reconcileSignatureCache). func (r *Runtime) warmSignatureCache(tools []*config.ToolMetadata) { if r.sigCache == nil { return } + r.sigReconcileMu.Lock() + defer r.sigReconcileMu.Unlock() for _, tool := range tools { if tool == nil || tool.Hash == "" { continue @@ -2094,10 +2107,19 @@ func (r *Runtime) warmSignatureCache(tools []*config.ToolMetadata) { // would otherwise accumulate for the life of the process. Called after index // rebuilds and differential updates; enumerates the SHARED index only — // per-profile indexes hold subsets of the same tools, hence the same hashes. +// +// sigReconcileMu is held from the enumeration through RetainHashes. Discovery +// passes of different servers run concurrently, and a pass indexes a tool +// before it warms its hash; holding the lock across both steps means a warm +// either completes before the enumeration starts (so its tool is already +// indexed and enumerated) or waits until RetainHashes is done, so a live hash +// is never evicted by a reconcile that enumerated too early. func (r *Runtime) reconcileSignatureCache() { if r.sigCache == nil || r.indexManager == nil { return } + r.sigReconcileMu.Lock() + defer r.sigReconcileMu.Unlock() serverNames, err := r.indexManager.GetAllIndexedServerNames() if err != nil { r.logger.Debug("signature cache reconcile skipped: cannot enumerate indexed servers", zap.Error(err)) diff --git a/internal/runtime/runtime.go b/internal/runtime/runtime.go index 48322f5f3..81acb2f28 100644 --- a/internal/runtime/runtime.go +++ b/internal/runtime/runtime.go @@ -315,6 +315,28 @@ type Runtime struct { // review holds its read lock while it reads records and live state. The // lock is per server so one server's long pass never blocks another's. captureSnapLocks sync.Map + // inflightInventories (guarded by inflightInventoryMu) holds, per server, + // the peer-context projection of the inventory a discovery pass is + // applying right now (registerInflightInventory). A pass publishes its + // tools to the StateView only at its end, and with per-server capture + // locks two servers' passes overlap; collectPeerToolMetadata overlays + // these entries so the TPA shadowing/clone check of one pass sees a peer's + // not-yet-published tools. Each pass registers before its approval scan + // and unregisters when applyInventoryGuarded returns (after its publish), + // so of two overlapping passes at least the later-scanning one sees the + // other's new tools: a same-name near-duplicate added on two servers at + // once is held on the second, as it was when every pass was serialized. + // inflightInventoryMu is a leaf lock: nothing else is taken under it. + inflightInventoryMu sync.Mutex + inflightInventories map[string]*inflightInventory + // sigReconcileMu serializes reconcileSignatureCache (index enumeration + // through RetainHashes) with warmSignatureCache. Without it a reconcile + // that enumerated the index before another server's pass indexed and + // warmed hash H would evict H afterwards, and nothing re-warms it until + // that server's next pass. Lock order: a server's approval lock (warm runs + // under it) -> sigReconcileMu -> the index manager's and cache's own + // locks; nothing taken under sigReconcileMu calls back into the runtime. + sigReconcileMu sync.Mutex // reviewAfterRecords is a test hook run after GetServerReview lists records. reviewAfterRecords func() // captureAfterPrune is a test hook run after a capture persisted and diff --git a/internal/runtime/tool_quarantine.go b/internal/runtime/tool_quarantine.go index dc146f793..4aecce424 100644 --- a/internal/runtime/tool_quarantine.go +++ b/internal/runtime/tool_quarantine.go @@ -595,6 +595,51 @@ func (r *Runtime) refreshHeldContractSnapshot(serverName, toolName string, exist return true } +// inflightInventory is one server's in-flight inventory as peer context: the +// tools its running discovery pass will publish, already projected the way +// collectPeerToolMetadata reads the StateView. The pointer identifies the +// registration, so a finished pass never unregisters a newer one. +type inflightInventory struct { + tools []*config.ToolMetadata +} + +// registerInflightInventory makes the inventory a discovery pass of +// serverName is applying visible as peer context to other servers' scans +// (Runtime.inflightInventories) and returns its idempotent unregister. The +// tools are copied into the peer projection here, by the pass's own +// goroutine, so concurrent readers never touch the pass's metadata. Like the +// StateView publish, only tools whose ServerName is serverName are kept. +func (r *Runtime) registerInflightInventory(serverName string, tools []*config.ToolMetadata) func() { + metas := make([]*config.ToolMetadata, 0, len(tools)) + for _, t := range tools { + if t == nil || t.ServerName != serverName { + continue + } + metas = append(metas, &config.ToolMetadata{ + ServerName: serverName, + Name: t.Name, + RawName: config.RawToolName(t), + Description: t.Description, + ParamsJSON: t.ParamsJSON, + OutputSchemaJSON: t.OutputSchemaJSON, + }) + } + entry := &inflightInventory{tools: metas} + r.inflightInventoryMu.Lock() + if r.inflightInventories == nil { + r.inflightInventories = make(map[string]*inflightInventory) + } + r.inflightInventories[serverName] = entry + r.inflightInventoryMu.Unlock() + return func() { + r.inflightInventoryMu.Lock() + if r.inflightInventories[serverName] == entry { + delete(r.inflightInventories, serverName) + } + r.inflightInventoryMu.Unlock() + } +} + // collectPeerToolMetadata returns every OTHER connected server's current tools, // keyed by server name, projected onto config.ToolMetadata for the synchronous // scan gate. It is the cross-server context that lets the peer-dependent @@ -602,6 +647,15 @@ func (r *Runtime) refreshHeldContractSnapshot(serverName, toolName string, exist // scan). Sourced from the lock-free StateView snapshot; returns nil when no // supervisor/snapshot is available (best effort — a genuinely single-server // deployment has no peers to shadow anyway). +// +// Discovery passes of different servers run concurrently (the capture lock is +// per server) and publish to the StateView only at their end, so the +// in-flight inventories of OTHER servers (Runtime.inflightInventories) are +// overlaid on the snapshot: a peer's in-flight tool replaces its published +// tool of the same raw name, and published tools the pass does not carry +// stay. The union is deliberately conservative — it is what this scan would +// have seen with the peer's pass ordered either before or after it — so a +// pass that later fails cannot make a scan miss a still-published tool. func (r *Runtime) collectPeerToolMetadata(serverName string) map[string][]*config.ToolMetadata { if r.supervisor == nil { return nil @@ -637,12 +691,38 @@ func (r *Runtime) collectPeerToolMetadata(serverName string) map[string][]*confi peers[name] = metas } } + r.overlayInflightPeers(serverName, peers) if len(peers) == 0 { return nil } return peers } +// overlayInflightPeers merges the in-flight inventories of every server but +// serverName into peers (see collectPeerToolMetadata): an in-flight tool +// replaces a published tool of the same raw name, other published tools stay. +func (r *Runtime) overlayInflightPeers(serverName string, peers map[string][]*config.ToolMetadata) { + r.inflightInventoryMu.Lock() + defer r.inflightInventoryMu.Unlock() + for name, entry := range r.inflightInventories { + if name == serverName || entry == nil || len(entry.tools) == 0 { + continue + } + inflight := make(map[string]struct{}, len(entry.tools)) + merged := make([]*config.ToolMetadata, 0, len(entry.tools)+len(peers[name])) + for _, t := range entry.tools { + inflight[t.RawName] = struct{}{} + merged = append(merged, t) + } + for _, t := range peers[name] { + if _, ok := inflight[config.RawToolName(t)]; !ok { + merged = append(merged, t) + } + } + peers[name] = merged + } +} + // ToolApprovalResult contains the result of checking tool approvals for a server. type ToolApprovalResult struct { // BlockedTools is the set of tool names that should not be indexed (pending, changed, or disabled). From ec5876b056fc5fd145ad0c69f872a7da2daa9793 Mon Sep 17 00:00:00 2001 From: Algis Dumbris Date: Mon, 2 Nov 2026 19:44:30 +0200 Subject: [PATCH 3/3] fix(runtime): read in-flight peers before the StateView snapshot; keep published contracts in the overlay collectPeerToolMetadata copied the in-flight registrations after taking the StateView snapshot, so a peer that published and unregistered between the two reads was seen by neither. The in-flight copy is taken first: a pass unregisters only after it published, so what the copy misses is in the snapshot taken next. The overlay replaced a published tool with an in-flight tool of the same raw name. An in-flight replacement is not yet published and its pass may fail, so the published contract is still what a clone would shadow. The overlay is now a union; only an identical published contract is dropped. Co-Authored-By: Claude Fable 5.1 --- .../capture_lock_per_server_peer_test.go | 28 ++++++-- internal/runtime/tool_quarantine.go | 70 +++++++++++++++---- 2 files changed, 81 insertions(+), 17 deletions(-) diff --git a/internal/runtime/capture_lock_per_server_peer_test.go b/internal/runtime/capture_lock_per_server_peer_test.go index 77e55886b..94e5ae1da 100644 --- a/internal/runtime/capture_lock_per_server_peer_test.go +++ b/internal/runtime/capture_lock_per_server_peer_test.go @@ -203,7 +203,7 @@ func TestReconcileSignatureCache_DoesNotEvictConcurrentlyWarmedHash(t *testing.T // An in-flight tool replaces the peer's published tool of the same raw name; // published tools the in-flight inventory does not carry stay (the union is // what the scan would see with the peer's pass ordered before or after it). -func TestOverlayInflightPeers_ReplacesSameNameKeepsOthers(t *testing.T) { +func TestOverlayInflightPeers_UnionKeepsPublishedContract(t *testing.T) { rt := setupQuarantineRuntime(t, nil, []*config.ServerConfig{scanServer("a"), scanServer("b")}) unregister := rt.registerInflightInventory("b", []*config.ToolMetadata{ {ServerName: "b", Name: "changed", RawName: "changed", Description: "new"}, @@ -219,12 +219,30 @@ func TestOverlayInflightPeers_ReplacesSameNameKeepsOthers(t *testing.T) { }, "c": {{ServerName: "c", Name: "c_tool", RawName: "c_tool"}}, } - rt.overlayInflightPeers("a", peers) + overlayInflightPeers(rt.inflightPeerInventories("a"), peers) - byName := map[string]string{} + descs := map[string][]string{} for _, tool := range peers["b"] { - byName[tool.RawName] = tool.Description + descs[tool.RawName] = append(descs[tool.RawName], tool.Description) } - assert.Equal(t, map[string]string{"changed": "new", "added": "added", "kept": "kept"}, byName) + // A replacement is not yet published and its pass may fail: BOTH contracts + // stay in the peer context, so a clone of either is seen. + assert.Equal(t, []string{"old", "new"}, descs["changed"], "published contract first: the detector reads a server's first entry per name") + assert.Equal(t, []string{"added"}, descs["added"]) + assert.Equal(t, []string{"kept"}, descs["kept"]) assert.Len(t, peers["c"], 1, "servers without an in-flight inventory are untouched") + assert.Empty(t, rt.inflightPeerInventories("b"), "a server never sees its own in-flight inventory") + + // An unchanged in-flight tool is listed once. + peers = map[string][]*config.ToolMetadata{ + "b": {{ServerName: "b", Name: "added", RawName: "added", Description: "added"}}, + } + overlayInflightPeers(rt.inflightPeerInventories("a"), peers) + n := 0 + for _, tool := range peers["b"] { + if tool.RawName == "added" { + n++ + } + } + assert.Equal(t, 1, n, "identical published and in-flight contracts are deduplicated") } diff --git a/internal/runtime/tool_quarantine.go b/internal/runtime/tool_quarantine.go index 4aecce424..955ab2b34 100644 --- a/internal/runtime/tool_quarantine.go +++ b/internal/runtime/tool_quarantine.go @@ -660,6 +660,13 @@ func (r *Runtime) collectPeerToolMetadata(serverName string) map[string][]*confi if r.supervisor == nil { return nil } + // The in-flight registrations are read BEFORE the StateView snapshot. A + // pass unregisters only after it published (or failed without publishing), + // so a peer missing from this copy because it already unregistered is + // either visible in the snapshot taken next or never became a tool. Read + // the other way round, a peer could publish and unregister between the two + // reads and be seen by neither. + inflight := r.inflightPeerInventories(serverName) snapshot := r.supervisor.StateView().Snapshot() if snapshot == nil { return nil @@ -691,38 +698,77 @@ func (r *Runtime) collectPeerToolMetadata(serverName string) map[string][]*confi peers[name] = metas } } - r.overlayInflightPeers(serverName, peers) + overlayInflightPeers(inflight, peers) if len(peers) == 0 { return nil } return peers } -// overlayInflightPeers merges the in-flight inventories of every server but -// serverName into peers (see collectPeerToolMetadata): an in-flight tool -// replaces a published tool of the same raw name, other published tools stay. -func (r *Runtime) overlayInflightPeers(serverName string, peers map[string][]*config.ToolMetadata) { +// inflightPeerInventories copies the in-flight inventories of every server but +// serverName (see collectPeerToolMetadata). The entries are immutable once +// registered, so only the map is copied under the lock. +func (r *Runtime) inflightPeerInventories(serverName string) map[string][]*config.ToolMetadata { r.inflightInventoryMu.Lock() defer r.inflightInventoryMu.Unlock() + if len(r.inflightInventories) == 0 { + return nil + } + out := make(map[string][]*config.ToolMetadata, len(r.inflightInventories)) for name, entry := range r.inflightInventories { if name == serverName || entry == nil || len(entry.tools) == 0 { continue } - inflight := make(map[string]struct{}, len(entry.tools)) - merged := make([]*config.ToolMetadata, 0, len(entry.tools)+len(peers[name])) - for _, t := range entry.tools { - inflight[t.RawName] = struct{}{} - merged = append(merged, t) + out[name] = entry.tools + } + return out +} + +// overlayInflightPeers merges in-flight peer inventories into peers (see +// collectPeerToolMetadata). It is a union: every published tool stays, +// including one with the same raw name as an in-flight tool whose contract +// differs (an in-flight replacement is not yet published and its pass may +// still fail; the published contract is still what a clone would shadow), and +// every in-flight tool is added after it. Published first: the detector's +// per-name checks read the first entry of a server (sameTokens), and the +// published contract is the one every pass before #1813's per-server locks +// would have scanned against. Only an in-flight tool whose contract equals the +// published one is skipped, so an unchanged tool is listed once. +// +// Known limitation (fails closed): a hold taken against a peer's in-flight +// tool stays if that peer's pass then fails before publishing; the operator +// reviews it like any other hold. +func overlayInflightPeers(inflight, peers map[string][]*config.ToolMetadata) { + for name, tools := range inflight { + if len(tools) == 0 { + continue } + published := make(map[string][]*config.ToolMetadata, len(peers[name])) for _, t := range peers[name] { - if _, ok := inflight[config.RawToolName(t)]; !ok { - merged = append(merged, t) + raw := config.RawToolName(t) + published[raw] = append(published[raw], t) + } + merged := make([]*config.ToolMetadata, 0, len(tools)+len(peers[name])) + merged = append(merged, peers[name]...) + next: + for _, t := range tools { + for _, p := range published[t.RawName] { + if samePeerContract(t, p) { + continue next + } } + merged = append(merged, t) } peers[name] = merged } } +// samePeerContract reports whether two peer projections carry the same +// contract as far as the peer-dependent scan checks read it. +func samePeerContract(a, b *config.ToolMetadata) bool { + return a.Description == b.Description && a.ParamsJSON == b.ParamsJSON && a.OutputSchemaJSON == b.OutputSchemaJSON +} + // ToolApprovalResult contains the result of checking tool approvals for a server. type ToolApprovalResult struct { // BlockedTools is the set of tool names that should not be indexed (pending, changed, or disabled).