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