feat: update memo when its source message is edited

README.md

@@ -115,4 +115,5 @@ Or you can start the service with Docker Compose. This can be combined with the
 - `/start <access_token>`: Start the bot with your Memos access token.
 - Send text messages: Save the message content as a memo.
 - Send files (photos, documents): Save the files as resources in a memo.
+- Edit a sent message: Update the memo that message originally created. Edits to messages Memogram never captured (for example, messages sent before this feature was available) are ignored.
 - `/search <words>`: Search for the memos.

memogram.go

@@ -167,6 +167,10 @@ func (s *Service) handler(ctx context.Context, b *bot.Bot, m *models.Update) {
 		fmt.Println("Service or config is nil")
 		return
 	}
+	if m != nil && m.EditedMessage != nil {
+		s.editedMessageHandler(ctx, b, m)
+		return
+	}
 	if m == nil || m.Message == nil || m.Message.From == nil {
 		s.sendError(b, 0, errors.New("invalid message structure: missing required fields"))
 		return
@@ -272,6 +276,7 @@ func (s *Service) handler(ctx context.Context, b *bot.Bot, m *models.Update) {
 		})
 		return
 	}
+	s.store.SetMemoForMessage(message.Chat.ID, int64(message.ID), memo.Name)
 
 	if message.Document != nil {
 		s.processFileMessage(ctx, authClient, b, m, message.Document.FileID, memo)
@@ -341,6 +346,54 @@ func (s *Service) startHandler(ctx context.Context, b *bot.Bot, m *models.Update
 	})
 }
 
