package db import ( "context" "database/sql" "encoding/json/jsontext" "encoding/json/v2" "errors" "log/slog" "reflect" "testing" "time" "uuid" "github.com/mattn/go-sqlite3" ) func TestQueries(t *testing.T) { var ( ctx = t.Context() db = createDB(t) ) defer db.Close() if err := db.ApplySchema(ctx, slog.Default()); err != nil { t.Fatalf("ApplySchema: %v", err) } t.Run("User", func(t *testing.T) { testUser(t, ctx, db) }) t.Run("Conversation", func(t *testing.T) { testConversation(t, ctx, db) }) t.Run("Messaging", func(t *testing.T) { testMessaging(t, ctx, db) }) t.Run("SequenceShift", func(t *testing.T) { testSequenceShift(t, ctx, db) }) t.Run("ConversationTimestamp", func(t *testing.T) { testConversationTimestamp(t, ctx, db) }) t.Run("Executions", func(t *testing.T) { testExecutions(t, ctx, db) }) t.Run("Delete", func(t *testing.T) { testDelete(t, ctx, db) }) } func testUser(t *testing.T, ctx context.Context, db *SQLite) { var ( qr = New(db.Reader) qw = New(db.Writer) ) u, err := qw.InsertUser(ctx, InsertUserParams{ Name: "user", Password: "test", Role: UserRoleOwner, }) if err != nil { t.Fatalf("InsertUser: %v", err) } res, err := qr.SelectUserId(ctx, u.Id) if err != nil { t.Fatalf("SelectUserId: %v", err) } if res == nil { t.Fatal("SelectUserId returned nil") } if res.Name != u.Name { t.Errorf( "expected name %s, got %s", u.Name, res.Name, ) } } func testConversation(t *testing.T, ctx context.Context, db *SQLite) { var ( qr = New(db.Reader) qw = New(db.Writer) title = "test" ) u, err := qw.InsertUser(ctx, InsertUserParams{ Name: "conversation", Password: "test", Role: UserRoleOwner, }) if err != nil { t.Fatalf("InsertUser: %v", err) } c, err := qw.InsertConversation(ctx, InsertConversationParams{ UserId: u.Id, Title: &title, }) if err != nil { t.Fatalf("InsertConversation: %v", err) } res, err := qr.SelectConversationId(ctx, c.Id) if err != nil { t.Fatalf("SelectConversationId: %v", err) } if res == nil { t.Fatal("SelectConversationId returned nil") } if res.Title == nil || *res.Title != title { t.Errorf("expected title %s, got %v", title, res.Title) } owned, err := qr.SelectConversationIdUserId(ctx, c.Id, u.Id) if err != nil { t.Fatalf("SelectConversationIdUserId: %v", err) } if owned.Id != c.Id { t.Errorf("expected conversation %s, got %s", c.Id, owned.Id) } if _, err := qr.SelectConversationIdUserId(ctx, c.Id, "unknown"); err == nil { t.Fatal("expected error selecting conversation for wrong user") } if _, err := qr.SelectConversationIdUserId(ctx, "unknown", u.Id); err == nil { t.Fatal("expected error selecting unknown conversation") } // Test conversation list. conversations, err := qr.SelectConversationsUserId(ctx, u.Id) if err != nil { t.Fatalf("SelectConversationsUserId: %v", err) } if len(conversations) != 1 { t.Fatalf("expected 1 conversation, got %d", len(conversations)) } // Rename the conversation. The rename must refresh updated_at: // plant a known-old timestamp, then check that the statement // moved it. const staleUpdatedAt = "2020-01-01T00:00:00Z" if _, err := db.Writer.ExecContext( ctx, "UPDATE conversations SET updated_at = ? WHERE id = ?", staleUpdatedAt, c.Id, ); err != nil { t.Fatalf("plant stale updated_at: %v", err) } renamed, err := qw.UpdateConversationTitle(ctx, UpdateConversationTitleParams{ Id: c.Id, Title: "renamed", }) if err != nil { t.Fatalf("UpdateConversationTitle: %v", err) } if renamed.Title == nil || *renamed.Title != "renamed" { t.Fatalf("Title = %v, want renamed", renamed.Title) } stale, err := time.Parse(time.RFC3339, staleUpdatedAt) if err != nil { t.Fatalf("parse staleUpdatedAt: %v", err) } if !renamed.UpdatedAt.After(stale) { t.Fatalf( "UpdatedAt = %v, want a time after %v", renamed.UpdatedAt, stale, ) } // A rename of a missing conversation reports no rows. if _, err := qw.UpdateConversationTitle(ctx, UpdateConversationTitleParams{ Id: "missing", Title: "renamed", }); !errors.Is(err, sql.ErrNoRows) { t.Fatalf("UpdateConversationTitle(missing) = %v, want sql.ErrNoRows", err) } // The guarded update sets the title of an untitled // conversation, and does not touch one that is titled. It must // refresh updated_at: plant the stale timestamp, then check // that the statement moved it. blank, err := qw.InsertConversation(ctx, InsertConversationParams{ UserId: u.Id, }) if err != nil { t.Fatalf("InsertConversation: %v", err) } if _, err := db.Writer.ExecContext( ctx, "UPDATE conversations SET updated_at = ? WHERE id = ?", staleUpdatedAt, blank.Id, ); err != nil { t.Fatalf("plant stale updated_at: %v", err) } titled, err := qw.UpdateConversationTitleIfUnset( ctx, UpdateConversationTitleIfUnsetParams{ Id: blank.Id, Title: "derived", }, ) if err != nil { t.Fatalf("UpdateConversationTitleIfUnset: %v", err) } if titled.Title == nil || *titled.Title != "derived" { t.Fatalf("Title = %v, want derived", titled.Title) } if !titled.UpdatedAt.After(stale) { t.Fatalf( "UpdatedAt = %v, want a time after %v", titled.UpdatedAt, stale, ) } for _, tt := range []struct { name string id string }{ {name: "titled twice", id: blank.Id}, {name: "renamed conversation", id: c.Id}, {name: "missing conversation", id: "missing"}, } { if _, err := qw.UpdateConversationTitleIfUnset( ctx, UpdateConversationTitleIfUnsetParams{ Id: tt.id, Title: "second", }, ); !errors.Is(err, sql.ErrNoRows) { t.Fatalf( "UpdateConversationTitleIfUnset(%s) = %v, want sql.ErrNoRows", tt.name, err, ) } } } func testMessaging(t *testing.T, ctx context.Context, db *SQLite) { var qw = New(db.Writer) u, err := qw.InsertUser(ctx, InsertUserParams{ Name: "messaging", Password: "test", Role: UserRoleOwner, }) if err != nil { t.Fatalf("InsertUser: %v", err) } c, err := qw.InsertConversation(ctx, InsertConversationParams{ UserId: u.Id, }) if err != nil { t.Fatalf("InsertConversation: %v", err) } // Insert a system message. seq, err := qw.NextSequence(ctx, c.Id) if err != nil { t.Fatalf("NextSequence: %v", err) } _, err = qw.InsertMessage(ctx, InsertMessageParams{ ConversationId: c.Id, Sequence: seq, Role: MessageRoleSystem, Content: jsontext.Value(`{"role":"system","content":"You are helpful"}`), }) if err != nil { t.Fatalf("InsertMessage: %v", err) } // Insert a user message. seq, err = qw.NextSequence(ctx, c.Id) if err != nil { t.Fatalf("NextSequence: %v", err) } _, err = qw.InsertMessage(ctx, InsertMessageParams{ ConversationId: c.Id, Sequence: seq, Role: MessageRoleUser, Content: jsontext.Value(`{"role":"user","content":"Hello"}`), }) if err != nil { t.Fatalf("InsertMessage: %v", err) } metadata := `{"id":"completion","usage":{"prompt_tokens":0}}` // Insert an assistant message with tool calls. seq, err = qw.NextSequence(ctx, c.Id) if err != nil { t.Fatalf("NextSequence: %v", err) } _, err = qw.InsertMessage(ctx, InsertMessageParams{ ConversationId: c.Id, Sequence: seq, Role: MessageRoleAssistant, Metadata: jsontext.Value(metadata), Content: jsontext.Value(`{"role":"assistant","content":null,"tool_calls":[{"id":"call_abc","type":"function","function":{"name":"get_weather","arguments":"{\"location\":\"NYC\"}"}}]}`), }) if err != nil { t.Fatalf("InsertMessage: %v", err) } // Insert a tool result message. seq, err = qw.NextSequence(ctx, c.Id) if err != nil { t.Fatalf("NextSequence: %v", err) } toolMsg, err := qw.InsertMessage(ctx, InsertMessageParams{ ConversationId: c.Id, Sequence: seq, Role: MessageRoleTool, Content: jsontext.Value(`{"role":"tool","content":"{\"temperature\":72}","tool_call_id":"call_abc"}`), }) if err != nil { t.Fatalf("InsertMessage: %v", err) } // Test SelectMessages. qr := New(db.Reader) messages, err := qr.SelectMessagesConversationId(ctx, c.Id) if err != nil { t.Fatalf("SelectMessagesConversationId: %v", err) } if len(messages) != 4 { t.Fatalf("expected 4 messages, got %d", len(messages)) } if messages[0].Role != MessageRoleSystem { t.Errorf("expected first message role system, got %s", messages[0].Role) } if messages[3].Role != MessageRoleTool { t.Errorf("expected last message role tool, got %s", messages[3].Role) } for i, message := range messages { if i == 2 { if message.Metadata == nil || string(message.Metadata) != metadata { t.Errorf("assistant metadata = %v", message.Metadata) } } else if message.Metadata != nil { t.Errorf("message %d metadata = %v, want nil", i, message.Metadata) } } if _, err := qw.UpdateMessageContent(ctx, UpdateMessageContentParams{ Id: messages[2].Id, Content: jsontext.Value(`{"role":"assistant","content":"edited"}`), }); err != nil { t.Fatalf("UpdateMessageContent: %v", err) } // Test UpdateMessageContent. if _, err := qw.UpdateMessageContent(ctx, UpdateMessageContentParams{ Id: toolMsg.Id, Content: jsontext.Value(`{"role":"tool","content":"{\"temperature\":75}","tool_call_id":"call_abc"}`), }); err != nil { t.Fatalf("UpdateMessageContent: %v", err) } updated, err := qr.SelectMessagesConversationId(ctx, c.Id) if err != nil { t.Fatalf("SelectMessagesConversationId: %v", err) } if updated[2].Metadata == nil || string(updated[2].Metadata) != metadata { t.Errorf("edited assistant metadata = %v", updated[2].Metadata) } if string(updated[3].Content) != `{"role":"tool","content":"{\"temperature\":75}","tool_call_id":"call_abc"}` { t.Errorf("updated content mismatch: %s", updated[3].Content) } for _, tt := range []struct { name string metadata jsontext.Value wantErr bool }{ {name: "absent"}, {name: "JSON null", metadata: jsontext.Value(`null`)}, {name: "empty object", metadata: jsontext.Value(`{}`)}, {name: "empty array", metadata: jsontext.Value(`[]`)}, {name: "malformed", metadata: jsontext.Value(`{`), wantErr: true}, {name: "empty", metadata: jsontext.Value{}, wantErr: true}, {name: "reported zero", metadata: jsontext.Value(`{"usage":{"total_tokens":0}}`)}, {name: "opaque JSON", metadata: jsontext.Value(`{"other":{"value":42}}`)}, } { t.Run(tt.name, func(t *testing.T) { seq, err := qw.NextSequence(ctx, c.Id) if err != nil { t.Fatal(err) } inserted, err := qw.InsertMessage(ctx, InsertMessageParams{ Content: jsontext.Value(`{"role":"assistant","content":"answer"}`), ConversationId: c.Id, Metadata: tt.metadata, Role: MessageRoleAssistant, Sequence: seq, }) if tt.wantErr { if err == nil { t.Fatal("invalid metadata was accepted") } return } if err != nil { t.Fatal(err) } var storedType string if err := db.Reader.QueryRowContext(ctx, "SELECT typeof(metadata) FROM messages WHERE id = ?", inserted.Id, ).Scan(&storedType); err != nil { t.Fatal(err) } wantType := "blob" if tt.metadata == nil { wantType = "null" } if storedType != wantType { t.Errorf("storage type = %s, want %s", storedType, wantType) } selected, err := qr.SelectMessagesConversationId(ctx, c.Id) if err != nil { t.Fatal(err) } for _, message := range []Message{*inserted, selected[len(selected)-1]} { if (message.Metadata == nil) != (tt.metadata == nil) || (tt.metadata != nil && string(message.Metadata) != string(tt.metadata)) { t.Fatalf("metadata = %v, want %v", message.Metadata, tt.metadata) } encoded, err := json.Marshal(message) if err != nil { t.Fatal(err) } var row map[string]any if err := json.Unmarshal(encoded, &row); err != nil { t.Fatal(err) } value, exists := row["metadata"] var wantValue any if tt.metadata != nil { if err := json.Unmarshal(tt.metadata, &wantValue); err != nil { t.Fatal(err) } } if !exists || !reflect.DeepEqual(value, wantValue) { t.Errorf("row JSON = %s", encoded) } } if err := qw.DeleteMessageId(ctx, inserted.Id); err != nil { t.Fatal(err) } }) } suppliedID := uuid.NewV4() for _, tt := range []struct { name string id *uuid.UUID wantErr bool }{ {name: "generated ID"}, {name: "supplied v4 ID", id: &suppliedID}, {name: "duplicate ID", id: &suppliedID, wantErr: true}, } { t.Run(tt.name, func(t *testing.T) { seq, err := qw.NextSequence(ctx, c.Id) if err != nil { t.Fatal(err) } inserted, err := qw.InsertMessage(ctx, InsertMessageParams{ Content: jsontext.Value(`{"role":"assistant","content":"answer"}`), ConversationId: c.Id, Id: tt.id, Role: MessageRoleAssistant, Sequence: seq, }) if tt.wantErr { if err == nil { t.Fatal("duplicate ID was accepted") } return } if err != nil { t.Fatal(err) } if _, err := uuid.Parse(inserted.Id); err != nil { t.Fatal(err) } if tt.id != nil && inserted.Id != tt.id.String() { t.Errorf("inserted ID = %s, want %s", inserted.Id, tt.id.String()) } if tt.id == nil { if err := qw.DeleteMessageId(ctx, inserted.Id); err != nil { t.Fatal(err) } } }) } if err := qw.DeleteMessageId(ctx, suppliedID.String()); err != nil { t.Fatal(err) } tx, err := db.Writer.BeginTx(ctx, nil) if err != nil { t.Fatal(err) } _, err = New(tx).InsertMessage(ctx, InsertMessageParams{ Content: jsontext.Value(`{"role":"assistant","content":"rolled back"}`), ConversationId: c.Id, Metadata: jsontext.Value(metadata), Role: MessageRoleAssistant, Sequence: 100, }) if err != nil { tx.Rollback() t.Fatal(err) } if err := tx.Rollback(); err != nil { t.Fatal(err) } selected, err := qr.SelectMessagesConversationId(ctx, c.Id) if err != nil || len(selected) != 4 { t.Fatalf("rollback messages = %d, error = %v", len(selected), err) } } func testSequenceShift(t *testing.T, ctx context.Context, db *SQLite) { var ( qr = New(db.Reader) qw = New(db.Writer) ) u, err := qw.InsertUser(ctx, InsertUserParams{ Name: "sequence-shift", Password: "test", Role: UserRoleOwner, }) if err != nil { t.Fatalf("InsertUser: %v", err) } c, err := qw.InsertConversation(ctx, InsertConversationParams{ UserId: u.Id, }) if err != nil { t.Fatalf("InsertConversation: %v", err) } insert := func(text string) { t.Helper() seq, err := qw.NextSequence(ctx, c.Id) if err != nil { t.Fatalf("NextSequence: %v", err) } if _, err := qw.InsertMessage(ctx, InsertMessageParams{ ConversationId: c.Id, Sequence: seq, Role: MessageRoleUser, Content: jsontext.Value(`{"role":"user","content":"` + text + `"}`), }); err != nil { t.Fatalf("InsertMessage: %v", err) } } for _, text := range []string{"one", "two", "three", "four"} { insert(text) } // updated_at has millisecond precision, so sleep to put the // shift in a later millisecond than the inserts. time.Sleep(10 * time.Millisecond) before, err := qr.SelectMessagesConversationId(ctx, c.Id) if err != nil { t.Fatalf("SelectMessagesConversationId: %v", err) } // Deleting "two" and shifting closes the gap: "three" moves to // 2 and "four" to 3, while "one" stays put. if err := qw.DeleteMessageId(ctx, before[1].Id); err != nil { t.Fatalf("DeleteMessageId: %v", err) } if err := qw.UpdateMessagesSequence(ctx, c.Id, before[1].Sequence); err != nil { t.Fatalf("UpdateMessagesSequence: %v", err) } moved, err := qr.SelectMessagesConversationId(ctx, c.Id) if err != nil { t.Fatalf("SelectMessagesConversationId: %v", err) } if len(moved) != 3 { t.Fatalf("expected 3 messages, got %d", len(moved)) } if moved[0].Sequence != 1 { t.Errorf("unshifted row sequence = %d, want 1", moved[0].Sequence) } if !moved[0].UpdatedAt.Equal(before[0].UpdatedAt) { t.Errorf("unshifted row updated_at = %s, want %s", moved[0].UpdatedAt, before[0].UpdatedAt) } for i, wantSeq := range []int{2, 3} { row, prev := moved[i+1], before[i+2] if row.Id != prev.Id { t.Fatalf("row %d id = %s, want %s", i+1, row.Id, prev.Id) } if row.Sequence != wantSeq { t.Errorf("row %d sequence = %d, want %d", i+1, row.Sequence, wantSeq) } if !row.UpdatedAt.After(prev.UpdatedAt) { t.Errorf("row %d updated_at = %s, want later than %s", i+1, row.UpdatedAt, prev.UpdatedAt) } } // Free sequence 2 by incrementing rows 2 and 3, inside a // transaction as the doc requires. time.Sleep(10 * time.Millisecond) tx, err := db.Writer.BeginTx(ctx, nil) if err != nil { t.Fatalf("BeginTx: %v", err) } if err := New(tx).UpdateMessagesSequenceIncrement(ctx, c.Id, 2); err != nil { tx.Rollback() t.Fatalf("UpdateMessagesSequenceIncrement: %v", err) } if err := tx.Commit(); err != nil { t.Fatalf("Commit: %v", err) } after, err := qr.SelectMessagesConversationId(ctx, c.Id) if err != nil { t.Fatalf("SelectMessagesConversationId: %v", err) } if len(after) != 3 { t.Fatalf("expected 3 messages, got %d", len(after)) } if after[0].Sequence != 1 { t.Errorf("unshifted row sequence = %d, want 1", after[0].Sequence) } if !after[0].UpdatedAt.Equal(before[0].UpdatedAt) { t.Errorf("unshifted row updated_at = %s, want %s", after[0].UpdatedAt, before[0].UpdatedAt) } for i, wantSeq := range []int{3, 4} { row, prev := after[i+1], moved[i+1] if row.Id != prev.Id { t.Fatalf("row %d id = %s, want %s", i+1, row.Id, prev.Id) } if row.Sequence != wantSeq { t.Errorf("row %d sequence = %d, want %d", i+1, row.Sequence, wantSeq) } if !row.UpdatedAt.After(prev.UpdatedAt) { t.Errorf("row %d updated_at = %s, want later than %s", i+1, row.UpdatedAt, prev.UpdatedAt) } } } func testConversationTimestamp(t *testing.T, ctx context.Context, database *SQLite) { var ( qr = New(database.Reader) qw = New(database.Writer) ) u, err := qw.InsertUser(ctx, InsertUserParams{ Name: "conversation-timestamp", Password: "test", Role: UserRoleOwner, }) if err != nil { t.Fatalf("InsertUser: %v", err) } c, err := qw.InsertConversation(ctx, InsertConversationParams{UserId: u.Id}) if err != nil { t.Fatalf("InsertConversation: %v", err) } const oldTimestamp = "2000-01-01T00:00:00.000Z" if _, err := database.Writer.ExecContext( ctx, "UPDATE conversations SET updated_at = ? WHERE id = ?", oldTimestamp, c.Id, ); err != nil { t.Fatalf("set conversation timestamp before insert: %v", err) } msg, err := qw.InsertMessage(ctx, InsertMessageParams{ ConversationId: c.Id, Sequence: 0, Role: MessageRoleUser, Content: jsontext.Value(`{"role":"user","content":"Hello"}`), }) if err != nil { t.Fatalf("InsertMessage: %v", err) } conv, err := qr.SelectConversationId(ctx, c.Id) if err != nil { t.Fatalf("SelectConversationId after message insert: %v", err) } if conv.UpdatedAt.Format("2006-01-02T15:04:05.000Z") == oldTimestamp { t.Error("message insert did not update conversation timestamp") } if _, err := database.Writer.ExecContext( ctx, "UPDATE conversations SET updated_at = ? WHERE id = ?", oldTimestamp, c.Id, ); err != nil { t.Fatalf("set conversation timestamp before update: %v", err) } if _, err := qw.UpdateMessageContent(ctx, UpdateMessageContentParams{ Id: msg.Id, Content: jsontext.Value(`{"role":"user","content":"Updated"}`), }); err != nil { t.Fatalf("UpdateMessageContent: %v", err) } conv, err = qr.SelectConversationId(ctx, c.Id) if err != nil { t.Fatalf("SelectConversationId after message update: %v", err) } if conv.UpdatedAt.Format("2006-01-02T15:04:05.000Z") == oldTimestamp { t.Error("message update did not update conversation timestamp") } if _, err := database.Writer.ExecContext( ctx, "UPDATE conversations SET updated_at = ? WHERE id = ?", oldTimestamp, c.Id, ); err != nil { t.Fatalf("set conversation timestamp before delete: %v", err) } if err := qw.DeleteMessageId(ctx, msg.Id); err != nil { t.Fatalf("DeleteMessageId: %v", err) } conv, err = qr.SelectConversationId(ctx, c.Id) if err != nil { t.Fatalf("SelectConversationId after message delete: %v", err) } if conv.UpdatedAt.Format("2006-01-02T15:04:05.000Z") == oldTimestamp { t.Error("message delete did not update conversation timestamp") } } func testDelete(t *testing.T, ctx context.Context, db *SQLite) { var ( qr = New(db.Reader) qw = New(db.Writer) ) u, err := qw.InsertUser(ctx, InsertUserParams{ Name: "delete", Password: "test", Role: UserRoleOwner, }) if err != nil { t.Fatalf("InsertUser: %v", err) } c, err := qw.InsertConversation(ctx, InsertConversationParams{ UserId: u.Id, }) if err != nil { t.Fatalf("InsertConversation: %v", err) } msg, err := qw.InsertMessage(ctx, InsertMessageParams{ ConversationId: c.Id, Sequence: 1, Role: MessageRoleUser, Content: jsontext.Value(`{"role":"user","content":"Hello"}`), }) if err != nil { t.Fatalf("InsertMessage: %v", err) } // Delete message. if err := qw.DeleteMessageId(ctx, msg.Id); err != nil { t.Fatalf("DeleteMessageId: %v", err) } messages, err := qr.SelectMessagesConversationId(ctx, c.Id) if err != nil { t.Fatalf("SelectMessagesConversationId: %v", err) } if len(messages) != 0 { t.Fatalf("expected 0 messages after delete, got %d", len(messages)) } // Delete conversation. if err := qw.DeleteConversationId(ctx, c.Id); err != nil { t.Fatalf("DeleteConversationId: %v", err) } _, err = qr.SelectConversationId(ctx, c.Id) if err == nil { t.Fatal("expected error selecting deleted conversation") } } func testExecutions(t *testing.T, ctx context.Context, db *SQLite) { var ( qr = New(db.Reader) qw = New(db.Writer) ) u, err := qw.InsertUser(ctx, InsertUserParams{ Name: "approvals", Password: "test", Role: UserRoleOwner, }) if err != nil { t.Fatalf("InsertUser: %v", err) } c, err := qw.InsertConversation(ctx, InsertConversationParams{ UserId: u.Id, }) if err != nil { t.Fatalf("InsertConversation: %v", err) } msg, err := qw.InsertMessage(ctx, InsertMessageParams{ ConversationId: c.Id, Sequence: 1, Role: MessageRoleAssistant, Content: jsontext.Value(`{"role":"assistant","content":null,"tool_calls":[{"id":"call_abc","type":"function","function":{"name":"host_run","arguments":"{\"command\":\"rm -rf /\"}"}}]}`), }) if err != nil { t.Fatalf("InsertMessage: %v", err) } // Insert an execution. execution, err := qw.InsertExecution(ctx, InsertExecutionParams{ MessageId: msg.Id, Status: ExecutionStatusPending, ToolCallId: "call_abc", }) if err != nil { t.Fatalf("InsertExecution: %v", err) } if execution.Status != ExecutionStatusPending { t.Fatalf("expected pending, got %s", execution.Status) } if execution.ToolCallId != "call_abc" { t.Fatalf("expected tool_call_id call_abc, got %s", execution.ToolCallId) } // Select by ID. res, err := qr.SelectExecutionId(ctx, execution.Id) if err != nil { t.Fatalf("SelectExecutionId: %v", err) } if res == nil { t.Fatal("SelectExecutionId returned nil") } if res.MessageId != msg.Id { t.Errorf("expected message_id %s, got %s", msg.Id, res.MessageId) } // Select by message ID. executions, err := qr.SelectExecutionsMessageId(ctx, msg.Id) if err != nil { t.Fatalf("SelectExecutionsMessageId: %v", err) } if len(executions) != 1 { t.Fatalf("expected 1 execution, got %d", len(executions)) } // Select pending. pending, err := qr.SelectPendingExecutions(ctx) if err != nil { t.Fatalf("SelectPendingExecutions: %v", err) } if len(pending) != 1 { t.Fatalf("expected 1 pending execution, got %d", len(pending)) } // Transition to approved. if err := qw.UpdateExecutionStatus(ctx, UpdateExecutionStatusParams{ Id: execution.Id, Status: ExecutionStatusApproved, }); err != nil { t.Fatalf("UpdateExecutionStatus: %v", err) } res, err = qr.SelectExecutionId(ctx, execution.Id) if err != nil { t.Fatalf("SelectExecutionId after update: %v", err) } if res.Status != ExecutionStatusApproved { t.Fatalf("expected approved, got %s", res.Status) } // Pending should now be empty. pending, err = qr.SelectPendingExecutions(ctx) if err != nil { t.Fatalf("SelectPendingExecutions after update: %v", err) } if len(pending) != 0 { t.Fatalf("expected 0 pending, got %d", len(pending)) } // Insert a second execution and deny it. execution2, err := qw.InsertExecution(ctx, InsertExecutionParams{ MessageId: msg.Id, Status: ExecutionStatusPending, ToolCallId: "call_def", }) if err != nil { t.Fatalf("InsertExecution second: %v", err) } if err := qw.UpdateExecutionStatus(ctx, UpdateExecutionStatusParams{ Id: execution2.Id, Status: ExecutionStatusDenied, }); err != nil { t.Fatalf("UpdateExecutionStatus denied: %v", err) } res2, err := qr.SelectExecutionId(ctx, execution2.Id) if err != nil { t.Fatalf("SelectExecutionId after deny: %v", err) } if res2.Status != ExecutionStatusDenied { t.Fatalf("expected denied, got %s", res2.Status) } // Transition to running, then completed. if err := qw.UpdateExecutionStatus(ctx, UpdateExecutionStatusParams{ Id: execution.Id, Status: ExecutionStatusRunning, }); err != nil { t.Fatalf("UpdateExecutionStatus running: %v", err) } if err := qw.UpdateExecutionStatus(ctx, UpdateExecutionStatusParams{ Id: execution.Id, Status: ExecutionStatusCompleted, }); err != nil { t.Fatalf("UpdateExecutionStatus completed: %v", err) } // Both executions for the message, ordered by created_at. all, err := qr.SelectExecutionsMessageId(ctx, msg.Id) if err != nil { t.Fatalf("SelectExecutionsMessageId final: %v", err) } if len(all) != 2 { t.Fatalf("expected 2 executions, got %d", len(all)) } if all[0].Status != ExecutionStatusCompleted { t.Fatalf("expected first execution completed, got %s", all[0].Status) } if all[1].Status != ExecutionStatusDenied { t.Fatalf("expected second execution denied, got %s", all[1].Status) } // Parameter values and the schema default, on a second message. msg2, err := qw.InsertMessage(ctx, InsertMessageParams{ ConversationId: c.Id, Sequence: 2, Role: MessageRoleAssistant, Content: jsontext.Value(`{"role":"assistant","content":"second"}`), }) if err != nil { t.Fatalf("InsertMessage second: %v", err) } execRunning, err := qw.InsertExecution(ctx, InsertExecutionParams{ MessageId: msg2.Id, Status: ExecutionStatusRunning, ToolCallId: "call_running", }) if err != nil { t.Fatalf("InsertExecution running: %v", err) } if execRunning.Status != ExecutionStatusRunning { t.Fatalf("expected running, got %s", execRunning.Status) } // A status other than running is stored as passed. execPending, err := qw.InsertExecution(ctx, InsertExecutionParams{ MessageId: msg2.Id, Status: ExecutionStatusPending, ToolCallId: "call_pending", }) if err != nil { t.Fatalf("InsertExecution pending: %v", err) } if execPending.Status != ExecutionStatusPending { t.Fatalf("expected pending, got %s", execPending.Status) } // The zero value is rejected by the table's status check. _, err = qw.InsertExecution(ctx, InsertExecutionParams{ MessageId: msg2.Id, ToolCallId: "call_empty", }) var sqliteErr sqlite3.Error if !errors.As(err, &sqliteErr) { t.Fatalf("InsertExecution with zero status: expected sqlite3.Error, got %v", err) } if sqliteErr.ExtendedCode != sqlite3.ErrConstraintCheck { t.Fatalf("InsertExecution with zero status: expected ErrConstraintCheck, got %v", err) } // A second row for the same message and tool call is rejected. _, err = qw.InsertExecution(ctx, InsertExecutionParams{ MessageId: msg2.Id, Status: ExecutionStatusPending, ToolCallId: "call_running", }) if !errors.As(err, &sqliteErr) { t.Fatalf("InsertExecution duplicate: expected sqlite3.Error, got %v", err) } if sqliteErr.ExtendedCode != sqlite3.ErrConstraintUnique { t.Fatalf("InsertExecution duplicate: expected ErrConstraintUnique, got %v", err) } // A row inserted without a status takes the schema default, running. if _, err := db.Writer.ExecContext( ctx, `INSERT INTO executions (id, message_id, tool_call_id) VALUES (?, ?, ?)`, "exec_default", msg2.Id, "call_default", ); err != nil { t.Fatalf("insert execution without status: %v", err) } res, err = qr.SelectExecutionId(ctx, "exec_default") if err != nil { t.Fatalf("SelectExecutionId default: %v", err) } if res == nil { t.Fatal("SelectExecutionId default returned nil") } if res.Status != ExecutionStatusRunning { t.Fatalf("expected default running, got %s", res.Status) } // The rejected insert stored no row: the message's rows are exactly the // three successful inserts. all2, err := qr.SelectExecutionsMessageId(ctx, msg2.Id) if err != nil { t.Fatalf("SelectExecutionsMessageId second: %v", err) } seen := map[string]bool{} for _, e := range all2 { seen[e.ToolCallId] = true } for _, id := range []string{"call_default", "call_pending", "call_running"} { if !seen[id] { t.Fatalf("execution %s not found", id) } } if len(seen) != 3 { t.Fatalf("expected 3 executions, got %d", len(seen)) } }