atlas-iac/services/hermes/router/telegram.go

package main

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

type telegramConfig struct {
	BotToken string
	RelayKey string
}

type telegramBot struct {
	config        telegramConfig
	router        *tenantRouter
	apiBase       string
	fileBase      string
	client        *http.Client
	mediaClient   *http.Client
	agentClient   *http.Client
	mu            sync.RWMutex
	botUsername   string
	ready         bool
	lastError     string
	workLimit     chan struct{}
	activityEvery time.Duration
	slotLocks     map[int]*sync.Mutex
	slotLocksMux  sync.Mutex
}

type telegramUpdate struct {
	UpdateID int64            `json:"update_id"`
	Message  *telegramMessage `json:"message"`
}

type telegramMessage struct {
	MessageID int64               `json:"message_id"`
	From      *telegramUser       `json:"from"`
	Chat      telegramChat        `json:"chat"`
	Text      string              `json:"text"`
	Caption   string              `json:"caption"`
	Photo     []telegramPhotoSize `json:"photo"`
}

type telegramPhotoSize struct {
	FileID   string `json:"file_id"`
	FileSize int64  `json:"file_size"`
	Width    int    `json:"width"`
	Height   int    `json:"height"`
}

type telegramUser struct {
	ID int64 `json:"id"`
}

type telegramChat struct {
	ID   int64  `json:"id"`
	Type string `json:"type"`
}

func readTelegramConfig(path string) (telegramConfig, error) {
	if path == "" {
		path = "/vault/secrets/telegram-config"
	}
	content, err := os.ReadFile(path)
	if err != nil {
		return telegramConfig{}, err
	}
	config := telegramConfig{}
	for _, line := range strings.Split(string(content), "\n") {
		key, value, found := strings.Cut(strings.TrimSpace(line), "=")
		if !found {
			continue
		}
		switch strings.TrimSpace(key) {
		case "bot_token":
			config.BotToken = strings.TrimSpace(value)
		case "relay_key":
			config.RelayKey = strings.TrimSpace(value)
		}
	}
	return config, nil
}

func newTelegramBot(config telegramConfig, router *tenantRouter) *telegramBot {
	return &telegramBot{
		config:        config,
		router:        router,
		apiBase:       "https://api.telegram.org/bot" + config.BotToken,
		fileBase:      "https://api.telegram.org/file/bot" + config.BotToken,
		client:        &http.Client{Timeout: 70 * time.Second},
		mediaClient:   newTenantMediaClient(),
		agentClient:   &http.Client{Timeout: 15 * time.Minute},
		workLimit:     make(chan struct{}, 4),
		activityEvery: 4 * time.Second,
		slotLocks:     map[int]*sync.Mutex{},
	}
}

func (bot *telegramBot) status() (string, bool, string) {
	bot.mu.RLock()
	defer bot.mu.RUnlock()
	return bot.botUsername, bot.ready, bot.lastError
}

func (bot *telegramBot) setStatus(username string, ready bool, lastError string) bool {
	bot.mu.Lock()
	defer bot.mu.Unlock()
	previousUsername := bot.botUsername
	previousReady := bot.ready
	previousError := bot.lastError
	if username != "" {
		bot.botUsername = strings.TrimPrefix(strings.TrimSpace(username), "@")
	}
	bot.ready = ready
	bot.lastError = strings.TrimSpace(lastError)
	return previousUsername != bot.botUsername || previousReady != bot.ready || previousError != bot.lastError
}

func (bot *telegramBot) slotLock(slot int) *sync.Mutex {
	bot.slotLocksMux.Lock()
	defer bot.slotLocksMux.Unlock()
	if bot.slotLocks[slot] == nil {
		bot.slotLocks[slot] = &sync.Mutex{}
	}
	return bot.slotLocks[slot]
}

