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

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"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
	client       *http.Client
	agentClient  *http.Client
	mu           sync.RWMutex
	botUsername  string
	workLimit    chan struct{}
	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"`
}

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,
		client:      &http.Client{Timeout: 70 * time.Second},
		agentClient: &http.Client{Timeout: 15 * time.Minute},
		workLimit:   make(chan struct{}, 4),
		slotLocks:   map[int]*sync.Mutex{},
	}
}

func (bot *telegramBot) username() string {
	bot.mu.RLock()
	defer bot.mu.RUnlock()
	return bot.botUsername
}

func (bot *telegramBot) setUsername(username string) {
	bot.mu.Lock()
	defer bot.mu.Unlock()
	bot.botUsername = strings.TrimPrefix(strings.TrimSpace(username), "@")
}

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"`
	}
	if response.StatusCode != http.StatusOK || json.Unmarshal(body, &envelope) != nil || !envelope.OK {
		return fmt.Errorf("Telegram API returned status %d", response.StatusCode)
	}
	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 {
		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 {
			bot.setUsername(me.Username)
			break
		}
		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 {
			time.Sleep(3 * time.Second)
			continue
		}
		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 commandParts(text string) (string, []string) {
	fields := strings.Fields(strings.TrimSpace(text))
	if len(fields) == 0 || !strings.HasPrefix(fields[0], "/") {
		return "", nil
	}
	command := strings.TrimPrefix(strings.ToLower(fields[0]), "/")
	command, _, _ = strings.Cut(command, "@")
	return command, fields[1:]
}

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.hermes.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 any text to chat with Hermes. Use /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.hermes.bstein.dev and open Telegram to connect it.")
		return
	}
	text := strings.TrimSpace(message.Text)
	if text == "" {
		_ = bot.sendText(message.Chat.ID, "Text messages are supported now; attachment support will be added separately.")
		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()
	_ = bot.sendAction(message.Chat.ID, "typing")
	reply, err := bot.askTenant(slot, text, update.UpdateID)
	if err != nil {
		_ = bot.sendText(message.Chat.ID, "Hermes could not answer right now. Please try again shortly.")
		return
	}
	_ = bot.sendText(message.Chat.ID, reply)
}

func (bot *telegramBot) askTenant(slot int, text string, updateID int64) (string, error) {
	if bot.router.backendAPIURL == nil {
		return "", errors.New("tenant API unavailable")
	}
	payload, _ := json.Marshal(map[string]any{
		"input":        text,
		"conversation": "telegram",
		"store":        true,
	})
	ctx, cancel := context.WithTimeout(context.Background(), 14*time.Minute)
	defer cancel()
	request, err := http.NewRequestWithContext(
		ctx,
		http.MethodPost,
		bot.router.backendAPIURL(slot)+"/v1/responses",
		bytes.NewReader(payload),
	)
	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", "telegram")
	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.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()
	delete(router.state.Telegram, router.identityHash("telegram", userID))
	return router.saveLocked()
}