odidere/internal/conversation/conversation_test.go

package conversation

import (
	"context"
	"encoding/json/v2"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"reflect"
	"sync"
	"testing"
	"time"

	"code.chimeric.al/chimerical/odidere/internal/config"
	"code.chimeric.al/chimerical/odidere/internal/db"
	"code.chimeric.al/chimerical/odidere/internal/llm"
	"code.chimeric.al/chimerical/odidere/internal/tool"
)

// newTestConversation returns a Conversation over a fresh saved
// conversation, using tools (an empty registry when nil).
func newTestConversation(t *testing.T, tools *tool.Registry) *Conversation {
	t.Helper()
	var (
		ctx = t.Context()
		log = slog.New(slog.DiscardHandler)
	)
	database, err := db.Connect(ctx, db.Config{
		Path: filepath.Join(
			t.TempDir(), "conversation.db",
		),
	})
	if err != nil {
		t.Fatalf("Connect: %v", err)
	}
	t.Cleanup(func() { database.Close() })
	if err := database.ApplySchema(ctx, log); err != nil {
		t.Fatalf("ApplySchema: %v", err)
	}
	saved, err := db.New(database.Writer).InsertConversation(
		ctx, db.InsertConversationParams{UserId: "default"},
	)
	if err != nil {
		t.Fatalf("InsertConversation: %v", err)
	}
	if tools == nil {
		tools, err = tool.NewRegistry(nil)
		if err != nil {
			t.Fatalf("NewRegistry: %v", err)
		}
	}
	c, err := New(Config{Id: saved.Id}, database, log, tools)
	if err != nil {
		t.Fatalf("New: %v", err)
	}
	return c
}

// userMessage returns a minimal user message holding text.
func userMessage(text string) *llm.Message {
	return &llm.Message{
		Role: llm.RoleUser,
		ContentParts: []llm.ContentPart{{
			Type: llm.ContentTypeText,
			Text: text,
		}},
	}
}

// infer reserves c and runs one inference call to completion.
func infer(
	ctx context.Context,
	c *Conversation,
	client *llm.Client,
	req InferenceRequest,
) error {
	inf, err := c.Reserve(ctx, client, req)
	if err != nil {
		return err
	}
	return inf.Run()
}

// newTestModelClient returns an llm client pointed at a server that
// answers with handler.
func newTestModelClient(
	t *testing.T,
	handler http.HandlerFunc,
) *llm.Client {
	t.Helper()
	srv := httptest.NewServer(handler)
	t.Cleanup(srv.Close)
	client, err := llm.NewClient(
		llm.Config{Provider: "test", URL: srv.URL},
		slog.New(slog.DiscardHandler),
	)
	if err != nil {
		t.Fatalf("NewClient: %v", err)
	}
	return client
}

// collectEvents drains ch until it closes and returns the events.
// It fails the test if the stream does not close in time.
func collectEvents(t *testing.T, ch <-chan Event) []Event {
	t.Helper()
	var events []Event
	for {
		select {
		case e, ok := <-ch:
			if !ok {
				return events
			}
			events = append(events, e)
		case <-time.After(5 * time.Second):
			t.Fatalf(
				"timed out waiting for the stream to close; got %v so far",
				events,
			)
			return nil
		}
	}
}

// TestNew checks dependency validation in New.
func TestNew(t *testing.T) {
	var (
		ctx = t.Context()
		log = slog.New(slog.DiscardHandler)
	)
	database, err := db.Connect(ctx, db.Config{
		Path: filepath.Join(t.TempDir(), "conversation.db"),
	})
	if err != nil {
		t.Fatalf("Connect: %v", err)
	}
	t.Cleanup(func() { database.Close() })
	tools, err := tool.NewRegistry(nil)
	if err != nil {
		t.Fatalf("NewRegistry: %v", err)
	}
	validCfg := Config{Id: "conv_1"}

	tests := []struct {
		name    string
		cfg     Config
		sdb     *db.SQLite
		log     *slog.Logger
		tools   *tool.Registry
		wantErr bool
	}{
		{
			name:    "valid",
			cfg:     validCfg,
			sdb:     database,
			log:     log,
			tools:   tools,
			wantErr: false,
		},
		{
			name:    "invalid config",
			cfg:     Config{},
			sdb:     database,
			log:     log,
			tools:   tools,
			wantErr: true,
		},
		{
			name:    "missing database",
			cfg:     validCfg,
			sdb:     nil,
			log:     log,
			tools:   tools,
			wantErr: true,
		},
		{
			name:    "missing logger",
			cfg:     validCfg,
			sdb:     database,
			log:     nil,
			tools:   tools,
			wantErr: true,
		},
		{
			name:    "missing tools",
			cfg:     validCfg,
			sdb:     database,
			log:     log,
			tools:   nil,
			wantErr: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			c, err := New(tt.cfg, tt.sdb, tt.log, tt.tools)
			if (err != nil) != tt.wantErr {
				t.Fatalf(
					"New() error = %v, wantErr %v",
					err, tt.wantErr,
				)
			}
			if !tt.wantErr && c == nil {
				t.Fatal("New() returned nil Conversation")
			}
		})
	}
}

// TestDecodeMessages checks decoding and validating saved message rows.
func TestDecodeMessages(t *testing.T) {
	validContent, err := json.Marshal(userMessage("hello"))
	if err != nil {
		t.Fatalf("Marshal: %v", err)
	}
	invalidContent, err := json.Marshal(llm.Message{
		Role:       llm.RoleUser,
		ToolCallID: "call_1",
	})
	if err != nil {
		t.Fatalf("Marshal: %v", err)
	}

	tests := []struct {
		name    string
		rows    []db.Message
		wantLen int
		wantErr bool
	}{
		{
			name:    "empty",
			rows:    nil,
			wantLen: 0,
			wantErr: false,
		},
		{
			name: "valid rows",
			rows: []db.Message{
				{Id: "m1", Content: validContent},
				{Id: "m2", Content: validContent},
			},
			wantLen: 2,
			wantErr: false,
		},
		{
			name: "malformed json",
			rows: []db.Message{
				{Id: "m1", Content: []byte("{not json")},
			},
			wantErr: true,
		},
		{
			name: "invalid message",
			rows: []db.Message{
				{Id: "m1", Content: invalidContent},
			},
			wantErr: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			msgs, err := DecodeMessages(tt.rows)
			if (err != nil) != tt.wantErr {
				t.Fatalf(
					"DecodeMessages() error = %v, wantErr %v",
					err, tt.wantErr,
				)
			}
			if tt.wantErr && msgs != nil {
				t.Fatalf("msgs = %v, want nil on error", msgs)
			}
			if !tt.wantErr && len(msgs) != tt.wantLen {
				t.Fatalf(
					"len(msgs) = %d, want %d",
					len(msgs), tt.wantLen,
				)
			}
		})
	}
}

// TestInferenceRequestValidate checks InferenceRequest validation rules.
func TestInferenceRequestValidate(t *testing.T) {
	negTwo := -2
	negOne := -1
	zero := 0

	tests := []struct {
		name    string
		req     InferenceRequest
		wantErr bool
	}{
		{
			name:    "empty request",
			req:     InferenceRequest{},
			wantErr: false,
		},
		{
			name: "valid user message and limits",
			req: InferenceRequest{
				Message:              userMessage("hi"),
				MaxIterations:        5,
				ThinkingBudgetTokens: &negOne,
				Timeout:              time.Minute,
			},
			wantErr: false,
		},
		{
			name: "zero thinking budget",
			req: InferenceRequest{
				ThinkingBudgetTokens: &zero,
			},
			wantErr: false,
		},
		{
			name: "negative max iterations",
			req: InferenceRequest{
				MaxIterations: -1,
			},
			wantErr: true,
		},
		{
			name: "thinking budget below -1",
			req: InferenceRequest{
				ThinkingBudgetTokens: &negTwo,
			},
			wantErr: true,
		},
		{
			name: "negative timeout",
			req: InferenceRequest{
				Timeout: -time.Second,
			},
			wantErr: true,
		},
		{
			name: "non-user message role",
			req: InferenceRequest{
				Message: &llm.Message{
					Role: llm.RoleAssistant,
					ContentParts: []llm.ContentPart{{
						Type: llm.ContentTypeText,
						Text: "assistant",
					}},
				},
			},
			wantErr: true,
		},
		{
			name: "user message with tool calls",
			req: InferenceRequest{
				Message: &llm.Message{
					Role: llm.RoleUser,
					ToolCalls: []llm.ToolCall{{
						ID: "call_1",
					}},
				},
			},
			wantErr: true,
		},
		{
			name: "user message with refusal",
			req: InferenceRequest{
				Message: &llm.Message{
					Role:    llm.RoleUser,
					Refusal: "nope",
				},
			},
			wantErr: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			err := tt.req.Validate()
			if (err != nil) != tt.wantErr {
				t.Fatalf(
					"Validate() error = %v, wantErr %v",
					err, tt.wantErr,
				)
			}
		})
	}
}

