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         for _, batch := range s.file.Batches {
 174                 if batch.Closed {
 175                         closed = append(closed, batch)
 176                 }
 177         }
 178
 179         return closed
 180 }
 181
 182 // Remove removes the batch with number, which is in the log.
 183 func (s *WhatsAppState) Remove(number int) error {
 184         s.mutex.Lock()
 185         defer s.mutex.Unlock()
 186
 187         s.file.Batches = slices.DeleteFunc(s.file.Batches, func(batch WhatsAppBatch) bool {
 188                 return batch.Number == number
 189         })
 190
 191         return s.save()
 192 }
 193
 194 // Wait returns how long after now the first open batch closes, or false when no
 195 // batch is open.
 196 func (s *WhatsAppState) Wait(now time.Time) (time.Duration, bool) {
 197         s.mutex.Lock()
 198         defer s.mutex.Unlock()
 199
 200         var first time.Time
 201         for _, batch := range s.file.Batches {
 202                 deadline := s.deadline(batch)
 203                 if !batch.Closed && (first.IsZero() || deadline.Before(first)) {
 204                         first = deadline
 205                 }
 206         }
 207
 208         if first.IsZero() {
 209                 return 0, false
 210         }
 211
 212         return max(first.Sub(now), 0), true
 213 }
 214
 215 // Send returns the send record id, or false when there is none.
 216 func (s *WhatsAppState) Send(id string) (WhatsAppSend, bool) {
 217         s.mutex.Lock()
 218         defer s.mutex.Unlock()
 219
 220         index := slices.IndexFunc(s.file.Sends, func(record WhatsAppSend) bool {
 221                 return record.ID == id
 222         })
 223
 224         if index < 0 {
 225                 return WhatsAppSend{}, false
 226         }
 227
 228         return s.file.Sends[index], true
 229 }
 230
 231 // PutSend writes record in the place of the record with its ID, or after the
 232 // last record, and drops the oldest records past whatsappSendRecords.
 233 func (s *WhatsAppState) PutSend(record WhatsAppSend) error {
 234         s.mutex.Lock()
 235         defer s.mutex.Unlock()
 236
 237         index := slices.IndexFunc(s.file.Sends, func(held WhatsAppSend) bool {
 238                 return held.ID == record.ID
 239         })
 240
 241         if index >= 0 {
 242                 s.file.Sends[index] = record
 243         } else {
 244                 s.file.Sends = append(s.file.Sends, record)
 245         }
 246
 247         s.file.Sends = s.file.Sends[max(len(s.file.Sends)-whatsappSendRecords, 0):]
 248
 249         return s.save()
 250 }
 251
 252 // open returns the index of the open batch of chat, or -1. The caller holds the
 253 // mutex.
 254 func (s *WhatsAppState) open(chat string) int {
 255         return slices.IndexFunc(s.file.Batches, func(batch WhatsAppBatch) bool {
 256                 return !batch.Closed && batch.Chat.ID == chat
 257         })
 258 }
 259
 260 // deadline returns when batch closes in real time.
 261 func (s *WhatsAppState) deadline(batch WhatsAppBatch) time.Time {
 262         quiet := batch.Last.Add(s.quiet)
 263         window := batch.First.Add(s.window)
 264         if window.Before(quiet) {
 265                 return window
 266         }
 267
 268         return quiet
 269 }
 270
 271 // save writes the state file whole. The caller holds the mutex.
 272 func (s *WhatsAppState) save() error {
 273         data, err := json.Marshal(s.file)
 274         if err != nil {
 275                 return fmt.Errorf("encode WhatsApp state: %w", err)
 276         }
 277
 278         err = os.MkdirAll(filepath.Dir(s.path), 0o700)
 279         if err == nil {
 280                 err = writeFileAtomic(s.path, append(data, '\n'))
 281         }
 282
 283         if err != nil {
 284                 return fmt.Errorf("save WhatsApp state: %w", err)
 285         }
 286
 287         return nil
 288 }
 289
 290 // WhatsAppStateFile is the content of the state file: the number of the last
 291 // batch, the batches that are not in the log yet, and the send records, oldest
 292 // first.
 293 type WhatsAppStateFile struct {
 294         Next    int             `json:"next"`
 295         Batches []WhatsAppBatch `json:"batches"`
 296         Sends   []WhatsAppSend  `json:"sends"`
 297 }
 298
 299 // WhatsAppSend is the record of one send: the ID that the agent gave, the chat,
 300 // the SHA-256 of the text, the WhatsApp message ID, the state, which is
 301 // sending, sent, or failed, the answer of the tool, and the time of its start.
 302 type WhatsAppSend struct {
 303         ID        string    `json:"id"`
 304         Chat      string    `json:"chat"`
 305         Digest    string    `json:"digest"`
 306         MessageID string    `json:"message_id"`
 307         State     string    `json:"state"`
 308         Answer    string    `json:"answer,omitempty"`
 309         Time      time.Time `json:"time"`
 310 }