Skip to content

Commit ee80a4b

Browse files
pawbanapawbana
committed
feat: add tool usage recording for blocking responses interceptor
1 parent 61a792b commit ee80a4b

File tree

6 files changed

+369
-4
lines changed

6 files changed

+369
-4
lines changed

‎fixtures/fixtures.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,9 @@ var (
5454
//go:embed openai/responses/blocking/builtin_tool.txtar
5555
OaiResponsesBlockingBuiltinTool []byte
5656

57+
//go:embed openai/responses/blocking/custom_tool.txtar
58+
OaiResponsesBlockingCustomTool []byte
59+
5760
//go:embed openai/responses/blocking/conversation.txtar
5861
OaiResponsesBlockingConversation []byte
5962

‎fixtures/openai/responses/blocking/custom_tool.txtar‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
-- request --
2+
{
3+
"input": "Use the code_exec tool to print hello world to the console.",
4+
"model": "gpt-5",
5+
"tools": [
6+
{
7+
"type": "custom",
8+
"name": "code_exec",
9+
"description": "Executes arbitrary Python code."
10+
}
11+
]
12+
}
13+
14+
-- non-streaming --
15+
{
16+
"id": "resp_09c614364030cdf000696942589da081a0af07f5859acb7308",
17+
"object": "response",
18+
"created_at": 1768505944,
19+
"status": "completed",
20+
"background": false,
21+
"billing": {
22+
"payer": "developer"
23+
},
24+
"completed_at": 1768505948,
25+
"error": null,
26+
"frequency_penalty": 0.0,
27+
"incomplete_details": null,
28+
"instructions": null,
29+
"max_output_tokens": null,
30+
"max_tool_calls": null,
31+
"model": "gpt-5-2025-08-07",
32+
"output": [
33+
{
34+
"id": "rs_09c614364030cdf00069694258e45881a0b8d5f198cde47d58",
35+
"type": "reasoning",
36+
"summary": []
37+
},
38+
{
39+
"id": "ctc_09c614364030cdf0006969425bf33481a09cc0f9522af2d980",
40+
"type": "custom_tool_call",
41+
"status": "completed",
42+
"call_id": "call_haf8njtwrVZ1754Gm6fjAtuA",
43+
"input": "print(\"hello world\")",
44+
"name": "code_exec"
45+
}
46+
],
47+
"parallel_tool_calls": true,
48+
"presence_penalty": 0.0,
49+
"previous_response_id": null,
50+
"prompt_cache_key": null,
51+
"prompt_cache_retention": null,
52+
"reasoning": {
53+
"effort": "medium",
54+
"summary": null
55+
},
56+
"safety_identifier": null,
57+
"service_tier": "default",
58+
"store": true,
59+
"temperature": 1.0,
60+
"text": {
61+
"format": {
62+
"type": "text"
63+
},
64+
"verbosity": "medium"
65+
},
66+
"tool_choice": "auto",
67+
"tools": [
68+
{
69+
"type": "custom",
70+
"description": "Executes arbitrary Python code.",
71+
"format": {
72+
"type": "text"
73+
},
74+
"name": "code_exec"
75+
}
76+
],
77+
"top_logprobs": 0,
78+
"top_p": 1.0,
79+
"truncation": "disabled",
80+
"usage": {
81+
"input_tokens": 64,
82+
"input_tokens_details": {
83+
"cached_tokens": 0
84+
},
85+
"output_tokens": 148,
86+
"output_tokens_details": {
87+
"reasoning_tokens": 128
88+
},
89+
"total_tokens": 212
90+
},
91+
"user": null,
92+
"metadata": {}
93+
}

‎intercept/responses/base.go‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -220,6 +220,47 @@ func (i *responsesInterceptionBase) recordUserPrompt(ctx context.Context, respon
220220
}
221221
}
222222