// TestInsertMessage checks saved rows and concurrent appends.
func TestInsertMessage(t *testing.T) {
	tests := []struct {
		name  string
		role  llm.Role
		count int
	}{
		{
			name:  "user message",
			role:  llm.RoleUser,
			count: 1,
		},
		{
			name:  "assistant message",
			role:  llm.RoleAssistant,
			count: 1,
		},
		{
			name:  "concurrent appends",
			role:  llm.RoleUser,
			count: 8,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			var (
				ctx = t.Context()
				c   = newTestConversation(t, nil)
				msg = llm.Message{
					Role: tt.role,
					ContentParts: []llm.ContentPart{{
						Type: llm.ContentTypeText,
						Text: "manually added",
					}},
				}
			)

			// Start callers together; each must commit once.
			// Their ordering is not guaranteed.
			var (
				wg      sync.WaitGroup
				start   = make(chan struct{})
				results = make(chan *db.Message, tt.count)
			)
			for range tt.count {
				wg.Go(func() {
					<-start
					saved, err := c.InsertMessage(ctx, msg)
					if err != nil {
						t.Errorf(
							"InsertMessage: %v",
							err,
						)
						return
					}
					results <- saved
				})
			}
			close(start)
			wg.Wait()
			close(results)

			rows, err := c.dbr.SelectMessagesConversationId(
				ctx, c.id,
			)
			if err != nil {
				t.Fatalf(
					"SelectMessagesConversationId: %v", err,
				)
			}
			if len(rows) != tt.count {
				t.Fatalf(
					"saved rows = %d, want %d",
					len(rows), tt.count,
				)
			}
			content, err := json.Marshal(msg)
			if err != nil {
				t.Fatalf("Marshal: %v", err)
			}
			for i, row := range rows {
				if row.Id == "" {
					t.Errorf("row %d has no Id", i)
				}
				if row.ConversationId != c.id {
					t.Errorf(
						"row %d conversation = %q, want %q",
						i, row.ConversationId, c.id,
					)
				}
				if row.Sequence != i+1 {
					t.Errorf(
						"row %d sequence = %d, want %d",
						i, row.Sequence, i+1,
					)
				}
				if row.Role != db.MessageRole(tt.role) {
					t.Errorf(
						"row %d role = %q, want %q",
						i, row.Role, tt.role,
					)
				}
				if string(row.Content) != string(content) {
					t.Errorf(
						"row %d content = %q, want %s",
						i, row.Content, content,
					)
				}
				if row.Metadata != nil {
					t.Errorf(
						"row %d metadata = %s, want nil",
						i, row.Metadata,
					)
				}
				if row.CreatedAt.IsZero() {
					t.Errorf("row %d has no created_at", i)
				}
				if row.UpdatedAt.IsZero() {
					t.Errorf("row %d has no updated_at", i)
				}
			}
			var returned int
			for saved := range results {
				returned++
				if saved == nil {
					t.Error("append returned nil")
					continue
				}
				index := saved.Sequence - 1
				if index < 0 || index >= len(rows) {
					t.Errorf(
						"sequence %d outside history",
						saved.Sequence,
					)
					continue
				}
				if !reflect.DeepEqual(rows[index], *saved) {
					t.Errorf(
						"row %q differs from saved row",
						saved.Id,
					)
				}
			}
			if returned != tt.count {
				t.Errorf(
					"returned rows = %d, want %d",
					returned, tt.count,
				)
			}
		})
	}
}

// TestInsertMessageErrors checks rejected appends leave no saved message.
func TestInsertMessageErrors(t *testing.T) {
	tests := []struct {
		name    string
		role    llm.Role
		failure string
		wantErr error
	}{
		{
			name:    "inference running",
			role:    llm.RoleUser,
			failure: "busy",
			wantErr: ErrBusy,
		},
		{
			name:    "inference stopped but not finished",
			role:    llm.RoleUser,
			failure: "stopped",
			wantErr: ErrBusy,
		},
		{
			name:    "invalid role",
			role:    llm.Role("invalid"),
			failure: "role",
		},

		{
			name:    "user with tool calls",
			role:    llm.RoleUser,
			failure: "calls",
		},
		{
			name:    "tool result in empty history",
			role:    llm.RoleTool,
			failure: "tool",
		},
		{
			name:    "missing conversation",
			role:    llm.RoleUser,
			failure: "missing",
		},
		{
			name:    "canceled context",
			role:    llm.RoleUser,
			failure: "canceled",
			wantErr: context.Canceled,
		},
		{
			name:    "closed database",
			role:    llm.RoleUser,
			failure: "closed",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			var (
				ctx = t.Context()
				c   = newTestConversation(t, nil)
				msg = llm.Message{
					Role: tt.role,
					ContentParts: []llm.ContentPart{{
						Type: llm.ContentTypeText,
						Text: "manually added",
					}},
				}
			)

			switch tt.failure {
			case "calls":
				msg.ToolCalls = []llm.ToolCall{{ID: "call"}}
			case "busy", "stopped":
				// Register work without calling Infer's stub.
				_, cancel := context.WithCancel(ctx)
				defer cancel()
				c.cancel = cancel
				if tt.failure == "stopped" {
					// Stop must leave the call registered.
					c.Stop()
				}
			case "missing":
				c.id = "missing"
			case "canceled":
				canceled, cancel := context.WithCancel(ctx)
				cancel()
				ctx = canceled
			case "closed":
				if err := c.db.Close(); err != nil {
					t.Fatalf("Close: %v", err)
				}
			}

			saved, err := c.InsertMessage(ctx, msg)
			if saved != nil {
				t.Fatalf("returned row = %+v, want nil", saved)
			}
			if err == nil {
				t.Fatal("append succeeded, want an error")
			}
			if tt.wantErr != nil && !errors.Is(err, tt.wantErr) {
				t.Fatalf("error = %v, want %v", err, tt.wantErr)
			}
			if tt.failure == "closed" {
				return
			}

			rows, err := c.dbr.SelectMessagesConversationId(
				t.Context(), c.id,
			)
			if err != nil {
				t.Fatalf(
					"SelectMessagesConversationId: %v", err,
				)
			}
			if len(rows) != 0 {
				t.Fatalf("saved rows = %d, want 0", len(rows))
			}
			if tt.failure != "role" {
				return
			}

			// Retry after failure: the transaction must release the
			// writer without consuming a sequence number.
			msg.Role = llm.RoleUser
			saved, err = c.InsertMessage(ctx, msg)
			if err != nil {
				t.Fatalf("append after rollback: %v", err)
			}
			if saved.Sequence != 1 {
				t.Fatalf(
					"sequence = %d, want 1", saved.Sequence,
				)
			}
		})
	}
}

