From 77dde5d133569799657238f6acd6175ec1460661 Mon Sep 17 00:00:00 2001 From: prosolis <5590409+prosolis@users.noreply.github.com> Date: Sun, 26 Jul 2026 10:16:53 -0700 Subject: [PATCH] dm: stop creating a duplicate DM room on every restart GetDMRoom's lookup chain could never hit. The in-memory cache died with the process, and the m.direct account data it fell back on was never written by the bot -- room creation only ever populated the map. As an appservice user there is no /sync to populate it either, so after each restart the first DM to a user fell through both checks and created a fresh room. Persist the mapping in a new dm_rooms table and read it before anything else. Creation now also publishes m.direct, so clients label the room as a DM. A stored room is checked for liveness before reuse, counting the user as present while merely invited -- they often never accept, and treating that as gone would recreate the room on every send, which is the original bug wearing a different hat. The same check now guards the m.direct path, which previously took the last entry on faith. Two adoption paths keep existing rooms from being orphaned: user-initiated DM invites are claimed on join, and a room the user is already talking in is adopted on first message (negative-cached so group rooms cost one member lookup, not one per message). For users predating the table, a one-shot sweep of joined rooms runs only on the path that would otherwise create a duplicate. The sweep cannot rank duplicates -- without /sync there are no timestamps, so it takes the newest by room-list order. First message from the user corrects any wrong pick. --- internal/db/db.go | 10 ++ internal/plugin/plugin.go | 228 ++++++++++++++++++++++++++++++++++++-- main.go | 12 ++ 3 files changed, 239 insertions(+), 11 deletions(-) diff --git a/internal/db/db.go b/internal/db/db.go index 22d919f..84e2bb1 100644 --- a/internal/db/db.go +++ b/internal/db/db.go @@ -1255,6 +1255,16 @@ CREATE TABLE IF NOT EXISTS presence ( updated_at INTEGER DEFAULT (unixepoch()) ); +-- DM rooms. The bot's m.direct account data is not a reliable store for an +-- appservice user (no /sync, and nothing writes it back), so the mapping lives +-- here. Without it every restart lost the cache and the bot created a fresh DM +-- room per user. +CREATE TABLE IF NOT EXISTS dm_rooms ( + user_id TEXT PRIMARY KEY, + room_id TEXT NOT NULL, + updated_at INTEGER DEFAULT (unixepoch()) +); + -- Markov CREATE TABLE IF NOT EXISTS markov_corpus ( id INTEGER PRIMARY KEY AUTOINCREMENT, diff --git a/internal/plugin/plugin.go b/internal/plugin/plugin.go index f4d3b0d..e15c4c0 100644 --- a/internal/plugin/plugin.go +++ b/internal/plugin/plugin.go @@ -83,9 +83,17 @@ func PluginVersion(p Plugin) string { } // dmCache maps user IDs to their DM room IDs to avoid creating duplicate rooms. +// It fronts the dm_rooms table; notDMCache is the negative half, marking rooms +// LearnDMRoom has already ruled out so group rooms cost one member lookup ever. var ( - dmCache = make(map[id.UserID]id.RoomID) - dmCacheMu sync.Mutex + dmCache = make(map[id.UserID]id.RoomID) + dmMapped = make(map[id.UserID]bool) + notDMCache = make(map[id.RoomID]bool) + dmCacheMu sync.Mutex + + // One-shot index of two-person rooms, built by findExistingDMRoom. + dmSweepOnce sync.Once + dmSweepIndex = make(map[id.UserID]id.RoomID) ) // Base provides common helpers for plugin implementations. @@ -686,6 +694,180 @@ func (b *Base) SendReact(roomID id.RoomID, eventID id.EventID, emoji string) err return err } +// rememberDMRoom pins a user's DM room in both the in-process cache and the +// database, so a restart doesn't send the bot off creating a duplicate room. +func rememberDMRoom(userID id.UserID, roomID id.RoomID) { + dmCacheMu.Lock() + dmCache[userID] = roomID + dmMapped[userID] = true + dmCacheMu.Unlock() + + if d := db.Get(); d != nil { + _, err := d.Exec(`INSERT INTO dm_rooms (user_id, room_id, updated_at) + VALUES (?, ?, unixepoch()) + ON CONFLICT(user_id) DO UPDATE SET room_id = excluded.room_id, updated_at = excluded.updated_at`, + string(userID), string(roomID)) + if err != nil { + slog.Error("persist dm room", "user", userID, "room", roomID, "err", err) + } + } +} + +// forgetDMRoom drops a mapping that no longer resolves to a live shared room. +func forgetDMRoom(userID id.UserID) { + dmCacheMu.Lock() + delete(dmCache, userID) + delete(dmMapped, userID) + dmCacheMu.Unlock() + + if d := db.Get(); d != nil { + if _, err := d.Exec(`DELETE FROM dm_rooms WHERE user_id = ?`, string(userID)); err != nil { + slog.Error("forget dm room", "user", userID, "err", err) + } + } +} + +// storedDMRoom reads the persisted DM room for a user, if any. +func storedDMRoom(userID id.UserID) (id.RoomID, bool) { + d := db.Get() + if d == nil { + return "", false + } + var roomID string + err := d.QueryRow(`SELECT room_id FROM dm_rooms WHERE user_id = ?`, string(userID)).Scan(&roomID) + if err != nil || roomID == "" { + return "", false + } + return id.RoomID(roomID), true +} + +// dmRoomUsable reports whether the bot and the user still share the room. The +// user counts as present while merely invited — they often never accept, and +// treating that as "gone" is what would recreate the room on every send. +func (b *Base) dmRoomUsable(roomID id.RoomID, userID id.UserID) bool { + ctx := context.Background() + + var self event.MemberEventContent + if err := b.Client.StateEvent(ctx, roomID, event.StateMember, string(b.Client.UserID), &self); err != nil { + return false + } + if self.Membership != event.MembershipJoin { + return false + } + + var other event.MemberEventContent + if err := b.Client.StateEvent(ctx, roomID, event.StateMember, string(userID), &other); err != nil { + return false + } + return other.Membership == event.MembershipJoin || other.Membership == event.MembershipInvite +} + +// publishDirect appends the room to the bot's m.direct account data, so the +// room is labelled as a DM rather than a nameless private room. +func (b *Base) publishDirect(userID id.UserID, roomID id.RoomID) { + ctx := context.Background() + + dmRooms := map[id.UserID][]id.RoomID{} + // A missing m.direct is a 404 — start from empty rather than bailing. + _ = b.Client.GetAccountData(ctx, "m.direct", &dmRooms) + + for _, existing := range dmRooms[userID] { + if existing == roomID { + return + } + } + dmRooms[userID] = append(dmRooms[userID], roomID) + + if err := b.Client.SetAccountData(ctx, "m.direct", dmRooms); err != nil { + slog.Warn("publish m.direct", "user", userID, "room", roomID, "err", err) + } +} + +// RecordDMRoom claims a room as the user's DM room. Used for user-initiated +// DM invites, where the invite itself is the intent — it overwrites any older +// mapping, since the user just told us which room they want to talk in. +func (b *Base) RecordDMRoom(userID id.UserID, roomID id.RoomID) { + slog.Info("recorded user-initiated DM room", "user", userID, "room", roomID) + rememberDMRoom(userID, roomID) + b.publishDirect(userID, roomID) +} + +// LearnDMRoom records a room the user messaged the bot in as their DM room, +// when it really is a two-person room. This adopts DM rooms that predate the +// database mapping instead of leaving them orphaned beside a freshly created one. +func (b *Base) LearnDMRoom(userID id.UserID, roomID id.RoomID) { + if b == nil || b.Client == nil { + return + } + + dmCacheMu.Lock() + _, known := dmCache[userID] + mapped := dmMapped[userID] + notDM := notDMCache[roomID] + dmCacheMu.Unlock() + if known || mapped || notDM { + return + } + if _, ok := storedDMRoom(userID); ok { + // Note it as mapped, not as resolved: caching the room here would let + // GetDMRoom skip its liveness check on a possibly-stale room. + dmCacheMu.Lock() + dmMapped[userID] = true + dmCacheMu.Unlock() + return + } + + members, err := b.Client.JoinedMembers(context.Background(), roomID) + if err != nil { + return + } + _, botIn := members.Joined[b.Client.UserID] + _, userIn := members.Joined[userID] + if len(members.Joined) != 2 || !botIn || !userIn { + // Group room — remember that so every later message here is free. + dmCacheMu.Lock() + notDMCache[roomID] = true + dmCacheMu.Unlock() + return + } + + slog.Info("adopted existing DM room", "user", userID, "room", roomID) + rememberDMRoom(userID, roomID) +} + +// findExistingDMRoom scans the bot's joined rooms for a two-person room shared +// with userID. The scan is expensive (one member lookup per room), so it runs +// at most once per process and only on the path that would otherwise create a +// duplicate room. +func (b *Base) findExistingDMRoom(userID id.UserID) (id.RoomID, bool) { + dmSweepOnce.Do(func() { + ctx := context.Background() + joined, err := b.Client.JoinedRooms(ctx) + if err != nil { + slog.Warn("DM sweep: list joined rooms", "err", err) + return + } + for _, roomID := range joined.JoinedRooms { + members, err := b.Client.JoinedMembers(ctx, roomID) + if err != nil || len(members.Joined) != 2 { + continue + } + if _, ok := members.Joined[b.Client.UserID]; !ok { + continue + } + for member := range members.Joined { + if member != b.Client.UserID { + dmSweepIndex[member] = roomID + } + } + } + slog.Info("DM sweep complete", "rooms", len(joined.JoinedRooms), "dms", len(dmSweepIndex)) + }) + + roomID, ok := dmSweepIndex[userID] + return roomID, ok +} + // GetDMRoom returns the DM room for a user, creating one if needed. func (b *Base) GetDMRoom(userID id.UserID) (id.RoomID, error) { dmCacheMu.Lock() @@ -695,17 +877,42 @@ func (b *Base) GetDMRoom(userID id.UserID) (id.RoomID, error) { } dmCacheMu.Unlock() - // Check account data for existing DM rooms - var dmRooms map[id.UserID][]id.RoomID - err := b.Client.GetAccountData(context.Background(), "m.direct", &dmRooms) - if err == nil { - if rooms, ok := dmRooms[userID]; ok && len(rooms) > 0 { - roomID := rooms[len(rooms)-1] // use most recent + // Persisted mapping — the authoritative store across restarts. + if roomID, ok := storedDMRoom(userID); ok { + if b.dmRoomUsable(roomID, userID) { dmCacheMu.Lock() dmCache[userID] = roomID dmCacheMu.Unlock() return roomID, nil } + slog.Info("stored DM room no longer usable, recreating", "user", userID, "room", roomID) + forgetDMRoom(userID) + } + + // Check account data for existing DM rooms + var dmRooms map[id.UserID][]id.RoomID + err := b.Client.GetAccountData(context.Background(), "m.direct", &dmRooms) + if err == nil { + for i := len(dmRooms[userID]) - 1; i >= 0; i-- { + roomID := dmRooms[userID][i] // most recent first + if b.dmRoomUsable(roomID, userID) { + rememberDMRoom(userID, roomID) + return roomID, nil + } + } + } + + // Last resort before creating: sweep the rooms the bot is already in for a + // two-person room shared with this user. Users who predate the dm_rooms + // table have such a room and nothing pointing at it — without this sweep + // their first post-upgrade DM would open yet another duplicate. If several + // duplicates exist we cannot tell which is liveliest (no /sync, so no + // timestamps), so the newest by room-list order wins. + if roomID, ok := b.findExistingDMRoom(userID); ok { + slog.Info("recovered pre-existing DM room by member sweep", "user", userID, "room", roomID) + rememberDMRoom(userID, roomID) + b.publishDirect(userID, roomID) + return roomID, nil } // No existing DM room — create one @@ -728,9 +935,8 @@ func (b *Base) GetDMRoom(userID id.UserID) (id.RoomID, error) { return "", fmt.Errorf("create DM room: %w", err) } - dmCacheMu.Lock() - dmCache[userID] = resp.RoomID - dmCacheMu.Unlock() + rememberDMRoom(userID, resp.RoomID) + b.publishDirect(userID, resp.RoomID) return resp.RoomID, nil } diff --git a/main.go b/main.go index c406f41..c442e51 100644 --- a/main.go +++ b/main.go @@ -255,6 +255,9 @@ func main() { // ---- Set up event handlers ---- + // Minimal Base used only for DM-room bookkeeping from the event handlers. + dmLearner := &plugin.Base{Client: client} + // Auto-join on invite + moderation member tracking sess.OnEventType(event.StateMember, func(ctx context.Context, evt *event.Event) { defer func() { @@ -277,6 +280,11 @@ func main() { slog.Error("failed to join room", "room", evt.RoomID, "err", err) } else { slog.Info("joined room", "room", evt.RoomID) + // A user-initiated DM invite: claim it now, before any + // outbound DM has a chance to create a rival room. + if mem.IsDirect { + dmLearner.RecordDMRoom(evt.Sender, evt.RoomID) + } } } return @@ -325,6 +333,10 @@ func main() { body = strings.TrimSpace(body[idx+2:]) } } + // Adopt the room the user is talking in if it's their DM room, so the + // bot replies there instead of opening a second one. + dmLearner.LearnDMRoom(evt.Sender, evt.RoomID) + msgCtx := plugin.MessageContext{ RoomID: evt.RoomID, EventID: evt.ID,