Compare commits

...
3 Commits
Author SHA1 Message Date
HugeFrog24 564f96c97a README 2026-07-17 11:53:45 +02:00
HugeFrog24 7fd7818142 Add web search capabilities 2026-07-17 10:49:24 +02:00
HugeFrog24 0543283b8a Adaptive thinking 2026-07-16 17:10:50 +02:00
36 changed files with 6377 additions and 2 deletions
+77
View File
@@ -0,0 +1,77 @@
# Git files
.git
.gitignore
.gitattributes
# Documentation
README.md
*.md
# Docker files
Dockerfile
docker-compose.yml
.dockerignore
# Environment files
.env
.env.*
# Log files
*.log
logs/
# Database files
*.db
*.sqlite
*.sqlite3
bot.db
# Config files (except default template)
config/*
!config/default.json
# Test files
*_test.go
test/
tests/
# Build artifacts
telegram-bot
*.exe
*.dll
*.so
*.dylib
# IDE files
.vscode/
.idea/
*.swp
*.swo
*~
# OS files
.DS_Store
Thumbs.db
# Go specific
vendor/
*.mod.backup
*.sum.backup
# Temporary files
tmp/
temp/
*.tmp
# Coverage files
*.out
coverage.html
# CI/CD files
.github/
.gitlab-ci.yml
.travis.yml
# Examples and documentation
examples/
docs/
+13
View File
@@ -0,0 +1,13 @@
# Enforce LF line endings for all files
* text eol=lf
# Specific file types that should always have LF line endings
*.go text eol=lf
*.json text eol=lf
*.sh text eol=lf
*.md text eol=lf
# Example: Binary files should not be modified
*.jpg binary
*.png binary
*.gif binary
+37
View File
@@ -0,0 +1,37 @@
name: CI
on:
push:
branches: [ main ]
pull_request:
branches: [ main ]
jobs:
# Lint job
lint:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v6
- uses: golangci/golangci-lint-action@v9
with:
version: v2.12.2
args: --timeout 5m
# Test job
test:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v6
- uses: actions/setup-go@v6
with:
go-version: '1.26.0'
- run: go test ./... -v
# Security scan job
security:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v6
- uses: securego/[email protected]
with:
args: ./...
+18
View File
@@ -0,0 +1,18 @@
# Local IDE config & user settings
.vscode/
# Go vendor directory
vendor/
# Environment variables
.env
# Any log files
*.log
# Database file
bot.db
# All config files except for the default
config/*
!config/default.json
+55
View File
@@ -0,0 +1,55 @@
# Multi-stage build for Go Telegram Bot
# Build stage
FROM golang:1.26-alpine AS builder
# Install build dependencies including C compiler for CGO
RUN apk add --no-cache git ca-certificates tzdata gcc musl-dev
# Set working directory
WORKDIR /build
# Copy go mod files first for better caching
COPY go.mod go.sum ./
# Download dependencies
RUN go mod download
# Copy source code
COPY . .
# Build the application
RUN CGO_ENABLED=1 GOOS=linux go build -a -installsuffix cgo -o telegram-bot .
# Runtime stage
FROM alpine:latest
# Merged into a single RUN to minimise image layers (docker:S7031).
# Order matters: packages must be installed before adduser/addgroup,
# and the app directory must exist before chown runs.
RUN apk --no-cache add ca-certificates tzdata sqlite && \
addgroup -g 1001 -S appgroup && \
adduser -u 1001 -S appuser -G appgroup && \
mkdir -p /app/config /app/data /app/logs && \
chown -R appuser:appgroup /app
# Set working directory
WORKDIR /app
# Copy binary from builder stage
COPY --from=builder /build/telegram-bot /app/telegram-bot
# Copy default config as template
COPY --chown=appuser:appgroup config/default.json /app/config/
# Switch to non-root user
USER appuser
# Expose any ports if needed (not required for this bot)
# EXPOSE 8080
# Health check
HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \
CMD pgrep telegram-bot || exit 1
# Run the application
CMD ["/app/telegram-bot"]
+229
View File
@@ -0,0 +1,229 @@
# Go Telegram Multibot
A scalable, multi-bot solution for Telegram using Go, GORM, and the Anthropic API.
## Design Considerations
- AI-powered (Anthropic Claude)
- Voice message support (ElevenLabs STT + TTS) — optional, enabled per bot via config
- Supports multiple bot profiles
- Uses SQLite for persistence
- Implements rate limiting and user management
- Modular architecture
- Comprehensive unit tests
## Usage
### Docker Deployment (Recommended)
1. Clone the repository:
```bash
git clone https://github.com/HugeFrog24/go-telegram-bot.git
cd go-telegram-bot
```
2. Copy the default config template and edit it:
```bash
cp config/default.json config/mybot.json
nano config/mybot.json
```
> [!IMPORTANT]
> Keep your config files secret and do not commit them to version control.
3. Create data directory and run:
```bash
mkdir -p data
docker-compose up -d
```
### Native Deployment
1. Install using `go get`:
```bash
go get -u github.com/HugeFrog24/go-telegram-bot
cd go-telegram-bot
```
2. Configure as above, then build:
```bash
go build -o telegram-bot
```
## Trying Out New Behavior Safely
Want to experiment with a different personality, tone, or set of instructions without disturbing the bot your users already talk to? Run a second, separate bot just for testing.
Each bot profile is its own config file with its own Telegram token, and bots are fully independent — separate identity, separate chat history, separate settings. So a "test twin" is quick to set up:
1. Create a new bot with [@BotFather](https://t.me/BotFather) and copy its token.
2. Copy your existing config to a new file, e.g. `cp config/mybot.json config/mybot-test.json`.
3. In the new file, paste the new token, give it a different `id`, and edit `system_prompts` to try your changes.
4. Start it alongside your main bot. Chat with the test bot, tweak its prompt, and restart the test bot to try again — your real users never see the experiments.
5. Happy with the result? Copy the same change into your main bot's config and restart it.
> [!NOTE]
> A test bot always needs its **own** token. Telegram only lets one running bot listen on a given token, so you can't point a second copy at your live bot — give the twin its own @BotFather bot instead.
## Configuration
Each bot is one JSON file in `config/` (see `config/default.json` for the template). Keys of note:
| Key | Type | Default | Description |
| ------------------ | ------ | ----------- | ----------- |
| `max_tokens` | number | `1000` | Maximum output tokens per reply. **Thinking tokens count toward this limit** — raise it (e.g. `4000`+) whenever `thinking` is `"adaptive"`, or a turn can spend the whole budget on reasoning and produce no text. |
| `thinking` | string | *(omitted)* | Reasoning mode: `"adaptive"` (the model decides when and how much to think) or `"disabled"`. Omit the key entirely to use the model's own API default. If the configured model doesn't support the chosen mode, the API rejects the request with a 400 — owners/admins see the raw error, regular users get the generic fallback. Check [Anthropic's model docs](https://platform.claude.com/docs/en/build-with-claude/adaptive-thinking) for per-model support. |
| `thinking_display` | string | *(omitted)* | `"summarized"` or `"omitted"`. Only valid together with `"thinking": "adaptive"`. Controls whether the API returns a readable summary of the reasoning (logged, never sent to chat). Thinking is billed the same either way; when omitted, the API's per-model default applies. |
| `web_search` | object | *(omitted)* | Enables Anthropic's server-side web search (and, optionally, web fetch), sandboxed to a domain allowlist. Omit the key entirely to leave both tools off — an absent block sends byte-identical requests to before. See [Web search / fetch](#web-search--fetch) below. |
Every reply logs one accounting line — `[usage] model=... in=... out=... thinking=... stop=...` — so thinking spend (billed even when its text is omitted) stays visible in `journalctl`/`docker compose logs`. A `stop=max_tokens` line is accompanied by an error-level warning that the reply was truncated.
### Web search / fetch
The `web_search` block wires in Anthropic's server-side `web_search` (and optionally `web_fetch`) tools, mirroring how `mcp_servers` wires in MCP toolsets: the engine provides the generic mechanism, and the per-bot config carries the policy. **No domains are hardcoded in the engine** — each profile declares its own allowlist. When the model decides a question needs current information, it runs a search server-side; results and any fetch targets are confined to the domains you list.
```json
"web_search": {
"allowed_domains": [
"example.com/help",
"docs.example.com"
],
"max_uses": 3,
"fetch": true,
"max_content_tokens": 50000
}
```
| Field | Type | Description |
| ---------------------- | -------- | ----------- |
| `allowed_domains` | string[] | Restricts **`web_search`** results (and, by default, `web_fetch` targets) to these. Each entry is a **domain with an optional path prefix and no scheme** — `example.com` or `example.com/help`. Mutually exclusive with `blocked_domains`. **The path only scopes `web_search`** — see the asymmetry note below. |
| `blocked_domains` | string[] | Excludes these domains instead of allowlisting. Mutually exclusive with `allowed_domains` (setting both fails validation at boot, since the API 400s). |
| `fetch_allowed_domains`| string[] | **Host-only** allowlist for `web_fetch`, independent of the search list. Omit it and `web_fetch` reuses the *hosts* of `allowed_domains`. Set it to fetch a **narrower** set than you search — e.g. to keep a shared host (a social platform) search-only. Only meaningful with `fetch: true`. |
| `max_uses` | number | Caps how many searches the model may run per turn. Omit for no cap. |
| `fetch` | bool | Also enable `web_fetch`, which pulls full page content for URLs already surfaced by a search or pasted by the user (it cannot fetch model-invented URLs). Citations are always on when fetch is enabled, so fetched passages are sourceable. |
| `max_content_tokens` | number | Caps the tokens a single `web_fetch` may pull into context. Only meaningful with `fetch: true`. |
**Search/fetch path asymmetry (important).** Anthropic filters the two tools differently: `web_search` honors a path prefix (`example.com/help` matches only `example.com/help/...`), but **`web_fetch` matches on the host only — an entry that includes a path never matches any fetch URL**. The engine bridges this: `web_search` gets your entries verbatim, while `web_fetch` gets the host portion of each entry (path stripped, deduped). So a path-scoped `allowed_domains` still permits fetching across the whole host. If that host isn't wholly trusted, list the fetchable hosts explicitly in `fetch_allowed_domains` instead.
**Search-only pattern for shared hosts (e.g. social).** To let the model *search* a single account on a shared platform without ever *fetching* the wider platform, put the account path in `allowed_domains` and leave its host out of `fetch_allowed_domains`:
```json
"web_search": {
"allowed_domains": ["helpcenter.example/hc", "x.com/youraccount"],
"fetch_allowed_domains": ["helpcenter.example"],
"fetch": true
}
```
Here search is confined to your help center *and* your one social account, but fetch can only ever reach `helpcenter.example` — a pasted `x.com/someone-else` URL is not fetchable. Note the search side only sandboxes cleanly on platforms that keep an account's content under its handle path (X/Twitter, TikTok `@handle`, Facebook); use the `handle/` or `handle/*` form to avoid a prefix bleed. Instagram is an exception — posts live at `instagram.com/p/<code>`, not under the handle — so it can't be account-sandboxed by path.
The allowlist is a **hard, server-side sandbox**, not a prompt request — off-list results are dropped and an off-list fetch target returns `url_not_allowed`, which is the exfiltration mitigation Anthropic recommends for bots processing untrusted input. Keep the list tight. A fetch can still fail for site-side reasons (bot protection such as Cloudflare returns `url_not_accessible`); the search snippet plus citations remain the reliable signal, so a full fetch is a bonus, not a dependency.
`web_fetch` can only open a URL that a `web_search` result surfaced (or the user pasted) — never one the model invents or rebuilds from a page title; those return `url_not_in_prior_context`. So a fetch-enabled bot's prompt should steer it to **search first, then fetch a returned URL**, and to **re-search with a more specific query** (rather than guess a URL) when the exact page it wants isn't in the results.
Web search is **not free** — Anthropic bills per search (plus the tokens the results add) — so scope the allowlist and `max_uses` deliberately, and lean on a prompt that tells the bot to search only when a question actually turns on current or authoritative information. Search activity logs as `[web] ...` lines, and long multi-search turns are handled transparently (the engine follows Anthropic's `pause_turn` continuations up to a small cap, so a slow search doesn't truncate the reply).
> [!TIP]
> For deep request/response debugging, the Anthropic Go SDK ships `option.WithDebugLog(...)` (dumps full HTTP bodies with auth headers redacted). It is not wired into the bot — dev-only, add it temporarily to the client constructor if you ever need wire-level traces.
### Future: persistent memory
The Anthropic memory tool (`memory_20250818`) is a candidate future feature for cross-conversation recall (a self-hosted analog of ChatGPT's "memory"). The Go SDK already ships the types (`BetaMemoryTool20250818Param` and its tool-union slot plus the six-command union: `view`/`create`/`str_replace`/`insert`/`delete`/`rename`), but — unlike the Python/TypeScript/Java SDKs — provides **no handler helper**: the bot would have to hand-write client-side command dispatch against per-chat storage, including strict path validation (canonicalize and confine every model-supplied path under a fixed memory root; reject `..`/symlink traversal) and a no-secrets policy for stored content. Not implemented yet.
## Systemd Unit Setup
To enable the bot to start automatically on system boot and run in the background, set up a systemd unit.
1. Copy the systemd unit template and edit it:
```bash
sudo cp examples/systemd/telegram-bot.service /etc/systemd/system/telegram-bot.service
```
Edit the service file:
```bash
sudo nano /etc/systemd/system/telegram-bot.service
```
Adjust the following parameters:
- WorkingDirectory
- ExecStart
- User
2. Enable and start the service:
```bash
sudo systemctl daemon-reload
```
```bash
sudo systemctl enable telegram-bot
```
```bash
sudo systemctl start telegram-bot
```
3. Check the status:
```bash
sudo systemctl status telegram-bot
```
For more details on the systemd setup, refer to the [demo service file](examples/systemd/telegram-bot.service).
## Logs
### Docker
```bash
docker-compose logs -f telegram-bot
```
### Systemd
```bash
journalctl -u telegram-bot -f
```
## Commands
| Command | Access | Description |
| --------------------------------- | ----------- | ------------------------------------------------------------ |
| `/stats` | All users | Show global bot statistics (total users and messages) |
| `/stats user` | All users | Show your own message statistics |
| `/stats user <user_id>` | Admin/Owner | Show statistics for a specific user |
| `/whoami` | All users | Show your Telegram ID, username, and role |
| `/clear` | All users | Soft-delete your own chat history |
| `/clear <user_id>` | Admin/Owner | Soft-delete all messages for a user across every chat |
| `/clear <user_id> <chat_id>` | Admin/Owner | Soft-delete a user's messages in a specific chat |
| `/clear_hard` | All users | Permanently delete your own chat history |
| `/clear_hard <user_id>` | Admin/Owner | Permanently delete all messages for a user across every chat |
| `/clear_hard <user_id> <chat_id>` | Admin/Owner | Permanently delete a user's messages in a specific chat |
| `/set_model <model-id>` | Admin/Owner | Switch the AI model live without restarting |
> **Note:** In private DMs each user's `chat_id` equals their `user_id`. The scoped `<chat_id>` form is mainly useful for group chat moderation.
## Testing
The GitHub actions workflow already runs tests on every commit:
> [![CI](https://github.com/HugeFrog24/go-telegram-bot/actions/workflows/go-ci.yaml/badge.svg?branch=main)](https://github.com/HugeFrog24/go-telegram-bot/actions/workflows/go-ci.yaml)
However, you can run the tests locally using:
```bash
go test -race -v ./...
```
## Storage
At the moment, a SQLite database (`./data/bot.db`) is used for persistent storage.
Remember to back it up regularly.
Future versions will support more robust storage backends.
+88
View File
@@ -0,0 +1,88 @@
package main
import (
"context"
"sort"
"time"
"github.com/go-telegram/bot/models"
)
const albumFlushWindow = 1 * time.Second
type pendingAlbum struct {
items []*models.Message
chatID, userID int64
username, firstName, lastName, languageCode string
isPremium bool
messageTime int
businessConnectionID string
timer *time.Timer
}
func (b *Bot) bufferAlbumItem(
ctx context.Context,
msg *models.Message,
chatID, userID int64,
username, firstName, lastName string,
isPremium bool,
languageCode string,
messageTime int,
businessConnectionID string,
) {
b.albumBuffersMu.Lock()
defer b.albumBuffersMu.Unlock()
album, exists := b.albumBuffers[msg.MediaGroupID]
if !exists {
album = &pendingAlbum{
chatID: chatID,
userID: userID,
username: username,
firstName: firstName,
lastName: lastName,
isPremium: isPremium,
languageCode: languageCode,
messageTime: messageTime,
businessConnectionID: businessConnectionID,
}
b.albumBuffers[msg.MediaGroupID] = album
}
album.items = append(album.items, msg)
if album.timer != nil {
album.timer.Stop()
}
mediaGroupID := msg.MediaGroupID
album.timer = time.AfterFunc(albumFlushWindow, func() {
b.flushAlbum(ctx, mediaGroupID)
})
}
func (b *Bot) flushAlbum(ctx context.Context, mediaGroupID string) {
b.albumBuffersMu.Lock()
album, exists := b.albumBuffers[mediaGroupID]
if !exists {
b.albumBuffersMu.Unlock()
return
}
delete(b.albumBuffers, mediaGroupID)
items := album.items
captured := *album
b.albumBuffersMu.Unlock()
sort.Slice(items, func(i, j int) bool { return items[i].ID < items[j].ID })
if !b.checkRateLimits(captured.userID) {
b.sendRateLimitExceededMessage(ctx, captured.chatID, captured.businessConnectionID)
return
}
b.handlePhotoMessage(
ctx, items,
captured.chatID, captured.userID,
captured.username, captured.firstName, captured.lastName,
captured.isPremium, captured.languageCode, captured.messageTime,
captured.businessConnectionID,
)
}
+435
View File
@@ -0,0 +1,435 @@
package main
import (
"context"
"errors"
"fmt"
"net/http"
"strings"
"sync/atomic"
"time"
"github.com/anthropics/anthropic-sdk-go"
"github.com/anthropics/anthropic-sdk-go/packages/param"
)
var ErrModelNotFound = errors.New("model not found or deprecated")
const maxFileNotFoundRetries = 3
const maxPauseTurnContinuations = 5
const defaultMaxTokens = 1000
const mcpUnsupportedSentinel = "format not currently supported by the Anthropic API"
var mcpUnsupportedCount atomic.Uint64
type mcpCall struct{ server, name, input string }
func (b *Bot) getAnthropicResponse(ctx context.Context, chatID int64, messages []anthropic.BetaMessageParam, isEmojiOnly bool, username string, firstName string, lastName string, isPremium bool, languageCode string, messageTime int, onSegment func(string) error) (string, error) {
staticPrompt := strings.TrimSpace(b.config.SystemPrompts["custom_instructions"])
InfoLogger.Printf("Sending %d messages to Anthropic", len(messages))
maxTokens := int64(defaultMaxTokens)
if b.config.MaxTokens > 0 {
maxTokens = int64(b.config.MaxTokens)
}
params := anthropic.BetaMessageNewParams{
Model: b.config.Model,
MaxTokens: maxTokens,
Messages: messages,
Betas: []anthropic.AnthropicBeta{anthropic.AnthropicBetaFilesAPI2025_04_14},
}
if staticPrompt != "" {
blocks := []anthropic.BetaTextBlockParam{
{Text: staticPrompt, CacheControl: anthropic.NewBetaCacheControlEphemeralParam()},
}
tail := buildUserContext(username, firstName, lastName, isPremium, languageCode, messageTime)
if isEmojiOnly {
if rule := strings.TrimSpace(b.config.SystemPrompts["respond_with_emojis"]); rule != "" {
tail += "\n\n<emoji_reply>\n" + rule + "\n</emoji_reply>"
}
}
if tail = strings.TrimSpace(tail); tail != "" {
blocks = append(blocks, anthropic.BetaTextBlockParam{Text: tail})
}
params.System = blocks
}
if b.config.Temperature != nil {
params.Temperature = param.NewOpt(float64(*b.config.Temperature))
}
if thinking, ok := thinkingParamFromConfig(b.config.Thinking, b.config.ThinkingDisplay); ok {
params.Thinking = thinking
}
var tools []anthropic.BetaToolUnionParam
if len(b.config.MCPServers) > 0 {
mcpServers := make([]anthropic.BetaRequestMCPServerURLDefinitionParam, 0, len(b.config.MCPServers))
for _, s := range b.config.MCPServers {
srv := anthropic.BetaRequestMCPServerURLDefinitionParam{
Name: s.Name,
URL: s.URL,
}
if s.AuthorizationToken != "" {
srv.AuthorizationToken = param.NewOpt(s.AuthorizationToken)
}
mcpServers = append(mcpServers, srv)
toolset := &anthropic.BetaMCPToolsetParam{
MCPServerName: s.Name,
}
if len(s.AllowedTools) > 0 {
toolset.DefaultConfig = anthropic.BetaMCPToolDefaultConfigParam{
Enabled: param.NewOpt(false),
}
toolset.Configs = make(map[string]anthropic.BetaMCPToolConfigParam, len(s.AllowedTools))
for _, tool := range s.AllowedTools {
toolset.Configs[tool] = anthropic.BetaMCPToolConfigParam{
Enabled: param.NewOpt(true),
}
}
}
tools = append(tools, anthropic.BetaToolUnionParam{OfMCPToolset: toolset})
}
params.MCPServers = mcpServers
params.Betas = append(params.Betas, anthropic.AnthropicBetaMCPClient2025_11_20)
}
tools = append(tools, webSearchTools(b.config.WebSearch)...)
if len(tools) > 0 {
params.Tools = tools
}
var fullText strings.Builder
var lastMsg anthropic.BetaMessage
fileRetries, pauseContinuations := 0, 0
for {
joined, msg, streamErr := b.streamMessages(ctx, params, onSegment)
if streamErr != nil {
var apiErr *anthropic.Error
if !errors.As(streamErr, &apiErr) || apiErr.StatusCode != http.StatusNotFound {
return "", fmt.Errorf("error creating Anthropic message: %w", streamErr)
}
missingFileID := extractMissingFileID(streamErr)
if missingFileID == "" {
return "", fmt.Errorf("%w: %s", ErrModelNotFound, b.config.Model)
}
fileRetries++
if fileRetries > maxFileNotFoundRetries {
return "", fmt.Errorf("max self-heal retries (%d) exceeded: too many file_ids gone from anthropic", maxFileNotFoundRetries)
}
ErrorLogger.Printf("[%s] self-heal: stripping dead file_id %s from chat %d (attempt %d/%d)",
b.config.ID, missingFileID, chatID, fileRetries, maxFileNotFoundRetries)
b.stripDeadFileIDFromMemory(chatID, missingFileID)
if _, cleanupErr := b.markFilesPendingCleanup(ctx, chatID, []string{missingFileID}); cleanupErr != nil {
ErrorLogger.Printf("[%s] mark files pending cleanup: %v", b.config.ID, cleanupErr)
}
params.Messages = b.prepareContextMessages(b.getOrCreateChatMemory(chatID))
continue
}
lastMsg = msg
if joined != "" {
if fullText.Len() > 0 {
fullText.WriteString("\n\n")
}
fullText.WriteString(joined)
}
if msg.StopReason == anthropic.BetaStopReasonPauseTurn {
pauseContinuations++
if pauseContinuations > maxPauseTurnContinuations {
ErrorLogger.Printf("[%s] pause_turn continuations exceeded (%d); returning partial answer",
b.config.ID, maxPauseTurnContinuations)
break
}
params.Messages = append(params.Messages, msg.ToParam())
continue
}
break
}
if fullText.Len() == 0 {
return "", emptyStreamError(string(lastMsg.StopReason),
lastMsg.Usage.OutputTokensDetails.ThinkingTokens, params.MaxTokens)
}
return fullText.String(), nil
}
func webSearchTools(cfg *WebSearchConfig) []anthropic.BetaToolUnionParam {
if cfg == nil {
return nil
}
search := &anthropic.BetaWebSearchTool20250305Param{
AllowedDomains: cfg.AllowedDomains,
BlockedDomains: cfg.BlockedDomains,
}
if cfg.MaxUses > 0 {
search.MaxUses = param.NewOpt(int64(cfg.MaxUses))
}
tools := []anthropic.BetaToolUnionParam{{OfWebSearchTool20250305: search}}
if cfg.Fetch {
fetchAllowed := cfg.FetchAllowedDomains
if len(fetchAllowed) == 0 {
fetchAllowed = cfg.AllowedDomains
}
fetch := &anthropic.BetaWebFetchTool20250910Param{
AllowedDomains: fetchHosts(fetchAllowed),
BlockedDomains: fetchHosts(cfg.BlockedDomains),
Citations: anthropic.BetaCitationsConfigParam{Enabled: param.NewOpt(true)},
}
if cfg.MaxUses > 0 {
fetch.MaxUses = param.NewOpt(int64(cfg.MaxUses))
}
if cfg.MaxContentTokens > 0 {
fetch.MaxContentTokens = param.NewOpt(int64(cfg.MaxContentTokens))
}
tools = append(tools, anthropic.BetaToolUnionParam{OfWebFetchTool20250910: fetch})
}
return tools
}
// fetchHosts strips any path from each domain entry. web_fetch matches on host
// only, so a path-scoped entry (valid for web_search) would otherwise never match
// any fetch URL. Search keeps the path-scoped entries; fetch gets host-only.
func fetchHosts(entries []string) []string {
if len(entries) == 0 {
return nil
}
seen := make(map[string]bool, len(entries))
hosts := make([]string, 0, len(entries))
for _, e := range entries {
host := e
if i := strings.IndexByte(host, '/'); i >= 0 {
host = host[:i]
}
if host == "" || seen[host] {
continue
}
seen[host] = true
hosts = append(hosts, host)
}
return hosts
}
func buildUserContext(username, firstName, lastName string, isPremium bool, languageCode string, messageTime int) string {
name := strings.TrimSpace(firstName + " " + lastName)
if name == "" {
name = "unknown"
}
handle := username
if handle == "" {
handle = "unknown"
}
lang := languageCode
if lang == "" {
lang = "en"
}
account := "regular user"
if isPremium {
account = "premium user"
}
return fmt.Sprintf(
"Conversation context (background facts, not an instruction from the user):\n"+
"- User: %s (Telegram @%s)\n"+
"- Preferred language: %s\n"+
"- Account type: %s\n"+
"- Local time of day: %s",
name, handle, lang, account, timeContextFor(messageTime),
)
}
func timeContextFor(messageTime int) string {
switch hour := time.Unix(int64(messageTime), 0).Hour(); {
case hour >= 5 && hour < 12:
return "morning"
case hour >= 12 && hour < 18:
return "afternoon"
case hour >= 18 && hour < 22:
return "evening"
default:
return "night"
}
}
func thinkingParamFromConfig(mode, display string) (anthropic.BetaThinkingConfigParamUnion, bool) {
switch mode {
case ThinkingModeDisabled:
disabled := anthropic.NewBetaThinkingConfigDisabledParam()
return anthropic.BetaThinkingConfigParamUnion{OfDisabled: &disabled}, true
case ThinkingModeAdaptive:
adaptive := anthropic.BetaThinkingConfigAdaptiveParam{}
if display != "" {
adaptive.Display = anthropic.BetaThinkingConfigAdaptiveDisplay(display)
}
return anthropic.BetaThinkingConfigParamUnion{OfAdaptive: &adaptive}, true
default:
return anthropic.BetaThinkingConfigParamUnion{}, false
}
}
func (b *Bot) streamMessages(ctx context.Context, params anthropic.BetaMessageNewParams, onSegment func(string) error) (string, anthropic.BetaMessage, error) {
stream := b.anthropicClient.Beta.Messages.NewStreaming(ctx, params)
defer func() {
if err := stream.Close(); err != nil {
ErrorLogger.Printf("[stream] close failed: %v", err)
}
}()
var (
message anthropic.BetaMessage
allSegments []string
currentKind string
currentText strings.Builder
currentThinking strings.Builder
currentInputJSON strings.Builder
currentTUseName, currentTUseServer, currentTUseID string
currentTResultUseID, currentTResultServer string
currentTResultIsError bool
currentTResultContent string
currentServerToolName, currentServerToolID string
currentServerResult string
mcpCalls = map[string]mcpCall{}
)
for stream.Next() {
e := stream.Current()
if accErr := message.Accumulate(e); accErr != nil {
ErrorLogger.Printf("[stream] accumulate failed: %v", accErr)
}
switch e.Type {
case "content_block_start":
cbs := e.AsContentBlockStart()
currentKind = cbs.ContentBlock.Type
currentText.Reset()
currentThinking.Reset()
currentInputJSON.Reset()
currentServerResult = ""
switch currentKind {
case "mcp_tool_use":
currentTUseName = cbs.ContentBlock.Name
currentTUseServer = cbs.ContentBlock.ServerName
currentTUseID = cbs.ContentBlock.ID
case "mcp_tool_result":
currentTResultUseID = cbs.ContentBlock.ToolUseID
currentTResultServer = cbs.ContentBlock.ServerName
currentTResultIsError = cbs.ContentBlock.IsError
currentTResultContent = cbs.ContentBlock.JSON.Content.Raw()
case "server_tool_use":
currentServerToolName = cbs.ContentBlock.Name
currentServerToolID = cbs.ContentBlock.ID
case "web_search_tool_result", "web_fetch_tool_result":
currentServerResult = cbs.ContentBlock.JSON.Content.Raw()
}
case "content_block_delta":
cbd := e.AsContentBlockDelta()
switch cbd.Delta.Type {
case "text_delta":
if currentKind == "text" {
currentText.WriteString(cbd.Delta.Text)
}
case "thinking_delta":
if currentKind == "thinking" {
currentThinking.WriteString(cbd.Delta.Thinking)
}
case "input_json_delta":
if currentKind == "mcp_tool_use" || currentKind == "server_tool_use" {
currentInputJSON.WriteString(cbd.Delta.PartialJSON)
}
}
case "content_block_stop":
switch currentKind {
case "text":
seg := strings.TrimSpace(currentText.String())
if seg != "" {
allSegments = append(allSegments, seg)
if onSegment != nil {
if cbErr := onSegment(seg); cbErr != nil {
ErrorLogger.Printf("[stream] onSegment failed: %v", cbErr)
}
}
}
case "mcp_tool_use":
mcpCalls[currentTUseID] = mcpCall{
server: currentTUseServer,
name: currentTUseName,
input: currentInputJSON.String(),
}
InfoLogger.Printf("[mcp] tool_use server=%q name=%q id=%q input=%s",
currentTUseServer, currentTUseName, currentTUseID, currentInputJSON.String())
case "mcp_tool_result":
preview := currentTResultContent
if len(preview) > 500 {
preview = preview[:500] + "...(truncated)"
}
InfoLogger.Printf("[mcp] tool_result tool_use_id=%q server=%q is_error=%v content=%s",
currentTResultUseID, currentTResultServer, currentTResultIsError, preview)
if strings.Contains(currentTResultContent, mcpUnsupportedSentinel) {
total := mcpUnsupportedCount.Add(1)
call := mcpCalls[currentTResultUseID]
ErrorLogger.Printf("[%s][mcp][unsupported] connector could not serialize result "+
"(total=%d): server=%q tool=%q input=%s tool_use_id=%q",
b.config.ID, total, call.server, call.name, call.input, currentTResultUseID)
}
case "server_tool_use":
InfoLogger.Printf("[web] %s id=%q input=%s",
currentServerToolName, currentServerToolID, currentInputJSON.String())
case "web_search_tool_result", "web_fetch_tool_result":
preview := currentServerResult
if len(preview) > 500 {
preview = preview[:500] + "...(truncated)"
}
InfoLogger.Printf("[web] %s content=%s", currentKind, preview)
case "thinking", "redacted_thinking":
if summary := strings.TrimSpace(currentThinking.String()); summary != "" {
if len(summary) > 500 {
summary = summary[:500] + "...(truncated)"
}
InfoLogger.Printf("[thinking] block complete: %s", summary)
} else {
InfoLogger.Printf("[thinking] block complete (content omitted)")
}
default:
if currentKind != "" {
InfoLogger.Printf("[stream] block type=%q (unhandled)", currentKind)
}
}
currentKind = ""
}
}
if err := stream.Err(); err != nil {
return "", message, err
}
stopReason := string(message.StopReason)
if stopReason != "" || message.Usage.OutputTokens > 0 {
InfoLogger.Printf("[usage] model=%s in=%d out=%d thinking=%d stop=%s",
params.Model, message.Usage.InputTokens, message.Usage.OutputTokens,
message.Usage.OutputTokensDetails.ThinkingTokens, stopReason)
if message.StopReason == anthropic.BetaStopReasonMaxTokens {
ErrorLogger.Printf("[usage] response truncated at max_tokens=%d - raise max_tokens (thinking counts toward it)",
params.MaxTokens)
}
}
return strings.Join(allSegments, "\n\n"), message, nil
}
func emptyStreamError(stopReason string, thinkingTokens, maxTokens int64) error {
if stopReason == "max_tokens" {
return fmt.Errorf("output budget exhausted before any text (thinking used %d of %d max_tokens) - raise max_tokens",
thinkingTokens, maxTokens)
}
return fmt.Errorf("unexpected response format from Anthropic")
}
+172
View File
@@ -0,0 +1,172 @@
package main
import (
"bytes"
"context"
"errors"
"fmt"
"net/http"
"strings"
"time"
"github.com/anthropics/anthropic-sdk-go"
)
const fileNotFoundPrefix = "File not found: "
func formatUploadFilename(botID uint, chatID int64, tgMessageID int, ext string) string {
return fmt.Sprintf("tg-%d-%d-%d.%s", botID, chatID, tgMessageID, ext)
}
func (b *Bot) uploadImageToAnthropic(ctx context.Context, data []byte, filename, contentType string) (string, error) {
resp, err := b.anthropicClient.Beta.Files.Upload(ctx, anthropic.BetaFileUploadParams{
File: anthropic.File(bytes.NewReader(data), filename, contentType),
Betas: []anthropic.AnthropicBeta{anthropic.AnthropicBetaFilesAPI2025_04_14},
})
if err != nil {
return "", fmt.Errorf("anthropic files upload: %w", err)
}
return resp.ID, nil
}
func (b *Bot) deleteFileFromAnthropic(ctx context.Context, fileID string) error {
_, err := b.anthropicClient.Beta.Files.Delete(ctx, fileID, anthropic.BetaFileDeleteParams{
Betas: []anthropic.AnthropicBeta{anthropic.AnthropicBetaFilesAPI2025_04_14},
})
if err == nil {
return nil
}
var apiErr *anthropic.Error
if errors.As(err, &apiErr) && apiErr.StatusCode == http.StatusNotFound {
return nil
}
return fmt.Errorf("anthropic files delete %s: %w", fileID, err)
}
func (b *Bot) compensatingDelete(ctx context.Context, fileIDs []string) {
for _, fid := range fileIDs {
if err := b.deleteFileFromAnthropic(ctx, fid); err != nil {
ErrorLogger.Printf("[%s] compensating delete for %s: %v", b.config.ID, fid, err)
}
}
}
func extractMissingFileID(err error) string {
if err == nil {
return ""
}
var apiErr *anthropic.Error
if !errors.As(err, &apiErr) {
return ""
}
if apiErr.StatusCode != http.StatusNotFound {
return ""
}
return parseMissingFileIDFromBody(apiErr.RawJSON())
}
func parseMissingFileIDFromBody(raw string) string {
idx := strings.Index(raw, fileNotFoundPrefix)
if idx == -1 {
return ""
}
rest := raw[idx+len(fileNotFoundPrefix):]
end := strings.IndexFunc(rest, func(r rune) bool {
return (r < 'a' || r > 'z') &&
(r < 'A' || r > 'Z') &&
(r < '0' || r > '9') &&
r != '_'
})
if end == -1 {
return rest
}
return rest[:end]
}
func (b *Bot) hardDeleteScope(ctx context.Context, query string, args ...interface{}) error {
var rows []Message
if err := b.db.Unscoped().Where(query, args...).Find(&rows).Error; err != nil {
return fmt.Errorf("scan rows: %w", err)
}
if len(rows) == 0 {
return nil
}
if err := b.db.Where(query, args...).Delete(&Message{}).Error; err != nil {
return fmt.Errorf("soft delete: %w", err)
}
hardDeletable := make([]uint, 0, len(rows))
for _, row := range rows {
if b.deleteRowFiles(ctx, row) {
hardDeletable = append(hardDeletable, row.ID)
}
}
if len(hardDeletable) == 0 {
return nil
}
if err := b.db.Unscoped().Where("id IN ?", hardDeletable).Delete(&Message{}).Error; err != nil {
return fmt.Errorf("hard delete: %w", err)
}
return nil
}
func (b *Bot) deleteRowFiles(ctx context.Context, row Message) bool {
if len(row.ImageFileIDs) == 0 {
return true
}
allOk := true
for _, fid := range row.ImageFileIDs {
if err := b.deleteFileFromAnthropic(ctx, fid); err != nil {
ErrorLogger.Printf("[%s] anthropic delete %s (row %d): %v", b.config.ID, fid, row.ID, err)
allOk = false
}
}
return allOk
}
func stripDeadFileIDs(src []string, deadSet map[string]struct{}) (survivors []string, dirty bool) {
survivors = make([]string, 0, len(src))
for _, fid := range src {
if _, dead := deadSet[fid]; dead {
dirty = true
continue
}
survivors = append(survivors, fid)
}
return survivors, dirty
}
func (b *Bot) markFilesPendingCleanup(ctx context.Context, chatID int64, deadFileIDs []string) (int, error) {
if len(deadFileIDs) == 0 {
return 0, nil
}
deadSet := make(map[string]struct{}, len(deadFileIDs))
for _, id := range deadFileIDs {
deadSet[id] = struct{}{}
}
var rows []Message
if err := b.db.WithContext(ctx).
Where("bot_id = ? AND chat_id = ? AND image_file_ids IS NOT NULL", b.botID, chatID).
Find(&rows).Error; err != nil {
return 0, fmt.Errorf("scan rows for cleanup: %w", err)
}
now := time.Now()
updated := 0
for _, row := range rows {
survivors, dirty := stripDeadFileIDs(row.ImageFileIDs, deadSet)
if !dirty {
continue
}
if len(survivors) == 0 {
row.ImageFileIDs = nil
row.FilesCleanedAt = &now
} else {
row.ImageFileIDs = survivors
}
if err := b.db.WithContext(ctx).Save(&row).Error; err != nil {
return updated, fmt.Errorf("update row %d: %w", row.ID, err)
}
updated++
}
return updated, nil
}
+191
View File
@@ -0,0 +1,191 @@
package main
import (
"testing"
"time"
"github.com/stretchr/testify/assert"
)
func TestFormatUploadFilename(t *testing.T) {
cases := []struct {
botID uint
chatID int64
tgMessageID int
ext string
want string
}{
{1, 12345, 42, "jpg", "tg-1-12345-42.jpg"},
{7, -1001234567890, 1, "png", "tg-7--1001234567890-1.png"},
{0, 0, 0, "webp", "tg-0-0-0.webp"},
}
for _, tc := range cases {
got := formatUploadFilename(tc.botID, tc.chatID, tc.tgMessageID, tc.ext)
assert.Equal(t, tc.want, got)
}
}
func TestParseMissingFileIDFromBody(t *testing.T) {
cases := []struct {
name string
body string
want string
}{
{
name: "canonical Anthropic file-not-found body",
body: `{"type":"error","error":{"type":"invalid_request_error","message":"File not found: file_011CNha8iCJcU1wXNR6q4V8w"}}`,
want: "file_011CNha8iCJcU1wXNR6q4V8w",
},
{
name: "trailing punctuation after the id is excluded",
body: `something File not found: file_abc123! more text`,
want: "file_abc123",
},
{
name: "body without the prefix yields empty",
body: `{"type":"error","error":{"message":"Model not found: claude-foo"}}`,
want: "",
},
{
name: "id at the very end of the buffer",
body: `File not found: file_xyz789`,
want: "file_xyz789",
},
{
name: "empty body",
body: "",
want: "",
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
assert.Equal(t, tc.want, parseMissingFileIDFromBody(tc.body))
})
}
}
func TestStripDeadFileIDs(t *testing.T) {
dead := map[string]struct{}{
"file_a": {},
"file_b": {},
}
cases := []struct {
name string
input []string
wantSurvivors []string
wantDirty bool
}{
{
name: "no overlap returns input verbatim",
input: []string{"file_x", "file_y"},
wantSurvivors: []string{"file_x", "file_y"},
wantDirty: false,
},
{
name: "partial overlap returns survivors and reports dirty",
input: []string{"file_a", "file_x", "file_b", "file_y"},
wantSurvivors: []string{"file_x", "file_y"},
wantDirty: true,
},
{
name: "all dead returns empty survivors and dirty",
input: []string{"file_a", "file_b"},
wantSurvivors: []string{},
wantDirty: true,
},
{
name: "empty input is not dirty",
input: []string{},
wantSurvivors: []string{},
wantDirty: false,
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
survivors, dirty := stripDeadFileIDs(tc.input, dead)
assert.Equal(t, tc.wantSurvivors, survivors)
assert.Equal(t, tc.wantDirty, dirty)
})
}
}
func TestMarkFilesPendingCleanup(t *testing.T) {
b, _ := setupBotForTest(t, 123)
chatID := int64(555)
row1 := Message{
BotID: b.botID,
ChatID: chatID,
UserID: 777,
Username: "u",
UserRole: "user",
Text: "look at these",
Timestamp: time.Now(),
IsUser: true,
ImageFileIDs: []string{"file_a", "file_x"},
}
assert.NoError(t, b.db.Create(&row1).Error)
row2 := Message{
BotID: b.botID,
ChatID: chatID,
UserID: 777,
Username: "u",
UserRole: "user",
Text: "screenshot",
Timestamp: time.Now(),
IsUser: true,
ImageFileIDs: []string{"file_a", "file_b"},
}
assert.NoError(t, b.db.Create(&row2).Error)
row3 := Message{
BotID: b.botID,
ChatID: chatID,
UserID: 777,
Username: "u",
UserRole: "user",
Text: "another",
Timestamp: time.Now(),
IsUser: true,
ImageFileIDs: []string{"file_x", "file_y"},
}
assert.NoError(t, b.db.Create(&row3).Error)
row4 := Message{
BotID: b.botID,
ChatID: 999,
UserID: 777,
Username: "u",
UserRole: "user",
Text: "other chat",
Timestamp: time.Now(),
IsUser: true,
ImageFileIDs: []string{"file_a"},
}
assert.NoError(t, b.db.Create(&row4).Error)
updated, err := b.markFilesPendingCleanup(t.Context(), chatID, []string{"file_a", "file_b"})
assert.NoError(t, err)
assert.Equal(t, 2, updated, "rows 1 and 2 should have been updated")
var r1 Message
assert.NoError(t, b.db.First(&r1, row1.ID).Error)
assert.Equal(t, []string{"file_x"}, r1.ImageFileIDs)
assert.Nil(t, r1.FilesCleanedAt)
var r2 Message
assert.NoError(t, b.db.First(&r2, row2.ID).Error)
assert.Empty(t, r2.ImageFileIDs)
assert.NotNil(t, r2.FilesCleanedAt)
var r3 Message
assert.NoError(t, b.db.First(&r3, row3.ID).Error)
assert.Equal(t, []string{"file_x", "file_y"}, r3.ImageFileIDs)
assert.Nil(t, r3.FilesCleanedAt)
var r4 Message
assert.NoError(t, b.db.First(&r4, row4.ID).Error)
assert.Equal(t, []string{"file_a"}, r4.ImageFileIDs)
assert.Nil(t, r4.FilesCleanedAt)
}
+310
View File
@@ -0,0 +1,310 @@
package main
import (
"encoding/json"
"strings"
"testing"
"time"
"github.com/anthropics/anthropic-sdk-go"
)
func TestTimeContextFor(t *testing.T) {
cases := []struct {
hour int
expected string
}{
{3, "night"},
{5, "morning"},
{11, "morning"},
{12, "afternoon"},
{17, "afternoon"},
{18, "evening"},
{21, "evening"},
{22, "night"},
{23, "night"},
}
for _, tc := range cases {
ts := int(time.Date(2025, 5, 15, tc.hour, 0, 0, 0, time.Local).Unix())
if got := timeContextFor(ts); got != tc.expected {
t.Errorf("timeContextFor(hour=%d) = %q, want %q", tc.hour, got, tc.expected)
}
}
}
func TestBuildUserContext(t *testing.T) {
noon := int(time.Date(2025, 5, 15, 12, 0, 0, 0, time.Local).Unix())
got := buildUserContext("alice", "Alice", "Smith", true, "de", noon)
for _, want := range []string{"Alice Smith", "@alice", "Preferred language: de", "premium user", "afternoon"} {
if !strings.Contains(got, want) {
t.Errorf("buildUserContext premium: missing %q in:\n%s", want, got)
}
}
got = buildUserContext("", "", "", false, "", noon)
for _, want := range []string{"User: unknown (Telegram @unknown)", "Preferred language: en", "regular user"} {
if !strings.Contains(got, want) {
t.Errorf("buildUserContext fallback: missing %q in:\n%s", want, got)
}
}
got = buildUserContext("bob", "Bob", "", false, "en", noon)
if !strings.Contains(got, "User: Bob (Telegram @bob)") {
t.Errorf("buildUserContext firstname-only: got:\n%s", got)
}
}
func TestThinkingParamFromConfig(t *testing.T) {
cases := []struct {
name string
mode string
display string
ok bool
want map[string]any
}{
{"unset omits param", "", "", false, nil},
{"unknown value omits param", "bogus", "", false, nil},
{"adaptive no display", ThinkingModeAdaptive, "", true,
map[string]any{"type": "adaptive"}},
{"adaptive summarized", ThinkingModeAdaptive, ThinkingDisplaySummarized, true,
map[string]any{"type": "adaptive", "display": "summarized"}},
{"adaptive omitted", ThinkingModeAdaptive, ThinkingDisplayOmitted, true,
map[string]any{"type": "adaptive", "display": "omitted"}},
{"disabled", ThinkingModeDisabled, "", true,
map[string]any{"type": "disabled"}},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
union, ok := thinkingParamFromConfig(tc.mode, tc.display)
if ok != tc.ok {
t.Fatalf("ok = %v, want %v", ok, tc.ok)
}
if !tc.ok {
return
}
raw, err := json.Marshal(union)
if err != nil {
t.Fatalf("marshal: %v", err)
}
var got map[string]any
if err := json.Unmarshal(raw, &got); err != nil {
t.Fatalf("unmarshal %s: %v", raw, err)
}
if len(got) != len(tc.want) {
t.Fatalf("wire shape %s: got %d keys, want %d (%v)", raw, len(got), len(tc.want), tc.want)
}
for k, v := range tc.want {
if got[k] != v {
t.Errorf("wire shape %s: key %q = %v, want %v", raw, k, got[k], v)
}
}
})
}
}
func TestBackwardCompatibleParams(t *testing.T) {
params := anthropic.BetaMessageNewParams{
Model: "claude-test",
MaxTokens: defaultMaxTokens,
Messages: []anthropic.BetaMessageParam{
anthropic.NewBetaUserMessage(anthropic.NewBetaTextBlock("hi")),
},
}
raw, err := json.Marshal(params)
if err != nil {
t.Fatalf("marshal: %v", err)
}
var got map[string]any
if err := json.Unmarshal(raw, &got); err != nil {
t.Fatalf("unmarshal: %v", err)
}
if _, present := got["thinking"]; present {
t.Errorf("zero Thinking union must omit the key; body: %s", raw)
}
if mt, ok := got["max_tokens"].(float64); !ok || int(mt) != defaultMaxTokens {
t.Errorf("max_tokens = %v, want %d; body: %s", got["max_tokens"], defaultMaxTokens, raw)
}
}
func TestWebSearchTools(t *testing.T) {
t.Run("nil config yields no tools", func(t *testing.T) {
if tools := webSearchTools(nil); tools != nil {
t.Errorf("webSearchTools(nil) = %v, want nil", tools)
}
})
t.Run("search only when fetch off", func(t *testing.T) {
tools := webSearchTools(&WebSearchConfig{
AllowedDomains: []string{"example.com/hc"},
MaxUses: 3,
})
if len(tools) != 1 {
t.Fatalf("got %d tools, want 1 (search only)", len(tools))
}
if tools[0].OfWebSearchTool20250305 == nil {
t.Fatalf("tools[0] is not a web_search tool")
}
if tools[0].OfWebFetchTool20250910 != nil {
t.Error("web_fetch tool present but fetch is off")
}
})
t.Run("search + fetch with allowlist and citations", func(t *testing.T) {
tools := webSearchTools(&WebSearchConfig{
AllowedDomains: []string{"example.com/hc", "docs.example.com"},
MaxUses: 3,
Fetch: true,
MaxContentTokens: 50000,
})
if len(tools) != 2 {
t.Fatalf("got %d tools, want 2 (search + fetch)", len(tools))
}
search := tools[0].OfWebSearchTool20250305
if search == nil {
t.Fatalf("tools[0] is not a web_search tool")
}
if !sameStrings(search.AllowedDomains, []string{"example.com/hc", "docs.example.com"}) {
t.Errorf("search AllowedDomains = %v, want the path-scoped list unchanged", search.AllowedDomains)
}
if search.MaxUses.Value != 3 {
t.Errorf("search MaxUses = %d, want 3", search.MaxUses.Value)
}
fetch := tools[1].OfWebFetchTool20250910
if fetch == nil {
t.Fatalf("tools[1] is not a web_fetch tool")
}
if !sameStrings(fetch.AllowedDomains, []string{"example.com", "docs.example.com"}) {
t.Errorf("fetch AllowedDomains = %v, want host-only [example.com docs.example.com]", fetch.AllowedDomains)
}
if fetch.MaxContentTokens.Value != 50000 {
t.Errorf("fetch MaxContentTokens = %d, want 50000", fetch.MaxContentTokens.Value)
}
if fetch.Citations.Enabled.Value != true {
t.Error("fetch citations not enabled")
}
})
t.Run("fetch hosts are deduped", func(t *testing.T) {
tools := webSearchTools(&WebSearchConfig{
AllowedDomains: []string{"a.com/x", "a.com/y", "b.com"},
Fetch: true,
})
fetch := tools[1].OfWebFetchTool20250910
if fetch == nil {
t.Fatalf("tools[1] is not a web_fetch tool")
}
if !sameStrings(fetch.AllowedDomains, []string{"a.com", "b.com"}) {
t.Errorf("fetch AllowedDomains = %v, want deduped [a.com b.com]", fetch.AllowedDomains)
}
})
t.Run("social host is search-only via fetch_allowed_domains", func(t *testing.T) {
tools := webSearchTools(&WebSearchConfig{
AllowedDomains: []string{"helpshift.example/hc", "x.com/thatskygame"},
FetchAllowedDomains: []string{"helpshift.example"},
Fetch: true,
})
search := tools[0].OfWebSearchTool20250305
if search == nil {
t.Fatalf("tools[0] is not a web_search tool")
}
var searchHasSocial bool
for _, d := range search.AllowedDomains {
if d == "x.com/thatskygame" {
searchHasSocial = true
}
}
if !searchHasSocial {
t.Errorf("search AllowedDomains = %v, want it to include x.com/thatskygame", search.AllowedDomains)
}
fetch := tools[1].OfWebFetchTool20250910
if fetch == nil {
t.Fatalf("tools[1] is not a web_fetch tool")
}
if !sameStrings(fetch.AllowedDomains, []string{"helpshift.example"}) {
t.Errorf("fetch AllowedDomains = %v, want only [helpshift.example]", fetch.AllowedDomains)
}
for _, d := range fetch.AllowedDomains {
if strings.HasPrefix(d, "x.com") {
t.Errorf("fetch AllowedDomains leaked the social host: %v", fetch.AllowedDomains)
}
}
})
t.Run("wire shape carries allowed_domains", func(t *testing.T) {
tools := webSearchTools(&WebSearchConfig{
AllowedDomains: []string{"thatgamecompany.helpshift.com/hc"},
MaxUses: 2,
Fetch: true,
})
raw, err := json.Marshal(tools)
if err != nil {
t.Fatalf("marshal: %v", err)
}
body := string(raw)
for _, want := range []string{
"web_search_20250305",
"web_fetch_20250910",
"thatgamecompany.helpshift.com/hc",
"allowed_domains",
} {
if !strings.Contains(body, want) {
t.Errorf("wire body missing %q:\n%s", want, body)
}
}
})
}
func sameStrings(got, want []string) bool {
if len(got) != len(want) {
return false
}
for i := range got {
if got[i] != want[i] {
return false
}
}
return true
}
func TestFetchHosts(t *testing.T) {
cases := []struct {
name string
in []string
want []string
}{
{"nil in nil out", nil, nil},
{"empty in nil out", []string{}, nil},
{"host passthrough", []string{"example.com"}, []string{"example.com"}},
{"strip path", []string{"example.com/hc/en"}, []string{"example.com"}},
{"dedup after strip", []string{"a.com/x", "a.com/y"}, []string{"a.com"}},
{"preserve order and subdomains", []string{"docs.example.com/a", "example.com"}, []string{"docs.example.com", "example.com"}},
{"drop empty leading slash", []string{"/oops", "ok.com"}, []string{"ok.com"}},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
if got := fetchHosts(tc.in); !sameStrings(got, tc.want) {
t.Errorf("fetchHosts(%v) = %v, want %v", tc.in, got, tc.want)
}
})
}
}
func TestEmptyStreamError(t *testing.T) {
err := emptyStreamError("max_tokens", 3900, 4000)
for _, want := range []string{"output budget exhausted", "3900", "4000"} {
if !strings.Contains(err.Error(), want) {
t.Errorf("max_tokens case: %q missing %q", err.Error(), want)
}
}
if got := emptyStreamError("end_turn", 0, 1000).Error(); got != "unexpected response format from Anthropic" {
t.Errorf("generic case = %q", got)
}
if got := emptyStreamError("", 0, 1000).Error(); got != "unexpected response format from Anthropic" {
t.Errorf("no-stop-reason case = %q", got)
}
}
+750
View File
@@ -0,0 +1,750 @@
package main
import (
"context"
"errors"
"fmt"
"strings"
"sync"
"time"
"github.com/anthropics/anthropic-sdk-go"
"github.com/anthropics/anthropic-sdk-go/option"
"github.com/go-telegram/bot"
"github.com/go-telegram/bot/models"
"gorm.io/gorm"
)
type Bot struct {
tgBot TelegramClient
db *gorm.DB
anthropicClient anthropic.Client
chatMemories map[int64]*ChatMemory
memorySize int
chatMemoriesMu sync.RWMutex
config BotConfig
userLimiters map[int64]*userLimiter
userLimitersMu sync.RWMutex
clock Clock
botID uint
albumBuffers map[string]*pendingAlbum
albumBuffersMu sync.Mutex
}
func messageType(msg *models.Message) string {
if msg.Sticker != nil {
return "sticker"
}
return "text"
}
func NewBot(db *gorm.DB, config BotConfig, clock Clock, tgClient TelegramClient) (*Bot, error) {
var botEntry BotModel
err := db.Where("identifier = ?", config.ID).First(&botEntry).Error
if errors.Is(err, gorm.ErrRecordNotFound) {
botEntry = BotModel{Identifier: config.ID, Name: config.ID}
if err := db.Create(&botEntry).Error; err != nil {
return nil, err
}
} else if err != nil {
return nil, err
}
var owner User
err = db.Where("telegram_id = ? AND bot_id = ?", config.OwnerTelegramID, botEntry.ID).First(&owner).Error
if errors.Is(err, gorm.ErrRecordNotFound) {
var ownerRole Role
err := db.Where("name = ?", "owner").First(&ownerRole).Error
if err != nil {
return nil, fmt.Errorf("owner role not found: %w", err)
}
owner = User{
BotID: botEntry.ID,
TelegramID: config.OwnerTelegramID,
Username: "",
RoleID: ownerRole.ID,
IsOwner: true,
}
if err := db.Create(&owner).Error; err != nil {
if strings.Contains(err.Error(), "unique index") {
return nil, fmt.Errorf("an owner already exists for this bot")
}
return nil, fmt.Errorf("failed to create owner user: %w", err)
}
} else if err != nil {
return nil, err
}
anthropicClient := anthropic.NewClient(option.WithAPIKey(config.AnthropicAPIKey))
b := &Bot{
db: db,
anthropicClient: anthropicClient,
chatMemories: make(map[int64]*ChatMemory),
memorySize: config.MemorySize,
config: config,
userLimiters: make(map[int64]*userLimiter),
clock: clock,
botID: botEntry.ID,
tgBot: tgClient,
albumBuffers: make(map[string]*pendingAlbum),
}
if tgClient == nil {
var err error
tgClient, err = initTelegramBot(config.TelegramToken, b)
if err != nil {
return nil, fmt.Errorf("failed to initialize Telegram bot: %w", err)
}
b.tgBot = tgClient
}
return b, nil
}
func (b *Bot) Start(ctx context.Context) {
b.tgBot.Start(ctx)
}
func (b *Bot) getOrCreateUser(userID int64, username string, isOwner bool) (User, error) {
var user User
err := b.db.Preload("Role").Where("telegram_id = ? AND bot_id = ?", userID, b.botID).First(&user).Error
if err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
if isOwner {
var existingOwner User
err := b.db.Where("bot_id = ? AND is_owner = ?", b.botID, true).First(&existingOwner).Error
if err == nil {
return User{}, fmt.Errorf("an owner already exists for this bot")
} else if !errors.Is(err, gorm.ErrRecordNotFound) {
return User{}, fmt.Errorf("failed to check existing owner: %w", err)
}
}
var role Role
var roleName string
if isOwner {
roleName = "owner"
} else {
roleName = "user"
}
err := b.db.Where("name = ?", roleName).First(&role).Error
if err != nil {
return User{}, fmt.Errorf("failed to get role: %w", err)
}
user = User{
BotID: b.botID,
TelegramID: userID,
Username: username,
RoleID: role.ID,
Role: role,
IsOwner: isOwner,
}
if err := b.db.Create(&user).Error; err != nil {
if strings.Contains(err.Error(), "unique index") {
return User{}, fmt.Errorf("an owner already exists for this bot")
}
return User{}, fmt.Errorf("failed to create user: %w", err)
}
} else {
return User{}, err
}
} else {
if isOwner && !user.IsOwner {
return User{}, fmt.Errorf("cannot change existing user to owner")
}
}
return user, nil
}
func (b *Bot) getRoleByName(roleName string) (Role, error) {
var role Role
err := b.db.Where("name = ?", roleName).First(&role).Error
return role, err
}
func (b *Bot) createMessage(chatID, userID int64, username, userRole, text string, isUser bool) Message {
message := Message{
ChatID: chatID,
UserRole: userRole,
Text: text,
Timestamp: time.Now(),
IsUser: isUser,
}
if isUser {
message.UserID = userID
message.Username = username
} else {
message.UserID = 0
message.Username = "AI Assistant"
}
return message
}
func (b *Bot) storeMessage(message *Message) error {
message.BotID = b.botID
return b.db.Create(message).Error
}
func (b *Bot) getOrCreateChatMemory(chatID int64) *ChatMemory {
b.chatMemoriesMu.RLock()
chatMemory, exists := b.chatMemories[chatID]
b.chatMemoriesMu.RUnlock()
if !exists {
b.chatMemoriesMu.Lock()
defer b.chatMemoriesMu.Unlock()
chatMemory, exists = b.chatMemories[chatID]
if !exists {
var count int64
b.db.Model(&Message{}).Where("chat_id = ? AND bot_id = ?", chatID, b.botID).Count(&count)
isNewChat := count == 0
var messages []Message
if !isNewChat {
err := b.db.Where("chat_id = ? AND bot_id = ?", chatID, b.botID).
Order("timestamp desc").
Limit(b.memorySize * 2).
Find(&messages).Error
if err != nil {
ErrorLogger.Printf("Error fetching messages from database: %v", err)
messages = []Message{}
} else {
for i, j := 0, len(messages)-1; i < j; i, j = i+1, j-1 {
messages[i], messages[j] = messages[j], messages[i]
}
}
} else {
messages = []Message{}
}
chatMemory = &ChatMemory{
Messages: messages,
Size: b.memorySize * 2,
}
b.chatMemories[chatID] = chatMemory
}
}
return chatMemory
}
func (b *Bot) stripDeadFileIDFromMemory(chatID int64, deadFileID string) {
b.chatMemoriesMu.Lock()
defer b.chatMemoriesMu.Unlock()
cm, exists := b.chatMemories[chatID]
if !exists {
return
}
for i := range cm.Messages {
if len(cm.Messages[i].ImageFileIDs) == 0 {
continue
}
survivors := make([]string, 0, len(cm.Messages[i].ImageFileIDs))
for _, fid := range cm.Messages[i].ImageFileIDs {
if fid != deadFileID {
survivors = append(survivors, fid)
}
}
cm.Messages[i].ImageFileIDs = survivors
}
}
func (b *Bot) addMessageToChatMemory(chatMemory *ChatMemory, message Message) {
b.chatMemoriesMu.Lock()
defer b.chatMemoriesMu.Unlock()
chatMemory.Messages = append(chatMemory.Messages, message)
if len(chatMemory.Messages) > chatMemory.Size {
chatMemory.Messages = chatMemory.Messages[len(chatMemory.Messages)-chatMemory.Size:]
}
}
func (b *Bot) prepareContextMessages(chatMemory *ChatMemory) []anthropic.BetaMessageParam {
b.chatMemoriesMu.RLock()
defer b.chatMemoriesMu.RUnlock()
InfoLogger.Printf("Chat memory contains %d messages", len(chatMemory.Messages))
for i, msg := range chatMemory.Messages {
InfoLogger.Printf("Message %d: IsUser=%v, Text=%q Images=%d", i, msg.IsUser, msg.Text, len(msg.ImageFileIDs))
}
var contextMessages []anthropic.BetaMessageParam
for _, msg := range chatMemory.Messages {
blocks := contentBlocksForMessage(msg)
if len(blocks) == 0 {
continue
}
var param anthropic.BetaMessageParam
if msg.IsUser {
param = anthropic.NewBetaUserMessage(blocks...)
} else {
param = anthropic.BetaMessageParam{
Role: anthropic.BetaMessageParamRoleAssistant,
Content: blocks,
}
}
contextMessages = append(contextMessages, param)
}
return contextMessages
}
func contentBlocksForMessage(msg Message) []anthropic.BetaContentBlockParamUnion {
var blocks []anthropic.BetaContentBlockParamUnion
if msg.IsUser && len(msg.ImageFileIDs) > 0 {
multi := len(msg.ImageFileIDs) > 1
for i, fileID := range msg.ImageFileIDs {
if multi {
blocks = append(blocks, anthropic.NewBetaTextBlock(fmt.Sprintf("Image %d:", i+1)))
}
blocks = append(blocks, anthropic.NewBetaImageBlock(anthropic.BetaFileImageSourceParam{FileID: fileID}))
}
}
if textContent := strings.TrimSpace(msg.Text); textContent != "" {
blocks = append(blocks, anthropic.NewBetaTextBlock(textContent))
}
return blocks
}
func roleHasScope(role Role, scope string) bool {
for _, s := range role.Scopes {
if s.Name == scope {
return true
}
}
return false
}
func (b *Bot) hasScope(userID int64, scope string) bool {
var user User
if err := b.db.Preload("Role.Scopes").
Where("telegram_id = ? AND bot_id = ?", userID, b.botID).
First(&user).Error; err != nil {
return false
}
if user.IsOwner {
return true
}
return roleHasScope(user.Role, scope)
}
var publicBotCommands = []models.BotCommand{
{Command: "stats", Description: "Get bot statistics. Usage: /stats or /stats user [user_id]"},
{Command: "whoami", Description: "Get your user information"},
{Command: "clear", Description: "Clear chat history (soft delete). Admins: /clear [user_id]"},
}
var adminBotCommands = []models.BotCommand{
{Command: "clear_hard", Description: "Clear chat history (permanently delete). Admins: /clear_hard [user_id]"},
{Command: "set_model", Description: "Switch the AI model (admin/owner only). Usage: /set_model <model-id>"},
}
func (b *Bot) registerAdminCommandsForUser(ctx context.Context, telegramID int64) {
allCommands := make([]models.BotCommand, 0, len(publicBotCommands)+len(adminBotCommands))
allCommands = append(allCommands, publicBotCommands...)
allCommands = append(allCommands, adminBotCommands...)
_, err := b.tgBot.SetMyCommands(ctx, &bot.SetMyCommandsParams{
Commands: allCommands,
Scope: &models.BotCommandScopeChat{ChatID: telegramID},
})
if err != nil {
ErrorLogger.Printf("Failed to register admin commands for user %d: %v", telegramID, err)
}
}
func setElevatedCommands(tgBot TelegramClient, users []User) {
allCommands := make([]models.BotCommand, 0, len(publicBotCommands)+len(adminBotCommands))
allCommands = append(allCommands, publicBotCommands...)
allCommands = append(allCommands, adminBotCommands...)
for _, u := range users {
if u.TelegramID == 0 {
continue
}
if !u.IsOwner && !roleHasScope(u.Role, ScopeModelSet) {
continue
}
_, err := tgBot.SetMyCommands(context.Background(), &bot.SetMyCommandsParams{
Commands: allCommands,
Scope: &models.BotCommandScopeChat{ChatID: u.TelegramID},
})
if err != nil {
ErrorLogger.Printf("Warning: could not set admin commands for user %d: %v", u.TelegramID, err)
}
}
}
func initTelegramBot(token string, b *Bot) (TelegramClient, error) {
opts := []bot.Option{
bot.WithDefaultHandler(b.handleUpdate),
}
tgBot, err := bot.New(token, opts...)
if err != nil {
return nil, err
}
_, err = tgBot.SetMyCommands(context.Background(), &bot.SetMyCommandsParams{
Commands: publicBotCommands,
Scope: &models.BotCommandScopeDefault{},
})
if err != nil {
ErrorLogger.Printf("Error setting default bot commands: %v", err)
return nil, err
}
var allUsers []User
if err := b.db.Preload("Role.Scopes").Where("bot_id = ?", b.botID).Find(&allUsers).Error; err != nil {
ErrorLogger.Printf("Warning: could not query users for command scoping: %v", err)
} else {
setElevatedCommands(tgBot, allUsers)
}
return tgBot, nil
}
func (b *Bot) sendResponse(ctx context.Context, chatID int64, text string, businessConnectionID string) error {
_, err := b.screenOutgoingMessage(chatID, text)
if err != nil {
ErrorLogger.Printf("Error storing assistant message: %v", err)
return err
}
params := &bot.SendMessageParams{
ChatID: chatID,
Text: text,
}
if businessConnectionID != "" {
params.BusinessConnectionID = businessConnectionID
}
_, err = b.tgBot.SendMessage(ctx, params)
if err != nil {
ErrorLogger.Printf("[%s] Error sending message to chat %d with BusinessConnectionID %s: %v",
b.config.ID, chatID, businessConnectionID, err)
return err
}
return nil
}
func (b *Bot) sendOneSegment(ctx context.Context, chatID int64, text, businessConnectionID string) error {
params := &bot.SendMessageParams{
ChatID: chatID,
Text: text,
}
if businessConnectionID != "" {
params.BusinessConnectionID = businessConnectionID
}
if _, err := b.tgBot.SendMessage(ctx, params); err != nil {
ErrorLogger.Printf("[%s] Error sending segment to chat %d with BusinessConnectionID %s: %v",
b.config.ID, chatID, businessConnectionID, err)
return err
}
return nil
}
func (b *Bot) sendStats(ctx context.Context, chatID int64, userID int64, targetUserID int64, businessConnectionID string) {
if targetUserID == 0 {
totalUsers, totalMessages, err := b.getStats()
if err != nil {
ErrorLogger.Printf("Error fetching stats: %v\n", err)
if err := b.sendResponse(ctx, chatID, "Sorry, I couldn't retrieve the stats at this time.", businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending response: %v", err)
}
return
}
statsMessage := fmt.Sprintf(
"📊 Bot Statistics:\n\n"+
"- Total Users: %d\n"+
"- Total Messages: %d",
totalUsers,
totalMessages,
)
if b.hasScope(userID, ScopeStatsViewAny) {
type topEntry struct {
UserID int64
MsgCount int64
}
var top []topEntry
if err := b.db.Model(&Message{}).
Select("user_id, COUNT(*) as msg_count").
Where("bot_id = ? AND is_user = ? AND deleted_at IS NULL", b.botID, true).
Group("user_id").
Order("msg_count DESC").
Limit(3).
Scan(&top).Error; err != nil {
ErrorLogger.Printf("Error fetching top users: %v", err)
} else if len(top) > 0 {
statsMessage += "\n\n🏆 Most Active Users:"
for i, entry := range top {
var u User
if err := b.db.Select("username").Where("telegram_id = ? AND bot_id = ?", entry.UserID, b.botID).First(&u).Error; err != nil {
u.Username = fmt.Sprintf("ID:%d", entry.UserID)
}
name := u.Username
if name == "" {
name = fmt.Sprintf("ID:%d", entry.UserID)
}
statsMessage += fmt.Sprintf("\n%d. @%s — %d messages", i+1, name, entry.MsgCount)
}
}
}
if err := b.sendResponse(ctx, chatID, statsMessage, businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending stats message: %v", err)
}
return
}
if targetUserID != userID {
if !b.hasScope(userID, ScopeStatsViewAny) {
InfoLogger.Printf("User %d attempted to view stats for user %d without permission", userID, targetUserID)
if err := b.sendResponse(ctx, chatID, "Permission denied. Only admins and owners can view other users' statistics.", businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending response: %v", err)
}
return
}
}
username, messagesIn, messagesOut, totalMessages, err := b.getUserStats(targetUserID)
if err != nil {
ErrorLogger.Printf("Error fetching user stats: %v\n", err)
if err := b.sendResponse(ctx, chatID, fmt.Sprintf("Sorry, I couldn't retrieve statistics for user ID %d.", targetUserID), businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending response: %v", err)
}
return
}
userInfo := fmt.Sprintf("@%s (ID: %d)", username, targetUserID)
if username == "" {
userInfo = fmt.Sprintf("User ID: %d", targetUserID)
}
statsMessage := fmt.Sprintf(
"👤 User Statistics for %s:\n\n"+
"- Messages Sent: %d\n"+
"- Messages Received: %d\n"+
"- Total Messages: %d",
userInfo,
messagesIn,
messagesOut,
totalMessages,
)
if err := b.sendResponse(ctx, chatID, statsMessage, businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending user stats message: %v", err)
}
}
func (b *Bot) getStats() (int64, int64, error) {
var totalUsers int64
if err := b.db.Model(&User{}).Where("bot_id = ?", b.botID).Count(&totalUsers).Error; err != nil {
return 0, 0, err
}
var totalMessages int64
if err := b.db.Model(&Message{}).Where("bot_id = ?", b.botID).Count(&totalMessages).Error; err != nil {
return 0, 0, err
}
return totalUsers, totalMessages, nil
}
func (b *Bot) getUserStats(userID int64) (string, int64, int64, int64, error) {
var user User
err := b.db.Where("telegram_id = ? AND bot_id = ?", userID, b.botID).First(&user).Error
if err != nil {
return "", 0, 0, 0, fmt.Errorf("user not found: %w", err)
}
var messagesIn int64
if err := b.db.Model(&Message{}).Where("user_id = ? AND bot_id = ? AND is_user = ?",
userID, b.botID, true).Count(&messagesIn).Error; err != nil {
return "", 0, 0, 0, err
}
var messagesOut int64
if err := b.db.Model(&Message{}).Where("chat_id IN (SELECT DISTINCT chat_id FROM messages WHERE user_id = ? AND bot_id = ? AND deleted_at IS NULL) AND bot_id = ? AND is_user = ?",
userID, b.botID, b.botID, false).Count(&messagesOut).Error; err != nil {
return "", 0, 0, 0, err
}
totalMessages := messagesIn + messagesOut
return user.Username, messagesIn, messagesOut, totalMessages, nil
}
func isOnlyEmojis(s string) bool {
for _, r := range s {
if !isEmoji(r) {
return false
}
}
return true
}
func isEmoji(r rune) bool {
return (r >= 0x1F600 && r <= 0x1F64F) ||
(r >= 0x1F300 && r <= 0x1F5FF) ||
(r >= 0x1F680 && r <= 0x1F6FF) ||
(r >= 0x2600 && r <= 0x26FF) ||
(r >= 0x2700 && r <= 0x27BF)
}
func (b *Bot) sendWhoAmI(ctx context.Context, chatID int64, userID int64, username string, businessConnectionID string) {
user, err := b.getOrCreateUser(userID, username, false)
if err != nil {
ErrorLogger.Printf("Error getting or creating user: %v", err)
if err := b.sendResponse(ctx, chatID, "Sorry, I couldn't retrieve your information.", businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending response: %v", err)
}
return
}
role, err := b.getRoleByName(user.Role.Name)
if err != nil {
ErrorLogger.Printf("Error getting role by name: %v", err)
if err := b.sendResponse(ctx, chatID, "Sorry, I couldn't retrieve your role information.", businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending response: %v", err)
}
return
}
whoAmIMessage := fmt.Sprintf(
"👤 Your Information:\n\n"+
"- Username: %s\n"+
"- Role: %s",
user.Username,
role.Name,
)
if err := b.sendResponse(ctx, chatID, whoAmIMessage, businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending /whoami message: %v", err)
}
}
func (b *Bot) screenIncomingMessage(message *models.Message) (Message, error) {
if b.config.DebugScreening {
start := time.Now()
defer func() {
InfoLogger.Printf(
"[Screen] Incoming: chat=%d user=%d type=%s memory_size=%d duration=%v",
message.Chat.ID,
message.From.ID,
messageType(message),
len(b.getOrCreateChatMemory(message.Chat.ID).Messages),
time.Since(start),
)
}()
}
userRole := "user"
messageText := message.Text
if message.Sticker != nil {
if message.Sticker.Emoji != "" {
messageText = fmt.Sprintf("Sent a sticker: %s", message.Sticker.Emoji)
} else {
messageText = "Sent a sticker."
}
}
if message.Voice != nil {
messageText = "[Voice message]"
}
userMessage := b.createMessage(message.Chat.ID, message.From.ID, message.From.Username, userRole, messageText, true)
if message.Sticker != nil {
userMessage.StickerFileID = message.Sticker.FileID
userMessage.StickerEmoji = message.Sticker.Emoji
if message.Sticker.Thumbnail != nil {
userMessage.StickerPNGFile = message.Sticker.Thumbnail.FileID
}
}
chatMemory := b.getOrCreateChatMemory(message.Chat.ID)
if err := b.storeMessage(&userMessage); err != nil {
return Message{}, err
}
b.addMessageToChatMemory(chatMemory, userMessage)
return userMessage, nil
}
func (b *Bot) screenOutgoingMessage(chatID int64, response string) (Message, error) {
if b.config.DebugScreening {
start := time.Now()
defer func() {
InfoLogger.Printf(
"[Screen] Outgoing: chat=%d len=%d memory_size=%d duration=%v",
chatID,
len(response),
len(b.getOrCreateChatMemory(chatID).Messages),
time.Since(start),
)
}()
}
assistantMessage := b.createMessage(chatID, 0, "", "assistant", response, false)
if err := b.storeMessage(&assistantMessage); err != nil {
return Message{}, err
}
now := time.Now()
err := b.db.Model(&Message{}).
Where("chat_id = ? AND bot_id = ? AND is_user = ? AND answered_on IS NULL",
chatID, b.botID, true).
Order("timestamp DESC").
Limit(1).
Update("answered_on", now).Error
if err != nil {
ErrorLogger.Printf("Error marking user message as answered: %v", err)
}
chatMemory := b.getOrCreateChatMemory(chatID)
b.addMessageToChatMemory(chatMemory, assistantMessage)
return assistantMessage, nil
}
func (b *Bot) promoteUserToAdmin(promoterID, userToPromoteID int64) error {
if !b.hasScope(promoterID, ScopeUserPromote) {
return errors.New("only admins or owners can promote users to admin")
}
userToPromote, err := b.getOrCreateUser(userToPromoteID, "", false)
if err != nil {
return err
}
var adminRole Role
if err := b.db.Where("name = ?", "admin").First(&adminRole).Error; err != nil {
return err
}
userToPromote.RoleID = adminRole.ID
userToPromote.Role = adminRole
if err := b.db.Save(&userToPromote).Error; err != nil {
return err
}
b.registerAdminCommandsForUser(context.Background(), userToPromoteID)
return nil
}
+105
View File
@@ -0,0 +1,105 @@
package main
import (
"testing"
"github.com/stretchr/testify/assert"
)
func TestContentBlocksForMessage(t *testing.T) {
t.Run("empty message yields no blocks", func(t *testing.T) {
blocks := contentBlocksForMessage(Message{IsUser: true})
assert.Empty(t, blocks)
})
t.Run("user text only yields one text block", func(t *testing.T) {
blocks := contentBlocksForMessage(Message{IsUser: true, Text: "hello"})
assert.Len(t, blocks, 1)
assert.NotNil(t, blocks[0].OfText)
assert.Equal(t, "hello", blocks[0].OfText.Text)
})
t.Run("user single image without caption — no label, no text", func(t *testing.T) {
blocks := contentBlocksForMessage(Message{
IsUser: true,
ImageFileIDs: []string{"file_solo"},
})
assert.Len(t, blocks, 1)
assert.NotNil(t, blocks[0].OfImage)
assert.NotNil(t, blocks[0].OfImage.Source.OfFile)
assert.Equal(t, "file_solo", blocks[0].OfImage.Source.OfFile.FileID)
})
t.Run("user single image with caption — image before text", func(t *testing.T) {
blocks := contentBlocksForMessage(Message{
IsUser: true,
Text: "is this right?",
ImageFileIDs: []string{"file_solo"},
})
assert.Len(t, blocks, 2)
assert.NotNil(t, blocks[0].OfImage, "image block must come before text per Anthropic guidance")
assert.Equal(t, "file_solo", blocks[0].OfImage.Source.OfFile.FileID)
assert.NotNil(t, blocks[1].OfText)
assert.Equal(t, "is this right?", blocks[1].OfText.Text)
})
t.Run("user album (multi-image) labels each with Image N:", func(t *testing.T) {
blocks := contentBlocksForMessage(Message{
IsUser: true,
Text: "compare these",
ImageFileIDs: []string{"file_a", "file_b", "file_c"},
})
assert.Len(t, blocks, 7)
assert.Equal(t, "Image 1:", blocks[0].OfText.Text)
assert.Equal(t, "file_a", blocks[1].OfImage.Source.OfFile.FileID)
assert.Equal(t, "Image 2:", blocks[2].OfText.Text)
assert.Equal(t, "file_b", blocks[3].OfImage.Source.OfFile.FileID)
assert.Equal(t, "Image 3:", blocks[4].OfText.Text)
assert.Equal(t, "file_c", blocks[5].OfImage.Source.OfFile.FileID)
assert.Equal(t, "compare these", blocks[6].OfText.Text)
})
t.Run("assistant message with images-set is text-only (defensive)", func(t *testing.T) {
blocks := contentBlocksForMessage(Message{
IsUser: false,
Text: "I see your screenshot",
ImageFileIDs: []string{"file_should_be_ignored"},
})
assert.Len(t, blocks, 1)
assert.NotNil(t, blocks[0].OfText)
assert.Equal(t, "I see your screenshot", blocks[0].OfText.Text)
})
t.Run("whitespace-only text is skipped but images survive", func(t *testing.T) {
blocks := contentBlocksForMessage(Message{
IsUser: true,
Text: " \n ",
ImageFileIDs: []string{"file_x"},
})
assert.Len(t, blocks, 1)
assert.NotNil(t, blocks[0].OfImage)
})
}
func TestStripDeadFileIDFromMemory(t *testing.T) {
b, _ := setupBotForTest(t, 100)
chatID := int64(42)
cm := b.getOrCreateChatMemory(chatID)
cm.Messages = []Message{
{IsUser: true, Text: "first", ImageFileIDs: []string{"file_a", "file_b"}},
{IsUser: false, Text: "reply"},
{IsUser: true, Text: "third", ImageFileIDs: []string{"file_b", "file_c"}},
}
b.stripDeadFileIDFromMemory(chatID, "file_b")
assert.Equal(t, []string{"file_a"}, cm.Messages[0].ImageFileIDs, "file_b should be removed from message 1")
assert.Empty(t, cm.Messages[1].ImageFileIDs, "assistant message untouched")
assert.Equal(t, []string{"file_c"}, cm.Messages[2].ImageFileIDs, "file_b should be removed from message 3")
}
func TestStripDeadFileIDFromMemory_UnknownChatIsNoop(t *testing.T) {
b, _ := setupBotForTest(t, 100)
b.stripDeadFileIDFromMemory(99999, "file_anything")
}
+25
View File
@@ -0,0 +1,25 @@
package main
import "time"
type Clock interface {
Now() time.Time
}
type RealClock struct{}
func (RealClock) Now() time.Time {
return time.Now()
}
type MockClock struct {
currentTime time.Time
}
func (mc *MockClock) Now() time.Time {
return mc.currentTime
}
func (mc *MockClock) Advance(d time.Duration) {
mc.currentTime = mc.currentTime.Add(d)
}
+295
View File
@@ -0,0 +1,295 @@
package main
import (
"encoding/json"
"fmt"
"os"
"path/filepath"
"strings"
)
type MCPServer struct {
Name string `json:"name"`
URL string `json:"url"`
AuthorizationToken string `json:"authorization_token,omitempty"`
AllowedTools []string `json:"allowed_tools,omitempty"`
}
type WebSearchConfig struct {
AllowedDomains []string `json:"allowed_domains,omitempty"`
BlockedDomains []string `json:"blocked_domains,omitempty"`
FetchAllowedDomains []string `json:"fetch_allowed_domains,omitempty"`
MaxUses int `json:"max_uses,omitempty"`
Fetch bool `json:"fetch,omitempty"`
MaxContentTokens int `json:"max_content_tokens,omitempty"`
}
const (
ThinkingModeAdaptive = "adaptive"
ThinkingModeDisabled = "disabled"
ThinkingDisplaySummarized = "summarized"
ThinkingDisplayOmitted = "omitted"
)
type BotConfig struct {
ID string `json:"id"`
TelegramToken string `json:"telegram_token"`
MemorySize int `json:"memory_size"`
MessagePerHour int `json:"messages_per_hour"`
MessagePerDay int `json:"messages_per_day"`
TempBanDuration string `json:"temp_ban_duration"`
Model string `json:"model"`
Temperature *float32 `json:"temperature,omitempty"`
MaxTokens int `json:"max_tokens,omitempty"`
Thinking string `json:"thinking,omitempty"`
ThinkingDisplay string `json:"thinking_display,omitempty"`
SystemPrompts map[string]string `json:"system_prompts"`
Active bool `json:"active"`
OwnerTelegramID int64 `json:"owner_telegram_id"`
AnthropicAPIKey string `json:"anthropic_api_key"`
ElevenLabsAPIKey string `json:"elevenlabs_api_key"`
ElevenLabsVoiceID string `json:"elevenlabs_voice_id"`
ElevenLabsModel string `json:"elevenlabs_model"`
DebugScreening bool `json:"debug_screening"`
MCPServers []MCPServer `json:"mcp_servers,omitempty"`
WebSearch *WebSearchConfig `json:"web_search,omitempty"`
ConfigFilePath string `json:"-"`
}
func validateConfigPath(configDir, filename string) (string, error) {
configDir = filepath.Clean(configDir)
filename = filepath.Clean(filename)
absConfigDir, err := filepath.Abs(configDir)
if err != nil {
return "", fmt.Errorf("failed to get absolute path for config directory: %w", err)
}
fullPath := filepath.Join(absConfigDir, filename)
absPath, err := filepath.Abs(fullPath)
if err != nil {
return "", fmt.Errorf("failed to get absolute path for config file: %w", err)
}
rel, err := filepath.Rel(absConfigDir, absPath)
if err != nil || strings.HasPrefix(rel, "..") || strings.Contains(rel, "..") {
return "", fmt.Errorf("invalid config path: file must be within the config directory")
}
if filepath.Ext(absPath) != ".json" {
return "", fmt.Errorf("invalid file extension: must be .json")
}
return absPath, nil
}
func loadAllConfigs(dir string) ([]BotConfig, error) {
var configs []BotConfig
ids := make(map[string]bool)
tokens := make(map[string]bool)
files, err := os.ReadDir(dir)
if err != nil {
return nil, fmt.Errorf("failed to read config directory: %w", err)
}
for _, file := range files {
if filepath.Ext(file.Name()) == ".json" {
validPath, err := validateConfigPath(dir, file.Name())
if err != nil {
InfoLogger.Printf("Invalid config path for %s: %v", file.Name(), err)
continue
}
config, err := loadConfig(validPath)
if err != nil {
InfoLogger.Printf("Failed to load config %s: %v", validPath, err)
continue
}
if !config.Active {
InfoLogger.Printf("Skipping inactive bot: %s", config.ID)
continue
}
if err := validateConfig(&config, ids, tokens); err != nil {
InfoLogger.Printf("Config validation failed for %s: %v", validPath, err)
continue
}
if config.Thinking == ThinkingModeAdaptive && config.MaxTokens > 0 && config.MaxTokens < 4000 {
InfoLogger.Printf("[%s] thinking=adaptive with max_tokens=%d: thinking tokens count toward max_tokens; consider >= 4000",
config.ID, config.MaxTokens)
}
if ws := config.WebSearch; ws != nil && len(ws.AllowedDomains) == 0 && len(ws.BlockedDomains) == 0 {
InfoLogger.Printf("[%s] web_search enabled with no allowed_domains/blocked_domains: the model may search the open web",
config.ID)
}
if ws := config.WebSearch; ws != nil && len(ws.FetchAllowedDomains) > 0 {
if !ws.Fetch {
InfoLogger.Printf("[%s] web_search.fetch_allowed_domains is set but fetch is disabled: it has no effect",
config.ID)
}
for _, d := range ws.FetchAllowedDomains {
if strings.Contains(d, "/") {
InfoLogger.Printf("[%s] web_search.fetch_allowed_domains entry %q includes a path: web_fetch matches host-only, so the whole host is fetchable",
config.ID, d)
}
}
}
config.ConfigFilePath = validPath
configs = append(configs, config)
}
}
if len(configs) == 0 {
return nil, fmt.Errorf("no valid configs found")
}
return configs, nil
}
func validateConfig(config *BotConfig, ids, tokens map[string]bool) error {
if config.ID == "" {
return fmt.Errorf("missing 'id' field")
}
if _, exists := ids[config.ID]; exists {
return fmt.Errorf("duplicate bot id '%s'", config.ID)
}
ids[config.ID] = true
if config.TelegramToken == "" {
return fmt.Errorf("missing 'telegram_token' field")
}
if _, exists := tokens[config.TelegramToken]; exists {
return fmt.Errorf("duplicate telegram_token")
}
tokens[config.TelegramToken] = true
if config.Model == "" {
return fmt.Errorf("missing 'model' field")
}
switch config.Thinking {
case "", ThinkingModeAdaptive, ThinkingModeDisabled:
default:
return fmt.Errorf("invalid 'thinking' value %q: must be %q or %q (or omitted)",
config.Thinking, ThinkingModeAdaptive, ThinkingModeDisabled)
}
switch config.ThinkingDisplay {
case "":
case ThinkingDisplaySummarized, ThinkingDisplayOmitted:
if config.Thinking != ThinkingModeAdaptive {
return fmt.Errorf("'thinking_display' requires 'thinking': %q (the API rejects display with thinking disabled)",
ThinkingModeAdaptive)
}
default:
return fmt.Errorf("invalid 'thinking_display' value %q: must be %q or %q (or omitted)",
config.ThinkingDisplay, ThinkingDisplaySummarized, ThinkingDisplayOmitted)
}
if config.MaxTokens < 0 {
return fmt.Errorf("'max_tokens' must be greater than 0 when set")
}
if ws := config.WebSearch; ws != nil {
if len(ws.AllowedDomains) > 0 && len(ws.BlockedDomains) > 0 {
return fmt.Errorf("'web_search' cannot set both allowed_domains and blocked_domains (the API rejects that)")
}
if ws.MaxUses < 0 {
return fmt.Errorf("'web_search.max_uses' must be greater than 0 when set")
}
if ws.MaxContentTokens < 0 {
return fmt.Errorf("'web_search.max_content_tokens' must be greater than 0 when set")
}
}
if config.MessagePerHour <= 0 {
return fmt.Errorf("'messages_per_hour' must be greater than 0")
}
if config.MessagePerDay <= 0 {
return fmt.Errorf("'messages_per_day' must be greater than 0")
}
return nil
}
func loadConfig(filename string) (BotConfig, error) {
var config BotConfig
file, err := os.OpenFile(filepath.Clean(filename), os.O_RDONLY, 0)
if err != nil {
return config, fmt.Errorf("failed to open config file %s: %w", filename, err)
}
defer func() {
if err := file.Close(); err != nil {
InfoLogger.Printf("Failed to close config file: %v", err)
}
}()
decoder := json.NewDecoder(file)
if err := decoder.Decode(&config); err != nil {
return config, fmt.Errorf("failed to decode JSON from %s: %w", filename, err)
}
return config, nil
}
func (c *BotConfig) Reload(configDir, filename string) error {
validPath, err := validateConfigPath(configDir, filename)
if err != nil {
return fmt.Errorf("invalid config path: %w", err)
}
cleanPath := filepath.Clean(validPath)
file, err := os.OpenFile(cleanPath, os.O_RDONLY, 0)
if err != nil {
return fmt.Errorf("failed to open config file %s: %w", cleanPath, err)
}
defer func() {
if err := file.Close(); err != nil {
InfoLogger.Printf("Failed to close config file: %v", err)
}
}()
decoder := json.NewDecoder(file)
if err := decoder.Decode(c); err != nil {
return fmt.Errorf("failed to decode JSON from %s: %w", validPath, err)
}
return nil
}
func (c *BotConfig) PersistModel(newModel string) error {
if c.ConfigFilePath == "" {
return fmt.Errorf("config file path not set; cannot persist model")
}
data, err := os.ReadFile(c.ConfigFilePath)
if err != nil {
return fmt.Errorf("failed to read config for update: %w", err)
}
var raw map[string]any
if err := json.Unmarshal(data, &raw); err != nil {
return fmt.Errorf("failed to parse config for update: %w", err)
}
raw["model"] = newModel
updated, err := json.MarshalIndent(raw, "", "\t")
if err != nil {
return fmt.Errorf("failed to re-encode config: %w", err)
}
if err := os.WriteFile(c.ConfigFilePath, updated, 0600); err != nil {
return fmt.Errorf("failed to write config: %w", err)
}
c.Model = newModel
return nil
}
+22
View File
@@ -0,0 +1,22 @@
{
"id": "default_bot",
"active": false,
"telegram_token": "YOUR_TELEGRAM_BOT_TOKEN",
"owner_telegram_id": 111111111,
"anthropic_api_key": "YOUR_SPECIFIC_ANTHROPIC_API_KEY",
"elevenlabs_api_key": "",
"elevenlabs_voice_id": "",
"elevenlabs_model": "",
"memory_size": 10,
"messages_per_hour": 20,
"messages_per_day": 100,
"temp_ban_duration": "24h",
"model": "claude-haiku-4-5",
"temperature": 0.7,
"max_tokens": 1000,
"debug_screening": false,
"system_prompts": {
"custom_instructions": "You are Atom, a helpful assistant texting through a limited Telegram interface with a 15-word maximum. Write like texting a friend - use shorthand, skip grammar, use slang/abbreviations. The system cuts off anything longer than 15 words.\n\n- Address the user by their first name, and reply in their preferred language (both are in the conversation context).\n- Use time-appropriate greetings based on the user's local time of day.\n- If a user asks about buying apples, inform them that we don't sell apples.\n- When asked for a joke, tell a clean, family-friendly joke about programming or technology.\n- If someone inquires about our services, explain that we offer AI-powered chatbot solutions.\n- For any questions about pricing, direct users to contact our sales team at [email protected].\n- If asked about your capabilities, be honest about what you can and cannot do.\nAlways maintain a friendly and professional tone.",
"respond_with_emojis": "The user's message contains only emoji. Reply using only emoji."
}
}
+958
View File
@@ -0,0 +1,958 @@
package main
import (
"encoding/json"
"fmt"
"os"
"path/filepath"
"strings"
"testing"
)
func TestMain(m *testing.M) {
initLoggers()
os.Exit(m.Run())
}
func TestBotConfig_UnmarshalJSON(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
jsonData := `{
"id": "bot123",
"telegram_token": "token123",
"memory_size": 1024,
"messages_per_hour": 10,
"messages_per_day": 100,
"temp_ban_duration": "1h",
"model": "claude-v1",
"temperature": 0.7,
"system_prompts": {"welcome": "Hello!"},
"active": true,
"owner_telegram_id": 123456789,
"anthropic_api_key": "api_key_123"
}`
var config BotConfig
if err := json.Unmarshal([]byte(jsonData), &config); err != nil {
t.Fatalf("Failed to unmarshal JSON: %v", err)
}
expectedModel := "claude-v1"
if config.Model != expectedModel {
t.Errorf("Expected model %s, got %s", expectedModel, config.Model)
}
expectedID := "bot123"
if config.ID != expectedID {
t.Errorf("Expected ID %s, got %s", expectedID, config.ID)
}
}
func TestValidateConfigPath(t *testing.T) {
execDir, err := os.Getwd()
if err != nil {
t.Fatalf("Failed to get current directory: %v", err)
}
tests := []struct {
name string
configDir string
filename string
wantErr bool
}{
{
name: "Valid Path",
configDir: execDir,
filename: "config.json",
wantErr: false,
},
{
name: "Invalid Extension",
configDir: execDir,
filename: "config.yaml",
wantErr: true,
},
{
name: "Path Traversal",
configDir: execDir,
filename: "../config.json",
wantErr: true,
},
{
name: "Absolute Path Outside",
configDir: execDir,
filename: "/etc/passwd",
wantErr: true,
},
{
name: "Nested Valid Path",
configDir: execDir,
filename: "subdir/config.json",
wantErr: false,
},
}
subDir := filepath.Join(execDir, "subdir")
if err := os.MkdirAll(subDir, 0755); err != nil {
t.Fatalf("Failed to create subdir: %v", err)
}
defer func() {
if err := os.RemoveAll(subDir); err != nil {
t.Errorf("Failed to remove test subdirectory: %v", err)
}
}()
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
configDir := tt.configDir
filename := tt.filename
if tt.name == "Nested Valid Path" {
configDir = subDir
}
_, err := validateConfigPath(configDir, filename)
if (err != nil) != tt.wantErr {
t.Errorf("validateConfigPath() error = %v, wantErr %v", err, tt.wantErr)
}
})
}
}
func TestLoadConfig(t *testing.T) {
tempDir, err := os.MkdirTemp("", "config_test")
if err != nil {
t.Fatalf("Failed to create temp dir: %v", err)
}
defer func() {
if err := os.RemoveAll(tempDir); err != nil {
t.Errorf("Failed to remove temp directory: %v", err)
}
}()
validConfig := `{
"id": "bot123",
"telegram_token": "token123",
"memory_size": 1024,
"messages_per_hour": 10,
"messages_per_day": 100,
"temp_ban_duration": "1h",
"model": "claude-v1",
"temperature": 0.7,
"system_prompts": {"welcome": "Hello!"},
"active": true,
"owner_telegram_id": 123456789,
"anthropic_api_key": "api_key_123"
}`
invalidConfig := `{
"id": "bot123",
"telegram_token": "token123",
"memory_size": "should be int",
"model": "claude-v1"
}`
validPath := filepath.Join(tempDir, "valid_config.json")
if err := os.WriteFile(validPath, []byte(validConfig), 0644); err != nil {
t.Fatalf("Failed to write valid config: %v", err)
}
invalidPath := filepath.Join(tempDir, "invalid_config.json")
if err := os.WriteFile(invalidPath, []byte(invalidConfig), 0644); err != nil {
t.Fatalf("Failed to write invalid config: %v", err)
}
tests := []struct {
name string
filename string
wantErr bool
expectID string
expectErr string
}{
{
name: "Load Valid Config",
filename: validPath,
wantErr: false,
expectID: "bot123",
},
{
name: "Load Invalid Config",
filename: invalidPath,
wantErr: true,
expectErr: "failed to decode JSON",
},
{
name: "Non-existent File",
filename: filepath.Join(tempDir, "nonexistent.json"),
wantErr: true,
expectErr: "failed to open config file",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
config, err := loadConfig(tt.filename)
if (err != nil) != tt.wantErr {
t.Errorf("loadConfig() error = %v, wantErr %v", err, tt.wantErr)
return
}
if tt.wantErr && err != nil && tt.expectErr != "" {
if !contains(err.Error(), tt.expectErr) {
t.Errorf("loadConfig() error = %v, expected to contain %v", err, tt.expectErr)
}
return
}
if config.ID != tt.expectID {
t.Errorf("Expected ID %s, got %s", tt.expectID, config.ID)
}
})
}
}
func TestValidateConfig(t *testing.T) {
tests := []struct {
name string
config BotConfig
ids map[string]bool
tokens map[string]bool
wantErr bool
expectedError string
}{
{
name: "Valid Config",
config: BotConfig{
ID: "bot123",
TelegramToken: "token123",
Model: "claude-v1",
Active: true,
OwnerTelegramID: 123456789,
MessagePerHour: 10,
MessagePerDay: 100,
},
ids: make(map[string]bool),
tokens: make(map[string]bool),
wantErr: false,
},
{
name: "Missing ID",
config: BotConfig{
TelegramToken: "token123",
Model: "claude-v1",
Active: true,
},
ids: make(map[string]bool),
tokens: make(map[string]bool),
wantErr: true,
expectedError: "missing 'id' field",
},
{
name: "Duplicate ID",
config: BotConfig{
ID: "bot123",
TelegramToken: "token123",
Model: "claude-v1",
Active: true,
},
ids: map[string]bool{"bot123": true},
tokens: make(map[string]bool),
wantErr: true,
expectedError: "duplicate bot id",
},
{
name: "Missing Telegram Token",
config: BotConfig{
ID: "bot123",
Model: "claude-v1",
Active: true,
},
ids: make(map[string]bool),
tokens: make(map[string]bool),
wantErr: true,
expectedError: "missing 'telegram_token' field",
},
{
name: "Duplicate Telegram Token",
config: BotConfig{
ID: "bot123",
TelegramToken: "token123",
Model: "claude-v1",
Active: true,
},
ids: make(map[string]bool),
tokens: map[string]bool{"token123": true},
wantErr: true,
expectedError: "duplicate telegram_token",
},
{
name: "Missing Model",
config: BotConfig{
ID: "bot123",
TelegramToken: "token123",
Active: true,
},
ids: make(map[string]bool),
tokens: make(map[string]bool),
wantErr: true,
expectedError: "missing 'model' field",
},
{
name: "Zero MessagePerHour",
config: BotConfig{
ID: "bot123",
TelegramToken: "token123",
Model: "claude-v1",
MessagePerHour: 0,
MessagePerDay: 100,
},
ids: make(map[string]bool),
tokens: make(map[string]bool),
wantErr: true,
expectedError: "'messages_per_hour' must be greater than 0",
},
{
name: "Zero MessagePerDay",
config: BotConfig{
ID: "bot123",
TelegramToken: "token123",
Model: "claude-v1",
MessagePerHour: 10,
MessagePerDay: 0,
},
ids: make(map[string]bool),
tokens: make(map[string]bool),
wantErr: true,
expectedError: "'messages_per_day' must be greater than 0",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
err := validateConfig(&tt.config, tt.ids, tt.tokens)
if (err != nil) != tt.wantErr {
t.Errorf("validateConfig() error = %v, wantErr %v", err, tt.wantErr)
return
}
if tt.wantErr && err != nil && tt.expectedError != "" {
if !contains(err.Error(), tt.expectedError) {
t.Errorf("validateConfig() error = %v, expected to contain %v", err, tt.expectedError)
}
}
})
}
}
func TestLoadAllConfigs(t *testing.T) {
tempDir, err := os.MkdirTemp("", "load_all_configs_test")
if err != nil {
t.Fatalf("Failed to create temp dir: %v", err)
}
defer func() {
if err := os.RemoveAll(tempDir); err != nil {
t.Errorf("Failed to remove temp directory: %v", err)
}
}()
tests := []struct {
name string
setupFiles map[string]string
expectConfigs int
expectError bool
expectErrorMsg string
}{
{
name: "Load All Valid Configs",
setupFiles: map[string]string{
"valid_config.json": `{
"id": "bot123",
"telegram_token": "token123",
"memory_size": 1024,
"messages_per_hour": 10,
"messages_per_day": 100,
"temp_ban_duration": "1h",
"model": "claude-v1",
"temperature": 0.7,
"system_prompts": {"welcome": "Hello!"},
"active": true,
"owner_telegram_id": 123456789,
"anthropic_api_key": "api_key_123"
}`,
},
expectConfigs: 1,
expectError: false,
},
{
name: "Skip Inactive Config",
setupFiles: map[string]string{
"valid_config.json": `{
"id": "bot123",
"telegram_token": "token123",
"memory_size": 1024,
"messages_per_hour": 10,
"messages_per_day": 100,
"temp_ban_duration": "1h",
"model": "claude-v1",
"system_prompts": {"welcome": "Hello!"},
"active": true,
"owner_telegram_id": 123456789,
"anthropic_api_key": "api_key_123"
}`,
"inactive_config.json": `{
"id": "bot124",
"telegram_token": "token124",
"memory_size": 512,
"messages_per_hour": 5,
"messages_per_day": 50,
"temp_ban_duration": "30m",
"model": "claude-v2",
"temperature": 0.5,
"system_prompts": {"welcome": "Hi!"},
"active": false,
"owner_telegram_id": 987654321,
"anthropic_api_key": "api_key_124"
}`,
},
expectConfigs: 1,
expectError: false,
},
{
name: "Duplicate Bot ID",
setupFiles: map[string]string{
"valid_config.json": `{
"id": "bot123",
"telegram_token": "token123",
"memory_size": 1024,
"messages_per_hour": 10,
"messages_per_day": 100,
"temp_ban_duration": "1h",
"model": "claude-v1",
"system_prompts": {"welcome": "Hello!"},
"active": true,
"owner_telegram_id": 123456789,
"anthropic_api_key": "api_key_123"
}`,
"duplicate_id_config.json": `{
"id": "bot123",
"telegram_token": "token125",
"memory_size": 256,
"messages_per_hour": 2,
"messages_per_day": 20,
"temp_ban_duration": "15m",
"model": "claude-v3",
"temperature": 0.3,
"system_prompts": {"welcome": "Hey!"},
"active": true,
"owner_telegram_id": 1122334455,
"anthropic_api_key": "api_key_125"
}`,
},
expectConfigs: 1,
expectError: false,
},
{
name: "Duplicate Telegram Token",
setupFiles: map[string]string{
"valid_config.json": `{
"id": "bot123",
"telegram_token": "token123",
"memory_size": 1024,
"messages_per_hour": 10,
"messages_per_day": 100,
"temp_ban_duration": "1h",
"model": "claude-v1",
"system_prompts": {"welcome": "Hello!"},
"active": true,
"owner_telegram_id": 123456789,
"anthropic_api_key": "api_key_123"
}`,
"duplicate_token_config.json": `{
"id": "bot126",
"telegram_token": "token123",
"memory_size": 128,
"messages_per_hour": 1,
"messages_per_day": 10,
"temp_ban_duration": "5m",
"model": "claude-v4",
"temperature": 0.2,
"system_prompts": {"welcome": "Greetings!"},
"active": true,
"owner_telegram_id": 5566778899,
"anthropic_api_key": "api_key_126"
}`,
},
expectConfigs: 1,
expectError: false,
},
{
name: "Invalid Config",
setupFiles: map[string]string{
"valid_config.json": `{
"id": "bot123",
"telegram_token": "token123",
"memory_size": 1024,
"messages_per_hour": 10,
"messages_per_day": 100,
"temp_ban_duration": "1h",
"model": "claude-v1",
"system_prompts": {"welcome": "Hello!"},
"active": true,
"owner_telegram_id": 123456789,
"anthropic_api_key": "api_key_123"
}`,
"invalid_config.json": `{
"id": "bot127",
"telegram_token": "token127",
"model": "",
"active": true
}`,
},
expectConfigs: 1,
expectError: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if err := os.RemoveAll(tempDir); err != nil {
t.Fatalf("Failed to remove temp dir: %v", err)
}
if err := os.MkdirAll(tempDir, 0755); err != nil {
t.Fatalf("Failed to create temp dir: %v", err)
}
for filename, content := range tt.setupFiles {
err := os.WriteFile(filepath.Join(tempDir, filename), []byte(content), 0644)
if err != nil {
t.Fatalf("Failed to write file %s: %v", filename, err)
}
}
configs, err := loadAllConfigs(tempDir)
if (err != nil) != tt.expectError {
t.Errorf("loadAllConfigs() error = %v, wantErr %v", err, tt.expectError)
return
}
if len(configs) != tt.expectConfigs {
t.Errorf("Expected %d configs, got %d", tt.expectConfigs, len(configs))
}
})
}
}
func TestBotConfig_Reload(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
tempDir, err := os.MkdirTemp("", "reload_test")
if err != nil {
t.Fatalf("Failed to create temp dir: %v", err)
}
defer func() {
if err := os.RemoveAll(tempDir); err != nil {
t.Errorf("Failed to remove temp directory: %v", err)
}
}()
config1 := `{
"id": "bot123",
"telegram_token": "token123",
"memory_size": 1024,
"messages_per_hour": 10,
"messages_per_day": 100,
"temp_ban_duration": "1h",
"model": "claude-v1",
"temperature": 0.7,
"system_prompts": {"welcome": "Hello!"},
"active": true,
"owner_telegram_id": 123456789,
"anthropic_api_key": "api_key_123"
}`
configPath := filepath.Join(tempDir, "config.json")
if err := os.WriteFile(configPath, []byte(config1), 0644); err != nil {
t.Fatalf("Failed to write initial config: %v", err)
}
var config BotConfig
if err := config.Reload(tempDir, "config.json"); err != nil {
t.Fatalf("Failed to reload config: %v", err)
}
if config.ID != "bot123" {
t.Errorf("Expected ID 'bot123', got '%s'", config.ID)
}
if config.Model != "claude-v1" {
t.Errorf("Expected Model 'claude-v1', got '%s'", config.Model)
}
config2 := `{
"id": "bot123",
"telegram_token": "token123_updated",
"memory_size": 2048,
"messages_per_hour": 20,
"messages_per_day": 200,
"temp_ban_duration": "2h",
"model": "claude-v2",
"temperature": 0.3,
"system_prompts": {"welcome": "Hi there!"},
"active": true,
"owner_telegram_id": 987654321,
"anthropic_api_key": "api_key_456"
}`
if err := os.WriteFile(configPath, []byte(config2), 0644); err != nil {
t.Fatalf("Failed to write updated config: %v", err)
}
if err := config.Reload(tempDir, "config.json"); err != nil {
t.Fatalf("Failed to reload updated config: %v", err)
}
if config.TelegramToken != "token123_updated" {
t.Errorf("Expected TelegramToken 'token123_updated', got '%s'", config.TelegramToken)
}
if config.MemorySize != 2048 {
t.Errorf("Expected MemorySize 2048, got %d", config.MemorySize)
}
if config.Model != "claude-v2" {
t.Errorf("Expected Model 'claude-v2', got '%s'", config.Model)
}
if config.OwnerTelegramID != 987654321 {
t.Errorf("Expected OwnerTelegramID 987654321, got %d", config.OwnerTelegramID)
}
}
func TestBotConfig_UnmarshalJSON_Invalid(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
jsonData := `{
"id": "bot123",
"telegram_token": "token123",
"memory_size": 1024,
"messages_per_hour": 10,
"messages_per_day": 100,
"temp_ban_duration": "1h",
"model": "",
"temperature": 0.7,
"system_prompts": {"welcome": "Hello!"},
"active": true,
"owner_telegram_id": 123456789,
"anthropic_api_key": "api_key_123"
}`
var config BotConfig
err := json.Unmarshal([]byte(jsonData), &config)
if err != nil {
t.Fatalf("Failed to unmarshal JSON: %v", err)
}
if config.Model != "" {
t.Errorf("Expected empty model, got %s", config.Model)
}
}
func contains(s, substr string) bool {
return strings.Contains(s, substr)
}
func TestTemperatureConfig(t *testing.T) {
tempDir, err := os.MkdirTemp("", "temperature_test")
if err != nil {
t.Fatalf("Failed to create temp dir: %v", err)
}
defer func() {
if err := os.RemoveAll(tempDir); err != nil {
t.Errorf("Failed to remove temp directory: %v", err)
}
}()
configWithTemp := `{
"id": "bot123",
"telegram_token": "token123",
"memory_size": 1024,
"messages_per_hour": 10,
"messages_per_day": 100,
"temp_ban_duration": "1h",
"model": "claude-v1",
"temperature": 0.42,
"system_prompts": {"welcome": "Hello!"},
"active": true,
"owner_telegram_id": 123456789,
"anthropic_api_key": "api_key_123"
}`
configWithoutTemp := `{
"id": "bot124",
"telegram_token": "token124",
"memory_size": 1024,
"messages_per_hour": 10,
"messages_per_day": 100,
"temp_ban_duration": "1h",
"model": "claude-v1",
"system_prompts": {"welcome": "Hello!"},
"active": true,
"owner_telegram_id": 123456789,
"anthropic_api_key": "api_key_123"
}`
withTempPath := filepath.Join(tempDir, "with_temp.json")
if err := os.WriteFile(withTempPath, []byte(configWithTemp), 0644); err != nil {
t.Fatalf("Failed to write config with temperature: %v", err)
}
withoutTempPath := filepath.Join(tempDir, "without_temp.json")
if err := os.WriteFile(withoutTempPath, []byte(configWithoutTemp), 0644); err != nil {
t.Fatalf("Failed to write config without temperature: %v", err)
}
configWithTempObj, err := loadConfig(withTempPath)
if err != nil {
t.Fatalf("Failed to load config with temperature: %v", err)
}
if configWithTempObj.Temperature == nil {
t.Errorf("Expected Temperature to be set, got nil")
} else if *configWithTempObj.Temperature != 0.42 {
t.Errorf("Expected Temperature 0.42, got %f", *configWithTempObj.Temperature)
}
configWithoutTempObj, err := loadConfig(withoutTempPath)
if err != nil {
t.Fatalf("Failed to load config without temperature: %v", err)
}
if configWithoutTempObj.Temperature != nil {
t.Errorf("Expected Temperature to be nil, got %f", *configWithoutTempObj.Temperature)
}
}
func TestBotConfig_PersistModel(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
tempDir, err := os.MkdirTemp("", "persist_model_test")
if err != nil {
t.Fatalf("Failed to create temp dir: %v", err)
}
defer func() {
if err := os.RemoveAll(tempDir); err != nil {
t.Errorf("Failed to remove temp directory: %v", err)
}
}()
initialJSON := `{
"id": "bot1",
"telegram_token": "token1",
"model": "claude-v1",
"messages_per_hour": 10,
"messages_per_day": 100
}`
configPath := filepath.Join(tempDir, "config.json")
if err := os.WriteFile(configPath, []byte(initialJSON), 0600); err != nil {
t.Fatalf("Failed to write config file: %v", err)
}
config := BotConfig{
ID: "bot1",
Model: "claude-v1",
ConfigFilePath: configPath,
}
if err := config.PersistModel("claude-sonnet-4-6"); err != nil {
t.Fatalf("PersistModel() unexpected error: %v", err)
}
if string(config.Model) != "claude-sonnet-4-6" {
t.Errorf("in-memory model: got %q, want %q", config.Model, "claude-sonnet-4-6")
}
data, err := os.ReadFile(configPath)
if err != nil {
t.Fatalf("Failed to read updated config file: %v", err)
}
var raw map[string]any
if err := json.Unmarshal(data, &raw); err != nil {
t.Fatalf("Failed to unmarshal updated config: %v", err)
}
if raw["model"] != "claude-sonnet-4-6" {
t.Errorf("on-disk model: got %v, want %q", raw["model"], "claude-sonnet-4-6")
}
if raw["id"] != "bot1" {
t.Errorf("on-disk id should be preserved: got %v, want %q", raw["id"], "bot1")
}
noPath := BotConfig{Model: "claude-v1"}
if err := noPath.PersistModel("claude-sonnet-4-6"); err == nil {
t.Error("PersistModel with empty ConfigFilePath: expected error, got nil")
}
}
func thinkingTestConfig(id string) BotConfig {
return BotConfig{
ID: id,
TelegramToken: "token-" + id,
MemorySize: 10,
MessagePerHour: 10,
MessagePerDay: 100,
TempBanDuration: "1h",
Model: "claude-test",
}
}
func TestThinkingConfig(t *testing.T) {
cases := []struct {
name string
thinking string
display string
wantErr string
}{
{"absent", "", "", ""},
{"adaptive", ThinkingModeAdaptive, "", ""},
{"disabled", ThinkingModeDisabled, "", ""},
{"adaptive summarized", ThinkingModeAdaptive, ThinkingDisplaySummarized, ""},
{"adaptive omitted", ThinkingModeAdaptive, ThinkingDisplayOmitted, ""},
{"legacy enabled rejected", "enabled", "", "invalid 'thinking'"},
{"case sensitive", "Adaptive", "", "invalid 'thinking'"},
{"unknown display", ThinkingModeAdaptive, "verbose", "invalid 'thinking_display'"},
{"display without thinking", "", ThinkingDisplaySummarized, "'thinking_display' requires"},
{"display with disabled", ThinkingModeDisabled, ThinkingDisplayOmitted, "'thinking_display' requires"},
}
for i, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
cfg := thinkingTestConfig(fmt.Sprintf("bot-think-%d", i))
cfg.Thinking = tc.thinking
cfg.ThinkingDisplay = tc.display
err := validateConfig(&cfg, map[string]bool{}, map[string]bool{})
if tc.wantErr == "" {
if err != nil {
t.Fatalf("validateConfig(thinking=%q display=%q) = %v, want nil", tc.thinking, tc.display, err)
}
return
}
if err == nil {
t.Fatalf("validateConfig(thinking=%q display=%q) = nil, want error containing %q", tc.thinking, tc.display, tc.wantErr)
}
if !contains(err.Error(), tc.wantErr) {
t.Errorf("error %q does not contain %q", err.Error(), tc.wantErr)
}
})
}
}
func TestMaxTokensConfig(t *testing.T) {
var withValue BotConfig
if err := json.Unmarshal([]byte(`{"max_tokens": 4000}`), &withValue); err != nil {
t.Fatalf("unmarshal: %v", err)
}
if withValue.MaxTokens != 4000 {
t.Errorf("MaxTokens = %d, want 4000", withValue.MaxTokens)
}
var withoutValue BotConfig
if err := json.Unmarshal([]byte(`{}`), &withoutValue); err != nil {
t.Fatalf("unmarshal: %v", err)
}
if withoutValue.MaxTokens != 0 {
t.Errorf("MaxTokens = %d, want 0 when absent", withoutValue.MaxTokens)
}
neg := thinkingTestConfig("bot-maxtok-neg")
neg.MaxTokens = -1
if err := validateConfig(&neg, map[string]bool{}, map[string]bool{}); err == nil {
t.Error("validateConfig(max_tokens=-1) = nil, want error")
}
zero := thinkingTestConfig("bot-maxtok-zero")
zero.MaxTokens = 0
if err := validateConfig(&zero, map[string]bool{}, map[string]bool{}); err != nil {
t.Errorf("validateConfig(max_tokens=0) = %v, want nil", err)
}
}
func TestThinkingConfigLoad(t *testing.T) {
jsonData := `{
"id": "bot-think-load",
"thinking": "adaptive",
"thinking_display": "omitted",
"max_tokens": 4096
}`
var cfg BotConfig
if err := json.Unmarshal([]byte(jsonData), &cfg); err != nil {
t.Fatalf("unmarshal: %v", err)
}
if cfg.Thinking != ThinkingModeAdaptive {
t.Errorf("Thinking = %q, want %q", cfg.Thinking, ThinkingModeAdaptive)
}
if cfg.ThinkingDisplay != ThinkingDisplayOmitted {
t.Errorf("ThinkingDisplay = %q, want %q", cfg.ThinkingDisplay, ThinkingDisplayOmitted)
}
if cfg.MaxTokens != 4096 {
t.Errorf("MaxTokens = %d, want 4096", cfg.MaxTokens)
}
}
func TestWebSearchConfig(t *testing.T) {
t.Run("absent leaves nil", func(t *testing.T) {
var cfg BotConfig
if err := json.Unmarshal([]byte(`{}`), &cfg); err != nil {
t.Fatalf("unmarshal: %v", err)
}
if cfg.WebSearch != nil {
t.Errorf("WebSearch = %+v, want nil when absent", cfg.WebSearch)
}
})
t.Run("loads allowlist", func(t *testing.T) {
jsonData := `{
"web_search": {
"allowed_domains": ["example.com/hc", "docs.example.com"],
"max_uses": 3,
"fetch": true,
"max_content_tokens": 50000
}
}`
var cfg BotConfig
if err := json.Unmarshal([]byte(jsonData), &cfg); err != nil {
t.Fatalf("unmarshal: %v", err)
}
if cfg.WebSearch == nil {
t.Fatalf("WebSearch = nil, want populated")
}
if len(cfg.WebSearch.AllowedDomains) != 2 {
t.Errorf("AllowedDomains = %v, want 2 entries", cfg.WebSearch.AllowedDomains)
}
if cfg.WebSearch.MaxUses != 3 {
t.Errorf("MaxUses = %d, want 3", cfg.WebSearch.MaxUses)
}
if !cfg.WebSearch.Fetch {
t.Error("Fetch = false, want true")
}
if cfg.WebSearch.MaxContentTokens != 50000 {
t.Errorf("MaxContentTokens = %d, want 50000", cfg.WebSearch.MaxContentTokens)
}
})
cases := []struct {
name string
ws *WebSearchConfig
wantErr string
}{
{"nil ok", nil, ""},
{"allowlist ok", &WebSearchConfig{AllowedDomains: []string{"example.com"}, MaxUses: 3}, ""},
{"blocklist ok", &WebSearchConfig{BlockedDomains: []string{"evil.com"}}, ""},
{"empty ok", &WebSearchConfig{}, ""},
{"both allow and block rejected",
&WebSearchConfig{AllowedDomains: []string{"a.com"}, BlockedDomains: []string{"b.com"}},
"cannot set both allowed_domains and blocked_domains"},
{"negative max_uses rejected",
&WebSearchConfig{AllowedDomains: []string{"a.com"}, MaxUses: -1},
"'web_search.max_uses' must be greater than 0"},
{"negative max_content_tokens rejected",
&WebSearchConfig{AllowedDomains: []string{"a.com"}, MaxContentTokens: -1},
"'web_search.max_content_tokens' must be greater than 0"},
}
for i, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
cfg := thinkingTestConfig(fmt.Sprintf("bot-web-%d", i))
cfg.WebSearch = tc.ws
err := validateConfig(&cfg, map[string]bool{}, map[string]bool{})
if tc.wantErr == "" {
if err != nil {
t.Fatalf("validateConfig(web_search=%+v) = %v, want nil", tc.ws, err)
}
return
}
if err == nil {
t.Fatalf("validateConfig(web_search=%+v) = nil, want error containing %q", tc.ws, tc.wantErr)
}
if !contains(err.Error(), tc.wantErr) {
t.Errorf("error %q does not contain %q", err.Error(), tc.wantErr)
}
})
}
}
+119
View File
@@ -0,0 +1,119 @@
package main
import (
"fmt"
"log"
"os"
"time"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
"gorm.io/gorm/logger"
)
func initDB() (*gorm.DB, error) {
if err := os.MkdirAll("data", 0750); err != nil {
return nil, fmt.Errorf("failed to create data directory: %w", err)
}
newLogger := logger.New(
log.New(log.Writer(), "\r\n", log.LstdFlags),
logger.Config{
SlowThreshold: time.Second,
LogLevel: logger.Info,
Colorful: false,
},
)
db, err := gorm.Open(sqlite.Open("data/bot.db?_journal_mode=WAL&_busy_timeout=5000&_foreign_keys=on"), &gorm.Config{
Logger: newLogger,
})
if err != nil {
return nil, fmt.Errorf("failed to connect to database: %w", err)
}
sqlDB, err := db.DB()
if err != nil {
return nil, fmt.Errorf("failed to get underlying sql.DB: %w", err)
}
sqlDB.SetMaxOpenConns(1)
err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{}, &Scope{})
if err != nil {
return nil, fmt.Errorf("failed to migrate database schema: %w", err)
}
err = db.Exec(`CREATE UNIQUE INDEX IF NOT EXISTS idx_bot_owner ON users (bot_id, is_owner) WHERE is_owner = 1;`).Error
if err != nil {
return nil, fmt.Errorf("failed to create unique index for bot owners: %w", err)
}
err = createDefaultRoles(db)
if err != nil {
return nil, err
}
if err := createDefaultScopes(db); err != nil {
return nil, fmt.Errorf("createDefaultScopes: %w", err)
}
return db, nil
}
func createDefaultScopes(db *gorm.DB) error {
all := []string{
ScopeStatsViewOwn, ScopeStatsViewAny,
ScopeHistoryClearOwn, ScopeHistoryClearAny,
ScopeHistoryClearHardOwn, ScopeHistoryClearHardAny,
ScopeModelSet, ScopeUserPromote, ScopeTTSUse,
}
for _, name := range all {
if err := db.FirstOrCreate(&Scope{}, Scope{Name: name}).Error; err != nil {
return fmt.Errorf("failed to create scope %s: %w", name, err)
}
}
userScopes := []string{
ScopeStatsViewOwn,
ScopeHistoryClearOwn,
ScopeHistoryClearHardOwn,
}
elevatedScopes := []string{
ScopeStatsViewOwn, ScopeStatsViewAny,
ScopeHistoryClearOwn, ScopeHistoryClearAny,
ScopeHistoryClearHardOwn, ScopeHistoryClearHardAny,
ScopeModelSet, ScopeUserPromote, ScopeTTSUse,
}
assignments := map[string][]string{
"user": userScopes,
"admin": elevatedScopes,
"owner": elevatedScopes,
}
for roleName, scopes := range assignments {
var role Role
if err := db.Where("name = ?", roleName).First(&role).Error; err != nil {
return fmt.Errorf("role %s not found: %w", roleName, err)
}
var scopeModels []Scope
if err := db.Where("name IN ?", scopes).Find(&scopeModels).Error; err != nil {
return fmt.Errorf("failed to find scopes for %s: %w", roleName, err)
}
if err := db.Model(&role).Association("Scopes").Replace(scopeModels); err != nil {
return fmt.Errorf("failed to assign scopes to %s: %w", roleName, err)
}
}
return nil
}
func createDefaultRoles(db *gorm.DB) error {
roles := []string{"user", "admin", "owner"}
for _, roleName := range roles {
var role Role
if err := db.FirstOrCreate(&role, Role{Name: roleName}).Error; err != nil {
ErrorLogger.Printf("Failed to create default role %s: %v", roleName, err)
return fmt.Errorf("failed to create default role %s: %w", roleName, err)
}
InfoLogger.Printf("Created or confirmed default role: %s", roleName)
}
return nil
}
+39
View File
@@ -0,0 +1,39 @@
services:
telegram-bot:
image: bogerserge/go-telegram-bot:latest
build:
context: .
dockerfile: Dockerfile
platforms:
- linux/amd64
- linux/arm64
container_name: go-telegram-bot
restart: unless-stopped
# Optional: Environment variables (can be overridden with .env file)
# environment:
# - BOT_LOG_LEVEL=info
# Volume mounts
volumes:
# Bind mount config directory for live configuration updates
- ./config:/app/config:ro
# Named volume for persistent database storage
- ./data:/app/data
# Optional: Bind mount for log access (uncomment if needed)
# - ./logs:/app/logs
# Health check
healthcheck:
test: ["CMD", "pgrep", "telegram-bot"]
interval: 30s
timeout: 10s
retries: 3
start_period: 10s
# Logging configuration
logging:
driver: "json-file"
options:
max-size: "10m"
max-file: "3"
+99
View File
@@ -0,0 +1,99 @@
package main
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"mime/multipart"
"net/http"
)
const (
elevenLabsTTSURL = "https://api.elevenlabs.io/v1/text-to-speech/"
elevenLabsSTTURL = "https://api.elevenlabs.io/v1/speech-to-text"
elevenLabsDefaultModel = "eleven_multilingual_v2"
)
func (b *Bot) generateSpeech(ctx context.Context, text string) (io.Reader, error) {
model := b.config.ElevenLabsModel
if model == "" {
model = elevenLabsDefaultModel
}
body, err := json.Marshal(map[string]string{
"text": text,
"model_id": model,
})
if err != nil {
return nil, fmt.Errorf("elevenlabs TTS marshal error: %w", err)
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost,
elevenLabsTTSURL+b.config.ElevenLabsVoiceID, bytes.NewReader(body))
if err != nil {
return nil, fmt.Errorf("elevenlabs TTS request error: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("xi-api-key", b.config.ElevenLabsAPIKey)
resp, err := http.DefaultClient.Do(req)
if err != nil {
return nil, fmt.Errorf("elevenlabs TTS error: %w", err)
}
if resp.StatusCode != http.StatusOK {
defer func() { _ = resp.Body.Close() }()
errBody, _ := io.ReadAll(resp.Body)
return nil, fmt.Errorf("elevenlabs TTS error: status %d: %s", resp.StatusCode, errBody)
}
return resp.Body, nil
}
func (b *Bot) transcribeVoice(ctx context.Context, fileID string) (string, error) {
audioBytes, err := b.downloadTelegramFile(ctx, fileID)
if err != nil {
return "", err
}
var buf bytes.Buffer
mw := multipart.NewWriter(&buf)
if err := mw.WriteField("model_id", "scribe_v1"); err != nil {
return "", fmt.Errorf("multipart write error: %w", err)
}
part, err := mw.CreateFormFile("file", "audio.ogg")
if err != nil {
return "", fmt.Errorf("multipart create file error: %w", err)
}
if _, err := io.Copy(part, bytes.NewReader(audioBytes)); err != nil {
return "", fmt.Errorf("multipart copy error: %w", err)
}
if err := mw.Close(); err != nil {
return "", fmt.Errorf("multipart close error: %w", err)
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost,
elevenLabsSTTURL, &buf)
if err != nil {
return "", fmt.Errorf("create STT request error: %w", err)
}
req.Header.Set("Content-Type", mw.FormDataContentType())
req.Header.Set("xi-api-key", b.config.ElevenLabsAPIKey)
sttResp, err := http.DefaultClient.Do(req)
if err != nil {
return "", fmt.Errorf("elevenlabs STT request error: %w", err)
}
defer func() { _ = sttResp.Body.Close() }()
if sttResp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(sttResp.Body)
return "", fmt.Errorf("elevenlabs STT error: status %d: %s", sttResp.StatusCode, body)
}
var result struct {
Text string `json:"text"`
}
if err := json.NewDecoder(sttResp.Body).Decode(&result); err != nil {
return "", fmt.Errorf("elevenlabs STT decode error: %w", err)
}
return result.Text, nil
}
+35
View File
@@ -0,0 +1,35 @@
[Unit]
# A concise description of the service
Description=Telegram Bot Service
# Postpone starting until network is available
After=network.target
[Service]
# The user that runs the bot
User=tibik
# The directory where the bot is located
WorkingDirectory=/home/tibik/go-telegram-bot
# The command to start the bot
ExecStart=/home/tibik/go-telegram-bot/telegram-bot
# Restart if crashed
Restart=always
# Delay between restarts to avoid resource exhaustion
RestartSec=5
# Capture stdout (INFO logs)
StandardOutput=journal
# Capture stderr (ERROR logs)
StandardError=journal
# Identifier for journalctl filtering
SyslogIdentifier=telegram-bot
[Install]
# The bot will start automatically at system boot
WantedBy=multi-user.target
# NOTE:
# New line comments: good
# Inline comments: no good, they mess up the service file
# View logs: journalctl -u telegram-bot
# Follow logs: journalctl -u telegram-bot -f
# View errors: journalctl -u telegram-bot -p err
Binary file not shown.
+38 -2
View File
@@ -1,3 +1,39 @@
module github.com/HugeFrog24/thatsky-telegram-bot
module github.com/HugeFrog24/go-telegram-bot
go 1.23.2
go 1.26.0
require (
github.com/anthropics/anthropic-sdk-go v1.57.0
github.com/go-telegram/bot v1.22.0
github.com/stretchr/testify v1.11.1
golang.org/x/sync v0.22.0
golang.org/x/time v0.15.0
gorm.io/driver/sqlite v1.6.0
gorm.io/gorm v1.31.2
)
require (
github.com/bahlo/generic-list-go v0.2.0 // indirect
github.com/buger/jsonparser v1.2.0 // indirect
github.com/davecgh/go-spew v1.1.1 // indirect
github.com/invopop/jsonschema v0.14.0 // indirect
github.com/jinzhu/inflection v1.0.0 // indirect
github.com/jinzhu/now v1.1.5 // indirect
github.com/kr/pretty v0.3.1 // indirect
github.com/mailru/easyjson v0.9.2 // indirect
github.com/mattn/go-sqlite3 v1.14.48 // indirect
github.com/pb33f/ordered-map/v2 v2.3.1 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect
github.com/rogpeppe/go-internal v1.14.1 // indirect
github.com/standard-webhooks/standard-webhooks/libraries v0.0.1 // indirect
github.com/stretchr/objx v0.5.3 // indirect
github.com/tidwall/gjson v1.19.0 // indirect
github.com/tidwall/match v1.2.0 // indirect
github.com/tidwall/pretty v1.2.1 // indirect
github.com/tidwall/sjson v1.2.5 // indirect
github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect
go.yaml.in/yaml/v4 v4.0.0-rc.6 // indirect
golang.org/x/text v0.40.0 // indirect
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
gopkg.in/yaml.v3 v3.0.1 // indirect
)
+100
View File
@@ -0,0 +1,100 @@
github.com/anthropics/anthropic-sdk-go v1.45.0 h1:rWnpyBpm9OAm97jyH5bi6W4SRCwJeNY/RyhaJ7CHSUI=
github.com/anthropics/anthropic-sdk-go v1.45.0/go.mod h1:bx5vWuHFuGPkELH8Z4KUiNSohFnUwScdpTyr+50myPo=
github.com/anthropics/anthropic-sdk-go v1.52.0 h1:1TB9jt4DN87VMwS/hB1VK26tYzK0ipEOtqPaPGFtJQg=
github.com/anthropics/anthropic-sdk-go v1.52.0/go.mod h1:3EfIfmFqxH6rbiLcIP4tPFyXL/IHakx2wDG4OU+TIEI=
github.com/anthropics/anthropic-sdk-go v1.57.0 h1:iEAcPbUKfJ2Iqz9uN/jEndCNW2+x7OYLHDidXDhPjI0=
github.com/anthropics/anthropic-sdk-go v1.57.0/go.mod h1:3EfIfmFqxH6rbiLcIP4tPFyXL/IHakx2wDG4OU+TIEI=
github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk=
github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg=
github.com/buger/jsonparser v1.2.0 h1:4EFcvK1kD4jyj6YqNK6skK6w+y7FHHBR+XBCtxwu/6g=
github.com/buger/jsonparser v1.2.0/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0=
github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/dnaeon/go-vcr v1.2.0 h1:zHCHvJYTMh1N7xnV7zf1m1GPBF9Ad0Jk/whtQ1663qI=
github.com/dnaeon/go-vcr v1.2.0/go.mod h1:R4UdLID7HZT3taECzJs4YgbbH6PIGXB6W/sc5OLb6RQ=
github.com/go-telegram/bot v1.20.0 h1:4Pea/qTidSspr4WBJw9FbHUMNhYeqszBqQUfsQEyFbc=
github.com/go-telegram/bot v1.20.0/go.mod h1:i2TRs7fXWIeaceF3z7KzsMt/he0TwkVC680mvdTFYeM=
github.com/go-telegram/bot v1.21.0 h1:Va/PbGc2vBDdv57GCUEEVV6ROlHWiC6SklJY9Hvhzps=
github.com/go-telegram/bot v1.21.0/go.mod h1:i2TRs7fXWIeaceF3z7KzsMt/he0TwkVC680mvdTFYeM=
github.com/go-telegram/bot v1.22.0 h1:zK29OoTYMmR5emJrCtGa2SjaGleeZiUB/C1i7kc2lXE=
github.com/go-telegram/bot v1.22.0/go.mod h1:i2TRs7fXWIeaceF3z7KzsMt/he0TwkVC680mvdTFYeM=
github.com/invopop/jsonschema v0.13.0 h1:KvpoAJWEjR3uD9Kbm2HWJmqsEaHt8lBUpd0qHcIi21E=
github.com/invopop/jsonschema v0.13.0/go.mod h1:ffZ5Km5SWWRAIN6wbDXItl95euhFz2uON45H2qjYt+0=
github.com/invopop/jsonschema v0.14.0 h1:MHQqLhvpNUZfw+hM3AZDYK7jxO8FZoQeQM77g8iyZjg=
github.com/invopop/jsonschema v0.14.0/go.mod h1:ygm6C2EaVNMBDPpaPlnOA2pFAxBnxGjFlMZABxm9n2I=
github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E=
github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc=
github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ=
github.com/jinzhu/now v1.1.5/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
github.com/mailru/easyjson v0.9.2 h1:dX8U45hQsZpxd80nLvDGihsQ/OxlvTkVUXH2r/8cb2M=
github.com/mailru/easyjson v0.9.2/go.mod h1:1+xMtQp2MRNVL/V1bOzuP3aP8VNwRW55fQUto+XFtTU=
github.com/mattn/go-sqlite3 v1.14.44 h1:3VSe+xafpbzsLbdr2AWlAZk9yRHiBhTBakioXaCKTF8=
github.com/mattn/go-sqlite3 v1.14.44/go.mod h1:pjEuOr8IwzLJP2MfGeTb0A35jauH+C2kbHKBr7yXKVQ=
github.com/mattn/go-sqlite3 v1.14.47 h1:jOBI62gS7nKeZv+as1oGEy0+1qISgXwH/QBlR6KbfIo=
github.com/mattn/go-sqlite3 v1.14.47/go.mod h1:6JTjA44L93a0QCyJef5YvlPoKXntQPjzWv5gtm9sB6w=
github.com/mattn/go-sqlite3 v1.14.48 h1:7XHIgl0a8HwOaiK4E47ozLkST78rR9+OtNGx27D/TFs=
github.com/mattn/go-sqlite3 v1.14.48/go.mod h1:6JTjA44L93a0QCyJef5YvlPoKXntQPjzWv5gtm9sB6w=
github.com/pb33f/ordered-map/v2 v2.3.1 h1:5319HDO0aw4DA4gzi+zv4FXU9UlSs3xGZ40wcP1nBjY=
github.com/pb33f/ordered-map/v2 v2.3.1/go.mod h1:qxFQgd0PkVUtOMCkTapqotNgzRhMPL7VvaHKbd1HnmQ=
github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs=
github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ=
github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc=
github.com/standard-webhooks/standard-webhooks/libraries v0.0.1 h1:uOfcYT+3QungH6tIGSVCR/Y3KJmgJiHcojJbMTPDZAI=
github.com/standard-webhooks/standard-webhooks/libraries v0.0.1/go.mod h1:L1MQhA6x4dn9r007T033lsaZMv9EmBAdXyU/+EF40fo=
github.com/stretchr/objx v0.5.3 h1:jmXUvGomnU1o3W/V5h2VEradbpJDwGrzugQQvL0POH4=
github.com/stretchr/objx v0.5.3/go.mod h1:rDQraq+vQZU7Fde9LOZLr8Tax6zZvy4kuNKF+QYS+U0=
github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
github.com/tidwall/gjson v1.19.0 h1:xwxm7n691Uf3u5OFjzngavjGTh55KX5q/9w9xHW88JU=
github.com/tidwall/gjson v1.19.0/go.mod h1:V37/opeE/JbLUOfH0QTXiNez2l0RUjYUhpT4szFQAfc=
github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM=
github.com/tidwall/match v1.2.0 h1:0pt8FlkOwjN2fPt4bIl4BoNxb98gGHN2ObFEDkrfZnM=
github.com/tidwall/match v1.2.0/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM=
github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU=
github.com/tidwall/pretty v1.2.1 h1:qjsOFOWWQl+N3RsoF5/ssm1pHmJJwhjlSbZ51I6wMl4=
github.com/tidwall/pretty v1.2.1/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU=
github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY=
github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28=
github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc=
github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw=
go.yaml.in/yaml/v4 v4.0.0-rc.6 h1:1h7H1ohdUh93/FyE4YaDa1Zh64K6VVbjF4K6WUxMtH4=
go.yaml.in/yaml/v4 v4.0.0-rc.6/go.mod h1:aZqd9kCMsGL7AuUv/m/PvWLdg5sjJsZ4oHDEnfPPfY0=
golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4=
golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
golang.org/x/sync v0.21.0 h1:HLII4xRRTtCRkxYp4HNFF0Js/Og6q2i++KXbg0gHCwM=
golang.org/x/sync v0.21.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek=
golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
golang.org/x/text v0.37.0 h1:Cqjiwd9eSg8e0QAkyCaQTNHFIIzWtidPahFWR83rTrc=
golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38=
golang.org/x/text v0.38.0 h1:sXmwo9DwP3OK9EZ7PqAdaooSGozfl/3a6/xJcbzPRhE=
golang.org/x/text v0.38.0/go.mod h1:YXZt3QhHUKYT53r2lLKFIVi6Ao1jdzrTR/KQ09qyxF4=
golang.org/x/text v0.40.0 h1:Ub2Z6/xjgF1WrYQz2nuITOEegKFtiIy+rieRJ5lHZKs=
golang.org/x/text v0.40.0/go.mod h1:hpnzDAfGV753zIKo+wk3u1bVKCGPbrnF7+7LBF/UHVY=
golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U=
golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
gopkg.in/yaml.v2 v2.2.8 h1:obN1ZagJSUGI0Ek/LBmuj4SNLPfIny3KsKFopxRdj10=
gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
gorm.io/driver/sqlite v1.6.0 h1:WHRRrIiulaPiPFmDcod6prc4l2VGVWHz80KspNsxSfQ=
gorm.io/driver/sqlite v1.6.0/go.mod h1:AO9V1qIQddBESngQUKWL9yoH93HIeA1X6V633rBwyT8=
gorm.io/gorm v1.31.1 h1:7CA8FTFz/gRfgqgpeKIBcervUn3xSyPUmr6B2WXJ7kg=
gorm.io/gorm v1.31.1/go.mod h1:XyQVbO2k6YkOis7C2437jSit3SsDK72s7n7rsSHd+Gs=
gorm.io/gorm v1.31.2 h1:3o8FXNo9v9S858gil+3LlZA1LkCOzgb4g5BL64FgaCo=
gorm.io/gorm v1.31.2/go.mod h1:XyQVbO2k6YkOis7C2437jSit3SsDK72s7n7rsSHd+Gs=
+598
View File
@@ -0,0 +1,598 @@
package main
import (
"context"
"errors"
"fmt"
"strconv"
"strings"
"github.com/anthropics/anthropic-sdk-go"
"github.com/go-telegram/bot"
"github.com/go-telegram/bot/models"
"golang.org/x/sync/errgroup"
)
func (b *Bot) handleVoiceMessage(ctx context.Context, message *models.Message, userMsg Message, chatID, userID int64, username, firstName, lastName string, isPremium bool, languageCode string, messageTime int, businessConnectionID string) {
if b.config.ElevenLabsAPIKey == "" {
if err := b.sendResponse(ctx, chatID, "I don't understand voice messages.", businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending voice-unsupported message: %v", err)
}
return
}
if !b.hasScope(userID, ScopeTTSUse) {
if err := b.sendResponse(ctx, chatID, "You don't have permission to use voice features.", businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending permission denied message: %v", err)
}
return
}
transcript, err := b.transcribeVoice(ctx, message.Voice.FileID)
if err != nil {
ErrorLogger.Printf("Error transcribing voice message from user %d: %v", userID, err)
if err := b.sendResponse(ctx, chatID, "Sorry, I couldn't understand your voice message.", businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending transcription error message: %v", err)
}
return
}
if err := b.db.Model(&userMsg).Update("text", transcript).Error; err != nil {
ErrorLogger.Printf("Error updating voice transcript in DB: %v", err)
}
b.chatMemoriesMu.Lock()
if mem, exists := b.chatMemories[chatID]; exists {
for i := len(mem.Messages) - 1; i >= 0; i-- {
if mem.Messages[i].ID == userMsg.ID {
mem.Messages[i].Text = transcript
break
}
}
}
b.chatMemoriesMu.Unlock()
chatMemory := b.getOrCreateChatMemory(chatID)
contextMessages := b.prepareContextMessages(chatMemory)
response, err := b.getAnthropicResponse(ctx, chatID, contextMessages, false, username, firstName, lastName, isPremium, languageCode, messageTime, nil)
if err != nil {
ErrorLogger.Printf("Error getting Anthropic response for voice: %v", err)
if err := b.sendResponse(ctx, chatID, b.anthropicErrorResponse(err, userID), businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending anthropic error response: %v", err)
}
return
}
audioReader, err := b.generateSpeech(ctx, response)
if err != nil {
ErrorLogger.Printf("Error generating speech, falling back to text: %v", err)
if err := b.sendResponse(ctx, chatID, response, businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending text fallback: %v", err)
}
return
}
if _, err := b.screenOutgoingMessage(chatID, response); err != nil {
ErrorLogger.Printf("Error storing assistant voice response: %v", err)
}
params := &bot.SendAudioParams{
ChatID: chatID,
Audio: &models.InputFileUpload{Filename: "response.mp3", Data: audioReader},
}
if businessConnectionID != "" {
params.BusinessConnectionID = businessConnectionID
}
if _, err := b.tgBot.SendAudio(ctx, params); err != nil {
ErrorLogger.Printf("Error sending audio to chat %d: %v", chatID, err)
}
}
func (b *Bot) uploadPhotoFromItem(ctx context.Context, item *models.Message, chatID int64) (string, error) {
photo := largestPhotoSize(item.Photo)
data, err := b.downloadTelegramFile(ctx, photo.FileID)
if err != nil {
return "", fmt.Errorf("download %s: %w", photo.FileID, err)
}
filename := formatUploadFilename(b.botID, chatID, item.ID, "jpg")
return b.uploadImageToAnthropic(ctx, data, filename, "image/jpeg")
}
func (b *Bot) handlePhotoMessage(
ctx context.Context,
items []*models.Message,
chatID, userID int64,
username, firstName, lastName string,
isPremium bool,
languageCode string,
messageTime int,
businessConnectionID string,
) {
if len(items) == 0 {
return
}
uploaded := make([]string, len(items))
caption := ""
g, gctx := errgroup.WithContext(ctx)
for i, item := range items {
if item.Caption != "" {
caption = item.Caption
}
if len(item.Photo) == 0 {
continue
}
i, item := i, item
g.Go(func() error {
fileID, err := b.uploadPhotoFromItem(gctx, item, chatID)
if err != nil {
return err
}
uploaded[i] = fileID
return nil
})
}
if err := g.Wait(); err != nil {
ErrorLogger.Printf("[%s] photo upload failed: %v", b.config.ID, err)
var successful []string
for _, fid := range uploaded {
if fid != "" {
successful = append(successful, fid)
}
}
b.compensatingDelete(ctx, successful)
if sendErr := b.sendResponse(ctx, chatID, "Sorry, I couldn't process one of your images.", businessConnectionID); sendErr != nil {
ErrorLogger.Printf("Error sending photo failure message: %v", sendErr)
}
return
}
finalUploaded := make([]string, 0, len(uploaded))
for _, fid := range uploaded {
if fid != "" {
finalUploaded = append(finalUploaded, fid)
}
}
if len(finalUploaded) == 0 {
return
}
chatMemory := b.getOrCreateChatMemory(chatID)
userMessage := b.createMessage(chatID, userID, username, "user", caption, true)
userMessage.ImageFileIDs = finalUploaded
if err := b.storeMessage(&userMessage); err != nil {
b.compensatingDelete(ctx, finalUploaded)
ErrorLogger.Printf("[%s] store photo message failed: %v", b.config.ID, err)
if sendErr := b.sendResponse(ctx, chatID, "Sorry, I had trouble saving your message.", businessConnectionID); sendErr != nil {
ErrorLogger.Printf("Error sending store failure message: %v", sendErr)
}
return
}
b.addMessageToChatMemory(chatMemory, userMessage)
contextMessages := b.prepareContextMessages(chatMemory)
joined, err := b.getAnthropicResponse(
ctx, chatID, contextMessages, false,
username, firstName, lastName, isPremium, languageCode, messageTime,
func(seg string) error {
return b.sendOneSegment(ctx, chatID, seg, businessConnectionID)
},
)
if err != nil {
ErrorLogger.Printf("Error getting Anthropic response for photo: %v", err)
if sendErr := b.sendResponse(ctx, chatID, b.anthropicErrorResponse(err, userID), businessConnectionID); sendErr != nil {
ErrorLogger.Printf("Error sending anthropic error response: %v", sendErr)
}
return
}
if _, storeErr := b.screenOutgoingMessage(chatID, joined); storeErr != nil {
ErrorLogger.Printf("Error recording assistant turn: %v", storeErr)
}
}
func (b *Bot) anthropicErrorResponse(err error, userID int64) string {
isElevated := b.hasScope(userID, ScopeModelSet)
if errors.Is(err, ErrModelNotFound) && isElevated {
return fmt.Sprintf(
"⚠️ Model `%s` is no longer available (deprecated or removed by Anthropic).\n"+
"Use /set_model <model-id> to switch. Current models: https://platform.claude.com/docs/en/about-claude/models/overview",
b.config.Model,
)
}
if isElevated {
var apiErr *anthropic.Error
if errors.As(err, &apiErr) {
body := apiErr.RawJSON()
if len(body) > 800 {
body = body[:800] + "...(truncated)"
}
out := fmt.Sprintf("⚠️ Anthropic API error %d:\n%s", apiErr.StatusCode, body)
if apiErr.RequestID != "" {
out += fmt.Sprintf("\nRequest-ID: %s", apiErr.RequestID)
}
return out
}
return fmt.Sprintf("⚠️ Anthropic call failed: %v", err)
}
return "I'm sorry, I'm having trouble processing your request right now."
}
func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.Update) {
var message *models.Message
if update.Message != nil {
message = update.Message
} else if update.BusinessMessage != nil {
message = update.BusinessMessage
} else {
return
}
var businessConnectionID string
if update.BusinessConnection != nil {
businessConnectionID = update.BusinessConnection.ID
} else if message.BusinessConnectionID != "" {
businessConnectionID = message.BusinessConnectionID
}
if message.From == nil {
return
}
chatID := message.Chat.ID
userID := message.From.ID
username := message.From.Username
firstName := message.From.FirstName
lastName := message.From.LastName
languageCode := message.From.LanguageCode
isPremium := message.From.IsPremium
messageTime := message.Date
text := message.Text
var isOwner bool
if b.db.Where("telegram_id = ? AND bot_id = ? AND is_owner = ?", userID, b.botID, true).First(&User{}).Error == nil {
isOwner = true
}
user, err := b.getOrCreateUser(userID, username, isOwner)
if err != nil {
ErrorLogger.Printf("Error getting or creating user: %v", err)
return
}
if user.Username != username {
user.Username = username
if err := b.db.Save(&user).Error; err != nil {
ErrorLogger.Printf("Error updating user username: %v", err)
}
}
if message.MediaGroupID != "" && len(message.Photo) > 0 {
b.bufferAlbumItem(ctx, message, chatID, userID, username, firstName, lastName,
isPremium, languageCode, messageTime, businessConnectionID)
return
}
if len(message.Photo) > 0 {
if !b.checkRateLimits(userID) {
b.sendRateLimitExceededMessage(ctx, chatID, businessConnectionID)
return
}
b.handlePhotoMessage(ctx, []*models.Message{message},
chatID, userID, username, firstName, lastName,
isPremium, languageCode, messageTime,
businessConnectionID)
return
}
userMsg, err := b.screenIncomingMessage(message)
if err != nil {
ErrorLogger.Printf("Error storing user message: %v", err)
return
}
if message.Entities != nil {
for _, entity := range message.Entities {
if entity.Type == "bot_command" {
command := strings.TrimSpace(message.Text[entity.Offset : entity.Offset+entity.Length])
switch command {
case "/stats":
parts := strings.Fields(message.Text)
if len(parts) == 1 {
b.sendStats(ctx, chatID, userID, 0, businessConnectionID)
return
}
if len(parts) >= 2 && parts[1] == "user" {
targetUserID := userID
if len(parts) >= 3 {
var parseErr error
targetUserID, parseErr = strconv.ParseInt(parts[2], 10, 64)
if parseErr != nil {
InfoLogger.Printf("User %d provided invalid user ID format: %s", userID, parts[2])
if err := b.sendResponse(ctx, chatID, "Invalid user ID format. Usage: /stats user [user_id]", businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending response: %v", err)
}
return
}
}
b.sendStats(ctx, chatID, userID, targetUserID, businessConnectionID)
return
}
if err := b.sendResponse(ctx, chatID, "Invalid command format. Usage: /stats or /stats user [user_id]", businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending response: %v", err)
}
return
case "/whoami":
b.sendWhoAmI(ctx, chatID, userID, username, businessConnectionID)
return
case "/clear":
parts := strings.Fields(message.Text)
var targetUserID, targetChatID int64
if len(parts) > 1 {
var parseErr error
targetUserID, parseErr = strconv.ParseInt(parts[1], 10, 64)
if parseErr != nil {
InfoLogger.Printf("User %d provided invalid user ID format: %s", userID, parts[1])
if err := b.sendResponse(ctx, chatID, "Invalid user ID format. Usage: /clear [user_id] [chat_id]", businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending response: %v", err)
}
return
}
}
if len(parts) > 2 {
var parseErr error
targetChatID, parseErr = strconv.ParseInt(parts[2], 10, 64)
if parseErr != nil {
InfoLogger.Printf("User %d provided invalid chat ID format: %s", userID, parts[2])
if err := b.sendResponse(ctx, chatID, "Invalid chat ID format. Usage: /clear [user_id] [chat_id]", businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending response: %v", err)
}
return
}
}
b.clearChatHistory(ctx, chatID, userID, targetUserID, targetChatID, businessConnectionID, false)
return
case "/set_model":
if !b.hasScope(userID, ScopeModelSet) {
if err := b.sendResponse(ctx, chatID, "Permission denied. Only admins and owners can change the model.", businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending response: %v", err)
}
return
}
parts := strings.Fields(message.Text)
if len(parts) < 2 || strings.TrimSpace(parts[1]) == "" {
if err := b.sendResponse(ctx, chatID, "Usage: /set_model <model-id>", businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending response: %v", err)
}
return
}
newModel := strings.TrimSpace(parts[1])
if err := b.config.PersistModel(newModel); err != nil {
ErrorLogger.Printf("Failed to persist model change: %v", err)
if err := b.sendResponse(ctx, chatID, fmt.Sprintf("Model updated in memory to `%s`, but failed to save to config file: %v", newModel, err), businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending response: %v", err)
}
return
}
InfoLogger.Printf("Model changed to %s by user %d", newModel, userID)
if err := b.sendResponse(ctx, chatID, fmt.Sprintf("✅ Model updated to `%s` and saved to config.", newModel), businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending response: %v", err)
}
return
case "/clear_hard":
parts := strings.Fields(message.Text)
var targetUserID, targetChatID int64
if len(parts) > 1 {
var parseErr error
targetUserID, parseErr = strconv.ParseInt(parts[1], 10, 64)
if parseErr != nil {
InfoLogger.Printf("User %d provided invalid user ID format: %s", userID, parts[1])
if err := b.sendResponse(ctx, chatID, "Invalid user ID format. Usage: /clear_hard [user_id] [chat_id]", businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending response: %v", err)
}
return
}
}
if len(parts) > 2 {
var parseErr error
targetChatID, parseErr = strconv.ParseInt(parts[2], 10, 64)
if parseErr != nil {
InfoLogger.Printf("User %d provided invalid chat ID format: %s", userID, parts[2])
if err := b.sendResponse(ctx, chatID, "Invalid chat ID format. Usage: /clear_hard [user_id] [chat_id]", businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending response: %v", err)
}
return
}
}
b.clearChatHistory(ctx, chatID, userID, targetUserID, targetChatID, businessConnectionID, true)
return
}
}
}
}
if !b.checkRateLimits(userID) {
b.sendRateLimitExceededMessage(ctx, chatID, businessConnectionID)
return
}
if message.Voice != nil {
b.handleVoiceMessage(ctx, message, userMsg, chatID, userID, username, firstName, lastName, isPremium, languageCode, messageTime, businessConnectionID)
return
}
chatMemory := b.getOrCreateChatMemory(chatID)
contextMessages := b.prepareContextMessages(chatMemory)
if message.Sticker != nil {
b.handleStickerMessage(ctx, chatID, userMsg, message, contextMessages, businessConnectionID)
return
}
if text == "" {
InfoLogger.Printf("Received a non-text message from user %d in chat %d", userID, chatID)
return
}
isEmojiOnly := isOnlyEmojis(text)
joined, err := b.getAnthropicResponse(
ctx, chatID, contextMessages, isEmojiOnly,
username, firstName, lastName, isPremium, languageCode, messageTime,
func(seg string) error {
return b.sendOneSegment(ctx, chatID, seg, businessConnectionID)
},
)
if err != nil {
ErrorLogger.Printf("Error getting Anthropic response: %v", err)
if sendErr := b.sendResponse(ctx, chatID, b.anthropicErrorResponse(err, userID), businessConnectionID); sendErr != nil {
ErrorLogger.Printf("Error sending response: %v", sendErr)
}
return
}
if _, storeErr := b.screenOutgoingMessage(chatID, joined); storeErr != nil {
ErrorLogger.Printf("Error recording assistant turn: %v", storeErr)
}
}
func (b *Bot) sendRateLimitExceededMessage(ctx context.Context, chatID int64, businessConnectionID string) {
if err := b.sendResponse(ctx, chatID, "Rate limit exceeded. Please try again later.", businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending rate limit exceeded message: %v", err)
}
}
func (b *Bot) handleStickerMessage(ctx context.Context, chatID int64, userMessage Message, message *models.Message, contextMessages []anthropic.BetaMessageParam, businessConnectionID string) {
response, err := b.generateStickerResponse(ctx, userMessage, contextMessages)
if err != nil {
ErrorLogger.Printf("Error generating sticker response: %v", err)
if message.Sticker.IsAnimated {
response = "Wow, that's a cool animated sticker!"
} else if message.Sticker.IsVideo {
response = "Interesting video sticker!"
} else {
response = "That's a cool sticker!"
}
}
if err := b.sendResponse(ctx, chatID, response, businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending response: %v", err)
return
}
}
func (b *Bot) generateStickerResponse(ctx context.Context, message Message, contextMessages []anthropic.BetaMessageParam) (string, error) {
if message.StickerFileID != "" {
messageTime := int(message.Timestamp.Unix())
response, err := b.getAnthropicResponse(ctx, message.ChatID, contextMessages, true, message.Username, "", "", false, "", messageTime, nil)
if err != nil {
return "", err
}
return response, nil
}
return "Hmm, that's interesting!", nil
}
func (b *Bot) clearChatHistory(ctx context.Context, chatID int64, currentUserID int64, targetUserID int64, targetChatID int64, businessConnectionID string, hardDelete bool) {
if targetUserID != 0 && targetUserID != currentUserID {
requiredScope := ScopeHistoryClearAny
if hardDelete {
requiredScope = ScopeHistoryClearHardAny
}
if !b.hasScope(currentUserID, requiredScope) {
InfoLogger.Printf("User %d attempted to clear history for user %d without permission", currentUserID, targetUserID)
if err := b.sendResponse(ctx, chatID, "Permission denied. Only admins and owners can clear other users' histories.", businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending response: %v", err)
}
return
}
var targetUser User
err := b.db.Where("telegram_id = ? AND bot_id = ?", targetUserID, b.botID).First(&targetUser).Error
if err != nil {
ErrorLogger.Printf("Error finding target user %d: %v", targetUserID, err)
if err := b.sendResponse(ctx, chatID, fmt.Sprintf("User with ID %d not found.", targetUserID), businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending response: %v", err)
}
return
}
} else {
targetUserID = currentUserID
}
var err error
if hardDelete {
if targetUserID == currentUserID {
err = b.hardDeleteScope(ctx, "chat_id = ? AND bot_id = ?", chatID, b.botID)
InfoLogger.Printf("User %d permanently deleted their own chat history in chat %d", currentUserID, chatID)
} else {
if targetChatID != 0 {
err = b.hardDeleteScope(ctx, "chat_id = ? AND bot_id = ?", targetChatID, b.botID)
InfoLogger.Printf("Admin/owner %d permanently deleted chat history for user %d in chat %d", currentUserID, targetUserID, targetChatID)
} else {
err = b.hardDeleteScope(ctx,
"bot_id = ? AND (user_id = ? OR (chat_id = ? AND is_user = ?))",
b.botID, targetUserID, targetUserID, false)
InfoLogger.Printf("Admin/owner %d permanently deleted all chat history for user %d", currentUserID, targetUserID)
}
}
} else {
if targetUserID == currentUserID {
err = b.db.Where("chat_id = ? AND bot_id = ?", chatID, b.botID).Delete(&Message{}).Error
InfoLogger.Printf("User %d soft deleted their own chat history in chat %d", currentUserID, chatID)
} else {
if targetChatID != 0 {
err = b.db.Where("chat_id = ? AND bot_id = ?", targetChatID, b.botID).Delete(&Message{}).Error
InfoLogger.Printf("Admin/owner %d soft deleted chat history for user %d in chat %d", currentUserID, targetUserID, targetChatID)
} else {
err = b.db.Where("bot_id = ? AND user_id = ?", b.botID, targetUserID).Delete(&Message{}).Error
if err == nil {
err = b.db.Where("chat_id = ? AND bot_id = ? AND is_user = ?", targetUserID, b.botID, false).Delete(&Message{}).Error
}
InfoLogger.Printf("Admin/owner %d soft deleted all chat history for user %d", currentUserID, targetUserID)
}
}
}
if err != nil {
ErrorLogger.Printf("Error clearing chat history: %v", err)
if err := b.sendResponse(ctx, chatID, "Sorry, I couldn't clear the chat history.", businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending response: %v", err)
}
return
}
b.chatMemoriesMu.Lock()
if targetUserID == currentUserID {
delete(b.chatMemories, chatID)
} else if targetChatID != 0 {
delete(b.chatMemories, targetChatID)
} else {
delete(b.chatMemories, targetUserID)
}
b.chatMemoriesMu.Unlock()
var confirmationMessage string
if targetUserID == currentUserID {
confirmationMessage = "Your chat history has been cleared."
} else {
var targetUser User
err := b.db.Where("telegram_id = ? AND bot_id = ?", targetUserID, b.botID).First(&targetUser).Error
if err == nil && targetUser.Username != "" {
confirmationMessage = fmt.Sprintf("Chat history for user @%s (ID: %d) has been cleared.", targetUser.Username, targetUserID)
} else {
confirmationMessage = fmt.Sprintf("Chat history for user with ID %d has been cleared.", targetUserID)
}
}
if err := b.sendResponse(ctx, chatID, confirmationMessage, businessConnectionID); err != nil {
ErrorLogger.Printf("Error sending response: %v", err)
}
}
+830
View File
@@ -0,0 +1,830 @@
package main
import (
"context"
"errors"
"fmt"
"os"
"path/filepath"
"testing"
"time"
"github.com/go-telegram/bot"
"github.com/go-telegram/bot/models"
"github.com/stretchr/testify/assert"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
)
func TestHandleUpdate_NewChat(t *testing.T) {
db := setupTestDB(t)
mockClock := &MockClock{
currentTime: time.Now(),
}
config := BotConfig{
ID: "test_bot",
OwnerTelegramID: 123,
TelegramToken: "test_token",
MemorySize: 10,
MessagePerHour: 5,
MessagePerDay: 10,
TempBanDuration: "1h",
SystemPrompts: make(map[string]string),
Active: true,
}
mockTgClient := &MockTelegramClient{}
botModel := &BotModel{
Identifier: config.ID,
Name: config.ID,
}
err := db.Create(botModel).Error
assert.NoError(t, err)
configModel := &ConfigModel{
BotID: botModel.ID,
MemorySize: config.MemorySize,
MessagePerHour: config.MessagePerHour,
MessagePerDay: config.MessagePerDay,
TempBanDuration: config.TempBanDuration,
SystemPrompts: "{}",
TelegramToken: config.TelegramToken,
Active: config.Active,
}
err = db.Create(configModel).Error
assert.NoError(t, err)
b, err := NewBot(db, config, mockClock, mockTgClient)
assert.NoError(t, err)
testCases := []struct {
name string
userID int64
wantSubstr string
}{
{
name: "Owner First Message",
userID: 123,
wantSubstr: "Anthropic call failed:",
},
{
name: "Regular User First Message",
userID: 456,
wantSubstr: "I'm sorry, I'm having trouble processing your request right now.",
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
mockTgClient.SendMessageFunc = func(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error) {
assert.Equal(t, tc.userID, params.ChatID)
assert.Contains(t, params.Text, tc.wantSubstr)
return &models.Message{}, nil
}
update := &models.Update{
Message: &models.Message{
Chat: models.Chat{ID: tc.userID},
From: &models.User{
ID: tc.userID,
Username: "testuser",
},
Text: "Hello",
},
}
b.handleUpdate(context.Background(), nil, update)
var storedMsg Message
err := db.Where("chat_id = ? AND user_id = ? AND text = ?", tc.userID, tc.userID, "Hello").First(&storedMsg).Error
assert.NoError(t, err)
var respMsg Message
err = db.Where("chat_id = ? AND is_user = ?", tc.userID, false).
Order("timestamp DESC").
First(&respMsg).Error
assert.NoError(t, err)
assert.Contains(t, respMsg.Text, tc.wantSubstr)
})
}
}
func TestClearChatHistory(t *testing.T) {
db := setupTestDB(t)
mockClock := &MockClock{
currentTime: time.Now(),
}
config := BotConfig{
ID: "test_bot",
OwnerTelegramID: 123,
TelegramToken: "test_token",
MemorySize: 10,
MessagePerHour: 5,
MessagePerDay: 10,
TempBanDuration: "1h",
SystemPrompts: make(map[string]string),
Active: true,
}
mockTgClient := &MockTelegramClient{}
botModel := &BotModel{
Identifier: config.ID,
Name: config.ID,
}
err := db.Create(botModel).Error
assert.NoError(t, err)
configModel := &ConfigModel{
BotID: botModel.ID,
MemorySize: config.MemorySize,
MessagePerHour: config.MessagePerHour,
MessagePerDay: config.MessagePerDay,
TempBanDuration: config.TempBanDuration,
SystemPrompts: "{}",
TelegramToken: config.TelegramToken,
Active: config.Active,
}
err = db.Create(configModel).Error
assert.NoError(t, err)
b, err := NewBot(db, config, mockClock, mockTgClient)
assert.NoError(t, err)
ownerID := int64(123)
adminID := int64(456)
regularUserID := int64(789)
nonExistentUserID := int64(999)
chatID := int64(1000)
adminRole, err := b.getRoleByName("admin")
assert.NoError(t, err)
adminUser := User{
BotID: b.botID,
TelegramID: adminID,
Username: "admin",
RoleID: adminRole.ID,
Role: adminRole,
IsOwner: false,
}
err = db.Create(&adminUser).Error
assert.NoError(t, err)
regularRole, err := b.getRoleByName("user")
assert.NoError(t, err)
regularUser := User{
BotID: b.botID,
TelegramID: regularUserID,
Username: "regular",
RoleID: regularRole.ID,
Role: regularRole,
IsOwner: false,
}
err = db.Create(&regularUser).Error
assert.NoError(t, err)
for _, userID := range []int64{ownerID, adminID, regularUserID} {
for i := 0; i < 5; i++ {
message := Message{
BotID: b.botID,
ChatID: userID,
UserID: userID,
Username: "test",
UserRole: "user",
Text: "Test message",
Timestamp: time.Now(),
IsUser: true,
}
err = db.Create(&message).Error
assert.NoError(t, err)
}
}
testCases := []struct {
name string
currentUserID int64
targetUserID int64
hardDelete bool
expectedError bool
expectedCount int64
expectedMsg string
targetChatID int64
businessConnID string
}{
{
name: "Owner clears own history",
currentUserID: ownerID,
targetUserID: ownerID,
hardDelete: false,
expectedError: false,
expectedCount: 0,
expectedMsg: "Your chat history has been cleared.",
},
{
name: "Admin clears own history",
currentUserID: adminID,
targetUserID: adminID,
hardDelete: false,
expectedError: false,
expectedCount: 0,
expectedMsg: "Your chat history has been cleared.",
},
{
name: "Regular user clears own history",
currentUserID: regularUserID,
targetUserID: regularUserID,
hardDelete: false,
expectedError: false,
expectedCount: 0,
expectedMsg: "Your chat history has been cleared.",
},
{
name: "Owner clears admin's history",
currentUserID: ownerID,
targetUserID: adminID,
hardDelete: false,
expectedError: false,
expectedCount: 0,
expectedMsg: "Chat history for user @admin (ID: 456) has been cleared.",
},
{
name: "Admin clears regular user's history",
currentUserID: adminID,
targetUserID: regularUserID,
hardDelete: false,
expectedError: false,
expectedCount: 0,
expectedMsg: "Chat history for user @regular (ID: 789) has been cleared.",
},
{
name: "Regular user attempts to clear admin's history",
currentUserID: regularUserID,
targetUserID: adminID,
hardDelete: false,
expectedError: true,
expectedCount: 5,
expectedMsg: "Permission denied. Only admins and owners can clear other users' histories.",
},
{
name: "Admin attempts to clear non-existent user's history",
currentUserID: adminID,
targetUserID: nonExistentUserID,
hardDelete: false,
expectedError: true,
expectedCount: 5,
expectedMsg: "User with ID 999 not found.",
},
{
name: "Owner hard deletes regular user's history",
currentUserID: ownerID,
targetUserID: regularUserID,
hardDelete: true,
expectedError: false,
expectedCount: 0,
expectedMsg: "Chat history for user @regular (ID: 789) has been cleared.",
},
{
name: "Admin clears regular user's history scoped to non-matching chat",
currentUserID: adminID,
targetUserID: regularUserID,
targetChatID: int64(9999),
hardDelete: false,
expectedError: false,
expectedCount: 5,
expectedMsg: "Chat history for user @regular (ID: 789) has been cleared.",
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
if tc.name != "Owner hard deletes regular user's history" {
err = db.Where("user_id = ?", tc.targetUserID).Delete(&Message{}).Error
assert.NoError(t, err)
for i := 0; i < 5; i++ {
message := Message{
BotID: b.botID,
ChatID: chatID,
UserID: tc.targetUserID,
Username: "test",
UserRole: "user",
Text: "Test message",
Timestamp: time.Now(),
IsUser: true,
}
err = db.Create(&message).Error
assert.NoError(t, err)
}
}
var sentMessage string
mockTgClient.SendMessageFunc = func(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error) {
sentMessage = params.Text
return &models.Message{}, nil
}
b.clearChatHistory(context.Background(), chatID, tc.currentUserID, tc.targetUserID, tc.targetChatID, tc.businessConnID, tc.hardDelete)
assert.Equal(t, tc.expectedMsg, sentMessage)
var count int64
if tc.hardDelete {
db.Unscoped().Model(&Message{}).Where("user_id = ? AND chat_id = ?", tc.targetUserID, chatID).Count(&count)
} else {
db.Model(&Message{}).Where("user_id = ? AND chat_id = ?", tc.targetUserID, chatID).Count(&count)
}
assert.Equal(t, tc.expectedCount, count)
})
}
}
func TestStatsCommand(t *testing.T) {
db := setupTestDB(t)
mockClock := &MockClock{
currentTime: time.Now(),
}
config := BotConfig{
ID: "test_bot",
OwnerTelegramID: 123,
TelegramToken: "test_token",
MemorySize: 10,
MessagePerHour: 5,
MessagePerDay: 10,
TempBanDuration: "1h",
SystemPrompts: make(map[string]string),
Active: true,
}
mockTgClient := &MockTelegramClient{}
botModel := &BotModel{
Identifier: config.ID,
Name: config.ID,
}
err := db.Create(botModel).Error
assert.NoError(t, err)
configModel := &ConfigModel{
BotID: botModel.ID,
MemorySize: config.MemorySize,
MessagePerHour: config.MessagePerHour,
MessagePerDay: config.MessagePerDay,
TempBanDuration: config.TempBanDuration,
SystemPrompts: "{}",
TelegramToken: config.TelegramToken,
Active: config.Active,
}
err = db.Create(configModel).Error
assert.NoError(t, err)
b, err := NewBot(db, config, mockClock, mockTgClient)
assert.NoError(t, err)
ownerID := int64(123)
adminID := int64(456)
regularUserID := int64(789)
chatID := int64(1000)
adminRole, err := b.getRoleByName("admin")
assert.NoError(t, err)
adminUser := User{
BotID: b.botID,
TelegramID: adminID,
Username: "admin",
RoleID: adminRole.ID,
Role: adminRole,
IsOwner: false,
}
err = db.Create(&adminUser).Error
assert.NoError(t, err)
regularRole, err := b.getRoleByName("user")
assert.NoError(t, err)
regularUser := User{
BotID: b.botID,
TelegramID: regularUserID,
Username: "regular",
RoleID: regularRole.ID,
Role: regularRole,
IsOwner: false,
}
err = db.Create(&regularUser).Error
assert.NoError(t, err)
for _, userID := range []int64{ownerID, adminID, regularUserID} {
for i := 0; i < 5; i++ {
userMessage := Message{
BotID: b.botID,
ChatID: chatID,
UserID: userID,
Username: "test",
UserRole: "user",
Text: "Test message",
Timestamp: time.Now(),
IsUser: true,
}
err = db.Create(&userMessage).Error
assert.NoError(t, err)
botMessage := Message{
BotID: b.botID,
ChatID: chatID,
UserID: 0,
Username: "AI Assistant",
UserRole: "assistant",
Text: "Test response",
Timestamp: time.Now(),
IsUser: false,
}
err = db.Create(&botMessage).Error
assert.NoError(t, err)
}
}
testCases := []struct {
name string
command string
currentUserID int64
expectedError bool
expectedMsg string
businessConnID string
}{
{
name: "Global stats",
command: "/stats",
currentUserID: regularUserID,
expectedError: false,
expectedMsg: "📊 Bot Statistics:",
},
{
name: "User requests own stats",
command: "/stats user",
currentUserID: regularUserID,
expectedError: false,
expectedMsg: "👤 User Statistics for @regular (ID: 789):",
},
{
name: "Admin requests another user's stats",
command: "/stats user 789",
currentUserID: adminID,
expectedError: false,
expectedMsg: "👤 User Statistics for @regular (ID: 789):",
},
{
name: "Owner requests another user's stats",
command: "/stats user 456",
currentUserID: ownerID,
expectedError: false,
expectedMsg: "👤 User Statistics for @admin (ID: 456):",
},
{
name: "Regular user attempts to request another user's stats",
command: "/stats user 456",
currentUserID: regularUserID,
expectedError: true,
expectedMsg: "Permission denied. Only admins and owners can view other users' statistics.",
},
{
name: "User provides invalid user ID format",
command: "/stats user abc",
currentUserID: adminID,
expectedError: true,
expectedMsg: "Invalid user ID format. Usage: /stats user [user_id]",
},
{
name: "User provides invalid command format",
command: "/stats invalid",
currentUserID: adminID,
expectedError: true,
expectedMsg: "Invalid command format. Usage: /stats or /stats user [user_id]",
},
{
name: "User requests non-existent user's stats",
command: "/stats user 999",
currentUserID: adminID,
expectedError: true,
expectedMsg: "Sorry, I couldn't retrieve statistics for user ID 999.",
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
var sentMessage string
mockTgClient.SendMessageFunc = func(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error) {
sentMessage = params.Text
return &models.Message{}, nil
}
update := &models.Update{
Message: &models.Message{
Chat: models.Chat{ID: chatID},
From: &models.User{
ID: tc.currentUserID,
Username: getUsernameByID(tc.currentUserID),
},
Text: tc.command,
Entities: []models.MessageEntity{
{
Type: "bot_command",
Offset: 0,
Length: 6,
},
},
},
}
b.handleUpdate(context.Background(), nil, update)
assert.Contains(t, sentMessage, tc.expectedMsg)
})
}
}
func getUsernameByID(id int64) string {
switch id {
case 123:
return "owner"
case 456:
return "admin"
case 789:
return "regular"
default:
return "unknown"
}
}
func setupTestDB(t *testing.T) *gorm.DB {
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
if err != nil {
t.Fatalf("Failed to open test database: %v", err)
}
err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{}, &Scope{})
if err != nil {
t.Fatalf("Failed to migrate database schema: %v", err)
}
err = createDefaultRoles(db)
if err != nil {
t.Fatalf("Failed to create default roles: %v", err)
}
if err := createDefaultScopes(db); err != nil {
t.Fatalf("Failed to create default scopes: %v", err)
}
return db
}
func setupBotForTest(t *testing.T, ownerID int64) (*Bot, *MockTelegramClient) {
t.Helper()
db := setupTestDB(t)
mockClock := &MockClock{currentTime: time.Now()}
config := BotConfig{
ID: "test_bot",
OwnerTelegramID: ownerID,
TelegramToken: "test_token",
MemorySize: 10,
MessagePerHour: 5,
MessagePerDay: 10,
TempBanDuration: "1h",
Model: "claude-3-5-haiku-latest",
SystemPrompts: make(map[string]string),
Active: true,
}
mockTgClient := &MockTelegramClient{}
botModel := &BotModel{Identifier: config.ID, Name: config.ID}
assert.NoError(t, db.Create(botModel).Error)
assert.NoError(t, db.Create(&ConfigModel{
BotID: botModel.ID,
MemorySize: config.MemorySize,
MessagePerHour: config.MessagePerHour,
MessagePerDay: config.MessagePerDay,
TempBanDuration: config.TempBanDuration,
SystemPrompts: "{}",
TelegramToken: config.TelegramToken,
Active: config.Active,
}).Error)
b, err := NewBot(db, config, mockClock, mockTgClient)
assert.NoError(t, err)
return b, mockTgClient
}
func TestAnthropicErrorResponse(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
b, _ := setupBotForTest(t, 123)
adminRole, err := b.getRoleByName("admin")
assert.NoError(t, err)
assert.NoError(t, b.db.Create(&User{
BotID: b.botID, TelegramID: 456, Username: "admin",
RoleID: adminRole.ID, Role: adminRole,
}).Error)
userRole, err := b.getRoleByName("user")
assert.NoError(t, err)
assert.NoError(t, b.db.Create(&User{
BotID: b.botID, TelegramID: 789, Username: "regular",
RoleID: userRole.ID, Role: userRole,
}).Error)
modelErr := fmt.Errorf("%w: claude-3-5-haiku-latest", ErrModelNotFound)
otherErr := errors.New("network error")
tests := []struct {
name string
err error
userID int64
wantSubstr string
wantMissing string
}{
{
name: "owner receives actionable model-not-found message",
err: modelErr,
userID: 123,
wantSubstr: "/set_model",
},
{
name: "admin receives actionable model-not-found message",
err: modelErr,
userID: 456,
wantSubstr: "/set_model",
},
{
name: "regular user receives generic message for model-not-found",
err: modelErr,
userID: 789,
wantSubstr: "I'm sorry",
wantMissing: "/set_model",
},
{
name: "owner receives elevated detail for non-API error",
err: otherErr,
userID: 123,
wantSubstr: "Anthropic call failed:",
wantMissing: "I'm sorry",
},
{
name: "regular user receives generic message for non-model error",
err: otherErr,
userID: 789,
wantSubstr: "I'm sorry",
wantMissing: "Anthropic call failed",
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
resp := b.anthropicErrorResponse(tc.err, tc.userID)
assert.Contains(t, resp, tc.wantSubstr)
if tc.wantMissing != "" {
assert.NotContains(t, resp, tc.wantMissing)
}
})
}
}
func TestSetModelCommand(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
b, mockTgClient := setupBotForTest(t, 123)
tempDir, err := os.MkdirTemp("", "set_model_cmd_test")
assert.NoError(t, err)
defer func() { _ = os.RemoveAll(tempDir) }()
configPath := filepath.Join(tempDir, "config.json")
initialJSON := `{"id":"test_bot","telegram_token":"test_token","model":"claude-3-5-haiku-latest","messages_per_hour":5,"messages_per_day":10}`
assert.NoError(t, os.WriteFile(configPath, []byte(initialJSON), 0600))
b.config.ConfigFilePath = configPath
adminRole, err := b.getRoleByName("admin")
assert.NoError(t, err)
assert.NoError(t, b.db.Create(&User{
BotID: b.botID, TelegramID: 456, Username: "admin",
RoleID: adminRole.ID, Role: adminRole,
}).Error)
userRole, err := b.getRoleByName("user")
assert.NoError(t, err)
assert.NoError(t, b.db.Create(&User{
BotID: b.botID, TelegramID: 789, Username: "regular",
RoleID: userRole.ID, Role: userRole,
}).Error)
chatID := int64(1000)
assert.NoError(t, b.db.Create(&Message{
BotID: b.botID, ChatID: chatID, UserID: 789, Username: "regular",
UserRole: "user", Text: "hello", IsUser: true,
}).Error)
makeUpdate := func(userID int64, text string, cmdLen int) *models.Update {
return &models.Update{
Message: &models.Message{
Chat: models.Chat{ID: chatID},
From: &models.User{ID: userID, Username: getUsernameByID(userID)},
Text: text,
Entities: []models.MessageEntity{
{Type: "bot_command", Offset: 0, Length: cmdLen},
},
},
}
}
tests := []struct {
name string
userID int64
text string
wantSubstr string
}{
{
name: "regular user is denied",
userID: 789,
text: "/set_model claude-sonnet-4-6",
wantSubstr: "Permission denied",
},
{
name: "admin missing argument shows usage",
userID: 456,
text: "/set_model",
wantSubstr: "Usage:",
},
{
name: "owner missing argument shows usage",
userID: 123,
text: "/set_model",
wantSubstr: "Usage:",
},
{
name: "admin sets model successfully",
userID: 456,
text: "/set_model claude-sonnet-4-6",
wantSubstr: "✅",
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
var sentMessage string
mockTgClient.SendMessageFunc = func(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error) {
sentMessage = params.Text
return &models.Message{}, nil
}
b.handleUpdate(context.Background(), nil, makeUpdate(tc.userID, tc.text, 10))
assert.Contains(t, sentMessage, tc.wantSubstr)
})
}
t.Run("model change persisted in memory and on disk", func(t *testing.T) {
assert.Equal(t, "claude-sonnet-4-6", string(b.config.Model))
data, err := os.ReadFile(configPath)
assert.NoError(t, err)
assert.Contains(t, string(data), `"claude-sonnet-4-6"`)
})
}
func TestHasScope(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
const ownerID int64 = 100
b, _ := setupBotForTest(t, ownerID)
adminRole, err := b.getRoleByName("admin")
assert.NoError(t, err)
assert.NoError(t, b.db.Create(&User{
BotID: b.botID, TelegramID: 200, Username: "admin_user",
RoleID: adminRole.ID, Role: adminRole,
}).Error)
userRole, err := b.getRoleByName("user")
assert.NoError(t, err)
assert.NoError(t, b.db.Create(&User{
BotID: b.botID, TelegramID: 300, Username: "regular_user",
RoleID: userRole.ID, Role: userRole,
}).Error)
tests := []struct {
name string
userID int64
scope string
want bool
}{
{"owner bypass: model:set", ownerID, ScopeModelSet, true},
{"owner bypass: stats:view:any", ownerID, ScopeStatsViewAny, true},
{"admin: model:set", 200, ScopeModelSet, true},
{"admin: stats:view:any", 200, ScopeStatsViewAny, true},
{"admin: history:clear:any", 200, ScopeHistoryClearAny, true},
{"user: model:set denied", 300, ScopeModelSet, false},
{"user: stats:view:any denied", 300, ScopeStatsViewAny, false},
{"user: history:clear:any denied", 300, ScopeHistoryClearAny, false},
{"user: stats:view:own allowed", 300, ScopeStatsViewOwn, true},
{"user: history:clear:own allowed", 300, ScopeHistoryClearOwn, true},
{"unknown telegram_id", 999, ScopeModelSet, false},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
assert.Equal(t, tc.want, b.hasScope(tc.userID, tc.scope))
})
}
}
+17
View File
@@ -0,0 +1,17 @@
package main
import (
"log"
"os"
)
var (
InfoLogger *log.Logger
ErrorLogger *log.Logger
)
func initLoggers() {
InfoLogger = log.New(os.Stdout, "INFO: ", log.Ldate|log.Ltime|log.Lshortfile)
ErrorLogger = log.New(os.Stderr, "ERROR: ", log.Ldate|log.Ltime|log.Lshortfile)
}
+53
View File
@@ -0,0 +1,53 @@
package main
import (
"context"
"os"
"os/signal"
"sync"
)
func main() {
initLoggers()
InfoLogger.Println("Starting Telegram Bot Application")
db, err := initDB()
if err != nil {
ErrorLogger.Fatalf("Error initializing database: %v", err)
}
configs, err := loadAllConfigs("config")
if err != nil {
ErrorLogger.Fatalf("Error loading configurations: %v", err)
}
var wg sync.WaitGroup
ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
defer cancel()
for _, config := range configs {
wg.Add(1)
go func(cfg BotConfig) {
defer wg.Done()
realClock := RealClock{}
bot, err := NewBot(db, cfg, realClock, nil)
if err != nil {
ErrorLogger.Printf("Error creating bot %s: %v", cfg.ID, err)
return
}
go bot.Start(ctx)
<-ctx.Done()
InfoLogger.Printf("Bot %s stopped", cfg.ID)
}(config)
}
wg.Wait()
InfoLogger.Println("All bots have stopped. Exiting application.")
}
+90
View File
@@ -0,0 +1,90 @@
package main
import (
"time"
"gorm.io/gorm"
)
type BotModel struct {
gorm.Model
Identifier string `gorm:"uniqueIndex"`
Name string
Configs []ConfigModel `gorm:"foreignKey:BotID;constraint:OnDelete:CASCADE"`
Users []User `gorm:"foreignKey:BotID;constraint:OnDelete:CASCADE"`
Messages []Message `gorm:"foreignKey:BotID;constraint:OnDelete:CASCADE"`
}
type ConfigModel struct {
gorm.Model
BotID uint `gorm:"index"`
MemorySize int `json:"memory_size"`
MessagePerHour int `json:"messages_per_hour"`
MessagePerDay int `json:"messages_per_day"`
TempBanDuration string `json:"temp_ban_duration"`
SystemPrompts string `json:"system_prompts"`
TelegramToken string `json:"telegram_token"`
Active bool `json:"active"`
}
type Message struct {
gorm.Model
BotID uint `gorm:"index"`
ChatID int64 `gorm:"index"`
UserID int64 `gorm:"index"`
Username string `gorm:"index"`
UserRole string
Text string `gorm:"type:text"`
Timestamp time.Time `gorm:"index"`
IsUser bool
StickerFileID string
StickerPNGFile string
StickerEmoji string
DeletedAt gorm.DeletedAt `gorm:"index"`
AnsweredOn *time.Time `gorm:"index"`
ImageFileIDs []string `gorm:"type:text;serializer:json"`
FilesCleanedAt *time.Time `gorm:"index"`
}
type ChatMemory struct {
Messages []Message
Size int
BusinessConnectionID string
}
const (
ScopeStatsViewOwn = "stats:view:own"
ScopeStatsViewAny = "stats:view:any"
ScopeHistoryClearOwn = "history:clear:own"
ScopeHistoryClearAny = "history:clear:any"
ScopeHistoryClearHardOwn = "history:clear_hard:own"
ScopeHistoryClearHardAny = "history:clear_hard:any"
ScopeModelSet = "model:set"
ScopeUserPromote = "user:promote"
ScopeTTSUse = "tts:use"
)
type Scope struct {
gorm.Model
Name string `gorm:"uniqueIndex"`
}
type Role struct {
gorm.Model
Name string `gorm:"uniqueIndex"`
Scopes []Scope `gorm:"many2many:role_scopes;"`
}
type User struct {
gorm.Model
BotID uint `gorm:"uniqueIndex:idx_user_bot;index"`
TelegramID int64 `gorm:"uniqueIndex:idx_user_bot;not null"`
Username string
RoleID uint
Role Role `gorm:"foreignKey:RoleID"`
IsOwner bool `gorm:"default:false"`
}
func (User) TableName() string {
return "users"
}
+64
View File
@@ -0,0 +1,64 @@
package main
import (
"time"
"golang.org/x/time/rate"
)
type userLimiter struct {
hourlyLimiter *rate.Limiter
dailyLimiter *rate.Limiter
lastHourlyReset time.Time
lastDailyReset time.Time
banUntil time.Time
clock Clock
}
func (b *Bot) checkRateLimits(userID int64) bool {
b.userLimitersMu.Lock()
defer b.userLimitersMu.Unlock()
limiter, exists := b.userLimiters[userID]
if !exists {
limiter = &userLimiter{
hourlyLimiter: rate.NewLimiter(rate.Every(time.Hour/time.Duration(b.config.MessagePerHour)), b.config.MessagePerHour),
dailyLimiter: rate.NewLimiter(rate.Every(24*time.Hour/time.Duration(b.config.MessagePerDay)), b.config.MessagePerDay),
lastHourlyReset: b.clock.Now(),
lastDailyReset: b.clock.Now(),
clock: b.clock,
}
b.userLimiters[userID] = limiter
}
now := limiter.clock.Now()
if now.Before(limiter.banUntil) {
return false
}
if now.Sub(limiter.lastHourlyReset) >= time.Hour {
limiter.hourlyLimiter = rate.NewLimiter(rate.Every(time.Hour/time.Duration(b.config.MessagePerHour)), b.config.MessagePerHour)
limiter.lastHourlyReset = now
}
if now.Sub(limiter.lastDailyReset) >= 24*time.Hour {
limiter.dailyLimiter = rate.NewLimiter(rate.Every(24*time.Hour/time.Duration(b.config.MessagePerDay)), b.config.MessagePerDay)
limiter.lastDailyReset = now
}
dailyRes := limiter.dailyLimiter.ReserveN(now, 1)
hourlyRes := limiter.hourlyLimiter.ReserveN(now, 1)
if dailyRes.DelayFrom(now) > 0 || hourlyRes.DelayFrom(now) > 0 {
dailyRes.CancelAt(now)
hourlyRes.CancelAt(now)
banDuration, err := time.ParseDuration(b.config.TempBanDuration)
if err != nil {
banDuration = 24 * time.Hour
}
limiter.banUntil = now.Add(banDuration)
return false
}
return true
}
+67
View File
@@ -0,0 +1,67 @@
package main
import (
"testing"
"time"
)
func TestCheckRateLimits(t *testing.T) {
mockClock := &MockClock{
currentTime: time.Date(2023, 10, 1, 0, 0, 0, 0, time.UTC),
}
config := BotConfig{
ID: "bot1",
MemorySize: 10,
MessagePerHour: 5,
MessagePerDay: 10,
TempBanDuration: "1m",
SystemPrompts: make(map[string]string),
TelegramToken: "YOUR_TELEGRAM_BOT_TOKEN",
OwnerTelegramID: 123456789,
}
bot := &Bot{
config: config,
userLimiters: make(map[int64]*userLimiter),
clock: mockClock,
}
userID := int64(12345)
sendMessage := func() bool {
return bot.checkRateLimits(userID)
}
for i := 0; i < config.MessagePerHour; i++ {
if !sendMessage() {
t.Errorf("Expected message %d to be allowed", i+1)
}
}
if sendMessage() {
t.Errorf("Expected message to be denied due to hourly limit exceeded")
}
if sendMessage() {
t.Errorf("Expected message to be denied while user is banned")
}
mockClock.Advance(time.Minute)
mockClock.Advance(time.Hour)
if !sendMessage() {
t.Errorf("Expected message to be allowed after ban duration")
}
for i := 0; i < config.MessagePerDay-config.MessagePerHour-1; i++ {
if !sendMessage() {
t.Errorf("Expected message %d to be allowed towards daily limit", i+1)
}
}
if sendMessage() {
t.Errorf("Expected message to be denied due to daily limit exceeded")
}
}
+17
View File
@@ -0,0 +1,17 @@
package main
import (
"context"
"github.com/go-telegram/bot"
"github.com/go-telegram/bot/models"
)
type TelegramClient interface {
SendMessage(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error)
SendAudio(ctx context.Context, params *bot.SendAudioParams) (*models.Message, error)
SetMyCommands(ctx context.Context, params *bot.SetMyCommandsParams) (bool, error)
GetFile(ctx context.Context, params *bot.GetFileParams) (*models.File, error)
FileDownloadLink(f *models.File) string
Start(ctx context.Context)
}
+66
View File
@@ -0,0 +1,66 @@
package main
import (
"context"
"github.com/go-telegram/bot"
"github.com/go-telegram/bot/models"
"github.com/stretchr/testify/mock"
)
type MockTelegramClient struct {
mock.Mock
SendMessageFunc func(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error)
SendAudioFunc func(ctx context.Context, params *bot.SendAudioParams) (*models.Message, error)
SetMyCommandsFunc func(ctx context.Context, params *bot.SetMyCommandsParams) (bool, error)
GetFileFunc func(ctx context.Context, params *bot.GetFileParams) (*models.File, error)
FileDownloadLinkFunc func(f *models.File) string
StartFunc func(ctx context.Context)
}
func (m *MockTelegramClient) SendMessage(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error) {
if m.SendMessageFunc != nil {
return m.SendMessageFunc(ctx, params)
}
args := m.Called(ctx, params)
if msg, ok := args.Get(0).(*models.Message); ok {
return msg, args.Error(1)
}
return nil, args.Error(1)
}
func (m *MockTelegramClient) SetMyCommands(ctx context.Context, params *bot.SetMyCommandsParams) (bool, error) {
if m.SetMyCommandsFunc != nil {
return m.SetMyCommandsFunc(ctx, params)
}
return true, nil
}
func (m *MockTelegramClient) SendAudio(ctx context.Context, params *bot.SendAudioParams) (*models.Message, error) {
if m.SendAudioFunc != nil {
return m.SendAudioFunc(ctx, params)
}
return nil, nil
}
func (m *MockTelegramClient) GetFile(ctx context.Context, params *bot.GetFileParams) (*models.File, error) {
if m.GetFileFunc != nil {
return m.GetFileFunc(ctx, params)
}
return &models.File{}, nil
}
func (m *MockTelegramClient) FileDownloadLink(f *models.File) string {
if m.FileDownloadLinkFunc != nil {
return m.FileDownloadLinkFunc(f)
}
return ""
}
func (m *MockTelegramClient) Start(ctx context.Context) {
if m.StartFunc != nil {
m.StartFunc(ctx)
return
}
m.Called(ctx)
}
+52
View File
@@ -0,0 +1,52 @@
package main
import (
"context"
"fmt"
"io"
"net/http"
tgbot "github.com/go-telegram/bot"
"github.com/go-telegram/bot/models"
)
func largestPhotoSize(photos []models.PhotoSize) models.PhotoSize {
if len(photos) == 0 {
return models.PhotoSize{}
}
largest := photos[0]
largestArea := largest.Width * largest.Height
for i := 1; i < len(photos); i++ {
area := photos[i].Width * photos[i].Height
if area > largestArea {
largest = photos[i]
largestArea = area
}
}
return largest
}
func (b *Bot) downloadTelegramFile(ctx context.Context, fileID string) ([]byte, error) {
fileInfo, err := b.tgBot.GetFile(ctx, &tgbot.GetFileParams{FileID: fileID})
if err != nil {
return nil, fmt.Errorf("telegram GetFile %s: %w", fileID, err)
}
downloadURL := b.tgBot.FileDownloadLink(fileInfo)
req, err := http.NewRequestWithContext(ctx, http.MethodGet, downloadURL, nil)
if err != nil {
return nil, fmt.Errorf("telegram download request %s: %w", fileID, err)
}
resp, err := http.DefaultClient.Do(req)
if err != nil {
return nil, fmt.Errorf("telegram download %s: %w", fileID, err)
}
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("telegram download %s: status %d", fileID, resp.StatusCode)
}
data, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("telegram download read %s: %w", fileID, err)
}
return data, nil
}
+53
View File
@@ -0,0 +1,53 @@
package main
import (
"testing"
"github.com/go-telegram/bot/models"
"github.com/stretchr/testify/assert"
)
func TestLargestPhotoSize(t *testing.T) {
cases := []struct {
name string
photos []models.PhotoSize
wantFileID string
}{
{
name: "ascending sizes — last is largest",
photos: []models.PhotoSize{
{FileID: "thumb", Width: 90, Height: 90},
{FileID: "small", Width: 320, Height: 320},
{FileID: "full", Width: 1280, Height: 720},
},
wantFileID: "full",
},
{
name: "descending sizes — first is largest",
photos: []models.PhotoSize{
{FileID: "full", Width: 1280, Height: 720},
{FileID: "small", Width: 320, Height: 320},
{FileID: "thumb", Width: 90, Height: 90},
},
wantFileID: "full",
},
{
name: "single photo",
photos: []models.PhotoSize{
{FileID: "solo", Width: 800, Height: 600},
},
wantFileID: "solo",
},
{
name: "empty slice returns zero value (caller guards upstream)",
photos: []models.PhotoSize{},
wantFileID: "",
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
got := largestPhotoSize(tc.photos)
assert.Equal(t, tc.wantFileID, got.FileID)
})
}
}
+260
View File
@@ -0,0 +1,260 @@
package main
import (
"context"
"fmt"
"testing"
"time"
"github.com/go-telegram/bot"
"github.com/go-telegram/bot/models"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
)
const (
errOpenDB = "Failed to open in-memory database: %v"
errMigrateSchema = "Failed to migrate database schema: %v"
errCreateRoles = "Failed to create default roles: %v"
errCreateScopes = "Failed to create default scopes: %v"
errCreateBot = "Failed to create bot: %v"
memoryDSN = ":memory:"
)
func TestOwnerAssignment(t *testing.T) {
initLoggers()
db, err := gorm.Open(sqlite.Open(memoryDSN), &gorm.Config{})
if err != nil {
t.Fatalf(errOpenDB, err)
}
err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{}, &Scope{})
if err != nil {
t.Fatalf(errMigrateSchema, err)
}
err = createDefaultRoles(db)
if err != nil {
t.Fatalf(errCreateRoles, err)
}
if err := createDefaultScopes(db); err != nil {
t.Fatalf(errCreateScopes, err)
}
config := BotConfig{
ID: "test_bot",
TelegramToken: "TEST_TELEGRAM_TOKEN",
MemorySize: 10,
MessagePerHour: 5,
MessagePerDay: 10,
TempBanDuration: "1m",
SystemPrompts: make(map[string]string),
Active: true,
OwnerTelegramID: 111111111,
}
mockClock := &MockClock{
currentTime: time.Now(),
}
mockTGClient := &MockTelegramClient{
SendMessageFunc: func(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error) {
chatID, ok := params.ChatID.(int64)
if !ok {
return nil, fmt.Errorf("ChatID is not of type int64")
}
return &models.Message{ID: 1, Chat: models.Chat{ID: chatID}}, nil
},
}
bot, err := NewBot(db, config, mockClock, mockTGClient)
if err != nil {
t.Fatalf(errCreateBot, err)
}
var owner User
err = db.Where("telegram_id = ? AND bot_id = ? AND is_owner = ?", config.OwnerTelegramID, bot.botID, true).First(&owner).Error
if err != nil {
t.Fatalf("Owner was not created: %v", err)
}
_, err = bot.getOrCreateUser(222222222, "AnotherOwner", true)
if err == nil {
t.Fatalf("Expected error when creating a second owner, but got none")
}
expectedErrorMsg := "an owner already exists for this bot"
if err.Error() != expectedErrorMsg {
t.Fatalf("Unexpected error message: %v", err)
}
regularUser, err := bot.getOrCreateUser(333333333, "RegularUser", false)
if err != nil {
t.Fatalf("Failed to create regular user: %v", err)
}
if regularUser.Role.Name != "user" {
t.Fatalf("Expected role 'user', got '%s'", regularUser.Role.Name)
}
_, err = bot.getOrCreateUser(333333333, "AdminUser", true)
if err == nil {
t.Fatalf("Expected error when changing existing user to owner, but got none")
}
expectedErrorMsg = "cannot change existing user to owner"
if err.Error() != expectedErrorMsg {
t.Fatalf("Unexpected error message: %v", err)
}
}
func TestPromoteUserToAdmin(t *testing.T) {
initLoggers()
db, err := gorm.Open(sqlite.Open(memoryDSN), &gorm.Config{})
if err != nil {
t.Fatalf(errOpenDB, err)
}
err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{}, &Scope{})
if err != nil {
t.Fatalf(errMigrateSchema, err)
}
err = createDefaultRoles(db)
if err != nil {
t.Fatalf(errCreateRoles, err)
}
if err := createDefaultScopes(db); err != nil {
t.Fatalf(errCreateScopes, err)
}
config := BotConfig{
ID: "test_bot",
TelegramToken: "TEST_TELEGRAM_TOKEN",
MemorySize: 10,
MessagePerHour: 5,
MessagePerDay: 10,
TempBanDuration: "1m",
SystemPrompts: make(map[string]string),
Active: true,
OwnerTelegramID: 111111111,
}
mockClock := &MockClock{currentTime: time.Now()}
mockTGClient := &MockTelegramClient{}
bot, err := NewBot(db, config, mockClock, mockTGClient)
if err != nil {
t.Fatalf(errCreateBot, err)
}
owner, err := bot.getOrCreateUser(config.OwnerTelegramID, "OwnerUser", true)
if err != nil {
t.Fatalf("Failed to create owner: %v", err)
}
regularUser, err := bot.getOrCreateUser(444444444, "RegularUser", false)
if err != nil {
t.Fatalf("Failed to create regular user: %v", err)
}
err = bot.promoteUserToAdmin(owner.TelegramID, regularUser.TelegramID)
if err != nil {
t.Fatalf("Failed to promote user to admin: %v", err)
}
promotedUser, err := bot.getOrCreateUser(444444444, "RegularUser", false)
if err != nil {
t.Fatalf("Failed to get promoted user: %v", err)
}
if promotedUser.Role.Name != "admin" {
t.Fatalf("Expected role 'admin', got '%s'", promotedUser.Role.Name)
}
}
func TestGetOrCreateUser(t *testing.T) {
initLoggers()
db, err := gorm.Open(sqlite.Open(memoryDSN), &gorm.Config{})
if err != nil {
t.Fatalf(errOpenDB, err)
}
err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{}, &Scope{})
if err != nil {
t.Fatalf(errMigrateSchema, err)
}
err = createDefaultRoles(db)
if err != nil {
t.Fatalf(errCreateRoles, err)
}
if err := createDefaultScopes(db); err != nil {
t.Fatalf(errCreateScopes, err)
}
mockClock := &MockClock{
currentTime: time.Date(2023, 10, 1, 0, 0, 0, 0, time.UTC),
}
config := BotConfig{
ID: "bot1",
MemorySize: 10,
MessagePerHour: 5,
MessagePerDay: 10,
TempBanDuration: "1m",
SystemPrompts: make(map[string]string),
TelegramToken: "YOUR_TELEGRAM_BOT_TOKEN",
OwnerTelegramID: 123456789,
}
mockTGClient := &MockTelegramClient{
SendMessageFunc: func(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error) {
chatID, ok := params.ChatID.(int64)
if !ok {
return nil, fmt.Errorf("ChatID is not of type int64")
}
return &models.Message{ID: 1, Chat: models.Chat{ID: chatID}}, nil
},
}
bot, err := NewBot(db, config, mockClock, mockTGClient)
if err != nil {
t.Fatalf(errCreateBot, err)
}
var owner User
err = db.Where("telegram_id = ? AND bot_id = ? AND is_owner = ?", config.OwnerTelegramID, bot.botID, true).First(&owner).Error
if err != nil {
t.Fatalf("Owner was not created: %v", err)
}
_, err = bot.getOrCreateUser(222222222, "AnotherOwner", true)
if err == nil {
t.Fatalf("Expected error when creating a second owner, but got none")
}
newUser, err := bot.getOrCreateUser(987654321, "TestUser", false)
if err != nil {
t.Fatalf("Failed to create a new user: %v", err)
}
var userInDB User
err = db.Where("telegram_id = ?", newUser.TelegramID).First(&userInDB).Error
if err != nil {
t.Fatalf("New user was not created in the database: %v", err)
}
existingUser, err := bot.getOrCreateUser(987654321, "TestUser", false)
if err != nil {
t.Fatalf("Failed to get existing user: %v", err)
}
if existingUser.ID != userInDB.ID {
t.Fatalf("Expected to get the existing user, but got a different user")
}
}