// TestDeleteMessage checks row removal, sequence compaction, and
// execution-row cascade.
func TestDeleteMessage(t *testing.T) {
	var (
		ctx = t.Context()
		c   = newTestConversation(t, nil)
	)

	// Seed: user, assistant with one tool call, the call's result,
	// user.
	user := llm.Message{
		Role: llm.RoleUser,
		ContentParts: []llm.ContentPart{{
			Type: llm.ContentTypeText,
			Text: "hello",
		}},
	}
	seed := []struct {
		msg  llm.Message
		role db.MessageRole
	}{
		{user, db.MessageRoleUser},
		{
			llm.Message{
				Role: llm.RoleAssistant,
				ToolCalls: []llm.ToolCall{{
					ID:   "call_1",
					Type: llm.ToolTypeFunction,
					Function: llm.Function{
						Name:      "search",
						Arguments: `{}`,
					},
				}},
			},
			db.MessageRoleAssistant,
		},
		{
			llm.Message{
				Role:       llm.RoleTool,
				Name:       "search",
				ToolCallID: "call_1",
				ContentParts: []llm.ContentPart{{
					Type: llm.ContentTypeText,
					Text: "results",
				}},
			},
			db.MessageRoleTool,
		},
		{user, db.MessageRoleUser},
	}
	var rows []db.Message
	for i, s := range seed {
		content, err := json.Marshal(s.msg)
		if err != nil {
			t.Fatalf("Marshal: %v", err)
		}
		row, err := c.dbw.InsertMessage(ctx, db.InsertMessageParams{
			Content:        content,
			ConversationId: c.id,
			Role:           s.role,
			Sequence:       i + 1,
		})
		if err != nil {
			t.Fatalf("InsertMessage: %v", err)
		}
		rows = append(rows, *row)
	}
	if _, err := c.dbw.InsertExecution(ctx, db.InsertExecutionParams{
		MessageId:  rows[1].Id,
		ToolCallId: "call_1",
		Status:     db.ExecutionStatusCompleted,
	}); err != nil {
		t.Fatalf("InsertExecution: %v", err)
	}

	// Deleting the assistant cascades its execution row.
	if err := c.DeleteMessage(
		ctx, rows[1].Id, rows[1].UpdatedAt,
	); err != nil {
		t.Fatalf("DeleteMessage: %v", err)
	}
	remaining, err := c.dbr.SelectMessagesConversationId(ctx, c.id)
	if err != nil {
		t.Fatalf("SelectMessagesConversationId: %v", err)
	}
	if len(remaining) != 3 {
		t.Fatalf("saved rows = %d, want 3", len(remaining))
	}
	for i, row := range remaining {
		if row.Sequence != i+1 {
			t.Fatalf(
				"row %d sequence = %d, want %d",
				i, row.Sequence, i+1,
			)
		}
	}
	if executions, err := c.dbr.SelectExecutionsMessageId(
		ctx, rows[1].Id,
	); err != nil {
		t.Fatalf("SelectExecutionsMessageId: %v", err)
	} else if len(executions) != 0 {
		t.Fatalf("execution rows survived: %d", len(executions))
	}

	// The tool result deletes too. The first delete's shift gave it a
	// new updated_at, so the revision is the re-read one.
	if err := c.DeleteMessage(
		ctx, rows[2].Id, remaining[1].UpdatedAt,
	); err != nil {
		t.Fatalf("DeleteMessage: %v", err)
	}
	remaining, err = c.dbr.SelectMessagesConversationId(ctx, c.id)
	if err != nil {
		t.Fatalf("SelectMessagesConversationId: %v", err)
	}
	if len(remaining) != 2 {
		t.Fatalf("saved rows = %d, want 2", len(remaining))
	}
	for i, row := range remaining {
		if row.Sequence != i+1 {
			t.Fatalf(
				"row %d sequence = %d, want %d",
				i, row.Sequence, i+1,
			)
		}
	}
}

// TestDeleteMessageErrors checks revision and busy rejection.
func TestDeleteMessageErrors(t *testing.T) {
	tests := []struct {
		name    string
		failure string
		wantErr error
	}{
		{
			name:    "inference running",
			failure: "busy",
			wantErr: ErrBusy,
		},
		{
			name:    "inference stopped but not finished",
			failure: "stopped",
			wantErr: ErrBusy,
		},
		{
			name:    "missing message",
			failure: "missing",
		},
		{
			name:    "missing conversation",
			failure: "conversation",
		},
		{
			name:    "stale revision",
			failure: "stale",
			wantErr: ErrRevisionConflict,
		},
		{
			name:    "canceled context",
			failure: "canceled",
			wantErr: context.Canceled,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			var (
				ctx = t.Context()
				c   = newTestConversation(t, nil)
			)
			conversationId := c.id
			row, err := c.InsertMessage(ctx, *userMessage("hello"))
			if err != nil {
				t.Fatalf("InsertMessage: %v", err)
			}

			messageId := row.Id
			revision := row.UpdatedAt
			switch tt.failure {
			case "busy", "stopped":
				// Register work without calling Infer's stub.
				_, cancel := context.WithCancel(ctx)
				defer cancel()
				c.cancel = cancel
				if tt.failure == "stopped" {
					// Stop must leave the call registered.
					c.Stop()
				}
			case "missing":
				messageId = "missing"
			case "conversation":
				c.id = "missing"
			case "stale":
				revision = row.UpdatedAt.Add(time.Hour)
			case "canceled":
				canceled, cancel := context.WithCancel(ctx)
				cancel()
				ctx = canceled
			}

			err = c.DeleteMessage(ctx, messageId, revision)
			if err == nil {
				t.Fatal("delete succeeded, want an error")
			}
			if tt.wantErr != nil && !errors.Is(err, tt.wantErr) {
				t.Fatalf("error = %v, want %v", err, tt.wantErr)
			}

			rows, err := c.dbr.SelectMessagesConversationId(
				t.Context(), conversationId,
			)
			if err != nil {
				t.Fatalf(
					"SelectMessagesConversationId: %v", err,
				)
			}
			if len(rows) != 1 {
				t.Fatalf("saved rows = %d, want 1", len(rows))
			}
		})
	}
}

// TestUpdateToolResult replaces a saved result's output and completes
// the call's execution, which an interrupted call leaves running.
func TestUpdateToolResult(t *testing.T) {
	var (
		ctx = t.Context()
		c   = newTestConversation(t, nil)
	)

	// Seed: assistant with one tool call, the call's result, user.
	seed := []struct {
		msg  llm.Message
		role db.MessageRole
	}{
		{
			llm.Message{
				Role: llm.RoleAssistant,
				ToolCalls: []llm.ToolCall{{
					ID:   "call_1",
					Type: llm.ToolTypeFunction,
					Function: llm.Function{
						Name:      "search",
						Arguments: `{}`,
					},
				}},
			},
			db.MessageRoleAssistant,
		},
		{
			llm.Message{
				Role:       llm.RoleTool,
				Name:       "search",
				ToolCallID: "call_1",
				ContentParts: []llm.ContentPart{{
					Type: llm.ContentTypeText,
					Text: "results",
				}},
			},
			db.MessageRoleTool,
		},
		{
			llm.Message{
				Role: llm.RoleUser,
				ContentParts: []llm.ContentPart{{
					Type: llm.ContentTypeText,
					Text: "hello",
				}},
			},
			db.MessageRoleUser,
		},
	}
	var rows []db.Message
	for i, s := range seed {
		content, err := json.Marshal(s.msg)
		if err != nil {
			t.Fatalf("Marshal: %v", err)
		}
		row, err := c.dbw.InsertMessage(ctx, db.InsertMessageParams{
			Content:        content,
			ConversationId: c.id,
			Role:           s.role,
			Sequence:       i + 1,
		})
		if err != nil {
			t.Fatalf("InsertMessage: %v", err)
		}
		rows = append(rows, *row)
	}
	if _, err := c.dbw.InsertExecution(ctx, db.InsertExecutionParams{
		MessageId:  rows[0].Id,
		ToolCallId: "call_1",
		Status:     db.ExecutionStatusRunning,
	}); err != nil {
		t.Fatalf("InsertExecution: %v", err)
	}
	original := append([]db.Message(nil), rows...)

	want := llm.Message{
		Role:       llm.RoleTool,
		Name:       "search",
		ToolCallID: "call_1",
		ContentParts: []llm.ContentPart{{
			Type: llm.ContentTypeText,
			Text: "edited",
		}},
	}
	updated, err := c.UpdateToolResult(
		ctx, rows[0].Id, "call_1", rows[1].UpdatedAt, "edited",
	)
	if err != nil {
		t.Fatalf("UpdateToolResult: %v", err)
	}

	// The replacement keeps the result row's identity and position.
	if updated.Id != rows[1].Id {
		t.Errorf("returned row %q, want %q", updated.Id, rows[1].Id)
	}
	if updated.Sequence != rows[1].Sequence {
		t.Errorf(
			"sequence = %d, want %d",
			updated.Sequence, rows[1].Sequence,
		)
	}
	if updated.Role != db.MessageRoleTool {
		t.Errorf("role = %q, want tool", updated.Role)
	}
	wantContent, err := json.Marshal(want)
	if err != nil {
		t.Fatalf("Marshal: %v", err)
	}
	if string(updated.Content) != string(wantContent) {
		t.Errorf("content = %q, want %s", updated.Content, wantContent)
	}

	// The saved row matches; the other rows are untouched.
	rows, err = c.dbr.SelectMessagesConversationId(ctx, c.id)
	if err != nil {
		t.Fatalf("SelectMessagesConversationId: %v", err)
	}
	if len(rows) != 3 {
		t.Fatalf("saved rows = %d, want 3", len(rows))
	}
	if !reflect.DeepEqual(rows[1], *updated) {
		t.Errorf("saved result row differs from returned row")
	}
	for i := range original {
		if i == 1 {
			continue
		}
		if !reflect.DeepEqual(rows[i], original[i]) {
			t.Errorf("row %d changed", i)
		}
	}

	// The running execution is completed.
	executions, err := c.dbr.SelectExecutionsMessageId(
		ctx, rows[0].Id,
	)
	if err != nil {
		t.Fatalf("SelectExecutionsMessageId: %v", err)
	}
	if len(executions) != 1 {
		t.Fatalf("execution rows = %d, want 1", len(executions))
	}
	if executions[0].Status != db.ExecutionStatusCompleted {
		t.Errorf(
			"status = %q, want %q",
			executions[0].Status, db.ExecutionStatusCompleted,
		)
	}
}

