fix: send model thoughts with tool usage recording

Signed-off-by: Danny Kopping <danny@coder.com>

Author: dannykopping

bridge_integration_test.go

@@ -216,20 +216,17 @@ func TestAnthropicMessagesModelThoughts(t *testing.T) {
 					assert.Contains(t, sp.AllEvents(), "message_stop")
 				}
 
-				// Verify model thoughts were captured and associated with the tool call.
-				thoughts := recorderClient.RecordedModelThoughts()
-				require.Len(t, thoughts, 1)
-				assert.Contains(t, thoughts[0].Content, "The user wants me to read")
-				assert.Contains(t, thoughts[0].Content, tc.expectedThinkingSubstr)
-				assert.NotEmpty(t, thoughts[0].InterceptionID)
-				assert.Equal(t, tc.expectedToolCallID, thoughts[0].ProviderToolCallID)
-
-				// Verify tool usage was also recorded.
+				// Verify tool usage was recorded with associated model thoughts.
 				toolUsages := recorderClient.RecordedToolUsages()
 				require.Len(t, toolUsages, 1)
 				assert.Equal(t, "Read", toolUsages[0].Tool)
 				assert.Equal(t, tc.expectedToolCallID, toolUsages[0].ToolCallID)
 
+				// Model thoughts should be embedded in the tool usage record.
+				require.Len(t, toolUsages[0].ModelThoughts, 1)
+				assert.Contains(t, toolUsages[0].ModelThoughts[0].Content, "The user wants me to read")
+				assert.Contains(t, toolUsages[0].ModelThoughts[0].Content, tc.expectedThinkingSubstr)
+
 				recorderClient.VerifyAllInterceptionsEnded(t)
 			})
 		}
@@ -271,9 +268,10 @@ func TestAnthropicMessagesModelThoughts(t *testing.T) {
 		sp := aibridge.NewSSEParser()
 		require.NoError(t, sp.Parse(resp.Body))
 
-		// No thoughts should be recorded when there are no tool calls.
-		thoughts := recorderClient.RecordedModelThoughts()
-		assert.Empty(t, thoughts)
+		// No tool usages (and therefore no thoughts) should be recorded
+		// when there are no tool calls.
+		toolUsages := recorderClient.RecordedToolUsages()
+		assert.Empty(t, toolUsages)
 
 		recorderClient.VerifyAllInterceptionsEnded(t)
 	})

intercept/messages/blocking.go

@@ -137,18 +137,16 @@ func (i *BlockingInterception) ProcessRequest(w http.ResponseWriter, r *http.Req
 
 		// Capture any thinking blocks that were returned.
 		var thoughtRecords []*recorder.ModelThoughtRecord
-		if !i.isSmallFastModel() {
-			for _, block := range resp.Content {
-				switch variant := block.AsAny().(type) {
-				case anthropic.ThinkingBlock:
-					thoughtRecords = append(thoughtRecords, &recorder.ModelThoughtRecord{
-						InterceptionID: i.ID().String(),
-						Content:        variant.Thinking,
-					})
-				case anthropic.RedactedThinkingBlock:
-					// For redacted thinking, there's nothing useful we can capture.
-					continue
-				}
+		for _, block := range resp.Content {
+			switch variant := block.AsAny().(type) {
+			case anthropic.ThinkingBlock:
+				thoughtRecords = append(thoughtRecords, &recorder.ModelThoughtRecord{
+					Content:   variant.Thinking,
+					CreatedAt: time.Now(),
+				})
+			case anthropic.RedactedThinkingBlock:
+				// For redacted thinking, there's nothing useful we can capture.
+				continue
 			}
 		}
 
@@ -173,22 +171,15 @@ func (i *BlockingInterception) ProcessRequest(w http.ResponseWriter, r *http.Req
 				Tool:           toolUse.Name,
 				Args:           toolUse.Input,
 				Injected:       false,
+				ModelThoughts:  thoughtRecords,
 			})
-
-			// Associate the model thoughts with this tool call.
-			for _, thought := range thoughtRecords {
-				thought.ProviderToolCallID = toolUse.ID
-			}
+			// Clear after first use to avoid duplicating across
+			// multiple tool calls in the same message.
+			thoughtRecords = nil
 		}
 
-		// If no injected tool calls, persist thoughts and we're done.
+		// If no injected tool calls, we're done.
 		if len(pendingToolCalls) == 0 {
-			for _, thought := range thoughtRecords {
-				if thought.ProviderToolCallID == "" {
-					continue
-				}
-				_ = i.recorder.RecordModelThought(ctx, thought)
-			}
 			break
 		}
 
@@ -223,12 +214,11 @@ func (i *BlockingInterception) ProcessRequest(w http.ResponseWriter, r *http.Req
 				Args:            tc.Input,
 				Injected:        true,
 				InvocationError: err,
+				ModelThoughts:   thoughtRecords,
 			})
