chore: refactoring

5 files changed, 31 insertions, 34 deletions

diff --git a/internal/llm/agent/agent.go b/internal/llm/agent/agent.go
index d1dba429f..4c838b731 100644
--- a/internal/llm/agent/agent.go
+++ b/internal/llm/agent/agent.go
@@ -670,7 +670,7 @@ func (a *agent) CompactSession(ctx context.Context, sessionID string, force bool
 		// Filter messages that were created after the last summarization
 		var newMessages []message.Message
 		for _, msg := range sessionMessages {
-			if msg.CreatedAt > session.SummarizedAt.UnixMilli() {
+			if msg.CreatedAt.After(session.SummarizedAt) {
 				newMessages = append(newMessages, msg)
 			}
 		}
diff --git a/internal/message/content.go b/internal/message/content.go
index 0a1ba714b..d7b76f8ce 100644
--- a/internal/message/content.go
+++ b/internal/message/content.go
@@ -106,21 +106,11 @@ func (ToolResult) isPart() {}
 
 type Finish struct {
 	Reason FinishReason `json:"reason"`
-	Time   int64        `json:"time"`
+	Time   time.Time    `json:"time"`
 }
 
 func (Finish) isPart() {}
 
-type Message struct {
-	ID        string
-	Role      MessageRole
-	SessionID string
-	Parts     []ContentPart
-	Model     models.ModelID
-	CreatedAt int64
-	UpdatedAt int64
-}
-
 func (m *Message) Content() *TextContent {
 	for _, part := range m.Parts {
 		if c, ok := part.(TextContent); ok {
@@ -318,7 +308,7 @@ func (m *Message) AddFinish(reason FinishReason) {
 			break
 		}
 	}
-	m.Parts = append(m.Parts, Finish{Reason: reason, Time: time.Now().UnixMilli()})
+	m.Parts = append(m.Parts, Finish{Reason: reason, Time: time.Now()})
 }
 
 func (m *Message) AddImageURL(url, detail string) {
diff --git a/internal/message/message.go b/internal/message/message.go
index 178547868..197c2f0f5 100644
--- a/internal/message/message.go
+++ b/internal/message/message.go
@@ -16,6 +16,16 @@ import (
 	"github.com/sst/opencode/internal/pubsub"
 )
 
+type Message struct {
+	ID        string
+	Role      MessageRole
+	SessionID string
+	Parts     []ContentPart
+	Model     models.ModelID
+	CreatedAt time.Time
+	UpdatedAt time.Time
+}
+
 const (
 	EventMessageCreated pubsub.EventType = "message_created"
 	EventMessageUpdated pubsub.EventType = "message_updated"
@@ -81,7 +91,7 @@ func (s *service) Create(ctx context.Context, sessionID string, params CreateMes
 		}
 	}
 	if params.Role == User && !isFinished {
-		params.Parts = append(params.Parts, Finish{Reason: FinishReasonEndTurn, Time: time.Now().UnixMilli()})
+		params.Parts = append(params.Parts, Finish{Reason: FinishReasonEndTurn, Time: time.Now()})
 	}
 
 	partsJSON, err := marshallParts(params.Parts)
@@ -126,9 +136,9 @@ func (s *service) Update(ctx context.Context, message Message) (Message, error)
 
 	var dbFinishedAt sql.NullInt64
 	finishPart := message.FinishPart()
-	if finishPart != nil && finishPart.Time > 0 {
+	if finishPart != nil && !finishPart.Time.IsZero() {
 		dbFinishedAt = sql.NullInt64{
-			Int64: finishPart.Time,
+			Int64: finishPart.Time.UnixMilli(),
 			Valid: true,
 		}
 	}
@@ -290,8 +300,8 @@ func (s *service) fromDBItem(item db.Message) (Message, error) {
 		Role:      MessageRole(item.Role),
 		Parts:     parts,
 		Model:     models.ModelID(item.Model.String),
-		CreatedAt: item.CreatedAt * 1000,
-		UpdatedAt: item.UpdatedAt * 1000,
+		CreatedAt: time.UnixMilli(item.CreatedAt),
+		UpdatedAt: time.UnixMilli(item.UpdatedAt),
 	}
 
 	return msg, nil
@@ -400,14 +410,7 @@ func marshallParts(parts []ContentPart) ([]byte, error) {
 func unmarshallParts(data []byte) ([]ContentPart, error) {
 	var rawMessages []json.RawMessage
 	if err := json.Unmarshal(data, &rawMessages); err != nil {
-		// Handle case where 'parts' might be a single object if not an array initially
-		// This was a fallback, if your DB always stores an array, this might not be needed.
-		var singleRawMessage json.RawMessage
-		if errSingle := json.Unmarshal(data, &singleRawMessage); errSingle == nil {
-			rawMessages = []json.RawMessage{singleRawMessage}
-		} else {
-			return nil, fmt.Errorf("failed to unmarshal parts data as array: %w. Data: %s", err, string(data))
-		}
+		return nil, fmt.Errorf("failed to unmarshal parts data as array: %w. Data: %s", err, string(data))
 	}
 
 	parts := make([]ContentPart, 0, len(rawMessages))
@@ -461,11 +464,15 @@ func unmarshallParts(data []byte) ([]ContentPart, error) {
 			}
 			parts = append(parts, p)
 		case finishType:
-			var p Finish
+			type dbFinish struct {
+				Reason FinishReason `json:"reason"`
+				Time   int64        `json:"time"`
+			}
+			var p dbFinish
 			if err := json.Unmarshal(wrapper.Data, &p); err != nil {
 				return nil, fmt.Errorf("unmarshal Finish: %w. Data: %s", err, string(wrapper.Data))
 			}
-			parts = append(parts, p)
+			parts = append(parts, Finish{Reason: FinishReason(p.Reason), Time: time.UnixMilli(p.Time)})
 		default:
 			slog.Warn("Unknown part type during unmarshalling, attempting to parse as TextContent", "type", wrapper.Type, "data", string(wrapper.Data))
 			// Fallback: if type is unknown or empty, try to parse data as TextContent directly
diff --git a/internal/session/session.go b/internal/session/session.go
index aea4fdaa5..b961959fe 100644
--- a/internal/session/session.go
+++ b/internal/session/session.go
@@ -157,7 +157,7 @@ func (s *service) Update(ctx context.Context, session Session) (Session, error)
 	if !session.SummarizedAt.IsZero() {
 		summarizedAt = sql.NullInt64{Int64: session.SummarizedAt.UnixMilli(), Valid: true}
 	}
-	
+
 	params := db.UpdateSessionParams{
 		ID:               session.ID,
 		Title:            session.Title,
@@ -208,7 +208,7 @@ func (s *service) fromDBItem(item db.Session) Session {
 	if item.SummarizedAt.Valid {
 		summarizedAt = time.UnixMilli(item.SummarizedAt.Int64)
 	}
-	
+
 	return Session{
 		ID:               item.ID,
 		ParentSessionID:  item.ParentSessionID.String,
@@ -219,8 +219,8 @@ func (s *service) fromDBItem(item db.Session) Session {
 		Cost:             item.Cost,
 		Summary:          item.Summary.String,
 		SummarizedAt:     summarizedAt,
-		CreatedAt:        time.UnixMilli(item.CreatedAt * 1000),
-		UpdatedAt:        time.UnixMilli(item.UpdatedAt * 1000),
+		CreatedAt:        time.UnixMilli(item.CreatedAt),
+		UpdatedAt:        time.UnixMilli(item.UpdatedAt),
 	}
 }
 
diff --git a/internal/tui/components/chat/message.go b/internal/tui/components/chat/message.go
index 66a8c3541..e7a6f7237 100644
--- a/internal/tui/components/chat/message.go
+++ b/internal/tui/components/chat/message.go
@@ -100,7 +100,7 @@ func renderUserMessage(msg message.Message, isFocused bool, width int, position
 
 	// Add timestamp info
 	info := []string{}
-	timestamp := time.UnixMilli(msg.CreatedAt).Local().Format("02 Jan 2006 03:04 PM")
+	timestamp := msg.CreatedAt.Local().Format("02 Jan 2006 03:04 PM")
 	username, _ := config.GetUsername()
 	info = append(info, baseStyle.
 		Width(width-1).
@@ -148,7 +148,7 @@ func renderAssistantMessage(
 	baseStyle := styles.BaseStyle()
 
 	// Always add timestamp info
-	timestamp := time.UnixMilli(msg.CreatedAt).Local().Format("02 Jan 2006 03:04 PM")
+	timestamp := msg.CreatedAt.Local().Format("02 Jan 2006 03:04 PM")
 	modelName := "Assistant"
 	if msg.Model != "" {
 		modelName = models.SupportedModels[msg.Model].Name