Skip to content

Commit 5c3469e

Browse files
dannykoppingdannykopping
committed
OpenAI tool call test, overriding IDs of returned messages
Signed-off-by: Danny Kopping <dannykopping@gmail.com>
1 parent bd90c83 commit 5c3469e

File tree

9 files changed

+579
-136
lines changed

9 files changed

+579
-136
lines changed

‎aibridged/bridge.go

Lines changed: 115 additions & 46 deletions
Large diffs are not rendered by default.

‎aibridged/bridge_test.go

Lines changed: 134 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,9 @@ import (
2424
"github.com/stretchr/testify/require"
2525
"github.com/tidwall/sjson"
2626

27+
"github.com/openai/openai-go"
28+
oai_ssestream "github.com/openai/openai-go/packages/ssestream"
29+
2730
"github.com/mark3labs/mcp-go/mcp"
2831
"github.com/mark3labs/mcp-go/server"
2932

@@ -44,6 +47,8 @@ var (
4447

4548
//go:embed fixtures/openai/single_builtin_tool.txtar
4649
oaiSingleBuiltinTool []byte
50+
//go:embed fixtures/openai/single_injected_tool.txtar
51+
oaiSingleInjectedTool []byte
4752

4853
//go:embed fixtures/anthropic/simple.txtar
4954
antSimple []byte
@@ -162,6 +167,7 @@ func TestAnthropicMessages(t *testing.T) {
162167
}
163168
})
164169

170+
// TODO: fixture contains hardcoded tool name; this is flimsy since our naming convention may change for injected tools.
165171
t.Run("single injected tool", func(t *testing.T) {
166172
t.Parallel()
167173

@@ -392,6 +398,133 @@ func TestOpenAIChatCompletions(t *testing.T) {
392398
})
393399
}
394400
})
401+
402+
// TODO: fixture contains hardcoded tool name; this is flimsy since our naming convention may change for injected tools.
403+
t.Run("single injected tool", func(t *testing.T) {
404+
t.Parallel()
405+
406+
cases := []struct {
407+
streaming bool
408+
}{
409+
{
410+
streaming: true,
411+
},
412+
// {
413+
// streaming: false,
414+
// },
415+
}
416+
417+
for _, tc := range cases {
418+
t.Run(fmt.Sprintf("%s/streaming=%v", t.Name(), tc.streaming), func(t *testing.T) {
419+
t.Parallel()
420+
421+
arc := txtar.Parse(oaiSingleInjectedTool)
422+
t.Logf("%s: %s", t.Name(), arc.Comment)
423+
424+
files := filesMap(arc)
425+
require.Len(t, files, 5)
426+
require.Contains(t, files, fixtureRequest)
427+
require.Contains(t, files, fixtureStreamingResponse)
428+
require.Contains(t, files, fixtureNonStreamingResponse)
429+
require.Contains(t, files, fixtureStreamingToolResponse)
430+
require.Contains(t, files, fixtureNonStreamingToolResponse)
431+
432+
reqBody := files[fixtureRequest]
433+
434+
// Add the stream param to the request.
435+
newBody, err := sjson.SetBytes(reqBody, "stream", tc.streaming)
436+
require.NoError(t, err)
437+
reqBody = newBody
438+
439+
ctx := testutil.Context(t, testutil.WaitLong)
440+
// Conditionally return fixtures based on request count.
441+
// First request: halts with tool call instruction.
442+
// Second request: tool call invocation.
443+
mockSrv := newMockServer(ctx, t, files, func(reqCount uint32, resp []byte) []byte {
444+
if reqCount == 1 {
445+
return resp
446+
}
447+
448+
if reqCount > 2 {
449+
t.Fatalf("did not expect more than 2 calls; received %d", reqCount)
450+
}
451+
452+
if !tc.streaming {
453+
return files[fixtureNonStreamingToolResponse]
454+
}
455+
return files[fixtureStreamingToolResponse]
456+
})
457+
t.Cleanup(mockSrv.Close)
458+
459+
coderdClient := &fakeBridgeDaemonClient{}
460+
logger := testutil.Logger(t) // slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
461+
462+
// Setup Coder MCP integration.
463+
mcpSrv := httptest.NewServer(createMockMCPSrv(t))
464+
mcpBridge, err := aibridged.NewMCPToolBridge("coder", mcpSrv.URL, map[string]string{}, logger)
465+
require.NoError(t, err)
466+
467+
// Initialize MCP client, fetch tools, and inject into bridge.
468+
require.NoError(t, mcpBridge.Init(testutil.Context(t, testutil.WaitShort)))
469+
tools := mcpBridge.ListTools()
470+
require.NotEmpty(t, tools)
471+
472+
b, err := aibridged.NewBridge(codersdk.AIBridgeConfig{
473+
Daemons: 1,
474+
OpenAI: codersdk.AIBridgeOpenAIConfig{
475+
BaseURL: serpent.String(mockSrv.URL),
476+
Key: serpent.String(sessionToken),
477+
},
478+
}, logger, func() (proto.DRPCAIBridgeDaemonClient, bool) {
479+
return coderdClient, true
480+
}, tools)
481+
require.NoError(t, err)
482+
483+
// Invoke request to mocked API via aibridge.
484+
bridgeSrv := httptest.NewServer(b.Handler())
485+
req := createOpenAIChatCompletionsReq(t, bridgeSrv.URL, reqBody)
486+
client := &http.Client{}
487+
resp, err := client.Do(req)
488+
require.NoError(t, err)
489+
require.Equal(t, http.StatusOK, resp.StatusCode)
490+
defer resp.Body.Close()
491+
492+
// We must ALWAYS have 2 calls to the bridge.
493+
require.Eventually(t, func() bool { return mockSrv.callCount.Load() == 2 }, testutil.WaitLong, testutil.IntervalFast)
494+
495+
// TODO: this is a bit flimsy since this API won't be in beta forever.
496+
var content *openai.ChatCompletionChoice
497+
if tc.streaming {
498+
// Parse the response stream.
499+
decoder := oai_ssestream.NewDecoder(resp)
500+
stream := oai_ssestream.NewStream[openai.ChatCompletionChunk](decoder, nil)
501+
var message openai.ChatCompletionAccumulator
502+
for stream.Next() {
503+
chunk := stream.Current()
504+
message.AddChunk(chunk)
505+
}
506+
507+
require.NoError(t, stream.Err())
508+
require.Len(t, message.Choices, 1)
509+
content = &message.Choices[0]
510+
} else {
511+
// Parse & unmarshal the response.
512+
out, err := io.ReadAll(resp.Body)
513+
require.NoError(t, err)
514+
515+
// TODO: this is a bit flimsy since this API won't be in beta forever.
516+
var message openai.ChatCompletion
517+
require.NoError(t, json.Unmarshal(out, &message))
518+
require.NotNil(t, message)
519+
require.Len(t, message.Choices, 1)
520+
content = &message.Choices[0]
521+
}
522+
523+
require.NotNil(t, content)
524+
require.Contains(t, content.Message.Content, "admin")
525+
})
526+
}
527+
})
395528
}
396529

