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