refactor: consolidate duplicated tool/prompt logic into entityRegistry - #1148
refactor: consolidate duplicated tool/prompt logic into entityRegistry#1148aniketpandey05 wants to merge 4 commits into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughIntroduces a generic registry for upstream tools and prompts, rewires ChangesGeneric entityRegistry extraction and MCPManager wiring
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
coderabbitai
Bot
added
the
review-effort/large
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (3)
internal/broker/upstream/registry.go (1)
109-112: 💤 Low valueWasteful
toServercall for removals.
toServerbuilds a full server item (mutates name, allocates meta, creates handler closure) just to extract the prefixed name. Consider adding agetPrefixedName func(E) stringto avoid the allocation overhead on removals.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/broker/upstream/registry.go` around lines 109 - 112, The removal loop in the registry.go file calls `r.toServer(item)` just to extract the prefixed name via `r.getServerName()`, which is wasteful since `toServer` performs full server item construction (name mutation, meta allocation, handler closure creation). Add a new method `getPrefixedName(item E) string` that directly extracts the prefixed name from an item without the overhead of full conversion, then replace the `r.getServerName(r.toServer(item))` call in the removal loop with `r.getPrefixedName(item)` to eliminate unnecessary allocations.internal/broker/upstream/manager_test.go (1)
381-381: ⚡ Quick winConsider adding a test helper for serverItems access.
Tests directly manipulate and read the unexported
serverItemsfield. While legal within the same package, this couples tests to registry internals. AGetServerItemsForTesting()method onentityRegistrywould improve encapsulation and maintainability.Suggested test helper pattern
Add to
entityRegistryinregistry.go:+// GetServerItemsForTesting returns a copy of serverItems for test verification +func (r *entityRegistry[E, S]) GetServerItemsForTesting() []S { + return append([]S(nil), r.serverItems...) +}Then update tests:
-manager.tools.serverItems = make([]server.ServerTool, tc.numServerTools) +manager.tools.SetServerItemsForTesting(make([]server.ServerTool, tc.numServerTools)) -manager.tools.serverItems = []server.ServerTool{{Tool: mcp.Tool{Name: "existing_tool"}}} +manager.tools.SetServerItemsForTesting([]server.ServerTool{{Tool: mcp.Tool{Name: "existing_tool"}}}) -for i, st := range manager.tools.serverItems { +for i, st := range manager.tools.GetServerItemsForTesting() {Also applies to: 758-758, 962-963
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/broker/upstream/manager_test.go` at line 381, Tests directly access the unexported serverItems field on the manager, coupling the tests to internal implementation details. Add a GetServerItemsForTesting() method to the entityRegistry type in registry.go that returns the serverItems field for testing purposes. Then update all three locations where serverItems is accessed in manager_test.go (at line 381, line 758, and lines 962-963) to use this new test helper method instead of directly reading or manipulating the unexported field.internal/broker/upstream/manager.go (1)
400-414: ⚖️ Poor tradeoffRegistry state update could be encapsulated.
Direct manipulation of registry fields (
items,byName,byServedName,serverItems) breaks abstraction. The same pattern repeats for prompts (lines 455-467). Consider adding anapplyDifforupdatemethod toentityRegistrythat handles state transitions internally.This would reduce the duplication between the tools and prompts update blocks.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/broker/upstream/manager.go` around lines 400 - 414, The code directly manipulates entityRegistry fields (items, byName, byServedName, serverItems) within the manager's lock/unlock blocks, breaking encapsulation. Add an applyDiff or update method to the entityRegistry type that encapsulates all state transitions internally (recreating byName and byServedName maps, applying deletions with slices.DeleteFunc, and appending new items). Call this new method from the manager code instead of directly manipulating the registry fields. Apply this refactoring to both the tools update block and the prompts update block to eliminate the duplication between them.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@internal/broker/upstream/manager.go`:
- Around line 127-128: The comment for the toolsLock sync.RWMutex contains stale
references to serverTools and serverPrompts which were removed during
refactoring to registry-based fields. Update the comment to accurately reflect
which fields are now actually protected by this lock, removing the outdated
field names and replacing them with the current registry-based field names that
toolsLock guards.
In `@internal/broker/upstream/registry.go`:
- Line 1: The file has Go formatting violations detected by gofmt. Run the `make
fmt` command to automatically reformat the file according to Go style standards.
This will fix all formatting issues in the package upstream file.
- Around line 128-129: The loop iterating over gatewayItems does not validate
that existingPtr is non-nil before dereferencing it when calling
r.getMetaFields(*existingPtr). Add a nil check for existingPtr before
dereferencing it to prevent a panic if listFromGateway() returns map entries
with nil pointer values. Skip processing or handle the case appropriately when
existingPtr is nil within the for loop that iterates over existingName and
existingPtr from gatewayItems.
---
Nitpick comments:
In `@internal/broker/upstream/manager_test.go`:
- Line 381: Tests directly access the unexported serverItems field on the
manager, coupling the tests to internal implementation details. Add a
GetServerItemsForTesting() method to the entityRegistry type in registry.go that
returns the serverItems field for testing purposes. Then update all three
locations where serverItems is accessed in manager_test.go (at line 381, line
758, and lines 962-963) to use this new test helper method instead of directly
reading or manipulating the unexported field.
In `@internal/broker/upstream/manager.go`:
- Around line 400-414: The code directly manipulates entityRegistry fields
(items, byName, byServedName, serverItems) within the manager's lock/unlock
blocks, breaking encapsulation. Add an applyDiff or update method to the
entityRegistry type that encapsulates all state transitions internally
(recreating byName and byServedName maps, applying deletions with
slices.DeleteFunc, and appending new items). Call this new method from the
manager code instead of directly manipulating the registry fields. Apply this
refactoring to both the tools update block and the prompts update block to
eliminate the duplication between them.
In `@internal/broker/upstream/registry.go`:
- Around line 109-112: The removal loop in the registry.go file calls
`r.toServer(item)` just to extract the prefixed name via `r.getServerName()`,
which is wasteful since `toServer` performs full server item construction (name
mutation, meta allocation, handler closure creation). Add a new method
`getPrefixedName(item E) string` that directly extracts the prefixed name from
an item without the overhead of full conversion, then replace the
`r.getServerName(r.toServer(item))` call in the removal loop with
`r.getPrefixedName(item)` to eliminate unnecessary allocations.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: b004d583-9c82-4347-b1b7-17c97221aee6
📒 Files selected for processing (3)
internal/broker/upstream/manager.gointernal/broker/upstream/manager_test.gointernal/broker/upstream/registry.go
1d1240f to
57a83d6
Compare
|
@coderabbitai review |
✅ Action performedReview finished.
|
coderabbitai
Bot
added
the
high-risk
|
@coderabbitai review |
✅ Action performedReview finished.
|
|
@maleck13 please review this one , thank you |
|
@rogueslasher,Read this with my guard up... it's a refactor of the federation state, so all that matters is whether it preserves behaviour. Went looking for a dropped guard and couldn't find one: the generic One small thing: in Nice cleanup... the generics read well. |
05ac2d9 to
625eea4
Compare
|
@coderabbitai review |
✅ Action performedReview finished.
|
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
internal/broker/upstream/manager.go (2)
536-550: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winReset protocol validation before returning on errors.
Line 536 returns before lines 545-550 run, so failed validations can leave stale or empty
protocolValidationeven thoughExpectedVersionis always known. Move the protocol reset above the error branch.Proposed fix
man.status.InvalidPrompts = len(invalidPrompts) man.status.InvalidPromptList = invalidPrompts + man.status.ProtocolValidation = ProtocolValidation{ExpectedVersion: mcp.LATEST_PROTOCOL_VERSION} + if info := man.mcp.ProtocolInfo(); info != nil { + man.status.ProtocolValidation.IsValid = true + man.status.ProtocolValidation.SupportedVersion = info.ProtocolVersion + } if err != nil { man.status.Message = err.Error() man.status.Ready = false return } man.status.TotalTools = toolCount man.status.TotalPrompts = promptCount man.status.Ready = true man.status.Message = fmt.Sprintf("server added successfully. Total tools added %d. Total prompts added %d", toolCount, promptCount) - // always report the version we expect; fill in the negotiated version once it is known - man.status.ProtocolValidation = ProtocolValidation{ExpectedVersion: mcp.LATEST_PROTOCOL_VERSION} - if info := man.mcp.ProtocolInfo(); info != nil { - man.status.ProtocolValidation.IsValid = true - man.status.ProtocolValidation.SupportedVersion = info.ProtocolVersion - }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/broker/upstream/manager.go` around lines 536 - 550, The status update in manager.go leaves ProtocolValidation unset when the early error return is taken in the upstream add path. Update the logic around the add-server flow in the manager status handling so man.status.ProtocolValidation is reset before the err != nil branch returns, using mcp.LATEST_PROTOCOL_VERSION as the expected version. Keep the existing success-path negotiation logic that reads man.mcp.ProtocolInfo(), but ensure the reset happens unconditionally before any return from this block.
573-744: 🎯 Functional Correctness | 🔴 Critical | ⚡ Quick winResolve the merge conflict in
internal/broker/upstream/manager.go.<<<<<<<,=======, and>>>>>>>are still present, and the stale legacy tool/prompt helpers duplicate the registry-backed methods below. Drop the conflicted block and keep the registry-backed implementation so the file compiles.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/broker/upstream/manager.go` around lines 573 - 744, Resolve the merge conflict in MCPManager by removing all conflict markers and the stale legacy helper block. Keep the registry-backed implementations of GetManagedTools, GetServedManagedTool, and SetToolsForTesting that delegate to man.tools, and discard the duplicated findToolConflicts/getPrompts/findPromptConflicts versions that were superseded. Ensure the file retains only one consistent version of the tool/prompt management logic so manager.go compiles cleanly.Source: Linters/SAST tools
🧹 Nitpick comments (1)
internal/broker/upstream/manager_test.go (1)
989-990: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winAvoid copying
server.ServerToolin this loop.Line 990 ranges over
manager.tools.serverItemsby value, which copies each struct. The repo guideline prefers indexing here:proposed change
- for i, st := range manager.tools.serverItems { - serverToolNames[i] = st.Tool.Name + for i := range manager.tools.serverItems { + serverToolNames[i] = manager.tools.serverItems[i].Tool.Name }As per coding guidelines, "Use
for i := rangenotfor _, v := rangeon large structs".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/broker/upstream/manager_test.go` around lines 989 - 990, The loop over manager.tools.serverItems is copying each server.ServerTool value; update the iteration to use indexing instead of ranging by value so the struct is not copied. Keep the serverToolNames collection logic the same, but iterate with an index-based loop over manager.tools.serverItems and read each element through the slice index.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@internal/broker/upstream/manager.go`:
- Around line 536-550: The status update in manager.go leaves ProtocolValidation
unset when the early error return is taken in the upstream add path. Update the
logic around the add-server flow in the manager status handling so
man.status.ProtocolValidation is reset before the err != nil branch returns,
using mcp.LATEST_PROTOCOL_VERSION as the expected version. Keep the existing
success-path negotiation logic that reads man.mcp.ProtocolInfo(), but ensure the
reset happens unconditionally before any return from this block.
- Around line 573-744: Resolve the merge conflict in MCPManager by removing all
conflict markers and the stale legacy helper block. Keep the registry-backed
implementations of GetManagedTools, GetServedManagedTool, and SetToolsForTesting
that delegate to man.tools, and discard the duplicated
findToolConflicts/getPrompts/findPromptConflicts versions that were superseded.
Ensure the file retains only one consistent version of the tool/prompt
management logic so manager.go compiles cleanly.
---
Nitpick comments:
In `@internal/broker/upstream/manager_test.go`:
- Around line 989-990: The loop over manager.tools.serverItems is copying each
server.ServerTool value; update the iteration to use indexing instead of ranging
by value so the struct is not copied. Keep the serverToolNames collection logic
the same, but iterate with an index-based loop over manager.tools.serverItems
and read each element through the slice index.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 3f73e321-ba1c-49b9-a104-f94b12ad8848
📒 Files selected for processing (3)
internal/broker/upstream/manager.gointernal/broker/upstream/manager_test.gointernal/broker/upstream/registry.go
🚧 Files skipped from review as they are similar to previous changes (1)
- internal/broker/upstream/registry.go
|
@coderabbitai review |
✅ Action performedReview finished.
|
|
@coderabbitai review |
✅ Action performedReview finished.
|
|
@coderabbitai review |
✅ Action performedReview finished.
|
b52a1b3 to
45eca42
Compare
|
@maleck13 please review this thank you |
Patryk-Stefanski
left a comment
Patryk-Stefanski
left a comment
There was a problem hiding this comment.
The DRY consolidation is justified and the entityRegistry abstraction is clean. Three issues need fixing before merge: a silent behaviour change in setStatus, log level escalation in findConflicts, and an unnecessary allocation in diff.
| man.status.ProtocolValidation = ProtocolValidation{ExpectedVersion: mcp.LATEST_PROTOCOL_VERSION} | ||
| if info := man.mcp.ProtocolInfo(); info != nil { | ||
| man.status.ProtocolValidation.IsValid = true | ||
| man.status.ProtocolValidation.SupportedVersion = info.ProtocolVersion |
There was a problem hiding this comment.
Blocking: moves ProtocolValidation to the error path — unreviewed behaviour change not in scope for a refactor.
On main, ProtocolValidation was only written on the success path (after the if err != nil { return } guard). This PR moves it before the guard, so callers receiving Ready: false (connection failure, ping failure, disabled server) now also get a populated ProtocolValidation field.
Please revert to the original position (success path only), or if populating it on the error path is intentional, extract it as a separate commit with a rationale and a test.
| toRemove = append(toRemove, r.getServerName(r.toServer(item))) | ||
| } | ||
| } | ||
| return toAdd, toRemove |
There was a problem hiding this comment.
Blocking: r.toServer(item) is called here just to extract the prefixed name — unnecessary allocation on every reconcile.
r.toServer for a tool constructs a full server.ServerTool including metadata and a handler closure, when all that's needed is the served name. On main, diffTools used prefixedName(man.mcp.GetPrefix(), oldTool.Name) directly.
Consider storing prefix on the registry (already threaded through newToolRegistry/newPromptRegistry) and using it directly:
for name := range oldMap {
if _, exists := newMap[name]; !exists {
toRemove = append(toRemove, prefixedName(r.prefix, name))
}
}| r.logger.Error("unable to check conflict, meta is nil", | ||
| "upstream mcp server", r.serverID, r.entityTag, existingName) | ||
| continue | ||
| } |
There was a problem hiding this comment.
Blocking: log level escalated from Debug to Error for expected skip conditions.
On main, nil meta and missing ID were logged at Debug with "skipping conflict check" messages. These are expected for tools not managed by this gateway and will fire on every reconcile tick. Logging them as Error will produce spurious noise in operator logs.
Please revert the three skip/continue cases (nil gateway entry, meta is nil, id is missing) back to Debug. The actual conflict detection at the bottom of the loop is correctly at Debug and could reasonably stay or be promoted separately.
| func (r *entityRegistry[E, S]) toServerItem(item E) S { | ||
| return r.toServer(item) | ||
| } | ||
|
|
There was a problem hiding this comment.
Nit: toServerItem is a one-liner wrapper over r.toServer with no added value.
This only exists so tests can call it by a stable exported name. Either call r.toServer(item) directly in the tests, or remove this method.
Signed-off-by: rogueslasher <aniketpandey25092005@gmail.com>
Signed-off-by: rogueslasher <aniketpandey25092005@gmail.com>
Signed-off-by: rogueslasher <aniketpandey25092005@gmail.com>
Signed-off-by: Aniket Pandey <aniketpandey25092005@gmail.com>
45eca42 to
6ed288f
Compare
4 participants
Refactors the upstream manager by extracting a generic
entityRegistryto eliminate duplicated tool and prompt management logic. Previously, tools and prompts had separate implementations for several operations despite sharing nearly identical behavior. This change consolidates that logic into a single reusable registry abstraction while preserving the existing public API and behavior.Changes
internal/broker/upstream/registry.gocontaining the new genericentityRegistry.MCPManagerto use separate tool and prompt registries backed by the shared implementation.closes #1103
Summary by CodeRabbit