fix openAI tool arguments consistency

Author: pawbana

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, "", recorderClient.toolUsages[0].Args)
+				require.IsType(t, map[string]any{}, recorderClient.toolUsages[0].Args)
 				require.Contains(t, recorderClient.toolUsages[0].Args, "path")
-				assert.Equal(t, "README.md", gjson.Get(recorderClient.toolUsages[0].Args.(string), "path").Str)
+				assert.Equal(t, "README.md", recorderClient.toolUsages[0].Args.(map[string]any)["path"])
 
 				require.Len(t, recorderClient.userPrompts, 1)
 				assert.Equal(t, "how large is the README.md file in my current path", recorderClient.userPrompts[0].Prompt)
@@ -852,7 +852,7 @@ func TestOpenAIInjectedTools(t *testing.T) {
 			// Ensure expected tool was invoked with expected input.
 			require.Len(t, recorderClient.toolUsages, 1)
 			require.Equal(t, mockToolName, recorderClient.toolUsages[0].Tool)
-			expected := "{\"owner\":\"admin\"}"
+			expected := map[string]any{"owner": "admin"}
 			require.EqualValues(t, expected, recorderClient.toolUsages[0].Args)
 
 			var (

intercept_openai_chat_base.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"net/http"
+	"strings"
 
 	"github.com/coder/aibridge/mcp"
 	"github.com/coder/aibridge/tracing"
@@ -104,6 +105,18 @@ 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

@@ -119,7 +119,7 @@ func (i *OpenAIBlockingChatInterception) ProcessRequest(w http.ResponseWriter, r
 						InterceptionID: i.ID().String(),
 						MsgID:          completion.ID,
 						Tool:           toolCall.Function.Name,
-						Args:           toolCall.Function.Arguments,
+						Args:           i.unmarshalArgs(toolCall.Function.Arguments),
 						Injected:       false,
 					})
 				}
@@ -150,13 +150,14 @@ func (i *OpenAIBlockingChatInterception) ProcessRequest(w http.ResponseWriter, r
 				appendedPrevMsg = true
 			}
 
-			res, err := tool.Call(ctx, i.tracer, tc.Function.Arguments)
+			args := i.unmarshalArgs(tc.Function.Arguments)
+			res, err := tool.Call(ctx, i.tracer, args)
 			_ = i.recorder.RecordToolUsage(ctx, &ToolUsageRecord{
 				InterceptionID:  i.ID().String(),
 				MsgID:           completion.ID,
 				ServerURL:       &tool.ServerURL,
 				Tool:            tool.Name,
-				Args:            tc.Function.Arguments,
+				Args:            args,
 				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:           toolCall.Arguments,
+				Args:           i.unmarshalArgs(toolCall.Arguments),
 				Injected:       false,
 			})
 			toolCall = nil
@@ -241,13 +241,14 @@ func (i *OpenAIStreamingChatInterception) ProcessRequest(w http.ResponseWriter,
 		i.req.Messages = append(i.req.Messages, processor.getLastCompletion().ToParam())
 
 		id := toolCall.ID
-		toolRes, toolErr := tool.Call(streamCtx, i.tracer, toolCall.Arguments)
+		args := i.unmarshalArgs(toolCall.Arguments)
+		toolRes, toolErr := tool.Call(streamCtx, i.tracer, args)
 		_ = i.recorder.RecordToolUsage(streamCtx, &ToolUsageRecord{
 			InterceptionID:  i.ID().String(),
 			MsgID:           processor.getMsgID(),
 			ServerURL:       &tool.ServerURL,
 			Tool:            tool.Name,
-			Args:            toolCall.Arguments,
+			Args:            args,
 			Injected:        true,
 			InvocationError: toolErr,
 		})

mcp/tool.go

@@ -54,6 +54,9 @@ func (t *Tool) Call(ctx context.Context, tracer trace.Tracer, input any) (_ *mcp
 		attribute.String(tracing.MCPServerName, t.ServerName),
 		attribute.String(tracing.MCPServerURL, t.ServerURL),
 	)
+	ctx, span := tracer.Start(ctx, "Intercept.ProcessRequest.ToolCall", trace.WithAttributes(spanAttrs...))
+	defer tracing.EndSpanErr(span, &outErr)
+
 	inputJson, err := json.Marshal(input)
 	if err != nil {
 		t.Logger.Warn(ctx, "failed to marshal tool input, will be omitted from span attrs: %v", err)
@@ -62,12 +65,9 @@ func (t *Tool) Call(ctx context.Context, tracer trace.Tracer, input any) (_ *mcp
 		if len(strJson) > maxSpanInputAttrLen {
 			strJson = strJson[:100]
 		}
-		spanAttrs = append(spanAttrs, attribute.String(tracing.MCPInput, strJson))
+		span.SetAttributes(attribute.String(tracing.MCPInput, strJson))
 	}
 
-	ctx, span := tracer.Start(ctx, "Intercept.ProcessRequest.ToolCall", trace.WithAttributes(spanAttrs...))
-	defer tracing.EndSpanErr(span, &outErr)
-
 	return t.Client.CallTool(ctx, mcp.CallToolRequest{
 		Params: mcp.CallToolParams{
 			Name:      t.Name,

trace_integration_test.go

@@ -638,7 +638,7 @@ func TestOpenAIInjectedToolsTrace(t *testing.T) {
 					attribute.String(tracing.Provider, aibridge.ProviderOpenAI),
 					attribute.String(tracing.Model, gjson.Get(reqBody, "model").Str),
 					attribute.String(tracing.InitiatorID, userID),
-					attribute.String(tracing.MCPInput, "\"{\\\"owner\\\":\\\"admin\\\"}\""),
+					attribute.String(tracing.MCPInput, "{\"owner\":\"admin\"}"),
 					attribute.String(tracing.MCPToolName, "coder_list_workspaces"),
 					attribute.String(tracing.MCPServerName, tool.ServerName),
 					attribute.String(tracing.MCPServerURL, tool.ServerURL),