review 2

Author: pawbana

bridge_integration_test.go

@@ -974,11 +974,11 @@ func setupInjectedToolTest(t *testing.T, fixture []byte, streaming bool, configu
 
 	recorderClient := &mockRecorderClient{}
 
-	// Setup MCP tools.
-	tools := setupMCPServerProxiesForTest(t, testTracer)
+	// Setup MCP proxies.
+	proxies := setupMCPServerProxiesForTest(t, testTracer)
 
 	// Configure the bridge with injected tools.
-	mcpMgr := mcp.NewServerProxyManager(tools, testTracer)
+	mcpMgr := mcp.NewServerProxyManager(proxies, testTracer)
 	require.NoError(t, mcpMgr.Init(ctx))
 	b, err := configureFn(mockSrv.URL, recorderClient, mcpMgr)
 	require.NoError(t, err)
@@ -1005,7 +1005,7 @@ func setupInjectedToolTest(t *testing.T, fixture []byte, streaming bool, configu
 		return mockSrv.callCount.Load() == 2
 	}, time.Second*10, time.Millisecond*50)
 
-	return recorderClient, tools, resp
+	return recorderClient, proxies, resp
 }
 
 func TestErrorHandling(t *testing.T) {

go.mod

@@ -27,7 +27,6 @@ require (
 )
 
 require (
-	github.com/google/go-cmp v0.7.0
 	go.opentelemetry.io/otel v1.38.0
 	go.opentelemetry.io/otel/sdk v1.38.0
 	go.opentelemetry.io/otel/trace v1.38.0

intercept_anthropic_messages_blocking.go

@@ -88,7 +88,7 @@ func (i *AnthropicMessagesBlockingInterception) ProcessRequest(w http.ResponseWr
 
 	for {
 		// TODO add outer loop span (https://github.com/coder/aibridge/issues/67)
-		resp, err = i.traceNewMessage(ctx, svc, messages) // traces svc.New(ctx, msgParams) call
+		resp, err = i.newMessage(ctx, svc, messages)
 		if err != nil {
 			if isConnError(err) {
 				// Can't write a response, just error out.
@@ -298,7 +298,7 @@ func (i *AnthropicMessagesBlockingInterception) ProcessRequest(w http.ResponseWr
 	return nil
 }
 
-func (i *AnthropicMessagesBlockingInterception) traceNewMessage(ctx context.Context, svc anthropic.MessageService, msgParams anthropic.MessageNewParams) (_ *anthropic.Message, outErr error) {
+func (i *AnthropicMessagesBlockingInterception) newMessage(ctx context.Context, svc anthropic.MessageService, msgParams anthropic.MessageNewParams) (_ *anthropic.Message, outErr error) {
 	ctx, span := i.tracer.Start(ctx, "Intercept.ProcessRequest.Upstream", trace.WithAttributes(tracing.InterceptionAttributesFromContext(ctx)...))
 	defer tracing.EndSpanErr(span, &outErr)
 

intercept_openai_chat_blocking.go

@@ -76,7 +76,7 @@ func (i *OpenAIBlockingChatInterception) ProcessRequest(w http.ResponseWriter, r
 		var opts []option.RequestOption
 		opts = append(opts, option.WithRequestTimeout(time.Second*60)) // TODO: configurable timeout
 
-		completion, err = i.traceChatCompletionsNew(ctx, svc, opts) // traces svc.New(ctx, i.req.ChatCompletionNewParams, opts...) call
+		completion, err = i.newChatCompletion(ctx, svc, opts)
 		if err != nil {
 			break
 		}
@@ -233,7 +233,7 @@ func (i *OpenAIBlockingChatInterception) ProcessRequest(w http.ResponseWriter, r
 	return nil
 }
 
-func (i *OpenAIBlockingChatInterception) traceChatCompletionsNew(ctx context.Context, svc openai.ChatCompletionService, opts []option.RequestOption) (_ *openai.ChatCompletion, outErr error) {
+func (i *OpenAIBlockingChatInterception) newChatCompletion(ctx context.Context, svc openai.ChatCompletionService, opts []option.RequestOption) (_ *openai.ChatCompletion, outErr error) {
 	ctx, span := i.tracer.Start(ctx, "Intercept.ProcessRequest.Upstream", trace.WithAttributes(tracing.InterceptionAttributesFromContext(ctx)...))
 	defer tracing.EndSpanErr(span, &outErr)
 

mcp/tool.go

@@ -15,7 +15,7 @@ import (
 )
 
 const (
-	maxSpanInputAttrLen   = 100
+	maxSpanInputAttrLen   = 100    // truncates tool.Call span input attribute to first `maxSpanInputAttrLen` letters
 	injectedToolPrefix    = "bmcp" // "bridged MCP"
 	injectedToolDelimiter = "_"
 )
@@ -63,7 +63,7 @@ func (t *Tool) Call(ctx context.Context, input any, tracer trace.Tracer) (_ *mcp
 	} else {
 		strJson := string(inputJson)
 		if len(strJson) > maxSpanInputAttrLen {
-			strJson = strJson[:100]
+			strJson = strJson[:maxSpanInputAttrLen]
 		}
 		span.SetAttributes(attribute.String(tracing.MCPInput, strJson))
 	}

trace_integration_test.go

@@ -5,6 +5,8 @@ import (
 	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"slices"
+	"strings"
 	"testing"
 	"time"
 
@@ -13,8 +15,6 @@ import (
 	"github.com/coder/aibridge"
 	"github.com/coder/aibridge/mcp"
 	"github.com/coder/aibridge/tracing"
-	"github.com/google/go-cmp/cmp"
-	"github.com/google/go-cmp/cmp/cmpopts"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/tidwall/gjson"
@@ -670,13 +670,12 @@ func TestTracePassthrough(t *testing.T) {
 	require.Len(t, spans, 1)
 
 	assert.Equal(t, spans[0].Name(), "Passthrough")
-	attrs := []attribute.KeyValue{
-		attribute.String(tracing.PassthroughURL, "/v1/models"),
+	want := []attribute.KeyValue{
 		attribute.String(tracing.PassthroughMethod, "GET"),
+		attribute.String(tracing.PassthroughURL, "/v1/models"),
 	}
-	if attrDiff := cmp.Diff(spans[0].Attributes(), attrs, cmpopts.EquateComparable(attribute.KeyValue{}), cmpopts.SortSlices(cmpAttrKeyVal)); attrDiff != "" {
-		t.Errorf("unexpectet attrs diff: %s", attrDiff)
-	}
+	got := slices.SortedFunc(slices.Values(spans[0].Attributes()), cmpAttrKeyVal)
+	require.Equal(t, want, got)
 }
 
 func TestNewServerProxyManagerTraces(t *testing.T) {
@@ -714,8 +713,8 @@ func TestNewServerProxyManagerTraces(t *testing.T) {
 	verifyTraces(t, sr, []expectTrace{{"StreamableHTTPServerProxy.Init.fetchTools", 1, codes.Unset}}, attrs)
 }
 
-func cmpAttrKeyVal(a attribute.KeyValue, b attribute.KeyValue) bool {
-	return a.Key < b.Key
+func cmpAttrKeyVal(a attribute.KeyValue, b attribute.KeyValue) int {
+	return strings.Compare(string(a.Key), string(b.Key))
 }
 
 // checks counts of traces with given name, status and attributes
@@ -729,9 +728,9 @@ func verifyTraces(t *testing.T, spanRecorder *tracetest.SpanRecorder, expect []e
 				continue
 			}
 			found++
-			if attrDiff := cmp.Diff(s.Attributes(), attrs, cmpopts.EquateEmpty(), cmpopts.EquateComparable(attribute.KeyValue{}), cmpopts.SortSlices(cmpAttrKeyVal)); attrDiff != "" {
-				t.Errorf("unexpectet attrs for span named: %v, diff: %s", e.name, attrDiff)
-			}
+			want := slices.SortedFunc(slices.Values(attrs), cmpAttrKeyVal)
+			got := slices.SortedFunc(slices.Values(s.Attributes()), cmpAttrKeyVal)
+			require.Equal(t, want, got)
 			assert.Equalf(t, e.status, s.Status().Code, "unexpected status for trace naned: %v got: %v want: %v", e.name, s.Status().Code, e.status)
 		}
 		if found != e.count {

tracing/tracing.go

@@ -63,6 +63,17 @@ func RequestBridgeAttributesFromContext(ctx context.Context) []attribute.KeyValu
 	return attrs
 }
 
+// EndSpanErr ends given span and sets Error status if error is not nil
+// uses pointer to error because defer evaluates function arguments
+// when defer statement is executed not when deferred function is called
+//
+// example usage:
+//
+//	func Example() (result any, outErr error) {
+//	    _, span := tracer.Start(...)
+//	    defer tracing.EndSpanErr(span, &outErr)
+//
+// }
 func EndSpanErr(span trace.Span, err *error) {
 	if span == nil {
 		return