| 1 | package plugin |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "fmt" |
| 7 | "sync" |
| 8 | "testing" |
| 9 | |
| 10 | mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp" |
| 11 | ) |
| 12 | |
| 13 | // captureInitializeServer runs an in-memory server whose only tool records the |
| 14 | // client's initialize capabilities, so a test can assert exactly what the |
| 15 | // host declared on the wire. |
| 16 | type capturedInitialize struct { |
| 17 | mu sync.Mutex |
| 18 | caps json.RawMessage |
| 19 | } |
| 20 | |
| 21 | func (c *capturedInitialize) capabilities() json.RawMessage { |
| 22 | c.mu.Lock() |
| 23 | defer c.mu.Unlock() |
| 24 | return c.caps |
| 25 | } |
| 26 | |
| 27 | func newCapabilityCaptureServer(t *testing.T, captured *capturedInitialize) *mcpsdk.Server { |
| 28 | t.Helper() |
| 29 | server := mcpsdk.NewServer(&mcpsdk.Implementation{Name: "cap-fixture", Version: "1"}, nil) |
| 30 | server.AddTool(&mcpsdk.Tool{ |
| 31 | Name: "report_capabilities", |
| 32 | Description: "reports the client's declared capabilities", |
| 33 | InputSchema: map[string]any{"type": "object", "properties": map[string]any{}}, |
| 34 | }, func(ctx context.Context, req *mcpsdk.CallToolRequest) (*mcpsdk.CallToolResult, error) { |
| 35 | ss, ok := req.GetSession().(*mcpsdk.ServerSession) |
| 36 | if !ok || ss.InitializeParams() == nil { |
| 37 | return nil, fmt.Errorf("no initialize params on session") |
| 38 | } |
| 39 | raw, err := json.Marshal(ss.InitializeParams().Capabilities) |
| 40 | if err != nil { |
| 41 | return nil, err |
| 42 | } |
| 43 | captured.mu.Lock() |
| 44 | captured.caps = raw |
| 45 | captured.mu.Unlock() |
| 46 | return &mcpsdk.CallToolResult{Content: []mcpsdk.Content{&mcpsdk.TextContent{Text: "ok"}}}, nil |
| 47 | }) |
| 48 | return server |
| 49 | } |
| 50 | |
| 51 | // connectWithProfile drives one real SDK session through the profile-aware |
| 52 | // build path and calls the capture tool. |
| 53 | func connectWithProfile(t *testing.T, profile HostProfile) json.RawMessage { |
| 54 | t.Helper() |
| 55 | captured := &capturedInitialize{} |
| 56 | lifeCtx, cancelLife := context.WithCancel(context.Background()) |
| 57 | transport := &sdkSessionTransport{ |
| 58 | name: "cap-check", spec: Spec{Name: "cap-check", Type: "http"}, profile: profile, |
| 59 | lifeCtx: lifeCtx, cancel: cancelLife, state: SessionStateConnecting, |
| 60 | } |
| 61 | transport.endpointFactory = func(ctx context.Context) (sdkEndpoint, error) { |
| 62 | clientSide, serverSide := mcpsdk.NewInMemoryTransports() |
| 63 | server := newCapabilityCaptureServer(t, captured) |
| 64 | go func() { _ = server.Run(ctx, serverSide) }() |
| 65 | return sdkEndpoint{transport: clientSide}, nil |
| 66 | } |
| 67 | t.Cleanup(transport.close) |
| 68 | |
| 69 | managed, err := transport.acquire(t.Context()) |
| 70 | if err != nil { |
| 71 | t.Fatalf("acquire: %v", err) |
| 72 | } |
| 73 | if _, err := invokeSDKMethod(t.Context(), managed.session, "tools/call", map[string]any{ |
| 74 | "name": "report_capabilities", |
| 75 | "arguments": map[string]any{}, |
| 76 | }); err != nil { |
| 77 | t.Fatalf("tools/call: %v", err) |
| 78 | } |
| 79 | raw := captured.capabilities() |
| 80 | if raw == nil { |
| 81 | t.Fatal("capture tool never ran") |
| 82 | } |
| 83 | return raw |
| 84 | } |
| 85 | |
| 86 | func TestProfileDeclaresExpectedClientCapabilities(t *testing.T) { |
| 87 | cases := []struct { |
| 88 | profile HostProfile |
| 89 | wantElicitationForm bool |
| 90 | wantElicitationURL bool |
| 91 | wantAppsUI bool |
| 92 | }{ |
| 93 | {HostProfileCore, false, false, false}, |
| 94 | {HostProfileInteractive, true, true, false}, |
| 95 | {HostProfileDesktopApps, true, true, true}, |
| 96 | } |
| 97 | for _, tc := range cases { |
| 98 | t.Run(tc.profile.String(), func(t *testing.T) { |
| 99 | raw := connectWithProfile(t, tc.profile) |
| 100 | var caps struct { |
| 101 | Elicitation *struct { |
| 102 | Form json.RawMessage `json:"form"` |
| 103 | URL json.RawMessage `json:"url"` |
| 104 | } `json:"elicitation"` |
| 105 | Extensions map[string]json.RawMessage `json:"extensions"` |
| 106 | } |
| 107 | if err := json.Unmarshal(raw, &caps); err != nil { |
| 108 | t.Fatalf("parse capabilities %s: %v", raw, err) |
| 109 | } |
| 110 | if tc.wantElicitationForm && (caps.Elicitation == nil || len(caps.Elicitation.Form) == 0) { |
| 111 | t.Errorf("profile %s: elicitation.form not declared: %s", tc.profile, raw) |
| 112 | } |
| 113 | if tc.wantElicitationURL && (caps.Elicitation == nil || len(caps.Elicitation.URL) == 0) { |
| 114 | t.Errorf("profile %s: elicitation.url not declared: %s", tc.profile, raw) |
| 115 | } |
| 116 | if !tc.wantElicitationForm && !tc.wantElicitationURL && caps.Elicitation != nil { |
| 117 | t.Errorf("profile %s: elicitation declared: %s", tc.profile, raw) |
| 118 | } |
| 119 | ui, hasUI := caps.Extensions[AppsUIExtensionID] |
| 120 | if tc.wantAppsUI != hasUI { |
| 121 | t.Errorf("profile %s: extension %s presence = %v, want %v (%s)", tc.profile, AppsUIExtensionID, hasUI, tc.wantAppsUI, raw) |
| 122 | } |
| 123 | if tc.wantAppsUI { |
| 124 | var settings struct { |
| 125 | MimeTypes []string `json:"mimeTypes"` |
| 126 | } |
| 127 | if err := json.Unmarshal(ui, &settings); err != nil { |
| 128 | t.Fatalf("parse ui extension settings %s: %v", ui, err) |
| 129 | } |
| 130 | if len(settings.MimeTypes) != 1 || settings.MimeTypes[0] != AppsMimeType { |
| 131 | t.Errorf("ui extension mimeTypes = %v, want [%s]", settings.MimeTypes, AppsMimeType) |
| 132 | } |
| 133 | } |
| 134 | }) |
| 135 | } |
| 136 | } |
| 137 | |
| 138 | func TestCoreProfileMatchesLegacyCapabilityBytes(t *testing.T) { |
| 139 | raw := connectWithProfile(t, HostProfileCore) |
| 140 | var caps map[string]json.RawMessage |
| 141 | if err := json.Unmarshal(raw, &caps); err != nil { |
| 142 | t.Fatalf("parse %s: %v", raw, err) |
| 143 | } |
| 144 | if _, ok := caps["elicitation"]; ok { |
| 145 | t.Errorf("core profile declared elicitation: %s", raw) |
| 146 | } |
| 147 | if _, ok := caps["extensions"]; ok { |
| 148 | t.Errorf("core profile declared extensions: %s", raw) |
| 149 | } |
| 150 | } |
| 151 | |
| 152 | func TestNoBrokerElicitationCancels(t *testing.T) { |
| 153 | lifeCtx, cancelLife := context.WithCancel(context.Background()) |
| 154 | transport := &sdkSessionTransport{ |
| 155 | name: "no-broker", spec: Spec{Name: "no-broker", Type: "http"}, profile: HostProfileInteractive, |
| 156 | lifeCtx: lifeCtx, cancel: cancelLife, state: SessionStateConnecting, |
| 157 | } |
| 158 | t.Cleanup(transport.close) |
| 159 | res, err := transport.handleElicitation(t.Context(), &mcpsdk.ElicitRequest{Params: &mcpsdk.ElicitParams{ |
| 160 | Mode: "form", Message: "need input", |
| 161 | }}) |
| 162 | if err != nil { |
| 163 | t.Fatalf("handleElicitation: %v", err) |
| 164 | } |
| 165 | if res.Action != "cancel" { |
| 166 | t.Fatalf("no-broker action = %q, want cancel", res.Action) |
| 167 | } |
| 168 | } |
| 169 |