Skip to content

Commit e5b87c3

Browse files
MieszkoWasniowskiMieszkoWasniowski
authored
Merge pull request #16 from collibra/clean-session-handler
refactor: cleanup collibraHost and session context handling
2 parents 839e183 + a507c3e commit e5b87c3

8 files changed

Lines changed: 41 additions & 72 deletions

File tree

‎cmd/chip/http.go‎

Lines changed: 10 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -58,29 +58,26 @@ func newCollibraClient(config *Config) *http.Client {
5858
}
5959

6060
func (c *collibraClient) RoundTrip(request *http.Request) (*http.Response, error) {
61-
reqClone := request.Clone(request.Context())
62-
toolRequest, err := chip.GetCallToolRequest(reqClone.Context())
63-
if err != nil {
64-
return nil, err
65-
}
6661
if c.config.Api.Url == "" {
6762
return nil, fmt.Errorf("API URL is not configured")
6863
}
64+
baseURL, err := url.Parse(c.config.Api.Url)
65+
if err != nil {
66+
return nil, fmt.Errorf("invalid API URL configuration: %w", err)
67+
}
68+
reqClone := request.Clone(request.Context())
69+
toolRequest, ok := chip.GetCallToolRequest(reqClone.Context())
70+
if !ok {
71+
return nil, fmt.Errorf("toolRequest not found in ctx")
72+
}
6973
if c.config.Api.Username != "" && c.config.Api.Password != "" {
7074
reqClone.SetBasicAuth(c.config.Api.Username, c.config.Api.Password)
7175
} else {
7276
copyHeader(toolRequest, reqClone, "Authorization")
7377
}
74-
reqClone.Header.Set("X-MCP-Session-Id", chip.GetSessionId(toolRequest))
78+
reqClone.Header.Set("X-MCP-Session-Id", chip.GetSessionId(reqClone.Context()))
7579
reqClone.Header.Set("X-MCP-Tool-Name", toolRequest.Params.Name)
7680
reqClone.Header.Set("traceparent", generateTraceParent())
77-
baseURL, err := url.Parse(c.config.Api.Url)
78-
if err != nil {
79-
return nil, fmt.Errorf("invalid API URL configuration: %w", err)
80-
}
81-
if toolRequest.GetExtra() != nil {
82-
toolRequest.Extra.Header.Set("collibraUrl", c.config.Api.Url)
83-
}
8481
reqClone.URL.Scheme = baseURL.Scheme
8582
reqClone.URL.Host = baseURL.Host
8683
reqClone.URL.Path = path.Join(baseURL.Path, request.URL.Path)

‎cmd/chip/main.go‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -28,9 +28,8 @@ func main() {
2828
}
2929

3030
client := newCollibraClient(config)
31-
server := chip.NewServer(chip.WithToolMiddleware(chip.ToolMiddlewareFunc(logMiddleware)))
31+
server := chip.NewServer(chip.WithToolMiddleware(chip.ToolMiddlewareFunc(setCollibraHost(config.Api.Url))))
3232
toolConfig := &chip.ToolConfig{
33-
CollibraUrl: config.Api.Url,
3433
EnabledTools: config.Mcp.EnabledTools,
3534
DisabledTools: config.Mcp.DisabledTools,
3635
}
@@ -88,7 +87,10 @@ func runHttpServer(mode string, server *chip.Server, port int) {
8887
}
8988
}
9089

