Split long replies between Markdown blocks

[onechat] / telegram.go

package main

import (
        "bytes"
        "context"
        "encoding/json"
        "errors"
        "fmt"
        "io"
        "log"
        "net/http"
        "os"
        "path/filepath"
        "strings"
        "sync"
        "time"
        "unicode/utf8"
)

const (
        // telegramAPI is the Telegram Bot API service.
        telegramAPI = "https://api.telegram.org"

        // botConfigFileName is the bot config in the data directory. The daemon
        // runs the bot only when the file exists.
        botConfigFileName = "telegram.json"

        // offsetFileName is the update offset of the bot in the agent directory.
        offsetFileName = "telegram-offset.json"

        // telegramSource is the source of the user messages from Telegram.
        telegramSource = "telegram"

        // messageUnits is the most UTF-16 code units of one Telegram message.
        messageUnits = 4096

        // richMessageBytes is the most bytes of the Markdown of one rich
        // message. The Bot API allows 32768 characters of rendered text, and
        // the Markdown of a text has at least as many bytes as the text has
        // characters.
        richMessageBytes = 32768

        // typingInterval is how often the bot shows the typing action while a
        // turn runs. Telegram shows the action for about five seconds.
        typingInterval = 4 * time.Second

        // textOnly is the answer to a message that is not text.
        textOnly = "Only text messages are supported."
)

// botConfig holds the settings of telegram.json: the bot token and the user ID
// of the owner. The token is a secret, so the file must have mode 0600.
type botConfig struct {
        Token string `json:"token"`
        Owner int64  `json:"owner"`
}

// telegramBot connects the private chat of the owner with the agent. It polls
// the Bot API for new messages and sends each text message of the owner to the
// agent. It ignores every other update. Its sender sends the replies, the
// errors, the notices of the bot, and the user messages of the other clients to
// the chat, and shows the typing action while a turn runs. Only the replies
// have a format: their Markdown goes as message entities. The rest is plain
// text. Tool calls do not go to the chat.
//
// The offset file holds the ID of the next update. It moves past an update
// only after the message of the update is durable in the log, so a restart
// does not lose a message. A crash between the flush and the offset write
// makes Telegram send the update again, and the message repeats once.
//
// A failed Bot API call waits and tries again, from minWait up to maxWait, or
// as long as the Bot API asks. It never stops the agent.
type telegramBot struct {
        api        *botAPI
        owner      int64
        offsetPath string
        logger     *log.Logger
        minWait    time.Duration
        maxWait    time.Duration

        mutex sync.Mutex
        queue []Event
        ready chan struct{}
}

// botPart is the text of one message to the chat, at most messageUnits long,
// and its entities.
type botPart struct {
        text     string
        entities []botEntity
}

// offsetFile is the content of the offset file.
type offsetFile struct {
        Offset int64 `json:"offset"`
}

// loadBotConfig reads telegram.json in dir. It returns false when the file
// does not exist. An error never holds the token.
func loadBotConfig(dir string) (botConfig, bool, error) {
        path := filepath.Join(dir, botConfigFileName)
        info, err := os.Stat(path)
        if errors.Is(err, os.ErrNotExist) {
                return botConfig{}, false, nil
        }

        if err != nil {
                return botConfig{}, false, fmt.Errorf("read bot config: %w", err)
        }

        if info.Mode().Perm()&0o077 != 0 {
                return botConfig{}, false, fmt.Errorf("%s holds the bot token, so its mode must be 0600", botConfigFileName)
        }

        data, err := os.ReadFile(path)
        if err != nil {
                return botConfig{}, false, fmt.Errorf("read bot config: %w", err)
        }

        var config botConfig
        decoder := json.NewDecoder(bytes.NewReader(data))
        decoder.DisallowUnknownFields()
        err = decoder.Decode(&config)
        if err == nil && !errors.Is(decoder.Decode(&struct{}{}), io.EOF) {
                err = errors.New("data after the JSON object")
        }

        if err != nil {
                return botConfig{}, false, fmt.Errorf("%s is not valid JSON with the fields token and owner", botConfigFileName)
        }

        err = config.validate()
        if err != nil {
                return botConfig{}, false, fmt.Errorf("%s: %w", botConfigFileName, err)
        }

        return config, true, nil
}