// TestUpdateToolResultErrors checks busy, missing-result, stale-revision,
// and wrong-role rejection.
func TestUpdateToolResultErrors(t *testing.T) {
	tests := []struct {
		name    string
		failure string
		wantErr error
	}{
		{
			name:    "inference running",
			failure: "busy",
			wantErr: ErrBusy,
		},
		{
			name:    "no saved result",
			failure: "missing",
		},
		{
			name:    "stale revision",
			failure: "stale",
			wantErr: ErrRevisionConflict,
		},
		{
			name:    "not an assistant",
			failure: "role",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			var (
				ctx = t.Context()
				c   = newTestConversation(t, nil)
			)

			// Seed: assistant with two calls, only the first
			// answered, then a user message.
			call1 := llm.ToolCall{
				ID:   "call_1",
				Type: llm.ToolTypeFunction,
				Function: llm.Function{
					Name:      "search",
					Arguments: `{}`,
				},
			}
			call2 := llm.ToolCall{
				ID:   "call_2",
				Type: llm.ToolTypeFunction,
				Function: llm.Function{
					Name:      "search",
					Arguments: `{}`,
				},
			}
			seed := []struct {
				msg  llm.Message
				role db.MessageRole
			}{
				{
					llm.Message{
						Role:      llm.RoleAssistant,
						ToolCalls: []llm.ToolCall{call1, call2},
					},
					db.MessageRoleAssistant,
				},
				{
					llm.Message{
						Role:       llm.RoleTool,
						Name:       "search",
						ToolCallID: "call_1",
						ContentParts: []llm.ContentPart{{
							Type: llm.ContentTypeText,
							Text: "results",
						}},
					},
					db.MessageRoleTool,
				},
				{
					*userMessage("hello"),
					db.MessageRoleUser,
				},
			}
			var rows []db.Message
			for i, s := range seed {
				content, err := json.Marshal(s.msg)
				if err != nil {
					t.Fatalf("Marshal: %v", err)
				}
				row, err := c.dbw.InsertMessage(
					ctx,
					db.InsertMessageParams{
						Content:        content,
						ConversationId: c.id,
						Role:           s.role,
						Sequence:       i + 1,
					},
				)
				if err != nil {
					t.Fatalf("InsertMessage: %v", err)
				}
				rows = append(rows, *row)
			}

			mId := rows[0].Id
			callId := "call_1"
			revision := rows[1].UpdatedAt
			switch tt.failure {
			case "busy":
				// Register work without calling Infer's stub.
				_, cancel := context.WithCancel(ctx)
				defer cancel()
				c.cancel = cancel
			case "missing":
				callId = "call_2"
			case "stale":
				revision = rows[1].UpdatedAt.Add(time.Hour)
			case "role":
				mId = rows[2].Id
			}

			_, err := c.UpdateToolResult(
				ctx, mId, callId, revision, "edited",
			)
			if err == nil {
				t.Fatal("update succeeded, want an error")
			}
			if tt.wantErr != nil && !errors.Is(err, tt.wantErr) {
				t.Fatalf("error = %v, want %v", err, tt.wantErr)
			}

			// The failed update changes nothing.
			saved, err := c.dbr.SelectMessagesConversationId(
				t.Context(), c.id,
			)
			if err != nil {
				t.Fatalf(
					"SelectMessagesConversationId: %v", err,
				)
			}
			if len(saved) != 3 {
				t.Fatalf("saved rows = %d, want 3", len(saved))
			}
			for i := range rows {
				if string(saved[i].Content) !=
					string(rows[i].Content) {
					t.Fatalf("row %d changed", i)
				}
			}
		})
	}
}

// TestStopIdle checks Stop does nothing when nothing is running.
func TestStopIdle(t *testing.T) {
	c := &Conversation{}
	c.Stop()
	c.Stop()

	c.mu.Lock()
	defer c.mu.Unlock()
	if c.cancel != nil {
		t.Fatal("idle Stop registered work")
	}
}

// TestStopCancellation checks Stop cancels work without marking it finished.
func TestStopCancellation(t *testing.T) {
	ctx, cancel := context.WithCancel(t.Context())
	defer cancel()
	c := &Conversation{cancel: cancel}

	c.Stop()
	c.Stop()

	if ctx.Err() != context.Canceled {
		t.Fatalf("context error = %v, want context.Canceled", ctx.Err())
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.cancel == nil {
		t.Fatal("Stop released the inference reservation")
	}
}

// TestConfigValidate checks Config validation rules.
func TestConfigValidate(t *testing.T) {
	tests := []struct {
		name    string
		config  Config
		wantErr bool
	}{
		{
			name:    "missing conversation id",
			config:  Config{},
			wantErr: true,
		},
		{
			name: "valid conversation id",
			config: Config{
				Id: "35afb18e-3ee1-4053-9519-1efba4d2932c",
			},
			wantErr: false,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			err := tt.config.Validate()
			if (err != nil) != tt.wantErr {
				t.Errorf(
					"Validate() error = %v, wantErr %v",
					err, tt.wantErr,
				)
			}
		})
	}
}

// TestSubscribeIdle checks the diff with no ongoing Inference:
// missing, changed, and deleted messages, and no Done.
func TestSubscribeIdle(t *testing.T) {
	var (
		ctx = t.Context()
		c   = newTestConversation(t, nil)
	)
	var saved []*db.Message
	for _, text := range []string{"one", "two", "three"} {
		row, err := c.InsertMessage(ctx, *userMessage(text))
		if err != nil {
			t.Fatalf("InsertMessage: %v", err)
		}
		saved = append(saved, row)
	}
	manifest := func(rows ...*db.Message) []MessageManifest {
		m := make([]MessageManifest, len(rows))
		for i, row := range rows {
			m[i] = MessageManifest{
				Id:        row.Id,
				UpdatedAt: row.UpdatedAt,
			}
		}
		return m
	}

	tests := []struct {
		name     string
		prepare  func() []MessageManifest
		want     []EventType
		wantText []string
	}{
		{
			name:    "empty manifest",
			prepare: func() []MessageManifest { return nil },
			want: []EventType{
				EventMessageSaved,
				EventMessageSaved,
				EventMessageSaved,
			},
			wantText: []string{"one", "two", "three"},
		},
		{
			name: "up to date",
			prepare: func() []MessageManifest {
				return manifest(saved...)
			},
			want: nil,
		},
		{
			name: "missing middle",
			prepare: func() []MessageManifest {
				return manifest(saved[0], saved[2])
			},
			want:     []EventType{EventMessageSaved},
			wantText: []string{"two"},
		},
		{
			name: "stale entry",
			prepare: func() []MessageManifest {
				m := manifest(saved...)
				m[0].UpdatedAt = time.Time{}
				return m
			},
			want:     []EventType{EventMessageSaved},
			wantText: []string{"one"},
		},
		{
			name: "deleted",
			prepare: func() []MessageManifest {
				// Sleep so the shift's updated_at is a later
				// millisecond than the inserts' updated_at.
				time.Sleep(10 * time.Millisecond)
				if err := c.DeleteMessage(
					ctx, saved[1].Id, saved[1].UpdatedAt,
				); err != nil {
					t.Fatalf("DeleteMessage: %v", err)
				}
				return manifest(saved...)
			},
			// Deleting "two" moves "three" into its slot. The moved
			// row gets a new updated_at, so the diff reports it as
			// saved before the deletion.
			want: []EventType{
				EventMessageSaved,
				EventMessageDeleted,
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			ch, err := c.Subscribe(ctx, tt.prepare())
			if err != nil {
				t.Fatalf("Subscribe: %v", err)
			}
			got := collectEvents(t, ch)
			if len(got) != len(tt.want) {
				t.Fatalf(
					"got %d events, want %d: %v",
					len(got), len(tt.want), got,
				)
			}
			for i, e := range got {
				if e.Type() != tt.want[i] {
					t.Errorf(
						"event %d type = %v, want %v",
						i, e.Type(), tt.want[i],
					)
				}
			}
			for i, wantText := range tt.wantText {
				s, ok := got[i].(MessageSaved)
				if !ok {
					t.Fatalf(
						"event %d = %v, want MessageSaved",
						i, got[i],
					)
				}
				var m llm.Message
				if err := json.Unmarshal(
					s.Message.Content, &m,
				); err != nil {
					t.Fatalf(
						"unmarshal message %q: %v",
						s.Message.Id, err,
					)
				}
				if m.Text() != wantText {
					t.Errorf(
						"event %d text = %q, want %q",
						i, m.Text(), wantText,
					)
				}
			}
		})
	}
}

