initial agent setup

9 files changed, 312 insertions, 50 deletions

diff --git a/internal/db/db.go b/internal/db/db.go
index b97eb7108..e882106c4 100644
--- a/internal/db/db.go
+++ b/internal/db/db.go
@@ -24,15 +24,30 @@ func New(db DBTX) *Queries {
 func Prepare(ctx context.Context, db DBTX) (*Queries, error) {
 	q := Queries{db: db}
 	var err error
+	if q.createMessageStmt, err = db.PrepareContext(ctx, createMessage); err != nil {
+		return nil, fmt.Errorf("error preparing query CreateMessage: %w", err)
+	}
 	if q.createSessionStmt, err = db.PrepareContext(ctx, createSession); err != nil {
 		return nil, fmt.Errorf("error preparing query CreateSession: %w", err)
 	}
+	if q.deleteMessageStmt, err = db.PrepareContext(ctx, deleteMessage); err != nil {
+		return nil, fmt.Errorf("error preparing query DeleteMessage: %w", err)
+	}
 	if q.deleteSessionStmt, err = db.PrepareContext(ctx, deleteSession); err != nil {
 		return nil, fmt.Errorf("error preparing query DeleteSession: %w", err)
 	}
+	if q.deleteSessionMessagesStmt, err = db.PrepareContext(ctx, deleteSessionMessages); err != nil {
+		return nil, fmt.Errorf("error preparing query DeleteSessionMessages: %w", err)
+	}
+	if q.getMessageStmt, err = db.PrepareContext(ctx, getMessage); err != nil {
+		return nil, fmt.Errorf("error preparing query GetMessage: %w", err)
+	}
 	if q.getSessionByIDStmt, err = db.PrepareContext(ctx, getSessionByID); err != nil {
 		return nil, fmt.Errorf("error preparing query GetSessionByID: %w", err)
 	}
+	if q.listMessagesBySessionStmt, err = db.PrepareContext(ctx, listMessagesBySession); err != nil {
+		return nil, fmt.Errorf("error preparing query ListMessagesBySession: %w", err)
+	}
 	if q.listSessionsStmt, err = db.PrepareContext(ctx, listSessions); err != nil {
 		return nil, fmt.Errorf("error preparing query ListSessions: %w", err)
 	}
@@ -44,21 +59,46 @@ func Prepare(ctx context.Context, db DBTX) (*Queries, error) {
 
 func (q *Queries) Close() error {
 	var err error
+	if q.createMessageStmt != nil {
+		if cerr := q.createMessageStmt.Close(); cerr != nil {
+			err = fmt.Errorf("error closing createMessageStmt: %w", cerr)
+		}
+	}
 	if q.createSessionStmt != nil {
 		if cerr := q.createSessionStmt.Close(); cerr != nil {
 			err = fmt.Errorf("error closing createSessionStmt: %w", cerr)
 		}
 	}
+	if q.deleteMessageStmt != nil {
+		if cerr := q.deleteMessageStmt.Close(); cerr != nil {
+			err = fmt.Errorf("error closing deleteMessageStmt: %w", cerr)
+		}
+	}
 	if q.deleteSessionStmt != nil {
 		if cerr := q.deleteSessionStmt.Close(); cerr != nil {
 			err = fmt.Errorf("error closing deleteSessionStmt: %w", cerr)
 		}
 	}
+	if q.deleteSessionMessagesStmt != nil {
+		if cerr := q.deleteSessionMessagesStmt.Close(); cerr != nil {
+			err = fmt.Errorf("error closing deleteSessionMessagesStmt: %w", cerr)
+		}
+	}
+	if q.getMessageStmt != nil {
+		if cerr := q.getMessageStmt.Close(); cerr != nil {
+			err = fmt.Errorf("error closing getMessageStmt: %w", cerr)
+		}
+	}
 	if q.getSessionByIDStmt != nil {
 		if cerr := q.getSessionByIDStmt.Close(); cerr != nil {
 			err = fmt.Errorf("error closing getSessionByIDStmt: %w", cerr)
 		}
 	}
+	if q.listMessagesBySessionStmt != nil {
+		if cerr := q.listMessagesBySessionStmt.Close(); cerr != nil {
+			err = fmt.Errorf("error closing listMessagesBySessionStmt: %w", cerr)
+		}
+	}
 	if q.listSessionsStmt != nil {
 		if cerr := q.listSessionsStmt.Close(); cerr != nil {
 			err = fmt.Errorf("error closing listSessionsStmt: %w", cerr)
@@ -106,23 +146,33 @@ func (q *Queries) queryRow(ctx context.Context, stmt *sql.Stmt, query string, ar
 }
 
 type Queries struct {
-	db                 DBTX
-	tx                 *sql.Tx
-	createSessionStmt  *sql.Stmt
-	deleteSessionStmt  *sql.Stmt
-	getSessionByIDStmt *sql.Stmt
-	listSessionsStmt   *sql.Stmt
-	updateSessionStmt  *sql.Stmt
+	db                        DBTX
+	tx                        *sql.Tx
+	createMessageStmt         *sql.Stmt
+	createSessionStmt         *sql.Stmt
+	deleteMessageStmt         *sql.Stmt
+	deleteSessionStmt         *sql.Stmt
+	deleteSessionMessagesStmt *sql.Stmt
+	getMessageStmt            *sql.Stmt
+	getSessionByIDStmt        *sql.Stmt
+	listMessagesBySessionStmt *sql.Stmt
+	listSessionsStmt          *sql.Stmt
+	updateSessionStmt         *sql.Stmt
 }
 
 func (q *Queries) WithTx(tx *sql.Tx) *Queries {
 	return &Queries{
-		db:                 tx,
-		tx:                 tx,
-		createSessionStmt:  q.createSessionStmt,
-		deleteSessionStmt:  q.deleteSessionStmt,
-		getSessionByIDStmt: q.getSessionByIDStmt,
-		listSessionsStmt:   q.listSessionsStmt,
-		updateSessionStmt:  q.updateSessionStmt,
+		db:                        tx,
+		tx:                        tx,
+		createMessageStmt:         q.createMessageStmt,
+		createSessionStmt:         q.createSessionStmt,
+		deleteMessageStmt:         q.deleteMessageStmt,
+		deleteSessionStmt:         q.deleteSessionStmt,
+		deleteSessionMessagesStmt: q.deleteSessionMessagesStmt,
+		getMessageStmt:            q.getMessageStmt,
+		getSessionByIDStmt:        q.getSessionByIDStmt,
+		listMessagesBySessionStmt: q.listMessagesBySessionStmt,
+		listSessionsStmt:          q.listSessionsStmt,
+		updateSessionStmt:         q.updateSessionStmt,
 	}
 }
diff --git a/internal/db/messages.sql.go b/internal/db/messages.sql.go
new file mode 100644
index 000000000..d0f69458f
--- /dev/null
+++ b/internal/db/messages.sql.go
@@ -0,0 +1,117 @@
+// Code generated by sqlc. DO NOT EDIT.
+// versions:
+//   sqlc v1.27.0
+// source: messages.sql
+
+package db
+
+import (
+	"context"
+)
+
+const createMessage = `-- name: CreateMessage :one
+INSERT INTO messages (
+    id,
+    session_id,
+    message_data,
+    created_at,
+    updated_at
+) VALUES (
+    ?, ?, ?, strftime('%s', 'now'), strftime('%s', 'now')
+)
+RETURNING id, session_id, message_data, created_at, updated_at
+`
+
+type CreateMessageParams struct {
+	ID          string `json:"id"`
+	SessionID   string `json:"session_id"`
+	MessageData string `json:"message_data"`
+}
+
+func (q *Queries) CreateMessage(ctx context.Context, arg CreateMessageParams) (Message, error) {
+	row := q.queryRow(ctx, q.createMessageStmt, createMessage, arg.ID, arg.SessionID, arg.MessageData)
+	var i Message
+	err := row.Scan(
+		&i.ID,
+		&i.SessionID,
+		&i.MessageData,
+		&i.CreatedAt,
+		&i.UpdatedAt,
+	)
+	return i, err
+}
+
+const deleteMessage = `-- name: DeleteMessage :exec
+DELETE FROM messages
+WHERE id = ?
+`
+
+func (q *Queries) DeleteMessage(ctx context.Context, id string) error {
+	_, err := q.exec(ctx, q.deleteMessageStmt, deleteMessage, id)
+	return err
+}
+
+const deleteSessionMessages = `-- name: DeleteSessionMessages :exec
+DELETE FROM messages
+WHERE session_id = ?
+`
+
+func (q *Queries) DeleteSessionMessages(ctx context.Context, sessionID string) error {
+	_, err := q.exec(ctx, q.deleteSessionMessagesStmt, deleteSessionMessages, sessionID)
+	return err
+}
+
+const getMessage = `-- name: GetMessage :one
+SELECT id, session_id, message_data, created_at, updated_at
+FROM messages
+WHERE id = ? LIMIT 1
+`
+
+func (q *Queries) GetMessage(ctx context.Context, id string) (Message, error) {
+	row := q.queryRow(ctx, q.getMessageStmt, getMessage, id)
+	var i Message
+	err := row.Scan(
+		&i.ID,
+		&i.SessionID,
+		&i.MessageData,
+		&i.CreatedAt,
+		&i.UpdatedAt,
+	)
+	return i, err
+}
+
+const listMessagesBySession = `-- name: ListMessagesBySession :many
+SELECT id, session_id, message_data, created_at, updated_at
+FROM messages
+WHERE session_id = ?
+ORDER BY created_at ASC
+`
+
+func (q *Queries) ListMessagesBySession(ctx context.Context, sessionID string) ([]Message, error) {
+	rows, err := q.query(ctx, q.listMessagesBySessionStmt, listMessagesBySession, sessionID)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	items := []Message{}
+	for rows.Next() {
+		var i Message
+		if err := rows.Scan(
+			&i.ID,
+			&i.SessionID,
+			&i.MessageData,
+			&i.CreatedAt,
+			&i.UpdatedAt,
+		); err != nil {
+			return nil, err
+		}
+		items = append(items, i)
+	}
+	if err := rows.Close(); err != nil {
+		return nil, err
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	return items, nil
+}
diff --git a/internal/db/migrations/000001_initial.down.sql b/internal/db/migrations/000001_initial.down.sql
index 7324ddb6a..4f2712d86 100644
--- a/internal/db/migrations/000001_initial.down.sql
+++ b/internal/db/migrations/000001_initial.down.sql
@@ -1,4 +1,8 @@
--- sqlfluff:dialect:sqlite
 DROP TRIGGER IF EXISTS update_sessions_updated_at;
+DROP TRIGGER IF EXISTS update_messages_updated_at;
+
+DROP TRIGGER IF EXISTS update_session_message_count_on_delete;
+DROP TRIGGER IF EXISTS update_session_message_count_on_insert;
 
 DROP TABLE IF EXISTS sessions;
+DROP TABLE IF EXISTS messages;
diff --git a/internal/db/migrations/000001_initial.up.sql b/internal/db/migrations/000001_initial.up.sql
index 115b6145b..b87592bb5 100644
--- a/internal/db/migrations/000001_initial.up.sql
+++ b/internal/db/migrations/000001_initial.up.sql
@@ -1,9 +1,10 @@
--- sqlfluff:dialect:sqlite
+-- Sessions
 CREATE TABLE IF NOT EXISTS sessions (
     id TEXT PRIMARY KEY,
     title TEXT NOT NULL,
     message_count INTEGER NOT NULL DEFAULT 0 CHECK (message_count >= 0),
-    tokens INTEGER NOT NULL DEFAULT 0 CHECK (tokens >= 0),
+    prompt_tokens  INTEGER NOT NULL DEFAULT 0 CHECK (prompt_tokens >= 0),
+    completion_tokens  INTEGER NOT NULL DEFAULT 0 CHECK (completion_tokens>= 0),
     cost REAL NOT NULL DEFAULT 0.0 CHECK (cost >= 0.0),
     updated_at INTEGER NOT NULL,  -- Unix timestamp in milliseconds
     created_at INTEGER NOT NULL   -- Unix timestamp in milliseconds
@@ -15,3 +16,38 @@ BEGIN
 UPDATE sessions SET updated_at = strftime('%s', 'now')
 WHERE id = new.id;
 END;
+
+-- Messages
+CREATE TABLE IF NOT EXISTS messages (
+    id TEXT PRIMARY KEY,
+    session_id TEXT NOT NULL,
+    message_data TEXT NOT NULL,  -- JSON string of message content
+    created_at INTEGER NOT NULL,  -- Unix timestamp in milliseconds
+    updated_at INTEGER NOT NULL,  -- Unix timestamp in milliseconds
+    FOREIGN KEY (session_id) REFERENCES sessions (id) ON DELETE CASCADE
+);
+
+CREATE INDEX IF NOT EXISTS idx_messages_session_id ON messages (session_id);
+
+CREATE TRIGGER IF NOT EXISTS update_messages_updated_at
+AFTER UPDATE ON messages
+BEGIN
+UPDATE messages SET updated_at = strftime('%s', 'now')
+WHERE id = new.id;
+END;
+
+CREATE TRIGGER IF NOT EXISTS update_session_message_count_on_insert
+AFTER INSERT ON messages
+BEGIN
+UPDATE sessions SET
+    message_count = message_count + 1
+WHERE id = new.session_id;
+END;
+
+CREATE TRIGGER IF NOT EXISTS update_session_message_count_on_delete
+AFTER DELETE ON messages
+BEGIN
+UPDATE sessions SET
+    message_count = message_count - 1
+WHERE id = old.session_id;
+END;
diff --git a/internal/db/models.go b/internal/db/models.go
index 834aaf4dc..de9cd3cdb 100644
--- a/internal/db/models.go
+++ b/internal/db/models.go
@@ -4,12 +4,21 @@
 
 package db
 
+type Message struct {
+	ID          string `json:"id"`
+	SessionID   string `json:"session_id"`
+	MessageData string `json:"message_data"`
+	CreatedAt   int64  `json:"created_at"`
+	UpdatedAt   int64  `json:"updated_at"`
+}
+
 type Session struct {
-	ID           string  `json:"id"`
-	Title        string  `json:"title"`
-	MessageCount int64   `json:"message_count"`
-	Tokens       int64   `json:"tokens"`
-	Cost         float64 `json:"cost"`
-	UpdatedAt    int64   `json:"updated_at"`
-	CreatedAt    int64   `json:"created_at"`
+	ID               string  `json:"id"`
+	Title            string  `json:"title"`
+	MessageCount     int64   `json:"message_count"`
+	PromptTokens     int64   `json:"prompt_tokens"`
+	CompletionTokens int64   `json:"completion_tokens"`
+	Cost             float64 `json:"cost"`
+	UpdatedAt        int64   `json:"updated_at"`
+	CreatedAt        int64   `json:"created_at"`
 }
diff --git a/internal/db/querier.go b/internal/db/querier.go
index ac2374cf8..626d206a4 100644
--- a/internal/db/querier.go
+++ b/internal/db/querier.go
@@ -9,10 +9,14 @@ import (
 )
 
 type Querier interface {
-	// sqlfluff:dialect:sqlite
+	CreateMessage(ctx context.Context, arg CreateMessageParams) (Message, error)
 	CreateSession(ctx context.Context, arg CreateSessionParams) (Session, error)
+	DeleteMessage(ctx context.Context, id string) error
 	DeleteSession(ctx context.Context, id string) error
+	DeleteSessionMessages(ctx context.Context, sessionID string) error
+	GetMessage(ctx context.Context, id string) (Message, error)
 	GetSessionByID(ctx context.Context, id string) (Session, error)
+	ListMessagesBySession(ctx context.Context, sessionID string) ([]Message, error)
 	ListSessions(ctx context.Context) ([]Session, error)
 	UpdateSession(ctx context.Context, arg UpdateSessionParams) (Session, error)
 }
diff --git a/internal/db/sessions.sql.go b/internal/db/sessions.sql.go
index 3e0b6eea2..f3ee4ff42 100644
--- a/internal/db/sessions.sql.go
+++ b/internal/db/sessions.sql.go
@@ -14,7 +14,8 @@ INSERT INTO sessions (
     id,
     title,
     message_count,
-    tokens,
+    prompt_tokens,
+    completion_tokens,
     cost,
     updated_at,
     created_at
@@ -24,26 +25,28 @@ INSERT INTO sessions (
     ?,
     ?,
     ?,
+    ?,
     strftime('%s', 'now'),
     strftime('%s', 'now')
-) RETURNING id, title, message_count, tokens, cost, updated_at, created_at
+) RETURNING id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at
 `
 
 type CreateSessionParams struct {
-	ID           string  `json:"id"`
-	Title        string  `json:"title"`
-	MessageCount int64   `json:"message_count"`
-	Tokens       int64   `json:"tokens"`
-	Cost         float64 `json:"cost"`
+	ID               string  `json:"id"`
+	Title            string  `json:"title"`
+	MessageCount     int64   `json:"message_count"`
+	PromptTokens     int64   `json:"prompt_tokens"`
+	CompletionTokens int64   `json:"completion_tokens"`
+	Cost             float64 `json:"cost"`
 }
 
-// sqlfluff:dialect:sqlite
 func (q *Queries) CreateSession(ctx context.Context, arg CreateSessionParams) (Session, error) {
 	row := q.queryRow(ctx, q.createSessionStmt, createSession,
 		arg.ID,
 		arg.Title,
 		arg.MessageCount,
-		arg.Tokens,
+		arg.PromptTokens,
+		arg.CompletionTokens,
 		arg.Cost,
 	)
 	var i Session
@@ -51,7 +54,8 @@ func (q *Queries) CreateSession(ctx context.Context, arg CreateSessionParams) (S
 		&i.ID,
 		&i.Title,
 		&i.MessageCount,
-		&i.Tokens,
+		&i.PromptTokens,
+		&i.CompletionTokens,
 		&i.Cost,
 		&i.UpdatedAt,
 		&i.CreatedAt,
@@ -70,7 +74,7 @@ func (q *Queries) DeleteSession(ctx context.Context, id string) error {
 }
 
 const getSessionByID = `-- name: GetSessionByID :one
-SELECT id, title, message_count, tokens, cost, updated_at, created_at
+SELECT id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at
 FROM sessions
 WHERE id = ? LIMIT 1
 `
@@ -82,7 +86,8 @@ func (q *Queries) GetSessionByID(ctx context.Context, id string) (Session, error
 		&i.ID,
 		&i.Title,
 		&i.MessageCount,
-		&i.Tokens,
+		&i.PromptTokens,
+		&i.CompletionTokens,
 		&i.Cost,
 		&i.UpdatedAt,
 		&i.CreatedAt,
@@ -91,7 +96,7 @@ func (q *Queries) GetSessionByID(ctx context.Context, id string) (Session, error
 }
 
 const listSessions = `-- name: ListSessions :many
-SELECT id, title, message_count, tokens, cost, updated_at, created_at
+SELECT id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at
 FROM sessions
 ORDER BY created_at DESC
 `
@@ -109,7 +114,8 @@ func (q *Queries) ListSessions(ctx context.Context) ([]Session, error) {
 			&i.ID,
 			&i.Title,
 			&i.MessageCount,
-			&i.Tokens,
+			&i.PromptTokens,
+			&i.CompletionTokens,
 			&i.Cost,
 			&i.UpdatedAt,
 			&i.CreatedAt,
@@ -131,23 +137,26 @@ const updateSession = `-- name: UpdateSession :one
 UPDATE sessions
 SET
     title = ?,
-    tokens = ?,
+    prompt_tokens = ?,
+    completion_tokens = ?,
     cost = ?
 WHERE id = ?
-RETURNING id, title, message_count, tokens, cost, updated_at, created_at
+RETURNING id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at
 `
 
 type UpdateSessionParams struct {
-	Title  string  `json:"title"`
-	Tokens int64   `json:"tokens"`
-	Cost   float64 `json:"cost"`
-	ID     string  `json:"id"`
+	Title            string  `json:"title"`
+	PromptTokens     int64   `json:"prompt_tokens"`
+	CompletionTokens int64   `json:"completion_tokens"`
+	Cost             float64 `json:"cost"`
+	ID               string  `json:"id"`
 }
 
 func (q *Queries) UpdateSession(ctx context.Context, arg UpdateSessionParams) (Session, error) {
 	row := q.queryRow(ctx, q.updateSessionStmt, updateSession,
 		arg.Title,
-		arg.Tokens,
+		arg.PromptTokens,
+		arg.CompletionTokens,
 		arg.Cost,
 		arg.ID,
 	)
@@ -156,7 +165,8 @@ func (q *Queries) UpdateSession(ctx context.Context, arg UpdateSessionParams) (S
 		&i.ID,
 		&i.Title,
 		&i.MessageCount,
-		&i.Tokens,
+		&i.PromptTokens,
+		&i.CompletionTokens,
 		&i.Cost,
 		&i.UpdatedAt,
 		&i.CreatedAt,
diff --git a/internal/db/sql/messages.sql b/internal/db/sql/messages.sql
new file mode 100644
index 000000000..db5e192fc
--- /dev/null
+++ b/internal/db/sql/messages.sql
@@ -0,0 +1,30 @@
+-- name: GetMessage :one
+SELECT *
+FROM messages
+WHERE id = ? LIMIT 1;
+
+-- name: ListMessagesBySession :many
+SELECT *
+FROM messages
+WHERE session_id = ?
+ORDER BY created_at ASC;
+
+-- name: CreateMessage :one
+INSERT INTO messages (
+    id,
+    session_id,
+    message_data,
+    created_at,
+    updated_at
+) VALUES (
+    ?, ?, ?, strftime('%s', 'now'), strftime('%s', 'now')
+)
+RETURNING *;
+
+-- name: DeleteMessage :exec
+DELETE FROM messages
+WHERE id = ?;
+
+-- name: DeleteSessionMessages :exec
+DELETE FROM messages
+WHERE session_id = ?;
diff --git a/internal/db/sql/sessions.sql b/internal/db/sql/sessions.sql
index e90e5e328..2be8b7ccc 100644
--- a/internal/db/sql/sessions.sql
+++ b/internal/db/sql/sessions.sql
@@ -1,10 +1,10 @@
--- sqlfluff:dialect:sqlite
 -- name: CreateSession :one
 INSERT INTO sessions (
     id,
     title,
     message_count,
-    tokens,
+    prompt_tokens,
+    completion_tokens,
     cost,
     updated_at,
     created_at
@@ -14,6 +14,7 @@ INSERT INTO sessions (
     ?,
     ?,
     ?,
+    ?,
     strftime('%s', 'now'),
     strftime('%s', 'now')
 ) RETURNING *;
@@ -32,7 +33,8 @@ ORDER BY created_at DESC;
 UPDATE sessions
 SET
     title = ?,
-    tokens = ?,
+    prompt_tokens = ?,
+    completion_tokens = ?,
     cost = ?
 WHERE id = ?
 RETURNING *;