// validate makes sure that the token has the form of a bot token and that the
// owner is a user ID. The token goes into a URL path, so it can only hold the
// characters of a bot token.
func (c botConfig) validate() error {
        id, secret, found := strings.Cut(c.Token, ":")
        if !found || id == "" || secret == "" || strings.Trim(id, "0123456789") != "" || strings.Trim(secret, "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ_-") != "" {
                return errors.New("token does not have the form of a bot token from BotFather")
        }

        if c.Owner <= 0 {
                return errors.New("owner must be the numeric Telegram user ID of the owner")
        }

        return nil
}

// newTelegramBot returns the bot of owner on api. The offset file is at
// offsetPath, and logger gets the failures. No log line holds the text of a
// message.
func newTelegramBot(api *botAPI, owner int64, offsetPath string, logger *log.Logger) *telegramBot {
        return &telegramBot{
                api:        api,
                owner:      owner,
                offsetPath: offsetPath,
                logger:     logger,
                minWait:    time.Second,
                maxWait:    time.Minute,
                ready:      make(chan struct{}, 1),
        }
}

// Notify queues event for the chat. It never blocks, so a slow network does not
// stop the agent.
func (b *telegramBot) Notify(event Event) {
        switch event.Kind {
        case "busy", "idle", "reply", "error", "notice":
        case "user":
                if event.Source == telegramSource {
                        return
                }
        default:
                return
        }

        b.mutex.Lock()
        b.queue = append(b.queue, event)
        b.mutex.Unlock()

        select {
        case b.ready <- struct{}{}:
        default:
        }
}

// Run polls for the messages of the owner and sends the queued events to the
// chat until ctx ends. It returns an error only when a message or the offset
// cannot be written, because then the bot would lose or repeat messages.
func (b *telegramBot) Run(ctx context.Context, l *loop) error {
        ctx, cancel := context.WithCancel(ctx)
        defer cancel()

        var sender sync.WaitGroup
        sender.Add(1)
        go func() {
                defer sender.Done()

                b.send(ctx)
        }()

        err := b.poll(ctx, l)
        cancel()
        sender.Wait()

        return err
}

// send sends each queued event to the chat in order until ctx ends. While a
// turn runs, it shows the typing action every typingInterval and after each
// message, because a message ends the action.
func (b *telegramBot) send(ctx context.Context) {
        ticker := time.NewTicker(typingInterval)
        defer ticker.Stop()

        busy := false
        for ctx.Err() == nil {
                event, found := b.next()
                if !found {
                        select {
                        case <-b.ready:
                        case <-ticker.C:
                                if busy {
                                        b.typing(ctx)
                                }
                        case <-ctx.Done():
                        }

                        continue
                }

                switch event.Kind {
                case "busy":
                        busy = true
                        b.typing(ctx)
                        ticker.Reset(typingInterval)
                case "idle":
                        busy = false
                default:
                        for _, part := range botParts(event) {
                                b.deliver(ctx, part)
                        }

                        if busy {
                                b.typing(ctx)
                        }
                }
        }
}

// next removes the first queued event and returns it.
func (b *telegramBot) next() (Event, bool) {
        b.mutex.Lock()
        defer b.mutex.Unlock()

        if len(b.queue) == 0 {
                return Event{}, false
        }

        event := b.queue[0]
        b.queue = b.queue[1:]

        return event, true
}

// typing shows the typing action once. A failure is not logged, because the
// next message or poll logs the state of the network.
func (b *telegramBot) typing(ctx context.Context) {
        _ = b.api.SendTyping(ctx, b.owner)
}

// deliver sends part to the chat. When Telegram refuses a part with entities
// (HTTP 400), it sends the text of the part again once with no entities, so a
// reply is not lost to a bad entity. It tries again after each failure that a
// later try can fix, until ctx ends. It drops the part after any other failure.
func (b *telegramBot) deliver(ctx context.Context, part botPart) {
        wait := b.minWait
        for {
                err := b.api.SendMessage(ctx, b.owner, part.text, part.entities)
                if err == nil || ctx.Err() != nil {
                        return
                }

                var answer *botError
                if len(part.entities) > 0 && errors.As(err, &answer) && answer.Status == http.StatusBadRequest {
                        b.logger.Printf("entities refused, the message goes again as plain text: %v", err)
                        part.entities = nil

                        continue
                }

                retry, asked := retryable(err)
                if !retry {
                        b.logger.Printf("message dropped: %v", err)
                        return
                }

                wait = b.pause(ctx, err, wait, asked)
        }
}

