tatara-chat (phase 8) Implementation Plan¶

package chat

import "testing"

func sp(s string) *string { return &s }

func TestFilterVisible(t *testing.T) {
    me := "p-me"
    other := "p-other"
    msgs := []Message{
        {ID: 1, Body: "global", AuthorID: sp(other)},                 // global, visible
        {ID: 2, Body: "dm-to-me", AuthorID: sp(other), TargetID: sp(me)},   // dm to me, visible
        {ID: 3, Body: "dm-to-other", AuthorID: sp(me), TargetID: sp(other)}, // dm by me, visible
        {ID: 4, Body: "dm-elsewhere", AuthorID: sp(other), TargetID: sp(other)}, // hidden
        {ID: 5, Body: "system", Kind: "system"},                      // global system, visible
    }
    got := filterVisible(msgs, me)
    if len(got) != 4 {
        t.Fatalf("want 4 visible, got %d", len(got))
    }
    for _, m := range got {
        if m.ID == 4 {
            t.Fatalf("message 4 should be hidden from %s", me)
        }
    }
}

// Package chat provides the agent chat subsystem: named rooms, participants
// with server-generated UUID identity, and an append-only message log with
// per-participant read cursors.
package chat

import (
    "errors"
    "time"
)

// ErrNotFound is returned when a room or participant does not exist.
var ErrNotFound = errors.New("not found")

// ErrRoomArchived is returned when writing to an archived room.
var ErrRoomArchived = errors.New("room is archived")

// Room is a named chat instance.
type Room struct {
    ID        string    `json:"id"`
    Name      string    `json:"name"`
    CreatedBy string    `json:"created_by"`
    Status    string    `json:"status"`
    CreatedAt time.Time `json:"created_at"`
}

// Participant is an agent enrolled in a room. ID is server-generated and is
// the handle the agent uses to send and poll. Name is a display label only.
type Participant struct {
    ID         string     `json:"id"`
    RoomID     string     `json:"room_id"`
    Name       string     `json:"name"`
    Role       string     `json:"role"`
    JoinedAt   time.Time  `json:"joined_at"`
    LeftAt     *time.Time `json:"left_at,omitempty"`
    LastSeenID int64      `json:"last_seen_id"`
}

// Message is one entry in a room's append-only log. TargetID nil means the
// message went to the whole room; a non-nil TargetID is a direct message.
// AuthorID is nil for system-generated lifecycle notices.
type Message struct {
    ID         int64     `json:"id"`
    RoomID     string    `json:"room_id"`
    AuthorID   *string   `json:"author_id,omitempty"`
    AuthorName string    `json:"author_name"`
    TargetID   *string   `json:"target_id,omitempty"`
    Kind       string    `json:"kind"`
    Body       string    `json:"body"`
    CreatedAt  time.Time `json:"created_at"`
}

// filterVisible returns the messages visible to the participant id: global
// messages plus direct messages sent to or by the participant.
func filterVisible(msgs []Message, pid string) []Message {
    out := make([]Message, 0, len(msgs))
    for _, m := range msgs {
        switch {
        case m.TargetID == nil:
            out = append(out, m)
        case *m.TargetID == pid:
            out = append(out, m)
        case m.AuthorID != nil && *m.AuthorID == pid:
            out = append(out, m)
        }
    }
    return out
}

package chat

import (
    "context"
    "database/sql"
    _ "embed"
)

//go:embed migrations/0001_chat.sql
var migration0001 string

// Migrate applies the chat schema to db, creating tables if they do not exist.
func Migrate(ctx context.Context, db *sql.DB) error {
    _, err := db.ExecContext(ctx, migration0001)
    return err
}

//go:build integration

package chat_test

import (
    "context"
    "database/sql"
    "os"
    "testing"

    _ "github.com/jackc/pgx/v5/stdlib"
    "github.com/stretchr/testify/require"

    "github.com/szymonrychu/tatara-chat/internal/chat"
)

