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
 143                 index = len(s.file.Batches) - 1
 144         }
 145
 146         batch := &s.file.Batches[index]
 147         batch.Messages = append(batch.Messages, message)
 148         batch.Last = arrival
 149         if len(batch.Messages) >= whatsappBatchMessages {
 150                 batch.Closed = true
 151                 closed = true
 152         }
 153
 154         return closed, s.save()
 155 }
 156
 157 // CloseDue closes each open batch whose time is up at now and tells if one
 158 // closed.
 159 func (s *WhatsAppState) CloseDue(now time.Time) (bool, error) {
 160         s.mutex.Lock()
 161         defer s.mutex.Unlock()
 162
 163         closed := false
 164         for index := range s.file.Batches {
 165                 batch := &s.file.Batches[index]
 166                 if !batch.Closed && !now.Before(s.deadline(*batch)) {
 167                         batch.Closed = true
 168                         closed = true
 169                 }
 170         }
 171
 172         if !closed {
 173                 return false, nil
 174         }
 175
 176         return true, s.save()
 177 }
 178
 179 // CloseAll closes every open batch.
 180 func (s *WhatsAppState) CloseAll() error {
 181         s.mutex.Lock()
 182         defer s.mutex.Unlock()
 183
 184         for index := range s.file.Batches {
 185                 s.file.Batches[index].Closed = true
 186         }
 187
 188         return s.save()
 189 }
 190
 191 // Closed returns the closed batches in order.
 192 func (s *WhatsAppState) Closed() []WhatsAppBatch {
 193         s.mutex.Lock()
 194         defer s.mutex.Unlock()
 195
 196         var closed []WhatsAppBatch
 197
 198         for _, batch := range s.file.Batches {
 199                 if batch.Closed {
 200                         closed = append(closed, batch)
 201                 }
 202         }
 203
 204         return closed
 205 }
 206
 207 // Remove removes the batch with number, which is in the log.
 208 func (s *WhatsAppState) Remove(number int) error {
 209         s.mutex.Lock()
 210         defer s.mutex.Unlock()
 211
 212         s.file.Batches = slices.DeleteFunc(s.file.Batches, func(batch WhatsAppBatch) bool {
 213                 return batch.Number == number
 214         })
 215
 216         return s.save()
 217 }
 218
 219 // Wait returns how long after now the first open batch closes, or false when no
 220 // batch is open.
 221 func (s *WhatsAppState) Wait(now time.Time) (time.Duration, bool) {
 222         s.mutex.Lock()
 223         defer s.mutex.Unlock()
 224
 225         var first time.Time
 226
 227         for _, batch := range s.file.Batches {
 228                 deadline := s.deadline(batch)
 229                 if !batch.Closed && (first.IsZero() || deadline.Before(first)) {
 230                         first = deadline
 231                 }
 232         }
 233
 234         if first.IsZero() {
 235                 return 0, false
 236         }
 237
 238         return max(first.Sub(now), 0), true
 239 }
 240
 241 // Send returns the send record id, or false when there is none.
 242 func (s *WhatsAppState) Send(id string) (WhatsAppSend, bool) {
 243         s.mutex.Lock()
 244         defer s.mutex.Unlock()
 245
 246         index := slices.IndexFunc(s.file.Sends, func(record WhatsAppSend) bool {
 247                 return record.ID == id
 248         })
 249
 250         if index < 0 {
 251                 return WhatsAppSend{}, false
 252         }
 253
 254         return s.file.Sends[index], true
 255 }
 256
 257 // PutSend writes record in the place of the record with its ID, or after the
 258 // last record, and drops the oldest records past whatsappSendRecords.
 259 func (s *WhatsAppState) PutSend(record WhatsAppSend) error {
 260         s.mutex.Lock()
 261         defer s.mutex.Unlock()
 262
 263         index := slices.IndexFunc(s.file.Sends, func(held WhatsAppSend) bool {
 264                 return held.ID == record.ID
 265         })
 266
 267         if index >= 0 {
 268                 s.file.Sends[index] = record
 269         } else {
 270                 s.file.Sends = append(s.file.Sends, record)
 271         }
 272
 273         s.file.Sends = s.file.Sends[max(len(s.file.Sends)-whatsappSendRecords, 0):]
 274
 275         return s.save()
 276 }
 277
 278 // open returns the index of the open batch of chat, or -1. The caller holds the
 279 // mutex.
 280 func (s *WhatsAppState) open(chat string) int {
 281         return slices.IndexFunc(s.file.Batches, func(batch WhatsAppBatch) bool {
 282                 return !batch.Closed && batch.Chat.ID == chat
 283         })
 284 }
 285
 286 // deadline returns when batch closes in real time.
 287 func (s *WhatsAppState) deadline(batch WhatsAppBatch) time.Time {
 288         quiet := batch.Last.Add(whatsappBatchIdleTimeout)
 289         window := batch.First.Add(whatsappBatchTimeout)
 290         if window.Before(quiet) {
 291                 return window
 292         }
 293
 294         return quiet
 295 }
 296
 297 // save writes the state file whole. The caller holds the mutex.
 298 func (s *WhatsAppState) save() error {
 299         data, err := json.Marshal(s.file)
 300         if err != nil {
 301                 return fmt.Errorf("encode WhatsApp state: %w", err)
 302         }
 303
 304         err = os.MkdirAll(filepath.Dir(s.path), 0o700)
 305         if err == nil {
 306                 err = writeFileAtomic(s.path, append(data, '\n'))
 307         }
 308
 309         if err != nil {
 310                 return fmt.Errorf("save WhatsApp state: %w", err)
 311         }
 312
 313         return nil
 314 }