// pause logs err and waits the longer of wait and asked, or until ctx ends. It
// returns the wait of the next failure: twice wait, up to maxWait.
func (b *telegramBot) pause(ctx context.Context, err error, wait, asked time.Duration) time.Duration {
        delay := max(wait, asked)
        b.logger.Printf("%v (next try in %s)", err, delay)

        timer := time.NewTimer(delay)
        defer timer.Stop()

        select {
        case <-timer.C:
        case <-ctx.Done():
        }

        return min(2*wait, b.maxWait)
}

// poll gets the updates from the stored offset on and handles each one, until
// ctx ends. Each update is handled to its end also when ctx ends during it, so
// a stop does not repeat its message.
func (b *telegramBot) poll(ctx context.Context, l *loop) error {
        offset, err := loadOffset(b.offsetPath)
        if err != nil {
                return err
        }

        wait := b.minWait
        for ctx.Err() == nil {
                updates, err := b.api.GetUpdates(ctx, offset)
                if ctx.Err() != nil {
                        return nil
                }

                if err != nil {
                        _, asked := retryable(err)
                        wait = b.pause(ctx, err, wait, asked)

                        continue
                }

                wait = b.minWait
                for _, update := range updates {
                        if ctx.Err() != nil {
                                return nil
                        }

                        err = b.handle(context.WithoutCancel(ctx), l, update)
                        if err != nil {
                                return err
                        }

                        offset = update.UpdateID + 1
                        err = saveOffset(b.offsetPath, offset)
                        if err != nil {
                                return err
                        }
                }
        }

        return nil
}

// handle runs one update. A text message of the owner in the private chat goes
// to the agent, and handle returns when it is durable in the log. The line
// /stop interrupts the turn. Any other message of the owner gets the answer
// that only text is supported. Every other update is ignored, and its content
// is not logged.
func (b *telegramBot) handle(ctx context.Context, l *loop, update botUpdate) error {
        message := update.Message
        if message == nil || message.Chat.Type != "private" || message.Chat.ID != b.owner || message.From == nil || message.From.ID != b.owner {
                b.logger.Print("update ignored: not a message of the owner in the private chat")
                return nil
        }

        switch {
        case message.Text == "":
                b.notice(textOnly)
        case strings.TrimSpace(message.Text) == stopCommand:
                if !l.Interrupt() {
                        b.notice("No turn runs.")
                }
        default:
                _, err := l.Submit(telegramSource, message.Text)
                if err != nil {
                        return err
                }

                err = l.Flush(ctx)
                if err != nil {
                        return fmt.Errorf("flush the message from Telegram: %w", err)
                }
        }

        return nil
}

// notice queues a short answer of the bot for the chat.
func (b *telegramBot) notice(text string) {
        b.Notify(Event{
                Kind: "notice",
                Text: text,
        })
}

// loadOffset reads the update offset at path. A missing file gives 0, which
// asks for every update that Telegram keeps.
func loadOffset(path string) (int64, error) {
        data, err := os.ReadFile(path)
        if errors.Is(err, os.ErrNotExist) {
                return 0, nil
        }

        if err != nil {
                return 0, fmt.Errorf("read update offset: %w", err)
        }

        var file offsetFile
        err = json.Unmarshal(data, &file)
        if err != nil {
                return 0, fmt.Errorf("%s: %w", path, err)
        }

        return file.Offset, nil
}

// saveOffset replaces the update offset at path.
func saveOffset(path string, offset int64) error {
        data, err := json.Marshal(offsetFile{
                Offset: offset,
        })
        if err != nil {
                return fmt.Errorf("encode update offset: %w", err)
        }

        err = writeFileAtomic(path, append(data, '\n'))
        if err != nil {
                return fmt.Errorf("write update offset: %w", err)
        }

        return nil
}

// botParts returns the messages of event for the chat. The Markdown of a reply
// becomes text with entities. Every other event is plain text. A reply that has
// no text left after the conversion, such as an empty code block, goes as its
// Markdown.
func botParts(event Event) []botPart {
        text, entities := botText(event), []botEntity(nil)
        if event.Kind == "reply" {
                text, entities = markdownEntities(event.Text)
        }

        if strings.TrimSpace(text) == "" {
                text, entities = event.Text, nil
        }

        return splitEntities(text, entities, messageUnits)
}