// openPG returns a *sql.DB to a throwaway test Postgres, skipping when
// TATARA_TEST_PG_DSN is unset. Each test migrates and truncates.
func openPG(t *testing.T) *sql.DB {
    t.Helper()
    dsn := os.Getenv("TATARA_TEST_PG_DSN")
    if dsn == "" {
        t.Skip("TATARA_TEST_PG_DSN not set")
    }
    db, err := sql.Open("pgx", dsn)
    require.NoError(t, err)
    require.NoError(t, db.PingContext(context.Background()))
    require.NoError(t, chat.Migrate(context.Background(), db))
    _, err = db.Exec(`TRUNCATE chat_messages, chat_participants, chat_rooms RESTART IDENTITY CASCADE`)
    require.NoError(t, err)
    return db
}

func TestMigrateIdempotent(t *testing.T) {
    db := openPG(t)
    defer db.Close()
    require.NoError(t, chat.Migrate(context.Background(), db))
}

//go:build integration

package chat_test

import (
    "context"
    "testing"

    "github.com/stretchr/testify/require"

    "github.com/szymonrychu/tatara-chat/internal/chat"
)

func TestRoomLifecycle(t *testing.T) {
    ctx := context.Background()
    db := openPG(t)
    defer db.Close()
    s := chat.NewStore(db)

    rm, err := s.CreateRoom(ctx, "stream-x", "orchestrator")
    require.NoError(t, err)
    require.NotEmpty(t, rm.ID)
    require.Equal(t, "active", rm.Status)

    got, err := s.GetRoom(ctx, rm.ID)
    require.NoError(t, err)
    require.Equal(t, rm.ID, got.ID)

    rooms, err := s.ListRooms(ctx, "active")
    require.NoError(t, err)
    require.Len(t, rooms, 1)

    require.NoError(t, s.ArchiveRoom(ctx, rm.ID))
    got, err = s.GetRoom(ctx, rm.ID)
    require.NoError(t, err)
    require.Equal(t, "archived", got.Status)

    _, err = s.GetRoom(ctx, "00000000-0000-0000-0000-000000000000")
    require.ErrorIs(t, err, chat.ErrNotFound)
}

package chat

import (
    "context"
    "database/sql"
    "errors"
    "fmt"
)

// Store persists rooms, participants, and messages on a *sql.DB (pgx stdlib).
type Store struct {
    db *sql.DB
}

// NewStore creates a chat Store.
func NewStore(db *sql.DB) *Store { return &Store{db: db} }

// CreateRoom inserts a new active room.
func (s *Store) CreateRoom(ctx context.Context, name, createdBy string) (Room, error) {
    var rm Room
    err := s.db.QueryRowContext(ctx, `
        INSERT INTO chat_rooms (name, created_by)
        VALUES ($1, $2)
        RETURNING id, name, created_by, status, created_at`,
        name, createdBy).Scan(&rm.ID, &rm.Name, &rm.CreatedBy, &rm.Status, &rm.CreatedAt)
    if err != nil {
        return Room{}, fmt.Errorf("create room: %w", err)
    }
    return rm, nil
}

// ListRooms returns rooms, optionally filtered by status, newest first.
func (s *Store) ListRooms(ctx context.Context, status string) ([]Room, error) {
    q := `SELECT id, name, created_by, status, created_at FROM chat_rooms`
    var args []any
    if status != "" {
        q += ` WHERE status = $1`
        args = append(args, status)
    }
    q += ` ORDER BY created_at DESC`
    rows, err := s.db.QueryContext(ctx, q, args...)
    if err != nil {
        return nil, fmt.Errorf("list rooms: %w", err)
    }
    defer rows.Close()
    var out []Room
    for rows.Next() {
        var rm Room
        if err := rows.Scan(&rm.ID, &rm.Name, &rm.CreatedBy, &rm.Status, &rm.CreatedAt); err != nil {
            return nil, fmt.Errorf("scan room: %w", err)
        }
        out = append(out, rm)
    }
    return out, rows.Err()
}

