Part data model (#950)

4 files changed, 251 insertions, 100 deletions

diff --git a/packages/tui/internal/app/app.go b/packages/tui/internal/app/app.go
index 03eb361b7..fb8358d8c 100644
--- a/packages/tui/internal/app/app.go
+++ b/packages/tui/internal/app/app.go
@@ -16,11 +16,17 @@ import (
 	"github.com/sst/opencode/internal/commands"
 	"github.com/sst/opencode/internal/components/toast"
 	"github.com/sst/opencode/internal/config"
+	"github.com/sst/opencode/internal/id"
 	"github.com/sst/opencode/internal/styles"
 	"github.com/sst/opencode/internal/theme"
 	"github.com/sst/opencode/internal/util"
 )
 
+type Message struct {
+	Info  opencode.MessageUnion
+	Parts []opencode.PartUnion
+}
+
 type App struct {
 	Info             opencode.App
 	Modes            []opencode.Mode
@@ -35,7 +41,7 @@ type App struct {
 	Provider         *opencode.Provider
 	Model            *opencode.Model
 	Session          *opencode.Session
-	Messages         []opencode.MessageUnion
+	Messages         []Message
 	Commands         commands.CommandRegistry
 	InitialModel     *string
 	InitialPrompt    *string
@@ -158,7 +164,7 @@ func New(
 		ModeIndex:     modeIndex,
 		Mode:          mode,
 		Session:       &opencode.Session{},
-		Messages:      []opencode.MessageUnion{},
+		Messages:      []Message{},
 		Commands:      commands.LoadFromConfig(configInfo),
 		InitialModel:  initialModel,
 		InitialPrompt: initialPrompt,
@@ -351,7 +357,7 @@ func (a *App) IsBusy() bool {
 	}
 
 	lastMessage := a.Messages[len(a.Messages)-1]
-	if casted, ok := lastMessage.(opencode.AssistantMessage); ok {
+	if casted, ok := lastMessage.Info.(opencode.AssistantMessage); ok {
 		return casted.Time.Completed == 0
 	}
 	return false
@@ -452,54 +458,67 @@ func (a *App) SendChatMessage(
 		cmds = append(cmds, util.CmdHandler(SessionSelectedMsg(session)))
 	}
 
-	optimisticParts := []opencode.UserMessagePart{{
-		Type: opencode.UserMessagePartTypeText,
-		Text: text,
+	message := opencode.UserMessage{
+		ID:        id.Ascending(id.Message),
+		SessionID: a.Session.ID,
+		Role:      opencode.UserMessageRoleUser,
+		Time: opencode.UserMessageTime{
+			Created: float64(time.Now().UnixMilli()),
+		},
+	}
+
+	parts := []opencode.PartUnion{opencode.TextPart{
+		ID:        id.Ascending(id.Part),
+		MessageID: message.ID,
+		SessionID: a.Session.ID,
+		Type:      opencode.TextPartTypeText,
+		Text:      text,
 	}}
 	if len(attachments) > 0 {
 		for _, attachment := range attachments {
-			optimisticParts = append(optimisticParts, opencode.UserMessagePart{
-				Type:     opencode.UserMessagePartTypeFile,
-				Filename: attachment.Filename.Value,
-				Mime:     attachment.Mime.Value,
-				URL:      attachment.URL.Value,
+			parts = append(parts, opencode.FilePart{
+				ID:        id.Ascending(id.Part),
+				MessageID: message.ID,
+				SessionID: a.Session.ID,
+				Type:      opencode.FilePartTypeFile,
+				Filename:  attachment.Filename.Value,
+				Mime:      attachment.Mime.Value,
+				URL:       attachment.URL.Value,
 			})
 		}
 	}
 
-	optimisticMessage := opencode.UserMessage{
-		ID:        fmt.Sprintf("optimistic-%d", time.Now().UnixNano()),
-		Role:      opencode.UserMessageRoleUser,
-		Parts:     optimisticParts,
-		SessionID: a.Session.ID,
-		Time: opencode.UserMessageTime{
-			Created: float64(time.Now().Unix()),
-		},
-	}
-
-	a.Messages = append(a.Messages, optimisticMessage)
-	cmds = append(cmds, util.CmdHandler(OptimisticMessageAddedMsg{Message: optimisticMessage}))
+	a.Messages = append(a.Messages, Message{Info: message, Parts: parts})
+	cmds = append(cmds, util.CmdHandler(OptimisticMessageAddedMsg{Message: message}))
 
 	cmds = append(cmds, func() tea.Msg {
-		parts := []opencode.UserMessagePartUnionParam{
-			opencode.TextPartParam{
-				Type: opencode.F(opencode.TextPartTypeText),
-				Text: opencode.F(text),
-			},
-		}
-		if len(attachments) > 0 {
-			for _, attachment := range attachments {
-				parts = append(parts, opencode.FilePartParam{
-					Mime:     attachment.Mime,
-					Type:     attachment.Type,
-					URL:      attachment.URL,
-					Filename: attachment.Filename,
+		partsParam := []opencode.SessionChatParamsPartUnion{}
+		for _, part := range parts {
+			switch casted := part.(type) {
+			case opencode.TextPart:
+				partsParam = append(partsParam, opencode.TextPartParam{
+					ID:        opencode.F(casted.ID),
+					MessageID: opencode.F(casted.MessageID),
+					SessionID: opencode.F(casted.SessionID),
+					Type:      opencode.F(casted.Type),
+					Text:      opencode.F(casted.Text),
+				})
+			case opencode.FilePart:
+				partsParam = append(partsParam, opencode.FilePartParam{
+					ID:        opencode.F(casted.ID),
+					Mime:      opencode.F(casted.Mime),
+					MessageID: opencode.F(casted.MessageID),
+					SessionID: opencode.F(casted.SessionID),
+					Type:      opencode.F(casted.Type),
+					URL:       opencode.F(casted.URL),
+					Filename:  opencode.F(casted.Filename),
 				})
 			}
 		}
 
 		_, err := a.Client.Session.Chat(ctx, a.Session.ID, opencode.SessionChatParams{
-			Parts:      opencode.F(parts),
+			Parts:      opencode.F(partsParam),
+			MessageID:  opencode.F(message.ID),
 			ProviderID: opencode.F(a.Provider.ID),
 			ModelID:    opencode.F(a.Model.ID),
 			Mode:       opencode.F(a.Mode.Name),
@@ -557,15 +576,25 @@ func (a *App) DeleteSession(ctx context.Context, sessionID string) error {
 	return nil
 }
 
-func (a *App) ListMessages(ctx context.Context, sessionId string) ([]opencode.Message, error) {
+func (a *App) ListMessages(ctx context.Context, sessionId string) ([]Message, error) {
 	response, err := a.Client.Session.Messages(ctx, sessionId)
 	if err != nil {
 		return nil, err
 	}
 	if response == nil {
-		return []opencode.Message{}, nil
+		return []Message{}, nil
+	}
+	messages := []Message{}
+	for _, message := range *response {
+		msg := Message{
+			Info:  message.Info.AsUnion(),
+			Parts: []opencode.PartUnion{},
+		}
+		for _, part := range message.Parts {
+			msg.Parts = append(msg.Parts, part.AsUnion())
+		}
+		messages = append(messages, msg)
 	}
-	messages := *response
 	return messages, nil
 }
 
diff --git a/packages/tui/internal/components/chat/messages.go b/packages/tui/internal/components/chat/messages.go
index 7ecd9b21f..191432cc3 100644
--- a/packages/tui/internal/components/chat/messages.go
+++ b/packages/tui/internal/components/chat/messages.go
@@ -106,6 +106,13 @@ func (m *messagesComponent) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 				m.viewport.GotoBottom()
 			}
 		}
+	case opencode.EventListResponseEventMessagePartUpdated:
+		if msg.Properties.Part.SessionID == m.app.Session.ID {
+			m.renderView(m.width)
+			if m.tail {
+				m.viewport.GotoBottom()
+			}
+		}
 	}
 
 	viewport, cmd := m.viewport.Update(msg)
@@ -131,16 +138,16 @@ func (m *messagesComponent) renderView(width int) {
 		var content string
 		var cached bool
 
-		switch casted := message.(type) {
+		switch casted := message.Info.(type) {
 		case opencode.UserMessage:
 		userLoop:
-			for partIndex, part := range casted.Parts {
-				switch part := part.AsUnion().(type) {
+			for partIndex, part := range message.Parts {
+				switch part := part.(type) {
 				case opencode.TextPart:
-					remainingParts := casted.Parts[partIndex+1:]
+					remainingParts := message.Parts[partIndex+1:]
 					fileParts := make([]opencode.FilePart, 0)
 					for _, part := range remainingParts {
-						switch part := part.AsUnion().(type) {
+						switch part := part.(type) {
 						case opencode.FilePart:
 							fileParts = append(fileParts, part)
 						}
@@ -181,7 +188,7 @@ func (m *messagesComponent) renderView(width int) {
 					if !cached {
 						content = renderText(
 							m.app,
-							message,
+							message.Info,
 							part.Text,
 							m.app.Info.User,
 							m.showToolDetails,
@@ -202,12 +209,12 @@ func (m *messagesComponent) renderView(width int) {
 
 		case opencode.AssistantMessage:
 			hasTextPart := false
-			for partIndex, p := range casted.Parts {
-				switch part := p.AsUnion().(type) {
+			for partIndex, p := range message.Parts {
+				switch part := p.(type) {
 				case opencode.TextPart:
 					hasTextPart = true
 					finished := casted.Time.Completed > 0
-					remainingParts := casted.Parts[partIndex+1:]
+					remainingParts := message.Parts[partIndex+1:]
 					toolCallParts := make([]opencode.ToolPart, 0)
 
 					// sometimes tool calls happen without an assistant message
@@ -222,7 +229,7 @@ func (m *messagesComponent) renderView(width int) {
 						if !remaining {
 							break
 						}
-						switch part := part.AsUnion().(type) {
+						switch part := part.(type) {
 						case opencode.TextPart:
 							// we only want tool calls associated with the current text part.
 							// if we hit another text part, we're done.
@@ -238,13 +245,13 @@ func (m *messagesComponent) renderView(width int) {
 					}
 
 					if finished {
-						key := m.cache.GenerateKey(casted.ID, p.Text, width, m.showToolDetails, m.selectedPart == m.partCount)
+						key := m.cache.GenerateKey(casted.ID, part.Text, width, m.showToolDetails, m.selectedPart == m.partCount)
 						content, cached = m.cache.Get(key)
 						if !cached {
 							content = renderText(
 								m.app,
-								message,
-								p.Text,
+								message.Info,
+								part.Text,
 								casted.ModelID,
 								m.showToolDetails,
 								m.partCount == m.selectedPart,
@@ -257,8 +264,8 @@ func (m *messagesComponent) renderView(width int) {
 					} else {
 						content = renderText(
 							m.app,
-							message,
-							p.Text,
+							message.Info,
+							part.Text,
 							casted.ModelID,
 							m.showToolDetails,
 							m.partCount == m.selectedPart,
@@ -268,7 +275,7 @@ func (m *messagesComponent) renderView(width int) {
 						)
 					}
 					if content != "" {
-						m = m.updateSelected(content, p.Text)
+						m = m.updateSelected(content, part.Text)
 						blocks = append(blocks, content)
 					}
 				case opencode.ToolPart:
@@ -314,7 +321,7 @@ func (m *messagesComponent) renderView(width int) {
 		}
 
 		error := ""
-		if assistant, ok := message.(opencode.AssistantMessage); ok {
+		if assistant, ok := message.Info.(opencode.AssistantMessage); ok {
 			switch err := assistant.Error.AsUnion().(type) {
 			case nil:
 			case opencode.AssistantMessageErrorMessageOutputLengthError:
@@ -386,7 +393,7 @@ func (m *messagesComponent) header(width int) string {
 	contextWindow := m.app.Model.Limit.Context
 
 	for _, message := range m.app.Messages {
-		if assistant, ok := message.(opencode.AssistantMessage); ok {
+		if assistant, ok := message.Info.(opencode.AssistantMessage); ok {
 			cost += assistant.Cost
 			usage := assistant.Tokens
 			if usage.Output > 0 {
diff --git a/packages/tui/internal/id/id.go b/packages/tui/internal/id/id.go
new file mode 100644
index 000000000..0490b8f20
--- /dev/null
+++ b/packages/tui/internal/id/id.go
@@ -0,0 +1,96 @@
+package id
+
+import (
+	"crypto/rand"
+	"encoding/hex"
+	"fmt"
+	"strings"
+	"sync"
+	"time"
+)
+
+const (
+	PrefixSession = "ses"
+	PrefixMessage = "msg"
+	PrefixUser    = "usr"
+	PrefixPart    = "prt"
+)
+
+const length = 26
+
+var (
+	lastTimestamp int64
+	counter       int64
+	mu            sync.Mutex
+)
+
+type Prefix string
+
+const (
+	Session Prefix = PrefixSession
+	Message Prefix = PrefixMessage
+	User    Prefix = PrefixUser
+	Part    Prefix = PrefixPart
+)
+
+func ValidatePrefix(id string, prefix Prefix) bool {
+	return strings.HasPrefix(id, string(prefix))
+}
+
+func Ascending(prefix Prefix, given ...string) string {
+	return generateID(prefix, false, given...)
+}
+
+func Descending(prefix Prefix, given ...string) string {
+	return generateID(prefix, true, given...)
+}
+
+func generateID(prefix Prefix, descending bool, given ...string) string {
+	if len(given) > 0 && given[0] != "" {
+		if !strings.HasPrefix(given[0], string(prefix)) {
+			panic(fmt.Sprintf("ID %s does not start with %s", given[0], string(prefix)))
+		}
+		return given[0]
+	}
+	
+	return generateNewID(prefix, descending)
+}
+
+func randomBase62(length int) string {
+	const chars = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
+	result := make([]byte, length)
+	bytes := make([]byte, length)
+	rand.Read(bytes)
+	
+	for i := 0; i < length; i++ {
+		result[i] = chars[bytes[i]%62]
+	}
+	
+	return string(result)
+}
+
+func generateNewID(prefix Prefix, descending bool) string {
+	mu.Lock()
+	defer mu.Unlock()
+	
+	currentTimestamp := time.Now().UnixMilli()
+	
+	if currentTimestamp != lastTimestamp {
+		lastTimestamp = currentTimestamp
+		counter = 0
+	}
+	counter++
+	
+	now := uint64(currentTimestamp)*0x1000 + uint64(counter)
+	
+	if descending {
+		now = ^now
+	}
+	
+	timeBytes := make([]byte, 6)
+	for i := 0; i < 6; i++ {
+		timeBytes[i] = byte((now >> (40 - 8*i)) & 0xff)
+	}
+	
+	return string(prefix) + "_" + hex.EncodeToString(timeBytes) + randomBase62(length-12)
+}
\ No newline at end of file
diff --git a/packages/tui/internal/tui/tui.go b/packages/tui/internal/tui/tui.go
index 389dd64f1..0ebdd35ab 100644
--- a/packages/tui/internal/tui/tui.go
+++ b/packages/tui/internal/tui/tui.go
@@ -5,6 +5,7 @@ import (
 	"log/slog"
 	"os"
 	"os/exec"
+	"slices"
 	"strings"
 	"time"
 
@@ -364,55 +365,76 @@ func (a appModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	case opencode.EventListResponseEventSessionDeleted:
 		if a.app.Session != nil && msg.Properties.Info.ID == a.app.Session.ID {
 			a.app.Session = &opencode.Session{}
-			a.app.Messages = []opencode.MessageUnion{}
+			a.app.Messages = []app.Message{}
 		}
 		return a, toast.NewSuccessToast("Session deleted successfully")
 	case opencode.EventListResponseEventSessionUpdated:
 		if msg.Properties.Info.ID == a.app.Session.ID {
 			a.app.Session = &msg.Properties.Info
 		}
-	case opencode.EventListResponseEventMessageUpdated:
-		if msg.Properties.Info.SessionID == a.app.Session.ID {
-			exists := false
-			optimisticReplaced := false
-
-			// First check if this is replacing an optimistic message
-			if msg.Properties.Info.Role == opencode.MessageRoleUser {
-				// Look for optimistic messages to replace
-				for i, m := range a.app.Messages {
-					switch m := m.(type) {
-					case opencode.UserMessage:
-						if strings.HasPrefix(m.ID, "optimistic-") && m.Role == opencode.UserMessageRoleUser {
-							// Replace the optimistic message with the real one
-							a.app.Messages[i] = msg.Properties.Info.AsUnion()
-							exists = true
-							optimisticReplaced = true
-							break
-						}
+	case opencode.EventListResponseEventMessagePartUpdated:
+		slog.Info("message part updated", "message", msg.Properties.Part.MessageID, "part", msg.Properties.Part.ID)
+		if msg.Properties.Part.SessionID == a.app.Session.ID {
+			messageIndex := slices.IndexFunc(a.app.Messages, func(m app.Message) bool {
+				switch casted := m.Info.(type) {
+				case opencode.UserMessage:
+					return casted.ID == msg.Properties.Part.MessageID
+				case opencode.AssistantMessage:
+					return casted.ID == msg.Properties.Part.MessageID
+				}
+				return false
+			})
+			if messageIndex > -1 {
+				message := a.app.Messages[messageIndex]
+				partIndex := slices.IndexFunc(message.Parts, func(p opencode.PartUnion) bool {
+					switch casted := p.(type) {
+					case opencode.TextPart:
+						return casted.ID == msg.Properties.Part.ID
+					case opencode.FilePart:
+						return casted.ID == msg.Properties.Part.ID
+					case opencode.ToolPart:
+						return casted.ID == msg.Properties.Part.ID
+					case opencode.StepStartPart:
+						return casted.ID == msg.Properties.Part.ID
+					case opencode.StepFinishPart:
+						return casted.ID == msg.Properties.Part.ID
 					}
+					return false
+				})
+				if partIndex > -1 {
+					message.Parts[partIndex] = msg.Properties.Part.AsUnion()
 				}
+				if partIndex == -1 {
+					message.Parts = append(message.Parts, msg.Properties.Part.AsUnion())
+				}
+				a.app.Messages[messageIndex] = message
 			}
-
-			// If not replacing optimistic, check for existing message with same ID
-			if !optimisticReplaced {
-				for i, m := range a.app.Messages {
-					var id string
-					switch m := m.(type) {
-					case opencode.UserMessage:
-						id = m.ID
-					case opencode.AssistantMessage:
-						id = m.ID
-					}
-					if id == msg.Properties.Info.ID {
-						a.app.Messages[i] = msg.Properties.Info.AsUnion()
-						exists = true
-						break
-					}
+		}
+	case opencode.EventListResponseEventMessageUpdated:
+		if msg.Properties.Info.SessionID == a.app.Session.ID {
+			matchIndex := slices.IndexFunc(a.app.Messages, func(m app.Message) bool {
+				switch casted := m.Info.(type) {
+				case opencode.UserMessage:
+					return casted.ID == msg.Properties.Info.ID
+				case opencode.AssistantMessage:
+					return casted.ID == msg.Properties.Info.ID
+				}
+				return false
+			})
+
+			if matchIndex > -1 {
+				match := a.app.Messages[matchIndex]
+				a.app.Messages[matchIndex] = app.Message{
+					Info:  msg.Properties.Info.AsUnion(),
+					Parts: match.Parts,
 				}
 			}
 
-			if !exists {
-				a.app.Messages = append(a.app.Messages, msg.Properties.Info.AsUnion())
+			if matchIndex == -1 {
+				a.app.Messages = append(a.app.Messages, app.Message{
+					Info:  msg.Properties.Info.AsUnion(),
+					Parts: []opencode.PartUnion{},
+				})
 			}
 		}
 	case opencode.EventListResponseEventSessionError:
@@ -473,10 +495,7 @@ func (a appModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 			return a, toast.NewErrorToast("Failed to open session")
 		}
 		a.app.Session = msg
-		a.app.Messages = make([]opencode.MessageUnion, 0)
-		for _, message := range messages {
-			a.app.Messages = append(a.app.Messages, message.AsUnion())
-		}
+		a.app.Messages = messages
 		return a, util.CmdHandler(app.SessionLoadedMsg{})
 	case app.ModelSelectedMsg:
 		a.app.Provider = &msg.Provider
@@ -837,7 +856,7 @@ func (a appModel) executeCommand(command commands.Command) (tea.Model, tea.Cmd)
 			return a, nil
 		}
 		a.app.Session = &opencode.Session{}
-		a.app.Messages = []opencode.MessageUnion{}
+		a.app.Messages = []app.Message{}
 		cmds = append(cmds, util.CmdHandler(app.SessionClearedMsg{}))
 	case commands.SessionListCommand:
 		sessionDialog := dialog.NewSessionDialog(a.app)