// TestSubscribeInference checks a client joining an ongoing streaming
// Inference: the diff, the chunks, the saved assistant message, and
// Done with no error.
func TestSubscribeInference(t *testing.T) {
	var (
		ctx        = t.Context()
		firstChunk = make(chan struct{})
		release    = make(chan struct{})
	)
	client := newTestModelClient(
		t,
		func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Content-Type", "text/event-stream")
			flusher := w.(http.Flusher)
			fmt.Fprint(
				w,
				"data: {\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\"}}]}\n\n",
			)
			fmt.Fprint(
				w,
				"data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\"}}]}\n\n",
			)
			flusher.Flush()
			close(firstChunk)
			// End the handler if the test aborts, so server cleanup
			// cannot hang on this request.
			select {
			case <-release:
			case <-r.Context().Done():
			}
			fmt.Fprint(
				w,
				"data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\" world\"}}]}\n\n",
			)
			fmt.Fprint(
				w,
				"data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n",
			)
			fmt.Fprint(w, "data: [DONE]\n\n")
		})

	var (
		c    = newTestConversation(t, nil)
		errs = make(chan error, 1)
	)
	go func() {
		errs <- infer(ctx, c, client, InferenceRequest{
			Message: userMessage("hi"),
			Model:   "test-model",
			Stream:  true,
		})
	}()
	select {
	case <-firstChunk:
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for the first chunk")
	}

	// Join while the Inference is mid-stream.
	ch, err := c.Subscribe(ctx, nil)
	if err != nil {
		t.Fatalf("Subscribe: %v", err)
	}
	close(release)
	got := collectEvents(t, ch)

	// The diff delivers the user message first.
	saved, ok := got[0].(MessageSaved)
	if !ok || saved.Message.Role != db.MessageRoleUser {
		t.Fatalf("first event = %v, want the user MessageSaved", got[0])
	}
	var (
		chunkTexts []string
		chunkIDs   []string
		asst       *MessageSaved
	)
	for _, e := range got[1:] {
		switch e := e.(type) {
		case AssistantUpdate:
			if asst != nil {
				t.Fatalf("AssistantUpdate after the assistant MessageSaved: %v", e)
			}
			chunkIDs = append(chunkIDs, e.MessageId)
			chunkTexts = append(
				chunkTexts, e.Chunk.Choices[0].Delta.Content,
			)
		case MessageSaved:
			if e.Message.Role != db.MessageRoleAssistant {
				t.Fatalf("unexpected saved role %q", e.Message.Role)
			}
			asst = &e
		case Done:
			if e.Error != "" {
				t.Errorf("Done.Error = %q, want empty", e.Error)
			}
		default:
			t.Fatalf("unexpected event %v", e)
		}
	}
	if asst == nil {
		t.Fatalf("no assistant MessageSaved in %v", got)
	}
	var asstMsg llm.Message
	if err := json.Unmarshal(
		asst.Message.Content, &asstMsg,
	); err != nil {
		t.Fatalf("unmarshal assistant message: %v", err)
	}
	if asstMsg.Text() != "hello world" {
		t.Errorf("assistant text = %q, want %q",
			asstMsg.Text(), "hello world")
	}
	if len(chunkTexts) == 0 {
		t.Fatal("no AssistantUpdate events")
	}
	for i, id := range chunkIDs {
		if id != asst.Message.Id {
			t.Errorf("chunk %d message id = %q, want the assistant id %q",
				i, id, asst.Message.Id)
		}
	}
	if last := got[len(got)-1]; last.Type() != EventDone {
		t.Fatalf("last event type = %v, want done", last.Type())
	}
	if err := <-errs; err != nil {
		t.Fatalf("Inference: %v", err)
	}
}

// TestSubscribeInferenceError checks that a dropped stream ends the
// subscriber with Done carrying the error, and saves no assistant
// message.
func TestSubscribeInferenceError(t *testing.T) {
	var (
		ctx        = t.Context()
		firstChunk = make(chan struct{})
		release    = make(chan struct{})
	)
	client := newTestModelClient(
		t,
		func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Content-Type", "text/event-stream")
			fmt.Fprint(
				w,
				"data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"}}]}\n\n",
			)
			w.(http.Flusher).Flush()
			close(firstChunk)
			// End the handler if the test aborts, so server cleanup
			// cannot hang on this request.
			select {
			case <-release:
			case <-r.Context().Done():
			}
			// Returning without [DONE] drops the connection.
		})

	var (
		c    = newTestConversation(t, nil)
		errs = make(chan error, 1)
	)
	go func() {
		errs <- infer(ctx, c, client, InferenceRequest{
			Message: userMessage("hi"),
			Model:   "test-model",
			Stream:  true,
		})
	}()
	select {
	case <-firstChunk:
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for the first chunk")
	}

	ch, err := c.Subscribe(ctx, nil)
	if err != nil {
		t.Fatalf("Subscribe: %v", err)
	}
	close(release)
	got := collectEvents(t, ch)

	done, ok := got[len(got)-1].(Done)
	if !ok {
		t.Fatalf("last event = %v, want Done", got[len(got)-1])
	}
	if done.Error == "" {
		t.Fatal("Done.Error = empty, want the stream error")
	}
	rows, err := db.New(c.db.Reader).
		SelectMessagesConversationId(ctx, c.id)
	if err != nil {
		t.Fatalf("SelectMessagesConversationId: %v", err)
	}
	for _, row := range rows {
		if row.Role != db.MessageRoleUser {
			t.Fatalf("unexpected saved role %q", row.Role)
		}
	}
	if err := <-errs; err == nil {
		t.Fatal("Inference error = nil, want the stream error")
	}
}

// TestSubscribeToolResult checks that tool results reach subscribers.
func TestSubscribeToolResult(t *testing.T) {
	var (
		ctx        = t.Context()
		requested  = make(chan struct{})
		subscribed = make(chan struct{})
	)
	registry, err := tool.NewRegistry([]config.ToolConfig{{
		Name:        "echo",
		Description: "echoes its arguments",
		Command:     "/bin/echo",
		Arguments:   []string{"hi"},
	}})
	if err != nil {
		t.Fatalf("NewRegistry: %v", err)
	}

	var request int
	client := newTestModelClient(
		t,
		func(w http.ResponseWriter, r *http.Request) {
			request++
			w.Header().Set("Content-Type", "application/json")
			switch request {
			case 1:
				close(requested)
				// End the handler if the test aborts, so server
				// cleanup cannot hang on this request.
				select {
				case <-subscribed:
				case <-r.Context().Done():
				}
				fmt.Fprint(
					w,
					`{"choices":[{"message":{"role":"assistant","tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"echo","arguments":"{}"}}]}}]}`,
				)
			case 2:
				fmt.Fprint(
					w,
					`{"choices":[{"message":{"role":"assistant","content":"done"}}]}`,
				)
			}
		})

	var (
		c    = newTestConversation(t, registry)
		errs = make(chan error, 1)
	)
	go func() {
		errs <- infer(ctx, c, client, InferenceRequest{
			Message: userMessage("hi"),
			Model:   "test-model",
		})
	}()
	select {
	case <-requested:
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for the first request")
	}

	// Join while the Inference awaits the tool-call answer.
	ch, err := c.Subscribe(ctx, nil)
	if err != nil {
		t.Fatalf("Subscribe: %v", err)
	}
	close(subscribed)
	got := collectEvents(t, ch)

	var (
		aId     string
		sawTool bool
	)
	for i, e := range got {
		switch e := e.(type) {
		case MessageSaved:
			switch e.Message.Role {
			case db.MessageRoleAssistant:
				if aId == "" {
					aId = e.Message.Id
				}
			case db.MessageRoleTool:
				sawTool = true
				var m llm.Message
				if err := json.Unmarshal(
					e.Message.Content, &m,
				); err != nil {
					t.Fatalf(
						"unmarshal message %q: %v", e.Message.Id, err,
					)
				}
				if m.Text() == "" {
					t.Errorf("event %d: empty tool result", i)
				}
			default:
				if i != 0 {
					t.Errorf("event %d role = %q, want the user first", i, e.Message.Role)
				}
			}
		case Done:
			if e.Error != "" {
				t.Errorf("Done.Error = %q, want empty", e.Error)
			}
		default:
			t.Errorf("event %d = %v, want MessageSaved or Done", i, e)
		}
	}
	if aId == "" {
		t.Fatalf("no assistant MessageSaved in %v", got)
	}
	if !sawTool {
		t.Fatalf("no tool MessageSaved in %v", got)
	}
	if err := <-errs; err != nil {
		t.Fatalf("Inference: %v", err)
	}
}