// botText returns the text of event for the chat. A user message of another
// client has its source as a mark.
func botText(event Event) string {
        switch event.Kind {
        case "user":
                return "[" + event.Source + "] " + event.Text
        case "error":
                return "Error: " + event.Text
        default:
                return event.Text
        }
}

// splitEntities cuts text into parts with splitMessage, and gives each part the
// entities in it, with offsets from the start of the part. An entity that
// crosses a cut becomes one entity in each part.
func splitEntities(text string, entities []botEntity, limit int) []botPart {
        var parts []botPart
        start := 0
        for _, part := range splitMessage(text, limit) {
                end := start + utf16Units(part)
                var kept []botEntity
                for _, entity := range entities {
                        from := max(entity.Offset, start)
                        to := min(entity.Offset+entity.Length, end)
                        if to <= from {
                                continue
                        }

                        entity.Offset = from - start
                        entity.Length = to - from
                        kept = append(kept, entity)
                }

                parts = append(parts, botPart{
                        text:     part,
                        entities: kept,
                })

                start = end
        }

        return parts
}

// splitMessage cuts text into parts of at most limit UTF-16 code units, the
// unit of the Telegram message limit. It cuts after the last line break or
// space of a part when one is in its second half, and else at the limit. It
// never cuts a character.
func splitMessage(text string, limit int) []string {
        var parts []string
        for {
                end := prefixEnd(text, limit)
                if end == len(text) {
                        return append(parts, text)
                }

                cut := strings.LastIndexByte(text[:end], '\n') + 1
                if cut <= end/2 {
                        cut = strings.LastIndexByte(text[:end], ' ') + 1
                }

                if cut <= end/2 {
                        cut = end
                }

                parts = append(parts, text[:cut])
                text = text[cut:]
        }
}

// prefixEnd returns the byte length of the longest prefix of text with at most
// limit UTF-16 code units.
func prefixEnd(text string, limit int) int {
        units := 0
        for index, char := range text {
                size := 1
                if char > 0xFFFF {
                        size = 2
                }

                if units+size > limit {
                        return index
                }

                units += size
        }

        return len(text)
}

// splitRich cuts the Markdown of a reply into parts of at most limit bytes. It
// cuts after a blank line outside a code block, so a table, a code block, or a
// paragraph stays whole. When no such cut fits in a part, it cuts after the
// last line break outside a code block, then after the last line break in a
// code block, which it closes at the end of the part and opens again in the
// next part with the same fence line. A line longer than limit is cut at the
// limit. It never cuts a character.
func splitRich(markdown string, limit int) []string {
        var parts []string
        for len(markdown) > limit {
                cut, fence, opening := richCut(markdown, limit)
                if fence == "" {
                        parts = append(parts, markdown[:cut])
                        markdown = markdown[cut:]

                        continue
                }

                parts = append(parts, markdown[:cut]+fence+"\n")
                markdown = opening + "\n" + markdown[cut:]
        }

        return append(parts, markdown)
}

// richCut returns where splitRich cuts markdown, which is longer than limit.
// For a cut in a code block, it also returns the fence that closes the block
// and the line that opens it. Each part before a cut holds text other than
// space, and a cut in a code block leaves room for the fence, so each cut makes
// the rest shorter.
func richCut(markdown string, limit int) (int, string, string) {
        blank, line, code := 0, 0, 0
        var fence, opening, inFence, inOpening string
        for start := 0; start < limit; {
                end := strings.IndexByte(markdown[start:], '\n')
                if end < 0 || start+end+1 > limit {
                        break
                }

                end += start + 1
                text := strings.TrimRight(markdown[start:end], "\r\n")
                _, opens := codeFence(text)
                switch {
                case fence != "":
                        if closesFence(text, fence) {
                                fence = ""
                        } else if end > len(opening)+1 && end+len(fence)+1 <= limit {
                                code, inFence, inOpening = end, fence, opening
                        }
                case opens != "":
                        fence, opening = opens, text
                case strings.TrimSpace(markdown[:end]) == "":
                case strings.TrimSpace(text) == "":
                        blank = end
                }

                if fence == "" && strings.TrimSpace(markdown[:end]) != "" {
                        line = end
                }

                start = end
        }

        switch {
        case blank > 0:
                return blank, "", ""
        case line > 0:
                return line, "", ""
        case code > 0:
                return code, inFence, inOpening
        }

        cut := limit
        for !utf8.RuneStart(markdown[cut]) {
                cut--
        }

        return cut, "", ""
}