+// Edits of messages with no known memo (sent before this feature existed) are ignored.
+func (s *Service) editedMessageHandler(ctx context.Context, b *bot.Bot, m *models.Update) {
+	message := m.EditedMessage
+	if message == nil || message.From == nil || message.Chat.ID == 0 {
+		return
+	}
+	if !s.isUserAllowed(message.From.Username) {
+		return
+	}
+
+	userID := message.From.ID
+	accessToken, ok := s.store.GetUserAccessToken(userID)
+	if !ok {
+		return
+	}
+
+	memoName, ok := s.store.GetMemoForMessage(message.Chat.ID, int64(message.ID))
+	if !ok {
+		return
+	}
+
+	content := message.Text
+	contentEntities := message.Entities
+	if message.Caption != "" {
+		content = message.Caption
+		contentEntities = message.CaptionEntities
+	}
+	if len(contentEntities) > 0 {
+		content = formatContent(content, contentEntities)
+	}
+	if content == "" {
+		return
+	}
+
+	authClient := s.client.NewAuthenticatedClient(accessToken)
+	_, err := authClient.MemoService.UpdateMemo(ctx, connect.NewRequest(&v1pb.UpdateMemoRequest{
+		Memo: &v1pb.Memo{
+			Name:    memoName,
+			Content: content,
+		},
+		UpdateMask: &fieldmaskpb.FieldMask{Paths: []string{"content"}},
+	}))
+	if err != nil {
+		slog.Error("failed to update memo on edit", slog.Any("err", err))
+		s.sendError(b, message.Chat.ID, fmt.Errorf("failed to update memo: %w", err))
+	}
+}
+
 func (s *Service) keyboard(memo *v1pb.Memo) *models.InlineKeyboardMarkup {
 	// add inline keyboard to edit memo's visibility or pinned status.
 	return &models.InlineKeyboardMarkup{

store/message.go

@@ -0,0 +1,153 @@
+package store
+
+import (
+	"bufio"
+	"fmt"
+	"log/slog"
+	"os"
+	"path/filepath"
+	"sort"
+	"strconv"
+	"strings"
+)
+
+// messageKey scopes a Telegram message ID to its chat: message IDs are only
+// unique within a chat, not globally.
+type messageKey struct {
+	chatID    int64
+	messageID int64
+}
+
+func (s *Store) GetMemoForMessage(chatID, messageID int64) (string, bool) {
+	memoName, ok := s.messageMemoCache.Load(messageKey{chatID: chatID, messageID: messageID})
+	if !ok {
+		return "", false
+	}
+	return memoName.(string), true
+}
+
+func (s *Store) SetMemoForMessage(chatID, messageID int64, memoName string) {
+	s.saveMu.Lock()
+	defer s.saveMu.Unlock()
+
+	s.messageMemoCache.Store(messageKey{chatID: chatID, messageID: messageID}, memoName)
+	if err := s.saveMessageMemoMapToFile(); err != nil {
+		slog.Error("failed to save message memo map to file", "error", err)
+	}
+}
+
+func (s *Store) messageMemoMapDataPath() string {
+	return s.Data + ".messages"
+}
+
+func (s *Store) saveMessageMemoMapToFile() error {
+	entries := s.snapshotMessageMemoMap()
+	dataPath := s.messageMemoMapDataPath()
+	dataDir := filepath.Dir(dataPath)
+	tmpFile, err := os.CreateTemp(dataDir, "memogram-messages-*.tmp")
+	if err != nil {
+		return fmt.Errorf("create temp file: %w", err)
+	}
+	defer os.Remove(tmpFile.Name())
+
+	writer := bufio.NewWriter(tmpFile)
+	for _, entry := range entries {
+		if _, err := fmt.Fprintf(writer, "%d:%d:%s\n", entry.key.chatID, entry.key.messageID, entry.memoName); err != nil {
+			tmpFile.Close()
+			return fmt.Errorf("write data file: %w", err)
+		}
+	}
+	if err := writer.Flush(); err != nil {
+		tmpFile.Close()
+		return fmt.Errorf("flush data file: %w", err)
+	}
+	if err := tmpFile.Sync(); err != nil {
+		tmpFile.Close()
+		return fmt.Errorf("sync data file: %w", err)
+	}
+	if err := tmpFile.Close(); err != nil {
+		return fmt.Errorf("close data file: %w", err)
+	}
+
+	if err := os.Rename(tmpFile.Name(), dataPath); err != nil {
+		return fmt.Errorf("replace data file: %w", err)
+	}
+	return nil
+}
+
+func (s *Store) loadMessageMemoMapFromFile() error {
+	dataPath := s.messageMemoMapDataPath()
+	if _, err := os.Stat(dataPath); os.IsNotExist(err) {
+		file, err := os.Create(dataPath)
+		if err != nil {
+			return err
+		}
+		defer file.Close()
+	}
+
+	file, err := os.Open(dataPath)
+	if err != nil {
+		return err
+	}
+	defer file.Close()
+
+	scanner := bufio.NewScanner(file)
+	for scanner.Scan() {
+		line := strings.TrimSpace(scanner.Text())
+		if line == "" || strings.HasPrefix(line, "#") {
+			continue
+		}
+		key, memoName := parseMessageMemoLine(line)
+		if key.messageID == 0 || memoName == "" {
+			continue
+		}
+		s.messageMemoCache.Store(key, memoName)
+	}
+	return scanner.Err()
+}
+
+func parseMessageMemoLine(line string) (messageKey, string) {
+	parts := strings.SplitN(line, ":", 3)
+	if len(parts) != 3 {
+		return messageKey{}, ""
+	}
+	chatID, err := strconv.ParseInt(parts[0], 10, 64)
+	if err != nil {
+		return messageKey{}, ""
+	}
+	messageID, err := strconv.ParseInt(parts[1], 10, 64)
+	if err != nil {
+		return messageKey{}, ""
+	}
+	return messageKey{chatID: chatID, messageID: messageID}, parts[2]
+}
+
+type messageMemoEntry struct {
+	key      messageKey
+	memoName string
+}
+
+func (s *Store) snapshotMessageMemoMap() []messageMemoEntry {
+	entries := make([]messageMemoEntry, 0)
+	s.messageMemoCache.Range(func(key, value interface{}) bool {
+		messageKey, ok := key.(messageKey)
+		if !ok {
+			return true
+		}
+		memoName, ok := value.(string)
+		if !ok {
+			return true
+		}
+		entries = append(entries, messageMemoEntry{key: messageKey, memoName: memoName})
+		return true
+	})
+
+	sort.Slice(entries, func(i, j int) bool {
+		if entries[i].key.chatID != entries[j].key.chatID {
+			return entries[i].key.chatID < entries[j].key.chatID
+		}
+		return entries[i].key.messageID < entries[j].key.messageID
+	})
+
+	return entries
+}

store/message_test.go

@@ -0,0 +1,77 @@
+package store
+
+import (
+	"path/filepath"
+	"testing"
+)
+
+func TestSaveAndLoadMessageMemoMap(t *testing.T) {
+	dataPath := filepath.Join(t.TempDir(), "data.txt")
+
+	store := NewStore(dataPath)
+	if err := store.Init(); err != nil {
+		t.Fatalf("init store: %v", err)
+	}
+
+	store.SetMemoForMessage(1, 101, "memos/abc123")
+	store.SetMemoForMessage(1, 202, "memos/def456")
+
+	reloaded := NewStore(dataPath)
+	if err := reloaded.Init(); err != nil {
+		t.Fatalf("init reloaded store: %v", err)
+	}
+
+	memoName, ok := reloaded.GetMemoForMessage(1, 101)
+	if !ok || memoName != "memos/abc123" {
+		t.Fatalf("expected memos/abc123 for message 101, got %q", memoName)
+	}
+
+	memoName, ok = reloaded.GetMemoForMessage(1, 202)
+	if !ok || memoName != "memos/def456" {
+		t.Fatalf("expected memos/def456 for message 202, got %q", memoName)
+	}
+}
+
+func TestGetMemoForMessageMissing(t *testing.T) {
+	dataPath := filepath.Join(t.TempDir(), "data.txt")
+
+	store := NewStore(dataPath)
+	if err := store.Init(); err != nil {
+		t.Fatalf("init store: %v", err)
+	}
+
+	if _, ok := store.GetMemoForMessage(1, 999); ok {
+		t.Fatalf("expected no mapping for message 999")
+	}
+}
+
+func TestMessageMemoMapScopedByChat(t *testing.T) {
+	dataPath := filepath.Join(t.TempDir(), "data.txt")
+
+	store := NewStore(dataPath)
+	if err := store.Init(); err != nil {
+		t.Fatalf("init store: %v", err)
+	}
+
+	// Same message ID in two different chats must not collide: Telegram
+	// message IDs are only unique within a chat.
+	store.SetMemoForMessage(1, 5, "memos/chat-one")
+	store.SetMemoForMessage(2, 5, "memos/chat-two")
+
+	// Reload from disk to also cover the chatID:messageID:memoName persisted format,
+	// not just the in-memory cache.
+	reloaded := NewStore(dataPath)
+	if err := reloaded.Init(); err != nil {
+		t.Fatalf("init reloaded store: %v", err)
+	}
+
+	memoName, ok := reloaded.GetMemoForMessage(1, 5)
+	if !ok || memoName != "memos/chat-one" {
+		t.Fatalf("expected memos/chat-one for chat 1 message 5, got %q", memoName)
+	}
+
+	memoName, ok = reloaded.GetMemoForMessage(2, 5)
+	if !ok || memoName != "memos/chat-two" {
+		t.Fatalf("expected memos/chat-two for chat 2 message 5, got %q", memoName)
+	}
+}

store/store.go

@@ -8,21 +8,30 @@ import (
 type Store struct {
 	Data string
 
+	// saveMu serializes cache-mutation-then-file-save sequences so an older
+	// snapshot can't win the rename race and discard a newer entry.
+	saveMu sync.Mutex
+
 	userAccessTokenCache sync.Map // map[int64]string
+	messageMemoCache     sync.Map // map[messageKey]string
 }
 
 func NewStore(data string) *Store {
 	return &Store{
 		Data: data,
 
 		userAccessTokenCache: sync.Map{},
+		messageMemoCache:     sync.Map{},
 	}
 }
 
 func (s *Store) Init() error {
 	if err := s.loadUserAccessTokenMapFromFile(); err != nil {
 		return fmt.Errorf("failed to load user access token map from file: %w", err)
 	}
+	if err := s.loadMessageMemoMapFromFile(); err != nil {
+		return fmt.Errorf("failed to load message memo map from file: %w", err)
+	}
 
 	return nil
 }

store/user.go

@@ -22,6 +22,9 @@ func (s *Store) GetUserAccessToken(userID int64) (string, bool) {
 
 // SetUserAccessToken sets the access token for the user.
 func (s *Store) SetUserAccessToken(userID int64, accessToken string) {
+	s.saveMu.Lock()
+	defer s.saveMu.Unlock()
+
 	s.userAccessTokenCache.Store(userID, accessToken)
 	if err := s.SaveUserAccessTokenMapToFile(); err != nil {
 		slog.Error("failed to save user access token map to file", "error", err)