397530
func TestSimple(t *testing.T) {
@@ -580,8 +713,7 @@ func newMockServer(ctx context.Context, t *testing.T, files archiveFileMap, resp
580713

581714
ms := &mockServer{}
582715
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
583-
callCount := ms.callCount.Add(1)
584-
t.Logf("\n\nCALL COUNT: %d\n\n", callCount)
716+
ms.callCount.Add(1)
585717

586718
body, err := io.ReadAll(r.Body)
587719
defer r.Body.Close()

‎aibridged/fixtures/README.md

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
These fixtures were created by adding logging middleware to API calls to view the raw requests/responses.
2+
3+
```go
4+
...
5+
opts = append(opts, option.WithMiddleware(LoggingMiddleware))
6+
...
7+
8+
func LoggingMiddleware(req *http.Request, next option.MiddlewareNext) (res *http.Response, err error) {
9+
reqOut, _ := httputil.DumpRequest(req, true)
10+
11+
// Forward the request to the next handler
12+
res, err = next(req)
13+
fmt.Printf("[req] %s\n", reqOut)
14+
15+
// Handle stuff after the request
16+
if err != nil {
17+
return res, err
18+
}
19+
20+
respOut, _ := httputil.DumpResponse(res, true)
21+
fmt.Printf("[resp] %s\n", respOut)
22+
23+
return res, err
24+
}
25+
```

‎aibridged/fixtures/anthropic/single_injected_tool.txtar

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ Coder MCP tools automatically injected.
33
-- request --
44
{
55
"model": "claude-sonnet-4-20250514",
6-
"max_tokens": 1024,
6+
"max_tokens": 1,
77
"messages": [
88
{
99
"role": "user",
@@ -14,7 +14,7 @@ Coder MCP tools automatically injected.
1414

1515
-- streaming --
1616
event: message_start
17-
data: {"type":"message_start","message":{"id":"msg_01FHxVHRpKWz6dFfm9s1jkNY","type":"message","role":"assistant","model":"claude-sonnet-4-20250514","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":385,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"output_tokens":5,"service_tier":"standard"}} }
17+
data: {"type":"message_start","message":{"id":"1545b825-1943-43e6-90e0-2d76b7b51afd","type":"message","role":"assistant","model":"claude-sonnet-4-20250514","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":385,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"output_tokens":5,"service_tier":"standard"}} }
1818

1919
event: content_block_start
2020
data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""} }
@@ -49,7 +49,7 @@ data: {"type":"message_stop" }
4949

5050
-- streaming/tool-call --
5151
event: message_start
52-
data: {"type":"message_start","message":{"id":"msg_01Rsww53iPc1nW2uEBrMkxH6","type":"message","role":"assistant","model":"claude-sonnet-4-20250514","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":7,"cache_creation_input_tokens":1162,"cache_read_input_tokens":35876,"output_tokens":1,"service_tier":"standard"}} }
52+
data: {"type":"message_start","message":{"id":"0a2b5ae7-4973-41d1-b725-d724e7d85408","type":"message","role":"assistant","model":"claude-sonnet-4-20250514","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":7,"cache_creation_input_tokens":1162,"cache_read_input_tokens":35876,"output_tokens":1,"service_tier":"standard"}} }
5353

5454
event: content_block_start
5555
data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""} }
@@ -72,7 +72,7 @@ data: {"type":"message_stop" }
7272

7373
-- non-streaming --
7474
{
75-
"id": "msg_018NqsW4eMw8tw47xBby5sV6",
75+
"id": "e3d937ce-d871-4b01-9c9b-c20227c27553",
7676
"type": "message",
7777
"role": "assistant",
7878
"model": "claude-sonnet-4-20250514",
@@ -102,7 +102,7 @@ data: {"type":"message_stop" }
102102

103103
-- non-streaming/tool-call --
104104
{
105-
"id": "msg_01H7amafFFU8rkqPNNX7LoZ2",
105+
"id": "7f7f66fa-fbad-48ea-80dc-31d4a05debcc",
106106
"type": "message",
107107
"role": "assistant",
108108
"model": "claude-sonnet-4-20250514",

0 commit comments

Comments
 (0)