Record session id on init

Author: jboolean

contrib/mark3labs/mcp-go/hooks.go

@@ -66,9 +66,18 @@ func (h *Hooks) AddHooks(hooks *server.Hooks) {
 // onBeforeInitialize is called before initialization is executed.
 func (h *Hooks) onBeforeInitialize(ctx context.Context, id any, request *mcp.InitializeRequest) {
 	taskSpan, _ := llmobs.StartTaskSpan(ctx, "mcp.initialize")
+	tagWithSessionID(ctx, taskSpan)
 	h.spanCache.Set(id, taskSpan, ttlcache.DefaultTTL)
 }
 
+func tagWithSessionID(ctx context.Context, span llmobs.Span) {
+	session := server.ClientSessionFromContext(ctx)
+	if session != nil {
+		sessionID := session.SessionID()
+		span.Annotate(llmobs.WithAnnotatedTags(map[string]string{"mcp_session_id": sessionID}))
+	}
+}
+
 // finishSpanWithIO retrieves a span from cache, annotates it with input/output, and finishes it.
 func finishSpanWithIO[Req any, Res any](h *Hooks, id any, request Req, result Res) {
 	if item := h.spanCache.Get(id); item != nil {
@@ -93,12 +102,8 @@ func (h *Hooks) onAfterInitialize(ctx context.Context, id any, request *mcp.Init
 // onBeforeCallTool is called before a tool is executed.
 func (h *Hooks) onBeforeCallTool(ctx context.Context, id any, request *mcp.CallToolRequest) {
 	toolSpan, _ := llmobs.StartToolSpan(ctx, request.Params.Name)
-	session := server.ClientSessionFromContext(ctx)
-	if session != nil {
-		sessionID := session.SessionID()
-		// toolSpan.Annotate(llmobs.WithAnnotatedSessionID(sessionID))
-		toolSpan.Annotate(llmobs.WithAnnotatedTags(map[string]string{"mcp_session_id": sessionID}))
-	}
+
+	tagWithSessionID(ctx, toolSpan)
 
 	h.spanCache.Set(id, toolSpan, ttlcache.DefaultTTL)
 }

contrib/mark3labs/mcp-go/hooks_test.go

@@ -68,8 +68,13 @@ func TestIntegrationSessionInitialize(t *testing.T) {
 	srv := server.NewMCPServer("test-server", "1.0.0",
 		server.WithHooks(hooks))
 
-	// Create initialize request
+	// Create initialize request with session
 	ctx := context.Background()
+	sessionID := "test-session-init"
+	session := &mockSession{id: sessionID}
+	session.Initialize()
+	ctx = srv.WithContext(ctx, session)
+
 	initRequest := `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05","capabilities":{},"clientInfo":{"name":"test-client","version":"1.0.0"}}}`
 
 	// Send initialize request
@@ -96,6 +101,9 @@ func TestIntegrationSessionInitialize(t *testing.T) {
 	assert.Equal(t, "mcp.initialize", taskSpan.Name)
 	assert.Equal(t, "task", taskSpan.Meta["span.kind"])
 
+	// Verify session ID tag is set
+	assert.Contains(t, taskSpan.Tags, "mcp_session_id:test-session-init")
+
 	// Verify input/output annotations exist
 	assert.Contains(t, taskSpan.Meta, "input")
 	assert.Contains(t, taskSpan.Meta, "output")
@@ -166,11 +174,16 @@ func TestIntegrationToolCallSuccess(t *testing.T) {
 	session.Initialize()
 	ctx = srv.WithContext(ctx, session)
 
+	// Send initialize request first
+	initRequest := `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05","capabilities":{},"clientInfo":{"name":"test-client","version":"1.0.0"}}}`
+	response := srv.HandleMessage(ctx, []byte(initRequest))
+	assert.NotNil(t, response)
+
 	// Create the actual MCP tools/call request as JSON
 	toolCallRequest := `{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"calculator","arguments":{"operation":"add","x":5,"y":3}}}`
 
 	// Call the actual MCP server's HandleMessage method
-	response := srv.HandleMessage(ctx, []byte(toolCallRequest))
+	response = srv.HandleMessage(ctx, []byte(toolCallRequest))
 	assert.NotNil(t, response)
 
 	// Marshal response to check it
@@ -184,17 +197,31 @@ func TestIntegrationToolCallSuccess(t *testing.T) {
 	assert.Equal(t, "2.0", resp["jsonrpc"])
 	assert.NotNil(t, resp["result"])
 
-	// Verify LLMObs tool span was created
-	spans := tt.WaitForLLMObsSpans(t, 1)
-	require.Len(t, spans, 1)
+	// Verify LLMObs spans were created for both initialize and tool call
+	spans := tt.WaitForLLMObsSpans(t, 2)
+	require.Len(t, spans, 2)
+
+	// Find initialize and tool spans
+	var initSpan, toolSpan *testtracer.LLMObsSpan
+	for i := range spans {
+		if spans[i].Name == "mcp.initialize" {
+			initSpan = &spans[i]
+		} else if spans[i].Name == "calculator" {
+			toolSpan = &spans[i]
+		}
+	}
+
+	require.NotNil(t, initSpan, "initialize span not found")
+	require.NotNil(t, toolSpan, "tool span not found")
+
+	// Verify both spans have the same session ID tag
+	expectedTag := "mcp_session_id:test-session-123"
+	assert.Contains(t, initSpan.Tags, expectedTag)
+	assert.Contains(t, toolSpan.Tags, expectedTag)
 
-	toolSpan := spans[0]
 	assert.Equal(t, "calculator", toolSpan.Name)
 	assert.Equal(t, "tool", toolSpan.Meta["span.kind"])
 
-	// Verify session ID tag is set
-	assert.Contains(t, toolSpan.Tags, "mcp_session_id:test-session-123")
-
 	// Verify input/output annotations exist
 	assert.Contains(t, toolSpan.Meta, "input")
 	assert.Contains(t, toolSpan.Meta, "output")