fix: enforce tool-use lineage invariants before chat requests

Author: binarylogic

internal/chat/client.go

@@ -170,6 +170,7 @@ func (c *client) StreamSnapshots(ctx context.Context, req Request, onSnapshot fu
 	for _, summary := range summarizeRequestMessages(req.Messages) {
 		c.scope.Debug("chat request message", "summary", summary)
 	}
+	c.scope.Debug("chat request tool lineage", "summary", summarizeToolLineage(req.Messages))
 
 	// Get fresh token for this request
 	token, err := c.auth.GetAccessToken(ctx)
@@ -344,3 +345,19 @@ func summarizeRequestMessages(messages []domain.Message) []string {
 	}
 	return out
 }
+
+func summarizeToolLineage(messages []domain.Message) string {
+	toolUses := make([]string, 0)
+	toolResults := make([]string, 0)
+	for _, msg := range messages {
+		for _, b := range msg.Content {
+			if b.ToolUse != nil && b.ToolUse.ID != "" {
+				toolUses = append(toolUses, b.ToolUse.ID)
+			}
+			if b.ToolResult != nil && b.ToolResult.ToolUseID != "" {
+				toolResults = append(toolResults, b.ToolResult.ToolUseID)
+			}
+		}
+	}
+	return fmt.Sprintf("tool_use_ids=[%s] tool_result_ids=[%s]", strings.Join(toolUses, ","), strings.Join(toolResults, ","))
+}

internal/chat/protocol_mapping_test.go

@@ -119,27 +119,40 @@ func TestToWireRequest_SanitizesToolResultContent(t *testing.T) {
 
 	wireReq, err := toWireRequest(Request{
 		ConversationID: "00000000-0000-0000-0000-000000000001",
-		Messages: []domain.Message{{
-			Role: domain.RoleUser,
-			Content: []domain.Block{{
-				Type: domain.BlockTypeToolResult,
-				ToolResult: &domain.ToolResult{
-					ToolUseID: "tool-1",
-					Content: map[string]any{
-						"tool_use_id": "tool-1",
-						"rows":        []map[string]any{{"id": "svc-1"}},
+		Messages: []domain.Message{
+			{
+				Role: domain.RoleAssistant,
+				Content: []domain.Block{{
+					Type: domain.BlockTypeToolUse,
+					ToolUse: &domain.ToolUse{
+						ID:    "tool-1",
+						Name:  "query",
+						Input: json.RawMessage(`{"sql":"select 1"}`),
 					},
-				},
-			}},
-		}},
+				}},
+			},
+			{
+				Role: domain.RoleUser,
+				Content: []domain.Block{{
+					Type: domain.BlockTypeToolResult,
+					ToolResult: &domain.ToolResult{
+						ToolUseID: "tool-1",
+						Content: map[string]any{
+							"tool_use_id": "tool-1",
+							"rows":        []map[string]any{{"id": "svc-1"}},
+						},
+					},
+				}},
+			},
+		},
 	})
 	if err != nil {
 		t.Fatalf("toWireRequest() error = %v", err)
 	}
-	if len(wireReq.Messages) != 1 || len(wireReq.Messages[0].Content) != 1 {
+	if len(wireReq.Messages) != 2 || len(wireReq.Messages[1].Content) != 1 {
 		t.Fatalf("unexpected wire request shape: %#v", wireReq)
 	}
-	content := wireReq.Messages[0].Content[0].ToolResult.Content
+	content := wireReq.Messages[1].Content[0].ToolResult.Content
 	var parsed map[string]any
 	if err := json.Unmarshal(content, &parsed); err != nil {
 		t.Fatalf("tool_result.content should be JSON object, got error: %v", err)
@@ -196,3 +209,33 @@ func TestToWireRequest_RejectsInvalidContextEntity(t *testing.T) {
 		t.Fatalf("error = %v", err)
 	}
 }
+
+func TestToWireRequest_RejectsUnknownToolUseReference(t *testing.T) {
+	t.Parallel()
+
+	_, err := toWireRequest(Request{
+		ConversationID: "00000000-0000-0000-0000-000000000001",
+		Messages: []domain.Message{
+			{
+				Role: domain.RoleUser,
+				Content: []domain.Block{{
+					Type: domain.BlockTypeText,
+					Text: &domain.TextBlock{Content: "hello"},
+				}},
+			},
+			{
+				Role: domain.RoleUser,
+				Content: []domain.Block{{
+					Type: domain.BlockTypeToolResult,
+					ToolResult: &domain.ToolResult{
+						ToolUseID: "toolu_missing",
+						Content:   map[string]any{"rows": []any{}},
+					},
+				}},
+			},
+		},
+	})
+	if err == nil || !strings.Contains(err.Error(), `unknown tool_use_id "toolu_missing"`) {
+		t.Fatalf("error = %v", err)
+	}
+}

internal/chat/protocolv2/request.go

@@ -121,6 +121,9 @@ func Validate(req Request) error {
 			return fmt.Errorf("context_entities[%d]: %w", i, err)
 		}
 	}
+	if err := validateToolReferences(req.Messages); err != nil {
+		return err
+	}
 	return nil
 }
 
@@ -224,3 +227,27 @@ func validateContextEntity(entity ContextEntity) error {
 	}
 	return nil
 }
+
+func validateToolReferences(messages []Message) error {
+	seen := make(map[string]struct{})
+	for i, msg := range messages {
+		for j, block := range msg.Content {
+			switch block.Type {
+			case BlockTypeText, BlockTypeThinking:
+				// No cross-message linkage requirements.
+			case BlockTypeToolUse:
+				id := block.ToolUse.ID
+				if _, exists := seen[id]; exists {
+					return fmt.Errorf("messages[%d].content[%d]: duplicate tool_use id %q", i, j, id)
+				}
+				seen[id] = struct{}{}
+			case BlockTypeToolResult:
+				id := block.ToolResult.ToolUseID
+				if _, exists := seen[id]; !exists {
+					return fmt.Errorf("messages[%d].content[%d]: unknown tool_use_id %q", i, j, id)
+				}
+			}
+		}
+	}
+	return nil
+}

internal/chat/protocolv2/request_test.go

@@ -195,4 +195,34 @@ func TestValidate(t *testing.T) {
 			t.Fatalf("error = %v", err)
 		}
 	})
+
+	t.Run("unknown tool_use_id fails", func(t *testing.T) {
+		t.Parallel()
+		req := newValidRequest()
+		req.Messages[1].Content[0].ToolResult.ToolUseID = "toolu_missing"
+		err := Validate(req)
+		if err == nil || !strings.Contains(err.Error(), `unknown tool_use_id "toolu_missing"`) {
+			t.Fatalf("error = %v", err)
+		}
+	})
+
+	t.Run("duplicate tool_use id fails", func(t *testing.T) {
+		t.Parallel()
+		req := newValidRequest()
+		req.Messages = append(req.Messages, Message{
+			Role: RoleAssistant,
+			Content: []Block{{
+				Type: BlockTypeToolUse,
+				ToolUse: &ToolUse{
+					ID:    "toolu_1",
+					Name:  "query",
+					Input: json.RawMessage(`{"sql":"select 2"}`),
+				},
+			}},
+		})
+		err := Validate(req)
+		if err == nil || !strings.Contains(err.Error(), `duplicate tool_use id "toolu_1"`) {
+			t.Fatalf("error = %v", err)
+		}
+	})
 }