// GetRoom fetches a room by id. Returns ErrNotFound when missing.
func (s *Store) GetRoom(ctx context.Context, id string) (Room, error) {
    var rm Room
    err := s.db.QueryRowContext(ctx, `
        SELECT id, name, created_by, status, created_at FROM chat_rooms WHERE id = $1`,
        id).Scan(&rm.ID, &rm.Name, &rm.CreatedBy, &rm.Status, &rm.CreatedAt)
    if errors.Is(err, sql.ErrNoRows) {
        return Room{}, ErrNotFound
    }
    if err != nil {
        return Room{}, fmt.Errorf("get room: %w", err)
    }
    return rm, nil
}

// ArchiveRoom sets a room's status to archived. Returns ErrNotFound when missing.
func (s *Store) ArchiveRoom(ctx context.Context, id string) error {
    tag, err := s.db.ExecContext(ctx, `UPDATE chat_rooms SET status = 'archived' WHERE id = $1`, id)
    if err != nil {
        return fmt.Errorf("archive room: %w", err)
    }
    n, _ := tag.RowsAffected()
    if n == 0 {
        return ErrNotFound
    }
    return nil
}

func TestParticipantLifecycle(t *testing.T) {
    ctx := context.Background()
    db := openPG(t)
    defer db.Close()
    s := chat.NewStore(db)

    rm, err := s.CreateRoom(ctx, "stream-y", "orchestrator")
    require.NoError(t, err)

    p, err := s.AddParticipant(ctx, rm.ID, "impl-1", "implementer")
    require.NoError(t, err)
    require.NotEmpty(t, p.ID)
    require.Equal(t, rm.ID, p.RoomID)
    require.Equal(t, "implementer", p.Role)

    got, err := s.GetParticipant(ctx, rm.ID, p.ID)
    require.NoError(t, err)
    require.Equal(t, "impl-1", got.Name)

    parts, err := s.ListParticipants(ctx, rm.ID)
    require.NoError(t, err)
    require.Len(t, parts, 1)

    require.NoError(t, s.RemoveParticipant(ctx, rm.ID, p.ID))
    _, err = s.GetParticipant(ctx, rm.ID, p.ID)
    require.ErrorIs(t, err, chat.ErrNotFound) // left participants are not active

    _, err = s.AddParticipant(ctx, "00000000-0000-0000-0000-000000000000", "x", "implementer")
    require.ErrorIs(t, err, chat.ErrNotFound)
}

// AddParticipant enrolls an agent in a room and returns it with a
// server-generated UUID id. Its read cursor starts at the room's current max
// message id, so it sees only messages from join onward.
func (s *Store) AddParticipant(ctx context.Context, roomID, name, role string) (Participant, error) {
    if _, err := s.GetRoom(ctx, roomID); err != nil {
        return Participant{}, err
    }
    var p Participant
    err := s.db.QueryRowContext(ctx, `
        INSERT INTO chat_participants (room_id, name, role, last_seen_id)
        VALUES ($1, $2, $3, COALESCE((SELECT MAX(id) FROM chat_messages WHERE room_id = $1), 0))
        RETURNING id, room_id, name, role, joined_at, left_at, last_seen_id`,
        roomID, name, role).Scan(
        &p.ID, &p.RoomID, &p.Name, &p.Role, &p.JoinedAt, &p.LeftAt, &p.LastSeenID)
    if err != nil {
        return Participant{}, fmt.Errorf("add participant: %w", err)
    }
    return p, nil
}

// GetParticipant fetches an active (not left) participant of a room by id.
func (s *Store) GetParticipant(ctx context.Context, roomID, participantID string) (Participant, error) {
    var p Participant
    err := s.db.QueryRowContext(ctx, `
        SELECT id, room_id, name, role, joined_at, left_at, last_seen_id
        FROM chat_participants
        WHERE id = $1 AND room_id = $2 AND left_at IS NULL`,
        participantID, roomID).Scan(
        &p.ID, &p.RoomID, &p.Name, &p.Role, &p.JoinedAt, &p.LeftAt, &p.LastSeenID)
    if errors.Is(err, sql.ErrNoRows) {
        return Participant{}, ErrNotFound
    }
    if err != nil {
        return Participant{}, fmt.Errorf("get participant: %w", err)
    }
    return p, nil
}

