chore: refactoring

7 files changed, 87 insertions, 75 deletions

diff --git a/internal/db/logs.sql.go b/internal/db/logs.sql.go
index d227b472a..aa89d7d22 100644
--- a/internal/db/logs.sql.go
+++ b/internal/db/logs.sql.go
@@ -10,7 +10,7 @@ import (
 	"database/sql"
 )
 
-const createLog = `-- name: CreateLog :exec
+const createLog = `-- name: CreateLog :one
 INSERT INTO logs (
     id,
     session_id,
@@ -27,7 +27,7 @@ INSERT INTO logs (
     ?,
     ?,
     ?
-)
+) RETURNING id, session_id, timestamp, level, message, attributes, created_at
 `
 
 type CreateLogParams struct {
@@ -40,8 +40,8 @@ type CreateLogParams struct {
 	CreatedAt  int64          `json:"created_at"`
 }
 
-func (q *Queries) CreateLog(ctx context.Context, arg CreateLogParams) error {
-	_, err := q.exec(ctx, q.createLogStmt, createLog,
+func (q *Queries) CreateLog(ctx context.Context, arg CreateLogParams) (Log, error) {
+	row := q.queryRow(ctx, q.createLogStmt, createLog,
 		arg.ID,
 		arg.SessionID,
 		arg.Timestamp,
@@ -50,7 +50,17 @@ func (q *Queries) CreateLog(ctx context.Context, arg CreateLogParams) error {
 		arg.Attributes,
 		arg.CreatedAt,
 	)
-	return err
+	var i Log
+	err := row.Scan(
+		&i.ID,
+		&i.SessionID,
+		&i.Timestamp,
+		&i.Level,
+		&i.Message,
+		&i.Attributes,
+		&i.CreatedAt,
+	)
+	return i, err
 }
 
 const listAllLogs = `-- name: ListAllLogs :many
