add tool call trace test, fixed openAI tool call attrs?

Author: pawbana

aibtrace/aibtrace.go

@@ -18,7 +18,7 @@ const (
 	RequestPath = "request_path"
 
 	InterceptionID = "interception_id"
-	UserID         = "user_id"
+	InitiatorID    = "user_id"
 	Provider       = "provider"
 	Model          = "model"
 	Streaming      = "streaming"

bridge_integration_test.go

@@ -442,9 +442,9 @@ func TestOpenAIChatCompletions(t *testing.T) {
 
 				require.Len(t, recorderClient.toolUsages, 1)
 				assert.Equal(t, "read_file", recorderClient.toolUsages[0].Tool)
-				require.IsType(t, map[string]any{}, recorderClient.toolUsages[0].Args)
+				require.IsType(t, "", recorderClient.toolUsages[0].Args)
 				require.Contains(t, recorderClient.toolUsages[0].Args, "path")
-				assert.Equal(t, "README.md", recorderClient.toolUsages[0].Args.(map[string]any)["path"])
+				assert.Equal(t, "README.md", gjson.Get(recorderClient.toolUsages[0].Args.(string), "path").Str)
 
 				require.Len(t, recorderClient.userPrompts, 1)
 				assert.Equal(t, "how large is the README.md file in my current path", recorderClient.userPrompts[0].Prompt)
@@ -765,7 +765,7 @@ func TestAnthropicInjectedTools(t *testing.T) {
 			}
 
 			// Build the requirements & make the assertions which are common to all providers.
-			recorderClient, resp := setupInjectedToolTest(t, antSingleInjectedTool, streaming, configureFn, createAnthropicMessagesReq)
+			recorderClient, _, resp := setupInjectedToolTest(t, antSingleInjectedTool, streaming, configureFn, createAnthropicMessagesReq)
 
 			// Ensure expected tool was invoked with expected input.
 			require.Len(t, recorderClient.toolUsages, 1)
@@ -847,16 +847,13 @@ func TestOpenAIInjectedTools(t *testing.T) {
 			}
 
 			// Build the requirements & make the assertions which are common to all providers.
-			recorderClient, resp := setupInjectedToolTest(t, oaiSingleInjectedTool, streaming, configureFn, createOpenAIChatCompletionsReq)
+			recorderClient, _, resp := setupInjectedToolTest(t, oaiSingleInjectedTool, streaming, configureFn, createOpenAIChatCompletionsReq)
 
 			// Ensure expected tool was invoked with expected input.
 			require.Len(t, recorderClient.toolUsages, 1)
 			require.Equal(t, mockToolName, recorderClient.toolUsages[0].Tool)
-			expected, err := json.Marshal(map[string]any{"owner": "admin"})
-			require.NoError(t, err)
-			actual, err := json.Marshal(recorderClient.toolUsages[0].Args)
-			require.NoError(t, err)
-			require.EqualValues(t, expected, actual)
+			expected := "{\"owner\":\"admin\"}"
+			require.EqualValues(t, expected, recorderClient.toolUsages[0].Args)
 
 			var (
 				content *openai.ChatCompletionChoice
@@ -932,7 +929,7 @@ func TestOpenAIInjectedTools(t *testing.T) {
 
 // setupInjectedToolTest abstracts the common aspects required for the Test*InjectedTools tests.
 // Kinda fugly right now, we can refactor this later.
-func setupInjectedToolTest(t *testing.T, fixture []byte, streaming bool, configureFn func(addr string, client aibridge.Recorder, srvProxyMgr *mcp.ServerProxyManager) (*aibridge.RequestBridge, error), createRequestFn func(*testing.T, string, []byte) *http.Request) (*mockRecorderClient, *http.Response) {
+func setupInjectedToolTest(t *testing.T, fixture []byte, streaming bool, configureFn func(addr string, client aibridge.Recorder, srvProxyMgr *mcp.ServerProxyManager) (*aibridge.RequestBridge, error), createRequestFn func(*testing.T, string, []byte) *http.Request) (*mockRecorderClient, map[string]mcp.ServerProxier, *http.Response) {
 	t.Helper()
 
 	arc := txtar.Parse(fixture)
@@ -1008,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, resp
+	return recorderClient, tools, resp
 }
 
 func TestErrorHandling(t *testing.T) {

intercept_anthropic_messages_base.go

@@ -67,7 +67,7 @@ func (s *AnthropicMessagesInterceptionBase) baseTraceAttributes(r *http.Request,
 	return []attribute.KeyValue{
 		attribute.String(aibtrace.RequestPath, r.URL.Path),
 		attribute.String(aibtrace.InterceptionID, s.id.String()),
-		attribute.String(aibtrace.UserID, actorFromContext(r.Context()).id),
+		attribute.String(aibtrace.InitiatorID, actorFromContext(r.Context()).id),
 		attribute.String(aibtrace.Provider, ProviderAnthropic),
 		attribute.String(aibtrace.Model, s.Model()),
 		attribute.Bool(aibtrace.Streaming, streaming),

intercept_openai_chat_base.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"encoding/json"
 	"net/http"
-	"strings"
 
 	aibtrace "github.com/coder/aibridge/aibtrace"
 	"github.com/coder/aibridge/mcp"
@@ -51,7 +50,7 @@ func (s *OpenAIChatInterceptionBase) baseTraceAttributes(r *http.Request, stream
 	return []attribute.KeyValue{
 		attribute.String(aibtrace.RequestPath, r.URL.Path),
 		attribute.String(aibtrace.InterceptionID, s.id.String()),
-		attribute.String(aibtrace.UserID, actorFromContext(r.Context()).id),
+		attribute.String(aibtrace.InitiatorID, actorFromContext(r.Context()).id),
 		attribute.String(aibtrace.Provider, ProviderOpenAI),
 		attribute.String(aibtrace.Model, s.Model()),
 		attribute.Bool(aibtrace.Streaming, streaming),
@@ -105,18 +104,6 @@ func (i *OpenAIChatInterceptionBase) injectTools() {
 	}
 }
 
-func (i *OpenAIChatInterceptionBase) unmarshalArgs(in string) (args ToolArgs) {
-	if len(strings.TrimSpace(in)) == 0 {
-		return args // An empty string will fail JSON unmarshaling.
-	}
-
-	if err := json.Unmarshal([]byte(in), &args); err != nil {
-		i.logger.Warn(context.Background(), "failed to unmarshal tool args", slog.Error(err))
-	}
-
-	return args
-}
-
 // writeUpstreamError marshals and writes a given error.
 func (i *OpenAIChatInterceptionBase) writeUpstreamError(w http.ResponseWriter, oaiErr *OpenAIErrorResponse) {
 	if oaiErr == nil {

intercept_openai_chat_blocking.go

@@ -1,7 +1,6 @@
 package aibridge
 
 import (
-	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
@@ -120,7 +119,7 @@ func (i *OpenAIBlockingChatInterception) ProcessRequest(w http.ResponseWriter, r
 						InterceptionID: i.ID().String(),
 						MsgID:          completion.ID,
 						Tool:           toolCall.Function.Name,
-						Args:           i.unmarshalArgs(toolCall.Function.Arguments),
+						Args:           toolCall.Function.Arguments,
 						Injected:       false,
 					})
 				}
@@ -151,20 +150,13 @@ func (i *OpenAIBlockingChatInterception) ProcessRequest(w http.ResponseWriter, r
 				appendedPrevMsg = true
 			}
 
-			var (
-				args map[string]string
-				buf  bytes.Buffer
-			)
-			_ = json.NewEncoder(&buf).Encode(tc.Function.Arguments)
-			_ = json.NewDecoder(&buf).Decode(&args)
-			res, err := tool.Call(ctx, i.tracer, args)
-
+			res, err := tool.Call(ctx, i.tracer, tc.Function.Arguments)
 			_ = i.recorder.RecordToolUsage(ctx, &ToolUsageRecord{
 				InterceptionID:  i.ID().String(),
 				MsgID:           completion.ID,
 				ServerURL:       &tool.ServerURL,
 				Tool:            tool.Name,
-				Args:            i.unmarshalArgs(tc.Function.Arguments),
+				Args:            tc.Function.Arguments,
 				Injected:        true,
 				InvocationError: err,
 			})

intercept_openai_chat_streaming.go

@@ -154,7 +154,7 @@ func (i *OpenAIStreamingChatInterception) ProcessRequest(w http.ResponseWriter,
 				InterceptionID: i.ID().String(),
 				MsgID:          processor.getMsgID(),
 				Tool:           toolCall.Name,
-				Args:           i.unmarshalArgs(toolCall.Arguments),
+				Args:           toolCall.Arguments,
 				Injected:       false,
 			})
 			toolCall = nil
@@ -241,15 +241,13 @@ func (i *OpenAIStreamingChatInterception) ProcessRequest(w http.ResponseWriter,
 		i.req.Messages = append(i.req.Messages, processor.getLastCompletion().ToParam())
 
 		id := toolCall.ID
-		args := i.unmarshalArgs(toolCall.Arguments)
-		toolRes, toolErr := tool.Call(streamCtx, i.tracer, args)
-
+		toolRes, toolErr := tool.Call(streamCtx, i.tracer, toolCall.Arguments)
 		_ = i.recorder.RecordToolUsage(streamCtx, &ToolUsageRecord{
 			InterceptionID:  i.ID().String(),
 			MsgID:           processor.getMsgID(),
 			ServerURL:       &tool.ServerURL,
 			Tool:            tool.Name,
-			Args:            args,
+			Args:            toolCall.Arguments,
 			Injected:        true,
 			InvocationError: toolErr,
 		})

mcp/proxy_streamable_http.go

@@ -20,14 +20,16 @@ import (
 var _ ServerProxier = &StreamableHTTPServerProxy{}
 
 type StreamableHTTPServerProxy struct {
+	client *client.Client
+	logger slog.Logger
+	tracer trace.Tracer
+
+	allowlistPattern *regexp.Regexp
+	denylistPattern  *regexp.Regexp
+
 	serverName string
 	serverURL  string
-	client     *client.Client
-	logger     slog.Logger
-	tracer     trace.Tracer
 	tools      map[string]*Tool
-
-	allowlistPattern, denylistPattern *regexp.Regexp
 }
 
 func NewStreamableHTTPServerProxy(logger slog.Logger, tracer trace.Tracer, serverName, serverURL string, headers map[string]string, allowlist, denylist *regexp.Regexp) (*StreamableHTTPServerProxy, error) {