refactor: print trace id

Author: iFurySt

cmd/mcp-gateway/main.go

@@ -181,25 +181,14 @@ func run() {
 		logger.Fatal("Failed to initialize auth service", zap.Error(err))
 	}
 
-	// Create server instance
-	server, err := core.NewServer(
-		logger,
-		cfg.Port,
-		store,
-		sessionStore,
-		a,
-		core.WithForwardConfig(cfg.Forward),
-		core.WithTraceCapture(cfg.Tracing.Capture),
-	)
-	if err != nil {
-		logger.Fatal("Failed to create server", zap.Error(err))
-	}
-
-	// Enable gin tracing middleware if tracing is configured
+	// Initialize tracing BEFORE creating server if enabled
+	var tracingServiceName string
 	if cfg.Tracing.Enabled {
 		if cfg.Tracing.ServiceName == "" {
 			cfg.Tracing.ServiceName = "mcp-gateway"
 		}
+		tracingServiceName = cfg.Tracing.ServiceName
+
 		shutdown, err := trace.InitTracing(ctx, &cfg.Tracing, logger)
 		if err != nil {
 			logger.Error("Failed to initialize tracing", zap.Error(err))
@@ -217,10 +206,22 @@ func run() {
 				zap.String("endpoint", cfg.Tracing.Endpoint),
 				zap.String("protocol", cfg.Tracing.Protocol),
 			)
-
 		}
+	}
 
-		server.EnableTracing(cfg.Tracing.ServiceName)
+	// Create server instance with tracing enabled from the start
+	server, err := core.NewServer(
+		logger,
+		cfg.Port,
+		store,
+		sessionStore,
+		a,
+		core.WithForwardConfig(cfg.Forward),
+		core.WithTraceCapture(cfg.Tracing.Capture),
+		core.WithTracing(tracingServiceName), // Register OTel middleware early
+	)
+	if err != nil {
+		logger.Fatal("Failed to create server", zap.Error(err))
 	}
 
 	err = server.RegisterRoutes(ctx)

cmd/mcp-gateway/main_more_test.go

@@ -200,4 +200,4 @@ func TestCommandTimeout(t *testing.T) {
 	case <-time.After(timeout):
 		t.Error("command took too long to complete")
 	}
-}
+}

cmd/mock-server/backend/http-server_routes_test.go

