From 6cdc3e290a73b6c5d43905be0021e01785d08fc8 Mon Sep 17 00:00:00 2001 From: Mark Dai Date: Fri, 20 Jun 2025 22:05:24 +0800 Subject: [PATCH] forward client cookies --- pkg/server/server.go | 21 ++++++++++++--------- pkg/server/server_test.go | 16 ++++++++-------- 2 files changed, 20 insertions(+), 17 deletions(-) diff --git a/pkg/server/server.go b/pkg/server/server.go index 708b242..f80038c 100644 --- a/pkg/server/server.go +++ b/pkg/server/server.go @@ -435,7 +435,7 @@ func httpMethodPostHandler(w http.ResponseWriter, r *http.Request, toolSet *mcp. case "tools/list": respToSend = handleToolsListJSONRPC(connID, &req, toolSet) case "tools/call": - respToSend = handleToolCallJSONRPC(connID, &req, toolSet, cfg, r.Header) + respToSend = handleToolCallJSONRPC(connID, &req, toolSet, cfg, r.Header, r.Cookies()) default: log.Printf("Received unknown JSON-RPC method '%s' for %s", req.Method, connID) respToSend = createJSONRPCError(reqID, -32601, fmt.Sprintf("Method not found: %s", req.Method), nil) @@ -519,9 +519,9 @@ func handleToolsListJSONRPC(connID string, req *jsonRPCRequest, toolSet *mcp.Too } // executeToolCall performs the actual HTTP request based on the resolved operation and parameters. -// It now correctly handles API key injection based on the *cfg* parameter and -// accepts clientHeaders to forward request headers from the MCP client. -func executeToolCall(params *ToolCallParams, toolSet *mcp.ToolSet, cfg *config.Config, clientHeaders http.Header) (*http.Response, error) { +// It handles API key injection based on the *cfg* parameter and +// accepts clientHeaders and clientCookies to forward request headers and cookies from the MCP client. +func executeToolCall(params *ToolCallParams, toolSet *mcp.ToolSet, cfg *config.Config, clientHeaders http.Header, clientCookies []*http.Cookie) (*http.Response, error) { toolName := params.ToolName toolInput := params.Input // This is the map[string]interface{} from the client @@ -563,7 +563,10 @@ func executeToolCall(params *ToolCallParams, toolSet *mcp.ToolSet, cfg *config.C } } } - cookieParams := []*http.Cookie{} // For cookies to add + cookieParams := []*http.Cookie{} // For cookies to add + if clientCookies != nil { + cookieParams = append(cookieParams, clientCookies...) + } bodyData := make(map[string]interface{}) // For building the request body requestBodyRequired := operation.Method == "POST" || operation.Method == "PUT" || operation.Method == "PATCH" @@ -763,9 +766,9 @@ func executeToolCall(params *ToolCallParams, toolSet *mcp.ToolSet, cfg *config.C } // handleToolCallJSONRPC processes a 'tools/call' JSON-RPC request. It forwards -// any headers from the originating HTTP request via the clientHeaders argument -// so they can be applied to the outgoing API request. -func handleToolCallJSONRPC(connID string, req *jsonRPCRequest, toolSet *mcp.ToolSet, cfg *config.Config, clientHeaders http.Header) jsonRPCResponse { +// any headers and cookies from the originating HTTP request via the clientHeaders +// and clientCookies arguments so they can be applied to the outgoing API request. +func handleToolCallJSONRPC(connID string, req *jsonRPCRequest, toolSet *mcp.ToolSet, cfg *config.Config, clientHeaders http.Header, clientCookies []*http.Cookie) jsonRPCResponse { // req.Params is interface{}, but should contain json.RawMessage for tools/call rawParams, ok := req.Params.(json.RawMessage) if !ok { @@ -797,7 +800,7 @@ func handleToolCallJSONRPC(connID string, req *jsonRPCRequest, toolSet *mcp.Tool log.Printf("Executing tool '%s' for %s with input: %+v", params.ToolName, connID, params.Input) // --- Execute the actual tool call --- - httpResp, execErr := executeToolCall(¶ms, toolSet, cfg, clientHeaders) + httpResp, execErr := executeToolCall(¶ms, toolSet, cfg, clientHeaders, clientCookies) // --- Process Response --- var resultPayload ToolResultPayload diff --git a/pkg/server/server_test.go b/pkg/server/server_test.go index 3ea4e51..d86de1e 100644 --- a/pkg/server/server_test.go +++ b/pkg/server/server_test.go @@ -187,13 +187,13 @@ func TestHttpMethodPostHandler(t *testing.T) { expectedSyncStatus: http.StatusAccepted, expectedSyncBody: "Request accepted, response will be sent via SSE.\n", checkAsyncResponse: func(t *testing.T, resp jsonRPCResponse) { - assert.Equal(t, "call-post-err-1", resp.ID) - assert.Nil(t, resp.Error) - resultPayload, ok := resp.Result.(ToolResultPayload) - require.True(t, ok) - assert.True(t, resultPayload.IsError) - require.NotNil(t, resultPayload.Error) - assert.Contains(t, resultPayload.Error.Message, "operation details for tool 'nonexistent_tool' not found") + assert.Equal(t, "call-post-err-1", resp.ID) + assert.Nil(t, resp.Error) + resultPayload, ok := resp.Result.(ToolResultPayload) + require.True(t, ok) + assert.True(t, resultPayload.IsError) + require.NotNil(t, resultPayload.Error) + assert.Contains(t, resultPayload.Error.Message, "operation details for tool 'nonexistent_tool' not found") }, }, { @@ -665,7 +665,7 @@ func TestExecuteToolCall(t *testing.T) { } // --- Execute Function --- - httpResp, err := executeToolCall(&tc.params, toolSet, &testCfg, nil) // Use the potentially modified testCfg + httpResp, err := executeToolCall(&tc.params, toolSet, &testCfg, nil, nil) // Use the potentially modified testCfg // --- Assertions --- if tc.expectError {