// TestSubscribeCancel checks that canceling the subscriber's context
// closes the channel without Done, while Inference keeps running.
func TestSubscribeCancel(t *testing.T) {
	var (
		ctx        = t.Context()
		firstChunk = make(chan struct{})
	)
	client := newTestModelClient(
		t,
		func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Content-Type", "text/event-stream")
			fmt.Fprint(
				w,
				"data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\"}}]}\n\n",
			)
			w.(http.Flusher).Flush()
			close(firstChunk)
			// Hold the stream open until the request is canceled.
			<-r.Context().Done()
		},
	)

	var (
		c    = newTestConversation(t, nil)
		errs = make(chan error, 1)
	)
	go func() {
		errs <- infer(ctx, c, client, InferenceRequest{
			Message: userMessage("hi"),
			Model:   "test-model",
			Stream:  true,
		})
	}()
	select {
	case <-firstChunk:
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for the first chunk")
	}

	subCtx, cancel := context.WithCancel(ctx)
	ch, err := c.Subscribe(subCtx, nil)
	if err != nil {
		t.Fatalf("Subscribe: %v", err)
	}
	cancel()

	// Delivery stops and ch closes without Done.
	for _, e := range collectEvents(t, ch) {
		if e.Type() == EventDone {
			t.Fatalf("received Done after cancel: %v", e)
		}
	}

	// Inference kept running; Stop ends it.
	c.Stop()
	if err := <-errs; err == nil {
		t.Fatal("Inference error = nil, want the cancellation")
	}
}

// TestSubscribeStartupError checks that a failure resolving the saved
// history's tool calls ends subscribers with Done carrying the error,
// not a success.
func TestSubscribeStartupError(t *testing.T) {
	var (
		ctx     = t.Context()
		dir     = t.TempDir()
		started = filepath.Join(dir, "started")
		release = filepath.Join(dir, "release")
		script  = filepath.Join(dir, "slow.sh")
	)
	// The tool holds the Inference open until the test releases it,
	// so the subscriber can join before the failure is reported.
	if err := os.WriteFile(script, []byte(
		"#!/bin/sh\ntouch "+started+"\nwhile [ ! -f "+release+
			" ]; do sleep 0.1; done\n",
	), 0o755); err != nil {
		t.Fatalf("WriteFile: %v", err)
	}
	registry, err := tool.NewRegistry([]config.ToolConfig{{
		Name:        "slow",
		Description: "blocks until released",
		Command:     script,
	}})
	if err != nil {
		t.Fatalf("NewRegistry: %v", err)
	}

	var (
		c   = newTestConversation(t, registry)
		msg = llm.Message{
			Role: llm.RoleAssistant,
			ToolCalls: []llm.ToolCall{
				{
					ID:   "call_slow",
					Type: llm.ToolTypeFunction,
					Function: llm.Function{
						Name:      "slow",
						Arguments: "{}",
					},
				},
				{
					ID:   "call_boom",
					Type: llm.ToolTypeFunction,
					Function: llm.Function{
						Name:      "boom",
						Arguments: "{}",
					},
				},
			},
		}
	)
	row, err := c.InsertMessage(ctx, msg)
	if err != nil {
		t.Fatalf("InsertMessage: %v", err)
	}
	// call_boom "finished" earlier without a saved result, so it
	// must be refused; call_slow keeps the resolution open.
	if _, err := c.dbw.InsertExecution(
		ctx,
		db.InsertExecutionParams{
			MessageId:  row.Id,
			ToolCallId: "call_boom",
			Status:     db.ExecutionStatusCompleted,
		},
	); err != nil {
		t.Fatalf("InsertExecution: %v", err)
	}

	// The client is never used: the failure precedes any request.
	client := newTestModelClient(t, func(
		http.ResponseWriter, *http.Request,
	) {
	})
	errs := make(chan error, 1)
	go func() {
		errs <- infer(ctx, c, client, InferenceRequest{
			Model: "test-model",
		})
	}()

	deadline := time.Now().Add(5 * time.Second)
	for {
		if _, err := os.Stat(started); err == nil {
			break
		}
		if time.Now().After(deadline) {
			t.Fatal("timed out waiting for the slow tool to start")
		}
		time.Sleep(10 * time.Millisecond)
	}

	ch, err := c.Subscribe(ctx, nil)
	if err != nil {
		t.Fatalf("Subscribe: %v", err)
	}
	if err := os.WriteFile(release, nil, 0o644); err != nil {
		t.Fatalf("WriteFile: %v", err)
	}
	got := collectEvents(t, ch)

	first, ok := got[0].(MessageSaved)
	if !ok || first.Message.Role != db.MessageRoleAssistant {
		t.Fatalf("first event = %v, want the assistant MessageSaved", got[0])
	}
	var sawTool bool
	for _, e := range got {
		switch e := e.(type) {
		case MessageSaved:
			if e.Message.Role != db.MessageRoleTool {
				continue
			}
			sawTool = true
		case Done:
			if e.Error == "" {
				t.Error("Done.Error = empty, want the startup failure")
			}
		}
	}
	if !sawTool {
		t.Errorf("no tool MessageSaved in %v", got)
	}
	if last := got[len(got)-1]; last.Type() != EventDone {
		t.Fatalf("last event type = %v, want done", last.Type())
	}
	if err := <-errs; err == nil {
		t.Fatal("Inference error = nil, want the startup failure")
	}
}

// TestUpdateMessage checks replacing a saved message's content while
// preserving its identity, sequence, and inference metadata.
func TestUpdateMessage(t *testing.T) {
	model := "test-model"
	var (
		ctx = t.Context()
		c   = newTestConversation(t, nil)
		msg = llm.Message{
			Role: llm.RoleAssistant,
			ContentParts: []llm.ContentPart{{
				Type: llm.ContentTypeText,
				Text: "initial",
			}},
		}
		meta = llm.Metadata{Model: &model}
	)
	saved, err := c.insertMessage(ctx, nil, msg, &meta, nil)
	if err != nil {
		t.Fatalf("insertMessage: %v", err)
	}

	replacement := llm.Message{
		Role: llm.RoleAssistant,
		ContentParts: []llm.ContentPart{{
			Type: llm.ContentTypeText,
			Text: "updated",
		}},
	}
	updated, err := c.UpdateMessage(
		ctx, saved.Id, saved.UpdatedAt, replacement,
	)
	if err != nil {
		t.Fatalf("UpdateMessage: %v", err)
	}
	if updated.Id != saved.Id {
		t.Errorf("Id = %q, want %q", updated.Id, saved.Id)
	}
	if updated.Sequence != saved.Sequence {
		t.Errorf(
			"Sequence = %d, want %d",
			updated.Sequence, saved.Sequence,
		)
	}
	if string(updated.Metadata) != string(saved.Metadata) {
		t.Errorf(
			"Metadata = %s, want %s",
			updated.Metadata, saved.Metadata,
		)
	}
	var decoded llm.Message
	if err := json.Unmarshal(updated.Content, &decoded); err != nil {
		t.Fatalf("Unmarshal: %v", err)
	}
	if decoded.Text() != "updated" {
		t.Errorf("Text() = %q, want %q", decoded.Text(), "updated")
	}
}

// TestUpdateMessageErrors checks busy, revision, role-change, and
// validation rejection in UpdateMessage.
func TestUpdateMessageErrors(t *testing.T) {
	tests := []struct {
		name    string
		failure string
		wantErr error
	}{
		{
			name:    "inference running",
			failure: "busy",
			wantErr: ErrBusy,
		},
		{
			name:    "stale revision",
			failure: "stale",
			wantErr: ErrRevisionConflict,
		},
		{
			name:    "role change",
			failure: "role",
		},
		{
			name:    "invalid message",
			failure: "invalid",
		},
		{
			name:    "missing message",
			failure: "missing",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			var (
				ctx = t.Context()
				c   = newTestConversation(t, nil)
			)
			saved, err := c.InsertMessage(
				ctx, *userMessage("original"),
			)
			if err != nil {
				t.Fatalf("InsertMessage: %v", err)
			}

			mId := saved.Id
			revision := saved.UpdatedAt
			replacement := *userMessage("replacement")

			switch tt.failure {
			case "busy":
				_, cancel := context.WithCancel(ctx)
				defer cancel()
				c.cancel = cancel
			case "stale":
				revision = saved.UpdatedAt.Add(time.Hour)
			case "role":
				replacement.Role = llm.RoleAssistant
			case "invalid":
				replacement.ToolCallID = "call_1"
			case "missing":
				mId = "missing"
			}

			_, err = c.UpdateMessage(
				ctx, mId, revision, replacement,
			)
			if err == nil {
				t.Fatal("UpdateMessage succeeded, want error")
			}
			if tt.wantErr != nil && !errors.Is(err, tt.wantErr) {
				t.Fatalf("error = %v, want %v", err, tt.wantErr)
			}
		})
	}
}