@@ -1,71 +1,70 @@
 package backend
 
 import (
-    "bytes"
-    "encoding/json"
-    "net/http"
-    "net/http/httptest"
-    "testing"
+	"bytes"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"testing"
 )
 
 func TestHTTPServer_UserRoutes(t *testing.T) {
-    s := NewHTTPServer()
+	s := NewHTTPServer()
 
-    // Create user
-    body := map[string]any{
-        "username": "u1",
-        "email":    "[email protected]",
-    }
-    bb, _ := json.Marshal(body)
-    w := httptest.NewRecorder()
-    r := httptest.NewRequest(http.MethodPost, "/users", bytes.NewReader(bb))
-    r.Header.Set("Content-Type", "application/json")
-    s.router.ServeHTTP(w, r)
-    if w.Code != http.StatusCreated {
-        t.Fatalf("create user status = %d", w.Code)
-    }
+	// Create user
+	body := map[string]any{
+		"username": "u1",
+		"email":    "[email protected]",
+	}
+	bb, _ := json.Marshal(body)
+	w := httptest.NewRecorder()
+	r := httptest.NewRequest(http.MethodPost, "/users", bytes.NewReader(bb))
+	r.Header.Set("Content-Type", "application/json")
+	s.router.ServeHTTP(w, r)
+	if w.Code != http.StatusCreated {
+		t.Fatalf("create user status = %d", w.Code)
+	}
 
-    // Get by email
-    w = httptest.NewRecorder()
-    r = httptest.NewRequest(http.MethodGet, "/users/email/[email protected]", nil)
-    s.router.ServeHTTP(w, r)
-    if w.Code != http.StatusOK {
-        t.Fatalf("get user status = %d", w.Code)
-    }
+	// Get by email
+	w = httptest.NewRecorder()
+	r = httptest.NewRequest(http.MethodGet, "/users/email/[email protected]", nil)
+	s.router.ServeHTTP(w, r)
+	if w.Code != http.StatusOK {
+		t.Fatalf("get user status = %d", w.Code)
+	}
 
-    // Update preferences
-    prefs := map[string]any{
-        "isPublic":  true,
-        "showEmail": false,
-        "theme":     "dark",
-        "tags":      []string{"a", "b"},
-        "settings":  map[string]any{"k": "v"},
-        "notifications": []Notification{{Type: "email", Channel: "system", Enabled: true, Frequency: 0}},
-    }
-    pb, _ := json.Marshal(prefs)
-    w = httptest.NewRecorder()
-    r = httptest.NewRequest(http.MethodPut, "/users/[email protected]/preferences", bytes.NewReader(pb))
-    r.Header.Set("Content-Type", "application/json")
-    s.router.ServeHTTP(w, r)
-    if w.Code != http.StatusOK {
-        t.Fatalf("update prefs status = %d", w.Code)
-    }
+	// Update preferences
+	prefs := map[string]any{
+		"isPublic":      true,
+		"showEmail":     false,
+		"theme":         "dark",
+		"tags":          []string{"a", "b"},
+		"settings":      map[string]any{"k": "v"},
+		"notifications": []Notification{{Type: "email", Channel: "system", Enabled: true, Frequency: 0}},
+	}
+	pb, _ := json.Marshal(prefs)
+	w = httptest.NewRecorder()
+	r = httptest.NewRequest(http.MethodPut, "/users/[email protected]/preferences", bytes.NewReader(pb))
+	r.Header.Set("Content-Type", "application/json")
+	s.router.ServeHTTP(w, r)
+	if w.Code != http.StatusOK {
+		t.Fatalf("update prefs status = %d", w.Code)
+	}
 
-    // Upload avatar missing url -> 400
-    w = httptest.NewRecorder()
-    r = httptest.NewRequest(http.MethodPost, "/users/[email protected]/avatar", nil)
-    s.router.ServeHTTP(w, r)
-    if w.Code != http.StatusBadRequest {
-        t.Fatalf("avatar missing url status = %d", w.Code)
-    }
+	// Upload avatar missing url -> 400
+	w = httptest.NewRecorder()
+	r = httptest.NewRequest(http.MethodPost, "/users/[email protected]/avatar", nil)
+	s.router.ServeHTTP(w, r)
+	if w.Code != http.StatusBadRequest {
+		t.Fatalf("avatar missing url status = %d", w.Code)
+	}
 
-    // Upload avatar with url
-    w = httptest.NewRecorder()
-    r = httptest.NewRequest(http.MethodPost, "/users/[email protected]/avatar", bytes.NewBufferString("url=https://img"))
-    r.Header.Set("Content-Type", "application/x-www-form-urlencoded")
-    s.router.ServeHTTP(w, r)
-    if w.Code != http.StatusOK {
-        t.Fatalf("avatar ok status = %d", w.Code)
-    }
+	// Upload avatar with url
+	w = httptest.NewRecorder()
+	r = httptest.NewRequest(http.MethodPost, "/users/[email protected]/avatar", bytes.NewBufferString("url=https://img"))
+	r.Header.Set("Content-Type", "application/x-www-form-urlencoded")
+	s.router.ServeHTTP(w, r)
+	if w.Code != http.StatusOK {
+		t.Fatalf("avatar ok status = %d", w.Code)
+	}
 }
-

cmd/mock-server/backend/mcp-server_more_test.go