func (bot *telegramBot) call(ctx context.Context, method string, values url.Values, result any) error {
	request, err := http.NewRequestWithContext(
		ctx,
		http.MethodPost,
		bot.apiBase+"/"+method,
		strings.NewReader(values.Encode()),
	)
	if err != nil {
		return errors.New("create Telegram request")
	}
	request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	response, err := bot.client.Do(request)
	if err != nil {
		return errors.New("Telegram API unavailable")
	}
	defer response.Body.Close()
	body, err := io.ReadAll(io.LimitReader(response.Body, 2<<20))
	if err != nil {
		return errors.New("read Telegram response")
	}
	var envelope struct {
		OK          bool            `json:"ok"`
		Result      json.RawMessage `json:"result"`
		Description string          `json:"description"`
		ErrorCode   int             `json:"error_code"`
	}
	if err := json.Unmarshal(body, &envelope); err != nil {
		return fmt.Errorf("Telegram API returned undecodable status %d", response.StatusCode)
	}
	if response.StatusCode != http.StatusOK || !envelope.OK {
		description := strings.TrimSpace(envelope.Description)
		if description == "" {
			description = http.StatusText(response.StatusCode)
		}
		return fmt.Errorf("Telegram API error %d: %s", envelope.ErrorCode, description)
	}
	if result != nil && len(envelope.Result) > 0 {
		if err := json.Unmarshal(envelope.Result, result); err != nil {
			return errors.New("decode Telegram response")
		}
	}
	return nil
}

func (bot *telegramBot) run() {
	for {
		for {
			ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
			var me struct {
				Username string `json:"username"`
			}
			err := bot.call(ctx, "getMe", url.Values{}, &me)
			cancel()
			if err == nil {
				ctx, cancel = context.WithTimeout(context.Background(), 15*time.Second)
				err = bot.call(ctx, "deleteWebhook", url.Values{
					"drop_pending_updates": {"false"},
				}, nil)
				cancel()
				if err == nil {
					if bot.setStatus(me.Username, true, "") {
						log.Printf("Telegram bot ready as @%s", strings.TrimPrefix(me.Username, "@"))
					}
					break
				}
			}
			if bot.setStatus("", false, "Telegram rejected the configured bot token or is temporarily unavailable.") {
				log.Printf("Telegram bot readiness check failed: %v", err)
			}
			time.Sleep(10 * time.Second)
		}
		for {
			offset := bot.router.telegramOffset()
			ctx, cancel := context.WithTimeout(context.Background(), 65*time.Second)
			values := url.Values{
				"offset":          {strconv.FormatInt(offset, 10)},
				"timeout":         {"50"},
				"allowed_updates": {`["message"]`},
			}
			var updates []telegramUpdate
			err := bot.call(ctx, "getUpdates", values, &updates)
			cancel()
			if err != nil {
				if bot.setStatus("", false, "Hermes lost contact with the configured Telegram bot.") {
					log.Printf("Telegram update polling failed: %v", err)
				}
				time.Sleep(3 * time.Second)
				break
			}
			for _, update := range updates {
				_ = bot.router.setTelegramOffset(update.UpdateID + 1)
				bot.workLimit <- struct{}{}
				go func(update telegramUpdate) {
					defer func() { <-bot.workLimit }()
					bot.handleUpdate(update)
				}(update)
			}
		}
	}
}

