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..94e5ae1da --- /dev/null +++ b/internal/runtime/capture_lock_per_server_peer_test.go @@ -0,0 +1,248 @@ +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_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"}, + {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"}}, + } + overlayInflightPeers(rt.inflightPeerInventories("a"), peers) + + descs := map[string][]string{} + for _, tool := range peers["b"] { + descs[tool.RawName] = append(descs[tool.RawName], tool.Description) + } + // 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/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..b3d119e83 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. @@ -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)) @@ -2941,13 +2963,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..81acb2f28 100644 --- a/internal/runtime/runtime.go +++ b/internal/runtime/runtime.go @@ -309,10 +309,34 @@ 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 + // 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_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..955ab2b34 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,10 +647,26 @@ 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 } + // 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 @@ -637,12 +698,77 @@ func (r *Runtime) collectPeerToolMetadata(serverName string) map[string][]*confi peers[name] = metas } } + overlayInflightPeers(inflight, peers) if len(peers) == 0 { return nil } return peers } +// 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 + } + 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] { + 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). @@ -2690,7 +2816,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 +2833,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 {