@@ -1,47 +1,47 @@
 package backend
 
 import (
-    "context"
-    "encoding/json"
-    "testing"
+	"context"
+	"encoding/json"
+	"testing"
 
-    "github.com/mark3labs/mcp-go/mcp"
-    "github.com/stretchr/testify/assert"
+	"github.com/mark3labs/mcp-go/mcp"
+	"github.com/stretchr/testify/assert"
 )
 
 func TestNewMCPServer_And_GenerateResources(t *testing.T) {
-    srv := NewMCPServer()
-    if srv == nil {
-        t.Fatal("expected server instance")
-    }
+	srv := NewMCPServer()
+	if srv == nil {
+		t.Fatal("expected server instance")
+	}
 
-    rs := generateResources()
-    if len(rs) != 100 {
-        t.Fatalf("expected 100 resources, got %d", len(rs))
-    }
+	rs := generateResources()
+	if len(rs) != 100 {
+		t.Fatalf("expected 100 resources, got %d", len(rs))
+	}
 }
 
 func TestHandleSimpleAndComplexPrompts(t *testing.T) {
-    // simple
-    out, err := handleSimplePrompt(context.Background(), mcp.GetPromptRequest{})
-    assert.NoError(t, err)
-    assert.GreaterOrEqual(t, len(out.Messages), 1)
+	// simple
+	out, err := handleSimplePrompt(context.Background(), mcp.GetPromptRequest{})
+	assert.NoError(t, err)
+	assert.GreaterOrEqual(t, len(out.Messages), 1)
 
-    // complex with args
-    req := mcp.GetPromptRequest{Params: mcp.GetPromptParams{Arguments: map[string]string{
-        "temperature": "0.7",
-        "style":       "short",
-    }}}
-    out2, err := handleComplexPrompt(context.Background(), req)
-    assert.NoError(t, err)
-    assert.GreaterOrEqual(t, len(out2.Messages), 2)
+	// complex with args
+	req := mcp.GetPromptRequest{Params: mcp.GetPromptParams{Arguments: map[string]string{
+		"temperature": "0.7",
+		"style":       "short",
+	}}}
+	out2, err := handleComplexPrompt(context.Background(), req)
+	assert.NoError(t, err)
+	assert.GreaterOrEqual(t, len(out2.Messages), 2)
 }
 
 func TestHandleLongRunningOperation_ZeroSteps_NoProgress(t *testing.T) {
-    // Provide steps=0 to avoid progress notifications and timing
-    args, _ := json.Marshal(map[string]any{"duration": 0, "steps": 0})
-    req := mcp.CallToolRequest{Params: mcp.CallToolParams{Arguments: json.RawMessage(args), Meta: &mcp.Meta{}}}
-    out, err := handleLongRunningOperationTool(context.Background(), req)
-    assert.NoError(t, err)
-    assert.NotNil(t, out)
+	// Provide steps=0 to avoid progress notifications and timing
+	args, _ := json.Marshal(map[string]any{"duration": 0, "steps": 0})
+	req := mcp.CallToolRequest{Params: mcp.CallToolParams{Arguments: json.RawMessage(args), Meta: &mcp.Meta{}}}
+	out, err := handleLongRunningOperationTool(context.Background(), req)
+	assert.NoError(t, err)
+	assert.NotNil(t, out)
 }

internal/apiserver/handler/auth_missing_test.go

@@ -214,7 +214,7 @@ func TestHandler_UpdateUserTenants_Success(t *testing.T) {
 	req, _ := http.NewRequest(http.MethodPut, "/users/tenants", nil)
 	r.ServeHTTP(w, req)
 	assert.Equal(t, http.StatusOK, w.Code)
-	assert.Contains(t, f.removedIds, uint(1)) // removed tenant 1
+	assert.Contains(t, f.removedIds, uint(1))       // removed tenant 1
 	assert.Contains(t, f.addedPairs, [2]uint{5, 3}) // added tenant 3 to user 5
 }
 
@@ -272,4 +272,4 @@ func TestHandler_UpdateUserTenants_UserNotFound(t *testing.T) {
 	req, _ := http.NewRequest(http.MethodPut, "/users/tenants", nil)
 	r.ServeHTTP(w, req)
 	assert.Equal(t, http.StatusInternalServerError, w.Code)
-}
+}