func (bot *telegramBot) handleUpdate(update telegramUpdate) {
	message := update.Message
	if message == nil || message.From == nil || message.Chat.Type != "private" || message.Chat.ID != message.From.ID {
		return
	}
	userID := strconv.FormatInt(message.From.ID, 10)
	command, args := commandParts(message.Text)
	if command == "start" || command == "link" {
		if len(args) == 0 {
			_ = bot.sendText(message.Chat.ID, "Sign in to chat.bstein.dev, open Telegram, and create a one-time link code.")
			return
		}
		if _, err := bot.router.consumeLink(userID, args[0]); err != nil {
			_ = bot.sendText(message.Chat.ID, "That link code is invalid or expired. Create a new one in Hermes WebUI.")
			return
		}
		_ = bot.sendText(message.Chat.ID, "Telegram is linked to your private Hermes account. Send a message whenever you are ready.")
		return
	}
	if command == "unlink" {
		if err := bot.router.unlinkTelegramUser(userID); err != nil {
			_ = bot.sendText(message.Chat.ID, "I could not unlink Telegram right now. Try again shortly.")
			return
		}
		_ = bot.sendText(message.Chat.ID, "Telegram has been unlinked from Hermes.")
		return
	}
	if command == "help" {
		_ = bot.sendText(message.Chat.ID, "Send text or a photo to chat with Hermes. Photos are available for analysis on the message that carries them. Image edits can reuse the latest image Hermes generated; inbound Telegram photos are not stored as editable artifacts. Use /topic <name> to create or switch a durable topic, /topics to list topics, or /unlink to disconnect this Telegram account. Model and intensity controls are available in Hermes WebUI.")
		return
	}
	slot, linked := bot.router.telegramSlot(userID)
	if !linked {
		_ = bot.sendText(message.Chat.ID, "This Telegram account is not linked. Sign in to chat.bstein.dev and open Telegram to connect it.")
		return
	}
	if command == "topic" {
		if len(args) == 0 {
			active := bot.router.activeTelegramTopic(userID)
			_ = bot.sendText(message.Chat.ID, "Current topic: "+active.Label+". Use /topic <name> to switch or create one.")
			return
		}
		active, err := bot.router.selectTelegramTopic(userID, strings.Join(args, " "))
		if err != nil {
			_ = bot.sendText(message.Chat.ID, "That topic name is invalid. Use a short descriptive name, up to 48 characters.")
			return
		}
		_ = bot.sendText(message.Chat.ID, "Switched to topic: "+active.Label+". Messages here keep their own bounded context.")
		return
	}
	if command == "topics" {
		active := bot.router.activeTelegramTopic(userID)
		lines := []string{"Telegram topics (current: " + active.Label + "):"}
		for _, topic := range bot.router.telegramTopics(userID) {
			lines = append(lines, "• "+topic.Label)
		}
		lines = append(lines, "Use /topic <name> to switch or create one.")
		_ = bot.sendText(message.Chat.ID, strings.Join(lines, "\n"))
		return
	}
	text := strings.TrimSpace(message.Text)
	hasPhoto := len(message.Photo) > 0
	if hasPhoto {
		text = strings.TrimSpace(message.Caption)
		if text == "" {
			text = "Describe this image and answer any question it implies."
		}
	}
	if text == "" {
		_ = bot.sendText(message.Chat.ID, "Hermes currently accepts text and image attachments. Please send other file types through Hermes WebUI.")
		return
	}
	if utf8.RuneCountInString(text) > 12000 {
		_ = bot.sendText(message.Chat.ID, "That message is too long. Please split it into smaller parts.")
		return
	}
	lock := bot.slotLock(slot)
	lock.Lock()
	defer lock.Unlock()
	topic, err := bot.router.touchTelegramTopic(userID)
	if err != nil {
		_ = bot.sendText(message.Chat.ID, "Hermes could not save the active topic right now. Please try again shortly.")
		return
	}
	text = telegramInputPrompt(text, hasPhoto)
	stopActivity := bot.startTelegramActivity(message.Chat.ID)
	defer stopActivity()
	var reply string
	if hasPhoto {
		media, fetchErr := bot.fetchTelegramPhoto(message.Photo)
		if fetchErr != nil {
			_ = bot.sendText(message.Chat.ID, "Hermes could not safely read that image. Please try a smaller JPEG, PNG, or WebP image.")
			return
		}
		defer os.Remove(media.Path)
		reply, err = bot.askTenantImage(slot, text, media, update.UpdateID, topic)
	} else {
		reply, err = bot.askTenant(slot, text, update.UpdateID, topic)
	}
	if err != nil {
		_ = bot.sendText(message.Chat.ID, "Hermes could not answer right now. Please try again shortly.")
		return
	}
	_ = bot.sendReply(message.Chat.ID, slot, reply)
}

func (bot *telegramBot) askTenant(slot int, text string, updateID int64, topic activeTelegramTopic) (string, error) {
	if bot.router.backendAPIURL == nil {
		return "", errors.New("tenant API unavailable")
	}
	payload, _ := json.Marshal(map[string]any{
		"input":        text,
		"conversation": topic.Conversation,
		"store":        true,
		"truncation":   "auto",
	})
	return bot.askTenantRequest(slot, bytes.NewReader(payload), updateID, topic)
}

