git.giacomolli.uk - onechat/blob - whatsapp_state.go

   1 package main
   2
   3 import (
   4         "encoding/json"
   5         "errors"
   6         "fmt"
   7         "os"
   8         "path/filepath"
   9         "slices"
  10         "sync"
  11         "time"
  12 )
  13
  14 const (
  15         // whatsappStateFileName is the state file in the WhatsApp directory:
  16         // the batches that are not in the log yet and the records of the sends.
  17         whatsappStateFileName = "state.json"
  18
  19         // whatsappBatchQuiet is how long a chat is quiet before its batch
  20         // closes.
  21         whatsappBatchQuiet = 20 * time.Second
  22
  23         // whatsappBatchWindow is the longest time of a batch.
  24         whatsappBatchWindow = 2 * time.Minute
  25
  26         // whatsappBatchMessages is the most messages of a batch.
  27         whatsappBatchMessages = 50
  28
  29         // whatsappSendRecords is the most send records that the state file
  30         // keeps.
  31         whatsappSendRecords = 1000
  32 )
  33
  34 // openWhatsAppState reads the batches and the send records of the state file at
  35 // path. A missing file holds none. A batch closes after quiet with no message,
  36 // or after window at the latest.
  37 func openWhatsAppState(path string, quiet, window time.Duration) (*WhatsAppState, error) {
  38         s := &WhatsAppState{
  39                 path:   path,
  40                 quiet:  quiet,
  41                 window: window,
  42         }
  43
  44         data, err := os.ReadFile(path)
  45         if errors.Is(err, os.ErrNotExist) {
  46                 return s, nil
  47         }
  48
  49         if err != nil {
  50                 return nil, fmt.Errorf("read WhatsApp state: %w", err)
  51         }
  52
  53         err = json.Unmarshal(data, &s.file)
  54         if err != nil {
  55                 return nil, fmt.Errorf("%s: %w", path, err)
  56         }
  57
  58         return s, nil
  59 }
  60
  61 // WhatsAppState holds the state file: the send records, and the batches of the
  62 // WhatsApp messages that are not in the log yet, open and closed, in the order
  63 // of their first message. The messages of one chat form one batch while the
  64 // times of the messages are at most quiet apart, at most window after the
  65 // first, and the batch has fewer than whatsappBatchMessages messages. The times
  66 // are the WhatsApp times, so a catch-up after a stop gives the same batches as
  67 // live messages. A batch also closes when quiet passes after the arrival of its
  68 // last message, or window after the arrival of its first message, in real time.
  69 //
  70 // A send record keeps the result of a send by its ID, so that a repeated call
  71 // of the send tool does not send again. The state file keeps the last
  72 // whatsappSendRecords records.
  73 //
  74 // Each change writes the state file whole, so the batches and the records
  75 // survive a restart. The mutex protects them, because the event handler of the
  76 // library adds messages while Run closes and removes batches, and the turns
  77 // send.
  78 type WhatsAppState struct {
  79         path   string
  80         quiet  time.Duration
  81         window time.Duration
  82
  83         mutex sync.Mutex
  84         file  WhatsAppStateFile
  85 }
  86
  87 // Add puts message of chat into the open batch of chat, or into a new batch,
  88 // with arrival as its time of arrival. A message that the batch holds already,
  89 // which the servers sent again, changes nothing. It returns true when a batch
  90 // closed.
  91 func (s *WhatsAppState) Add(chat WhatsAppChat, message WhatsAppMessage, arrival time.Time) (bool, error) {
  92         s.mutex.Lock()
  93         defer s.mutex.Unlock()
  94
  95         closed := false
  96         index := s.open(chat.ID)
  97         if index >= 0 {
  98                 batch := &s.file.Batches[index]
  99                 if slices.ContainsFunc(batch.Messages, func(held WhatsAppMessage) bool {
 100                         return held.ID == message.ID
 101                 }) {
 102                         return false, nil
 103                 }
 104
 105                 if message.Time.Sub(batch.Messages[len(batch.Messages)-1].Time) > s.quiet || message.Time.Sub(batch.Messages[0].Time) > s.window {
 106                         batch.Closed = true
 107                         closed = true
 108                         index = -1
 109                 }
 110         }
 111
 112         if index < 0 {
 113                 s.file.Next++
 114                 s.file.Batches = append(s.file.Batches, WhatsAppBatch{
 115                         Number: s.file.Next,
 116                         Chat:   chat,
 117                         First:  arrival,
 118                 })
 119                 index = len(s.file.Batches) - 1
 120         }
 121
 122         batch := &s.file.Batches[index]
 123         batch.Messages = append(batch.Messages, message)
 124         batch.Last = arrival
 125         if len(batch.Messages) >= whatsappBatchMessages {
 126                 batch.Closed = true
 127                 closed = true
 128         }
 129
 130         return closed, s.save()
 131 }
 132
 133 // CloseDue closes each open batch whose time is up at now and tells if one
 134 // closed.
 135 func (s *WhatsAppState) CloseDue(now time.Time) (bool, error) {
 136         s.mutex.Lock()
 137         defer s.mutex.Unlock()
 138
 139         closed := false
 140         for index := range s.file.Batches {
 141                 batch := &s.file.Batches[index]
 142                 if !batch.Closed && !now.Before(s.deadline(*batch)) {
 143                         batch.Closed = true
 144                         closed = true
 145                 }
 146         }
 147
 148         if !closed {
 149                 return false, nil
 150         }
 151
 152         return true, s.save()
 153 }
 154
 155 // CloseAll closes every open batch.
 156 func (s *WhatsAppState) CloseAll() error {
 157         s.mutex.Lock()
 158         defer s.mutex.Unlock()
 159
 160         for index := range s.file.Batches {
 161                 s.file.Batches[index].Closed = true
 162         }
 163
 164         return s.save()
 165 }
 166
 167 // Closed returns the closed batches in order.
 168 func (s *WhatsAppState) Closed() []WhatsAppBatch {
 169         s.mutex.Lock()
 170         defer s.mutex.Unlock()
 171
 172         var closed []WhatsAppBatch
 173
 174         for _, batch := range s.file.Batches {
 175                 if batch.Closed {
 176                         closed = append(closed, batch)
 177                 }
 178         }
 179
 180         return closed
 181 }
 182
 183 // Remove removes the batch with number, which is in the log.
 184 func (s *WhatsAppState) Remove(number int) error {
 185         s.mutex.Lock()
 186         defer s.mutex.Unlock()
 187
 188         s.file.Batches = slices.DeleteFunc(s.file.Batches, func(batch WhatsAppBatch) bool {
 189                 return batch.Number == number
 190         })
 191
 192         return s.save()
 193 }
 194
 195 // Wait returns how long after now the first open batch closes, or false when no
 196 // batch is open.
 197 func (s *WhatsAppState) Wait(now time.Time) (time.Duration, bool) {
 198         s.mutex.Lock()
 199         defer s.mutex.Unlock()
 200
 201         var first time.Time
 202
 203         for _, batch := range s.file.Batches {
 204                 deadline := s.deadline(batch)
 205                 if !batch.Closed && (first.IsZero() || deadline.Before(first)) {
 206                         first = deadline
 207                 }
 208         }
 209
 210         if first.IsZero() {
 211                 return 0, false
 212         }
 213
 214         return max(first.Sub(now), 0), true
 215 }
 216
 217 // Send returns the send record id, or false when there is none.
 218 func (s *WhatsAppState) Send(id string) (WhatsAppSend, bool) {
 219         s.mutex.Lock()
 220         defer s.mutex.Unlock()
 221
 222         index := slices.IndexFunc(s.file.Sends, func(record WhatsAppSend) bool {
 223                 return record.ID == id
 224         })
 225
 226         if index < 0 {
 227                 return WhatsAppSend{}, false
 228         }
 229
 230         return s.file.Sends[index], true
 231 }
 232
 233 // PutSend writes record in the place of the record with its ID, or after the
 234 // last record, and drops the oldest records past whatsappSendRecords.
 235 func (s *WhatsAppState) PutSend(record WhatsAppSend) error {
 236         s.mutex.Lock()
 237         defer s.mutex.Unlock()
 238
 239         index := slices.IndexFunc(s.file.Sends, func(held WhatsAppSend) bool {
 240                 return held.ID == record.ID
 241         })
 242
 243         if index >= 0 {
 244                 s.file.Sends[index] = record
 245         } else {
 246                 s.file.Sends = append(s.file.Sends, record)
 247         }
 248
 249         s.file.Sends = s.file.Sends[max(len(s.file.Sends)-whatsappSendRecords, 0):]
 250
 251         return s.save()
 252 }
 253
 254 // open returns the index of the open batch of chat, or -1. The caller holds the
 255 // mutex.
 256 func (s *WhatsAppState) open(chat string) int {
 257         return slices.IndexFunc(s.file.Batches, func(batch WhatsAppBatch) bool {
 258                 return !batch.Closed && batch.Chat.ID == chat
 259         })
 260 }
 261
 262 // deadline returns when batch closes in real time.
 263 func (s *WhatsAppState) deadline(batch WhatsAppBatch) time.Time {
 264         quiet := batch.Last.Add(s.quiet)
 265         window := batch.First.Add(s.window)
 266         if window.Before(quiet) {
 267                 return window
 268         }
 269
 270         return quiet
 271 }
 272
 273 // save writes the state file whole. The caller holds the mutex.
 274 func (s *WhatsAppState) save() error {
 275         data, err := json.Marshal(s.file)
 276         if err != nil {
 277                 return fmt.Errorf("encode WhatsApp state: %w", err)
 278         }
 279
 280         err = os.MkdirAll(filepath.Dir(s.path), 0o700)
 281         if err == nil {
 282                 err = writeFileAtomic(s.path, append(data, '\n'))
 283         }
 284
 285         if err != nil {
 286                 return fmt.Errorf("save WhatsApp state: %w", err)
 287         }
 288
 289         return nil
 290 }
 291
 292 // WhatsAppStateFile is the content of the state file: the number of the last
 293 // batch, the batches that are not in the log yet, and the send records, oldest
 294 // first.
 295 type WhatsAppStateFile struct {
 296         Next    int             `json:"next"`
 297         Batches []WhatsAppBatch `json:"batches"`
 298         Sends   []WhatsAppSend  `json:"sends"`
 299 }
 300
 301 // WhatsAppSend is the record of one send: the ID that the agent gave, the chat,
 302 // the SHA-256 of the text, the WhatsApp message ID, the state, which is
 303 // sending, sent, or failed, the answer of the tool, and the time of its start.
 304 type WhatsAppSend struct {
 305         ID        string    `json:"id"`
 306         Chat      string    `json:"chat"`
 307         Digest    string    `json:"digest"`
 308         MessageID string    `json:"message_id"`
 309         State     string    `json:"state"`
 310         Answer    string    `json:"answer,omitempty"`
 311         Time      time.Time `json:"time"`
 312 }