// TestToolCallLifecycle checks AddToolCall, UpdateToolCall, and
// DeleteToolCall on a saved assistant message, including cascading
// removal of the call's tool result and execution row.
func TestToolCallLifecycle(t *testing.T) {
	var (
		ctx = t.Context()
		c   = newTestConversation(t, nil)
	)
	asst, err := c.InsertMessage(ctx, llm.Message{
		Role: llm.RoleAssistant,
		ContentParts: []llm.ContentPart{{
			Type: llm.ContentTypeText,
			Text: "checking",
		}},
	})
	if err != nil {
		t.Fatalf("InsertMessage: %v", err)
	}

	// Add a tool call.
	call := llm.ToolCall{
		ID:   "call_1",
		Type: llm.ToolTypeFunction,
		Function: llm.Function{
			Name:      "search",
			Arguments: `{"q":"initial"}`,
		},
	}
	asst, err = c.AddToolCall(ctx, asst.Id, asst.UpdatedAt, call)
	if err != nil {
		t.Fatalf("AddToolCall: %v", err)
	}

	// Update the tool call's arguments.
	call.Function.Arguments = `{"q":"updated"}`
	asst, err = c.UpdateToolCall(
		ctx, asst.Id, "call_1", asst.UpdatedAt, call,
	)
	if err != nil {
		t.Fatalf("UpdateToolCall: %v", err)
	}
	var decoded llm.Message
	if err := json.Unmarshal(asst.Content, &decoded); err != nil {
		t.Fatalf("Unmarshal: %v", err)
	}
	if len(decoded.ToolCalls) != 1 ||
		decoded.ToolCalls[0].Function.Arguments != `{"q":"updated"}` {
		t.Fatalf("ToolCalls = %+v", decoded.ToolCalls)
	}

	// Add a tool result and a trailing user message so DeleteToolCall
	// must remove the result and compact sequences.
	if _, err := c.AddToolResult(
		ctx, asst.Id, "call_1", asst.UpdatedAt, "search output",
	); err != nil {
		t.Fatalf("AddToolResult: %v", err)
	}
	if _, err := c.InsertMessage(
		ctx, *userMessage("follow-up"),
	); err != nil {
		t.Fatalf("InsertMessage: %v", err)
	}
	asstRow, err := c.dbr.SelectConversationMessageId(
		ctx, c.id, asst.Id,
	)
	if err != nil {
		t.Fatalf("SelectConversationMessageId: %v", err)
	}

	// Delete the tool call: its result and execution record must be
	// removed, and the trailing user message shifts to sequence 2.
	if err := c.DeleteToolCall(
		ctx, asst.Id, "call_1", asstRow.UpdatedAt,
	); err != nil {
		t.Fatalf("DeleteToolCall: %v", err)
	}
	rows, err := c.dbr.SelectMessagesConversationId(ctx, c.id)
	if err != nil {
		t.Fatalf("SelectMessagesConversationId: %v", err)
	}
	if len(rows) != 2 {
		t.Fatalf("saved rows = %d, want 2", len(rows))
	}
	if rows[1].Role != db.MessageRoleUser || rows[1].Sequence != 2 {
		t.Errorf(
			"row 1 = (role %q, seq %d), want (user, 2)",
			rows[1].Role, rows[1].Sequence,
		)
	}
	executions, err := c.dbr.SelectExecutionsMessageId(ctx, asst.Id)
	if err != nil {
		t.Fatalf("SelectExecutionsMessageId: %v", err)
	}
	if len(executions) != 0 {
		t.Errorf("executions = %d, want 0", len(executions))
	}
}

// TestToolCallErrors checks error cases for AddToolCall, UpdateToolCall,
// and DeleteToolCall.
func TestToolCallErrors(t *testing.T) {
	validCall := llm.ToolCall{
		ID:   "call_1",
		Type: llm.ToolTypeFunction,
		Function: llm.Function{
			Name:      "search",
			Arguments: `{}`,
		},
	}

	type runFunc func(
		ctx context.Context,
		c *Conversation,
		asst, user *db.Message,
	) error

	tests := []struct {
		name    string
		run     runFunc
		wantErr error
	}{
		{
			name: "AddToolCall busy",
			run: func(
				ctx context.Context,
				c *Conversation,
				asst, _ *db.Message,
			) error {
				_, cancel := context.WithCancel(ctx)
				defer cancel()
				c.cancel = cancel
				_, err := c.AddToolCall(
					ctx, asst.Id, asst.UpdatedAt, validCall,
				)
				return err
			},
			wantErr: ErrBusy,
		},
		{
			name: "AddToolCall stale revision",
			run: func(
				ctx context.Context,
				c *Conversation,
				asst, _ *db.Message,
			) error {
				_, err := c.AddToolCall(
					ctx,
					asst.Id,
					asst.UpdatedAt.Add(time.Hour),
					validCall,
				)
				return err
			},
			wantErr: ErrRevisionConflict,
		},
		{
			name: "AddToolCall non-assistant role",
			run: func(
				ctx context.Context,
				c *Conversation,
				_, user *db.Message,
			) error {
				_, err := c.AddToolCall(
					ctx, user.Id, user.UpdatedAt, validCall,
				)
				return err
			},
		},
		{
			name: "AddToolCall duplicate call id",
			run: func(
				ctx context.Context,
				c *Conversation,
				asst, _ *db.Message,
			) error {
				_, err := c.AddToolCall(
					ctx, asst.Id, asst.UpdatedAt, validCall,
				)
				return err
			},
		},
		{
			name: "UpdateToolCall stale revision",
			run: func(
				ctx context.Context,
				c *Conversation,
				asst, _ *db.Message,
			) error {
				_, err := c.UpdateToolCall(
					ctx,
					asst.Id,
					"call_1",
					asst.UpdatedAt.Add(time.Hour),
					validCall,
				)
				return err
			},
			wantErr: ErrRevisionConflict,
		},
		{
			name: "DeleteToolCall missing call",
			run: func(
				ctx context.Context,
				c *Conversation,
				asst, _ *db.Message,
			) error {
				return c.DeleteToolCall(
					ctx, asst.Id, "missing", asst.UpdatedAt,
				)
			},
		},
		{
			name: "DeleteToolCall stale revision",
			run: func(
				ctx context.Context,
				c *Conversation,
				asst, _ *db.Message,
			) error {
				return c.DeleteToolCall(
					ctx,
					asst.Id,
					"call_1",
					asst.UpdatedAt.Add(time.Hour),
				)
			},
			wantErr: ErrRevisionConflict,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			var (
				ctx = t.Context()
				c   = newTestConversation(t, nil)
			)
			asst, err := c.InsertMessage(ctx, llm.Message{
				Role:      llm.RoleAssistant,
				ToolCalls: []llm.ToolCall{validCall},
			})
			if err != nil {
				t.Fatalf("InsertMessage assistant: %v", err)
			}
			user, err := c.InsertMessage(ctx, *userMessage("hi"))
			if err != nil {
				t.Fatalf("InsertMessage user: %v", err)
			}

			err = tt.run(ctx, c, asst, user)
			if err == nil {
				t.Fatal("operation succeeded, want error")
			}
			if tt.wantErr != nil && !errors.Is(err, tt.wantErr) {
				t.Fatalf("error = %v, want %v", err, tt.wantErr)
			}
		})
	}
}

// TestAddToolResult checks inserting tool results right after an
// assistant message, shifting later messages down, and rejecting
// duplicate results or stale revisions.
func TestAddToolResult(t *testing.T) {
	var (
		ctx = t.Context()
		c   = newTestConversation(t, nil)
	)
	asst, err := c.InsertMessage(ctx, llm.Message{
		Role: llm.RoleAssistant,
		ToolCalls: []llm.ToolCall{{
			ID:   "call_1",
			Type: llm.ToolTypeFunction,
			Function: llm.Function{
				Name:      "search",
				Arguments: `{}`,
			},
		}},
	})
	if err != nil {
		t.Fatalf("InsertMessage assistant: %v", err)
	}
	user, err := c.InsertMessage(ctx, *userMessage("later"))
	if err != nil {
		t.Fatalf("InsertMessage user: %v", err)
	}

	res, err := c.AddToolResult(
		ctx, asst.Id, "call_1", asst.UpdatedAt, "found it",
	)
	if err != nil {
		t.Fatalf("AddToolResult: %v", err)
	}
	if res.Sequence != 2 {
		t.Errorf("result sequence = %d, want 2", res.Sequence)
	}

	// A second result for the same call is rejected.
	if _, err := c.AddToolResult(
		ctx, asst.Id, "call_1", asst.UpdatedAt, "duplicate",
	); err == nil {
		t.Fatal("duplicate AddToolResult succeeded, want error")
	}

	// Stale revision is rejected.
	if _, err := c.AddToolResult(
		ctx,
		asst.Id,
		"call_1",
		asst.UpdatedAt.Add(time.Hour),
		"stale",
	); !errors.Is(err, ErrRevisionConflict) {
		t.Fatalf("error = %v, want ErrRevisionConflict", err)
	}

	rows, err := c.dbr.SelectMessagesConversationId(ctx, c.id)
	if err != nil {
		t.Fatalf("SelectMessagesConversationId: %v", err)
	}
	if len(rows) != 3 {
		t.Fatalf("saved rows = %d, want 3", len(rows))
	}
	if rows[2].Id != user.Id || rows[2].Sequence != 3 {
		t.Errorf(
			"trailing user = (id %q, seq %d), want (%q, 3)",
			rows[2].Id, rows[2].Sequence, user.Id,
		)
	}
}