func (bot *telegramBot) askTenantRequest(slot int, bodyReader io.Reader, updateID int64, topic activeTelegramTopic) (string, error) {
	if bot.router.backendAPIURL == nil {
		return "", errors.New("tenant API unavailable")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 14*time.Minute)
	defer cancel()
	request, err := http.NewRequestWithContext(
		ctx,
		http.MethodPost,
		bot.router.backendAPIURL(slot)+"/v1/responses",
		bodyReader,
	)
	if err != nil {
		return "", err
	}
	request.Header.Set("Authorization", "Bearer "+bot.config.RelayKey)
	request.Header.Set("Content-Type", "application/json")
	request.Header.Set("X-Hermes-Session-Key", topic.Conversation)
	request.Header.Set("X-Hermes-Conversation-Platform", "telegram")
	request.Header.Set("X-Hermes-Conversation-Title", "Telegram · "+topic.Label)
	request.Header.Set("Idempotency-Key", "telegram-"+strconv.FormatInt(updateID, 10))
	response, err := bot.agentClient.Do(request)
	if err != nil {
		return "", errors.New("tenant request failed")
	}
	defer response.Body.Close()
	body, err := io.ReadAll(io.LimitReader(response.Body, 8<<20))
	if err != nil || response.StatusCode != http.StatusOK {
		return "", fmt.Errorf("tenant returned status %d", response.StatusCode)
	}
	var parsed struct {
		Output []struct {
			Type    string `json:"type"`
			Role    string `json:"role"`
			Content []struct {
				Type string `json:"type"`
				Text string `json:"text"`
			} `json:"content"`
		} `json:"output"`
	}
	if err := json.Unmarshal(body, &parsed); err != nil {
		return "", errors.New("decode tenant response")
	}
	var parts []string
	for _, item := range parsed.Output {
		if item.Type != "message" || item.Role != "assistant" {
			continue
		}
		for _, content := range item.Content {
			if (content.Type == "output_text" || content.Type == "text") && strings.TrimSpace(content.Text) != "" {
				parts = append(parts, strings.TrimSpace(content.Text))
			}
		}
	}
	answer := strings.Join(parts, "\n\n")
	if answer == "" {
		return "", errors.New("tenant returned no assistant text")
	}
	return answer, nil
}

func (bot *telegramBot) sendAction(chatID int64, action string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	return bot.sendActionContext(ctx, chatID, action)
}

func (bot *telegramBot) sendActionContext(ctx context.Context, chatID int64, action string) error {
	return bot.call(ctx, "sendChatAction", url.Values{
		"chat_id": {strconv.FormatInt(chatID, 10)},
		"action":  {action},
	}, nil)
}

func splitTelegramText(text string) []string {
	runes := []rune(strings.TrimSpace(text))
	if len(runes) == 0 {
		return []string{"Hermes completed the request without a text response."}
	}
	const limit = 3900
	var chunks []string
	for len(runes) > limit {
		cut := limit
		for index := limit; index > limit-500; index-- {
			if runes[index-1] == '\n' || runes[index-1] == ' ' {
				cut = index
				break
			}
		}
		chunks = append(chunks, strings.TrimSpace(string(runes[:cut])))
		runes = runes[cut:]
	}
	if tail := strings.TrimSpace(string(runes)); tail != "" {
		chunks = append(chunks, tail)
	}
	return chunks
}

func (bot *telegramBot) sendText(chatID int64, text string) error {
	for _, chunk := range splitTelegramText(text) {
		ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
		err := bot.call(ctx, "sendMessage", url.Values{
			"chat_id":                  {strconv.FormatInt(chatID, 10)},
			"text":                     {chunk},
			"disable_web_page_preview": {"true"},
		}, nil)
		cancel()
		if err != nil {
			return err
		}
	}
	return nil
}

func (router *tenantRouter) telegramOffset() int64 {
	router.mu.Lock()
	defer router.mu.Unlock()
	return router.state.TelegramOffset
}

func (router *tenantRouter) setTelegramOffset(offset int64) error {
	router.mu.Lock()
	defer router.mu.Unlock()
	if offset <= router.state.TelegramOffset {
		return nil
	}
	router.state.TelegramOffset = offset
	return router.saveLocked()
}

func (router *tenantRouter) unlinkTelegramUser(userID string) error {
	router.mu.Lock()
	defer router.mu.Unlock()
	identity := router.identityHash("telegram", userID)
	delete(router.state.Telegram, identity)
	delete(router.state.TelegramTopics, identity)
	return router.saveLocked()
}