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