Add db message metadata

internal/db/queries_test.go

@@ -2,8 +2,11 @@ package db
 
 import (
 	"context"
+	"encoding/json/jsontext"
+	"encoding/json/v2"
 	"errors"
 	"log/slog"
+	"reflect"
 	"testing"
 
 	"github.com/mattn/go-sqlite3"
@@ -177,6 +180,7 @@ func testMessaging(t *testing.T, ctx context.Context, db *SQLite) {
 		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 {
@@ -186,6 +190,7 @@ func testMessaging(t *testing.T, ctx context.Context, db *SQLite) {
 		ConversationId: c.Id,
 		Sequence:       seq,
 		Role:           MessageRoleAssistant,
+		Metadata:       jsontext.Value(metadata),
 		Content:        `{"role":"assistant","content":null,"tool_calls":[{"id":"call_abc","type":"function","function":{"name":"get_weather","arguments":"{\"location\":\"NYC\"}"}}]}`,
 	})
 	if err != nil {
@@ -223,6 +228,22 @@ func testMessaging(t *testing.T, ctx context.Context, db *SQLite) {
 		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: `{"role":"assistant","content":"edited"}`,
+	}); err != nil {
+		t.Fatalf("UpdateMessageContent: %v", err)
+	}
+
 	// Test UpdateMessageContent.
 	if err := qw.UpdateMessageContent(ctx, UpdateMessageContentParams{
 		Id:      toolMsg.Id,
@@ -235,9 +256,116 @@ func testMessaging(t *testing.T, ctx context.Context, db *SQLite) {
 	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 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:        `{"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)
+			}
+		})
+	}
+	tx, err := db.Writer.BeginTx(ctx, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	_, err = New(tx).InsertMessage(ctx, InsertMessageParams{
+		Content:        `{"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 testConversationTimestamp(t *testing.T, ctx context.Context, database *SQLite) {