diff --git a/cmd/thv/app/mcp.go b/cmd/thv/app/mcp.go index 876938f413..2606118ceb 100644 --- a/cmd/thv/app/mcp.go +++ b/cmd/thv/app/mcp.go @@ -13,9 +13,9 @@ import ( "text/tabwriter" "time" - "github.com/mark3labs/mcp-go/mcp" "github.com/spf13/cobra" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" thclient "github.com/stacklok/toolhive/pkg/mcp/client" "github.com/stacklok/toolhive/pkg/workloads" ) diff --git a/cmd/thv/app/mcp_call.go b/cmd/thv/app/mcp_call.go index 2932124844..78deceb96b 100644 --- a/cmd/thv/app/mcp_call.go +++ b/cmd/thv/app/mcp_call.go @@ -12,9 +12,9 @@ import ( "os" "time" - "github.com/mark3labs/mcp-go/mcp" "github.com/spf13/cobra" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" thclient "github.com/stacklok/toolhive/pkg/mcp/client" ) diff --git a/cmd/thv/app/mcp_call_test.go b/cmd/thv/app/mcp_call_test.go index 489db112dc..8f57e84454 100644 --- a/cmd/thv/app/mcp_call_test.go +++ b/cmd/thv/app/mcp_call_test.go @@ -9,9 +9,10 @@ import ( "strings" "testing" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) func TestReadToolArgs(t *testing.T) { diff --git a/go.mod b/go.mod index 5ad289a951..b2eeff8bab 100644 --- a/go.mod +++ b/go.mod @@ -32,7 +32,6 @@ require ( github.com/google/uuid v1.6.0 github.com/lestrrat-go/httprc/v3 v3.0.6 github.com/lestrrat-go/jwx/v3 v3.0.13 - github.com/mark3labs/mcp-go v0.55.1 github.com/moby/moby/client v0.4.1 github.com/modelcontextprotocol/registry v1.8.0 github.com/oauth2-proxy/mockoidc v0.0.0-20240214162133-caebfff84d25 @@ -50,7 +49,7 @@ require ( github.com/shirou/gopsutil/v4 v4.26.5 github.com/spf13/viper v1.21.0 github.com/stacklok/toolhive-catalog v0.20260706.0 - github.com/stacklok/toolhive-core v0.0.28 + github.com/stacklok/toolhive-core v0.0.29 github.com/stretchr/testify v1.11.1 github.com/swaggo/swag/v2 v2.0.0-rc5 github.com/tailscale/hujson v0.0.0-20260302212456-ecc657c15afd @@ -91,8 +90,9 @@ require ( github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.31 // indirect github.com/go-openapi/runtime/server-middleware v0.30.0 // indirect github.com/klauspost/cpuid/v2 v2.3.0 // indirect + github.com/modelcontextprotocol/go-sdk v1.6.1 // indirect github.com/oklog/ulid/v2 v2.1.1 // indirect - github.com/santhosh-tekuri/jsonschema/v6 v6.0.2 // indirect + github.com/segmentio/encoding v0.5.4 // indirect github.com/youmark/pkcs8 v0.0.0-20240726163527-a2c0da244d78 // indirect ) diff --git a/go.sum b/go.sum index f31ded8863..ef6720dbc3 100644 --- a/go.sum +++ b/go.sum @@ -183,8 +183,6 @@ github.com/digitorus/timestamp v0.0.0-20231217203849-220c5c2851b7 h1:lxmTCgmHE1G github.com/digitorus/timestamp v0.0.0-20231217203849-220c5c2851b7/go.mod h1:GvWntX9qiTlOud0WkQ6ewFm0LPy5JUR1Xo0Ngbd1w6Y= github.com/distribution/reference v0.6.0 h1:0IXCQ5g4/QMHHkarYzh5l+u8T3t73zM5QvfrDyIgxBk= github.com/distribution/reference v0.6.0/go.mod h1:BbU0aIcezP1/5jX/8MP0YiH4SdvB5Y4f/wlDRiLyi3E= -github.com/dlclark/regexp2 v1.11.0 h1:G/nrcoOa7ZXlpoa/91N3X7mM3r8eIlMBBJZvsz/mxKI= -github.com/dlclark/regexp2 v1.11.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= github.com/docker/cli v29.6.0+incompatible h1:nw9himxMMZ7eIeherJNlKQq+acnlzGgHd+4uf10QRSc= github.com/docker/cli v29.6.0+incompatible/go.mod h1:JLrzqnKDaYBop7H2jaqPtU4hHvMKP+vjCwu2uszcLI8= github.com/docker/docker-credential-helpers v0.9.3 h1:gAm/VtF9wgqJMoxzT3Gj5p4AqIjCBS4wrsOh9yRqcz8= @@ -550,8 +548,6 @@ github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0/go.mod h1:zJYVVT2 github.com/luna-duclos/instrumentedsql v1.1.3/go.mod h1:9J1njvFds+zN7y85EDhN9XNQLANWwZt2ULeIC8yMNYs= github.com/magiconair/properties v1.8.10 h1:s31yESBquKXCV9a/ScB3ESkOjUYYv+X0rg8SYxI99mE= github.com/magiconair/properties v1.8.10/go.mod h1:Dhd985XPs7jluiymwWYZ0G4Z61jb3vdS329zhj2hYo0= -github.com/mark3labs/mcp-go v0.55.1 h1:GLYqNm9qdMGPhCtK4g1t1y1vhAPfayOBuaibDi4mrSA= -github.com/mark3labs/mcp-go v0.55.1/go.mod h1:+8WclSK1ZUweCP3hvktSji8n8ABG/95QaEkeVE/Uwas= github.com/maruel/natural v1.1.1 h1:Hja7XhhmvEFhcByqDoHz9QZbkWey+COd9xWfCfn1ioo= github.com/maruel/natural v1.1.1/go.mod h1:v+Rfd79xlw1AgVBjbO0BEQmptqb5HvL/k9GRHB7ZKEg= github.com/mattn/go-colorable v0.1.1/go.mod h1:FuOcm+DKB9mbwrcAfNl7/TZVBZ6rcnceauSikq3lYCQ= @@ -607,6 +603,8 @@ github.com/moby/sys/userns v0.1.0 h1:tVLXkFOxVu9A64/yh59slHVv9ahO9UIev4JZusOLG/g github.com/moby/sys/userns v0.1.0/go.mod h1:IHUYgu/kao6N8YZlp9Cf444ySSvCmDlmzUcYfDHOl28= github.com/moby/term v0.5.2 h1:6qk3FJAFDs6i/q3W/pQ97SX192qKfZgGjCQqfCJkgzQ= github.com/moby/term v0.5.2/go.mod h1:d3djjFCrjnB+fl8NJux+EJzu0msscUP+f8it8hPkFLc= +github.com/modelcontextprotocol/go-sdk v1.6.1 h1:0zOSupjKUxPKSocPT1Wtago+mUHU2/uZ4xSOY0FGReU= +github.com/modelcontextprotocol/go-sdk v1.6.1/go.mod h1:kzm3kzFL1/+AziGOE0nUs3gvPoNxMCvkxokMkuFapXQ= github.com/modelcontextprotocol/registry v1.8.0 h1:x/seX0ji4iqRUpSovmkBcGbxfRiZZ7dgPwBcpCJrSTM= github.com/modelcontextprotocol/registry v1.8.0/go.mod h1:G6AUpTpZSekQvcLl5griUjijEE8vedARE/TyCaHEFdo= github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= @@ -718,8 +716,6 @@ github.com/ryanuber/go-glob v1.0.0 h1:iQh3xXAumdQ+4Ufa5b25cRpC5TYKlno6hsv6Cb3pkB github.com/ryanuber/go-glob v1.0.0/go.mod h1:807d1WSdnB0XRJzKNil9Om6lcp/3a0v4qIHxIXzX/Yc= github.com/sagikazarmark/locafero v0.11.0 h1:1iurJgmM9G3PA/I+wWYIOw/5SyBtxapeHDcg+AAIFXc= github.com/sagikazarmark/locafero v0.11.0/go.mod h1:nVIGvgyzw595SUSUE6tvCp3YYTeHs15MvlmU87WwIik= -github.com/santhosh-tekuri/jsonschema/v6 v6.0.2 h1:KRzFb2m7YtdldCEkzs6KqmJw4nqEVZGK7IN2kJkjTuQ= -github.com/santhosh-tekuri/jsonschema/v6 v6.0.2/go.mod h1:JXeL+ps8p7/KNMjDQk3TCwPpBy0wYklyWTfbkIzdIFU= github.com/sassoftware/relic v7.2.1+incompatible h1:Pwyh1F3I0r4clFJXkSI8bOyJINGqpgjJU3DYAZeI05A= github.com/sassoftware/relic v7.2.1+incompatible/go.mod h1:CWfAxv73/iLZ17rbyhIEq3K9hs5w6FpNMdUT//qR+zk= github.com/sassoftware/relic/v7 v7.6.2 h1:rS44Lbv9G9eXsukknS4mSjIAuuX+lMq/FnStgmZlUv4= @@ -731,6 +727,8 @@ github.com/secure-systems-lab/go-securesystemslib v0.11.0 h1:iuCR9kcMFD4QurdKrGv github.com/secure-systems-lab/go-securesystemslib v0.11.0/go.mod h1:+PMOTjUGwHj2vcZ+TFKlb1tXRbrdWE1LYDT5i9JC80Q= github.com/segmentio/asm v1.2.1 h1:DTNbBqs57ioxAD4PrArqftgypG4/qNpXoJx8TVXxPR0= github.com/segmentio/asm v1.2.1/go.mod h1:BqMnlJP91P8d+4ibuonYZw9mfnzI9HfxselHZr5aAcs= +github.com/segmentio/encoding v0.5.4 h1:OW1VRern8Nw6ITAtwSZ7Idrl3MXCFwXHPgqESYfvNt0= +github.com/segmentio/encoding v0.5.4/go.mod h1:HS1ZKa3kSN32ZHVZ7ZLPLXWvOVIiZtyJnO1gPH1sKt0= github.com/sergi/go-diff v1.2.0/go.mod h1:STckp+ISIX8hZLjrqAeVduY0gWCT9IjLuqbuNXdaHfM= github.com/sergi/go-diff v1.4.0 h1:n/SP9D5ad1fORl+llWyN+D6qoUETXNZARKjyY2/KVCw= github.com/sergi/go-diff v1.4.0/go.mod h1:A0bzQcvG0E7Rwjx0REVgAGH58e96+X0MeOfepqsbeW4= @@ -789,8 +787,8 @@ github.com/spf13/viper v1.21.0 h1:x5S+0EU27Lbphp4UKm1C+1oQO+rKx36vfCoaVebLFSU= github.com/spf13/viper v1.21.0/go.mod h1:P0lhsswPGWD/1lZJ9ny3fYnVqxiegrlNrEmgLjbTCAY= github.com/stacklok/toolhive-catalog v0.20260706.0 h1:ZgEtTogdPjl2Tc/zamYExRH609HDJxS+EbxgdjmMH/w= github.com/stacklok/toolhive-catalog v0.20260706.0/go.mod h1:J+l8hyYSbIhLCe3IsBJuVpP1wQbsuUXRS5o+B9M0GpU= -github.com/stacklok/toolhive-core v0.0.28 h1:+ajc8S/YdPm59Ik10pQxDYk75HjpyzwUU/Ru1WiLOO8= -github.com/stacklok/toolhive-core v0.0.28/go.mod h1:t8dqTOjyT67i+xidAS1cRWh3K5qYW38R/NgP7xlvirE= +github.com/stacklok/toolhive-core v0.0.29 h1:3djwGLG/+mav8zM9wKYFlBMmcNnyQ2rW3cS7tEJsIH0= +github.com/stacklok/toolhive-core v0.0.29/go.mod h1:t8dqTOjyT67i+xidAS1cRWh3K5qYW38R/NgP7xlvirE= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.2.0/go.mod h1:qt09Ya8vawLte6SNmTgCsAVtYtaKzEcn8ATUoHMkEqE= diff --git a/pkg/authz/annotation_cache.go b/pkg/authz/annotation_cache.go index 1637e96b29..bda00e21e1 100644 --- a/pkg/authz/annotation_cache.go +++ b/pkg/authz/annotation_cache.go @@ -6,8 +6,7 @@ package authz import ( "sync" - "github.com/mark3labs/mcp-go/mcp" - + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/authz/authorizers" ) diff --git a/pkg/authz/annotation_cache_test.go b/pkg/authz/annotation_cache_test.go index d4a4483d5c..a8c9930358 100644 --- a/pkg/authz/annotation_cache_test.go +++ b/pkg/authz/annotation_cache_test.go @@ -7,10 +7,10 @@ import ( "sync" "testing" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/authz/authorizers" ) diff --git a/pkg/authz/integration_test.go b/pkg/authz/integration_test.go index 79ca874820..217ee072b7 100644 --- a/pkg/authz/integration_test.go +++ b/pkg/authz/integration_test.go @@ -11,11 +11,11 @@ import ( "testing" "github.com/golang-jwt/jwt/v5" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "golang.org/x/exp/jsonrpc2" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/authz/authorizers/cedar" mcpparser "github.com/stacklok/toolhive/pkg/mcp" diff --git a/pkg/authz/middleware_test.go b/pkg/authz/middleware_test.go index 2c3a815561..9e93943485 100644 --- a/pkg/authz/middleware_test.go +++ b/pkg/authz/middleware_test.go @@ -14,12 +14,12 @@ import ( "testing" "github.com/golang-jwt/jwt/v5" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" "golang.org/x/exp/jsonrpc2" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/authz/authorizers" "github.com/stacklok/toolhive/pkg/authz/authorizers/cedar" diff --git a/pkg/authz/response_filter.go b/pkg/authz/response_filter.go index a713113a42..46be1731be 100644 --- a/pkg/authz/response_filter.go +++ b/pkg/authz/response_filter.go @@ -13,9 +13,9 @@ import ( "net/http" "strings" - "github.com/mark3labs/mcp-go/mcp" "golang.org/x/exp/jsonrpc2" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/authz/authorizers" "github.com/stacklok/toolhive/pkg/vmcp/optimizer" "github.com/stacklok/toolhive/pkg/vmcp/session/optimizerdec" diff --git a/pkg/authz/response_filter_test.go b/pkg/authz/response_filter_test.go index 37e09e68a0..a6d1f20f47 100644 --- a/pkg/authz/response_filter_test.go +++ b/pkg/authz/response_filter_test.go @@ -16,11 +16,11 @@ import ( "testing" "github.com/golang-jwt/jwt/v5" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "golang.org/x/exp/jsonrpc2" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/authz/authorizers/cedar" mcpparser "github.com/stacklok/toolhive/pkg/mcp" diff --git a/pkg/authz/tool_filter.go b/pkg/authz/tool_filter.go index 0001801c69..44c736714e 100644 --- a/pkg/authz/tool_filter.go +++ b/pkg/authz/tool_filter.go @@ -7,8 +7,7 @@ import ( "context" "log/slog" - "github.com/mark3labs/mcp-go/mcp" - + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/authz/authorizers" ) diff --git a/pkg/authz/tool_filter_test.go b/pkg/authz/tool_filter_test.go index 0043de1526..882cb9fd93 100644 --- a/pkg/authz/tool_filter_test.go +++ b/pkg/authz/tool_filter_test.go @@ -8,10 +8,10 @@ import ( "errors" "testing" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/authz/authorizers" "github.com/stacklok/toolhive/pkg/authz/authorizers/cedar" diff --git a/pkg/mcp/client/call.go b/pkg/mcp/client/call.go index 84e050a696..d51c191b73 100644 --- a/pkg/mcp/client/call.go +++ b/pkg/mcp/client/call.go @@ -8,7 +8,7 @@ import ( "fmt" "log/slog" - "github.com/mark3labs/mcp-go/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) // CallTool connects to the MCP server at serverURL, performs the initialize diff --git a/pkg/mcp/client/client.go b/pkg/mcp/client/client.go index 34db99dd3e..0ae9c77baf 100644 --- a/pkg/mcp/client/client.go +++ b/pkg/mcp/client/client.go @@ -16,9 +16,8 @@ import ( "net/url" "strings" - mcpclient "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/mcp" - + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/transport/ssecommon" "github.com/stacklok/toolhive/pkg/transport/streamable" "github.com/stacklok/toolhive/pkg/transport/types" diff --git a/pkg/mcp/parser_integration_test.go b/pkg/mcp/parser_integration_test.go index 398dba5f20..454a165fb1 100644 --- a/pkg/mcp/parser_integration_test.go +++ b/pkg/mcp/parser_integration_test.go @@ -10,11 +10,12 @@ import ( "testing" "time" - "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" ) // TestParsingMiddlewareWithRealMCPClients tests the parsing middleware with real MCP clients and servers diff --git a/pkg/script/executor.go b/pkg/script/executor.go index 0ebd9d8dbf..74d6f77da9 100644 --- a/pkg/script/executor.go +++ b/pkg/script/executor.go @@ -9,9 +9,9 @@ import ( "fmt" "strings" - "github.com/mark3labs/mcp-go/mcp" "go.starlark.net/starlark" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/script/internal/builtins" "github.com/stacklok/toolhive/pkg/script/internal/conversions" "github.com/stacklok/toolhive/pkg/script/internal/core" diff --git a/pkg/script/internal/builtins/builtins.go b/pkg/script/internal/builtins/builtins.go index 3cd94e220f..a7a2d655bf 100644 --- a/pkg/script/internal/builtins/builtins.go +++ b/pkg/script/internal/builtins/builtins.go @@ -8,9 +8,9 @@ import ( "context" "log/slog" - "github.com/mark3labs/mcp-go/mcp" "go.starlark.net/starlark" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/script/internal/conversions" ) diff --git a/pkg/script/internal/builtins/builtins_test.go b/pkg/script/internal/builtins/builtins_test.go index 2068004c77..07417926e0 100644 --- a/pkg/script/internal/builtins/builtins_test.go +++ b/pkg/script/internal/builtins/builtins_test.go @@ -10,10 +10,10 @@ import ( "testing" "time" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/require" "go.starlark.net/starlark" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/script/internal/core" ) diff --git a/pkg/script/internal/builtins/tools.go b/pkg/script/internal/builtins/tools.go index 59766c2d1a..a990653741 100644 --- a/pkg/script/internal/builtins/tools.go +++ b/pkg/script/internal/builtins/tools.go @@ -7,9 +7,9 @@ import ( "context" "fmt" - "github.com/mark3labs/mcp-go/mcp" "go.starlark.net/starlark" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/script/internal/conversions" ) diff --git a/pkg/script/internal/conversions/result.go b/pkg/script/internal/conversions/result.go index 2af2951b8f..80717eb348 100644 --- a/pkg/script/internal/conversions/result.go +++ b/pkg/script/internal/conversions/result.go @@ -8,7 +8,7 @@ import ( "fmt" "log/slog" - "github.com/mark3labs/mcp-go/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) // ParseToolResult converts an mcp.CallToolResult into a Go value suitable diff --git a/pkg/script/internal/conversions/result_test.go b/pkg/script/internal/conversions/result_test.go index 5f534fe04c..821e06f1bb 100644 --- a/pkg/script/internal/conversions/result_test.go +++ b/pkg/script/internal/conversions/result_test.go @@ -6,8 +6,9 @@ package conversions import ( "testing" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) func TestParseToolResult(t *testing.T) { diff --git a/pkg/script/script.go b/pkg/script/script.go index dc1091310b..7696e652e0 100644 --- a/pkg/script/script.go +++ b/pkg/script/script.go @@ -9,7 +9,7 @@ package script import ( "context" - "github.com/mark3labs/mcp-go/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) // DefaultStepLimit is the default maximum number of Starlark execution steps. diff --git a/pkg/script/script_test.go b/pkg/script/script_test.go index a9f65ec50b..f92b761e7f 100644 --- a/pkg/script/script_test.go +++ b/pkg/script/script_test.go @@ -9,8 +9,9 @@ import ( "fmt" "testing" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) func TestExecutor(t *testing.T) { diff --git a/pkg/telemetry/middleware.go b/pkg/telemetry/middleware.go index cc17e5f4db..1b9f5eec6e 100644 --- a/pkg/telemetry/middleware.go +++ b/pkg/telemetry/middleware.go @@ -15,7 +15,6 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/mcp" "go.opentelemetry.io/otel" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/codes" @@ -23,6 +22,7 @@ import ( "go.opentelemetry.io/otel/propagation" "go.opentelemetry.io/otel/trace" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpparser "github.com/stacklok/toolhive/pkg/mcp" "github.com/stacklok/toolhive/pkg/transport/types" ) diff --git a/pkg/transport/bridge.go b/pkg/transport/bridge.go index 796cd563c8..a97804818e 100644 --- a/pkg/transport/bridge.go +++ b/pkg/transport/bridge.go @@ -11,11 +11,10 @@ import ( "strings" "sync" - "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" - + "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/transport/types" "github.com/stacklok/toolhive/pkg/versions" ) @@ -105,6 +104,19 @@ func (b *StdioBridge) run(ctx context.Context) { params = map[string]any{} } + // On a *_list_changed notification, re-fetch the upstream capability set + // before forwarding the notification. forwardAll is additive only: + // AddTool/AddResource/AddPrompt upsert by name/URI, so re-running it adds + // or updates capabilities but does NOT prune ones the upstream removed — + // a stale capability stays advertised locally after an upstream removal. + // (SetTools exists if pruning is needed later.) + switch n.Method { + case "notifications/tools/list_changed", + "notifications/resources/list_changed", + "notifications/prompts/list_changed": + b.forwardAll(context.Background()) + } + b.srv.SendNotificationToAllClients(n.Method, params) }) diff --git a/pkg/transport/bridge_test.go b/pkg/transport/bridge_test.go new file mode 100644 index 0000000000..0d0b2a8e39 --- /dev/null +++ b/pkg/transport/bridge_test.go @@ -0,0 +1,284 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package transport + +import ( + "context" + "encoding/json" + "net/http/httptest" + "os" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" + "github.com/stacklok/toolhive/pkg/transport/types" +) + +// toolNamesOnServer returns the names of the tools currently registered on the +// given local MCP server, by issuing a synthetic tools/list request through the +// shim's HandleMessage dispatch path. The bridge's forwardAll registers upstream +// tools here, so this is the assertion seam for "what does the bridge advertise +// downstream". Call only after the bridge has fully stopped (run() returned), +// since StdioBridge exposes no synchronization for its srv field while running. +func toolNamesOnServer(t *testing.T, srv *server.MCPServer) []string { + t.Helper() + resp := srv.HandleMessage(context.Background(), []byte(`{"jsonrpc":"2.0","method":"tools/list","id":1}`)) + jr, ok := resp.(mcp.JSONRPCResponse) + require.True(t, ok, "tools/list should return a JSONRPCResponse, got %T", resp) + buf, err := json.Marshal(jr.Result) + require.NoError(t, err) + var lt mcp.ListToolsResult + require.NoError(t, json.Unmarshal(buf, <)) + names := make([]string, 0, len(lt.Tools)) + for _, tl := range lt.Tools { + names = append(names, tl.Name) + } + return names +} + +func containsTool(names []string, want string) bool { + for _, n := range names { + if n == want { + return true + } + } + return false +} + +// noopToolHandler is a stand-in tool handler; the bridge re-fetch tests never +// invoke tools, they only assert the advertised set. +func noopToolHandler(_ context.Context, _ mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return nil, nil +} + +// listToolsCounter is a server hook that atomically counts tools/list requests +// received by the backend. The bridge's forwardAll issues a tools/list upstream +// on startup and again whenever its OnNotification handler re-syncs on a +// *_list_changed notification, so the counter is the race-free readiness and +// re-fetch signal observed entirely on the backend side. +type listToolsCounter struct { + count atomic.Int64 +} + +func (c *listToolsCounter) hook() server.OnBeforeListToolsFunc { + return func(context.Context, any, *mcp.ListToolsRequest) { c.count.Add(1) } +} + +// sessionHolder captures the single client session that connects to the backend, +// guarded by a mutex because the hook fires on the backend's session goroutine. +type sessionHolder struct { + mu sync.Mutex + _session server.ClientSession +} + +func (h *sessionHolder) set(s server.ClientSession) { + h.mu.Lock() + defer h.mu.Unlock() + h._session = s +} + +func (h *sessionHolder) get() server.ClientSession { + h.mu.Lock() + defer h.mu.Unlock() + return h._session +} + +// TestBridge_ToolsListChanged_TriggersReSync verifies the bridge's +// notifications/tools/list_changed re-fetch fix: when the upstream backend's +// tool set changes after the bridge has connected, the upstream emits a +// tools/list_changed notification, the bridge's OnNotification handler re-runs +// forwardAll, and the newly added tool appears on the bridge's local stdio +// server. +// +// The backend is a real mcpcompat streamable-HTTP server. A tool-set change +// visible to an already-connected client is driven through the per-session +// overlay (SessionWithTools.SetSessionTools): the shim syncs the overlay onto the +// live go-sdk server, which emits notifications/tools/list_changed over the +// standalone SSE stream to the bridge's upstream client (the bridge connects +// with WithContinuousListening, so the SSE stream is established). The bridge's +// OnNotification closure is the fix under test. +// +// ServeStdio blocks reading os.Stdin; this test swaps os.Stdin for a pipe so +// ServeStdio returns cleanly on stdin EOF at teardown. The bridge exposes no +// synchronization for its srv field, so readiness/re-fetch is observed via the +// backend's tools/list counter (race-free) and bridge.srv is read only after +// Shutdown returns (b.wg.Wait() provides the happens-before edge). +// +//nolint:paralleltest // Swaps process-global os.Stdin; cannot run in parallel. +func TestBridge_ToolsListChanged_TriggersReSync(t *testing.T) { + // This test swaps process-global os.Stdin, so it cannot run in parallel + // with any other test touching stdio. + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + // --- upstream backend with an initial tool "alpha" --- + counter := &listToolsCounter{} + holder := &sessionHolder{} + hooks := &server.Hooks{} + hooks.AddOnRegisterSession(func(_ context.Context, s server.ClientSession) { holder.set(s) }) + hooks.AddBeforeListTools(counter.hook()) + + backend := server.NewMCPServer( + "backend", "1.0", + server.WithToolCapabilities(true), + server.WithHooks(hooks), + ) + backend.AddTool(mcp.NewTool("alpha"), noopToolHandler) + + httpSrv := server.NewStreamableHTTPServer(backend) + ts := httptest.NewServer(httpSrv) + t.Cleanup(ts.Close) + + // --- bridge pointing at the backend (streamable-http) --- + bridge, err := NewStdioBridge("test", ts.URL, types.TransportTypeStreamableHTTP) + require.NoError(t, err) + + // Swap os.Stdin for a pipe so ServeStdio (which hardcodes os.Stdin) does not + // touch the real process stdin and can be unblocked at teardown by closing + // the write end. Leave os.Stdout real; the assertions don't read it. + origIn := os.Stdin + pipeR, pipeW, err := os.Pipe() + require.NoError(t, err) + t.Cleanup(func() { + os.Stdin = origIn + _ = pipeR.Close() + _ = pipeW.Close() + }) + os.Stdin = pipeR + + bridge.Start(ctx) + t.Cleanup(func() { + // Close the stdin write end so the go-sdk stdio reader hits EOF and + // ServeStdio returns, letting run() exit and the WaitGroup drain. Then + // Shutdown closes the upstream client and waits for run() to finish. + _ = pipeW.Close() + bridge.Shutdown() + }) + + // Step 1: wait for the bridge to connect and run its initial forwardAll, + // which issues the first upstream tools/list (counter >= 1) after the + // session is registered. Observed entirely on the backend side: race-free. + require.Eventually(t, func() bool { + return holder.get() != nil && counter.count.Load() >= 1 + }, 5*time.Second, 50*time.Millisecond, + "bridge did not complete initial forwardAll (session=%v, listCalls=%d)", + holder.get() != nil, counter.count.Load()) + + // Step 2: mutate the upstream tool set — add "beta" via the per-session + // overlay. SetSessionTools syncs onto the live go-sdk server bound to this + // session, which emits notifications/tools/list_changed over the SSE stream + // to the bridge's upstream client. The bridge's OnNotification handler then + // re-runs forwardAll, issuing a SECOND upstream tools/list (counter >= 2). + backendSession := holder.get() + swt, ok := backendSession.(server.SessionWithTools) + require.True(t, ok, "backend session must implement SessionWithTools") + swt.SetSessionTools(map[string]server.ServerTool{ + "alpha": {Tool: mcp.NewTool("alpha"), Handler: noopToolHandler}, + "beta": {Tool: mcp.NewTool("beta"), Handler: noopToolHandler}, + }) + + // Step 3: wait for the re-fetch — the second tools/list proves the + // OnNotification handler fired and re-ran forwardAll (the fix). + require.Eventually(t, func() bool { + return counter.count.Load() >= 2 + }, 5*time.Second, 50*time.Millisecond, + "bridge did not re-run forwardAll after tools/list_changed (listCalls=%d)", + counter.count.Load()) + + // Step 4: stop the bridge, then read bridge.srv. Shutdown calls + // b.wg.Wait(), which synchronizes with run()'s defer b.wg.Done() (which is + // happens-after every field write run() made), so bridge.srv is safe to read + // here without a data race. + _ = pipeW.Close() // unblock ServeStdio (stdin EOF) + bridge.Shutdown() + + names := toolNamesOnServer(t, bridge.srv) + assert.True(t, containsTool(names, "alpha"), + "alpha must be present after the re-fetch (forwardAll is additive), got %v", names) + assert.True(t, containsTool(names, "beta"), + "beta must be present after the re-sync, got %v", names) +} + +// TestBridge_ProgressAndLoggingNotifications_ForwardedByShim guards the premise +// the bridge's notification forwarding depends on: the mcpcompat client the +// bridge uses must deliver upstream notifications/progress and +// notifications/message to registered OnNotification handlers. The bridge's +// OnNotification handler then relays every method downstream unconditionally +// (SendNotificationToAllClients, in run()), so client delivery is the load- +// bearing link. +// +// An earlier shim iteration did NOT install the go-sdk +// ProgressNotificationHandler/LoggingMessageHandler and dropped these +// notifications; they were wired in toolhive-core (issue #156). This test fails +// if that regresses — a client behind the bridge would silently stop seeing +// progress/logging on long-running tool calls. +func TestBridge_ProgressAndLoggingNotifications_ForwardedByShim(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + // Live mcpcompat backend; capture the connected session so we can push + // server->client notifications to it once the client is listening. + holder := &sessionHolder{} + hooks := &server.Hooks{} + hooks.AddOnRegisterSession(func(_ context.Context, s server.ClientSession) { holder.set(s) }) + backend := server.NewMCPServer("backend", "1.0", server.WithToolCapabilities(true), server.WithHooks(hooks)) + + httpSrv := server.NewStreamableHTTPServer(backend) + ts := httptest.NewServer(httpSrv) + t.Cleanup(ts.Close) + + // The bridge connects its upstream client with the standalone SSE stream on + // (WithContinuousListening); mirror that here so server->client notifications + // have a stream to arrive on. + c, err := client.NewStreamableHttpClient(ts.URL, transport.WithContinuousListening()) + require.NoError(t, err) + t.Cleanup(func() { _ = c.Close() }) + + var mu sync.Mutex + seen := map[string]int{} + c.OnNotification(func(n mcp.JSONRPCNotification) { + mu.Lock() + seen[n.Method]++ + mu.Unlock() + }) + + require.NoError(t, c.Start(ctx)) + _, err = c.Initialize(ctx, mcp.InitializeRequest{Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{Name: "test", Version: "1.0"}, + }}) + require.NoError(t, err) + + // Logging is level-gated per the MCP spec (and per mcp-go, whose SetLevel the + // shim aliases): the client must subscribe before the server delivers + // notifications/message. Progress is not gated. Subscribe so both paths are + // exercised. + require.NoError(t, c.SetLoggingLevel(ctx, "info")) + + // Wait for the session to register (and the standalone SSE stream to attach) + // before pushing notifications, so they are not emitted into the void. + require.Eventually(t, func() bool { return holder.get() != nil }, 5*time.Second, 50*time.Millisecond, + "client did not connect to the backend") + + backend.SendNotificationToAllClients("notifications/progress", + map[string]any{"progressToken": "tok", "progress": 0.5}) + backend.SendNotificationToAllClients("notifications/message", + map[string]any{"level": "info", "data": "hello"}) + + require.Eventually(t, func() bool { + mu.Lock() + defer mu.Unlock() + return seen["notifications/progress"] >= 1 && seen["notifications/message"] >= 1 + }, 5*time.Second, 50*time.Millisecond, + "shim client must forward progress + logging notifications to OnNotification (regression: they were dropped); got %v", seen) +} diff --git a/pkg/transport/proxy/streamable/streamable_proxy_mcp_client_integration_test.go b/pkg/transport/proxy/streamable/streamable_proxy_mcp_client_integration_test.go index f085b12ad8..475ead6d05 100644 --- a/pkg/transport/proxy/streamable/streamable_proxy_mcp_client_integration_test.go +++ b/pkg/transport/proxy/streamable/streamable_proxy_mcp_client_integration_test.go @@ -11,11 +11,12 @@ import ( "testing" "time" - "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "golang.org/x/exp/jsonrpc2" + + "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) const ( diff --git a/pkg/transport/proxy/transparent/host_header_regression_test.go b/pkg/transport/proxy/transparent/host_header_regression_test.go new file mode 100644 index 0000000000..819b663761 --- /dev/null +++ b/pkg/transport/proxy/transparent/host_header_regression_test.go @@ -0,0 +1,103 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package transparent + +import ( + "net/http" + "net/http/httptest" + "net/http/httputil" + "net/url" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestRegression_LocalhostProxy_NonLocalhostHostHeaderRewritten verifies that +// tracingTransport.RoundTrip rewrites a non-localhost Host header to match the +// target URL's host. Without this, an attacker could inject a Host header to +// bypass host-based validation on the upstream server. +func TestRegression_LocalhostProxy_NonLocalhostHostHeaderRewritten(t *testing.T) { + t.Parallel() + + // Capture the Host header the upstream server actually receives. + var receivedHost string + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedHost = r.Host + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(upstream.Close) + + targetURL, err := url.Parse(upstream.URL) + require.NoError(t, err) + + p := NewTransparentProxy("127.0.0.1", 0, "", nil, nil, nil, false, false, "streamable-http", nil, nil, "", false) + + proxy := &httputil.ReverseProxy{ + Rewrite: func(pr *httputil.ProxyRequest) { + pr.SetURL(targetURL) + pr.SetXForwarded() + }, + FlushInterval: -1, + Transport: newTracingTransport(http.DefaultTransport, p), + ModifyResponse: p.modifyResponse, + } + + // Send a request with a malicious Host header. + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "http://evil.example.com/some/path", nil) + req.Host = "evil.example.com" + + proxy.ServeHTTP(rec, req) + + // The outbound request's Host must be the target URL's host, + // not the attacker-supplied value. + assert.Equal(t, targetURL.Host, receivedHost, + "tracingTransport must rewrite Host header from attacker value to target URL host") + assert.NotEqual(t, "evil.example.com", receivedHost, + "attacker Host header must not reach the upstream server") +} + +// TestRegression_TracingTransport_RewritesHostHeader exercises the Host rewrite +// directly against tracingTransport.RoundTrip. The end-to-end test above routes +// through httputil.ReverseProxy, whose SetURL sets Out.Host="" so net/http would +// derive the sent Host from URL.Host even if the tracingTransport rewrite were +// deleted. Driving RoundTrip directly with a request whose req.Host differs from +// req.URL.Host gates the rewrite itself: if the RoundTrip Host-rewrite is removed, +// the attacker-supplied req.Host reaches the upstream and this test fails. +func TestRegression_TracingTransport_RewritesHostHeader(t *testing.T) { + t.Parallel() + + var receivedHost string + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedHost = r.Host + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(upstream.Close) + + targetURL, err := url.Parse(upstream.URL) + require.NoError(t, err) + + p := NewTransparentProxy("127.0.0.1", 0, "", nil, nil, nil, false, false, "streamable-http", nil, nil, "", false) + transport := newTracingTransport(http.DefaultTransport, p) + + // The request targets the real upstream (req.URL.Host) but carries an + // attacker-supplied Host header (req.Host) that differs from it. + req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, upstream.URL+"/some/path", nil) + require.NoError(t, err) + req.Host = "evil.example.com" + require.NotEqual(t, req.Host, req.URL.Host, "test precondition: req.Host must differ from req.URL.Host") + + resp, err := transport.RoundTrip(req) + require.NoError(t, err) + t.Cleanup(func() { _ = resp.Body.Close() }) + require.Equal(t, http.StatusOK, resp.StatusCode) + + // RoundTrip must have rewritten req.Host to the target URL host before + // forwarding, so the upstream never observes the attacker value. + assert.Equal(t, targetURL.Host, receivedHost, + "tracingTransport.RoundTrip must rewrite Host header to the target URL host") + assert.NotEqual(t, "evil.example.com", receivedHost, + "attacker Host header must not reach the upstream server") +} diff --git a/pkg/tui/inspector.go b/pkg/tui/inspector.go index bb431866f5..2b4bc9ce03 100644 --- a/pkg/tui/inspector.go +++ b/pkg/tui/inspector.go @@ -14,9 +14,9 @@ import ( "github.com/charmbracelet/bubbles/textinput" tea "github.com/charmbracelet/bubbletea" - mcpclient "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/mcp" + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/core" ) diff --git a/pkg/tui/inspector_test.go b/pkg/tui/inspector_test.go index 2bff458729..a249a68406 100644 --- a/pkg/tui/inspector_test.go +++ b/pkg/tui/inspector_test.go @@ -7,10 +7,10 @@ import ( "testing" "github.com/charmbracelet/bubbles/textinput" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/core" ) diff --git a/pkg/tui/model.go b/pkg/tui/model.go index d1b1f2258b..8318c9de55 100644 --- a/pkg/tui/model.go +++ b/pkg/tui/model.go @@ -10,9 +10,9 @@ import ( "github.com/charmbracelet/bubbles/textinput" "github.com/charmbracelet/bubbles/viewport" - mcpclient "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/mcp" + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" regtypes "github.com/stacklok/toolhive-core/registry/types" "github.com/stacklok/toolhive/pkg/core" "github.com/stacklok/toolhive/pkg/registry" diff --git a/pkg/tui/tools.go b/pkg/tui/tools.go index e0e3c2fa06..e28ba07fd2 100644 --- a/pkg/tui/tools.go +++ b/pkg/tui/tools.go @@ -8,9 +8,9 @@ import ( "errors" tea "github.com/charmbracelet/bubbletea" - mcpclient "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/mcp" + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/core" thclient "github.com/stacklok/toolhive/pkg/mcp/client" ) diff --git a/pkg/tui/update.go b/pkg/tui/update.go index efbdd3d186..c4614ebdb4 100644 --- a/pkg/tui/update.go +++ b/pkg/tui/update.go @@ -9,9 +9,9 @@ import ( "time" tea "github.com/charmbracelet/bubbletea" - mcpclient "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/mcp" + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/core" "github.com/stacklok/toolhive/pkg/runner" ) diff --git a/pkg/tui/update_inspector.go b/pkg/tui/update_inspector.go index eb41e4298d..79997b87b5 100644 --- a/pkg/tui/update_inspector.go +++ b/pkg/tui/update_inspector.go @@ -9,7 +9,8 @@ import ( "github.com/atotto/clipboard" "github.com/charmbracelet/bubbles/key" tea "github.com/charmbracelet/bubbletea" - "github.com/mark3labs/mcp-go/mcp" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) // handleInspectorKey handles key input when the inspector panel is active. diff --git a/pkg/tui/update_navigation.go b/pkg/tui/update_navigation.go index c8f4a6993e..4d07fa63ad 100644 --- a/pkg/tui/update_navigation.go +++ b/pkg/tui/update_navigation.go @@ -9,8 +9,8 @@ import ( "github.com/atotto/clipboard" "github.com/charmbracelet/bubbles/key" tea "github.com/charmbracelet/bubbletea" - mcpclient "github.com/mark3labs/mcp-go/client" + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" "github.com/stacklok/toolhive/pkg/core" "github.com/stacklok/toolhive/pkg/runner" types "github.com/stacklok/toolhive/pkg/transport/types" diff --git a/pkg/tui/view_inspector.go b/pkg/tui/view_inspector.go index 5ddec4ab30..e883eda658 100644 --- a/pkg/tui/view_inspector.go +++ b/pkg/tui/view_inspector.go @@ -9,8 +9,8 @@ import ( "strings" "github.com/charmbracelet/lipgloss" - "github.com/mark3labs/mcp-go/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/cmd/thv/app/ui" ) diff --git a/pkg/vmcp/client/auth_error_mapping_regression_test.go b/pkg/vmcp/client/auth_error_mapping_regression_test.go new file mode 100644 index 0000000000..323a8d79fa --- /dev/null +++ b/pkg/vmcp/client/auth_error_mapping_regression_test.go @@ -0,0 +1,384 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package client + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/client" + mcptransport "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// jsonRPCResponse is a generic JSON-RPC 2.0 response envelope. +type jsonRPCResponse struct { + JSONRPC string `json:"jsonrpc"` + ID json.RawMessage `json:"id"` + Result json.RawMessage `json:"result,omitempty"` + Error *jsonRPCError `json:"error,omitempty"` +} + +type jsonRPCError struct { + Code int `json:"code"` + Message string `json:"message"` +} + +// jsonRPCRequest is a generic JSON-RPC 2.0 request envelope (for method routing). +type jsonRPCRequest struct { + JSONRPC string `json:"jsonrpc"` + ID json.RawMessage `json:"id"` + Method string `json:"method"` +} + +// TestRegression_401_MapsToErrAuthenticationFailed verifies that a backend +// returning HTTP 401 on initialize is classified as ErrAuthenticationFailed. +func TestRegression_401_MapsToErrAuthenticationFailed(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusUnauthorized) + _ = json.NewEncoder(w).Encode(jsonRPCResponse{ + JSONRPC: "2.0", + Error: &jsonRPCError{Code: -32000, Message: "Unauthorized"}, + }) + })) + t.Cleanup(srv.Close) + + h := &httpBackendClient{ + clientFactory: func(ctx context.Context, target *vmcp.BackendTarget) (*client.Client, error) { + c, err := client.NewStreamableHttpClient( + target.BaseURL, + mcptransport.WithHTTPTimeout(30*time.Second), + ) + if err != nil { + return nil, err + } + if err := c.Start(ctx); err != nil { + return nil, err + } + return c, nil + }, + } + + target := &vmcp.BackendTarget{ + WorkloadID: "test-401-backend", + WorkloadName: "Test 401 Backend", + BaseURL: srv.URL, + TransportType: "streamable-http", + } + + _, err := h.ListCapabilities(context.Background(), target) + require.Error(t, err) + assert.True(t, errors.Is(err, vmcp.ErrAuthenticationFailed), + "expected ErrAuthenticationFailed, got: %v", err) +} + +// TestRegression_403OnInitialize_LegacySSEFallback verifies that a backend +// returning HTTP 403 on initialize is classified as ErrBackendUnavailable. +// +// NOTE: The mcp-go streamable-HTTP transport returns a generic HTTP error for +// 403 ("request failed with status 403"), not transport.ErrLegacySSEServer. +// The "legacy SSE" hint in wrapBackendError is only added when the origin error +// IS transport.ErrLegacySSEServer (returned by SSE transport, not streamable-HTTP). +// For streamable-HTTP, 403 falls through to string-based classification and +// correctly maps to ErrBackendUnavailable, but without the SSE-specific message. +func TestRegression_403OnInitialize_LegacySSEFallback(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusForbidden) + _ = json.NewEncoder(w).Encode(jsonRPCResponse{ + JSONRPC: "2.0", + Error: &jsonRPCError{Code: -32000, Message: "Forbidden"}, + }) + })) + t.Cleanup(srv.Close) + + h := &httpBackendClient{ + clientFactory: func(ctx context.Context, target *vmcp.BackendTarget) (*client.Client, error) { + c, err := client.NewStreamableHttpClient( + target.BaseURL, + mcptransport.WithHTTPTimeout(30*time.Second), + ) + if err != nil { + return nil, err + } + if err := c.Start(ctx); err != nil { + return nil, err + } + return c, nil + }, + } + + target := &vmcp.BackendTarget{ + WorkloadID: "test-403-backend", + WorkloadName: "Test 403 Backend", + BaseURL: srv.URL, + TransportType: "streamable-http", + } + + _, err := h.ListCapabilities(context.Background(), target) + require.Error(t, err) + assert.True(t, errors.Is(err, vmcp.ErrBackendUnavailable), + "expected ErrBackendUnavailable, got: %v", err) + assert.Contains(t, err.Error(), "403", + "error message should reference 403 status, got: %v", err) +} + +// TestRegression_403OnInitialize_MatchesSentinel verifies that +// transport.ErrLegacySSEServer is NOT in the error chain for 403 on +// initialize, because wrapBackendError uses %v (not %w) for the +// original error, AND the mcp-go streamable-HTTP transport does not +// return ErrLegacySSEServer for 403 (it returns a generic HTTP error). +// Regardless of which error type is at the origin, the sentinel should +// never be in the chain. +func TestRegression_403OnInitialize_MatchesSentinel(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusForbidden) + _ = json.NewEncoder(w).Encode(jsonRPCResponse{ + JSONRPC: "2.0", + Error: &jsonRPCError{Code: -32000, Message: "Forbidden"}, + }) + })) + t.Cleanup(srv.Close) + + h := &httpBackendClient{ + clientFactory: func(ctx context.Context, target *vmcp.BackendTarget) (*client.Client, error) { + c, err := client.NewStreamableHttpClient( + target.BaseURL, + mcptransport.WithHTTPTimeout(30*time.Second), + ) + if err != nil { + return nil, err + } + if err := c.Start(ctx); err != nil { + return nil, err + } + return c, nil + }, + } + + target := &vmcp.BackendTarget{ + WorkloadID: "test-403-sentinel-backend", + WorkloadName: "Test 403 Sentinel Backend", + BaseURL: srv.URL, + TransportType: "streamable-http", + } + + _, err := h.ListCapabilities(context.Background(), target) + require.Error(t, err) + + // wrapBackendError uses %v for the original error, so + // transport.ErrLegacySSEServer is NOT in the chain. + assert.False(t, errors.Is(err, mcptransport.ErrLegacySSEServer), + "transport.ErrLegacySSEServer should NOT be in the error chain (wrapBackendError uses %v)") +} + +// TestRegression_BackendToolErrorWith401_NotClassifiedAsAuthFailure verifies +// that an MCP tool error result (IsError=true on a 200 HTTP response) whose +// message contains "401 unauthorized" is NOT classified as an auth failure. +func TestRegression_BackendToolErrorWith401_NotClassifiedAsAuthFailure(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + + body, err := io.ReadAll(r.Body) + if err != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + + var req jsonRPCRequest + if err := json.Unmarshal(body, &req); err != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + + switch req.Method { + case "initialize": + resp := jsonRPCResponse{ + JSONRPC: "2.0", + ID: req.ID, + Result: json.RawMessage(`{ + "protocolVersion": "2024-11-05", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "test-backend", "version": "1.0.0"} + }`), + } + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(resp) + + case "tools/call": + resp := jsonRPCResponse{ + JSONRPC: "2.0", + ID: req.ID, + Result: json.RawMessage(`{ + "content": [{"type": "text", "text": "tool error: 401 unauthorized - permission denied"}], + "isError": true + }`), + } + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(resp) + + case "tools/list": + resp := jsonRPCResponse{ + JSONRPC: "2.0", + ID: req.ID, + Result: json.RawMessage(`{ + "tools": [{"name": "test-tool", "description": "A test tool", "inputSchema": {"type": "object"}}] + }`), + } + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(resp) + + default: + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(jsonRPCResponse{ + JSONRPC: "2.0", + ID: req.ID, + Result: json.RawMessage(`{}`), + }) + } + })) + t.Cleanup(srv.Close) + + h := &httpBackendClient{ + clientFactory: func(ctx context.Context, target *vmcp.BackendTarget) (*client.Client, error) { + c, err := client.NewStreamableHttpClient( + target.BaseURL, + mcptransport.WithHTTPTimeout(30*time.Second), + ) + if err != nil { + return nil, err + } + if err := c.Start(ctx); err != nil { + return nil, err + } + return c, nil + }, + } + + target := &vmcp.BackendTarget{ + WorkloadID: "test-tool-error-backend", + WorkloadName: "Test Tool Error Backend", + BaseURL: srv.URL, + TransportType: "streamable-http", + } + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + result, err := h.CallTool(ctx, target, "test-tool", map[string]any{"arg": "val"}, nil) + + if err != nil { + if errors.Is(err, vmcp.ErrAuthenticationFailed) { + t.Fatalf("UNEXPECTED: CallTool returned ErrAuthenticationFailed for a 200 response with IsError=true containing '401 unauthorized'") + } + t.Fatalf("unexpected transport error from CallTool: %v", err) + } + + assert.NotNil(t, result, "expected non-nil result for IsError=true") + assert.True(t, result.IsError, "expected IsError=true") + assert.NotEmpty(t, result.Content, "expected non-empty content") +} + +// TestRegression_WrapBackendError_ClassificationMatrix pins the remaining +// wrapBackendError classification arms that the httptest-driven cases above do +// not exercise. wrapBackendError maps several distinct failure shapes onto vmcp +// sentinels, and health monitoring (#4935, #5223) branches on those sentinels; +// a regression that reshuffles the arms would silently mis-route recovery. Each +// case reconstructs the error shape the transport layer produces in the wild and +// asserts the sentinel via errors.Is (the chain-aware check wrapBackendError's +// %w wrapping is designed to satisfy). +func TestRegression_WrapBackendError_ClassificationMatrix(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + err error + wantSentinel error + }{ + { + // A backend returning a 5xx surfaces as a generic HTTP status error + // (no transport sentinel). It must classify as ErrBackendUnavailable + // so the health monitor can recover once the backend is healthy again. + name: "5xx HTTP status maps to ErrBackendUnavailable", + err: fmt.Errorf("request failed with status code 503 service unavailable"), + wantSentinel: vmcp.ErrBackendUnavailable, + }, + { + // A network timeout arrives as a net.Error whose Timeout() reports true; + // wrapBackendError's errors.As(&netErr)+Timeout() arm must map it to + // ErrTimeout ahead of any string-based fallback. + name: "net.Error timeout maps to ErrTimeout", + err: &net.DNSError{Err: "i/o timeout", Name: "backend.example", IsTimeout: true}, + wantSentinel: vmcp.ErrTimeout, + }, + { + // An unexpectedly closed connection surfaces as io.EOF; it means the + // backend dropped the stream, so it maps to ErrBackendUnavailable. + name: "io.EOF maps to ErrBackendUnavailable", + err: io.EOF, + wantSentinel: vmcp.ErrBackendUnavailable, + }, + { + // A refused dial (backend down) is a net.OpError whose message contains + // "connection refused"; the Timeout() arm is skipped (not a timeout) and + // the connection-error string arm maps it to ErrBackendUnavailable. + name: "connection refused maps to ErrBackendUnavailable", + err: &net.OpError{ + Op: "dial", + Net: "tcp", + Err: errors.New("connect: connection refused"), + }, + wantSentinel: vmcp.ErrBackendUnavailable, + }, + { + // transport.ErrAuthorizationRequired is the migration-sensitive sentinel + // (mcp-go v0.49.0+) returned for 401 + WWW-Authenticate. It must map to + // ErrAuthenticationFailed so health monitoring engages the auth-aware + // branch (#4935) instead of marking the backend unhealthy (#5223). + name: "transport.ErrAuthorizationRequired maps to ErrAuthenticationFailed", + err: mcptransport.ErrAuthorizationRequired, + wantSentinel: vmcp.ErrAuthenticationFailed, + }, + { + // The production chain wraps the sentinel in *transport.Error via + // *AuthorizationRequiredError; both layers Unwrap to the sentinel so + // errors.Is must still classify it as ErrAuthenticationFailed. + name: "wrapped AuthorizationRequiredError maps to ErrAuthenticationFailed", + err: mcptransport.NewError(&mcptransport.AuthorizationRequiredError{}), + wantSentinel: vmcp.ErrAuthenticationFailed, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + err := wrapBackendError(tt.err, "test-backend", "initialize") + require.Error(t, err) + assert.True(t, errors.Is(err, tt.wantSentinel), + "expected %v, got: %v", tt.wantSentinel, err) + }) + } +} diff --git a/pkg/vmcp/client/auth_propagation_integration_test.go b/pkg/vmcp/client/auth_propagation_integration_test.go index 7844d1a8ae..1a69d145ac 100644 --- a/pkg/vmcp/client/auth_propagation_integration_test.go +++ b/pkg/vmcp/client/auth_propagation_integration_test.go @@ -11,10 +11,10 @@ import ( "testing" "time" - "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/vmcp" vmcpauth "github.com/stacklok/toolhive/pkg/vmcp/auth" "github.com/stacklok/toolhive/pkg/vmcp/auth/strategies" diff --git a/pkg/vmcp/client/client.go b/pkg/vmcp/client/client.go index 386418a343..f59558f3da 100644 --- a/pkg/vmcp/client/client.go +++ b/pkg/vmcp/client/client.go @@ -4,7 +4,7 @@ // Package client provides MCP protocol client implementation for communicating with backend servers. // // This package implements the BackendClient interface defined in the vmcp package, -// using the mark3labs/mcp-go SDK for protocol communication. +// using the stacklok/toolhive-core/mcpcompat SDK for protocol communication. package client import ( @@ -21,12 +21,12 @@ import ( "syscall" "time" - "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" "go.opentelemetry.io/otel" "go.opentelemetry.io/otel/propagation" + "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/secrets" "github.com/stacklok/toolhive/pkg/versions" @@ -100,7 +100,7 @@ func WithDialControl(control func(network, address string, c syscall.RawConn) er } } -// httpBackendClient implements vmcp.BackendClient using mark3labs/mcp-go HTTP client. +// httpBackendClient implements vmcp.BackendClient using stacklok/toolhive-core/mcpcompat HTTP client. // It supports streamable-HTTP and SSE transports for backend MCP servers. type httpBackendClient struct { // clientFactory creates MCP clients for backends. @@ -397,7 +397,7 @@ func (h *httpBackendClient) resolveAuthStrategy(target *vmcp.BackendTarget) (vmc return strategy, nil } -// defaultClientFactory creates mark3labs MCP clients for different transport types. +// defaultClientFactory creates mcpcompat MCP clients for different transport types. func (h *httpBackendClient) defaultClientFactory(ctx context.Context, target *vmcp.BackendTarget) (*client.Client, error) { // Build transport chain (outermost to innermost, request execution order): // size limit (response body) → trace propagation → identity propagation → authentication → HTTP diff --git a/pkg/vmcp/client/client_test.go b/pkg/vmcp/client/client_test.go index f274f864c7..fda9a3a9f8 100644 --- a/pkg/vmcp/client/client_test.go +++ b/pkg/vmcp/client/client_test.go @@ -28,9 +28,6 @@ import ( "testing" "time" - "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/client/transport" - mcpserver "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.opentelemetry.io/otel/propagation" @@ -38,6 +35,9 @@ import ( "go.opentelemetry.io/otel/trace" "go.uber.org/mock/gomock" + "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" pkgauth "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/auth" diff --git a/pkg/vmcp/client/meta_integration_test.go b/pkg/vmcp/client/meta_integration_test.go index 7c65d4d7f0..e8659709b9 100644 --- a/pkg/vmcp/client/meta_integration_test.go +++ b/pkg/vmcp/client/meta_integration_test.go @@ -13,11 +13,11 @@ import ( "testing" "time" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/auth" "github.com/stacklok/toolhive/pkg/vmcp/auth/strategies" diff --git a/pkg/vmcp/codemode/decorator.go b/pkg/vmcp/codemode/decorator.go index 4a643fd85d..32127f54e1 100644 --- a/pkg/vmcp/codemode/decorator.go +++ b/pkg/vmcp/codemode/decorator.go @@ -10,8 +10,7 @@ import ( "log/slog" "time" - "github.com/mark3labs/mcp-go/mcp" - + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/script" "github.com/stacklok/toolhive/pkg/vmcp" diff --git a/pkg/vmcp/composer/composer.go b/pkg/vmcp/composer/composer.go index 915442a3f7..37b8ebc782 100644 --- a/pkg/vmcp/composer/composer.go +++ b/pkg/vmcp/composer/composer.go @@ -428,7 +428,7 @@ type WorkflowStateStore interface { // ElicitationProtocolHandler handles MCP elicitation protocol interactions. // // This interface provides an SDK-agnostic abstraction for elicitation requests, -// enabling migration from mark3labs SDK to official SDK without changing workflow code. +// enabling migration from mcpcompat SDK to official SDK without changing workflow code. // // Per MCP 2025-06-18 spec: Elicitation is a synchronous request/response protocol // where the server sends a request and blocks until the client responds. diff --git a/pkg/vmcp/conversion/content.go b/pkg/vmcp/conversion/content.go index a1a4382a5d..0c46fa36e5 100644 --- a/pkg/vmcp/conversion/content.go +++ b/pkg/vmcp/conversion/content.go @@ -10,8 +10,7 @@ import ( "fmt" "log/slog" - "github.com/mark3labs/mcp-go/mcp" - + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/vmcp" ) diff --git a/pkg/vmcp/conversion/content_test.go b/pkg/vmcp/conversion/content_test.go index b7d700de80..8e38921d9f 100644 --- a/pkg/vmcp/conversion/content_test.go +++ b/pkg/vmcp/conversion/content_test.go @@ -6,10 +6,10 @@ package conversion import ( "testing" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/vmcp" ) diff --git a/pkg/vmcp/conversion/conversion_test.go b/pkg/vmcp/conversion/conversion_test.go index 41044e47b5..3713ae8ce8 100644 --- a/pkg/vmcp/conversion/conversion_test.go +++ b/pkg/vmcp/conversion/conversion_test.go @@ -6,10 +6,10 @@ package conversion_test import ( "testing" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/conversion" ) diff --git a/pkg/vmcp/conversion/errors.go b/pkg/vmcp/conversion/errors.go index 83058a514f..0bf6a2ad63 100644 --- a/pkg/vmcp/conversion/errors.go +++ b/pkg/vmcp/conversion/errors.go @@ -6,8 +6,7 @@ package conversion import ( "errors" - sdkmcp "github.com/mark3labs/mcp-go/mcp" - + sdkmcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" thvmcp "github.com/stacklok/toolhive/pkg/mcp" "github.com/stacklok/toolhive/pkg/vmcp" ) diff --git a/pkg/vmcp/conversion/meta.go b/pkg/vmcp/conversion/meta.go index c7b6d03b0b..2264a1c02f 100644 --- a/pkg/vmcp/conversion/meta.go +++ b/pkg/vmcp/conversion/meta.go @@ -6,7 +6,7 @@ package conversion import ( "maps" - "github.com/mark3labs/mcp-go/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) // FromMCPMeta converts MCP SDK meta to map[string]any for vmcp wrapper types. diff --git a/pkg/vmcp/health/checker_test.go b/pkg/vmcp/health/checker_test.go index 47d3881edc..49bd2b6453 100644 --- a/pkg/vmcp/health/checker_test.go +++ b/pkg/vmcp/health/checker_test.go @@ -10,11 +10,11 @@ import ( "testing" "time" - "github.com/mark3labs/mcp-go/client/transport" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" "github.com/stacklok/toolhive/pkg/vmcp" authtypes "github.com/stacklok/toolhive/pkg/vmcp/auth/types" "github.com/stacklok/toolhive/pkg/vmcp/mocks" diff --git a/pkg/vmcp/optimizer/internal/tokencounter/counter.go b/pkg/vmcp/optimizer/internal/tokencounter/counter.go index f0f6495f73..606414324b 100644 --- a/pkg/vmcp/optimizer/internal/tokencounter/counter.go +++ b/pkg/vmcp/optimizer/internal/tokencounter/counter.go @@ -7,7 +7,7 @@ package tokencounter import ( "encoding/json" - "github.com/mark3labs/mcp-go/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) // Counter estimates the number of tokens a tool definition would consume diff --git a/pkg/vmcp/optimizer/internal/tokencounter/counter_test.go b/pkg/vmcp/optimizer/internal/tokencounter/counter_test.go index d3abd6cfb5..9f720a03d5 100644 --- a/pkg/vmcp/optimizer/internal/tokencounter/counter_test.go +++ b/pkg/vmcp/optimizer/internal/tokencounter/counter_test.go @@ -7,8 +7,9 @@ import ( "encoding/json" "testing" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) func TestJSONByteDivisionCounter(t *testing.T) { diff --git a/pkg/vmcp/optimizer/internal/toolstore/sqlite_store.go b/pkg/vmcp/optimizer/internal/toolstore/sqlite_store.go index a3c6f67b28..933efa5b4d 100644 --- a/pkg/vmcp/optimizer/internal/toolstore/sqlite_store.go +++ b/pkg/vmcp/optimizer/internal/toolstore/sqlite_store.go @@ -19,11 +19,11 @@ import ( "sort" "strings" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" "golang.org/x/sync/errgroup" _ "modernc.org/sqlite" // registers the "sqlite" database/sql driver + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/similarity" "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/types" ) diff --git a/pkg/vmcp/optimizer/internal/toolstore/sqlite_store_bench_test.go b/pkg/vmcp/optimizer/internal/toolstore/sqlite_store_bench_test.go index 8e5bc2728b..c9853c9d0a 100644 --- a/pkg/vmcp/optimizer/internal/toolstore/sqlite_store_bench_test.go +++ b/pkg/vmcp/optimizer/internal/toolstore/sqlite_store_bench_test.go @@ -12,10 +12,10 @@ import ( "fmt" "testing" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/types" ) diff --git a/pkg/vmcp/optimizer/internal/toolstore/sqlite_store_test.go b/pkg/vmcp/optimizer/internal/toolstore/sqlite_store_test.go index 4894a45258..3acc254192 100644 --- a/pkg/vmcp/optimizer/internal/toolstore/sqlite_store_test.go +++ b/pkg/vmcp/optimizer/internal/toolstore/sqlite_store_test.go @@ -10,10 +10,10 @@ import ( "sync/atomic" "testing" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/types" ) diff --git a/pkg/vmcp/optimizer/internal/types/mocks/mock_types.go b/pkg/vmcp/optimizer/internal/types/mocks/mock_types.go index 1a344ee73e..f9418bbe3c 100644 --- a/pkg/vmcp/optimizer/internal/types/mocks/mock_types.go +++ b/pkg/vmcp/optimizer/internal/types/mocks/mock_types.go @@ -13,8 +13,8 @@ import ( context "context" reflect "reflect" - mcp "github.com/mark3labs/mcp-go/mcp" - server "github.com/mark3labs/mcp-go/server" + mcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" + server "github.com/stacklok/toolhive-core/mcpcompat/server" gomock "go.uber.org/mock/gomock" ) diff --git a/pkg/vmcp/optimizer/internal/types/types.go b/pkg/vmcp/optimizer/internal/types/types.go index 9c54f5a92c..cb5d7d14cc 100644 --- a/pkg/vmcp/optimizer/internal/types/types.go +++ b/pkg/vmcp/optimizer/internal/types/types.go @@ -10,8 +10,8 @@ import ( "context" "time" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" ) // ToolStore defines the interface for storing and searching tools. diff --git a/pkg/vmcp/optimizer/optimizer.go b/pkg/vmcp/optimizer/optimizer.go index 35389cf0f2..5c58b39bca 100644 --- a/pkg/vmcp/optimizer/optimizer.go +++ b/pkg/vmcp/optimizer/optimizer.go @@ -21,9 +21,8 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" - + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" httpval "github.com/stacklok/toolhive-core/validation/http" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/similarity" diff --git a/pkg/vmcp/optimizer/optimizer_test.go b/pkg/vmcp/optimizer/optimizer_test.go index 1c59228507..2d862bb42e 100644 --- a/pkg/vmcp/optimizer/optimizer_test.go +++ b/pkg/vmcp/optimizer/optimizer_test.go @@ -10,12 +10,12 @@ import ( "strings" "testing" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/tokencounter" "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/types" diff --git a/pkg/vmcp/server/adapter/capability_adapter.go b/pkg/vmcp/server/adapter/capability_adapter.go index d02b2d3c6c..68463f9c06 100644 --- a/pkg/vmcp/server/adapter/capability_adapter.go +++ b/pkg/vmcp/server/adapter/capability_adapter.go @@ -8,9 +8,8 @@ import ( "fmt" "log/slog" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" - + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/conversion" ) @@ -19,7 +18,7 @@ import ( // // This is the Anti-Corruption Layer between: // - Domain model (aggregator.AggregatedCapabilities) -// - External library (mark3labs/mcp-go SDK types) +// - External library (stacklok/toolhive-core/mcpcompat SDK types) // // The adapter: // 1. Converts aggregator types to SDK types diff --git a/pkg/vmcp/server/adapter/capability_adapter_annotations_test.go b/pkg/vmcp/server/adapter/capability_adapter_annotations_test.go index 1b85496f15..263fcea5fc 100644 --- a/pkg/vmcp/server/adapter/capability_adapter_annotations_test.go +++ b/pkg/vmcp/server/adapter/capability_adapter_annotations_test.go @@ -7,12 +7,12 @@ import ( "context" "testing" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/server/adapter" "github.com/stacklok/toolhive/pkg/vmcp/server/adapter/mocks" diff --git a/pkg/vmcp/server/adapter/capability_adapter_test.go b/pkg/vmcp/server/adapter/capability_adapter_test.go index 71e0bf13d9..4150e0239c 100644 --- a/pkg/vmcp/server/adapter/capability_adapter_test.go +++ b/pkg/vmcp/server/adapter/capability_adapter_test.go @@ -7,12 +7,12 @@ import ( "context" "testing" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/server/adapter" "github.com/stacklok/toolhive/pkg/vmcp/server/adapter/mocks" diff --git a/pkg/vmcp/server/adapter/handler_factory.go b/pkg/vmcp/server/adapter/handler_factory.go index ace09654d4..f8a3dfb7a5 100644 --- a/pkg/vmcp/server/adapter/handler_factory.go +++ b/pkg/vmcp/server/adapter/handler_factory.go @@ -13,8 +13,7 @@ import ( "fmt" "log/slog" - "github.com/mark3labs/mcp-go/mcp" - + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/conversion" "github.com/stacklok/toolhive/pkg/vmcp/internal/compositetools" diff --git a/pkg/vmcp/server/adapter/handler_factory_test.go b/pkg/vmcp/server/adapter/handler_factory_test.go index f6a0b599c7..d2d36856fb 100644 --- a/pkg/vmcp/server/adapter/handler_factory_test.go +++ b/pkg/vmcp/server/adapter/handler_factory_test.go @@ -8,11 +8,11 @@ import ( "errors" "testing" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/vmcp" vmcpmocks "github.com/stacklok/toolhive/pkg/vmcp/mocks" "github.com/stacklok/toolhive/pkg/vmcp/router" diff --git a/pkg/vmcp/server/adapter/mocks/mock_handler_factory.go b/pkg/vmcp/server/adapter/mocks/mock_handler_factory.go index 5d516c2177..53862acf60 100644 --- a/pkg/vmcp/server/adapter/mocks/mock_handler_factory.go +++ b/pkg/vmcp/server/adapter/mocks/mock_handler_factory.go @@ -13,7 +13,7 @@ import ( context "context" reflect "reflect" - mcp "github.com/mark3labs/mcp-go/mcp" + mcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" adapter "github.com/stacklok/toolhive/pkg/vmcp/server/adapter" gomock "go.uber.org/mock/gomock" ) diff --git a/pkg/vmcp/server/capability_regression_test.go b/pkg/vmcp/server/capability_regression_test.go new file mode 100644 index 0000000000..af346cadda --- /dev/null +++ b/pkg/vmcp/server/capability_regression_test.go @@ -0,0 +1,65 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// TestRegression_InitializeAdvertisesToolsAndResourcesCapabilities pins the +// capabilities advertised in the initialize response on the Serve path. +// +// BEHAVIOR (go-sdk bridge, toolhive-core v0.0.28): the initialize response +// advertises tools and resources capabilities alongside logging, matching +// mcp-go. Earlier releases advertised only {"logging":{}} on the Serve path; +// v0.0.28 surfaces the tool/resource capabilities in the initialize response. +// +// This test pins that behavior so a future regression is a deliberate, visible +// flip rather than a silent drift. It asserts on the RAW initialize response +// body parsed with encoding/json — not tools/list. +func TestRegression_InitializeAdvertisesToolsAndResourcesCapabilities(t *testing.T) { + t.Parallel() + + fc := &fakeCore{ + tools: []vmcp.Tool{{Name: "cap-tool", Description: "a capability test tool"}}, + resources: []vmcp.Resource{{Name: "cap-doc", URI: "file:///cap.txt"}}, + } + _, _, baseURL := registerServeSession(t, fc) + + initResp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": map[string]any{ + "protocolVersion": "2025-06-18", + "capabilities": map[string]any{}, + "clientInfo": map[string]any{"name": "test", "version": "1.0"}, + }, + }, "") + defer initResp.Body.Close() + require.Equal(t, 200, initResp.StatusCode, "initialize should succeed") + + env, _ := readServeJSONRPC(t, initResp) + result, ok := env["result"].(map[string]any) + require.True(t, ok, "initialize response must have a result object; env: %v", env) + + capabilities, ok := result["capabilities"].(map[string]any) + require.True(t, ok, "result.capabilities must be present; result: %v", result) + + // As of toolhive-core v0.0.28 the go-sdk bridge advertises tools and + // resources capabilities in the initialize response on the Serve path + // (alongside logging). This pins that behavior so a future regression is a + // deliberate, visible flip. + assert.Contains(t, capabilities, "logging", + "logging capability must be advertised; got %v", capabilities) + assert.Contains(t, capabilities, "tools", + "tools capability must be advertised in initialize on the Serve path; got %v", capabilities) + assert.Contains(t, capabilities, "resources", + "resources capability must be advertised in initialize on the Serve path; got %v", capabilities) +} diff --git a/pkg/vmcp/server/context_isolation_regression_test.go b/pkg/vmcp/server/context_isolation_regression_test.go new file mode 100644 index 0000000000..07c0e4a3c5 --- /dev/null +++ b/pkg/vmcp/server/context_isolation_regression_test.go @@ -0,0 +1,255 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" + "github.com/stacklok/toolhive/pkg/audit" + "github.com/stacklok/toolhive/pkg/auth" + transportsession "github.com/stacklok/toolhive/pkg/transport/session" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/server/sessionmanager" + vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" + sessionfactorymocks "github.com/stacklok/toolhive/pkg/vmcp/session/mocks" + sessionmocks "github.com/stacklok/toolhive/pkg/vmcp/session/types/mocks" +) + +// perToolCore is a fakeCore variant whose CallTool returns a result text equal +// to the called tool's name. This lets the context-isolation test detect +// cross-contamination: each concurrent call must observe its own tool name. +type perToolCore struct { + fakeCore +} + +func (p *perToolCore) CallTool( + _ context.Context, _ *auth.Identity, name string, _ map[string]any, _ map[string]any, +) (*vmcp.ToolCallResult, error) { + p.callToolCalls.Add(1) + p.lastCallToolName.Store(name) + return &vmcp.ToolCallResult{ + Content: []vmcp.Content{{Type: vmcp.ContentTypeText, Text: name}}, + }, nil +} + +// newPerToolSessionFactory mirrors newToolSessionFactory but wires a per-tool +// CallTool result on the mock MultiSession so the SDK tool handler routes +// through perToolCore (whose result carries the tool name). The mock session's +// CallTool is only a fallback; the Serve path routes through coreToolHandler → +// core.CallTool, so perToolCore.CallTool is the one that runs. +func newPerToolSessionFactory( + t *testing.T, ctrl *gomock.Controller, tools []vmcp.Tool, +) *sessionfactorymocks.MockMultiSessionFactory { + t.Helper() + factory := sessionfactorymocks.NewMockMultiSessionFactory(ctrl) + factory.EXPECT().MakeSessionWithID(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, id string, _ *auth.Identity, _ []*vmcp.Backend) (vmcpsession.MultiSession, error) { + mock := sessionmocks.NewMockMultiSession(ctrl) + mock.EXPECT().ID().Return(id).AnyTimes() + mock.EXPECT().UpdatedAt().Return(time.Time{}).AnyTimes() + mock.EXPECT().CreatedAt().Return(time.Time{}).AnyTimes() + mock.EXPECT().Type().Return(transportsession.SessionType("")).AnyTimes() + mock.EXPECT().GetData().Return(nil).AnyTimes() + mock.EXPECT().SetData(gomock.Any()).AnyTimes() + mock.EXPECT().GetMetadata().Return(map[string]string{ + vmcpsession.MetadataKeyIdentityBinding: "unauthenticated", + }).AnyTimes() + mock.EXPECT().GetMetadataValue(vmcpsession.MetadataKeyIdentityBinding). + Return("unauthenticated", true).AnyTimes() + mock.EXPECT().SetMetadata(gomock.Any(), gomock.Any()).AnyTimes() + toolsCopy := make([]vmcp.Tool, len(tools)) + copy(toolsCopy, tools) + mock.EXPECT().Tools().Return(toolsCopy).AnyTimes() + mock.EXPECT().AllTools().Return(toolsCopy).AnyTimes() + mock.EXPECT().Resources().Return(nil).AnyTimes() + mock.EXPECT().Prompts().Return(nil).AnyTimes() + mock.EXPECT().BackendSessions().Return(nil).AnyTimes() + rt := &vmcp.RoutingTable{Tools: make(map[string]*vmcp.BackendTarget, len(tools))} + for _, tool := range tools { + rt.Tools[tool.Name] = &vmcp.BackendTarget{WorkloadID: tool.Name} + } + mock.EXPECT().GetRoutingTable().Return(rt).AnyTimes() + mock.EXPECT().ReadResource(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() + mock.EXPECT().GetPrompt(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() + mock.EXPECT().CallTool(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + Return(&vmcp.ToolCallResult{Content: []vmcp.Content{{Type: "text", Text: "ok"}}}, nil).AnyTimes() + mock.EXPECT().Close().Return(nil).AnyTimes() + return mock, nil + }).AnyTimes() + return factory +} + +// TestRegression_ConcurrentToolCalls_NoAuditBleed fires two tools/call POSTs +// concurrently against the same session and asserts each response carries its +// OWN tool's result text — no cross-contamination. The Serve path sets the +// audit BackendInfo per-request via the handler closure (coreToolHandler writes +// the pre-resolved backend name), so concurrent calls must not bleed a backend +// name or result from one request into another. +// +// The audit BackendInfo is attached per-request by the audit middleware (a fresh +// *BackendInfo per request, see pkg/audit/auditor.go Middleware), and the Serve +// handler writes to it from the per-request context — so the isolation surface +// under test is the result text and the observed backend label. We assert each +// call's result text matches the tool it invoked. +func TestRegression_ConcurrentToolCalls_NoAuditBleed(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + tools := []vmcp.Tool{{Name: "tool-a"}, {Name: "tool-b"}} + fc := &perToolCore{fakeCore: fakeCore{tools: tools}} + factory := newPerToolSessionFactory(t, ctrl, tools) + + srv, err := Serve(context.Background(), fc, &ServerConfig{ + SessionTTL: time.Minute, + SessionManagerConfig: &sessionmanager.FactoryConfig{Base: factory}, + BackendRegistry: vmcp.NewImmutableRegistry([]vmcp.Backend{}), + }) + require.NoError(t, err) + t.Cleanup(func() { _ = srv.Stop(context.Background()) }) + + streamable := server.NewStreamableHTTPServer( + srv.mcpServer, + server.WithEndpointPath("/mcp"), + server.WithSessionIdManager(srv.vmcpSessionMgr), + ) + ts := httptest.NewServer(streamable) + t.Cleanup(ts.Close) + baseURL := ts.URL + + initResp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": map[string]any{ + "protocolVersion": "2025-06-18", + "capabilities": map[string]any{}, + "clientInfo": map[string]any{"name": "test", "version": "1.0"}, + }, + }, "") + defer initResp.Body.Close() + require.Equal(t, http.StatusOK, initResp.StatusCode) + sessionID := initResp.Header.Get("Mcp-Session-Id") + require.NotEmpty(t, sessionID) + require.Eventually(t, func() bool { + _, ok := srv.vmcpSessionMgr.GetMultiSession(context.Background(), sessionID) + return ok + }, 2*time.Second, 10*time.Millisecond, "session should be registered") + + type callOutcome struct { + toolName string + body string + err error + } + + var wg sync.WaitGroup + outcomes := make([]callOutcome, len(tools)) + wg.Add(len(tools)) + for i, tool := range tools { + i, tool := i, tool + go func() { + defer wg.Done() + // doServeMCP returns an error instead of calling require: FailNow from a + // worker goroutine only runs Goexit off the test goroutine and misreports. + // All assertions happen on the test goroutine after wg.Wait() below. + resp, doErr := doServeMCP(baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 2 + i, + "method": "tools/call", + "params": map[string]any{ + "name": tool.Name, + "arguments": map[string]any{}, + }, + }, sessionID) + if doErr != nil { + outcomes[i] = callOutcome{toolName: tool.Name, err: doErr} + return + } + defer resp.Body.Close() + body, readErr := io.ReadAll(resp.Body) + outcomes[i] = callOutcome{toolName: tool.Name, body: string(body), err: readErr} + }() + } + + // Fail-fast wait: never block indefinitely on a WaitGroup. + done := make(chan struct{}) + go func() { wg.Wait(); close(done) }() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("timeout waiting for concurrent tool calls to complete") + } + + // Each call must have returned its OWN tool name in the result text, proving + // no result cross-contamination between the concurrent requests. + for _, o := range outcomes { + require.NoError(t, o.err) + assert.Contains(t, o.body, o.toolName, + "concurrent call for %q must carry its own tool's result; got %s", o.toolName, o.body) + } + + // The core must have been reached exactly once per tool (2 total) — no + // duplicated or dropped calls. + assert.Equal(t, int32(len(tools)), fc.callToolCalls.Load(), + "core.CallTool must be called exactly once per concurrent request") + + // Verify coreToolHandler labels the request-scoped audit BackendInfo with the + // backend name it was constructed with. Each goroutine owns a freshly-allocated + // BackendInfo and passes its own backend name, so this is NOT a shared-state + // isolation test — it asserts the handler copies the resolved backend name into + // the per-request BackendInfo (rather than dropping or swapping it). Running it + // concurrently additionally lets the race detector flag any accidental sharing + // of the labelling surface across invocations. + type labelOutcome struct { + want string + got string + hasRes bool + err error + } + var bgWG sync.WaitGroup + bgWG.Add(len(tools)) + labelOutcomes := make([]labelOutcome, len(tools)) + for i, tool := range tools { + i, tool := i, tool + go func() { + defer bgWG.Done() + // No require/assert here: collect results and assert on the test + // goroutine after Wait, since FailNow is only safe there. + bi := &audit.BackendInfo{} + ctx := audit.WithBackendInfo(context.Background(), bi) + req := mcp.CallToolRequest{Params: mcp.CallToolParams{Name: tool.Name, Arguments: map[string]any{}}} + res, err := srv.coreToolHandler(sessionID, tool.Name, tool.Name)(ctx, req) + labelOutcomes[i] = labelOutcome{want: tool.Name, got: bi.BackendName, hasRes: res != nil, err: err} + }() + } + bgDone := make(chan struct{}) + go func() { bgWG.Wait(); close(bgDone) }() + select { + case <-bgDone: + case <-time.After(10 * time.Second): + t.Fatal("timeout waiting for backend-label goroutines") + } + + // Assert on the test goroutine (require/FailNow is only safe here). Each + // handler invocation must have labelled its own BackendInfo with the backend + // name it was given. + for _, o := range labelOutcomes { + require.NoError(t, o.err) + require.True(t, o.hasRes, "coreToolHandler must return a non-nil result") + assert.Equal(t, o.want, o.got, + "coreToolHandler must copy the resolved backend name %q into the request-scoped audit BackendInfo", + o.want) + } +} diff --git a/pkg/vmcp/server/derive.go b/pkg/vmcp/server/derive.go index a4a5c059ef..a1799babc7 100644 --- a/pkg/vmcp/server/derive.go +++ b/pkg/vmcp/server/derive.go @@ -76,6 +76,7 @@ func deriveServerConfig( Port: cfg.Port, // 0 means "OS-assigned". EndpointPath: cfg.EndpointPath, SessionTTL: cfg.SessionTTL, + HeartbeatInterval: cfg.HeartbeatInterval, AuthMiddleware: cfg.AuthMiddleware, AuthInfoHandler: cfg.AuthInfoHandler, PassthroughHeaders: cfg.PassthroughHeaders, diff --git a/pkg/vmcp/server/derive_test.go b/pkg/vmcp/server/derive_test.go index 26abece62e..645d8cddbc 100644 --- a/pkg/vmcp/server/derive_test.go +++ b/pkg/vmcp/server/derive_test.go @@ -39,6 +39,7 @@ func populatedLegacyConfig() *Config { Port: 7777, EndpointPath: "/custom", SessionTTL: 17 * time.Minute, + HeartbeatInterval: 5 * time.Second, AuthMiddleware: passthrough, AuthzMiddleware: passthrough, AuthInfoHandler: http.NewServeMux(), @@ -70,6 +71,7 @@ func TestDeriveServerConfigProjectsTransportFields(t *testing.T) { assert.Equal(t, 7777, got.Port) assert.Equal(t, "/custom", got.EndpointPath) assert.Equal(t, 17*time.Minute, got.SessionTTL) + assert.Equal(t, 5*time.Second, got.HeartbeatInterval) assert.Equal(t, 11*time.Second, got.StatusReportingInterval) // Func/handler/pointer fields projected by reference. diff --git a/pkg/vmcp/server/schema_fidelity_regression_test.go b/pkg/vmcp/server/schema_fidelity_regression_test.go new file mode 100644 index 0000000000..97c42723b8 --- /dev/null +++ b/pkg/vmcp/server/schema_fidelity_regression_test.go @@ -0,0 +1,146 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// TestRegression_ToolSchemaWithTypeObject_ProjectedIntact guards schema +// fidelity for the SUPPORTED case: a tool whose InputSchema includes a +// top-level "type":"object" is projected to the client UNCHANGED in tools/list, +// with its properties and required array preserved. The Serve path +// (coreSessionTools in serve_handlers.go) marshals the core's InputSchema into +// RawInputSchema; the go-sdk bridge's normalizeObjectSchema passes object-typed +// schemas through verbatim, so properties/required survive. +// +// All assertions use encoding/json and generic maps — not SDK types — so the +// check is against the wire representation the client sees. +func TestRegression_ToolSchemaWithTypeObject_ProjectedIntact(t *testing.T) { + t.Parallel() + + originalSchema := map[string]any{ + "type": "object", + "properties": map[string]any{ + "x": map[string]any{"type": "string"}, + }, + "required": []any{"x"}, + } + fc := &fakeCore{tools: []vmcp.Tool{{ + Name: "schema-tool", + Description: "a schema-fidelity test tool", + InputSchema: originalSchema, + }}} + _, sessionID, baseURL := registerServeSession(t, fc) + + resp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/list", + "params": map[string]any{}, + }, sessionID) + defer resp.Body.Close() + require.Equal(t, 200, resp.StatusCode, "tools/list should succeed") + + env, body := readServeJSONRPC(t, resp) + result, ok := env["result"].(map[string]any) + require.True(t, ok, "tools/list must have a result; body: %s", string(body)) + + tools, ok := result["tools"].([]any) + require.True(t, ok, "result.tools must be an array; result: %v", result) + + // Find the projected tool by name. + var found map[string]any + for _, tRaw := range tools { + tm, ok := tRaw.(map[string]any) + if !ok { + continue + } + if tm["name"] == "schema-tool" { + found = tm + break + } + } + require.NotNil(t, found, "schema-tool must be present in tools/list; tools: %v", tools) + + inputSchema, ok := found["inputSchema"].(map[string]any) + require.True(t, ok, "inputSchema must be an object; tool: %v", found) + + // The projected schema must equal the original map exactly: properties and + // required preserved, type unchanged. + assert.Equal(t, originalSchema, inputSchema, + "an object-typed inputSchema must be projected intact; got %v", inputSchema) + assert.Equal(t, "object", inputSchema["type"], + "the top-level type must remain \"object\"; got %v", inputSchema) +} + +// TestRegression_ToolSchemaWithoutTypeObject_ProjectedIntact guards the go-sdk +// bridge's handling of a schema that omits the top-level "type":"object". +// Earlier mcpcompat releases dropped properties/required for such schemas (a +// fidelity gap versus mcp-go); as of toolhive-core v0.0.28 normalizeObjectSchema +// preserves them and supplies the missing "type":"object". This test now pins +// that intact projection so any future regression is a deliberate, visible flip. +func TestRegression_ToolSchemaWithoutTypeObject_ProjectedIntact(t *testing.T) { + t.Parallel() + + originalSchema := map[string]any{ + "properties": map[string]any{ + "x": map[string]any{"type": "string"}, + }, + "required": []any{"x"}, + // NOTE: deliberately NO top-level "type": "object". + } + fc := &fakeCore{tools: []vmcp.Tool{{ + Name: "schema-tool", + Description: "a schema-fidelity test tool", + InputSchema: originalSchema, + }}} + _, sessionID, baseURL := registerServeSession(t, fc) + + resp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/list", + "params": map[string]any{}, + }, sessionID) + defer resp.Body.Close() + require.Equal(t, 200, resp.StatusCode, "tools/list should succeed") + + env, body := readServeJSONRPC(t, resp) + result, ok := env["result"].(map[string]any) + require.True(t, ok, "tools/list must have a result; body: %s", string(body)) + + tools, ok := result["tools"].([]any) + require.True(t, ok, "result.tools must be an array; result: %v", result) + + var found map[string]any + for _, tRaw := range tools { + tm, ok := tRaw.(map[string]any) + if !ok { + continue + } + if tm["name"] == "schema-tool" { + found = tm + break + } + } + require.NotNil(t, found, "schema-tool must be present in tools/list; tools: %v", tools) + + inputSchema, ok := found["inputSchema"].(map[string]any) + require.True(t, ok, "inputSchema must be an object; tool: %v", found) + + // The schema is projected intact: properties and required are preserved and + // the missing top-level "type":"object" is supplied (toolhive-core v0.0.28). + assert.Equal(t, "object", inputSchema["type"], + "a missing top-level type must be supplied as \"object\"; got %v", inputSchema) + assert.Equal(t, originalSchema["properties"], inputSchema["properties"], + "properties must be preserved; got %v", inputSchema) + assert.Equal(t, originalSchema["required"], inputSchema["required"], + "required must be preserved; got %v", inputSchema) +} diff --git a/pkg/vmcp/server/sdk_elicitation_adapter.go b/pkg/vmcp/server/sdk_elicitation_adapter.go index 3cf0eb7ebd..003aa00ea9 100644 --- a/pkg/vmcp/server/sdk_elicitation_adapter.go +++ b/pkg/vmcp/server/sdk_elicitation_adapter.go @@ -9,13 +9,12 @@ import ( "context" "maps" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" - + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/vmcp" ) -// sdkElicitationAdapter wraps mark3labs MCPServer to implement vmcp.ElicitationRequester. +// sdkElicitationAdapter wraps mcpcompat MCPServer to implement vmcp.ElicitationRequester. // // It is the sole point where mcp-go elicitation types appear: it translates the // domain ElicitationRequest/ElicitationResult to/from the SDK types, keeping the @@ -24,22 +23,22 @@ import ( // Per MCP 2025-06-18 spec: The SDK handles JSON-RPC ID correlation internally, // so the adapter does not manage IDs. // -// Thread-safety: Safe for concurrent calls. The mark3labs MCPServer is thread-safe. +// Thread-safety: Safe for concurrent calls. The mcpcompat MCPServer is thread-safe. type sdkElicitationAdapter struct { - // mcpServer is the mark3labs SDK server instance that handles elicitation protocol. + // mcpServer is the mcpcompat SDK server instance that handles elicitation protocol. // Typed as the minimal mcpElicitationRequester seam so the translation logic // can be unit-tested against a fake SDK; *server.MCPServer satisfies it. mcpServer mcpElicitationRequester } -// mcpElicitationRequester is the minimal slice of the mark3labs SDK that the +// mcpElicitationRequester is the minimal slice of the mcpcompat SDK that the // adapter depends on. *server.MCPServer satisfies it in production; tests // substitute a fake to verify domain ⇄ mcp-go translation without a live session. type mcpElicitationRequester interface { RequestElicitation(ctx context.Context, request mcp.ElicitationRequest) (*mcp.ElicitationResult, error) } -// NewSDKElicitationAdapter creates a new elicitation adapter that wraps the mark3labs SDK server. +// NewSDKElicitationAdapter creates a new elicitation adapter that wraps the mcpcompat SDK server. // // The returned adapter implements vmcp.ElicitationRequester by translating the domain // request to/from mcp-go types and delegating to the SDK's RequestElicitation method. @@ -59,11 +58,11 @@ func NewSDKElicitationAdapter(mcpServer *server.MCPServer) vmcp.ElicitationReque } // RequestElicitation translates the domain request to mcp-go, delegates to the -// mark3labs SDK's RequestElicitation method, and translates the response back. +// mcpcompat SDK's RequestElicitation method, and translates the response back. // // This is a synchronous blocking call that: // 1. Maps the domain ElicitationRequest to an mcp.ElicitationRequest -// 2. Forwards the request to the mark3labs SDK +// 2. Forwards the request to the mcpcompat SDK // 3. Blocks until the client responds or timeout occurs // 4. Maps the SDK's mcp.ElicitationResult back to the domain ElicitationResult // @@ -95,7 +94,7 @@ func (a *sdkElicitationAdapter) RequestElicitation( mcpReq.Params.Meta = mcp.NewMetaFromMap(maps.Clone(req.Meta)) } - // Delegate to the mark3labs SDK's RequestElicitation method. + // Delegate to the mcpcompat SDK's RequestElicitation method. // The SDK will: // 1. Extract session ID from context (set by SDK middleware) // 2. Generate JSON-RPC ID for the request diff --git a/pkg/vmcp/server/sdk_elicitation_adapter_test.go b/pkg/vmcp/server/sdk_elicitation_adapter_test.go index 8c28877593..587d989005 100644 --- a/pkg/vmcp/server/sdk_elicitation_adapter_test.go +++ b/pkg/vmcp/server/sdk_elicitation_adapter_test.go @@ -8,11 +8,11 @@ import ( "errors" "testing" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/vmcp" ) @@ -179,7 +179,7 @@ func TestSDKElicitationAdapter_Integration(t *testing.T) { } // TestServer_MCPServer_ReturnsSameInstance verifies that (*Server).MCPServer -// returns the exact mark3labs server pointer stored at construction time. +// returns the exact mcpcompat server pointer stored at construction time. // Identity matters because ClientSession correlation is keyed to the server // that received the initialize request; embedders building their own // elicitation requester must receive the authoritative instance. diff --git a/pkg/vmcp/server/serve.go b/pkg/vmcp/server/serve.go index 6fd16322c7..c9668c15e6 100644 --- a/pkg/vmcp/server/serve.go +++ b/pkg/vmcp/server/serve.go @@ -9,9 +9,8 @@ import ( "net/http" "time" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" - + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/audit" asrunner "github.com/stacklok/toolhive/pkg/authserver/runner" "github.com/stacklok/toolhive/pkg/telemetry" @@ -63,6 +62,10 @@ type ServerConfig struct { // SessionTTL is the session time-to-live duration (default: 30 minutes). SessionTTL time.Duration + // HeartbeatInterval configures the SSE keep-alive ping interval on GET + // connections (default: 30s when zero). + HeartbeatInterval time.Duration + // AuthMiddleware is the optional authentication middleware applied to MCP routes. // If nil, no authentication is required. AuthMiddleware func(http.Handler) http.Handler @@ -327,6 +330,7 @@ func buildServeConfig(cfg *ServerConfig) *Config { Port: cfg.Port, EndpointPath: cfg.EndpointPath, SessionTTL: cfg.SessionTTL, + HeartbeatInterval: cfg.HeartbeatInterval, AuthMiddleware: cfg.AuthMiddleware, AuthInfoHandler: cfg.AuthInfoHandler, PassthroughHeaders: cfg.PassthroughHeaders, diff --git a/pkg/vmcp/server/serve_handlers.go b/pkg/vmcp/server/serve_handlers.go index 384cec96b5..d8d4a2aacf 100644 --- a/pkg/vmcp/server/serve_handlers.go +++ b/pkg/vmcp/server/serve_handlers.go @@ -10,9 +10,8 @@ import ( "fmt" "log/slog" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" - + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/audit" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/vmcp" diff --git a/pkg/vmcp/server/serve_optimizer.go b/pkg/vmcp/server/serve_optimizer.go index f00ad34f5d..0f5a6cae00 100644 --- a/pkg/vmcp/server/serve_optimizer.go +++ b/pkg/vmcp/server/serve_optimizer.go @@ -9,9 +9,8 @@ import ( "fmt" "log/slog" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" - + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/conversion" diff --git a/pkg/vmcp/server/serve_optimizer_test.go b/pkg/vmcp/server/serve_optimizer_test.go index 1aa513e25e..50ecbf3376 100644 --- a/pkg/vmcp/server/serve_optimizer_test.go +++ b/pkg/vmcp/server/serve_optimizer_test.go @@ -13,12 +13,12 @@ import ( "testing" "time" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/ratelimit" "github.com/stacklok/toolhive/pkg/vmcp" diff --git a/pkg/vmcp/server/serve_session_test.go b/pkg/vmcp/server/serve_session_test.go index d19c68cf89..8e2e47bf34 100644 --- a/pkg/vmcp/server/serve_session_test.go +++ b/pkg/vmcp/server/serve_session_test.go @@ -16,12 +16,12 @@ import ( "testing" "time" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/audit" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/ratelimit" @@ -927,18 +927,30 @@ func TestServeOmitsPrompts(t *testing.T) { // package-server-internal analogue of the postMCP helper in the external suite. func postServeMCP(t *testing.T, baseURL string, body map[string]any, sessionID string) *http.Response { t.Helper() - rawBody, err := json.Marshal(body) + resp, err := doServeMCP(baseURL, body, sessionID) require.NoError(t, err) + return resp +} + +// doServeMCP performs a tools POST and returns the response or an error. It never +// calls require/assert, so it is safe to invoke from worker goroutines where +// FailNow/Goexit would run off the test goroutine and misreport. postServeMCP +// wraps it for the common test-goroutine case. +func doServeMCP(baseURL string, body map[string]any, sessionID string) (*http.Response, error) { + rawBody, err := json.Marshal(body) + if err != nil { + return nil, err + } req, err := http.NewRequestWithContext(context.Background(), http.MethodPost, baseURL+"/mcp", bytes.NewReader(rawBody)) - require.NoError(t, err) + if err != nil { + return nil, err + } req.Header.Set("Content-Type", "application/json") req.Header.Set("Accept", "application/json, text/event-stream") if sessionID != "" { req.Header.Set("Mcp-Session-Id", sessionID) } - resp, err := http.DefaultClient.Do(req) - require.NoError(t, err) - return resp + return http.DefaultClient.Do(req) } diff --git a/pkg/vmcp/server/serve_test.go b/pkg/vmcp/server/serve_test.go index 699601b123..07bce46cb7 100644 --- a/pkg/vmcp/server/serve_test.go +++ b/pkg/vmcp/server/serve_test.go @@ -345,6 +345,7 @@ func TestBuildServeConfigMapsSharedFields(t *testing.T) { Port: 1, EndpointPath: "/e", SessionTTL: time.Second, + HeartbeatInterval: time.Second, AuthMiddleware: func(h http.Handler) http.Handler { return h }, AuthInfoHandler: http.NewServeMux(), PassthroughHeaders: []string{"x-test"}, diff --git a/pkg/vmcp/server/server.go b/pkg/vmcp/server/server.go index 8e1e5d3179..4caa46cc77 100644 --- a/pkg/vmcp/server/server.go +++ b/pkg/vmcp/server/server.go @@ -21,8 +21,7 @@ import ( "sync" "time" - "github.com/mark3labs/mcp-go/server" - + "github.com/stacklok/toolhive-core/mcpcompat/server" tcredis "github.com/stacklok/toolhive-core/redis" "github.com/stacklok/toolhive/pkg/audit" "github.com/stacklok/toolhive/pkg/auth" @@ -96,6 +95,15 @@ const ( capabilityCacheTTL = 30 * time.Second ) +// heartbeatInterval returns the configured heartbeat interval, or the default +// when the configured value is zero or negative (unset or invalid). +func heartbeatInterval(d time.Duration) time.Duration { + if d <= 0 { + return defaultHeartbeatInterval + } + return d +} + //go:generate mockgen -destination=mocks/mock_watcher.go -package=mocks -source=server.go Watcher // Watcher is the interface for Kubernetes backend watcher integration. @@ -132,6 +140,17 @@ type Config struct { // Sessions inactive for this duration will be automatically cleaned up SessionTTL time.Duration + // HeartbeatInterval configures the SSE keep-alive ping interval on GET + // connections. When zero, the Handler defaults to defaultHeartbeatInterval (30s). + // Prevents proxies/load balancers from closing idle SSE connections. + // + // This is intentionally plumbed end-to-end (ServerConfig → Config → Handler → + // WithHeartbeatInterval) ahead of any CLI-flag or CRD wiring. No external entry + // point sets it yet, so in practice it is always the default unless an embedder + // assigns it programmatically; the plumbing exists so surfacing a flag/field + // later is a one-line change rather than a re-thread through the server. + HeartbeatInterval time.Duration + // AuthMiddleware is the optional authentication middleware to apply to MCP routes. // If nil, no authentication is required. // This should be a composed middleware chain (e.g., TokenValidator + MCP parser). @@ -265,7 +284,7 @@ type Server struct { // session manager; this is the resolved factory surfaced via Manager.OptimizerFactory. optimizerFactory func(context.Context, []server.ServerTool) (optimizer.Optimizer, error) - // MCP protocol server (mark3labs/mcp-go) + // MCP protocol server (stacklok/toolhive-core/mcpcompat) mcpServer *server.MCPServer // HTTP server for Streamable HTTP transport @@ -502,7 +521,7 @@ func (s *Server) Handler(_ context.Context) (http.Handler, error) { s.mcpServer, server.WithEndpointPath(s.config.EndpointPath), server.WithSessionIdManager(s.vmcpSessionMgr), - server.WithHeartbeatInterval(defaultHeartbeatInterval), + server.WithHeartbeatInterval(heartbeatInterval(s.config.HeartbeatInterval)), ) // Create HTTP mux with separated authenticated and unauthenticated routes @@ -907,7 +926,7 @@ func (s *Server) SessionManager() *transportsession.Manager { return s.sessionManager } -// MCPServer returns the underlying mark3labs *server.MCPServer instance +// MCPServer returns the underlying mcpcompat *server.MCPServer instance // servicing this vMCP server's /mcp endpoint. // // Intended for embedders that wrap the vMCP composer in their own pipeline and @@ -918,7 +937,7 @@ func (s *Server) SessionManager() *transportsession.Manager { // // Trust boundary: this accessor is in-process only; the returned pointer is // the same instance for the lifetime of the Server and is safe for concurrent -// use per mark3labs guarantees. +// use per mcpcompat guarantees. // // Safe operations include RequestElicitation against an active session, // registering observability hooks, and reading registered diff --git a/pkg/vmcp/server/session_lifecycle_regression_test.go b/pkg/vmcp/server/session_lifecycle_regression_test.go new file mode 100644 index 0000000000..22c75e2924 --- /dev/null +++ b/pkg/vmcp/server/session_lifecycle_regression_test.go @@ -0,0 +1,155 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "encoding/json" + "io" + "net/http" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive/pkg/auth" + "github.com/stacklok/toolhive/pkg/vmcp" +) + +// readServeJSONRPC reads an HTTP response body from the Serve-path streamable +// server and unmarshals the JSON-RPC envelope into a map. The streamable server +// is configured with JSONResponse:true, so the body is a single application/json +// document (not an SSE stream). It also returns the raw body bytes for +// substring assertions. +func readServeJSONRPC(t *testing.T, resp *http.Response) (map[string]any, []byte) { + t.Helper() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + var env map[string]any + require.NoError(t, json.Unmarshal(body, &env), "body: %s", string(body)) + return env, body +} + +// serveDelete sends an HTTP DELETE to /mcp with the given session ID. +func serveDelete(t *testing.T, baseURL, sessionID string) *http.Response { + t.Helper() + req, err := http.NewRequestWithContext( + context.Background(), http.MethodDelete, baseURL+"/mcp", nil, + ) + require.NoError(t, err) + if sessionID != "" { + req.Header.Set("Mcp-Session-Id", sessionID) + } + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + return resp +} + +// TestRegression_TerminatedSessionRejected verifies that a session terminated +// via the SDK's DELETE path (which drives vmcpSessionMgr.Terminate under the +// hood and forgets the SDK's local session) is no longer accepted for +// subsequent requests: a tools/list POST with the dead session ID must yield a +// 4xx response — the SDK's Validate rejects the unknown/terminated session +// before any handler runs — not a successful JSON-RPC result. +// +// Note: calling vmcpSessionMgr.Terminate alone does NOT immediately evict the +// SDK's in-memory local session (lazy eviction — the entry is closed when the +// cache is later re-checked). Only the DELETE path forgets the SDK session in +// lockstep with storage termination, so the rejection is observable on the +// next request. This mirrors the proven assertion in +// TestIntegration_SessionManagement_Termination. +func TestRegression_TerminatedSessionRejected(t *testing.T) { + t.Parallel() + + fc := &fakeCore{tools: []vmcp.Tool{{Name: "t"}}} + srv, sessionID, baseURL := registerServeSession(t, fc) + + // Terminate via the SDK DELETE path: this calls vmcpSessionMgr.Terminate + // (storage delete) AND forgets the SDK's local session, so the next request + // is rejected by Validate rather than served from a stale local copy. + delResp := serveDelete(t, baseURL, sessionID) + defer delResp.Body.Close() + require.Equal(t, http.StatusOK, delResp.StatusCode, "DELETE should return 200") + + // The termination must be reflected in storage before we assert. + require.Eventually(t, func() bool { + _, ok := srv.vmcpSessionMgr.GetMultiSession(context.Background(), sessionID) + return !ok + }, 2*time.Second, 10*time.Millisecond, "session must be gone from the manager after Terminate") + + resp := postServeMCP(t, baseURL, map[string]any{ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/list", + "params": map[string]any{}, + }, sessionID) + defer resp.Body.Close() + + // A terminated session must NOT be accepted: reject with a 4xx (the SDK's + // Validate returns 404 for an unknown/terminated session). A 200 here would + // mean a terminated session was still serving requests — a security regression. + assert.GreaterOrEqual(t, resp.StatusCode, 400, + "terminated session must be rejected, got status %d", resp.StatusCode) + assert.Less(t, resp.StatusCode, 500, + "rejection should be a client error (4xx), not a server error; got %d", resp.StatusCode) +} + +// TestRegression_DELETE_TerminationEvictsCache verifies that an HTTP DELETE to +// /mcp with a valid Mcp-Session-Id terminates the session: the response is 200 +// (the streamable server rewrites the SDK's 204 to 200 for mcp-go compatibility) +// and the session is subsequently absent from vmcpSessionMgr. +func TestRegression_DELETE_TerminationEvictsCache(t *testing.T) { + t.Parallel() + + fc := &fakeCore{tools: []vmcp.Tool{{Name: "t"}}} + srv, sessionID, baseURL := registerServeSession(t, fc) + + delResp := serveDelete(t, baseURL, sessionID) + defer delResp.Body.Close() + + // The streamable server rewrites 204 → 200 and forwards the termination to + // the SessionIdManager, so DELETE returns 200 OK on the happy path. + assert.Equal(t, http.StatusOK, delResp.StatusCode, + "DELETE on a live session should return 200 (got %d)", delResp.StatusCode) + + // The session must be evicted from the vMCP session manager. + require.Eventually(t, func() bool { + _, ok := srv.vmcpSessionMgr.GetMultiSession(context.Background(), sessionID) + return !ok + }, 2*time.Second, 10*time.Millisecond, "session must be evicted from vmcpSessionMgr after DELETE") +} + +// TestRegression_SessionIdentityBinding_SecondPrincipalRejected is a thin +// regression re-assertion that the Serve call path enforces the session's +// identity binding: a session bound to the anonymous identity rejects a caller +// presenting a different principal. The exhaustive test lives in +// TestServeEnforcesSessionBinding; this guards the regression at the tool-handler +// boundary by invoking coreToolHandler directly with a non-anonymous identity. +func TestRegression_SessionIdentityBinding_SecondPrincipalRejected(t *testing.T) { + t.Parallel() + + fc := &fakeCore{tools: []vmcp.Tool{{Name: "t"}}} + srv, sessionID, _ := registerServeSession(t, fc) + + // The registered session is bound to "unauthenticated" (anonymous) by + // newToolSessionFactory. A caller presenting a token is a session-upgrade + // attack and must be rejected before the core is reached. + ctx := auth.WithIdentity(context.Background(), &auth.Identity{Token: "attacker-token"}) + req := mcp.CallToolRequest{Params: mcp.CallToolParams{Name: "t", Arguments: map[string]any{}}} + + res, err := srv.coreToolHandler(sessionID, "t", "")(ctx, req) + require.NoError(t, err) + require.NotNil(t, res) + assert.True(t, res.IsError, "a second principal must be rejected with an error result") + + body, err := json.Marshal(res) + require.NoError(t, err) + assert.Contains(t, string(body), "Unauthorized", + "the rejection must carry an unauthorized message") + // The core must not have been reached on a binding failure. + assert.Equal(t, int32(0), fc.callToolCalls.Load(), + "core.CallTool must not be reached when the binding check fails") +} diff --git a/pkg/vmcp/server/session_management_integration_test.go b/pkg/vmcp/server/session_management_integration_test.go index 96f247f652..61779fbcf2 100644 --- a/pkg/vmcp/server/session_management_integration_test.go +++ b/pkg/vmcp/server/session_management_integration_test.go @@ -15,12 +15,12 @@ import ( "testing" "time" - mcpmcp "github.com/mark3labs/mcp-go/mcp" - mcpsdk "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" + mcpmcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpsdk "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/auth" transportsession "github.com/stacklok/toolhive/pkg/transport/session" "github.com/stacklok/toolhive/pkg/vmcp" diff --git a/pkg/vmcp/server/session_manager_interface.go b/pkg/vmcp/server/session_manager_interface.go index 70c6ab8ebe..547b2e0873 100644 --- a/pkg/vmcp/server/session_manager_interface.go +++ b/pkg/vmcp/server/session_manager_interface.go @@ -6,8 +6,7 @@ package server import ( "context" - mcpserver "github.com/mark3labs/mcp-go/server" - + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" sessiontypes "github.com/stacklok/toolhive/pkg/vmcp/session/types" ) diff --git a/pkg/vmcp/server/sessionmanager/factory.go b/pkg/vmcp/server/sessionmanager/factory.go index fb4bb159c7..38f1747e87 100644 --- a/pkg/vmcp/server/sessionmanager/factory.go +++ b/pkg/vmcp/server/sessionmanager/factory.go @@ -11,13 +11,13 @@ import ( "log/slog" "time" - "github.com/mark3labs/mcp-go/mcp" - mcpserver "github.com/mark3labs/mcp-go/server" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/codes" "go.opentelemetry.io/otel/metric" "go.opentelemetry.io/otel/trace" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/telemetry" "github.com/stacklok/toolhive/pkg/vmcp" diff --git a/pkg/vmcp/server/sessionmanager/horizontal_scaling_integration_test.go b/pkg/vmcp/server/sessionmanager/horizontal_scaling_integration_test.go index 102de1fd1e..961842b9b1 100644 --- a/pkg/vmcp/server/sessionmanager/horizontal_scaling_integration_test.go +++ b/pkg/vmcp/server/sessionmanager/horizontal_scaling_integration_test.go @@ -12,11 +12,11 @@ import ( "time" "github.com/alicebob/miniredis/v2" - mcpmcp "github.com/mark3labs/mcp-go/mcp" - mcpserver "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + mcpmcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" tcredis "github.com/stacklok/toolhive-core/redis" "github.com/stacklok/toolhive/pkg/auth" transportsession "github.com/stacklok/toolhive/pkg/transport/session" diff --git a/pkg/vmcp/server/sessionmanager/optimizer_gate_test.go b/pkg/vmcp/server/sessionmanager/optimizer_gate_test.go index 25bf9c1c8d..7b2bfacd18 100644 --- a/pkg/vmcp/server/sessionmanager/optimizer_gate_test.go +++ b/pkg/vmcp/server/sessionmanager/optimizer_gate_test.go @@ -7,11 +7,11 @@ import ( "context" "testing" - mcpserver "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/vmcp/optimizer" ) diff --git a/pkg/vmcp/server/sessionmanager/session_manager.go b/pkg/vmcp/server/sessionmanager/session_manager.go index 70687b3f9b..2d1161fd32 100644 --- a/pkg/vmcp/server/sessionmanager/session_manager.go +++ b/pkg/vmcp/server/sessionmanager/session_manager.go @@ -21,8 +21,8 @@ import ( "time" "github.com/google/uuid" - mcpserver "github.com/mark3labs/mcp-go/server" + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/cache" transportsession "github.com/stacklok/toolhive/pkg/transport/session" @@ -43,7 +43,7 @@ const ( ) // Manager bridges the domain session lifecycle (MultiSession / MultiSessionFactory) -// to the mark3labs SDK's SessionIdManager interface. +// to the mcpcompat SDK's SessionIdManager interface. // // It implements a two-phase session-creation pattern: // @@ -417,10 +417,25 @@ func (sm *Manager) cleanupFailedPlaceholder(sessionID string, metadata map[strin // Validate implements the SDK's SessionIdManager.Validate(). // -// Returns (isTerminated=true, nil) for explicitly terminated sessions. -// Returns (false, error) for unknown sessions — per the SDK interface contract, -// a lookup failure is signalled via err, not via isTerminated. +// Returns (isTerminated=true, nil) for a session that is definitively gone: +// explicitly terminated (placeholder marked terminated=true) OR absent from +// storage (a full MultiSession is deleted on Terminate, TTL-expired, or never +// existed). All of these must reject the request as a hard termination so the +// client re-initializes rather than retrying a dead session. +// +// Returns (false, error) only for a genuine, transient storage error (e.g. the +// backing store is unreachable) — the caller should treat this as retryable and +// NOT drop session state. +// // Returns (false, nil) for valid, active sessions. +// +// This distinction is load-bearing: the streamable transport maps +// (isTerminated=true) to HTTP 404 (-> client ErrSessionTerminated -> re-init) +// and a non-terminated error to HTTP 503 (retryable). Reporting a genuinely-gone +// session as an error would surface as 503 and make a client whose session was +// terminated on another replica retry the dead session forever. Terminate still +// DELETES the key (unchanged), so the resurrection-race guarantee is preserved; +// only the way an absent key is reported here changes. func (sm *Manager) Validate(sessionID string) (isTerminated bool, err error) { if sessionID == "" { return false, fmt.Errorf("Manager.Validate: empty session ID") @@ -431,8 +446,12 @@ func (sm *Manager) Validate(sessionID string) (isTerminated bool, err error) { metadata, err := sm.storage.Load(ctx, sessionID) if errors.Is(err, transportsession.ErrSessionNotFound) { - slog.Debug("Manager.Validate: session not found", "session_id", sessionID) - return false, fmt.Errorf("session not found") + // The session is gone (terminated + deleted, TTL-expired, or never + // existed). Report it as terminated so the transport answers 404 and the + // client re-initializes, rather than as an error (which the transport + // would treat as a transient 503 and the client would retry indefinitely). + slog.Debug("Manager.Validate: session not found; reporting as terminated", "session_id", sessionID) + return true, nil } if err != nil { return false, fmt.Errorf("Manager.Validate: storage error for session %q: %w", sessionID, err) @@ -455,15 +474,16 @@ func (sm *Manager) Validate(sessionID string) (isTerminated bool, err error) { // - MultiSession (Phase 2): the storage key is deleted. The node-local cache // self-heals on the next Get: checkSession detects ErrSessionNotFound, // evicts the entry, and onEvict closes backend connections. After deletion -// Validate() returns (false, error) — the same response as "never existed". +// Validate() reports the absent key as (isTerminated=true, nil), so the next +// request — on this or any other replica — is rejected with a definitive 404 +// and the client re-initializes instead of retrying the dead session. // // - Placeholder (Phase 1): the session is marked terminated=true and left // for TTL cleanup. This prevents CreateSession() from opening backend // connections for an already-terminated session (see fast-fail check in // CreateSession). The terminated flag also lets Validate() return // (isTerminated=true, nil) during the window between termination and TTL -// expiry, allowing the SDK to distinguish "actively terminated" from -// "never existed". +// expiry — the same terminated response the deleted case now gives. // // Returns (isNotAllowed=false, nil) on success; client termination is always permitted. func (sm *Manager) Terminate(sessionID string) (isNotAllowed bool, err error) { diff --git a/pkg/vmcp/server/sessionmanager/session_manager_test.go b/pkg/vmcp/server/sessionmanager/session_manager_test.go index bfd119785c..b281f5ddf2 100644 --- a/pkg/vmcp/server/sessionmanager/session_manager_test.go +++ b/pkg/vmcp/server/sessionmanager/session_manager_test.go @@ -503,7 +503,7 @@ func TestSessionManager_Validate(t *testing.T) { assert.Contains(t, err.Error(), "empty session ID") }) - t.Run("returns error for unknown session", func(t *testing.T) { + t.Run("reports unknown session as terminated", func(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) @@ -512,10 +512,15 @@ func TestSessionManager_Validate(t *testing.T) { registry := newFakeRegistry() sm, _ := newTestSessionManager(t, factory, registry) + // A session that is absent from storage (deleted on Terminate, TTL-expired, + // or never existed) must be reported as terminated (isTerminated=true, nil), + // NOT as an error. The transport maps terminated -> 404 so the client + // re-initializes; reporting an error would surface as a retryable 503 and a + // client whose session was terminated on another replica would retry the + // dead session forever (regression guarded here). isTerminated, err := sm.Validate("non-existent-id") - require.Error(t, err) - assert.False(t, isTerminated) - assert.Contains(t, err.Error(), "session not found") + require.NoError(t, err) + assert.True(t, isTerminated, "an absent session must report as terminated, not as a transient error") }) t.Run("returns false for active session", func(t *testing.T) { diff --git a/pkg/vmcp/server/sessionmanager/telemetry_test.go b/pkg/vmcp/server/sessionmanager/telemetry_test.go index d32acac96d..332b62d73c 100644 --- a/pkg/vmcp/server/sessionmanager/telemetry_test.go +++ b/pkg/vmcp/server/sessionmanager/telemetry_test.go @@ -8,14 +8,14 @@ import ( "fmt" "testing" - "github.com/mark3labs/mcp-go/mcp" - mcpserver "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" sdkmetric "go.opentelemetry.io/otel/sdk/metric" "go.opentelemetry.io/otel/sdk/metric/metricdata" tracenoop "go.opentelemetry.io/otel/trace/noop" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/vmcp/optimizer" ) diff --git a/pkg/vmcp/server/sse_keepalive_regression_test.go b/pkg/vmcp/server/sse_keepalive_regression_test.go new file mode 100644 index 0000000000..3f61560513 --- /dev/null +++ b/pkg/vmcp/server/sse_keepalive_regression_test.go @@ -0,0 +1,127 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package server + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/server/sessionmanager" +) + +// TestRegression_SSEKeepAlive_PeriodicBytesOnIdleStream verifies that an idle +// SSE stream opened against the Serve-path streamable server keeps producing +// bytes (a keep-alive ping or heartbeat) within a short window. A regression +// that drops the heartbeat wiring will cause this test to fail: no bytes arrive +// on an idle GET stream, so proxies/load balancers eventually close it. +// +// The streamable server mounts the SDK keep-alive over the configured +// HeartbeatInterval, so the test sets a short interval (100ms) and reads with a +// 500ms deadline — generous relative to the interval but tight enough that a +// missing keep-alive is caught rather than hanging the suite. +func TestRegression_SSEKeepAlive_PeriodicBytesOnIdleStream(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + testTool := vmcp.Tool{Name: "keepalive-tool", Description: "keep-alive regression anchor"} + factory, _ := newToolSessionFactory(t, ctrl, []vmcp.Tool{testTool}) + fc := &fakeCore{tools: []vmcp.Tool{testTool}} + + srv, err := Serve(context.Background(), fc, &ServerConfig{ + SessionTTL: time.Minute, + HeartbeatInterval: 100 * time.Millisecond, + SessionManagerConfig: &sessionmanager.FactoryConfig{Base: factory}, + BackendRegistry: vmcp.NewImmutableRegistry([]vmcp.Backend{}), + }) + require.NoError(t, err) + t.Cleanup(func() { _ = srv.Stop(context.Background()) }) + + // Serve through the real Handler wiring rather than a hand-built streamable + // server. This exercises the ServerConfig.HeartbeatInterval → Config → + // Handler → WithHeartbeatInterval path that server.go adds, so removing that + // wiring line in Handler regresses this test (no keep-alive bytes arrive). + handler, err := srv.Handler(context.Background()) + require.NoError(t, err) + ts := httptest.NewServer(handler) + t.Cleanup(ts.Close) + + // initialize → obtain a session ID for the subsequent GET stream. + initResp := postServeMCP(t, ts.URL, map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": map[string]any{ + "protocolVersion": "2025-06-18", + "capabilities": map[string]any{}, + "clientInfo": map[string]any{"name": "keepalive-test", "version": "1.0"}, + }, + }, "") + defer initResp.Body.Close() + require.Equal(t, http.StatusOK, initResp.StatusCode, "initialize should succeed") + sessionID := initResp.Header.Get("Mcp-Session-Id") + require.NotEmpty(t, sessionID, "session ID should be returned in Mcp-Session-Id header") + + // Open a long-lived SSE GET stream. The stream must stay open and emit + // keep-alive bytes even though no client request is in flight. + req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, ts.URL+"/mcp", nil) + require.NoError(t, err) + req.Header.Set("Accept", "text/event-stream") + req.Header.Set("Mcp-Session-Id", sessionID) + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode, "GET stream should open with 200") + + // Read with a 500ms deadline. The keep-alive interval is 100ms, so at least + // one ping/heartbeat should land within the deadline; a missing keep-alive + // surfaces as a read timeout with zero bytes — the regression. The streamable + // server's body is not a net.Conn, so SetReadDeadline is unavailable; drive + // the read on a goroutine and race it against a timer. + type readResult struct { + data []byte + err error + } + ch := make(chan readResult, 1) + go func() { + buf := make([]byte, 4096) + n, err := resp.Body.Read(buf) + ch <- readResult{data: buf[:n], err: err} + }() + + var got []byte + select { + case r := <-ch: + got = r.data + case <-time.After(500 * time.Millisecond): + } + + // At least some bytes must arrive — the keep-alive ping, an SSE comment + // (`:` prefix), or a JSON-RPC ping request frame. Zero bytes within the + // window means the keep-alive wiring is broken; fail loudly as a regression. + if len(got) == 0 { + t.Fatal("no bytes received on idle SSE stream within 500ms; keep-alive/heartbeat " + + "is not emitting periodic bytes (regression: proxies will close idle streams)") + } + + body := string(got) + t.Logf("received keep-alive bytes on idle SSE stream: %q", body) + // Accept any of: a JSON-RPC ping request, an SSE comment, or any SSE data + // frame. The point is that bytes flow; their exact form is shim-defined. + hasPing := strings.Contains(body, `"method":"ping"`) || + strings.Contains(body, `"method": "ping"`) + hasSSEComment := strings.Contains(body, ":") + hasDataFrame := strings.Contains(body, "data:") + assert.True(t, hasPing || hasSSEComment || hasDataFrame, + "keep-alive bytes should be an SSE comment, a data frame, or a JSON-RPC ping; got %q", body) +} diff --git a/pkg/vmcp/server/testutil_test.go b/pkg/vmcp/server/testutil_test.go index 9a29e20210..8f0dfa58a1 100644 --- a/pkg/vmcp/server/testutil_test.go +++ b/pkg/vmcp/server/testutil_test.go @@ -11,8 +11,8 @@ import ( "net/http/httptest" "testing" - mcpmcp "github.com/mark3labs/mcp-go/mcp" - mcpserver "github.com/mark3labs/mcp-go/server" + mcpmcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" ) // startRealMCPBackend creates a real in-process MCP server over streamable-HTTP diff --git a/pkg/vmcp/session/connector_integration_test.go b/pkg/vmcp/session/connector_integration_test.go index c66b1bdeaf..d6548568f3 100644 --- a/pkg/vmcp/session/connector_integration_test.go +++ b/pkg/vmcp/session/connector_integration_test.go @@ -11,11 +11,11 @@ import ( "testing" "github.com/google/uuid" - mcpmcp "github.com/mark3labs/mcp-go/mcp" - mcpserver "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + mcpmcp "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/vmcp" vmcpauth "github.com/stacklok/toolhive/pkg/vmcp/auth" diff --git a/pkg/vmcp/session/internal/backend/mcp_session.go b/pkg/vmcp/session/internal/backend/mcp_session.go index b8b5e3619a..986808ab12 100644 --- a/pkg/vmcp/session/internal/backend/mcp_session.go +++ b/pkg/vmcp/session/internal/backend/mcp_session.go @@ -12,10 +12,9 @@ import ( "net/http" "time" - mcpclient "github.com/mark3labs/mcp-go/client" - mcptransport "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" - + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + mcptransport "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/secrets" "github.com/stacklok/toolhive/pkg/versions" @@ -98,7 +97,7 @@ func (i *identityRoundTripper) RoundTrip(req *http.Request) (*http.Response, err // Compile-time assertion: mcpSession must implement Session. var _ Session = (*mcpSession)(nil) -// mcpSession wraps a persistent mark3labs MCP client for one backend. +// mcpSession wraps a persistent mcpcompat MCP client for one backend. // It is created once per backend during MakeSession and closed when the session ends. // // Phase 1 limitation — no reconnection: if the underlying transport drops @@ -251,7 +250,7 @@ func NewHTTPConnector(registry vmcpauth.OutgoingAuthRegistry) func( // Extract the backend-assigned session ID when the transport supports it. // Streamable-HTTP servers send an Mcp-Session-Id response header during - // Initialize; the mark3labs transport captures it internally and exposes + // Initialize; the mcpcompat transport captures it internally and exposes // it via GetSessionId(). SSE transports do not assign a session ID, so // the field remains empty for those backends. var backendSessionID string @@ -263,7 +262,7 @@ func NewHTTPConnector(registry vmcpauth.OutgoingAuthRegistry) func( } } -// createMCPClient builds and starts a mark3labs MCP client for target. +// createMCPClient builds and starts a mcpcompat MCP client for target. // The transport is started with context.Background() so its lifetime is bound // to client.Close(), not to any caller-supplied init context. // sessionHint, when non-empty, is passed as the initial Mcp-Session-Id for @@ -340,7 +339,7 @@ func createMCPClient( // request/response pair, so a per-response body size limit is safe and // correct. http.Client.Timeout provides a hard wall-clock deadline; // WithHTTPTimeout additionally wraps each SDK request in a - // context.WithTimeout so the mark3labs transport surfaces a descriptive + // context.WithTimeout so the mcpcompat transport surfaces a descriptive // error before the stdlib deadline fires. Both are set to // defaultBackendRequestTimeout: defense-in-depth. sizeLimited := httpRoundTripperFunc(func(req *http.Request) (*http.Response, error) { diff --git a/pkg/vmcp/session/internal/backend/mcp_session_capabilities_test.go b/pkg/vmcp/session/internal/backend/mcp_session_capabilities_test.go index f1c16f8e80..30c85d7c6f 100644 --- a/pkg/vmcp/session/internal/backend/mcp_session_capabilities_test.go +++ b/pkg/vmcp/session/internal/backend/mcp_session_capabilities_test.go @@ -12,12 +12,12 @@ import ( "sync" "testing" - mcpclient "github.com/mark3labs/mcp-go/client" - mcptransport "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + mcptransport "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/vmcp" ) @@ -234,7 +234,7 @@ func (f *fakeBackend) writeError(w http.ResponseWriter, id json.RawMessage, e *j } } -// newTestClient builds a streamable-HTTP mark3labs client pointing at url and +// newTestClient builds a streamable-HTTP mcpcompat client pointing at url and // runs Start() so the transport is ready for initAndQueryCapabilities. Cleanup // is registered via t.Cleanup. func newTestClient(t *testing.T, url string) *mcpclient.Client { diff --git a/pkg/vmcp/session/internal/backend/mcp_session_header_forward_test.go b/pkg/vmcp/session/internal/backend/mcp_session_header_forward_test.go index 9d58de9b76..d9db8e6230 100644 --- a/pkg/vmcp/session/internal/backend/mcp_session_header_forward_test.go +++ b/pkg/vmcp/session/internal/backend/mcp_session_header_forward_test.go @@ -9,10 +9,10 @@ import ( "testing" "time" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/secrets" "github.com/stacklok/toolhive/pkg/vmcp" vmcpauth "github.com/stacklok/toolhive/pkg/vmcp/auth" diff --git a/pkg/vmcp/session/internal/backend/mcp_session_identity_refresh_test.go b/pkg/vmcp/session/internal/backend/mcp_session_identity_refresh_test.go index ac0e8236a6..0a0ee2904e 100644 --- a/pkg/vmcp/session/internal/backend/mcp_session_identity_refresh_test.go +++ b/pkg/vmcp/session/internal/backend/mcp_session_identity_refresh_test.go @@ -8,10 +8,10 @@ import ( "testing" "time" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/vmcp" vmcpauth "github.com/stacklok/toolhive/pkg/vmcp/auth" diff --git a/pkg/vmcp/session/optimizerdec/decorator.go b/pkg/vmcp/session/optimizerdec/decorator.go index b3e4454e88..5c096d956d 100644 --- a/pkg/vmcp/session/optimizerdec/decorator.go +++ b/pkg/vmcp/session/optimizerdec/decorator.go @@ -10,8 +10,7 @@ import ( "encoding/json" "fmt" - "github.com/mark3labs/mcp-go/mcp" - + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/conversion" diff --git a/pkg/vmcp/session/optimizerdec/decorator_test.go b/pkg/vmcp/session/optimizerdec/decorator_test.go index 514c163df4..9166bb554d 100644 --- a/pkg/vmcp/session/optimizerdec/decorator_test.go +++ b/pkg/vmcp/session/optimizerdec/decorator_test.go @@ -9,11 +9,11 @@ import ( "reflect" "testing" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/optimizer" diff --git a/test/e2e/mcp_client_helpers.go b/test/e2e/mcp_client_helpers.go index f08b8a2e10..1b989e05e7 100644 --- a/test/e2e/mcp_client_helpers.go +++ b/test/e2e/mcp_client_helpers.go @@ -11,11 +11,12 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" //nolint:staticcheck // Standard practice for Ginkgo . "github.com/onsi/gomega" //nolint:staticcheck // Standard practice for Gomega + + "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) // MCPClientHelper provides high-level MCP client operations for e2e tests diff --git a/test/e2e/network_isolation_test.go b/test/e2e/network_isolation_test.go index 278a9c113d..7ea78487e5 100644 --- a/test/e2e/network_isolation_test.go +++ b/test/e2e/network_isolation_test.go @@ -15,10 +15,10 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/test/e2e" ) diff --git a/test/e2e/thv-operator/virtualmcp/helpers.go b/test/e2e/thv-operator/virtualmcp/helpers.go index 4f9f67420d..76c807eef8 100644 --- a/test/e2e/thv-operator/virtualmcp/helpers.go +++ b/test/e2e/thv-operator/virtualmcp/helpers.go @@ -21,9 +21,6 @@ import ( "time" "github.com/go-logr/logr" - mcpclient "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" "github.com/onsi/ginkgo/v2" "github.com/onsi/gomega" appsv1 "k8s.io/api/apps/v1" @@ -35,6 +32,9 @@ import ( "k8s.io/utils/ptr" "sigs.k8s.io/controller-runtime/pkg/client" + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/test/e2e/images" "github.com/stacklok/toolhive/test/e2e/thv-operator/testutil" diff --git a/test/e2e/thv-operator/virtualmcp/mcpremoteproxy_scaling_test.go b/test/e2e/thv-operator/virtualmcp/mcpremoteproxy_scaling_test.go index f226d99760..139587b817 100644 --- a/test/e2e/thv-operator/virtualmcp/mcpremoteproxy_scaling_test.go +++ b/test/e2e/thv-operator/virtualmcp/mcpremoteproxy_scaling_test.go @@ -9,9 +9,6 @@ import ( "fmt" "time" - mcpclient "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" "github.com/onsi/ginkgo/v2" "github.com/onsi/gomega" corev1 "k8s.io/api/core/v1" @@ -20,6 +17,9 @@ import ( "k8s.io/apimachinery/pkg/types" "sigs.k8s.io/controller-runtime/pkg/client" + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" "github.com/stacklok/toolhive/test/e2e/images" diff --git a/test/e2e/thv-operator/virtualmcp/mcpserver_scaling_test.go b/test/e2e/thv-operator/virtualmcp/mcpserver_scaling_test.go index 8d10d954ef..4f1dbb5610 100644 --- a/test/e2e/thv-operator/virtualmcp/mcpserver_scaling_test.go +++ b/test/e2e/thv-operator/virtualmcp/mcpserver_scaling_test.go @@ -13,9 +13,6 @@ import ( "strings" "time" - mcpclient "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" "github.com/onsi/ginkgo/v2" "github.com/onsi/gomega" "github.com/redis/go-redis/v9" @@ -27,6 +24,9 @@ import ( "k8s.io/apimachinery/pkg/util/intstr" "sigs.k8s.io/controller-runtime/pkg/client" + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/test/e2e/images" "github.com/stacklok/toolhive/test/e2e/thv-operator/testutil" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_aggregation_filtering_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_aggregation_filtering_test.go index a7ceca6611..8da104d667 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_aggregation_filtering_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_aggregation_filtering_test.go @@ -8,12 +8,12 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_aggregation_overrides_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_aggregation_overrides_test.go index 56cd88f0ec..a49a887e03 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_aggregation_overrides_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_aggregation_overrides_test.go @@ -7,12 +7,12 @@ import ( "fmt" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_auth_discovery_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_auth_discovery_test.go index 42e851d8a8..f0fc6a13c3 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_auth_discovery_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_auth_discovery_test.go @@ -12,9 +12,6 @@ import ( "strings" "time" - mcpclient "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" corev1 "k8s.io/api/core/v1" @@ -24,6 +21,9 @@ import ( "k8s.io/apimachinery/pkg/util/intstr" "sigs.k8s.io/controller-runtime/pkg/client" + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_codemode_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_codemode_test.go index 980fa5db84..26933ab179 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_codemode_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_codemode_test.go @@ -9,11 +9,11 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_defaultresults_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_defaultresults_test.go index 0d6de43cf0..d5d948f12c 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_defaultresults_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_defaultresults_test.go @@ -7,11 +7,11 @@ import ( "fmt" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" thvjson "github.com/stacklok/toolhive/pkg/json" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_hidden_tools_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_hidden_tools_test.go index c7eac4cc92..b07b02595f 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_hidden_tools_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_hidden_tools_test.go @@ -8,12 +8,12 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" thvjson "github.com/stacklok/toolhive/pkg/json" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_parallel_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_parallel_test.go index 3f8ed12e73..5dc1f6fe96 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_parallel_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_parallel_test.go @@ -7,12 +7,12 @@ import ( "fmt" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" thvjson "github.com/stacklok/toolhive/pkg/json" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_referenced_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_referenced_test.go index 38c40e03fb..1860e1f7aa 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_referenced_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_referenced_test.go @@ -7,12 +7,12 @@ import ( "fmt" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" thvjson "github.com/stacklok/toolhive/pkg/json" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_sequential_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_sequential_test.go index 8dc05d4ce3..67e74e3729 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_sequential_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_composite_sequential_test.go @@ -7,12 +7,12 @@ import ( "fmt" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" thvjson "github.com/stacklok/toolhive/pkg/json" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_conflict_resolution_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_conflict_resolution_test.go index 743ae9fb1b..c2bb5fbd37 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_conflict_resolution_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_conflict_resolution_test.go @@ -8,12 +8,12 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" vmcp "github.com/stacklok/toolhive/pkg/vmcp" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_discovered_mode_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_discovered_mode_test.go index b1b9e064b9..b94c39110b 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_discovered_mode_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_discovered_mode_test.go @@ -12,13 +12,13 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_excludeall_global_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_excludeall_global_test.go index 22b191fe30..04259511e1 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_excludeall_global_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_excludeall_global_test.go @@ -7,12 +7,12 @@ import ( "fmt" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_external_auth_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_external_auth_test.go index b727de8f04..ad0bd309bd 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_external_auth_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_external_auth_test.go @@ -8,13 +8,13 @@ import ( "fmt" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_composite_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_composite_test.go index 701502cc90..f3c6611da3 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_composite_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_composite_test.go @@ -8,11 +8,11 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" thvjson "github.com/stacklok/toolhive/pkg/json" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_multibackend_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_multibackend_test.go index 78d7d59809..15cc7a8903 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_multibackend_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_multibackend_test.go @@ -8,13 +8,13 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" vmcp "github.com/stacklok/toolhive/pkg/vmcp" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_test.go index 7772ed1f13..8e0281da51 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_optimizer_test.go @@ -8,11 +8,11 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" thvjson "github.com/stacklok/toolhive/pkg/json" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_rate_limiting_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_rate_limiting_test.go index 75a8f1f2ed..7a7c02134c 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_rate_limiting_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_rate_limiting_test.go @@ -12,9 +12,6 @@ import ( "strings" "time" - mcpclient "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" "github.com/onsi/ginkgo/v2" "github.com/onsi/gomega" dto "github.com/prometheus/client_model/go" @@ -23,6 +20,9 @@ import ( metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" "github.com/stacklok/toolhive/pkg/ratelimit" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_redis_session_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_redis_session_test.go index 97f48641cb..311b6efeef 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_redis_session_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_redis_session_test.go @@ -16,9 +16,6 @@ import ( "net/http" "time" - mcpclient "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" "github.com/onsi/ginkgo/v2" "github.com/onsi/gomega" corev1 "k8s.io/api/core/v1" @@ -27,6 +24,9 @@ import ( "k8s.io/apimachinery/pkg/runtime" "k8s.io/apimachinery/pkg/types" + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" "github.com/stacklok/toolhive/test/e2e/images" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_session_management_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_session_management_test.go index 936d94be5d..38c219cdfa 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_session_management_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_session_management_test.go @@ -13,9 +13,6 @@ import ( "strings" "time" - mcpclient "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" "github.com/onsi/ginkgo/v2" "github.com/onsi/gomega" appsv1 "k8s.io/api/apps/v1" @@ -24,6 +21,9 @@ import ( metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_toolconfig_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_toolconfig_test.go index 41ba4a3856..33699557b3 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_toolconfig_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_toolconfig_test.go @@ -8,12 +8,12 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" diff --git a/test/e2e/thv-operator/virtualmcp/virtualmcp_yardstick_base_test.go b/test/e2e/thv-operator/virtualmcp/virtualmcp_yardstick_base_test.go index d6121d6920..f17bb90e0b 100644 --- a/test/e2e/thv-operator/virtualmcp/virtualmcp_yardstick_base_test.go +++ b/test/e2e/thv-operator/virtualmcp/virtualmcp_yardstick_base_test.go @@ -10,12 +10,12 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/types" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1" "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1/v1beta1test" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" diff --git a/test/e2e/thv-operator/virtualmcp/wait_for_tools_helpers.go b/test/e2e/thv-operator/virtualmcp/wait_for_tools_helpers.go index 9699e21fc2..7b612625a6 100644 --- a/test/e2e/thv-operator/virtualmcp/wait_for_tools_helpers.go +++ b/test/e2e/thv-operator/virtualmcp/wait_for_tools_helpers.go @@ -9,11 +9,12 @@ import ( "strings" "time" - mcpclient "github.com/mark3labs/mcp-go/client" - "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" "github.com/onsi/ginkgo/v2" "github.com/onsi/gomega" + + mcpclient "github.com/stacklok/toolhive-core/mcpcompat/client" + "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) // WaitForExpectedTools creates MCP sessions with retry until the validateTools diff --git a/test/e2e/vmcp_cli_features_test.go b/test/e2e/vmcp_cli_features_test.go index ea0c2348c7..41dff597fb 100644 --- a/test/e2e/vmcp_cli_features_test.go +++ b/test/e2e/vmcp_cli_features_test.go @@ -16,11 +16,11 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" "gopkg.in/yaml.v3" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" thvjson "github.com/stacklok/toolhive/pkg/json" "github.com/stacklok/toolhive/pkg/vmcp" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" diff --git a/test/e2e/vmcp_optimizer_test.go b/test/e2e/vmcp_optimizer_test.go index e715c8413f..01fd2e46dd 100644 --- a/test/e2e/vmcp_optimizer_test.go +++ b/test/e2e/vmcp_optimizer_test.go @@ -27,10 +27,10 @@ import ( "strings" "time" - "github.com/mark3labs/mcp-go/mcp" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" thvjson "github.com/stacklok/toolhive/pkg/json" vmcp "github.com/stacklok/toolhive/pkg/vmcp" vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config" diff --git a/test/integration/vmcp/helpers/backend.go b/test/integration/vmcp/helpers/backend.go index c5594b625b..bef9376419 100644 --- a/test/integration/vmcp/helpers/backend.go +++ b/test/integration/vmcp/helpers/backend.go @@ -10,8 +10,8 @@ import ( "net/http/httptest" "testing" - "github.com/mark3labs/mcp-go/mcp" - "github.com/mark3labs/mcp-go/server" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + "github.com/stacklok/toolhive-core/mcpcompat/server" ) // BackendTool defines a tool for MCP backend servers. @@ -180,7 +180,7 @@ func WithCaptureHeaders() BackendServerOption { } } -// CreateBackendServer creates an MCP backend server using the mark3labs/mcp-go SDK. +// CreateBackendServer creates an MCP backend server using the stacklok/toolhive-core/mcpcompat SDK. // It returns an *httptest.Server ready to accept streamable-HTTP connections. // // The server automatically registers all provided tools with proper closure handling diff --git a/test/integration/vmcp/helpers/helpers_test.go b/test/integration/vmcp/helpers/helpers_test.go index 3d186c0ee5..28c19b3bb3 100644 --- a/test/integration/vmcp/helpers/helpers_test.go +++ b/test/integration/vmcp/helpers/helpers_test.go @@ -6,8 +6,9 @@ package helpers import ( "testing" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" + + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) // TestGetToolNames tests the GetToolNames helper function. diff --git a/test/integration/vmcp/helpers/mcp_client.go b/test/integration/vmcp/helpers/mcp_client.go index 75167cbdf6..ade50b5cd8 100644 --- a/test/integration/vmcp/helpers/mcp_client.go +++ b/test/integration/vmcp/helpers/mcp_client.go @@ -10,14 +10,15 @@ import ( "sync" "testing" - "github.com/mark3labs/mcp-go/client" - mcptransport "github.com/mark3labs/mcp-go/client/transport" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive-core/mcpcompat/client" + mcptransport "github.com/stacklok/toolhive-core/mcpcompat/client/transport" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" ) -// MCPClient wraps the mark3labs MCP client with test-friendly methods. +// MCPClient wraps the mcpcompat MCP client with test-friendly methods. // It automatically handles initialization and provides semantic assertion helpers // that integrate with Go's testing.TB interface. // diff --git a/test/integration/vmcp/helpers/vmcp_server.go b/test/integration/vmcp/helpers/vmcp_server.go index aecb2f4cad..7805e1047a 100644 --- a/test/integration/vmcp/helpers/vmcp_server.go +++ b/test/integration/vmcp/helpers/vmcp_server.go @@ -75,6 +75,7 @@ type vmcpServerConfig struct { workflowDefs map[string]*composer.WorkflowDefinition telemetryProvider *telemetry.Provider passthroughHeaders []string + sessionTTL time.Duration } // WithPrefixConflictResolution configures prefix-based conflict resolution. @@ -112,6 +113,15 @@ func WithPassthroughHeaders(headers ...string) VMCPServerOption { } } +// WithSessionTTL overrides the server's session time-to-live (default 30m). +// A short TTL is useful for sliding-TTL / eviction regression tests. A zero +// value leaves the default in place. +func WithSessionTTL(ttl time.Duration) VMCPServerOption { + return func(c *vmcpServerConfig) { + c.sessionTTL = ttl + } +} + // getFreePort returns an available TCP port on localhost. func getFreePort(tb testing.TB) int { tb.Helper() @@ -198,6 +208,9 @@ func NewVMCPServer( Base: sessionFactory, }, } + if config.sessionTTL > 0 { + serverCfg.SessionTTL = config.sessionTTL + } vmcpServer, err := vmcpserver.Serve(ctx, coreVMCP, serverCfg) require.NoError(tb, err, "failed to create vMCP server") diff --git a/test/integration/vmcp/pagination_regression_test.go b/test/integration/vmcp/pagination_regression_test.go new file mode 100644 index 0000000000..9d0e971c40 --- /dev/null +++ b/test/integration/vmcp/pagination_regression_test.go @@ -0,0 +1,72 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package vmcp_test + +import ( + "context" + "fmt" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/test/integration/vmcp/helpers" +) + +// TestRegression_Over1000Tools_CompleteSetReceived is a regression anchor for +// the vMCP pagination gap: a backend exposing more than the MCP page size +// (1000 tools) must surface its complete tool set across pagination cursors. +// +// Today the Serve-path tool query path issues a single ListTools with no +// cursor loop, so backends that paginate return only the first page. The test +// is intentionally skipped until the gap is closed (see gap-analysis V1 and the +// follow-up issue); the skeleton remains so the regression cannot regress +// silently — flipping the skip off re-runs the full assertion. +func TestRegression_Over1000Tools_CompleteSetReceived(t *testing.T) { + t.Parallel() + t.Skip("vMCP does not follow pagination cursors; see gap-analysis V1 / follow-up issue") + + ctx := context.Background() + + // Generate 1000+ tools on a single backend to exceed the MCP page size. + const toolCount = 1100 + tools := make([]helpers.BackendTool, 0, toolCount) + for i := 0; i < toolCount; i++ { + name := fmt.Sprintf("tool_%04d", i) + tools = append(tools, helpers.NewBackendTool( + name, + fmt.Sprintf("backend tool number %d", i), + func(_ context.Context, _ map[string]any) string { + return `{"ok": true}` + }, + )) + } + + backend := helpers.CreateBackendServer(t, tools, helpers.WithBackendName("many-tools")) + defer backend.Close() + + backends := []vmcp.Backend{ + helpers.NewBackend("many-tools", + helpers.WithURL(backend.URL+"/mcp"), + helpers.WithMetadata("group", "pagination-test"), + ), + } + + vmcpServer := helpers.NewVMCPServer(ctx, t, backends, + helpers.WithPrefixConflictResolution("{workload}_"), + ) + + vmcpURL := "http://" + vmcpServer.Address() + "/mcp" + client := helpers.NewMCPClient(ctx, t, vmcpURL) + defer client.Close() + + result := client.ListTools(ctx) + + // The complete tool set must be received; a single-page query returns at + // most the page size (1000), so fewer than toolCount indicates a dropped + // tail — the regression this test pins. + assert.GreaterOrEqual(t, len(result.Tools), toolCount, + "vMCP must surface the complete backend tool set across pagination cursors; "+ + "received %d of %d", len(result.Tools), toolCount) +} diff --git a/test/integration/vmcp/per_session_projection_regression_test.go b/test/integration/vmcp/per_session_projection_regression_test.go new file mode 100644 index 0000000000..6f4ac07141 --- /dev/null +++ b/test/integration/vmcp/per_session_projection_regression_test.go @@ -0,0 +1,424 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package vmcp_test + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "sort" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive-core/env" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" + mcpserver "github.com/stacklok/toolhive-core/mcpcompat/server" + "github.com/stacklok/toolhive/pkg/auth" + "github.com/stacklok/toolhive/pkg/authz/authorizers" + "github.com/stacklok/toolhive/pkg/authz/authorizers/cedar" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/aggregator" + vmcpauth "github.com/stacklok/toolhive/pkg/vmcp/auth/factory" + vmcpclient "github.com/stacklok/toolhive/pkg/vmcp/client" + "github.com/stacklok/toolhive/pkg/vmcp/mocks" + "github.com/stacklok/toolhive/pkg/vmcp/router" + "github.com/stacklok/toolhive/pkg/vmcp/server" + vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" + "github.com/stacklok/toolhive/test/integration/vmcp/helpers" +) + +// principalHeader is the request header the test identity middleware reads to +// determine the authenticated principal. The Cedar authorizer resolves the +// principal from the identity's "sub" claim, so the middleware sets both +// Subject and Claims["sub"] to the header value. +const principalHeader = "X-Test-Principal" + +// newCedarVMCPServer builds a Cedar-authz-backed vMCP server (via server.New, +// mirroring pkg/vmcp/server/authz_integration_test.go) backed by a real MCP +// backend at backendURL. An identity middleware injects a principal derived +// from the X-Test-Principal request header so each session binds to its +// caller and the Cedar admission seam can resolve Client::"". +// +// A priority conflict resolver is used so tool names are NOT prefixed — the +// Cedar policies can name tools by their raw backend names (e.g. +// Tool::"secret-tool"). +func newCedarVMCPServer(t *testing.T, backendURL string, policies ...string) *httptest.Server { + t.Helper() + + ctrl := gomock.NewController(t) + t.Cleanup(ctrl.Finish) + + mockBackendRegistry := mocks.NewMockBackendRegistry(ctrl) + + backend := vmcp.Backend{ + ID: "real-backend", + Name: "real-backend", + BaseURL: backendURL, + TransportType: "streamable-http", + } + mockBackendRegistry.EXPECT().List(gomock.Any()).Return([]vmcp.Backend{backend}).AnyTimes() + mockBackendRegistry.EXPECT().Get(gomock.Any(), gomock.Any()).Return(&backend).AnyTimes() + + authReg, err := vmcpauth.NewOutgoingAuthRegistry(context.Background(), &env.OSReader{}) + require.NoError(t, err) + factory := vmcpsession.NewSessionFactory(authReg) + + backendClient, err := vmcpclient.NewHTTPBackendClient(authReg) + require.NoError(t, err) + // Priority resolver keeps raw tool names so Cedar policies can name + // Tool::"secret-tool" rather than a prefixed variant. + resolver, err := aggregator.NewPriorityConflictResolver([]string{backend.Name}) + require.NoError(t, err) + agg := aggregator.NewDefaultAggregator(backendClient, resolver, nil, nil) + + // Identity middleware: derive the principal from the X-Test-Principal + // header so two sessions (alice, bob) bind to different identities. The + // Cedar authorizer resolves the principal from the "sub" claim, so both + // Subject and Claims["sub"] are set to the header value. + identityMiddleware := func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + principal := r.Header.Get(principalHeader) + if principal == "" { + principal = "anonymous" + } + id := &auth.Identity{PrincipalInfo: auth.PrincipalInfo{ + Subject: principal, + Name: principal, + Claims: map[string]any{"sub": principal, "name": principal}, + }} + next.ServeHTTP(w, r.WithContext(auth.WithIdentity(r.Context(), id))) + }) + } + + authzCfg, err := authorizers.NewConfig(cedar.Config{ + Version: "1.0", + Type: cedar.ConfigType, + Options: &cedar.ConfigOptions{Policies: policies, EntitiesJSON: "[]"}, + }) + require.NoError(t, err) + + srv, err := server.New( + context.Background(), + &server.Config{ + Name: "test-vmcp", + Host: "127.0.0.1", + Port: 0, + SessionTTL: 5 * time.Minute, + SessionFactory: factory, + Aggregator: agg, + AuthMiddleware: identityMiddleware, + Authz: authzCfg, + }, + router.NewSessionRouter(&vmcp.RoutingTable{}), + backendClient, + mockBackendRegistry, + nil, + ) + require.NoError(t, err) + + handler, err := srv.Handler(context.Background()) + require.NoError(t, err) + + ts := httptest.NewServer(handler) + t.Cleanup(ts.Close) + return ts +} + +// rawClient is a minimal JSON-RPC-over-HTTP client for the Cedar authz tests. +// It is needed because the helpers.MCPClient.CallTest helper uses require.NoError +// on the SDK error, which would abort the test on a policy-rejected call. This +// client surfaces the raw HTTP response so tests can inspect status codes and +// error bodies directly. +type rawClient struct { + baseURL string + sessionID string + nextID int +} + +func newRawClient(baseURL string) *rawClient { + return &rawClient{baseURL: baseURL, nextID: 1} +} + +func (c *rawClient) initialize(t *testing.T, principal string) { + t.Helper() + resp := c.postMCP(t, principal, map[string]any{ + "jsonrpc": "2.0", + "id": c.nextID, + "method": "initialize", + "params": map[string]any{ + "protocolVersion": "2025-06-18", + "capabilities": map[string]any{}, + "clientInfo": map[string]any{"name": "test", "version": "1.0"}, + }, + }, "") + c.nextID++ + defer resp.Body.Close() + c.sessionID = resp.Header.Get("Mcp-Session-Id") + require.NotEmpty(t, c.sessionID, "initialize response missing Mcp-Session-Id header") + + // Send the initialized notification (no id, no response expected). + notif := c.postMCP(t, principal, map[string]any{ + "jsonrpc": "2.0", + "method": "notifications/initialized", + }, c.sessionID) + notif.Body.Close() +} + +func (c *rawClient) listTools(t *testing.T, principal string) map[string]any { + t.Helper() + resp := c.postMCP(t, principal, map[string]any{ + "jsonrpc": "2.0", + "id": c.nextID, + "method": "tools/list", + "params": map[string]any{}, + }, c.sessionID) + c.nextID++ + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode, "tools/list failed: %s", string(body)) + var parsed map[string]any + require.NoError(t, json.Unmarshal(body, &parsed)) + return parsed +} + +func (c *rawClient) callTool(t *testing.T, principal, name string, args map[string]any) (int, map[string]any) { + t.Helper() + resp := c.postMCP(t, principal, map[string]any{ + "jsonrpc": "2.0", + "id": c.nextID, + "method": "tools/call", + "params": map[string]any{ + "name": name, + "arguments": args, + }, + }, c.sessionID) + c.nextID++ + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + var parsed map[string]any + _ = json.Unmarshal(body, &parsed) // best-effort; body may be plain text on 404 + return resp.StatusCode, parsed +} + +func (c *rawClient) postMCP(t *testing.T, principal string, body map[string]any, sessionID string) *http.Response { + t.Helper() + rawBody, err := json.Marshal(body) + require.NoError(t, err) + req, err := http.NewRequestWithContext(context.Background(), http.MethodPost, c.baseURL+"/mcp", bytes.NewReader(rawBody)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + req.Header.Set(principalHeader, principal) + if sessionID != "" { + req.Header.Set("Mcp-Session-Id", sessionID) + } + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + return resp +} + +// toolNamesFromResult extracts the tool names from a tools/list JSON-RPC result. +func toolNamesFromResult(t *testing.T, result map[string]any) []string { + t.Helper() + rawResult, ok := result["result"].(map[string]any) + require.True(t, ok, "response missing result object: %v", result) + toolsRaw, ok := rawResult["tools"].([]any) + require.True(t, ok, "result missing tools array: %v", rawResult) + names := make([]string, 0, len(toolsRaw)) + for _, t2 := range toolsRaw { + tool, ok := t2.(map[string]any) + require.True(t, ok) + name, ok := tool["name"].(string) + require.True(t, ok) + names = append(names, name) + } + sort.Strings(names) + return names +} + +// startAuthzBackend starts a real MCP backend exposing the named tools. Each +// tool echoes back its name and the input so a caller can confirm the backend +// was reached. +func startAuthzBackend(t *testing.T, toolNames ...string) string { + t.Helper() + tools := make([]helpers.BackendTool, 0, len(toolNames)) + for _, name := range toolNames { + tools = append(tools, helpers.NewBackendTool( + name, + "test tool "+name, + func(_ context.Context, args map[string]any) string { + return `{"tool":` + name + `,"args":` + toJSON(args) + `}` + }, + )) + } + backend := helpers.CreateBackendServer(t, tools, helpers.WithBackendName("real-backend")) + t.Cleanup(backend.Close) + return backend.URL + "/mcp" +} + +// toJSON is a tiny marshal helper that never fails the test (returns "null"). +func toJSON(v map[string]any) string { + b, err := json.Marshal(v) + if err != nil { + return "null" + } + return string(b) +} + +// TestRegression_TwoSessions_DifferentAuthz_DifferentToolsList proves the +// core admission seam projects a per-identity tool list: two sessions bound to +// different principals see different advertised tool sets under the same Cedar +// policy. Alice (permitted for secret-tool) sees both secret-tool and +// public-tool; Bob (only permitted for public-tool) sees only public-tool. +func TestRegression_TwoSessions_DifferentAuthz_DifferentToolsList(t *testing.T) { + t.Parallel() + + backendURL := startAuthzBackend(t, "secret-tool", "public-tool") + + ts := newCedarVMCPServer(t, backendURL, + `permit(principal == Client::"alice", action == Action::"call_tool", resource == Tool::"secret-tool");`, + `permit(principal, action == Action::"call_tool", resource == Tool::"public-tool");`, + ) + + // Alice's session. + alice := newRawClient(ts.URL) + alice.initialize(t, "alice") + aliceTools := toolNamesFromResult(t, alice.listTools(t, "alice")) + assert.Equal(t, []string{"public-tool", "secret-tool"}, aliceTools, + "alice (permitted for secret-tool) must see both tools") + + // Bob's session. + bob := newRawClient(ts.URL) + bob.initialize(t, "bob") + bobTools := toolNamesFromResult(t, bob.listTools(t, "bob")) + assert.Equal(t, []string{"public-tool"}, bobTools, + "bob (not permitted for secret-tool) must see only public-tool") +} + +// TestRegression_FilteredToolUncallableInSessionA_CallableInSessionB proves +// the list-filter/call-deny pairing: a tool filtered out of a session (bob's) +// is rejected as unknown when called, while the same tool is callable in a +// session where it was advertised (alice's). +func TestRegression_FilteredToolUncallableInSessionA_CallableInSessionB(t *testing.T) { + t.Parallel() + + backendURL := startAuthzBackend(t, "secret-tool", "public-tool") + + ts := newCedarVMCPServer(t, backendURL, + `permit(principal == Client::"alice", action == Action::"call_tool", resource == Tool::"secret-tool");`, + `permit(principal, action == Action::"call_tool", resource == Tool::"public-tool");`, + ) + + // Alice can call secret-tool (it is in her session's advertised set). + alice := newRawClient(ts.URL) + alice.initialize(t, "alice") + require.Contains(t, toolNamesFromResult(t, alice.listTools(t, "alice")), "secret-tool") + + status, result := alice.callTool(t, "alice", "secret-tool", map[string]any{"input": "hi"}) + require.Equal(t, http.StatusOK, status, "alice's permitted call must succeed: %v", result) + assert.False(t, isToolResultError(result), "alice's call to secret-tool must not be an error: %v", result) + + // Bob cannot call secret-tool: it was filtered out of his session, so the + // SDK rejects it as an unknown tool. Poll briefly because the session's + // tool injection runs in a hook that fires after initialize returns. + bob := newRawClient(ts.URL) + bob.initialize(t, "bob") + require.NotContains(t, toolNamesFromResult(t, bob.listTools(t, "bob")), "secret-tool") + + var lastResult map[string]any + require.Eventually(t, func() bool { + var s int + s, lastResult = bob.callTool(t, "bob", "secret-tool", map[string]any{"input": "hi"}) + if s != http.StatusOK { + return true // 404/400 also indicates rejection + } + return isToolResultError(lastResult) + }, 5*time.Second, 50*time.Millisecond, + "bob's call to secret-tool (filtered out) must be rejected") +} + +// TestRegression_SetSessionTools_FixedAtInitialize_NoMidSessionReconciliation +// asserts that a session's advertised tool set is fixed at initialize time +// and is not re-aggregated on each tools/list call. Calling tools/list twice +// on the same session must return the same set. +func TestRegression_SetSessionTools_FixedAtInitialize_NoMidSessionReconciliation(t *testing.T) { + t.Parallel() + + ctx := context.Background() + + backend := helpers.CreateBackendServer(t, []helpers.BackendTool{ + helpers.NewBackendTool("tool-a", "tool A", func(_ context.Context, _ map[string]any) string { + return `{"tool":"a"}` + }), + helpers.NewBackendTool("tool-b", "tool B", func(_ context.Context, _ map[string]any) string { + return `{"tool":"b"}` + }), + }, helpers.WithBackendName("ab-backend")) + defer backend.Close() + + backends := []vmcp.Backend{ + helpers.NewBackend("ab-backend", helpers.WithURL(backend.URL+"/mcp")), + } + + vmcpServer := helpers.NewVMCPServer(ctx, t, backends, + helpers.WithPrefixConflictResolution("{workload}_"), + ) + + vmcpURL := "http://" + vmcpServer.Address() + "/mcp" + client := helpers.NewMCPClient(ctx, t, vmcpURL) + defer client.Close() + + first := helpers.GetToolNames(client.ListTools(ctx)) + second := helpers.GetToolNames(client.ListTools(ctx)) + + sort.Strings(first) + sort.Strings(second) + + assert.Equal(t, []string{"ab-backend_tool-a", "ab-backend_tool-b"}, first, + "initial tools/list must show both tools") + assert.Equal(t, first, second, + "the session's tool set is fixed at initialize; a second tools/list must return the same set") +} + +// isToolResultError returns true when a tools/call JSON-RPC result has +// IsError=true (the SDK rejected the call) or the response carries a +// JSON-RPC error object. +func isToolResultError(result map[string]any) bool { + if _, ok := result["error"]; ok { + return true + } + if r, ok := result["result"].(map[string]any); ok { + if isErr, _ := r["isError"].(bool); isErr { + return true + } + } + return false +} + +// mustParseJSONRPC reads resp.Body and unmarshals it as a JSON-RPC envelope. +// If the body is not valid JSON (e.g. a plain-text 404), it returns an empty +// map so callers can fall back to status-code checks. +func mustParseJSONRPC(t *testing.T, resp *http.Response) map[string]any { + t.Helper() + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + var parsed map[string]any + _ = json.Unmarshal(body, &parsed) // best-effort; body may be plain text on 4xx + return parsed +} + +// Compile-time check that the mcpcompat imports are used (not mark3labs/mcp-go). +var ( + _ = mcp.LATEST_PROTOCOL_VERSION + _ = mcpserver.NewMCPServer +) diff --git a/test/integration/vmcp/sliding_ttl_regression_test.go b/test/integration/vmcp/sliding_ttl_regression_test.go new file mode 100644 index 0000000000..6a5d7a4cd9 --- /dev/null +++ b/test/integration/vmcp/sliding_ttl_regression_test.go @@ -0,0 +1,146 @@ +// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package vmcp_test + +import ( + "context" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/test/integration/vmcp/helpers" +) + +// TestRegression_SlidingSessionTTL_TrafficKeepsSessionAlive proves the +// session TTL is sliding: a session that receives traffic at least once per +// TTL window stays alive indefinitely, while an idle session is evicted once +// its TTL elapses. +// +// The transport session storage (LocalSessionDataStorage.Load) refreshes the +// last-access timestamp on every read, and the SDK's SessionIdManager.Validate +// is called on every request — so active traffic keeps the session alive while +// a session that goes idle for longer than the TTL is rejected on its next +// request. +// +// This test is timing-sensitive and must NOT run in parallel. +// +//nolint:paralleltest // timing-sensitive: relies on real TTL expiry and background cleanup +func TestRegression_SlidingSessionTTL_TrafficKeepsSessionAlive(t *testing.T) { + const sessionTTL = 2 * time.Second + + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + + backend := helpers.CreateBackendServer(t, []helpers.BackendTool{ + helpers.NewBackendTool("ping", "ping tool", func(_ context.Context, _ map[string]any) string { + return `{"pong":true}` + }), + }, helpers.WithBackendName("ping-backend")) + defer backend.Close() + + backends := []vmcp.Backend{ + helpers.NewBackend("ping-backend", helpers.WithURL(backend.URL+"/mcp")), + } + + vmcpServer := helpers.NewVMCPServer(ctx, t, backends, + helpers.WithPrefixConflictResolution("{workload}_"), + helpers.WithSessionTTL(sessionTTL), + ) + vmcpURL := "http://" + vmcpServer.Address() + "/mcp" + + // ── Active session: traffic every 500ms for 4s (2x the TTL) ────────────── + // + // If the TTL were fixed (not sliding), the session would expire at the 2s + // mark and the next tools/call would fail. A sliding TTL refreshes + // last-access on each request, so all calls must succeed. + // + // We use tools/call (not tools/list) here because tools/list returns the + // SDK's cached tool set even after the vMCP session storage has evicted the + // session (the go-sdk transport map is not TTL-gated), so it would pass + // whether the TTL slides or not. tools/call routes through + // enforceSessionBinding → storage.Load, so it fails once the session is gone + // — making it the assertion that actually distinguishes sliding from fixed TTL. + activeClient := helpers.NewMCPClient(ctx, t, vmcpURL) + defer activeClient.Close() + + const activeTool = "ping-backend_ping" + + const ( + tickInterval = 500 * time.Millisecond + totalWindow = 4 * time.Second + ) + ticks := int(totalWindow / tickInterval) + + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + ticker := time.NewTicker(tickInterval) + defer ticker.Stop() + for i := 0; i < ticks; i++ { + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + // tools/call routes through enforceSessionBinding → storage.Load, + // which refreshes last-access in the transport session storage. + _ = activeClient.CallTool(ctx, activeTool, map[string]any{}) + } + }() + wg.Wait() + + // One final call after the full window: the session must still be alive. + // CallTool fails the test if the session was evicted (the SDK surfaces the + // missing session as a transport error or an error result). + final := activeClient.CallTool(ctx, activeTool, map[string]any{}) + assert.False(t, final.IsError, + "active session must survive past its TTL while receiving traffic; got error result %v", final.Content) + + // ── Idle session: no traffic for 3s (> TTL) must be evicted ────────────── + // + // A second session that goes idle for longer than the TTL must be rejected + // on its next tool call. tools/list is NOT a reliable eviction signal here: + // the go-sdk StreamableHTTPHandler keeps its per-session transport in an + // in-memory map that is not TTL-gated, so tools/list continues to return + // the cached tool set even after the vMCP session storage has evicted the + // session. The eviction IS observable on tools/call, whose handler runs + // enforceSessionBinding → GetMultiSession → checkSession → storage.Load, + // which returns ErrSessionNotFound once the TTL has elapsed and the + // background cleanup sweep has removed the entry. + idleClient := newRawClient(vmcpURL) + idleClient.initialize(t, "") + + // Let the session go idle well past the TTL. The cleanup goroutine runs on + // a ttl/2 ticker, so waiting 2x the TTL after the idle period guarantees the + // background sweep has evicted the session. + idleFor := 3 * sessionTTL + time.Sleep(idleFor) + + // A tool call on the idle session must be rejected: enforceSessionBinding + // fails because the session is gone from storage. The SDK surfaces this as + // either an HTTP error (4xx) or a JSON-RPC error result — assert both. + resp := idleClient.postMCP(t, "", map[string]any{ + "jsonrpc": "2.0", + "id": idleClient.nextID, + "method": "tools/call", + "params": map[string]any{ + "name": "ping-backend_ping", + "arguments": map[string]any{}, + }, + }, idleClient.sessionID) + idleClient.nextID++ + defer resp.Body.Close() + + rejected := resp.StatusCode >= 400 || + isToolResultError(mustParseJSONRPC(t, resp)) + assert.True(t, rejected, + "idle session must be evicted after exceeding the TTL (status %d)", resp.StatusCode) +} + +// Compile-time check that this file stays in the vmcp_test package. +var _ vmcp.Backend diff --git a/test/integration/vmcp/vmcp_integration_test.go b/test/integration/vmcp/vmcp_integration_test.go index e8b37a6da4..91e296b810 100644 --- a/test/integration/vmcp/vmcp_integration_test.go +++ b/test/integration/vmcp/vmcp_integration_test.go @@ -12,10 +12,10 @@ import ( "testing" "time" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" thvjson "github.com/stacklok/toolhive/pkg/json" "github.com/stacklok/toolhive/pkg/telemetry" "github.com/stacklok/toolhive/pkg/vmcp" diff --git a/test/integration/vmcp/vmcp_typing_integration_test.go b/test/integration/vmcp/vmcp_typing_integration_test.go index d973cf738a..f34abf635f 100644 --- a/test/integration/vmcp/vmcp_typing_integration_test.go +++ b/test/integration/vmcp/vmcp_typing_integration_test.go @@ -9,9 +9,9 @@ import ( "testing" "time" - "github.com/mark3labs/mcp-go/mcp" "github.com/stretchr/testify/require" + "github.com/stacklok/toolhive-core/mcpcompat/mcp" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/composer" "github.com/stacklok/toolhive/test/integration/vmcp/helpers"