diff --git a/internal/db/querier.go b/internal/db/querier.go
index 08ad41ec7..7b089f021 100644
--- a/internal/db/querier.go
+++ b/internal/db/querier.go
@@ -11,7 +11,7 @@ import (
 
 type Querier interface {
 	CreateFile(ctx context.Context, arg CreateFileParams) (File, error)
-	CreateLog(ctx context.Context, arg CreateLogParams) error
+	CreateLog(ctx context.Context, arg CreateLogParams) (Log, error)
 	CreateMessage(ctx context.Context, arg CreateMessageParams) (Message, error)
 	CreateSession(ctx context.Context, arg CreateSessionParams) (Session, error)
 	DeleteFile(ctx context.Context, id string) error
diff --git a/internal/db/sql/logs.sql b/internal/db/sql/logs.sql
index 1a0655f7e..58414f88a 100644
--- a/internal/db/sql/logs.sql
+++ b/internal/db/sql/logs.sql
@@ -1,4 +1,4 @@
--- name: CreateLog :exec
+-- name: CreateLog :one
 INSERT INTO logs (
     id,
     session_id,
@@ -15,7 +15,7 @@ INSERT INTO logs (
     ?,
     ?,
     ?
-);
+) RETURNING *;
 
 -- name: ListLogsBySession :many
 SELECT * FROM logs
@@ -25,4 +25,4 @@ ORDER BY timestamp ASC;
 -- name: ListAllLogs :many
 SELECT * FROM logs
 ORDER BY timestamp DESC
-LIMIT ?;
\ No newline at end of file
+LIMIT ?;
diff --git a/internal/logging/logging.go b/internal/logging/logging.go
index 13ccd0b21..b42f7caf7 100644
--- a/internal/logging/logging.go
+++ b/internal/logging/logging.go
@@ -22,11 +22,11 @@ import (
 type Log struct {
 	ID         string
 	SessionID  string
-	Timestamp  int64
+	Timestamp  time.Time
 	Level      string
 	Message    string
 	Attributes map[string]string
-	CreatedAt  int64
+	CreatedAt  time.Time
 }
 
 const (
@@ -36,7 +36,7 @@ const (
 type Service interface {
 	pubsub.Subscriber[Log]
 
-	Create(ctx context.Context, log Log) error
+	Create(ctx context.Context, timestamp time.Time, level, message string, attributes map[string]string, sessionID string) error
 	ListBySession(ctx context.Context, sessionID string) ([]Log, error)
 	ListAll(ctx context.Context, limit int) ([]Log, error)
 }
@@ -69,42 +69,35 @@ func GetService() Service {
 	return globalLoggingService
 }
 
-func (s *service) Create(ctx context.Context, log Log) error {
-	if log.ID == "" {
-		log.ID = uuid.New().String()
-	}
-	if log.Timestamp == 0 {
-		log.Timestamp = time.Now().UnixMilli()
-	}
-	if log.CreatedAt == 0 {
-		log.CreatedAt = time.Now().UnixMilli()
-	}
-	if log.Level == "" {
-		log.Level = "info"
+func (s *service) Create(ctx context.Context, timestamp time.Time, level, message string, attributes map[string]string, sessionID string) error {
+	if level == "" {
+		level = "info"
 	}
 
 	var attributesJSON sql.NullString
-	if len(log.Attributes) > 0 {
-		attributesBytes, err := json.Marshal(log.Attributes)
+	if len(attributes) > 0 {
+		attributesBytes, err := json.Marshal(attributes)
 		if err != nil {
 			return fmt.Errorf("failed to marshal log attributes: %w", err)
 		}
 		attributesJSON = sql.NullString{String: string(attributesBytes), Valid: true}
 	}
 
-	err := s.db.CreateLog(ctx, db.CreateLogParams{
-		ID:         log.ID,
-		SessionID:  sql.NullString{String: log.SessionID, Valid: log.SessionID != ""},
-		Timestamp:  log.Timestamp,
-		Level:      log.Level,
-		Message:    log.Message,
+	dbLog, err := s.db.CreateLog(ctx, db.CreateLogParams{
+		ID:         uuid.New().String(),
+		SessionID:  sql.NullString{String: sessionID, Valid: sessionID != ""},
+		Timestamp:  timestamp.UnixMilli(),
+		Level:      level,
+		Message:    message,
 		Attributes: attributesJSON,
-		CreatedAt:  log.CreatedAt,
+		CreatedAt:  time.Now().UnixMilli(),
 	})
+
 	if err != nil {
 		return fmt.Errorf("db.CreateLog: %w", err)
 	}
 
+	log := s.fromDBItem(dbLog)
 	s.broker.Publish(EventLogCreated, log)
 	return nil
 }
@@ -114,7 +107,12 @@ func (s *service) ListBySession(ctx context.Context, sessionID string) ([]Log, e
 	if err != nil {
 		return nil, fmt.Errorf("db.ListLogsBySession: %w", err)
 	}
-	return s.fromDBItems(dbLogs)
+
+	logs := make([]Log, len(dbLogs))
+	for i, dbSess := range dbLogs {
+		logs[i] = s.fromDBItem(dbSess)
+	}
+	return logs, nil
 }
 
 func (s *service) ListAll(ctx context.Context, limit int) ([]Log, error) {
@@ -122,39 +120,41 @@ func (s *service) ListAll(ctx context.Context, limit int) ([]Log, error) {
 	if err != nil {
 		return nil, fmt.Errorf("db.ListAllLogs: %w", err)
 	}
-	return s.fromDBItems(dbLogs)
+	logs := make([]Log, len(dbLogs))
+	for i, dbSess := range dbLogs {
+		logs[i] = s.fromDBItem(dbSess)
+	}
+	return logs, nil
 }
 
 func (s *service) Subscribe(ctx context.Context) <-chan pubsub.Event[Log] {
 	return s.broker.Subscribe(ctx)
 }
 
-func (s *service) fromDBItems(items []db.Log) ([]Log, error) {
-	logs := make([]Log, len(items))
-	for i, item := range items {
-		log := Log{
-			ID:        item.ID,
-			SessionID: item.SessionID.String,
-			Timestamp: item.Timestamp * 1000,
-			Level:     item.Level,
-			Message:   item.Message,
-			CreatedAt: item.CreatedAt * 1000,
-		}
-		if item.Attributes.Valid && item.Attributes.String != "" {
-			if err := json.Unmarshal([]byte(item.Attributes.String), &log.Attributes); err != nil {
-				slog.Error("Failed to unmarshal log attributes", "log_id", item.ID, "error", err)
-				log.Attributes = make(map[string]string)
-			}
-		} else {
+func (s *service) fromDBItem(item db.Log) Log {
+	log := Log{
+		ID:        item.ID,
+		SessionID: item.SessionID.String,
+		Timestamp: time.UnixMilli(item.Timestamp),
+		Level:     item.Level,
+		Message:   item.Message,
+		CreatedAt: time.UnixMilli(item.CreatedAt),
+	}
+
+	if item.Attributes.Valid && item.Attributes.String != "" {
+		if err := json.Unmarshal([]byte(item.Attributes.String), &log.Attributes); err != nil {
+			slog.Error("Failed to unmarshal log attributes", "log_id", item.ID, "error", err)
 			log.Attributes = make(map[string]string)
 		}
-		logs[i] = log
+	} else {
+		log.Attributes = make(map[string]string)
 	}
-	return logs, nil
+
+	return log
 }
 
-func Create(ctx context.Context, log Log) error {
-	return GetService().Create(ctx, log)
+func Create(ctx context.Context, timestamp time.Time, level, message string, attributes map[string]string, sessionID string) error {
+	return GetService().Create(ctx, timestamp, level, message, attributes, sessionID)
 }
 
 func ListBySession(ctx context.Context, sessionID string) ([]Log, error) {
@@ -175,9 +175,13 @@ func (sw *slogWriter) Write(p []byte) (n int, err error) {
 	// Example: time=2024-05-09T12:34:56.789-05:00 level=INFO msg="User request" session=xyz foo=bar
 	d := logfmt.NewDecoder(bytes.NewReader(p))
 	for d.ScanRecord() {
-		logEntry := Log{
-			Attributes: make(map[string]string),
-		}
+		var timestamp time.Time
+		var level string
+		var message string
+		var sessionID string
+		var attributes map[string]string
+
+		attributes = make(map[string]string)
 		hasTimestamp := false
 
 		for d.ScanKeyval() {
@@ -191,45 +195,44 @@ func (sw *slogWriter) Write(p []byte) (n int, err error) {
 					parsedTime, timeErr = time.Parse(time.RFC3339, value)
 					if timeErr != nil {
 						slog.Error("Failed to parse time in slog writer", "value", value, "error", timeErr)
-						logEntry.Timestamp = time.Now().UnixMilli()
+						timestamp = time.Now()
 						hasTimestamp = true
 						continue
 					}
 				}
-				logEntry.Timestamp = parsedTime.UnixMilli()
+				timestamp = parsedTime
 				hasTimestamp = true
 			case "level":
-				logEntry.Level = strings.ToLower(value)
+				level = strings.ToLower(value)
 			case "msg", "message":
-				logEntry.Message = value
-			case "session_id", "session", "sid":
-				logEntry.SessionID = value
+				message = value
+			case "session_id":
+				sessionID = value
 			default:
-				logEntry.Attributes[key] = value
+				attributes[key] = value
 			}
 		}
-
 		if d.Err() != nil {
 			return len(p), fmt.Errorf("logfmt.ScanRecord: %w", d.Err())
 		}
 
 		if !hasTimestamp {
-			logEntry.Timestamp = time.Now().UnixMilli()
+			timestamp = time.Now()
 		}
 
 		// Create log entry via the service (non-blocking or handle error appropriately)
 		// Using context.Background() as this is a low-level logging write.
-		go func(le Log) { // Run in a goroutine to avoid blocking slog
+		go func(timestamp time.Time, level, message string, attributes map[string]string, sessionID string) { // Run in a goroutine to avoid blocking slog
 			if globalLoggingService == nil {
 				// If the logging service is not initialized, log the message to stderr
 				// fmt.Fprintf(os.Stderr, "ERROR [logging.slogWriter]: logging service not initialized\n")
 				return
 			}
-			if err := Create(context.Background(), le); err != nil {
+			if err := Create(context.Background(), timestamp, level, message, attributes, sessionID); err != nil {
 				// Log internal error using a more primitive logger to avoid loops
 				fmt.Fprintf(os.Stderr, "ERROR [logging.slogWriter]: failed to persist log: %v\n", err)
 			}
-		}(logEntry)
+		}(timestamp, level, message, attributes, sessionID)
 	}
 	if d.Err() != nil {
 		return len(p), fmt.Errorf("logfmt.ScanRecord final: %w", d.Err())
diff --git a/internal/message/content.go b/internal/message/content.go
index 4509460af..0a1ba714b 100644
--- a/internal/message/content.go
+++ b/internal/message/content.go
@@ -318,7 +318,7 @@ func (m *Message) AddFinish(reason FinishReason) {
 			break
 		}
 	}
-	m.Parts = append(m.Parts, Finish{Reason: reason, Time: time.Now().Unix()})
+	m.Parts = append(m.Parts, Finish{Reason: reason, Time: time.Now().UnixMilli()})
 }
 
 func (m *Message) AddImageURL(url, detail string) {
diff --git a/internal/tui/components/logs/details.go b/internal/tui/components/logs/details.go
index 5ad59b47a..631bb97b3 100644
--- a/internal/tui/components/logs/details.go
+++ b/internal/tui/components/logs/details.go
@@ -66,7 +66,7 @@ func (i *detailCmp) updateContent() {
 	levelStyle := getLevelStyle(i.currentLog.Level)
 
 	// Format timestamp
-	timeStr := time.UnixMilli(i.currentLog.Timestamp).Format(time.RFC3339)
+	timeStr := i.currentLog.Timestamp.Format(time.RFC3339)
 
 	header := lipgloss.JoinHorizontal(
 		lipgloss.Center,
diff --git a/internal/tui/components/logs/table.go b/internal/tui/components/logs/table.go
index 0c0070f42..dbfc9c9f9 100644
--- a/internal/tui/components/logs/table.go
+++ b/internal/tui/components/logs/table.go
@@ -2,7 +2,6 @@ package logs
 
 import (
 	"context"
-	"time"
 
 	"github.com/charmbracelet/bubbles/key"
 	"github.com/charmbracelet/bubbles/table"
@@ -161,7 +160,7 @@ func (i *tableCmp) updateRows() {
 
 	for _, log := range i.logs {
 		// Format timestamp as time
-		timeStr := time.UnixMilli(log.Timestamp).Format("15:04:05")
+		timeStr := log.Timestamp.Format("15:04:05")
 
 		// Include ID as hidden first column for selection
 		row := table.Row{