diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..7deffff --- /dev/null +++ b/.dockerignore @@ -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/ \ No newline at end of file diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..aafbd82 --- /dev/null +++ b/.gitattributes @@ -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 diff --git a/.github/workflows/go-ci.yaml b/.github/workflows/go-ci.yaml new file mode 100644 index 0000000..3311f2f --- /dev/null +++ b/.github/workflows/go-ci.yaml @@ -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/gosec@v2.26.1 + with: + args: ./... diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..3c26de9 --- /dev/null +++ b/.gitignore @@ -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 diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..a943c43 --- /dev/null +++ b/Dockerfile @@ -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"] \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..14a03e5 --- /dev/null +++ b/README.md @@ -0,0 +1,183 @@ +# 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. | + +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. + +> [!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 ` | 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 ` | Admin/Owner | Soft-delete all messages for a user across every chat | +| `/clear ` | Admin/Owner | Soft-delete a user's messages in a specific chat | +| `/clear_hard` | All users | Permanently delete your own chat history | +| `/clear_hard ` | Admin/Owner | Permanently delete all messages for a user across every chat | +| `/clear_hard ` | Admin/Owner | Permanently delete a user's messages in a specific chat | +| `/set_model ` | Admin/Owner | Switch the AI model live without restarting | + +> **Note:** In private DMs each user's `chat_id` equals their `user_id`. The scoped `` 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. diff --git a/album_buffer.go b/album_buffer.go new file mode 100644 index 0000000..d3bb17a --- /dev/null +++ b/album_buffer.go @@ -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, + ) +} diff --git a/anthropic.go b/anthropic.go new file mode 100644 index 0000000..f6fa8da --- /dev/null +++ b/anthropic.go @@ -0,0 +1,336 @@ +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 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\n" + rule + "\n" + } + } + 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 + } + + if len(b.config.MCPServers) > 0 { + mcpServers := make([]anthropic.BetaRequestMCPServerURLDefinitionParam, 0, len(b.config.MCPServers)) + tools := make([]anthropic.BetaToolUnionParam, 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.Tools = tools + params.Betas = append(params.Betas, anthropic.AnthropicBetaMCPClient2025_11_20) + } + + for attempt := 0; attempt < maxFileNotFoundRetries; attempt++ { + joined, streamErr := b.streamMessages(ctx, params, onSegment) + if streamErr == nil { + return joined, 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) + } + ErrorLogger.Printf("[%s] self-heal: stripping dead file_id %s from chat %d (attempt %d/%d)", + b.config.ID, missingFileID, chatID, attempt+1, 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)) + } + return "", fmt.Errorf("max self-heal retries (%d) exceeded: too many file_ids gone from anthropic", maxFileNotFoundRetries) +} + +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, 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 ( + allSegments []string + currentKind string + currentText strings.Builder + currentThinking strings.Builder + currentInputJSON strings.Builder + currentTUseName, currentTUseServer, currentTUseID string + currentTResultUseID, currentTResultServer string + currentTResultIsError bool + currentTResultContent string + mcpCalls = map[string]mcpCall{} + startInputTokens int64 + finalUsage anthropic.BetaMessageDeltaUsage + stopReason string + ) + + for stream.Next() { + e := stream.Current() + switch e.Type { + case "message_start": + startInputTokens = e.AsMessageStart().Message.Usage.InputTokens + + case "message_delta": + md := e.AsMessageDelta() + finalUsage = md.Usage + if md.Delta.StopReason != "" { + stopReason = string(md.Delta.StopReason) + } + + case "content_block_start": + cbs := e.AsContentBlockStart() + currentKind = cbs.ContentBlock.Type + currentText.Reset() + currentThinking.Reset() + currentInputJSON.Reset() + 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 "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" { + 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 "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("[mcp] block type=%q (unhandled)", currentKind) + } + } + currentKind = "" + } + } + + if err := stream.Err(); err != nil { + return "", err + } + + if stopReason != "" || finalUsage.OutputTokens > 0 { + inputTokens := finalUsage.InputTokens + if inputTokens == 0 { + inputTokens = startInputTokens + } + InfoLogger.Printf("[usage] model=%s in=%d out=%d thinking=%d stop=%s", + params.Model, inputTokens, finalUsage.OutputTokens, + finalUsage.OutputTokensDetails.ThinkingTokens, stopReason) + if stopReason == "max_tokens" { + ErrorLogger.Printf("[usage] response truncated at max_tokens=%d - raise max_tokens (thinking counts toward it)", + params.MaxTokens) + } + } + + if len(allSegments) == 0 { + return "", emptyStreamError(stopReason, finalUsage.OutputTokensDetails.ThinkingTokens, params.MaxTokens) + } + return strings.Join(allSegments, "\n\n"), 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") +} diff --git a/anthropic_files.go b/anthropic_files.go new file mode 100644 index 0000000..e37574f --- /dev/null +++ b/anthropic_files.go @@ -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 +} diff --git a/anthropic_files_test.go b/anthropic_files_test.go new file mode 100644 index 0000000..d93307c --- /dev/null +++ b/anthropic_files_test.go @@ -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) +} diff --git a/anthropic_test.go b/anthropic_test.go new file mode 100644 index 0000000..575fcae --- /dev/null +++ b/anthropic_test.go @@ -0,0 +1,143 @@ +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 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) + } +} diff --git a/bot.go b/bot.go new file mode 100644 index 0000000..5125d6d --- /dev/null +++ b/bot.go @@ -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 "}, +} + +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 +} diff --git a/bot_test.go b/bot_test.go new file mode 100644 index 0000000..bd45901 --- /dev/null +++ b/bot_test.go @@ -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") +} diff --git a/clock.go b/clock.go new file mode 100644 index 0000000..e4d21e8 --- /dev/null +++ b/clock.go @@ -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) +} diff --git a/config.go b/config.go new file mode 100644 index 0000000..7e04315 --- /dev/null +++ b/config.go @@ -0,0 +1,255 @@ +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"` +} + +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"` + 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) + } + + 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 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 +} diff --git a/config/default.json b/config/default.json new file mode 100644 index 0000000..d3d5e71 --- /dev/null +++ b/config/default.json @@ -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 sales@example.com.\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." + } +} diff --git a/config_test.go b/config_test.go new file mode 100644 index 0000000..6274927 --- /dev/null +++ b/config_test.go @@ -0,0 +1,877 @@ +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) + } +} diff --git a/database.go b/database.go new file mode 100644 index 0000000..71f938d --- /dev/null +++ b/database.go @@ -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 +} diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..9d0d689 --- /dev/null +++ b/docker-compose.yml @@ -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" diff --git a/elevenlabs.go b/elevenlabs.go new file mode 100644 index 0000000..705857f --- /dev/null +++ b/elevenlabs.go @@ -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 +} diff --git a/examples/systemd/telegram-bot.service b/examples/systemd/telegram-bot.service new file mode 100644 index 0000000..fa27d2b --- /dev/null +++ b/examples/systemd/telegram-bot.service @@ -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 \ No newline at end of file diff --git a/go-telegram-bot.exe b/go-telegram-bot.exe new file mode 100644 index 0000000..161f5d4 Binary files /dev/null and b/go-telegram-bot.exe differ diff --git a/go.mod b/go.mod index ee51142..7e02cbe 100644 --- a/go.mod +++ b/go.mod @@ -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 +) diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..ad5b3ed --- /dev/null +++ b/go.sum @@ -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= diff --git a/handlers.go b/handlers.go new file mode 100644 index 0000000..8ca1e6b --- /dev/null +++ b/handlers.go @@ -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 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 ", 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) + } +} diff --git a/handlers_test.go b/handlers_test.go new file mode 100644 index 0000000..3b94579 --- /dev/null +++ b/handlers_test.go @@ -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(®ularUser).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(®ularUser).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)) + }) + } +} diff --git a/logger.go b/logger.go new file mode 100644 index 0000000..f17eb44 --- /dev/null +++ b/logger.go @@ -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) +} diff --git a/main.go b/main.go new file mode 100644 index 0000000..11fe444 --- /dev/null +++ b/main.go @@ -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.") +} diff --git a/models.go b/models.go new file mode 100644 index 0000000..0e7a14a --- /dev/null +++ b/models.go @@ -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" +} diff --git a/rate_limiter.go b/rate_limiter.go new file mode 100644 index 0000000..7430a7e --- /dev/null +++ b/rate_limiter.go @@ -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 +} diff --git a/rate_limiter_test.go b/rate_limiter_test.go new file mode 100644 index 0000000..6543ebd --- /dev/null +++ b/rate_limiter_test.go @@ -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") + } +} diff --git a/telegram_client.go b/telegram_client.go new file mode 100644 index 0000000..7ee1791 --- /dev/null +++ b/telegram_client.go @@ -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) +} diff --git a/telegram_client_mock.go b/telegram_client_mock.go new file mode 100644 index 0000000..5642954 --- /dev/null +++ b/telegram_client_mock.go @@ -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) +} diff --git a/telegram_helpers.go b/telegram_helpers.go new file mode 100644 index 0000000..3bcf78d --- /dev/null +++ b/telegram_helpers.go @@ -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 +} diff --git a/telegram_helpers_test.go b/telegram_helpers_test.go new file mode 100644 index 0000000..b8b4ee2 --- /dev/null +++ b/telegram_helpers_test.go @@ -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) + }) + } +} diff --git a/user_management_test.go b/user_management_test.go new file mode 100644 index 0000000..389d2b2 --- /dev/null +++ b/user_management_test.go @@ -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") + } +}