review changes

Author: pawbana

intercept/responses/base.go

@@ -24,6 +24,7 @@ import (
 	"github.com/google/uuid"
 	"github.com/openai/openai-go/v3/option"
 	"github.com/openai/openai-go/v3/responses"
+	oaiconst "github.com/openai/openai-go/v3/shared/constant"
 	"github.com/tidwall/gjson"
 	"github.com/tidwall/sjson"
 	"go.opentelemetry.io/otel/attribute"
@@ -222,17 +223,18 @@ func (i *responsesInterceptionBase) recordUserPrompt(ctx context.Context, respon
 
 func (i *responsesInterceptionBase) recordToolUsage(ctx context.Context, response *responses.Response) {
 	if response == nil {
+		i.logger.Warn(ctx, "got empty response, skipping tool usage recording")
 		return
 	}
 
 	for _, item := range response.Output {
 		var args recorder.ToolArgs
 
-		// recodring other function types to be considered: https://github.com/coder/aibridge/issues/121
+		// recording other function types to be considered: https://github.com/coder/aibridge/issues/121
 		switch item.Type {
-		case "function_call":
-			args = i.parseJSONArgs(item.Arguments)
-		case "custom_tool_call":
+		case string(oaiconst.ValueOf[oaiconst.FunctionCall]()):
+			args = i.parseFunctionCallJSONArgs(ctx, item.Arguments)
+		case string(oaiconst.ValueOf[oaiconst.CustomToolCall]()):
 			args = item.Input
 		default:
 			continue
@@ -250,15 +252,17 @@ func (i *responsesInterceptionBase) recordToolUsage(ctx context.Context, respons
 	}
 }
 
-func (i *responsesInterceptionBase) parseJSONArgs(raw string) recorder.ToolArgs {
-	if trimmed := strings.TrimSpace(raw); trimmed != "" {
+func (i *responsesInterceptionBase) parseFunctionCallJSONArgs(ctx context.Context, raw string) recorder.ToolArgs {
+	trimmed := strings.TrimSpace(raw)
+	if trimmed != "" {
 		var args recorder.ToolArgs
-		if err := json.Unmarshal([]byte(trimmed), &args); err == nil {
+		if err := json.Unmarshal([]byte(trimmed), &args); err != nil {
+			i.logger.Warn(ctx, "failed to unmarshal tool args", slog.Error(err))
+		} else {
 			return args
 		}
 	}
-
-	return nil
+	return trimmed
 }
 
 // responseCopier helper struct to send original response to the client

intercept/responses/base_test.go

@@ -237,7 +237,7 @@ func TestRecordToolUsage(t *testing.T) {
 					InterceptionID: id.String(),
 					MsgID:          "resp_456",
 					Tool:           "get_weather",
-					Args:           nil,
+					Args:           "",
 					Injected:       false,
 				},
 			},
@@ -252,6 +252,11 @@ func TestRecordToolUsage(t *testing.T) {
 						Name:      "get_weather",
 						Arguments: `{"location": "NYC"}`,
 					},
+					{
+						Type:      "function_call",
+						Name:      "bad_json_args",
+						Arguments: `{"bad": args`,
+					},
 					{
 						Type: "message",
 						ID:   "msg_1",
@@ -277,6 +282,13 @@ func TestRecordToolUsage(t *testing.T) {
 					Args:           map[string]any{"location": "NYC"},
 					Injected:       false,
 				},
+				{
+					InterceptionID: id.String(),
+					MsgID:          "resp_789",
+					Tool:           "bad_json_args",
+					Args:           `{"bad": args`,
+					Injected:       false,
+				},
 				{
 					InterceptionID: id.String(),
 					MsgID:          "resp_789",
@@ -329,20 +341,20 @@ func TestParseJSONArgs(t *testing.T) {
 		{
 			name:     "empty_string",
 			raw:      "",
-			expected: nil,
+			expected: "",
 		},
 		{
 			name:     "whitespace_only",
 			raw:      "   \t\n  ",
-			expected: nil,
+			expected: "",
 		},
 		{
 			name:     "invalid_json",
 			raw:      "{not valid json}",
-			expected: nil,
+			expected: "{not valid json}",
 		},
 		{
-			name: "nested_object",
+			name: "nested_object_with_trailing_spaces",
 			raw:  ` {"user": {"name": "alice", "settings": {"theme": "dark", "notifications": true}}, "count": 42}   `,
 			expected: map[string]any{
 				"user": map[string]any{
@@ -362,7 +374,7 @@ func TestParseJSONArgs(t *testing.T) {
 			t.Parallel()
 
 			base := &responsesInterceptionBase{}
-			result := base.parseJSONArgs(tc.raw)
+			result := base.parseFunctionCallJSONArgs(t.Context(), tc.raw)
 			require.Equal(t, tc.expected, result)
 		})
 	}

responses_integration_test.go

@@ -201,7 +201,7 @@ func TestResponsesOutputMatchesUpstream(t *testing.T) {
 			recordedTools := mockRecorder.RecordedToolUsages()
 			if tc.expectToolRecorded != nil {
 				require.Len(t, recordedTools, 1)
-				recordedTools[0].InterceptionID = tc.expectToolRecorded.InterceptionID // ignore interception id
+				recordedTools[0].InterceptionID = tc.expectToolRecorded.InterceptionID // ignore interception id (interception id is not constant and response doesn't contain it)
 				recordedTools[0].CreatedAt = tc.expectToolRecorded.CreatedAt           // ignore time
 				require.Equal(t, tc.expectToolRecorded, recordedTools[0])
 			} else {