-
-			// Associate the model thoughts with this tool call.
-			for _, thought := range thoughtRecords {
-				thought.ProviderToolCallID = tc.ID
-			}
+			// Clear after first use to avoid duplicating across
+			// multiple tool calls in the same message.
+			thoughtRecords = nil
 
 			if err != nil {
 				// Always provide a tool_result even if the tool call failed
@@ -315,14 +305,6 @@ func (i *BlockingInterception) ProcessRequest(w http.ResponseWriter, r *http.Req
 			}
 		}
 
-		// Only persist thoughts that are associated to a tool call.
-		for _, thought := range thoughtRecords {
-			if thought.ProviderToolCallID == "" {
-				continue
-			}
-			_ = i.recorder.RecordModelThought(ctx, thought)
-		}
-
 		// Sync the raw payload with updated messages so that withBody()
 		// sends the updated payload on the next iteration.
 		if err := i.syncPayloadMessages(messages.Messages); err != nil {

intercept/messages/streaming.go

@@ -255,18 +255,16 @@ newStream:
 
 				// Capture any thinking blocks that were returned.
 				var thoughtRecords []*recorder.ModelThoughtRecord
-				if !i.isSmallFastModel() { // TODO: remove.
-					for _, block := range message.Content {
-						switch variant := block.AsAny().(type) {
-						case anthropic.ThinkingBlock:
-							thoughtRecords = append(thoughtRecords, &recorder.ModelThoughtRecord{
-								InterceptionID: i.ID().String(),
-								Content:        variant.Thinking,
-							})
-						case anthropic.RedactedThinkingBlock:
-							// For redacted thinking, there's nothing useful we can capture.
-							continue
-						}
+				for _, block := range message.Content {
+					switch variant := block.AsAny().(type) {
+					case anthropic.ThinkingBlock:
+						thoughtRecords = append(thoughtRecords, &recorder.ModelThoughtRecord{
+							Content:   variant.Thinking,
+							CreatedAt: time.Now(),
+						})
+					case anthropic.RedactedThinkingBlock:
+						// For redacted thinking, there's nothing useful we can capture.
+						continue
 					}
 				}
 
@@ -322,12 +320,11 @@ newStream:
 							Args:            input,
 							Injected:        true,
 							InvocationError: err,
+							ModelThoughts:   thoughtRecords,
 						})
-
-						// Associate the model thoughts with this tool call.
-						for _, thought := range thoughtRecords {
-							thought.ProviderToolCallID = id
-						}
+						// Clear after first use to avoid duplicating across
+						// multiple tool calls in the same message.
+						thoughtRecords = nil
 
 						if err != nil {
 							// Always provide a tool_result even if the tool call failed
@@ -413,15 +410,6 @@ newStream:
 						}
 					}
 
-					// Only persist thoughts that are associated to a tool call.
-					for _, thought := range thoughtRecords {
-						if thought.ProviderToolCallID == "" {
-							continue
-						}
-
-						_ = i.recorder.RecordModelThought(streamCtx, thought)
-					}
-
 					// Sync the raw payload with updated messages so that withBody()
 					// sends the updated payload on the next iteration.
 					if syncErr := i.syncPayloadMessages(messages.Messages); syncErr != nil {
@@ -448,23 +436,13 @@ newStream:
 								Tool:           variant.Name,
 								Args:           variant.Input,
 								Injected:       false,
+								ModelThoughts:  thoughtRecords,
 							})
-
-							// Associate the model thoughts with this tool call.
-							for _, thought := range thoughtRecords {
-								thought.ProviderToolCallID = variant.ID
-							}
+							// Clear after first use to avoid duplicating across
+							// multiple tool calls in the same message.
+							thoughtRecords = nil
 						}
 					}
-
-					// Only persist thoughts that are associated to a tool call.
-					for _, thought := range thoughtRecords {
-						if thought.ProviderToolCallID == "" {
-							continue
-						}
-
-						_ = i.recorder.RecordModelThought(streamCtx, thought)
-					}
 				}
 			}
 

internal/testutil/mock_recorder.go

@@ -20,7 +20,6 @@ type MockRecorder struct {
 	tokenUsages      []*recorder.TokenUsageRecord
 	userPrompts      []*recorder.PromptUsageRecord
 	toolUsages       []*recorder.ToolUsageRecord
-	modelThoughts    []*recorder.ModelThoughtRecord
 	interceptionsEnd map[string]*recorder.InterceptionRecordEnded
 }
 
@@ -65,13 +64,6 @@ func (m *MockRecorder) RecordToolUsage(ctx context.Context, req *recorder.ToolUs
 	return nil
 }
 
-func (m *MockRecorder) RecordModelThought(ctx context.Context, req *recorder.ModelThoughtRecord) error {
-	m.mu.Lock()
-	defer m.mu.Unlock()
-	m.modelThoughts = append(m.modelThoughts, req)
-	return nil
-}
-
 // RecordedTokenUsages returns a copy of recorded token usages in a thread-safe manner.
 // Note: This is a shallow clone - the slice is copied but the pointers reference the
 // same underlying records. This is sufficient for our test assertions which only read
