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