diff --git a/bridge_integration_test.go b/bridge_integration_test.go index a5df3e0..b0de5f8 100644 --- a/bridge_integration_test.go +++ b/bridge_integration_test.go @@ -651,6 +651,9 @@ func TestAnthropicInjectedTools(t *testing.T) { // Ensure tokens used during injected tool invocation are accounted for. assert.EqualValues(t, 15308, calculateTotalInputTokens(recorderClient.tokenUsages)) assert.EqualValues(t, 204, calculateTotalOutputTokens(recorderClient.tokenUsages)) + + // Ensure we received exactly one prompt. + require.Len(t, recorderClient.userPrompts, 1) }) } } @@ -744,6 +747,9 @@ func TestOpenAIInjectedTools(t *testing.T) { // Ensure tokens used during injected tool invocation are accounted for. require.EqualValues(t, 5047, calculateTotalInputTokens(recorderClient.tokenUsages)) require.EqualValues(t, 105, calculateTotalOutputTokens(recorderClient.tokenUsages)) + + // Ensure we received exactly one prompt. + require.Len(t, recorderClient.userPrompts, 1) }) } } diff --git a/intercept_anthropic_messages_blocking.go b/intercept_anthropic_messages_blocking.go index 86dad39..d80a25d 100644 --- a/intercept_anthropic_messages_blocking.go +++ b/intercept_anthropic_messages_blocking.go @@ -91,6 +91,7 @@ func (i *AnthropicMessagesBlockingInterception) ProcessRequest(w http.ResponseWr MsgID: resp.ID, Prompt: *prompt, }) + prompt = nil } _ = i.recorder.RecordTokenUsage(ctx, &TokenUsageRecord{ diff --git a/intercept_anthropic_messages_streaming.go b/intercept_anthropic_messages_streaming.go index ef5259a..c2ad7e0 100644 --- a/intercept_anthropic_messages_streaming.go +++ b/intercept_anthropic_messages_streaming.go @@ -405,6 +405,7 @@ newStream: MsgID: message.ID, Prompt: *prompt, }) + prompt = nil } // Check if the stream encountered any errors. diff --git a/intercept_openai_chat_blocking.go b/intercept_openai_chat_blocking.go index b6a0343..3b1fa7e 100644 --- a/intercept_openai_chat_blocking.go +++ b/intercept_openai_chat_blocking.go @@ -73,6 +73,7 @@ func (i *OpenAIBlockingChatInterception) ProcessRequest(w http.ResponseWriter, r MsgID: completion.ID, Prompt: *prompt, }) + prompt = nil } lastUsage := completion.Usage diff --git a/intercept_openai_chat_streaming.go b/intercept_openai_chat_streaming.go index f10c26e..931b131 100644 --- a/intercept_openai_chat_streaming.go +++ b/intercept_openai_chat_streaming.go @@ -157,6 +157,7 @@ func (i *OpenAIStreamingChatInterception) ProcessRequest(w http.ResponseWriter, MsgID: processor.getMsgID(), Prompt: *prompt, }) + prompt = nil } if lastUsage := processor.getLastUsage(); lastUsage.CompletionTokens > 0 {