fix(tui): attachment source is not stored when using message from message history (#1542)

2 files changed, 85 insertions, 0 deletions

diff --git a/packages/tui/internal/app/state.go b/packages/tui/internal/app/state.go
index e797d4220..bf2a602b3 100644
--- a/packages/tui/internal/app/state.go
+++ b/packages/tui/internal/app/state.go
@@ -120,5 +120,13 @@ func LoadState(filePath string) (*State, error) {
 		}
 		return nil, fmt.Errorf("failed to decode TOML from file %s: %w", filePath, err)
 	}
+
+	// Restore attachment sources types that were deserialized as map[string]any
+	for _, prompt := range state.MessageHistory {
+		for _, att := range prompt.Attachments {
+			att.RestoreSourceType()
+		}
+	}
+
 	return &state, nil
 }
diff --git a/packages/tui/internal/attachment/attachment.go b/packages/tui/internal/attachment/attachment.go
index 26bb91d49..038209ae8 100644
--- a/packages/tui/internal/attachment/attachment.go
+++ b/packages/tui/internal/attachment/attachment.go
@@ -75,3 +75,80 @@ func (a *Attachment) GetSymbolSource() (*SymbolSource, bool) {
 	ss, ok := a.Source.(*SymbolSource)
 	return ss, ok
 }
+
+// FromMap creates a TextSource from a map[string]any
+func (ts *TextSource) FromMap(sourceMap map[string]any) {
+	if value, ok := sourceMap["value"].(string); ok {
+		ts.Value = value
+	}
+}
+
+// FromMap creates a FileSource from a map[string]any
+func (fs *FileSource) FromMap(sourceMap map[string]any) {
+	if path, ok := sourceMap["path"].(string); ok {
+		fs.Path = path
+	}
+	if mime, ok := sourceMap["mime"].(string); ok {
+		fs.Mime = mime
+	}
+	if data, ok := sourceMap["data"].([]byte); ok {
+		fs.Data = data
+	}
+}
+
+// FromMap creates a SymbolSource from a map[string]any
+func (ss *SymbolSource) FromMap(sourceMap map[string]any) {
+	if path, ok := sourceMap["path"].(string); ok {
+		ss.Path = path
+	}
+	if name, ok := sourceMap["name"].(string); ok {
+		ss.Name = name
+	}
+	if kind, ok := sourceMap["kind"].(int); ok {
+		ss.Kind = kind
+	}
+	if rangeMap, ok := sourceMap["range"].(map[string]any); ok {
+		ss.Range = SymbolRange{}
+		if startMap, ok := rangeMap["start"].(map[string]any); ok {
+			if line, ok := startMap["line"].(int); ok {
+				ss.Range.Start.Line = line
+			}
+			if char, ok := startMap["char"].(int); ok {
+				ss.Range.Start.Char = char
+			}
+		}
+		if endMap, ok := rangeMap["end"].(map[string]any); ok {
+			if line, ok := endMap["line"].(int); ok {
+				ss.Range.End.Line = line
+			}
+			if char, ok := endMap["char"].(int); ok {
+				ss.Range.End.Char = char
+			}
+		}
+	}
+}
+
+// RestoreSourceType converts a map[string]any source back to the proper type
+func (a *Attachment) RestoreSourceType() {
+	if a.Source == nil {
+		return
+	}
+
+	// Check if Source is a map[string]any
+	if sourceMap, ok := a.Source.(map[string]any); ok {
+		switch a.Type {
+		case "text":
+			ts := &TextSource{}
+			ts.FromMap(sourceMap)
+			a.Source = ts
+		case "file":
+			fs := &FileSource{}
+			fs.FromMap(sourceMap)
+			a.Source = fs
+		case "symbol":
+			ss := &SymbolSource{}
+			ss.FromMap(sourceMap)
+			a.Source = ss
+		}
+	}
+}