@@ -114,13 +106,6 @@ func (m *MockRecorder) ToolUsages() []*recorder.ToolUsageRecord {
 	return m.toolUsages
 }
 
-// RecordedModelThoughts returns a copy of recorded model thoughts in a thread-safe manner.
-func (m *MockRecorder) RecordedModelThoughts() []*recorder.ModelThoughtRecord {
-	m.mu.Lock()
-	defer m.mu.Unlock()
-	return slices.Clone(m.modelThoughts)
-}
-
 // RecordedInterceptionEnd returns the stored InterceptionRecordEnded for the
 // given interception ID, or nil if not found.
 func (m *MockRecorder) RecordedInterceptionEnd(id string) *recorder.InterceptionRecordEnded {

recorder/recorder.go

@@ -116,24 +116,6 @@ func (r *RecorderWrapper) RecordToolUsage(ctx context.Context, req *ToolUsageRec
 	return err
 }
 
-func (r *RecorderWrapper) RecordModelThought(ctx context.Context, req *ModelThoughtRecord) (outErr error) {
-	ctx, span := r.tracer.Start(ctx, "Intercept.RecordModelThought", trace.WithAttributes(tracing.InterceptionAttributesFromContext(ctx)...))
-	defer tracing.EndSpanErr(span, &outErr)
-
-	client, err := r.clientFn()
-	if err != nil {
-		return fmt.Errorf("acquire client: %w", err)
-	}
-
-	req.CreatedAt = time.Now()
-	if err = client.RecordModelThought(ctx, req); err == nil {
-		return nil
-	}
-
-	r.logger.Warn(ctx, "failed to record model thought", slog.Error(err), slog.F("interception_id", req.InterceptionID))
-	return err
-}
-
 func NewRecorder(logger slog.Logger, tracer trace.Tracer, clientFn func() (Recorder, error)) *RecorderWrapper {
 	return &RecorderWrapper{
 		logger:   logger,
@@ -270,22 +252,6 @@ func (a *AsyncRecorder) RecordToolUsage(ctx context.Context, req *ToolUsageRecor
 	return nil // Caller is not interested in error.
 }
 
-func (a *AsyncRecorder) RecordModelThought(ctx context.Context, req *ModelThoughtRecord) error {
-	a.wg.Add(1)
-	go func() {
-		defer a.wg.Done()
-		timedCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), a.timeout)
-		defer cancel()
-
-		err := a.wrapped.RecordModelThought(timedCtx, req)
-		if err != nil {
-			a.logger.Warn(timedCtx, "failed to record model thought", slog.Error(err), slog.F("payload", req))
-		}
-	}()
-
-	return nil // Caller is not interested in error.
-}
-
 func (a *AsyncRecorder) Wait() {
 	a.wg.Wait()
 }

recorder/types.go

@@ -18,9 +18,8 @@ type Recorder interface {
 	// RecordPromptUsage records the prompts used in an interception with an upstream AI provider.
 	RecordPromptUsage(ctx context.Context, req *PromptUsageRecord) error
 	// RecordToolUsage records the tools used in an interception with an upstream AI provider.
+	// Any associated model thoughts should be included in the ToolUsageRecord.
 	RecordToolUsage(ctx context.Context, req *ToolUsageRecord) error
-	// RecordModelThought records the reasoning/thinking produced in an interception with an upstream AI provider.
-	RecordModelThought(ctx context.Context, req *ModelThoughtRecord) error
 }
 
 type ToolArgs any
@@ -74,12 +73,11 @@ type ToolUsageRecord struct {
 	InvocationError error
 	Metadata        Metadata
 	CreatedAt       time.Time
+	ModelThoughts   []*ModelThoughtRecord
 }
 
 type ModelThoughtRecord struct {
-	InterceptionID string
-	ProviderToolCallID string
-	Content        string
-	Metadata       Metadata
-	CreatedAt      time.Time
+	Content   string
+	Metadata  Metadata
+	CreatedAt time.Time
 }

trace_integration_test.go

@@ -46,7 +46,6 @@ func TestTraceAnthropic(t *testing.T) {
 		{"Intercept.RecordPromptUsage", 1, codes.Unset},
 		{"Intercept.RecordTokenUsage", 1, codes.Unset},
 		{"Intercept.RecordToolUsage", 1, codes.Unset},
-		{"Intercept.RecordModelThought", 1, codes.Unset},
 		{"Intercept.ProcessRequest.Upstream", 1, codes.Unset},
 	}
 
@@ -59,7 +58,6 @@ func TestTraceAnthropic(t *testing.T) {
 		{"Intercept.RecordPromptUsage", 1, codes.Unset},
 		{"Intercept.RecordTokenUsage", 2, codes.Unset},
 		{"Intercept.RecordToolUsage", 1, codes.Unset},
-		{"Intercept.RecordModelThought", 1, codes.Unset},
 		{"Intercept.ProcessRequest.Upstream", 1, codes.Unset},
 	}