// RemoveParticipant marks an active participant as left.
func (s *Store) RemoveParticipant(ctx context.Context, roomID, participantID string) error {
    tag, err := s.db.ExecContext(ctx, `
        UPDATE chat_participants SET left_at = NOW()
        WHERE id = $1 AND room_id = $2 AND left_at IS NULL`,
        participantID, roomID)
    if err != nil {
        return fmt.Errorf("remove participant: %w", err)
    }
    n, _ := tag.RowsAffected()
    if n == 0 {
        return ErrNotFound
    }
    return nil
}

// ListParticipants returns every participant of a room, join order.
func (s *Store) ListParticipants(ctx context.Context, roomID string) ([]Participant, error) {
    rows, err := s.db.QueryContext(ctx, `
        SELECT id, room_id, name, role, joined_at, left_at, last_seen_id
        FROM chat_participants WHERE room_id = $1 ORDER BY joined_at`, roomID)
    if err != nil {
        return nil, fmt.Errorf("list participants: %w", err)
    }
    defer rows.Close()
    var out []Participant
    for rows.Next() {
        var p Participant
        if err := rows.Scan(&p.ID, &p.RoomID, &p.Name, &p.Role, &p.JoinedAt, &p.LeftAt, &p.LastSeenID); err != nil {
            return nil, fmt.Errorf("scan participant: %w", err)
        }
        out = append(out, p)
    }
    return out, rows.Err()
}

func TestMessagesPollLog(t *testing.T) {
    ctx := context.Background()
    db := openPG(t)
    defer db.Close()
    s := chat.NewStore(db)

    rm, _ := s.CreateRoom(ctx, "stream-z", "orchestrator")
    a, _ := s.AddParticipant(ctx, rm.ID, "a", "implementer")
    b, _ := s.AddParticipant(ctx, rm.ID, "b", "reviewer")

    // a posts a global message and a DM to b.
    _, err := s.AddMessage(ctx, rm.ID, a.ID, nil, "message", "global hi")
    require.NoError(t, err)
    _, err = s.AddMessage(ctx, rm.ID, a.ID, &b.ID, "message", "dm to b")
    require.NoError(t, err)

    // b polls: sees both (global + dm-to-b).
    msgs, status, err := s.Poll(ctx, rm.ID, b.ID)
    require.NoError(t, err)
    require.Equal(t, "active", status)
    require.Len(t, msgs, 2)

    // second poll by b: cursor advanced, nothing new.
    msgs, _, err = s.Poll(ctx, rm.ID, b.ID)
    require.NoError(t, err)
    require.Len(t, msgs, 0)

    // a polls: sees its own global + its own dm (author==me), not filtered out.
    msgs, _, err = s.Poll(ctx, rm.ID, a.ID)
    require.NoError(t, err)
    require.Len(t, msgs, 2)

    // log returns everything regardless of cursor.
    all, err := s.Log(ctx, rm.ID)
    require.NoError(t, err)
    require.Len(t, all, 2)

    // invalid author rejected.
    _, err = s.AddMessage(ctx, rm.ID, "00000000-0000-0000-0000-000000000000", nil, "message", "x")
    require.ErrorIs(t, err, chat.ErrNotFound)

    // archived room rejects writes.
    require.NoError(t, s.ArchiveRoom(ctx, rm.ID))
    _, err = s.AddMessage(ctx, rm.ID, a.ID, nil, "message", "late")
    require.ErrorIs(t, err, chat.ErrRoomArchived)
}