91-
func logMiddleware(ctx context.Context, toolRequest *mcp.CallToolRequest, next chip.CallToolFunc) (*mcp.CallToolResult, error) {
92-
slog.InfoContext(ctx, fmt.Sprintf("Calling tool: %s", toolRequest.Params.Name), "tool_name", toolRequest.Params.Name)
93-
return next(ctx, toolRequest)
90+
func setCollibraHost(collibraHost string) func(ctx context.Context, toolRequest *mcp.CallToolRequest, next chip.CallToolFunc) (*mcp.CallToolResult, error) {
91+
return func(ctx context.Context, toolRequest *mcp.CallToolRequest, next chip.CallToolFunc) (*mcp.CallToolResult, error) {
92+
ctx = chip.SetCollibraHost(ctx, collibraHost)
93+
slog.InfoContext(ctx, fmt.Sprintf("Calling tool: %s", toolRequest.Params.Name), "tool_name", toolRequest.Params.Name)
94+
return next(ctx, toolRequest)
95+
}
9496
}

‎pkg/chip/context.go‎

Lines changed: 15 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@ package chip
22

33
import (
44
"context"
5-
"errors"
65

76
"github.com/google/uuid"
87
"github.com/modelcontextprotocol/go-sdk/mcp"
@@ -12,37 +11,31 @@ type contextKey int
1211

1312
const (
1413
callToolRequestKey contextKey = iota
15-
toolConfigKey
14+
collibraHostKey
1615
)
1716

18-
func SetToolConfig(ctx context.Context, toolConfig *ToolConfig) context.Context {
19-
return context.WithValue(ctx, toolConfigKey, toolConfig)
20-
}
21-
2217
func SetCallToolRequest(ctx context.Context, toolRequest *mcp.CallToolRequest) context.Context {
2318
return context.WithValue(ctx, callToolRequestKey, toolRequest)
2419
}
2520

26-
func GetCallToolRequest(ctx context.Context) (*mcp.CallToolRequest, error) {
21+
func GetCallToolRequest(ctx context.Context) (*mcp.CallToolRequest, bool) {
2722
toolRequest, ok := ctx.Value(callToolRequestKey).(*mcp.CallToolRequest)
28-
if !ok || toolRequest == nil {
29-
return nil, errors.New("CallToolRequest not found in ctx")
30-
}
31-
return toolRequest, nil
23+
return toolRequest, ok
3224
}
3325

34-
func GetToolConfig(ctx context.Context) (*ToolConfig, error) {
35-
config, ok := ctx.Value(toolConfigKey).(*ToolConfig)
36-
if !ok || config == nil {
37-
return nil, errors.New("ToolConfig not found in ctx")
38-
}
39-
return config, nil
26+
func SetCollibraHost(ctx context.Context, collibraHost string) context.Context {
27+
return context.WithValue(ctx, collibraHostKey, collibraHost)
28+
}
29+
30+
func GetCollibraHost(ctx context.Context) (string, bool) {
31+
collibraHost, ok := ctx.Value(collibraHostKey).(string)
32+
return collibraHost, ok
4033
}
4134

42-
func GetSessionId(toolRequest *mcp.CallToolRequest) string {
43-
sessionId := toolRequest.GetSession().ID()
44-
if sessionId == "" {
45-
return uuid.New().String()
35+
func GetSessionId(ctx context.Context) string {
36+
toolRequest, ok := GetCallToolRequest(ctx)
37+
if ok {
38+
return toolRequest.GetSession().ID()
4639
}
47-
return sessionId
40+
return uuid.New().String()
4841
}

‎pkg/chip/server.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package chip
22

33
import (
44
"context"
5+
"fmt"
56
"log/slog"
67
"slices"
78

@@ -45,7 +46,6 @@ func NewServer(opts ...ServerOption) *Server {
4546
}
4647

4748
type ToolConfig struct {
48-
CollibraUrl string
4949
EnabledTools []string
5050
DisabledTools []string
5151
}
@@ -73,7 +73,8 @@ type Tool[In, Out any] struct {
7373
ToolHandler ToolHandlerFunc[In, Out]
7474
}
7575

76-
func RegisterTool[In, Out any](s *Server, tool *Tool[In, Out], toolConfig *ToolConfig) {
76+
func RegisterTool[In, Out any](s *Server, tool *Tool[In, Out]) {
77+
slog.Info(fmt.Sprintf("Registering tool: %s", tool.Tool.Name))
7778
handler := func(ctx context.Context, toolRequest *mcp.CallToolRequest, input In) (*mcp.CallToolResult, Out, error) {
7879
var capturedOutput Out
7980

@@ -95,7 +96,6 @@ func RegisterTool[In, Out any](s *Server, tool *Tool[In, Out], toolConfig *ToolC
9596
}
9697

9798
ctx = SetCallToolRequest(ctx, toolRequest)
98-
ctx = SetToolConfig(ctx, toolConfig)
9999
res, err := middlewareChain(ctx, toolRequest)
100100

101101
return res, capturedOutput, err

‎pkg/chip/version.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,3 @@
11
package chip
22

3-
var Version = "0.0.21-SNAPSHOT"
3+
var Version = "0.0.22-SNAPSHOT"

‎pkg/tools/get_asset_details.go‎

Lines changed: 3 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -67,32 +67,15 @@ func handleAssetDetails(collibraClient *http.Client) chip.ToolHandlerFunc[AssetD
6767
}, nil
6868
}
6969

70-
collibraUrl, err := getCollibraUrl(ctx)
71-
if err != nil {
70+
collibraHost, ok := chip.GetCollibraHost(ctx)
71+
if !ok {
7272
slog.WarnContext(ctx, "Collibra instance URL unknown, links will be rendered without host")
7373
}
7474

7575
return AssetDetailsOutput{
7676
Asset: &assets[0],
7777
Found: true,
78-
Link: fmt.Sprintf("%s/asset/%s", collibraUrl, assetUUID),
78+
Link: fmt.Sprintf("%s/asset/%s", strings.TrimSuffix(collibraHost, "/"), assetUUID),
7979
}, nil
8080
}
8181
}
82-
83-
func getCollibraUrl(ctx context.Context) (string, error) {
84-
toolRequest, err := chip.GetCallToolRequest(ctx)
85-
if err != nil {
86-
return "", err
87-
}
88-
if toolRequest.GetExtra() != nil {
89-
if url := toolRequest.Extra.Header.Get("collibraUrl"); url != "" {
90-
return strings.TrimSuffix(url, "/"), nil
91-
}
92-
}
93-
config, err := chip.GetToolConfig(ctx)
94-
if err != nil {
95-
return "", err
96-
}
97-
return strings.TrimSuffix(config.CollibraUrl, "/"), nil
98-
}

‎pkg/tools/get_asset_details_test.go‎

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@ import (
66
"net/http/httptest"
77
"testing"
88

9-
"github.com/collibra/chip/pkg/chip"
109
"github.com/collibra/chip/pkg/clients"
1110
"github.com/collibra/chip/pkg/tools"
1211
"github.com/google/uuid"
@@ -32,12 +31,7 @@ func TestGetAssetDetails(t *testing.T) {
3231

3332
client := newClient(server)
3433

35-
config := &chip.ToolConfig{
36-
CollibraUrl: server.URL,
37-
}
38-
ctx := chip.SetToolConfig(context.Background(), config)
39-
40-
output, err := tools.NewAssetDetailsTool(client).ToolHandler(ctx, tools.AssetDetailsInput{
34+
output, err := tools.NewAssetDetailsTool(client).ToolHandler(context.Background(), tools.AssetDetailsInput{
4135
AssetID: assetId.String(),
4236
})
4337
if err != nil {

‎pkg/tools/tools_register.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,6 @@ func RegisterAll(server *chip.Server, client *http.Client, toolConfig *chip.Tool
2323

2424
func toolRegister[In, Out any](server *chip.Server, toolConfig *chip.ToolConfig, tool *chip.Tool[In, Out]) {
2525
if toolConfig.IsToolEnabled(tool.Tool.Name) {
26-
chip.RegisterTool(server, tool, toolConfig)
26+
chip.RegisterTool(server, tool)
2727
}
2828
}

0 commit comments

Comments
 (0)