223+
func (i *responsesInterceptionBase) recordToolUsage(ctx context.Context, response *responses.Response) {
224+
if response == nil {
225+
return
226+
}
227+
228+
for _, item := range response.Output {
229+
var args recorder.ToolArgs
230+
231+
// recodring other function types to be considered: https://github.com/coder/aibridge/issues/121
232+
switch item.Type {
233+
case "function_call":
234+
args = i.parseJSONArgs(item.Arguments)
235+
case "custom_tool_call":
236+
args = item.Input
237+
default:
238+
continue
239+
}
240+
241+
if err := i.recorder.RecordToolUsage(ctx, &recorder.ToolUsageRecord{
242+
InterceptionID: i.ID().String(),
243+
MsgID: response.ID,
244+
Tool: item.Name,
245+
Args: args,
246+
Injected: false,
247+
}); err != nil {
248+
i.logger.Warn(ctx, "failed to record tool usage", slog.Error(err), slog.F("tool", item.Name))
249+
}
250+
}
251+
}
252+
253+
func (i *responsesInterceptionBase) parseJSONArgs(raw string) recorder.ToolArgs {
254+
if trimmed := strings.TrimSpace(raw); trimmed != "" {
255+
var args recorder.ToolArgs
256+
if err := json.Unmarshal([]byte(trimmed), &args); err == nil {
257+
return args
258+
}
259+
}
260+
261+
return nil
262+
}
263+
223264
// responseCopier helper struct to send original response to the client
224265
type responseCopier struct {
225266
buff deltaBuffer

‎intercept/responses/base_test.go‎

Lines changed: 173 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,14 @@ package responses
22

33
import (
44
"testing"
5+
"time"
56

67
"cdr.dev/slog/v3"
78
"github.com/coder/aibridge/fixtures"
89
"github.com/coder/aibridge/internal/testutil"
10+
"github.com/coder/aibridge/recorder"
911
"github.com/google/uuid"
12+
oairesponses "github.com/openai/openai-go/v3/responses"
1013
"github.com/stretchr/testify/require"
1114
)
1215

@@ -194,3 +197,173 @@ func TestRecordPrompt(t *testing.T) {
194197
})
195198
}
196199
}
200+
201+
func TestRecordToolUsage(t *testing.T) {
202+
t.Parallel()
203+
204+
id := uuid.MustParse("11111111-1111-1111-1111-111111111111")
205+
206+
tests := []struct {
207+
name string
208+
response *oairesponses.Response
209+
expected []*recorder.ToolUsageRecord
210+
}{
211+
{
212+
name: "nil_response",
213+
response: nil,
214+
expected: nil,
215+
},
216+
{
217+
name: "empty_output",
218+
response: &oairesponses.Response{
219+
ID: "resp_123",
220+
},
221+
expected: nil,
222+
},
223+
{
224+
name: "empty_tool_args",
225+
response: &oairesponses.Response{
226+
ID: "resp_456",
227+
Output: []oairesponses.ResponseOutputItemUnion{
228+
{
229+
Type: "function_call",
230+
Name: "get_weather",
231+
Arguments: "",
232+
},
233+
},
234+
},
235+
expected: []*recorder.ToolUsageRecord{
236+
{
237+
InterceptionID: id.String(),
238+
MsgID: "resp_456",
239+
Tool: "get_weather",
240+
Args: nil,
241+
Injected: false,
242+
},
243+
},
244+
},
245+
{
246+
name: "multiple_tool_calls",
247+
response: &oairesponses.Response{
248+
ID: "resp_789",
249+
Output: []oairesponses.ResponseOutputItemUnion{
250+
{
251+
Type: "function_call",
252+
Name: "get_weather",
253+
Arguments: `{"location": "NYC"}`,
254+
},
255+
{
256+
Type: "message",
257+
ID: "msg_1",
258+
Role: "assistant",
259+
},
260+
{
261+
Type: "custom_tool_call",
262+
Name: "search",
263+
Input: `{\"query\": \"test\"}`,
264+
},
265+
{
266+
Type: "function_call",
267+
Name: "calculate",
268+
Arguments: `{"a": 1, "b": 2}`,
269+
},
270+
},
271+
},
272+
expected: []*recorder.ToolUsageRecord{
273+
{
274+
InterceptionID: id.String(),
275+
MsgID: "resp_789",
276+
Tool: "get_weather",
277+
Args: map[string]any{"location": "NYC"},
278+
Injected: false,
279+
},
280+
{
281+
InterceptionID: id.String(),
282+
MsgID: "resp_789",
283+
Tool: "search",
284+
Args: `{\"query\": \"test\"}`,
285+
Injected: false,
286+
},
287+
{
288+
InterceptionID: id.String(),
289+
MsgID: "resp_789",
290+
Tool: "calculate",
291+
Args: map[string]any{"a": float64(1), "b": float64(2)},
292+
Injected: false,
293+
},
294+
},
295+
},
296+
}
297+
298+
for _, tc := range tests {
299+
t.Run(tc.name, func(t *testing.T) {
300+
t.Parallel()
301+
302+
rec := &testutil.MockRecorder{}
303+
base := &responsesInterceptionBase{
304+
id: id,
305+
recorder: rec,
306+
logger: slog.Make(),
307+
}
308+
309+
base.recordToolUsage(t.Context(), tc.response)
310+
311+
tools := rec.RecordedToolUsages()
312+
require.Len(t, tools, len(tc.expected))
313+
for i, got := range tools {
314+
got.CreatedAt = time.Time{}
315+
require.Equal(t, tc.expected[i], got)
316+
}
317+
})
318+
}
319+
}
320+
321+
func TestParseJSONArgs(t *testing.T) {
322+
t.Parallel()
323+
324+
tests := []struct {
325+
name string
326+
raw string
327+
expected recorder.ToolArgs
328+
}{
329+
{
330+
name: "empty_string",
331+
raw: "",
332+
expected: nil,
333+
},
334+
{
335+
name: "whitespace_only",
336+
raw: " \t\n ",
337+
expected: nil,
338+
},
339+
{
340+
name: "invalid_json",
341+
raw: "{not valid json}",
342+
expected: nil,
343+
},
344+
{
345+
name: "nested_object",
346+
raw: ` {"user": {"name": "alice", "settings": {"theme": "dark", "notifications": true}}, "count": 42} `,
347+
expected: map[string]any{
348+
"user": map[string]any{
349+
"name": "alice",
350+
"settings": map[string]any{
351+
"theme": "dark",
352+
"notifications": true,
353+
},
354+
},
355+
"count": float64(42),
356+
},
357+
},
358+
}
359+
360+
for _, tc := range tests {
361+
t.Run(tc.name, func(t *testing.T) {
362+
t.Parallel()
363+
364+
base := &responsesInterceptionBase{}
365+
result := base.parseJSONArgs(tc.raw)
366+
require.Equal(t, tc.expected, result)
367+
})
368+
}
369+
}

‎intercept/responses/blocking.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ func (i *BlockingResponsesInterceptor) ProcessRequest(w http.ResponseWriter, r *
5555
// response could be nil eg. fixtures/openai/responses/blocking/wrong_response_format.txtar
5656
if response != nil {
5757
i.recordUserPrompt(ctx, response.ID)
58+
i.recordToolUsage(ctx, response)
5859
}
5960

6061
if upstreamErr != nil && !respCopy.responseReceived.Load() {

0 commit comments

Comments
 (0)