Skip to content

Commit c509b6d

Browse files
authored
feat: add support for text-only content in MCP tools (#20)
Cursor removed support for MCP Resources in version 0.46.9+, breaking compatibility with Firebolt’s MCP server which relied heavily on resource embedding. To restore functionality, introduced the `--disable-resources` CLI flag that allows `Docs` and `Connect` tools to return plain text responses instead of embedded resources.
1 parent 175ed4b commit c509b6d

File tree

7 files changed

+184
-31
lines changed

7 files changed

+184
-31
lines changed

cmd/firebolt-mcp-server/main.go

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,13 @@ func main() {
6363
Usage: "SSE transport listen address (used only if transport is set to sse)",
6464
Sources: cli.EnvVars("FIREBOLT_MCP_TRANSPORT_SSE_LISTEN_ADDRESS"),
6565
},
66+
&cli.BoolFlag{
67+
Name: "disable-resources",
68+
Category: "MCP Transport",
69+
Value: false,
70+
Usage: "Return text content instead of embedded resources (for clients that do not support resources)",
71+
Sources: cli.EnvVars("FIREBOLT_MCP_DISABLE_RESOURCES"),
72+
},
6673
&cli.StringFlag{
6774
Name: "client-id",
6875
Category: "Firebolt Authentication",
@@ -123,6 +130,7 @@ func run(ctx context.Context, cmd *cli.Command) error {
123130

124131
// Initialize MCP server
125132
docsProof := generateRandomSecret()
133+
disableResources := cmd.Bool("disable-resources")
126134
resourceDocs := resources.NewDocs(fireboltdocs.FS, docsProof)
127135
resourceAccounts := resources.NewAccounts(discoveryClient)
128136
resourceDatabases := resources.NewDatabases(dbPool)
@@ -133,8 +141,8 @@ func run(ctx context.Context, cmd *cli.Command) error {
133141
cmd.String("transport"),
134142
cmd.String("transport-sse-listen-address"),
135143
[]server.Tool{
136-
tools.NewConnect(resourceAccounts, resourceDatabases, resourceEngines, docsProof),
137-
tools.NewDocs(resourceDocs),
144+
tools.NewConnect(resourceAccounts, resourceDatabases, resourceEngines, docsProof, disableResources),
145+
tools.NewDocs(resourceDocs, disableResources),
138146
tools.NewQuery(dbPool),
139147
},
140148
[]server.Prompt{

pkg/clients/database/pool.go

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ func NewPoolWithConnectionFactory(
6363
}
6464

6565
type poolImpl struct {
66-
sync.RWMutex
66+
sync.Mutex
6767
isClosed bool
6868
logger *slog.Logger
6969
closers []func()
@@ -75,6 +75,9 @@ type poolImpl struct {
7575

7676
func (p *poolImpl) GetConnection(params PoolParams) (Connection, error) {
7777

78+
p.Lock()
79+
defer p.Unlock()
80+
7881
connectionParams := ConnectionParams{
7982
ClientID: p.clientID,
8083
ClientSecret: p.clientSecret,
@@ -84,29 +87,23 @@ func (p *poolImpl) GetConnection(params PoolParams) (Connection, error) {
8487
}
8588
hash := connectionParams.Hash()
8689

87-
// First, try to get an existing connection with a read lock
88-
p.RLock()
90+
// First, try to get an existing connection
8991
if p.isClosed {
90-
p.RUnlock()
9192
return nil, ErrPoolClosed
9293
}
9394
if conn, ok := p.connections[hash]; ok {
94-
p.RUnlock()
9595
return conn, nil
9696
}
97-
p.RUnlock()
9897

9998
// Create a new connection if one doesn't exist
10099
conn, closer, err := p.newConnectionFunc(p.logger, connectionParams)
101100
if err != nil {
102101
return nil, fmt.Errorf("failed to create connection: %w", err)
103102
}
104103

105-
// Store the new connection in the pool with a write lock
106-
p.Lock()
104+
// Store the new connection in the pool
107105
p.connections[hash] = conn
108106
p.closers = append(p.closers, closer)
109-
p.Unlock()
110107

111108
return conn, nil
112109
}

pkg/tools/connect.go

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ type Connect struct {
4444
databasesFetcher DatabaseResourcesFetcher // Fetches database resources
4545
enginesFetcher EngineResourcesFetcher // Fetches engine resources
4646
docsProof string // Shared with the docs resources
47+
disableResources bool // Return text content instead of embedded resources
4748
}
4849

4950
// NewConnect creates a new instance of the Connect tool with the provided resource fetchers.
@@ -53,12 +54,14 @@ func NewConnect(
5354
databasesFetcher DatabaseResourcesFetcher,
5455
enginesFetcher EngineResourcesFetcher,
5556
docsProof string,
57+
disableResources bool,
5658
) *Connect {
5759
return &Connect{
5860
accountsFetcher: accountsFetcher,
5961
databasesFetcher: databasesFetcher,
6062
enginesFetcher: enginesFetcher,
6163
docsProof: docsProof,
64+
disableResources: disableResources,
6265
}
6366
}
6467

@@ -178,7 +181,7 @@ func (t *Connect) Handler(ctx context.Context, request mcp.CallToolRequest) (*mc
178181
return &mcp.CallToolResult{
179182
Result: mcp.Result{},
180183
Content: itertools.Map(results, func(i mcp.ResourceContents) mcp.Content {
181-
return mcp.NewEmbeddedResource(i)
184+
return textOrResourceContent(t.disableResources, i)
182185
}),
183186
IsError: false,
184187
}, nil

pkg/tools/connect_test.go

Lines changed: 78 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -80,13 +80,13 @@ func createEngineResource(accountName, engineName string) mcp.ResourceContents {
8080

8181
func TestNewConnect(t *testing.T) {
8282
mock := &MockResourceFetcher{}
83-
connectTool := tools.NewConnect(mock, mock, mock, validProof)
83+
connectTool := tools.NewConnect(mock, mock, mock, validProof, false)
8484
assert.NotNil(t, connectTool)
8585
}
8686

8787
func TestConnect_Tool(t *testing.T) {
8888
mock := &MockResourceFetcher{}
89-
connectTool := tools.NewConnect(mock, mock, mock, validProof)
89+
connectTool := tools.NewConnect(mock, mock, mock, validProof, false)
9090

9191
tool := connectTool.Tool()
9292
assert.Equal(t, "firebolt_connect", tool.Name)
@@ -131,7 +131,7 @@ func TestConnect_Handler_Success(t *testing.T) {
131131
}
132132

133133
// Create the tool
134-
connectTool := tools.NewConnect(mock, mock, mock, validProof)
134+
connectTool := tools.NewConnect(mock, mock, mock, validProof, false)
135135

136136
// Execute the handler
137137
request := mcp.CallToolRequest{}
@@ -199,7 +199,7 @@ func TestConnect_Handler_AccountFetchFailure(t *testing.T) {
199199
},
200200
}
201201

202-
connectTool := tools.NewConnect(mock, mock, mock, validProof)
202+
connectTool := tools.NewConnect(mock, mock, mock, validProof, false)
203203
request := mcp.CallToolRequest{}
204204
request.Params.Arguments = map[string]any{
205205
"docs_proof": validProof,
@@ -225,7 +225,7 @@ func TestConnect_Handler_InvalidAccountResource(t *testing.T) {
225225
},
226226
}
227227

228-
connectTool := tools.NewConnect(mock, mock, mock, validProof)
228+
connectTool := tools.NewConnect(mock, mock, mock, validProof, false)
229229
request := mcp.CallToolRequest{}
230230
request.Params.Arguments = map[string]any{
231231
"docs_proof": validProof,
@@ -251,7 +251,7 @@ func TestConnect_Handler_InvalidAccountJSON(t *testing.T) {
251251
},
252252
}
253253

254-
connectTool := tools.NewConnect(mock, mock, mock, validProof)
254+
connectTool := tools.NewConnect(mock, mock, mock, validProof, false)
255255
request := mcp.CallToolRequest{}
256256
request.Params.Arguments = map[string]any{
257257
"docs_proof": validProof,
@@ -273,7 +273,7 @@ func TestConnect_Handler_DatabasesFetchFailure(t *testing.T) {
273273
},
274274
}
275275

276-
connectTool := tools.NewConnect(mock, mock, mock, validProof)
276+
connectTool := tools.NewConnect(mock, mock, mock, validProof, false)
277277
request := mcp.CallToolRequest{}
278278
request.Params.Arguments = map[string]any{
279279
"docs_proof": validProof,
@@ -298,7 +298,7 @@ func TestConnect_Handler_EnginesFetchFailure(t *testing.T) {
298298
},
299299
}
300300

301-
connectTool := tools.NewConnect(mock, mock, mock, validProof)
301+
connectTool := tools.NewConnect(mock, mock, mock, validProof, false)
302302
request := mcp.CallToolRequest{}
303303
request.Params.Arguments = map[string]any{
304304
"docs_proof": validProof,
@@ -309,3 +309,73 @@ func TestConnect_Handler_EnginesFetchFailure(t *testing.T) {
309309
assert.Contains(t, err.Error(), "failed to discover engine resources")
310310
assert.Nil(t, result)
311311
}
312+
313+
func TestConnect_Handler_DisableResources(t *testing.T) {
314+
// Create test data
315+
accounts := []string{"account1"}
316+
databases := map[string][]string{
317+
"account1": {"db1"},
318+
}
319+
engines := map[string][]string{
320+
"account1": {"engine1"},
321+
}
322+
323+
// Create mock fetcher
324+
mock := &MockResourceFetcher{
325+
AccountsFunc: func(ctx context.Context, accountName string) ([]mcp.ResourceContents, error) {
326+
var resources []mcp.ResourceContents
327+
for _, acc := range accounts {
328+
resources = append(resources, createAccountResource(acc))
329+
}
330+
return resources, nil
331+
},
332+
DatabasesFunc: func(ctx context.Context, accountName, databaseName string) ([]mcp.ResourceContents, error) {
333+
var resources []mcp.ResourceContents
334+
for _, db := range databases[accountName] {
335+
resources = append(resources, createDatabaseResource(accountName, db))
336+
}
337+
return resources, nil
338+
},
339+
EnginesFunc: func(ctx context.Context, accountName, engineName string) ([]mcp.ResourceContents, error) {
340+
var resources []mcp.ResourceContents
341+
for _, eng := range engines[accountName] {
342+
resources = append(resources, createEngineResource(accountName, eng))
343+
}
344+
return resources, nil
345+
},
346+
}
347+
348+
// Create the tool with disableResources set to true
349+
connectTool := tools.NewConnect(mock, mock, mock, validProof, true)
350+
351+
// Execute the handler
352+
request := mcp.CallToolRequest{}
353+
request.Params.Arguments = map[string]any{
354+
"docs_proof": validProof,
355+
}
356+
result, err := connectTool.Handler(t.Context(), request)
357+
358+
// Assertions
359+
require.NoError(t, err)
360+
require.NotNil(t, result)
361+
assert.False(t, result.IsError)
362+
363+
// Calculate expected total resources
364+
expectedCount := len(accounts) // accounts
365+
for _, dbs := range databases {
366+
expectedCount += len(dbs) // databases
367+
}
368+
for _, engs := range engines {
369+
expectedCount += len(engs) // engines
370+
}
371+
372+
// Check if we got the expected number of resources
373+
assert.Len(t, result.Content, expectedCount)
374+
375+
// Verify the content contains text content instead of embedded resources
376+
for _, content := range result.Content {
377+
textContent, ok := content.(mcp.TextContent)
378+
require.True(t, ok, "Expected TextContent when disableResources is true")
379+
assert.NotEmpty(t, textContent.Text)
380+
}
381+
}

pkg/tools/docs.go

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,14 +21,16 @@ type DocsResourcesFetcher interface {
2121
// Docs represents a tool for fetching and returning Firebolt documentation.
2222
// It provides access to documentation articles that explain Firebolt concepts and functionality.
2323
type Docs struct {
24-
docsFetcher DocsResourcesFetcher // Fetches documentation resources
24+
docsFetcher DocsResourcesFetcher // Fetches documentation resources
25+
disableResources bool // Return text content instead of embedded resources
2526
}
2627

2728
// NewDocs creates a new instance of the Docs tool with the provided documentation fetcher.
2829
// It requires an implementation for fetching documentation articles.
29-
func NewDocs(docsFetcher DocsResourcesFetcher) *Docs {
30+
func NewDocs(docsFetcher DocsResourcesFetcher, disableResources bool) *Docs {
3031
return &Docs{
31-
docsFetcher: docsFetcher,
32+
docsFetcher: docsFetcher,
33+
disableResources: disableResources,
3234
}
3335
}
3436

@@ -103,7 +105,7 @@ func (t *Docs) Handler(ctx context.Context, request mcp.CallToolRequest) (*mcp.C
103105
return &mcp.CallToolResult{
104106
Result: mcp.Result{},
105107
Content: itertools.Map(results, func(i mcp.ResourceContents) mcp.Content {
106-
return mcp.NewEmbeddedResource(i)
108+
return textOrResourceContent(t.disableResources, i)
107109
}),
108110
IsError: false,
109111
}, nil

0 commit comments

Comments
 (0)