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, }}, } } // 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 <- c.Inference(ctx, 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.Err != nil { t.Errorf("Done.Err = %v, want nil", e.Err) } 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 <- c.Inference(ctx, 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.Err == nil { t.Fatal("Done.Err = nil, 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 <- c.Inference(ctx, 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.Err != nil { t.Errorf("Done.Err = %v, want nil", e.Err) } 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 <- c.Inference(ctx, 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 <- c.Inference(ctx, 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.Err == nil { t.Error("Done.Err = nil, 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 = c.Inference(ctx, 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 = c.Inference(ctx, client, InferenceRequest{ Model: "test-model", MaxIterations: 1, }) if !errors.Is(err, ErrMaxIterations) { t.Fatalf("Inference error = %v, want ErrMaxIterations", err) } }