// AddMessage appends a message to a room. The author must be an active
// participant of the room; the author's display name is denormalized onto the
// row. Archived rooms are rejected.
func (s *Store) AddMessage(ctx context.Context, roomID, authorID string, target *string, kind, body string) (Message, error) {
    rm, err := s.GetRoom(ctx, roomID)
    if err != nil {
        return Message{}, err
    }
    if rm.Status == "archived" {
        return Message{}, ErrRoomArchived
    }
    var m Message
    err = s.db.QueryRowContext(ctx, `
        INSERT INTO chat_messages (room_id, author_id, author_name, target_id, kind, body)
        SELECT $1, p.id, p.name, $3, $4, $5
        FROM chat_participants p
        WHERE p.id = $2 AND p.room_id = $1 AND p.left_at IS NULL
        RETURNING id, room_id, author_id, author_name, target_id, kind, body, created_at`,
        roomID, authorID, target, kind, body).Scan(
        &m.ID, &m.RoomID, &m.AuthorID, &m.AuthorName, &m.TargetID, &m.Kind, &m.Body, &m.CreatedAt)
    if errors.Is(err, sql.ErrNoRows) {
        return Message{}, ErrNotFound // author not an active participant
    }
    if err != nil {
        return Message{}, fmt.Errorf("add message: %w", err)
    }
    return m, nil
}

// Poll returns the messages visible to the participant with id greater than
// the participant's read cursor, advances the cursor to the max fetched id,
// and returns the room status. A FOR UPDATE row lock on the participant row
// serializes concurrent polls, preventing cursor overwrites and duplicate
// delivery. The cursor advances past global messages the participant filtered
// out so they are not re-scanned.
func (s *Store) Poll(ctx context.Context, roomID, participantID string) ([]Message, string, error) {
    rm, err := s.GetRoom(ctx, roomID)
    if err != nil {
        return nil, "", err
    }

    tx, err := s.db.BeginTx(ctx, nil)
    if err != nil {
        return nil, "", fmt.Errorf("poll begin tx: %w", err)
    }
    defer func() { _ = tx.Rollback() }()

    var cursor int64
    err = tx.QueryRowContext(ctx, `
        SELECT last_seen_id FROM chat_participants
        WHERE id = $1 AND room_id = $2 AND left_at IS NULL
        FOR UPDATE`,
        participantID, roomID).Scan(&cursor)
    if errors.Is(err, sql.ErrNoRows) {
        return nil, "", ErrNotFound
    }
    if err != nil {
        return nil, "", fmt.Errorf("poll cursor: %w", err)
    }

    rows, err := tx.QueryContext(ctx, `
        SELECT id, room_id, author_id, author_name, target_id, kind, body, created_at
        FROM chat_messages WHERE room_id = $1 AND id > $2 ORDER BY id`,
        roomID, cursor)
    if err != nil {
        return nil, "", fmt.Errorf("poll messages: %w", err)
    }
    defer rows.Close()
    var all []Message
    for rows.Next() {
        var m Message
        if err := rows.Scan(&m.ID, &m.RoomID, &m.AuthorID, &m.AuthorName, &m.TargetID, &m.Kind, &m.Body, &m.CreatedAt); err != nil {
            return nil, "", fmt.Errorf("scan message: %w", err)
        }
        all = append(all, m)
    }
    if err := rows.Err(); err != nil {
        return nil, "", fmt.Errorf("poll rows: %w", err)
    }

    newCursor := cursor
    if n := len(all); n > 0 {
        newCursor = all[n-1].ID
    }
    if newCursor != cursor {
        if _, err := tx.ExecContext(ctx, `
            UPDATE chat_participants SET last_seen_id = $3
            WHERE id = $1 AND room_id = $2`,
            participantID, roomID, newCursor); err != nil {
            return nil, "", fmt.Errorf("advance cursor: %w", err)
        }
    }

    if err := tx.Commit(ctx); err != nil {
        return nil, "", fmt.Errorf("poll commit: %w", err)
    }
    return filterVisible(all, participantID), rm.Status, nil
}

// Log returns every message in a room in id order, without advancing any
// cursor. Used for review.
func (s *Store) Log(ctx context.Context, roomID string) ([]Message, error) {
    if _, err := s.GetRoom(ctx, roomID); err != nil {
        return nil, err
    }
    rows, err := s.db.QueryContext(ctx, `
        SELECT id, room_id, author_id, author_name, target_id, kind, body, created_at
        FROM chat_messages WHERE room_id = $1 ORDER BY id`, roomID)
    if err != nil {
        return nil, fmt.Errorf("log messages: %w", err)
    }
    defer rows.Close()
    var out []Message
    for rows.Next() {
        var m Message
        if err := rows.Scan(&m.ID, &m.RoomID, &m.AuthorID, &m.AuthorName, &m.TargetID, &m.Kind, &m.Body, &m.CreatedAt); err != nil {
            return nil, fmt.Errorf("scan message: %w", err)
        }
        out = append(out, m)
    }
    return out, rows.Err()
}