// TestDeleteConversation checks that Delete cancels active inference and
// removes the conversation and its messages.
func TestDeleteConversation(t *testing.T) {
	var (
		ctx    = t.Context()
		c      = newTestConversation(t, nil)
		subCtx context.Context
		cancel context.CancelFunc
	)
	if _, err := c.InsertMessage(ctx, *userMessage("hello")); err != nil {
		t.Fatalf("InsertMessage: %v", err)
	}
	subCtx, cancel = context.WithCancel(ctx)
	defer cancel()
	c.cancel = cancel

	if err := c.Delete(ctx); err != nil {
		t.Fatalf("Delete: %v", err)
	}
	if subCtx.Err() != context.Canceled {
		t.Errorf(
			"active inference context = %v, want Canceled",
			subCtx.Err(),
		)
	}
	rows, err := c.dbr.SelectMessagesConversationId(ctx, c.id)
	if err != nil {
		t.Fatalf("SelectMessagesConversationId: %v", err)
	}
	if len(rows) != 0 {
		t.Errorf("remaining messages = %d, want 0", len(rows))
	}
}

// TestInferenceApprovalAndMaxIterations checks that tools requiring
// approval pause with ErrApprovalRequired until approved, and that
// MaxIterations caps the model loop.
func TestInferenceApprovalAndMaxIterations(t *testing.T) {
	var (
		ctx = t.Context()
	)
	registry, err := tool.NewRegistry([]config.ToolConfig{{
		Name:             "guarded",
		Description:      "requires approval",
		Command:          "/bin/echo",
		Arguments:        []string{"approved-output"},
		RequiresApproval: true,
	}})
	if err != nil {
		t.Fatalf("NewRegistry: %v", err)
	}
	c := newTestConversation(t, registry)

	var request int
	client := newTestModelClient(
		t,
		func(w http.ResponseWriter, r *http.Request) {
			request++
			w.Header().Set("Content-Type", "application/json")
			if request == 1 {
				fmt.Fprint(
					w,
					`{"choices":[{"message":{"role":"assistant","tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"guarded","arguments":"{}"}}]}}]}`,
				)
				return
			}
			// Second request returns an unknown tool (which does
			// not require approval) so the 1-iteration loop hits
			// ErrMaxIterations.
			fmt.Fprint(
				w,
				`{"choices":[{"message":{"role":"assistant","tool_calls":[{"index":0,"id":"call_2","type":"function","function":{"name":"other","arguments":"{}"}}]}}]}`,
			)
		},
	)

	// First call: saves the assistant message and pending execution,
	// returning ErrApprovalRequired.
	err = infer(ctx, c, client, InferenceRequest{
		Message: userMessage("run guarded"),
		Model:   "test-model",
	})
	if !errors.Is(err, ErrApprovalRequired) {
		t.Fatalf("Inference error = %v, want ErrApprovalRequired", err)
	}

	// Approve the execution in the database, then run with
	// MaxIterations = 1 so the model's second tool call hits
	// ErrMaxIterations.
	rows, err := c.dbr.SelectMessagesConversationId(ctx, c.id)
	if err != nil || len(rows) != 2 {
		t.Fatalf("rows = %d (err %v), want 2", len(rows), err)
	}
	executions, err := c.dbr.SelectExecutionsMessageId(
		ctx, rows[1].Id,
	)
	if err != nil || len(executions) != 1 {
		t.Fatalf(
			"executions = %d (err %v), want 1",
			len(executions), err,
		)
	}
	if err := c.dbw.UpdateExecutionStatus(
		ctx,
		db.UpdateExecutionStatusParams{
			Id:     executions[0].Id,
			Status: db.ExecutionStatusApproved,
		},
	); err != nil {
		t.Fatalf("UpdateExecutionStatus: %v", err)
	}

	err = infer(ctx, c, client, InferenceRequest{
		Model:         "test-model",
		MaxIterations: 1,
	})
	if !errors.Is(err, ErrMaxIterations) {
		t.Fatalf("Inference error = %v, want ErrMaxIterations", err)
	}
}

// TestReserve checks that a reservation makes the conversation busy
// before Run starts, and that Run releases it.
func TestReserve(t *testing.T) {
	var (
		ctx    = t.Context()
		c      = newTestConversation(t, nil)
		client = newTestModelClient(
			t,
			func(w http.ResponseWriter, r *http.Request) {
				http.Error(w, "unavailable", http.StatusServiceUnavailable)
			})
		req = InferenceRequest{Model: "test-model"}
	)

	if _, err := c.Reserve(ctx, nil, req); err == nil {
		t.Fatal("Reserve(nil client) error = nil")
	}
	if _, err := c.Reserve(
		ctx, client, InferenceRequest{MaxIterations: -1},
	); err == nil {
		t.Fatal("Reserve(invalid request) error = nil")
	}

	inf, err := c.Reserve(ctx, client, req)
	if err != nil {
		t.Fatalf("Reserve: %v", err)
	}
	if _, err := c.Reserve(ctx, client, req); !errors.Is(err, ErrBusy) {
		t.Fatalf("second Reserve error = %v, want ErrBusy", err)
	}
	if _, err := c.InsertMessage(
		ctx, *userMessage("edit"),
	); !errors.Is(err, ErrBusy) {
		t.Fatalf("InsertMessage error = %v, want ErrBusy", err)
	}

	if err := inf.Run(); err == nil {
		t.Fatal("Run error = nil, want the model error")
	}
	if err := inf.Run(); err == nil {
		t.Fatal("second Run error = nil")
	}

	next, err := c.Reserve(ctx, client, req)
	if err != nil {
		t.Fatalf("Reserve after Run: %v", err)
	}
	next.Run()
}

// TestInferenceEventsFastFailure checks that a subscriber that holds
// the Inference receives Done even when Run fails before the subscriber
// reads its first event.
func TestInferenceEventsFastFailure(t *testing.T) {
	var (
		ctx    = t.Context()
		c      = newTestConversation(t, nil)
		client = newTestModelClient(
			t,
			func(w http.ResponseWriter, r *http.Request) {
				http.Error(w, "unavailable", http.StatusServiceUnavailable)
			})
	)
	inf, err := c.Reserve(ctx, client, InferenceRequest{Model: "test-model"})
	if err != nil {
		t.Fatalf("Reserve: %v", err)
	}
	runErr := inf.Run()
	if runErr == nil {
		t.Fatal("Run error = nil, want the model error")
	}

	got := collectEvents(t, inf.Events(ctx))
	if len(got) != 1 {
		t.Fatalf("events = %v, want only Done", got)
	}
	done, ok := got[0].(Done)
	if !ok {
		t.Fatalf("event = %v, want Done", got[0])
	}
	if done.Error != runErr.Error() {
		t.Fatalf("Done.Error = %q, want %q", done.Error, runErr.Error())
	}
}

// TestStopReserved checks that Stop cancels a reserved Inference
// before Run starts.
func TestStopReserved(t *testing.T) {
	var (
		ctx    = t.Context()
		c      = newTestConversation(t, nil)
		called = make(chan struct{}, 1)
		client = newTestModelClient(
			t,
			func(w http.ResponseWriter, r *http.Request) {
				called <- struct{}{}
			})
	)
	inf, err := c.Reserve(ctx, client, InferenceRequest{Model: "test-model"})
	if err != nil {
		t.Fatalf("Reserve: %v", err)
	}
	c.Stop()
	if err := inf.Run(); !errors.Is(err, context.Canceled) {
		t.Fatalf("Run error = %v, want context.Canceled", err)
	}
	select {
	case <-called:
		t.Fatal("model was called after Stop")
	default:
	}
}