Skip to content

Commit 6a70d23

Browse files
committed
P1: Wire FallbackURLs + DailyTokenBudget
- Bot.SetFallbackURLs() wraps http.Client with FallbackTransport - Bot.CheckDailyBudget() tracks daily token usage in ~/.odek/telegram_token_usage_<date> - Budget checked after each agent run; over-budget responses blocked - Both wired from resolved config in telegramCmd Subagents used: odek via delegate_task
1 parent d118635 commit 6a70d23

2 files changed

Lines changed: 104 additions & 4 deletions

File tree

cmd/odek/telegram.go

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,12 @@ func telegramCmd(args []string) error {
4848

4949
// 4. Create bot client.
5050
bot := telegram.NewBot(cfg.Token)
51+
bot.SetDailyTokenBudget(cfg.DailyTokenBudget)
52+
53+
// 4b. Configure fallback Telegram API endpoints if provided.
54+
if len(cfg.FallbackURLs) > 0 {
55+
bot.SetFallbackURLs(cfg.FallbackURLs)
56+
}
5157

5258
// 5. Create session store on disk (~/.odek/sessions/).
5359
store, err := session.NewStore()
@@ -233,6 +239,20 @@ func handleChatMessage(
233239
return
234240
}
235241

242+
// Check daily token budget.
243+
totalTokens := int64(agent.TotalInputTokens() + agent.TotalOutputTokens())
244+
if err := bot.CheckDailyBudget(totalTokens); err != nil {
245+
fmt.Fprintf(os.Stderr, "odek telegram: %v\n", err)
246+
reportError(bot, chatID, "Daily token budget exceeded. Usage for today has been tracked and will be enforced going forward.")
247+
// Still save the session so the conversation isn't lost.
248+
cs.Messages = updatedMessages
249+
cs.TurnCount++
250+
if err := sessionManager.Save(chatID, cs.Messages); err != nil {
251+
fmt.Fprintf(os.Stderr, "odek telegram: session save: %v\n", err)
252+
}
253+
return
254+
}
255+
236256
// Save the updated session messages.
237257
cs.Messages = updatedMessages
238258
cs.TurnCount++

internal/telegram/bot.go

Lines changed: 84 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,15 +9,17 @@ import (
99
"net/http"
1010
"os"
1111
"path/filepath"
12+
"strconv"
1213
"time"
1314
)
1415

1516
// Bot represents a Telegram Bot API client.
1617
type Bot struct {
17-
Token string
18-
BaseURL string
19-
FileBaseURL string
20-
Client *http.Client
18+
Token string
19+
BaseURL string
20+
FileBaseURL string
21+
Client *http.Client
22+
DailyTokenBudget int64
2123
}
2224

2325
// NewBot creates a new Bot with the given token and a default HTTP client
@@ -287,6 +289,84 @@ func (b *Bot) SetMyCommands(commands []BotCommand) error {
287289
return b.doJSON("setMyCommands", params, nil)
288290
}
289291

292+
// SetFallbackURLs configures fallback Telegram API endpoints to try if the
293+
// primary endpoint is unreachable. Each URL should be a base API URL such as
294+
// "https://api.telegram.org" (without the /bot<token> suffix). The fallback
295+
// transport rewrites the host on each request, keeping the original path
296+
// (which includes the token).
297+
func (b *Bot) SetFallbackURLs(urls []string) {
298+
if len(urls) == 0 {
299+
return
300+
}
301+
ft := NewFallbackTransport(urls)
302+
ft.WrapBot(b)
303+
}
304+
305+
// SetDailyTokenBudget sets the daily token usage budget for the bot.
306+
// When non-zero, CheckDailyBudget will reject token usage that exceeds
307+
// this limit within a calendar day.
308+
func (b *Bot) SetDailyTokenBudget(budget int64) {
309+
b.DailyTokenBudget = budget
310+
}
311+
312+
// budgetFilePath returns the path to the daily token usage tracking file.
313+
// The file is scoped to the current date so budgets reset each day.
314+
func budgetFilePath() string {
315+
home, err := os.UserHomeDir()
316+
if err != nil {
317+
home = "."
318+
}
319+
date := time.Now().Format("2006-01-02")
320+
return filepath.Join(home, ".odek", "telegram_token_usage_"+date)
321+
}
322+
323+
// CheckDailyBudget reads the current daily token usage tracking file,
324+
// adds the given number of tokens, and returns an error if the total
325+
// exceeds the configured DailyTokenBudget. If the budget is zero (unset),
326+
// no check is performed and nil is returned.
327+
func (b *Bot) CheckDailyBudget(tokens int64) error {
328+
if b.DailyTokenBudget <= 0 {
329+
return nil // budget not configured
330+
}
331+
if tokens <= 0 {
332+
return nil // nothing to track
333+
}
334+
335+
path := budgetFilePath()
336+
337+
// Ensure the parent .odek directory exists.
338+
dir := filepath.Dir(path)
339+
if err := os.MkdirAll(dir, 0755); err != nil {
340+
return fmt.Errorf("telegram: create budget dir: %w", err)
341+
}
342+
343+
// Read current usage (file may not exist yet — that's fine).
344+
var current int64
345+
data, err := os.ReadFile(path)
346+
if err == nil {
347+
if parsed, err := strconv.ParseInt(string(data), 10, 64); err == nil {
348+
current = parsed
349+
}
350+
} else if !os.IsNotExist(err) {
351+
return fmt.Errorf("telegram: read budget file: %w", err)
352+
}
353+
354+
total := current + tokens
355+
if total > b.DailyTokenBudget {
356+
return fmt.Errorf(
357+
"daily token budget exceeded: %d used + %d new = %d total, limit is %d",
358+
current, tokens, total, b.DailyTokenBudget,
359+
)
360+
}
361+
362+
// Write the updated count.
363+
if err := os.WriteFile(path, []byte(strconv.FormatInt(total, 10)), 0644); err != nil {
364+
return fmt.Errorf("telegram: write budget file: %w", err)
365+
}
366+
367+
return nil
368+
}
369+
290370
// GetMe returns basic information about the bot (useful as a health check).
291371
func (b *Bot) GetMe() (*User, error) {
292372
var user User

0 commit comments

Comments
 (0)