func TestConcurrentPollNoDup(t *testing.T) {
    ctx := context.Background()
    db := openPG(t)
    defer db.Close()
    s := chat.NewStore(db)
    rm, _ := s.CreateRoom(ctx, "race", "orchestrator")
    a, _ := s.AddParticipant(ctx, rm.ID, "a", "implementer")
    for i := 0; i < 20; i++ {
        _, _ = s.AddMessage(ctx, rm.ID, a.ID, nil, "message", "m")
    }
    type res struct{ n int }
    ch := make(chan res, 2)
    for i := 0; i < 2; i++ {
        go func() {
            msgs, _, err := s.Poll(ctx, rm.ID, a.ID)
            require.NoError(t, err)
            ch <- res{len(msgs)}
        }()
    }
    total := (<-ch).n + (<-ch).n
    require.Equal(t, 20, total) // each message delivered exactly once across both polls
}

func TestArchiveExpired(t *testing.T) {
    ctx := context.Background()
    db := openPG(t)
    defer db.Close()
    s := chat.NewStore(db)

    old, _ := s.CreateRoom(ctx, "old", "orchestrator")
    fresh, _ := s.CreateRoom(ctx, "fresh", "orchestrator")
    // backdate `old` past the TTL.
    _, err := db.Exec(`UPDATE chat_rooms SET created_at = NOW() - INTERVAL '48 hours' WHERE id = $1`, old.ID)
    require.NoError(t, err)

    n, err := s.ArchiveExpired(ctx, 24*time.Hour)
    require.NoError(t, err)
    require.Equal(t, 1, n)

    go1, _ := s.GetRoom(ctx, old.ID)
    require.Equal(t, "archived", go1.Status)
    go2, _ := s.GetRoom(ctx, fresh.ID)
    require.Equal(t, "active", go2.Status)

    // a system notice was written into the archived room.
    msgs, _ := s.Log(ctx, old.ID)
    require.Len(t, msgs, 1)
    require.Equal(t, "system", msgs[0].Kind)
    require.Nil(t, msgs[0].AuthorID)
}

import "time" // add to the existing import block

// ArchiveExpired archives every active room older than ttl (measured from
// created_at), writes a system notice into each, and returns the count
// archived. Runs in one transaction.
func (s *Store) ArchiveExpired(ctx context.Context, ttl time.Duration) (int, error) {
    tx, err := s.db.BeginTx(ctx, nil)
    if err != nil {
        return 0, fmt.Errorf("ttl begin tx: %w", err)
    }
    defer func() { _ = tx.Rollback() }()

    rows, err := tx.QueryContext(ctx, `
        UPDATE chat_rooms SET status = 'archived'
        WHERE status = 'active' AND created_at < NOW() - make_interval(secs => $1)
        RETURNING id`, ttl.Seconds())
    if err != nil {
        return 0, fmt.Errorf("ttl archive: %w", err)
    }
    var ids []string
    for rows.Next() {
        var id string
        if err := rows.Scan(&id); err != nil {
            rows.Close()
            return 0, fmt.Errorf("ttl scan: %w", err)
        }
        ids = append(ids, id)
    }
    rows.Close()
    if err := rows.Err(); err != nil {
        return 0, fmt.Errorf("ttl rows: %w", err)
    }

    for _, id := range ids {
        if _, err := tx.ExecContext(ctx, `
            INSERT INTO chat_messages (room_id, kind, body)
            VALUES ($1, 'system', 'room auto-archived after TTL')`, id); err != nil {
            return 0, fmt.Errorf("ttl notice: %w", err)
        }
    }

    if err := tx.Commit(); err != nil {
        return 0, fmt.Errorf("ttl commit: %w", err)
    }
    return len(ids), nil
}