From c0d7cdd62ce2f6450ed37a60f55aa74d2cdc6d7a Mon Sep 17 00:00:00 2001 From: Danny Kopping Date: Tue, 23 Dec 2025 16:24:03 +0200 Subject: [PATCH 1/2] fix: move cache breakpoint to final injected tool, if present Closes https://github.com/coder/aibridge/issues/69 Signed-off-by: Danny Kopping --- intercept_anthropic_message_internal_test.go | 323 +++++++++++++++++++ intercept_anthropic_messages_base.go | 55 +++- 2 files changed, 366 insertions(+), 12 deletions(-) create mode 100644 intercept_anthropic_message_internal_test.go diff --git a/intercept_anthropic_message_internal_test.go b/intercept_anthropic_message_internal_test.go new file mode 100644 index 0000000..600dba9 --- /dev/null +++ b/intercept_anthropic_message_internal_test.go @@ -0,0 +1,323 @@ +package aibridge + +import ( + "context" + "testing" + + "github.com/anthropics/anthropic-sdk-go" + "github.com/anthropics/anthropic-sdk-go/shared/constant" + "github.com/coder/aibridge/mcp" + mcpgo "github.com/mark3labs/mcp-go/mcp" + "github.com/stretchr/testify/require" +) + +func TestInjectTools_CacheBreakpoints(t *testing.T) { + t.Parallel() + + t.Run("cache control preserved when no tools to inject", func(t *testing.T) { + t.Parallel() + + // Request has existing tool with cache control, but no tools to inject. + i := &AnthropicMessagesInterceptionBase{ + req: &MessageNewParamsWrapper{ + MessageNewParams: anthropic.MessageNewParams{ + Tools: []anthropic.ToolUnionParam{ + { + OfTool: &anthropic.ToolParam{ + Name: "existing_tool", + CacheControl: anthropic.CacheControlEphemeralParam{ + Type: constant.ValueOf[constant.Ephemeral](), + }, + }, + }, + }, + }, + }, + mcpProxy: &mockServerProxier{tools: nil}, + } + + i.injectTools() + + // Cache control should remain untouched since no tools were injected. + require.Len(t, i.req.Tools, 1) + require.Equal(t, constant.ValueOf[constant.Ephemeral](), i.req.Tools[0].OfTool.CacheControl.Type) + }) + + t.Run("cache control breakpoint is preserved and moved to final tool", func(t *testing.T) { + t.Parallel() + + // Request has existing tool with cache control. + i := &AnthropicMessagesInterceptionBase{ + req: &MessageNewParamsWrapper{ + MessageNewParams: anthropic.MessageNewParams{ + Tools: []anthropic.ToolUnionParam{ + { + OfTool: &anthropic.ToolParam{ + Name: "existing_tool", + CacheControl: anthropic.CacheControlEphemeralParam{ + Type: constant.ValueOf[constant.Ephemeral](), + }, + }, + }, + }, + }, + }, + mcpProxy: &mockServerProxier{ + tools: []*mcp.Tool{ + {ID: "injected_tool", Name: "injected", Description: "Injected tool"}, + }, + }, + } + + i.injectTools() + + require.Len(t, i.req.Tools, 2) + // Original tool's cache control should be cleared. + require.Equal(t, "existing_tool", i.req.Tools[0].OfTool.Name) + require.Zero(t, i.req.Tools[0].OfTool.CacheControl) + // Cache control breakpoint should be moved to the final tool. + require.Equal(t, "injected_tool", i.req.Tools[1].OfTool.Name) + require.Equal(t, constant.ValueOf[constant.Ephemeral](), i.req.Tools[1].OfTool.CacheControl.Type) + }) + + // Multiple breakpoints should not be set, but if they are we should only move the first one to the end. + t.Run("only first cache control breakpoint is moved when multiple exist", func(t *testing.T) { + t.Parallel() + + // Request has multiple tools with cache control breakpoints. + i := &AnthropicMessagesInterceptionBase{ + req: &MessageNewParamsWrapper{ + MessageNewParams: anthropic.MessageNewParams{ + Tools: []anthropic.ToolUnionParam{ + { + OfTool: &anthropic.ToolParam{ + Name: "tool_with_cache_1", + CacheControl: anthropic.CacheControlEphemeralParam{ + Type: constant.ValueOf[constant.Ephemeral](), + }, + }, + }, + { + OfTool: &anthropic.ToolParam{ + Name: "tool_with_cache_2", + CacheControl: anthropic.CacheControlEphemeralParam{ + Type: constant.ValueOf[constant.Ephemeral](), + }, + }, + }, + }, + }, + }, + mcpProxy: &mockServerProxier{ + tools: []*mcp.Tool{ + {ID: "injected_tool", Name: "injected", Description: "Injected tool"}, + }, + }, + } + + i.injectTools() + + require.Len(t, i.req.Tools, 3) + // First tool's cache control should be cleared (it was captured). + require.Equal(t, "tool_with_cache_1", i.req.Tools[0].OfTool.Name) + require.Zero(t, i.req.Tools[0].OfTool.CacheControl) + // Second tool's cache control should remain (loop breaks after first match). + require.Equal(t, "tool_with_cache_2", i.req.Tools[1].OfTool.Name) + require.Equal(t, constant.ValueOf[constant.Ephemeral](), i.req.Tools[1].OfTool.CacheControl.Type) + // Only the first breakpoint is moved to the final tool. + require.Equal(t, "injected_tool", i.req.Tools[2].OfTool.Name) + require.Equal(t, constant.ValueOf[constant.Ephemeral](), i.req.Tools[2].OfTool.CacheControl.Type) + }) + + t.Run("no cache control added when none originally set", func(t *testing.T) { + t.Parallel() + + // Request has tools but none with cache control. + i := &AnthropicMessagesInterceptionBase{ + req: &MessageNewParamsWrapper{ + MessageNewParams: anthropic.MessageNewParams{ + Tools: []anthropic.ToolUnionParam{ + { + OfTool: &anthropic.ToolParam{ + Name: "existing_tool_no_cache", + }, + }, + }, + }, + }, + mcpProxy: &mockServerProxier{ + tools: []*mcp.Tool{ + {ID: "injected_tool", Name: "injected", Description: "Injected tool"}, + }, + }, + } + + i.injectTools() + + require.Len(t, i.req.Tools, 2) + // Neither tool should have cache control. + require.Equal(t, "existing_tool_no_cache", i.req.Tools[0].OfTool.Name) + require.Zero(t, i.req.Tools[0].OfTool.CacheControl) + require.Equal(t, "injected_tool", i.req.Tools[1].OfTool.Name) + require.Zero(t, i.req.Tools[1].OfTool.CacheControl) + }) +} + +func TestInjectTools_ParallelToolCalls(t *testing.T) { + t.Parallel() + + t.Run("disables parallel tool use for auto tool choice (default)", func(t *testing.T) { + t.Parallel() + + i := &AnthropicMessagesInterceptionBase{ + req: &MessageNewParamsWrapper{ + MessageNewParams: anthropic.MessageNewParams{ + // No tool choice set (default). + }, + }, + mcpProxy: &mockServerProxier{ + tools: []*mcp.Tool{{ID: "test_tool", Name: "test", Description: "Test"}}, + }, + } + + i.injectTools() + + require.NotNil(t, i.req.ToolChoice.OfAuto) + require.True(t, i.req.ToolChoice.OfAuto.DisableParallelToolUse.Valid()) + require.True(t, i.req.ToolChoice.OfAuto.DisableParallelToolUse.Value) + }) + + t.Run("disables parallel tool use for explicit auto tool choice", func(t *testing.T) { + t.Parallel() + + i := &AnthropicMessagesInterceptionBase{ + req: &MessageNewParamsWrapper{ + MessageNewParams: anthropic.MessageNewParams{ + ToolChoice: anthropic.ToolChoiceUnionParam{ + OfAuto: &anthropic.ToolChoiceAutoParam{ + Type: constant.ValueOf[constant.Auto](), + }, + }, + }, + }, + mcpProxy: &mockServerProxier{ + tools: []*mcp.Tool{{ID: "test_tool", Name: "test", Description: "Test"}}, + }, + } + + i.injectTools() + + require.NotNil(t, i.req.ToolChoice.OfAuto) + require.True(t, i.req.ToolChoice.OfAuto.DisableParallelToolUse.Valid()) + require.True(t, i.req.ToolChoice.OfAuto.DisableParallelToolUse.Value) + }) + + t.Run("disables parallel tool use for any tool choice", func(t *testing.T) { + t.Parallel() + + i := &AnthropicMessagesInterceptionBase{ + req: &MessageNewParamsWrapper{ + MessageNewParams: anthropic.MessageNewParams{ + ToolChoice: anthropic.ToolChoiceUnionParam{ + OfAny: &anthropic.ToolChoiceAnyParam{ + Type: constant.ValueOf[constant.Any](), + }, + }, + }, + }, + mcpProxy: &mockServerProxier{ + tools: []*mcp.Tool{{ID: "test_tool", Name: "test", Description: "Test"}}, + }, + } + + i.injectTools() + + require.NotNil(t, i.req.ToolChoice.OfAny) + require.True(t, i.req.ToolChoice.OfAny.DisableParallelToolUse.Valid()) + require.True(t, i.req.ToolChoice.OfAny.DisableParallelToolUse.Value) + }) + + t.Run("disables parallel tool use for tool choice type", func(t *testing.T) { + t.Parallel() + + i := &AnthropicMessagesInterceptionBase{ + req: &MessageNewParamsWrapper{ + MessageNewParams: anthropic.MessageNewParams{ + ToolChoice: anthropic.ToolChoiceUnionParam{ + OfTool: &anthropic.ToolChoiceToolParam{ + Type: constant.ValueOf[constant.Tool](), + Name: "specific_tool", + }, + }, + }, + }, + mcpProxy: &mockServerProxier{ + tools: []*mcp.Tool{{ID: "test_tool", Name: "test", Description: "Test"}}, + }, + } + + i.injectTools() + + require.NotNil(t, i.req.ToolChoice.OfTool) + require.True(t, i.req.ToolChoice.OfTool.DisableParallelToolUse.Valid()) + require.True(t, i.req.ToolChoice.OfTool.DisableParallelToolUse.Value) + }) + + t.Run("no-op for none tool choice type", func(t *testing.T) { + t.Parallel() + + i := &AnthropicMessagesInterceptionBase{ + req: &MessageNewParamsWrapper{ + MessageNewParams: anthropic.MessageNewParams{ + ToolChoice: anthropic.ToolChoiceUnionParam{ + OfNone: &anthropic.ToolChoiceNoneParam{ + Type: constant.ValueOf[constant.None](), + }, + }, + }, + }, + mcpProxy: &mockServerProxier{ + tools: []*mcp.Tool{{ID: "test_tool", Name: "test", Description: "Test"}}, + }, + } + + i.injectTools() + + // Tools are still injected. + require.Len(t, i.req.Tools, 1) + // But no parallel tool use modification for "none" type. + require.Nil(t, i.req.ToolChoice.OfAuto) + require.Nil(t, i.req.ToolChoice.OfAny) + require.NotNil(t, i.req.ToolChoice.OfNone) + }) +} + +// mockServerProxier is a test implementation of mcp.ServerProxier. +type mockServerProxier struct { + tools []*mcp.Tool +} + +func (m *mockServerProxier) Init(context.Context) error { + return nil +} + +func (m *mockServerProxier) Shutdown(context.Context) error { + return nil +} + +func (m *mockServerProxier) ListTools() []*mcp.Tool { + return m.tools +} + +func (m *mockServerProxier) GetTool(id string) *mcp.Tool { + for _, t := range m.tools { + if t.ID == id { + return t + } + } + return nil +} + +func (m *mockServerProxier) CallTool(context.Context, string, any) (*mcpgo.CallToolResult, error) { + return nil, nil +} diff --git a/intercept_anthropic_messages_base.go b/intercept_anthropic_messages_base.go index a9b8802..ced7d01 100644 --- a/intercept_anthropic_messages_base.go +++ b/intercept_anthropic_messages_base.go @@ -81,19 +81,50 @@ func (i *AnthropicMessagesInterceptionBase) injectTools() { return } - // Inject tools. - for _, tool := range i.mcpProxy.ListTools() { - i.req.Tools = append(i.req.Tools, anthropic.ToolUnionParam{ - OfTool: &anthropic.ToolParam{ - InputSchema: anthropic.ToolInputSchemaParam{ - Properties: tool.Params, - Required: tool.Required, + tools := i.mcpProxy.ListTools() + if len(tools) > 0 { + // Capture existing cache control breakpoint, if preset. + var cache *anthropic.CacheControlEphemeralParam + for _, t := range i.req.Tools { + if t.OfTool == nil { + continue + } + + if t.OfTool.CacheControl.Type != "" { + // Capture existing cache control breakpoint (copy values since we'll be clearing it in the next step). + cache = &anthropic.CacheControlEphemeralParam{ + TTL: t.OfTool.CacheControl.TTL, + Type: t.OfTool.CacheControl.Type, + } + // Reset it; we'll move this breakpoint to the final tool definition. + t.OfTool.CacheControl = anthropic.CacheControlEphemeralParam{} + break + } + } + + // Inject tools. + for _, tool := range tools { + i.req.Tools = append(i.req.Tools, anthropic.ToolUnionParam{ + OfTool: &anthropic.ToolParam{ + InputSchema: anthropic.ToolInputSchemaParam{ + Properties: tool.Params, + Required: tool.Required, + }, + Name: tool.ID, + Description: anthropic.String(tool.Description), + Type: anthropic.ToolTypeCustom, }, - Name: tool.ID, - Description: anthropic.String(tool.Description), - Type: anthropic.ToolTypeCustom, - }, - }) + }) + } + + // If there was a given cache control breakpoint, set it on the final tool as per the docs: + // https://platform.claude.com/docs/en/build-with-claude/prompt-caching#prompt-caching-examples (see "Caching tool definitions"). + count := len(i.req.Tools) + if cache != nil && count > 0 { + if i.req.Tools[count-1].OfTool != nil { + i.req.Tools[count-1].OfTool.CacheControl = *cache + } + } } // Note: Parallel tool calls are disabled to avoid tool_use/tool_result block mismatches. From 7f032141ea9c70f0bc8f32fadfdb75b370a8411d Mon Sep 17 00:00:00 2001 From: Danny Kopping Date: Tue, 23 Dec 2025 16:47:58 +0200 Subject: [PATCH 2/2] chore: only disable parallel tool calls when there are injected tools Signed-off-by: Danny Kopping --- bridge_integration_test.go | 45 +++++++++--- intercept_anthropic_message_internal_test.go | 23 ++++++ intercept_anthropic_messages_base.go | 75 ++++++++++---------- 3 files changed, 99 insertions(+), 44 deletions(-) diff --git a/bridge_integration_test.go b/bridge_integration_test.go index f73138e..ca0d55c 100644 --- a/bridge_integration_test.go +++ b/bridge_integration_test.go @@ -1369,7 +1369,8 @@ func TestStableRequestEncoding(t *testing.T) { } // TestAnthropicToolChoiceParallelDisabled verifies that parallel tool use is -// correctly disabled based on the tool_choice parameter in the request. +// correctly disabled based on the tool_choice parameter in the request, but +// only when MCP tools are injected. // See https://github.com/coder/aibridge/issues/2 func TestAnthropicToolChoiceParallelDisabled(t *testing.T) { t.Parallel() @@ -1384,39 +1385,61 @@ func TestAnthropicToolChoiceParallelDisabled(t *testing.T) { cases := []struct { name string toolChoice any // nil, or map with "type" key. + withInjectedTools bool expectDisableParallel bool expectToolChoiceTypeInRequest string }{ + // With injected tools - disable_parallel_tool_use should be set. { - name: "no tool_choice defined defaults to auto", + name: "with injected tools: no tool_choice defined defaults to auto", toolChoice: nil, + withInjectedTools: true, expectDisableParallel: true, expectToolChoiceTypeInRequest: toolChoiceAuto, }, { - name: "tool_choice auto", + name: "with injected tools: tool_choice auto", toolChoice: map[string]any{"type": toolChoiceAuto}, + withInjectedTools: true, expectDisableParallel: true, expectToolChoiceTypeInRequest: toolChoiceAuto, }, { - name: "tool_choice any", + name: "with injected tools: tool_choice any", toolChoice: map[string]any{"type": toolChoiceAny}, + withInjectedTools: true, expectDisableParallel: true, expectToolChoiceTypeInRequest: toolChoiceAny, }, { - name: "tool_choice tool", + name: "with injected tools: tool_choice tool", toolChoice: map[string]any{"type": toolChoiceTool, "name": "some_tool"}, + withInjectedTools: true, expectDisableParallel: true, expectToolChoiceTypeInRequest: toolChoiceTool, }, { - name: "tool_choice none", + name: "with injected tools: tool_choice none", toolChoice: map[string]any{"type": toolChoiceNone}, + withInjectedTools: true, expectDisableParallel: false, expectToolChoiceTypeInRequest: toolChoiceNone, }, + // Without injected tools - disable_parallel_tool_use should NOT be set. + { + name: "without injected tools: tool_choice auto", + toolChoice: map[string]any{"type": toolChoiceAuto}, + withInjectedTools: false, + expectDisableParallel: false, + expectToolChoiceTypeInRequest: toolChoiceAuto, + }, + { + name: "without injected tools: tool_choice any", + toolChoice: map[string]any{"type": toolChoiceAny}, + withInjectedTools: false, + expectDisableParallel: false, + expectToolChoiceTypeInRequest: toolChoiceAny, + }, } for _, tc := range cases { @@ -1426,8 +1449,14 @@ func TestAnthropicToolChoiceParallelDisabled(t *testing.T) { ctx, cancel := context.WithTimeout(t.Context(), time.Second*30) t.Cleanup(cancel) - // Configure the bridge. - mcpMgr := mcp.NewServerProxyManager(nil, testTracer) + // Setup MCP tools conditionally. + var mcpMgr *mcp.ServerProxyManager + if tc.withInjectedTools { + mcpProxiers, _ := setupMCPServerProxiesForTest(t, testTracer) + mcpMgr = mcp.NewServerProxyManager(mcpProxiers, testTracer) + } else { + mcpMgr = mcp.NewServerProxyManager(nil, testTracer) + } require.NoError(t, mcpMgr.Init(ctx)) arc := txtar.Parse(antSimple) diff --git a/intercept_anthropic_message_internal_test.go b/intercept_anthropic_message_internal_test.go index 600dba9..e1748db 100644 --- a/intercept_anthropic_message_internal_test.go +++ b/intercept_anthropic_message_internal_test.go @@ -166,6 +166,29 @@ func TestInjectTools_CacheBreakpoints(t *testing.T) { func TestInjectTools_ParallelToolCalls(t *testing.T) { t.Parallel() + t.Run("does not modify tool choice when no tools to inject", func(t *testing.T) { + t.Parallel() + + i := &AnthropicMessagesInterceptionBase{ + req: &MessageNewParamsWrapper{ + MessageNewParams: anthropic.MessageNewParams{ + ToolChoice: anthropic.ToolChoiceUnionParam{ + OfAuto: &anthropic.ToolChoiceAutoParam{ + Type: constant.ValueOf[constant.Auto](), + }, + }, + }, + }, + mcpProxy: &mockServerProxier{tools: nil}, // No tools to inject. + } + + i.injectTools() + + // Tool choice should remain unchanged - DisableParallelToolUse should not be set. + require.NotNil(t, i.req.ToolChoice.OfAuto) + require.False(t, i.req.ToolChoice.OfAuto.DisableParallelToolUse.Valid()) + }) + t.Run("disables parallel tool use for auto tool choice (default)", func(t *testing.T) { t.Parallel() diff --git a/intercept_anthropic_messages_base.go b/intercept_anthropic_messages_base.go index ced7d01..17f97bd 100644 --- a/intercept_anthropic_messages_base.go +++ b/intercept_anthropic_messages_base.go @@ -82,48 +82,51 @@ func (i *AnthropicMessagesInterceptionBase) injectTools() { } tools := i.mcpProxy.ListTools() - if len(tools) > 0 { - // Capture existing cache control breakpoint, if preset. - var cache *anthropic.CacheControlEphemeralParam - for _, t := range i.req.Tools { - if t.OfTool == nil { - continue - } + if len(tools) == 0 { + // No injected tools: no need to affect cache breakpoints or influence parallel tool calling. + return + } + + // Capture existing cache control breakpoint, if preset. + var cache *anthropic.CacheControlEphemeralParam + for _, t := range i.req.Tools { + if t.OfTool == nil { + continue + } - if t.OfTool.CacheControl.Type != "" { - // Capture existing cache control breakpoint (copy values since we'll be clearing it in the next step). - cache = &anthropic.CacheControlEphemeralParam{ - TTL: t.OfTool.CacheControl.TTL, - Type: t.OfTool.CacheControl.Type, - } - // Reset it; we'll move this breakpoint to the final tool definition. - t.OfTool.CacheControl = anthropic.CacheControlEphemeralParam{} - break + if t.OfTool.CacheControl.Type != "" { + // Capture existing cache control breakpoint (copy values since we'll be clearing it in the next step). + cache = &anthropic.CacheControlEphemeralParam{ + TTL: t.OfTool.CacheControl.TTL, + Type: t.OfTool.CacheControl.Type, } + // Reset it; we'll move this breakpoint to the final tool definition. + t.OfTool.CacheControl = anthropic.CacheControlEphemeralParam{} + break } + } - // Inject tools. - for _, tool := range tools { - i.req.Tools = append(i.req.Tools, anthropic.ToolUnionParam{ - OfTool: &anthropic.ToolParam{ - InputSchema: anthropic.ToolInputSchemaParam{ - Properties: tool.Params, - Required: tool.Required, - }, - Name: tool.ID, - Description: anthropic.String(tool.Description), - Type: anthropic.ToolTypeCustom, + // Inject tools. + for _, tool := range tools { + i.req.Tools = append(i.req.Tools, anthropic.ToolUnionParam{ + OfTool: &anthropic.ToolParam{ + InputSchema: anthropic.ToolInputSchemaParam{ + Properties: tool.Params, + Required: tool.Required, }, - }) - } + Name: tool.ID, + Description: anthropic.String(tool.Description), + Type: anthropic.ToolTypeCustom, + }, + }) + } - // If there was a given cache control breakpoint, set it on the final tool as per the docs: - // https://platform.claude.com/docs/en/build-with-claude/prompt-caching#prompt-caching-examples (see "Caching tool definitions"). - count := len(i.req.Tools) - if cache != nil && count > 0 { - if i.req.Tools[count-1].OfTool != nil { - i.req.Tools[count-1].OfTool.CacheControl = *cache - } + // If there was a given cache control breakpoint, set it on the final tool as per the docs: + // https://platform.claude.com/docs/en/build-with-claude/prompt-caching#prompt-caching-examples (see "Caching tool definitions"). + count := len(i.req.Tools) + if cache != nil && count > 0 { + if i.req.Tools[count-1].OfTool != nil { + i.req.Tools[count-1].OfTool.CacheControl = *cache } }