Skip to content

feat: add OAuth 2.0 support for MCP clients with discovery and PKCE #1430

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Pratham-Mishra04 wants to merge 5 commits into from
Closed

feat: add OAuth 2.0 support for MCP clients with discovery and PKCE #1430

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
e3b6ca5
feat: add oauth support to mcp
Pratham-Mishra04 Jan 23, 2026
cec3f0d
mcp auth server setup
akshaydeo Jan 23, 2026
7318932
code mode compact signatures
akshaydeo Jan 26, 2026
12425fa
mcp flow refactoring
akshaydeo Jan 27, 2026
47f0f8f
mcp tool groups
akshaydeo Jan 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
202 changes: 113 additions & 89 deletions core/bifrost.go
View file
Open in desktop

Large diffs are not rendered by default.

8 changes: 2 additions & 6 deletions core/go.mod
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,14 +13,13 @@ require (
github.com/aws/aws-sdk-go-v2/service/s3 v1.94.0
github.com/aws/smithy-go v1.24.0
github.com/bytedance/sonic v1.14.2
github.com/clarkmcc/go-typescript v0.7.0
github.com/dop251/goja v0.0.0-20251103141225-af2ceb9156d7
github.com/google/uuid v1.6.0
github.com/hajimehoshi/go-mp3 v0.3.4
github.com/mark3labs/mcp-go v0.43.2
github.com/rs/zerolog v1.34.0
github.com/stretchr/testify v1.11.1
github.com/valyala/fasthttp v1.68.0
go.starlark.net v0.0.0-20260102030733-3fee463870c9
golang.org/x/oauth2 v0.34.0
golang.org/x/text v0.32.0
)
Expand All @@ -29,7 +28,6 @@ require (
cloud.google.com/go/compute/metadata v0.9.0 // indirect
github.com/Azure/azure-sdk-for-go/sdk/internal v1.11.2 // indirect
github.com/AzureAD/microsoft-authentication-library-for-go v1.6.0 // indirect
github.com/Masterminds/semver/v3 v3.3.1 // indirect
github.com/andybalholm/brotli v1.2.0 // indirect
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.16 // indirect
github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.16 // indirect
Expand All @@ -50,10 +48,7 @@ require (
github.com/bytedance/sonic/loader v0.4.0 // indirect
github.com/cloudwego/base64x v0.1.6 // indirect
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
github.com/dlclark/regexp2 v1.11.4 // indirect
github.com/go-sourcemap/sourcemap v2.1.3+incompatible // indirect
github.com/golang-jwt/jwt/v5 v5.3.0 // indirect
github.com/google/pprof v0.0.0-20251213031049-b05bdaca462f // indirect
github.com/invopop/jsonschema v0.13.0 // indirect
github.com/klauspost/compress v1.18.2 // indirect
github.com/klauspost/cpuid/v2 v2.3.0 // indirect
Expand All @@ -73,5 +68,6 @@ require (
golang.org/x/crypto v0.46.0 // indirect
golang.org/x/net v0.48.0 // indirect
golang.org/x/sys v0.39.0 // indirect
google.golang.org/protobuf v1.36.11 // indirect
gopkg.in/yaml.v3 v3.0.1 // indirect
)
18 changes: 4 additions & 14 deletions core/go.sum
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,6 @@ github.com/AzureAD/microsoft-authentication-extensions-for-go/cache v0.1.1 h1:WJ
github.com/AzureAD/microsoft-authentication-extensions-for-go/cache v0.1.1/go.mod h1:tCcJZ0uHAmvjsVYzEFivsRTN00oz5BEsRgQHu5JZ9WE=
github.com/AzureAD/microsoft-authentication-library-for-go v1.6.0 h1:XRzhVemXdgvJqCH0sFfrBUTnUJSBrBf7++ypk+twtRs=
github.com/AzureAD/microsoft-authentication-library-for-go v1.6.0/go.mod h1:HKpQxkWaGLJ+D/5H8QRpyQXA1eKjxkFlOMwck5+33Jk=
github.com/Masterminds/semver/v3 v3.3.1 h1:QtNSWtVZ3nBfk8mAOu/B6v7FMJ+NHTIgUPi7rj+4nv4=
github.com/Masterminds/semver/v3 v3.3.1/go.mod h1:4V+yj/TJE1HU9XfppCwVMZq3I84lprf4nC11bSS5beM=
github.com/andybalholm/brotli v1.2.0 h1:ukwgCxwYrmACq68yiUqwIWnGY0cTPox/M94sVwToPjQ=
github.com/andybalholm/brotli v1.2.0/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY=
github.com/aws/aws-sdk-go-v2 v1.41.0 h1:tNvqh1s+v0vFYdA1xq0aOJH+Y5cRyZ5upu6roPgPKd4=
Expand Down Expand Up @@ -66,30 +64,20 @@ github.com/bytedance/sonic v1.14.2 h1:k1twIoe97C1DtYUo+fZQy865IuHia4PR5RPiuGPPII
github.com/bytedance/sonic v1.14.2/go.mod h1:T80iDELeHiHKSc0C9tubFygiuXoGzrkjKzX2quAx980=
github.com/bytedance/sonic/loader v0.4.0 h1:olZ7lEqcxtZygCK9EKYKADnpQoYkRQxaeY2NYzevs+o=
github.com/bytedance/sonic/loader v0.4.0/go.mod h1:AR4NYCk5DdzZizZ5djGqQ92eEhCCcdf5x77udYiSJRo=
github.com/clarkmcc/go-typescript v0.7.0 h1:3nVeaPYyTCWjX6Lf8GoEOTxME2bM5tLuWmwhSZ86uxg=
github.com/clarkmcc/go-typescript v0.7.0/go.mod h1:IZ/nzoVeydAmyfX7l6Jmp8lJDOEnae3jffoXwP4UyYg=
github.com/cloudwego/base64x v0.1.6 h1:t11wG9AECkCDk5fMSoxmufanudBtJ+/HemLstXDLI2M=
github.com/cloudwego/base64x v0.1.6/go.mod h1:OFcloc187FXDaYHvrNIjxSe8ncn0OOM8gEHfghB2IPU=
github.com/coreos/go-systemd/v22 v22.5.0/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/dlclark/regexp2 v1.11.4 h1:rPYF9/LECdNymJufQKmri9gV604RvvABwgOA8un7yAo=
github.com/dlclark/regexp2 v1.11.4/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
github.com/dop251/goja v0.0.0-20251103141225-af2ceb9156d7 h1:jxmXU5V9tXxJnydU5v/m9SG8TRUa/Z7IXODBpMs/P+U=
github.com/dop251/goja v0.0.0-20251103141225-af2ceb9156d7/go.mod h1:MxLav0peU43GgvwVgNbLAj1s/bSGboKkhuULvq/7hx4=
github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8=
github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0=
github.com/go-sourcemap/sourcemap v2.1.3+incompatible h1:W1iEw64niKVGogNgBN3ePyLFfuisuzeidWPMPWmECqU=
github.com/go-sourcemap/sourcemap v2.1.3+incompatible/go.mod h1:F8jJfvm2KbVjc5NqelyYJmf/v5J0dwNLS2mL4sNA1Jg=
github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
github.com/golang-jwt/jwt/v5 v5.3.0 h1:pv4AsKCKKZuqlgs5sUmn4x8UlGa0kEVt/puTpKx9vvo=
github.com/golang-jwt/jwt/v5 v5.3.0/go.mod h1:fxCRLWMO43lRc8nhHWY6LGqRcf+1gQWArsqaEUEa5bE=
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
github.com/google/pprof v0.0.0-20251213031049-b05bdaca462f h1:HU1RgM6NALf/KW9HEY6zry3ADbDKcmpQ+hJedoNGQYQ=
github.com/google/pprof v0.0.0-20251213031049-b05bdaca462f/go.mod h1:67FPmZWbr+KDT/VlpWtw6sO9XSjpJmLuHpoLmWiTGgY=
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/hajimehoshi/go-mp3 v0.3.4 h1:NUP7pBYH8OguP4diaTZ9wJbUbk3tC0KlfzsEpWmYj68=
Expand Down Expand Up @@ -155,6 +143,8 @@ github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZ
github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E=
github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=
github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4=
go.starlark.net v0.0.0-20260102030733-3fee463870c9 h1:nV1OyvU+0CYrp5eKfQ3rD03TpFYYhH08z31NK1HmtTk=
go.starlark.net v0.0.0-20260102030733-3fee463870c9/go.mod h1:YKMCv9b1WrfWmeqdV5MAuEHWsu5iC+fe6kYl2sQjdI8=
golang.org/x/arch v0.23.0 h1:lKF64A2jF6Zd8L0knGltUnegD62JMFBiCPBmQpToHhg=
golang.org/x/arch v0.23.0/go.mod h1:dNHoOeKiyja7GTvF9NJS1l3Z2yntpQNzgrjh1cU103A=
golang.org/x/crypto v0.46.0 h1:cKRW/pmt1pKAfetfu+RCEvjvZkA9RimPbh7bhFjGVBU=
Expand All @@ -172,11 +162,11 @@ golang.org/x/sys v0.39.0 h1:CvCKL8MeisomCi6qNZ+wbb0DN9E5AATixKsvNtMoMFk=
golang.org/x/sys v0.39.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
golang.org/x/text v0.32.0 h1:ZD01bjUt1FQ9WJ0ClOL5vxgxOI/sVCNgX1YtKwcY0mU=
golang.org/x/text v0.32.0/go.mod h1:o/rUWzghvpD5TXrTIBuJU77MTaN0ljMWE47kxGJQ7jY=
google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE=
google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
4 changes: 2 additions & 2 deletions core/internal/mcptests/agent_filtering_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -531,7 +531,7 @@ func TestAgent_FilteringWithMultipleClients(t *testing.T) {
tempConfig := GetTemperatureMCPClientConfig("")
tempConfig.ToolsToAutoExecute = []string{} // Not auto-executed

err = manager.AddClient(tempConfig)
err = manager.AddClient(&tempConfig)
if err != nil {
t.Skipf("Skipping test - temperature server not available: %v", err)
return
Expand Down Expand Up @@ -620,7 +620,7 @@ func TestAgent_ToolConflictInAgentMode(t *testing.T) {
tempConfig := GetTemperatureMCPClientConfig("")
tempConfig.ToolsToAutoExecute = []string{} // Not auto

err = manager.AddClient(tempConfig)
err = manager.AddClient(&tempConfig)
if err != nil {
t.Skipf("Skipping test - temperature server not available: %v", err)
return
Expand Down
20 changes: 16 additions & 4 deletions core/internal/mcptests/agent_request_id_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"testing"

"github.com/maximhq/bifrost/core/mcp"
"github.com/maximhq/bifrost/core/mcp/codemode/starlark"
"github.com/maximhq/bifrost/core/schemas"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
Expand All @@ -20,15 +21,26 @@ import (
func setupMCPManagerWithRequestIDFunc(t *testing.T, fetchNewRequestIDFunc func(ctx *schemas.BifrostContext) string, clientConfigs ...schemas.MCPClientConfig) *mcp.MCPManager {
t.Helper()

logger := &testLogger{t: t}

// Convert to pointer slice for MCPConfig
clientConfigPtrs := make([]*schemas.MCPClientConfig, len(clientConfigs))
for i := range clientConfigs {
clientConfigPtrs[i] = &clientConfigs[i]
}

// Create MCP config with request ID function
mcpConfig := &schemas.MCPConfig{
ClientConfigs: clientConfigs,
ClientConfigs: clientConfigPtrs,
FetchNewRequestIDFunc: fetchNewRequestIDFunc,
}

// Create MCP manager
logger := &testLogger{t: t}
manager := mcp.NewMCPManager(context.Background(), *mcpConfig, logger)
// Create Starlark CodeMode
starlark.SetLogger(logger)
codeMode := starlark.NewStarlarkCodeMode(nil)

// Create MCP manager - dependencies are injected automatically
manager := mcp.NewMCPManager(context.Background(), *mcpConfig, nil, logger, codeMode)

// Cleanup
t.Cleanup(func() {
Expand Down
14 changes: 7 additions & 7 deletions core/internal/mcptests/client_management_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,11 +25,11 @@ func TestAddClientDuplicate(t *testing.T) {

// Add client
clientConfig := GetSampleHTTPClientConfig(config.HTTPServerURL)
err := manager.AddClient(clientConfig)
err := manager.AddClient(&clientConfig)
require.NoError(t, err, "should add client first time")

// Try to add same client again
err = manager.AddClient(clientConfig)
err = manager.AddClient(&clientConfig)
// Should either return error or be idempotent
if err == nil {
clients := manager.GetClients()
Expand Down Expand Up @@ -183,7 +183,7 @@ func TestEditClient(t *testing.T) {
updatedConfig.Name = "UpdatedName"
updatedConfig.ToolsToExecute = []string{"calculator", "echo"}

err := manager.EditClient(clientID, updatedConfig)
err := manager.EditClient(clientID, &updatedConfig)
require.NoError(t, err, "should edit client")

// Verify changes
Expand All @@ -199,7 +199,7 @@ func TestEditClientInvalidID(t *testing.T) {

// Try to edit non-existent client
clientConfig := GetSampleHTTPClientConfig("http://example.com")
err := manager.EditClient("non-existent-id", clientConfig)
err := manager.EditClient("non-existent-id", &clientConfig)
assert.Error(t, err, "should error when editing non-existent client")
}

Expand All @@ -225,7 +225,7 @@ func TestEditClientInvalidConfig(t *testing.T) {
// Missing ConnectionString
}

err := manager.EditClient(clientID, invalidConfig)
err := manager.EditClient(clientID, &invalidConfig)
// Should return error or leave client unchanged
if err == nil {
clients = manager.GetClients()
Expand Down Expand Up @@ -257,7 +257,7 @@ func TestEditClientChangeConnectionType(t *testing.T) {
updatedConfig := clientConfig
updatedConfig.ConnectionType = schemas.MCPConnectionTypeSSE

err := manager.EditClient(clientID, updatedConfig)
err := manager.EditClient(clientID, &updatedConfig)
assert.Error(t, err, "should not allow connection type change")
clients = manager.GetClients()
if len(clients) > 0 {
Expand Down Expand Up @@ -428,7 +428,7 @@ func TestConcurrentClientOperations(t *testing.T) {
clientConfig := GetSampleHTTPClientConfig(config.HTTPServerURL)
clientConfig.ID = string(rune('a'+id)) + "-concurrent-client"

err := manager.AddClient(clientConfig)
err := manager.AddClient(&clientConfig)
if err != nil {
errors <- err
}
Expand Down
2 changes: 1 addition & 1 deletion core/internal/mcptests/concurrency_advanced_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -163,7 +163,7 @@ func TestConcurrent_AddRemoveClients(t *testing.T) {
ToolsToAutoExecute: []string{},
}

err := manager.AddClient(clientConfig)
err := manager.AddClient(&clientConfig)
if err != nil {
// InProcess connections without a server instance will fail
// This is expected - we're just testing that the operations are concurrent and don't deadlock
Expand Down
2 changes: 1 addition & 1 deletion core/internal/mcptests/concurrency_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -437,7 +437,7 @@ func TestConcurrent_EditClientDuringExecution(t *testing.T) {
// Edit client - update name (must not contain spaces)
updatedConfig := clientConfig
updatedConfig.Name = "UpdatedClientName"
err := manager.EditClient(clientConfig.ID, updatedConfig)
err := manager.EditClient(clientConfig.ID, &updatedConfig)
if err != nil {
errors <- fmt.Errorf("failed to edit client: %v", err)
} else {
Expand Down
8 changes: 4 additions & 4 deletions core/internal/mcptests/error_handling_protocol_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -289,7 +289,7 @@ func TestErrorHandling_STDIO_MCPErrorResponse(t *testing.T) {
errorServerConfig := GetErrorTestServerConfig(bifrostRoot)

manager := setupMCPManager(t)
err := manager.AddClient(errorServerConfig)
err := manager.AddClient(&errorServerConfig)
if err != nil {
t.Skipf("error-test-server not available: %v (build with: cd examples/mcps/error-test-server && go build -o bin/error-test-server)", err)
}
Expand Down Expand Up @@ -367,7 +367,7 @@ func TestErrorHandling_STDIO_TimeoutScenario(t *testing.T) {
ToolExecutionTimeout: 2 * time.Second, // 2 second timeout
})

err := manager.AddClient(errorServerConfig)
err := manager.AddClient(&errorServerConfig)
if err != nil {
t.Skipf("error-test-server not available: %v", err)
}
Expand Down Expand Up @@ -420,7 +420,7 @@ func TestErrorHandling_STDIO_MalformedJSON(t *testing.T) {
errorServerConfig := GetErrorTestServerConfig(bifrostRoot)

manager := setupMCPManager(t)
err := manager.AddClient(errorServerConfig)
err := manager.AddClient(&errorServerConfig)
if err != nil {
t.Skipf("error-test-server not available: %v", err)
}
Expand Down Expand Up @@ -469,7 +469,7 @@ func TestErrorHandling_STDIO_IntermittentFailures(t *testing.T) {
errorServerConfig := GetErrorTestServerConfig(bifrostRoot)

manager := setupMCPManager(t)
err := manager.AddClient(errorServerConfig)
err := manager.AddClient(&errorServerConfig)
if err != nil {
t.Skipf("error-test-server not available: %v", err)
}
Expand Down
23 changes: 17 additions & 6 deletions core/internal/mcptests/fixtures.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ import (

bifrost "github.com/maximhq/bifrost/core"
"github.com/maximhq/bifrost/core/mcp"
mcpcore "github.com/maximhq/bifrost/core/mcp"
"github.com/maximhq/bifrost/core/mcp/codemode/starlark"
"github.com/maximhq/bifrost/core/schemas"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
Expand Down Expand Up @@ -1461,17 +1461,28 @@ func setupBifrost(t *testing.T) *bifrost.Bifrost {
}

// setupMCPManager creates an MCP manager for testing
func setupMCPManager(t *testing.T, clientConfigs ...schemas.MCPClientConfig) *mcpcore.MCPManager {
func setupMCPManager(t *testing.T, clientConfigs ...schemas.MCPClientConfig) *mcp.MCPManager {
t.Helper()

logger := &testLogger{t: t}

// Convert to pointer slice for MCPConfig
clientConfigPtrs := make([]*schemas.MCPClientConfig, len(clientConfigs))
for i := range clientConfigs {
clientConfigPtrs[i] = &clientConfigs[i]
}

// Create MCP config
mcpConfig := &schemas.MCPConfig{
ClientConfigs: clientConfigs,
ClientConfigs: clientConfigPtrs,
}

// Create MCP manager
logger := &testLogger{t: t}
manager := mcpcore.NewMCPManager(context.Background(), *mcpConfig, logger)
// Create Starlark CodeMode
starlark.SetLogger(logger)
codeMode := starlark.NewStarlarkCodeMode(nil)

// Create MCP manager - dependencies are injected automatically
manager := mcp.NewMCPManager(context.Background(), *mcpConfig, nil, logger, codeMode)

// Cleanup
t.Cleanup(func() {
Expand Down
6 changes: 3 additions & 3 deletions core/internal/mcptests/health_monitoring_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ func TestHealthCheckSTDIOServerDropAndRecoverIn20Seconds(t *testing.T) {
t.Logf("✅ Health monitor detected server drop")

// 5. Restart STDIO process (re-add client)
err = manager.AddClient(clientConfig)
err = manager.AddClient(&clientConfig)
require.NoError(t, err, "should re-add client to simulate server recovery")
t.Logf("🔄 Simulated STDIO server recovery by re-adding client")

Expand Down Expand Up @@ -181,7 +181,7 @@ func TestHealthCheckStateTransitions(t *testing.T) {
assert.Len(t, clients, 0, "client should be removed")

// Re-add client (simulates reconnection)
err = manager.AddClient(clientConfig)
err = manager.AddClient(&clientConfig)
require.NoError(t, err, "should re-add client")

// Verify client is connected again
Expand Down Expand Up @@ -392,7 +392,7 @@ func TestHealthCheckReconnectAfterFailure(t *testing.T) {
time.Sleep(2 * time.Second)

// Re-add client (manual reconnection)
err = manager.AddClient(clientConfig)
err = manager.AddClient(&clientConfig)
require.NoError(t, err, "should re-add client")

// Wait for health monitoring to stabilize
Expand Down
4 changes: 2 additions & 2 deletions core/internal/mcptests/integration_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ func TestIntegration_FullChatWorkflow(t *testing.T) {
httpConfig := GetSampleHTTPClientConfig(config.HTTPServerURL)
httpConfig.ID = "http-integration-test"
applyTestConfigHeaders(t, &httpConfig)
err := manager.AddClient(httpConfig)
err := manager.AddClient(&httpConfig)
if err != nil {
t.Logf("Could not add HTTP client: %v", err)
}
Expand Down Expand Up @@ -341,7 +341,7 @@ func TestIntegration_ReconnectDuringExecution(t *testing.T) {
httpConfig := GetSampleHTTPClientConfig(config.HTTPServerURL)
httpConfig.ID = "reconnect-test-client"
applyTestConfigHeaders(t, &httpConfig)
err := manager.AddClient(httpConfig)
err := manager.AddClient(&httpConfig)
require.NoError(t, err, "should add HTTP client")

// Wait for client to connect
Expand Down
2 changes: 1 addition & 1 deletion core/internal/mcptests/tool_filtering_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -335,7 +335,7 @@ func TestFilteringChangesAfterClientEdit(t *testing.T) {

// Edit client to only allow second tool
clientConfig.ToolsToExecute = []string{"hash"}
err := manager.EditClient(clientConfig.ID, clientConfig)
err := manager.EditClient(clientConfig.ID, &clientConfig)
require.NoError(t, err, "edit should succeed")

// Verify configuration changed
Expand Down
Loading