mirror of
https://github.com/HugeFrog24/go-telegram-bot.git
synced 2026-10-28 05:45:26 +00:00
Compare commits
9
Commits
main
..
e1a9261699
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e1a9261699 | ||
|
|
6e2d2fce2f | ||
|
|
37d6242c06 | ||
|
|
d8d0da4704 | ||
|
|
c8af457af1 | ||
|
|
e5532df7f9 | ||
|
|
0ab56448c7 | ||
|
|
9f2b3df4c8 | ||
|
|
41c9b8075b |
@@ -0,0 +1,14 @@
|
||||
---
|
||||
alwaysApply: true
|
||||
always_on: true
|
||||
trigger: always_on
|
||||
applyTo: "**"
|
||||
description: Snyk Security At Inception
|
||||
---
|
||||
|
||||
# Project security best practices
|
||||
|
||||
- Always run snyk_code_scan tool for new first party code that is generated in a Snyk-supported language.
|
||||
- If any security issues are found based on newly introduced or modified code or dependencies, attempt to fix the issues using the results context from Snyk.
|
||||
- Rescan the code after fixing the issues to ensure that the issues were fixed and that there are no newly introduced issues.
|
||||
- Repeat this process until no new issues are found.
|
||||
@@ -6,44 +6,53 @@ on:
|
||||
pull_request:
|
||||
branches: [ main ]
|
||||
|
||||
env:
|
||||
GO_VERSION: '1.27.3'
|
||||
|
||||
jobs:
|
||||
# Lint job
|
||||
lint:
|
||||
# Common setup job that other jobs can depend on
|
||||
setup:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v7
|
||||
- uses: actions/setup-go@v7
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version: ${{ env.GO_VERSION }}
|
||||
go-version: '1.26.0'
|
||||
- uses: actions/cache@v4
|
||||
with:
|
||||
path: |
|
||||
~/.cache/go-build
|
||||
~/go/pkg/mod
|
||||
key: ${{ runner.os }}-go-${{ hashFiles('**/go.sum') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-go-
|
||||
- run: go mod tidy
|
||||
|
||||
# Lint job
|
||||
lint:
|
||||
needs: setup
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: golangci/golangci-lint-action@v9
|
||||
with:
|
||||
version: v2.15.0
|
||||
version: v2.10
|
||||
args: --timeout 5m
|
||||
|
||||
# Test job
|
||||
test:
|
||||
needs: setup
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v7
|
||||
- uses: actions/setup-go@v7
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/setup-go@v6
|
||||
with:
|
||||
go-version: ${{ env.GO_VERSION }}
|
||||
go-version: '1.26.0'
|
||||
- run: go test ./... -v
|
||||
|
||||
# Security scan job
|
||||
security:
|
||||
needs: setup
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v7
|
||||
- uses: actions/setup-go@v7
|
||||
- uses: actions/checkout@v6
|
||||
- uses: securego/gosec@master
|
||||
with:
|
||||
go-version: ${{ env.GO_VERSION }}
|
||||
# gosec is installed rather than run as its container action: the
|
||||
# published image ships its own Go, which lags behind and fails to load
|
||||
# the module ("go.mod requires go >= X"). Installing it here builds gosec
|
||||
# with the toolchain above, so CI's Go version is ours to choose.
|
||||
- run: go install github.com/securego/gosec/v2/cmd/[email protected]
|
||||
- run: gosec ./...
|
||||
args: ./...
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
{
|
||||
"mcpServers": {}
|
||||
}
|
||||
+2
-2
@@ -1,6 +1,6 @@
|
||||
# Multi-stage build for Go Telegram Bot
|
||||
# Build stage
|
||||
FROM golang:1.27-alpine AS builder
|
||||
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
|
||||
@@ -18,7 +18,7 @@ RUN go mod download
|
||||
COPY . .
|
||||
|
||||
# Build the application
|
||||
RUN CGO_ENABLED=1 GOOS=linux go build -o telegram-bot .
|
||||
RUN CGO_ENABLED=1 GOOS=linux go build -a -installsuffix cgo -o telegram-bot .
|
||||
|
||||
# Runtime stage
|
||||
FROM alpine:latest
|
||||
|
||||
@@ -4,8 +4,7 @@ A scalable, multi-bot solution for Telegram using Go, GORM, and the Anthropic AP
|
||||
|
||||
## Design Considerations
|
||||
|
||||
- AI-powered (Anthropic Claude)
|
||||
- Voice message support (ElevenLabs STT + TTS) — optional, enabled per bot via config
|
||||
- AI-powered
|
||||
- Supports multiple bot profiles
|
||||
- Uses SQLite for persistence
|
||||
- Implements rate limiting and user management
|
||||
@@ -40,157 +39,18 @@ A scalable, multi-bot solution for Telegram using Go, GORM, and the Anthropic AP
|
||||
|
||||
### Native Deployment
|
||||
|
||||
1. Clone the repository as above. Building needs Go 1.27 or later and a C compiler, because the SQLite driver uses cgo.
|
||||
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. Claude 5-generation models (Haiku 5.5, Sonnet 5 and later, Opus 5 and later, Fable) think even when `thinking` is omitted; on those, keep `effort` low or raise this limit. |
|
||||
| `temperature` | number | *(ignored)* | No longer sent: since API version `2026-10-21` the Claude API rejects `temperature`, `top_p` and `top_k` on every model. Old configs still load, and the bot logs a startup notice; remove the key. |
|
||||
| `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. |
|
||||
| `effort` | string | *(omitted)* | `"low"`, `"medium"`, `"high"`, `"xhigh"` or `"max"`. Caps how much reasoning and token spend a turn may use; the API default is `high` on most models and `medium` on Haiku 5.5 and Opus 5.5. For short-reply chat bots `"low"` or `"medium"` cuts cost noticeably. Rejected with a 400 on models older than Claude 4.6 (Haiku 4.5, Sonnet 4.5) — omit it there. |
|
||||
| `debounce_ms` | number | *(omitted)* | Quiet window, in milliseconds, for coalescing rapid follow-up messages into a single turn. Omit or set `0` to disable. See [Coalescing rapid messages](#coalescing-rapid-messages) below. |
|
||||
| `stream_drafts` | bool | `false` | Streams each reply into one Telegram draft that grows in place, with a native Stop button, and sends it as a single message when done. Private chats only. See [Draft streaming and the Stop button](#draft-streaming-and-the-stop-button) below. |
|
||||
| `cache_history` | bool | `true` | Places a prompt-cache breakpoint on the trailing conversation block, so each turn reads the prior history from cache instead of reprocessing it at full price. Set `false` to cache only the system prompt. |
|
||||
| `web_search` | object | *(omitted)* | Enables Anthropic's server-side web search (and, optionally, web fetch), sandboxed to a domain allowlist. Omit the key entirely to leave both tools off — an absent block sends byte-identical requests to before. See [Web search / fetch](#web-search--fetch) below. |
|
||||
|
||||
Every reply logs one accounting line — `[usage] model=... in=... out=... thinking=... cache_read=... cache_write=... 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]
|
||||
> Watch `cache_read`/`cache_write` after changing prompts or models. A cache breakpoint below the model's minimum cacheable prefix fails **silently** — no error, just `cache_write=0` forever. The minimum is model-specific and not monotonic across generations (Haiku 4.5 needs 4096 tokens; Sonnet 4.6 needs 1024), so a short system prompt that caches fine on one model may never cache on another. Note also that chat memory is a sliding window: once it is full, each turn evicts the oldest message and changes the prefix, so `cache_read` on long-running chats will be lower than on fresh ones.
|
||||
|
||||
### Coalescing rapid messages
|
||||
|
||||
A user who sends "how do I do X", then "sorry typo", then "lmao" in five seconds would otherwise get three separate replies — the bot starts a full turn per message, because each Telegram update independently drives one. `debounce_ms` holds text messages in a per-chat buffer and resets the window on every new message, dispatching a single turn once the user stops typing.
|
||||
|
||||
```json
|
||||
"debounce_ms": 2500
|
||||
```
|
||||
|
||||
Reasonable windows are 1500–3000ms for ordinary chat and up to 8000ms for Telegram Business, where a person writing to a business account tends to send longer bursts. The maximum is 30000ms.
|
||||
|
||||
Nothing is discarded while buffering. Each message is still persisted and added to chat memory the moment it arrives, so the single coalesced turn sees all of them — the buffer only decides *when* to answer, never *what* the model reads. Reply metadata (language, premium status, business connection) follows the most recent message in the batch.
|
||||
|
||||
What deliberately does **not** wait:
|
||||
|
||||
- **Commands** (`/stats`, `/clear`, …) dispatch immediately.
|
||||
- **Photos, albums, voice, and stickers** dispatch immediately and cancel any pending text window. The buffered text is not lost — it is already in memory, so the media turn answers it too.
|
||||
- **`/clear` and `/clear_hard`** cancel the buffer outright. Without this, the window would fire seconds after the wipe and replay the just-deleted messages back into memory.
|
||||
|
||||
**Edits count as typing.** Mobile clients deliver autocorrect fixes as `edited_message` updates a few seconds after the send. The stored text (or photo caption) is replaced with the edit. If the message is still inside its quiet window, the window restarts. If a turn is already reading the old wording and has shown nothing yet, that turn restarts with the correction. An edit to a message that was already answered only corrects the record.
|
||||
|
||||
#### Messages that arrive mid-turn
|
||||
|
||||
No window catches every slow typist, and no LLM API lets you append to a request already in flight, so there is also a rule for a message that lands *after* the window closed, while the model is still working. At most one turn runs per chat, and a text flush that finds one running is handled in one of three ways:
|
||||
|
||||
- **Dropped** if the running turn's context already includes the message. Its reply answers the new message too.
|
||||
- **Restarted** if nothing has been shown to the user yet. The running request is cancelled and rerun with the new message included. This is the only case that wastes tokens, and it is bounded by the time to first output; with `cache_history` on, the rerun reads the conversation prefix from cache.
|
||||
- **Collected** if output is already on screen. The message gets its own turn as soon as the running one finishes, instead of a second concurrent one. If the finished turn's context turns out to have included the message after all, the extra turn is skipped.
|
||||
|
||||
Photos, voice, and stickers wait for a running turn to finish and are never restarted. Photo uploads still start immediately and overlap the running turn; only the model call queues behind it.
|
||||
|
||||
While a turn is running the bot shows Telegram's "typing…" indicator (or "recording audio" while synthesising a voice reply), refreshed every 4 seconds because Telegram expires the status after 5. Without it, a debounce window reads as the bot ignoring you — which is what prompts users to send more messages in the first place.
|
||||
|
||||
> [!NOTE]
|
||||
> A cancelled request is not free. Anthropic bills the input tokens and any output generated before the connection closed, and a web search the turn already ran is billed again on the rerun. That is why a restart only happens before anything was shown, and why a message that never dispatched is still the cheapest message of all.
|
||||
|
||||
### Draft streaming and the Stop button
|
||||
|
||||
```json
|
||||
"stream_drafts": true
|
||||
```
|
||||
|
||||
With `stream_drafts` on, a reply in a private chat is streamed into one Telegram draft that grows in place (`sendMessageDraft`, available to every bot since Bot API 9.5), then sent as a single real message when the model finishes. Long replies split at 4096 characters. The draft is updated at most once per second, matching Telegram's per-chat send guidance, and re-sent every 10 seconds through silent stretches such as a web search so the preview does not expire.
|
||||
|
||||
The draft carries Telegram's native **Stop** button (Bot API 10.3). Tapping it cancels the Anthropic stream; whatever the user had watched being written is sent as a normal message and stored in memory, so the next turn continues from the reply the user actually saw. A message the user sent before tapping Stop still gets its own turn afterwards. The draft text is capped at 4096 characters; the final message is split instead.
|
||||
|
||||
What falls back to the previous one-message-per-text-block delivery:
|
||||
|
||||
- **Business connections and group chats.** Telegram documents drafts for private chats only, and `sendMessageDraft` has no `business_connection_id` parameter.
|
||||
- **A rejected draft call.** After the first error the bot stops touching the draft and just sends the final message.
|
||||
|
||||
> [!IMPORTANT]
|
||||
> While a draft is on screen, Telegram Android disables the composer: the user cannot send text or commands until the final message lands, and Telegram has said this is intended. The Stop button is their way out. For that reason the draft only appears once the model starts producing text; tool calls before that are covered by the typing indicator, which locks nothing. A March 2026 report described drafts rendering under "Pinned Messages" on iOS, so check the behaviour on your own client before enabling this for users.
|
||||
|
||||
### Web search / fetch
|
||||
|
||||
The `web_search` block wires in Anthropic's server-side `web_search` (and optionally `web_fetch`) tools, mirroring how `mcp_servers` wires in MCP toolsets: the engine provides the generic mechanism, and the per-bot config carries the policy. **No domains are hardcoded in the engine** — each profile declares its own allowlist. When the model decides a question needs current information, it runs a search server-side; results and any fetch targets are confined to the domains you list.
|
||||
|
||||
```json
|
||||
"web_search": {
|
||||
"allowed_domains": [
|
||||
"example.com/help",
|
||||
"docs.example.com"
|
||||
],
|
||||
"max_uses": 3,
|
||||
"fetch": true,
|
||||
"max_content_tokens": 50000
|
||||
}
|
||||
```
|
||||
|
||||
| Field | Type | Description |
|
||||
| ---------------------- | -------- | ----------- |
|
||||
| `allowed_domains` | string[] | Restricts **`web_search`** results (and, by default, `web_fetch` targets) to these. Each entry is a **domain with an optional path prefix and no scheme** — `example.com` or `example.com/help`. Mutually exclusive with `blocked_domains`. **The path only scopes `web_search`** — see the asymmetry note below. |
|
||||
| `blocked_domains` | string[] | Excludes these domains instead of allowlisting. Mutually exclusive with `allowed_domains` (setting both fails validation at boot, since the API 400s). |
|
||||
| `fetch_allowed_domains`| string[] | **Host-only** allowlist for `web_fetch`, independent of the search list. Omit it and `web_fetch` reuses the *hosts* of `allowed_domains`. Set it to fetch a **narrower** set than you search — e.g. to keep a shared host (a social platform) search-only. Only meaningful with `fetch: true`. |
|
||||
| `max_uses` | number | Caps how many searches the model may run per turn. Omit for no cap. |
|
||||
| `fetch` | bool | Also enable `web_fetch`, which pulls full page content for URLs already surfaced by a search or pasted by the user (it cannot fetch model-invented URLs). Citations are always on when fetch is enabled, so fetched passages are sourceable. |
|
||||
| `max_content_tokens` | number | Caps the tokens a single `web_fetch` may pull into context. Only meaningful with `fetch: true`. |
|
||||
| `dynamic_filtering` | bool | Defaults to `true`. Lets the tools filter results inside Anthropic's sandbox before they enter context — see below. Set `false` for models older than Claude 4.6, which reject the tool unless it is called directly. |
|
||||
|
||||
**Search/fetch path asymmetry (important).** Anthropic filters the two tools differently: `web_search` honors a path prefix (`example.com/help` matches only `example.com/help/...`), but **`web_fetch` matches on the host only — an entry that includes a path never matches any fetch URL**. The engine bridges this: `web_search` gets your entries verbatim, while `web_fetch` gets the host portion of each entry (path stripped, deduped). So a path-scoped `allowed_domains` still permits fetching across the whole host. If that host isn't wholly trusted, list the fetchable hosts explicitly in `fetch_allowed_domains` instead.
|
||||
|
||||
**Search-only pattern for shared hosts (e.g. social).** To let the model *search* a single account on a shared platform without ever *fetching* the wider platform, put the account path in `allowed_domains` and leave its host out of `fetch_allowed_domains`:
|
||||
|
||||
```json
|
||||
"web_search": {
|
||||
"allowed_domains": ["helpcenter.example/hc", "x.com/youraccount"],
|
||||
"fetch_allowed_domains": ["helpcenter.example"],
|
||||
"fetch": true
|
||||
}
|
||||
```
|
||||
|
||||
Here search is confined to your help center *and* your one social account, but fetch can only ever reach `helpcenter.example` — a pasted `x.com/someone-else` URL is not fetchable. Note the search side only sandboxes cleanly on platforms that keep an account's content under its handle path (X/Twitter, TikTok `@handle`, Facebook); use the `handle/` or `handle/*` form to avoid a prefix bleed. Instagram is an exception — posts live at `instagram.com/p/<code>`, not under the handle — so it can't be account-sandboxed by path.
|
||||
|
||||
The allowlist is a **hard, server-side sandbox**, not a prompt request — off-list results are dropped and an off-list fetch target returns `url_not_allowed`, which is the exfiltration mitigation Anthropic recommends for bots processing untrusted input. Keep the list tight. A fetch can still fail for site-side reasons (bot protection such as Cloudflare returns `url_not_accessible`); the search snippet plus citations remain the reliable signal, so a full fetch is a bonus, not a dependency.
|
||||
|
||||
`web_fetch` can only open a URL that a `web_search` result surfaced (or the user pasted) — never one the model invents or rebuilds from a page title; those return `url_not_in_prior_context`. So a fetch-enabled bot's prompt should steer it to **search first, then fetch a returned URL**, and to **re-search with a more specific query** (rather than guess a URL) when the exact page it wants isn't in the results.
|
||||
|
||||
**Dynamic filtering.** The engine uses the `web_search_20260318` / `web_fetch_20260318` tools, which run the search inside Anthropic's code-execution sandbox and let the model discard irrelevant parts of a result before it reaches the context window. On token-heavy sources (a long docs page, a GitHub README) this is the difference between paying for the whole page and paying for the part that answered the question; there is no extra charge for the code execution it uses. It needs a **Claude 4.6 or later** model — on anything older (Haiku 4.5, Sonnet 4.5) the API rejects the request unless the tool is called directly, so set `"dynamic_filtering": false` on those profiles. Disabling it also restores Zero Data Retention eligibility, which the sandbox otherwise forfeits.
|
||||
|
||||
Web search is **not free** — Anthropic bills per search (plus the tokens the results add) — so scope the allowlist and `max_uses` deliberately, and lean on a prompt that tells the bot to search only when a question actually turns on current or authoritative information. Search activity logs as `[web] ...` lines, and long multi-search turns are handled transparently (the engine follows Anthropic's `pause_turn` continuations up to a small cap, so a slow search doesn't truncate the reply).
|
||||
|
||||
> [!TIP]
|
||||
> For deep request/response debugging, the Anthropic Go SDK ships `option.WithDebugLog(...)` (dumps full HTTP bodies with auth headers redacted). It is not wired into the bot — dev-only, add it temporarily to the client constructor if you ever need wire-level traces.
|
||||
|
||||
### Future: persistent memory
|
||||
|
||||
The Anthropic memory tool (`memory_20250818`) is a candidate future feature for cross-conversation recall (a self-hosted analog of ChatGPT's "memory"). The Go SDK already ships the types (`BetaMemoryTool20250818Param` and its tool-union slot plus the six-command union: `view`/`create`/`str_replace`/`insert`/`delete`/`rename`), but — unlike the Python/TypeScript/Java SDKs — provides **no handler helper**: the bot would have to hand-write client-side command dispatch against per-chat storage, including strict path validation (canonicalize and confine every model-supplied path under a fixed memory root; reject `..`/symlink traversal) and a no-secrets policy for stored content. Not implemented yet.
|
||||
|
||||
## Systemd Unit Setup
|
||||
|
||||
To enable the bot to start automatically on system boot and run in the background, set up a systemd unit.
|
||||
@@ -263,7 +123,6 @@ journalctl -u telegram-bot -f
|
||||
| `/clear_hard` | All users | Permanently delete your own chat history |
|
||||
| `/clear_hard <user_id>` | Admin/Owner | Permanently delete all messages for a user across every chat |
|
||||
| `/clear_hard <user_id> <chat_id>` | Admin/Owner | Permanently delete a user's messages in a specific chat |
|
||||
| `/set_model <model-id>` | Admin/Owner | Switch the AI model live without restarting |
|
||||
|
||||
> **Note:** In private DMs each user's `chat_id` equals their `user_id`. The scoped `<chat_id>` form is mainly useful for group chat moderation.
|
||||
|
||||
|
||||
@@ -1,88 +0,0 @@
|
||||
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,
|
||||
)
|
||||
}
|
||||
+105
-452
@@ -2,476 +2,129 @@ package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/anthropics/anthropic-sdk-go/v2"
|
||||
"github.com/anthropics/anthropic-sdk-go/v2/packages/param"
|
||||
"github.com/liushuangls/go-anthropic/v2"
|
||||
)
|
||||
|
||||
var ErrModelNotFound = errors.New("model not found or deprecated")
|
||||
|
||||
const maxFileNotFoundRetries = 3
|
||||
|
||||
const maxPauseTurnContinuations = 5
|
||||
|
||||
const defaultMaxTokens = 1000
|
||||
|
||||
const mcpUnsupportedSentinel = "format not currently supported by the Anthropic API"
|
||||
|
||||
var mcpUnsupportedCount atomic.Uint64
|
||||
|
||||
type mcpCall struct{ server, name, input string }
|
||||
|
||||
func (b *Bot) getAnthropicResponse(ctx context.Context, chatID int64, messages []anthropic.BetaMessageParam, isEmojiOnly bool, username string, firstName string, lastName string, isPremium bool, languageCode string, messageTime int, onSegment func(string) error, onProgress func(string)) (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,
|
||||
}
|
||||
|
||||
if staticPrompt != "" {
|
||||
blocks := []anthropic.BetaTextBlockParam{
|
||||
{Text: staticPrompt, CacheControl: anthropic.NewBetaCacheControlEphemeralParam()},
|
||||
}
|
||||
tail := buildUserContext(username, firstName, lastName, isPremium, languageCode, messageTime)
|
||||
if isEmojiOnly {
|
||||
if rule := strings.TrimSpace(b.config.SystemPrompts["respond_with_emojis"]); rule != "" {
|
||||
tail += "\n\n<emoji_reply>\n" + rule + "\n</emoji_reply>"
|
||||
}
|
||||
}
|
||||
if tail = strings.TrimSpace(tail); tail != "" {
|
||||
blocks = append(blocks, anthropic.BetaTextBlockParam{Text: tail})
|
||||
}
|
||||
params.System = blocks
|
||||
}
|
||||
|
||||
if thinking, ok := thinkingParamFromConfig(b.config.Thinking, b.config.ThinkingDisplay); ok {
|
||||
params.Thinking = thinking
|
||||
}
|
||||
|
||||
if b.config.Effort != "" {
|
||||
params.OutputConfig = anthropic.BetaOutputConfigParam{
|
||||
Effort: anthropic.BetaOutputConfigEffort(b.config.Effort),
|
||||
}
|
||||
}
|
||||
|
||||
var tools []anthropic.BetaToolUnionParam
|
||||
|
||||
if len(b.config.MCPServers) > 0 {
|
||||
mcpServers := make([]anthropic.BetaRequestMCPServerURLDefinitionParam, 0, len(b.config.MCPServers))
|
||||
for _, s := range b.config.MCPServers {
|
||||
srv := anthropic.BetaRequestMCPServerURLDefinitionParam{
|
||||
Name: s.Name,
|
||||
URL: s.URL,
|
||||
}
|
||||
if s.AuthorizationToken != "" {
|
||||
srv.AuthorizationToken = param.NewOpt(s.AuthorizationToken)
|
||||
}
|
||||
mcpServers = append(mcpServers, srv)
|
||||
|
||||
toolset := &anthropic.BetaMCPToolsetParam{
|
||||
MCPServerName: s.Name,
|
||||
}
|
||||
if len(s.AllowedTools) > 0 {
|
||||
toolset.DefaultConfig = anthropic.BetaMCPToolDefaultConfigParam{
|
||||
Enabled: param.NewOpt(false),
|
||||
}
|
||||
toolset.Configs = make(map[string]anthropic.BetaMCPToolConfigParam, len(s.AllowedTools))
|
||||
for _, tool := range s.AllowedTools {
|
||||
toolset.Configs[tool] = anthropic.BetaMCPToolConfigParam{
|
||||
Enabled: param.NewOpt(true),
|
||||
}
|
||||
}
|
||||
}
|
||||
tools = append(tools, anthropic.BetaToolUnionParam{OfMCPToolset: toolset})
|
||||
}
|
||||
params.MCPServers = mcpServers
|
||||
params.Betas = append(params.Betas, anthropic.AnthropicBetaMCPClient2025_11_20)
|
||||
}
|
||||
|
||||
tools = append(tools, webSearchTools(b.config.WebSearch)...)
|
||||
|
||||
if len(tools) > 0 {
|
||||
params.Tools = tools
|
||||
}
|
||||
|
||||
var full string
|
||||
var lastMsg anthropic.BetaMessage
|
||||
fileRetries, pauseContinuations := 0, 0
|
||||
for {
|
||||
// full is passed as the progress prefix so that, across pause_turn
|
||||
// continuations, the user watches the whole reply so far.
|
||||
joined, msg, streamErr := b.streamMessages(ctx, params, onSegment, onProgress, full)
|
||||
if streamErr != nil {
|
||||
var apiErr *anthropic.Error
|
||||
if !errors.As(streamErr, &apiErr) || apiErr.StatusCode != http.StatusNotFound {
|
||||
return "", fmt.Errorf("error creating Anthropic message: %w", streamErr)
|
||||
}
|
||||
missingFileID := extractMissingFileID(streamErr)
|
||||
if missingFileID == "" {
|
||||
return "", fmt.Errorf("%w: %s", ErrModelNotFound, b.config.Model)
|
||||
}
|
||||
fileRetries++
|
||||
if fileRetries > maxFileNotFoundRetries {
|
||||
return "", fmt.Errorf("max self-heal retries (%d) exceeded: too many file_ids gone from anthropic", maxFileNotFoundRetries)
|
||||
}
|
||||
ErrorLogger.Printf("[%s] self-heal: stripping dead file_id %s from chat %d (attempt %d/%d)",
|
||||
b.config.ID, missingFileID, chatID, fileRetries, maxFileNotFoundRetries)
|
||||
b.stripDeadFileIDFromMemory(chatID, missingFileID)
|
||||
if _, cleanupErr := b.markFilesPendingCleanup(ctx, chatID, []string{missingFileID}); cleanupErr != nil {
|
||||
ErrorLogger.Printf("[%s] mark files pending cleanup: %v", b.config.ID, cleanupErr)
|
||||
}
|
||||
params.Messages = b.prepareContextMessages(b.getOrCreateChatMemory(chatID))
|
||||
continue
|
||||
}
|
||||
|
||||
lastMsg = msg
|
||||
full = joinNonEmpty(full, joined)
|
||||
|
||||
if msg.StopReason == anthropic.BetaStopReasonPauseTurn {
|
||||
pauseContinuations++
|
||||
if pauseContinuations > maxPauseTurnContinuations {
|
||||
ErrorLogger.Printf("[%s] pause_turn continuations exceeded (%d); returning partial answer",
|
||||
b.config.ID, maxPauseTurnContinuations)
|
||||
break
|
||||
}
|
||||
params.Messages = append(params.Messages, msg.ToParam())
|
||||
continue
|
||||
}
|
||||
break
|
||||
}
|
||||
|
||||
if full == "" {
|
||||
return "", emptyStreamError(string(lastMsg.StopReason), string(lastMsg.StopDetails.Category),
|
||||
lastMsg.Usage.OutputTokensDetails.ThinkingTokens, params.MaxTokens)
|
||||
}
|
||||
return full, nil
|
||||
}
|
||||
|
||||
func webSearchTools(cfg *WebSearchConfig) []anthropic.BetaToolUnionParam {
|
||||
if cfg == nil {
|
||||
return nil
|
||||
}
|
||||
// Dynamic filtering runs the tool inside code execution so results are
|
||||
// filtered before they reach the context window. It requires a Claude 4.6+
|
||||
// model; older ones must call the tool directly.
|
||||
var callers []string
|
||||
if !cfg.DynamicFilteringEnabled() {
|
||||
callers = []string{"direct"}
|
||||
}
|
||||
search := &anthropic.BetaWebSearchTool20260318Param{
|
||||
AllowedDomains: cfg.AllowedDomains,
|
||||
BlockedDomains: cfg.BlockedDomains,
|
||||
AllowedCallers: callers,
|
||||
}
|
||||
if cfg.MaxUses > 0 {
|
||||
search.MaxUses = param.NewOpt(int64(cfg.MaxUses))
|
||||
}
|
||||
tools := []anthropic.BetaToolUnionParam{{OfWebSearchTool20260318: search}}
|
||||
|
||||
if cfg.Fetch {
|
||||
fetchAllowed := cfg.FetchAllowedDomains
|
||||
if len(fetchAllowed) == 0 {
|
||||
fetchAllowed = cfg.AllowedDomains
|
||||
}
|
||||
fetch := &anthropic.BetaWebFetchTool20260318Param{
|
||||
AllowedDomains: fetchHosts(fetchAllowed),
|
||||
BlockedDomains: fetchHosts(cfg.BlockedDomains),
|
||||
AllowedCallers: callers,
|
||||
Citations: anthropic.BetaCitationsConfigParam{Enabled: param.NewOpt(true)},
|
||||
}
|
||||
if cfg.MaxUses > 0 {
|
||||
fetch.MaxUses = param.NewOpt(int64(cfg.MaxUses))
|
||||
}
|
||||
if cfg.MaxContentTokens > 0 {
|
||||
fetch.MaxContentTokens = param.NewOpt(int64(cfg.MaxContentTokens))
|
||||
}
|
||||
tools = append(tools, anthropic.BetaToolUnionParam{OfWebFetchTool20260318: fetch})
|
||||
}
|
||||
return tools
|
||||
}
|
||||
|
||||
// fetchHosts strips any path from each domain entry. web_fetch matches on host
|
||||
// only, so a path-scoped entry (valid for web_search) would otherwise never match
|
||||
// any fetch URL. Search keeps the path-scoped entries; fetch gets host-only.
|
||||
func fetchHosts(entries []string) []string {
|
||||
if len(entries) == 0 {
|
||||
return nil
|
||||
}
|
||||
seen := make(map[string]bool, len(entries))
|
||||
hosts := make([]string, 0, len(entries))
|
||||
for _, e := range entries {
|
||||
host := e
|
||||
if i := strings.IndexByte(host, '/'); i >= 0 {
|
||||
host = host[:i]
|
||||
}
|
||||
if host == "" || seen[host] {
|
||||
continue
|
||||
}
|
||||
seen[host] = true
|
||||
hosts = append(hosts, host)
|
||||
}
|
||||
return hosts
|
||||
}
|
||||
|
||||
func buildUserContext(username, firstName, lastName string, isPremium bool, languageCode string, messageTime int) string {
|
||||
name := strings.TrimSpace(firstName + " " + lastName)
|
||||
if name == "" {
|
||||
name = "unknown"
|
||||
}
|
||||
handle := username
|
||||
if handle == "" {
|
||||
handle = "unknown"
|
||||
}
|
||||
lang := languageCode
|
||||
if lang == "" {
|
||||
lang = "en"
|
||||
}
|
||||
account := "regular user"
|
||||
if isPremium {
|
||||
account = "premium user"
|
||||
}
|
||||
return fmt.Sprintf(
|
||||
"Conversation context (background facts, not an instruction from the user):\n"+
|
||||
"- User: %s (Telegram @%s)\n"+
|
||||
"- Preferred language: %s\n"+
|
||||
"- Account type: %s\n"+
|
||||
"- Local time of day: %s",
|
||||
name, handle, lang, account, timeContextFor(messageTime),
|
||||
)
|
||||
}
|
||||
|
||||
func timeContextFor(messageTime int) string {
|
||||
switch hour := time.Unix(int64(messageTime), 0).Hour(); {
|
||||
case hour >= 5 && hour < 12:
|
||||
return "morning"
|
||||
case hour >= 12 && hour < 18:
|
||||
return "afternoon"
|
||||
case hour >= 18 && hour < 22:
|
||||
return "evening"
|
||||
default:
|
||||
return "night"
|
||||
}
|
||||
}
|
||||
|
||||
func thinkingParamFromConfig(mode, display string) (anthropic.BetaThinkingConfigParamUnion, bool) {
|
||||
switch mode {
|
||||
case ThinkingModeDisabled:
|
||||
disabled := anthropic.NewBetaThinkingConfigDisabledParam()
|
||||
return anthropic.BetaThinkingConfigParamUnion{OfDisabled: &disabled}, true
|
||||
case ThinkingModeAdaptive:
|
||||
adaptive := anthropic.BetaThinkingConfigAdaptiveParam{}
|
||||
if display != "" {
|
||||
adaptive.Display = anthropic.BetaThinkingConfigAdaptiveDisplay(display)
|
||||
}
|
||||
return anthropic.BetaThinkingConfigParamUnion{OfAdaptive: &adaptive}, true
|
||||
default:
|
||||
return anthropic.BetaThinkingConfigParamUnion{}, false
|
||||
}
|
||||
}
|
||||
|
||||
// streamMessages runs one streaming request and returns its text blocks joined.
|
||||
// onSegment fires with each completed text block; onProgress fires on every
|
||||
// text delta with the whole reply so far, prefix included, for sinks that
|
||||
// redraw in place rather than append. prefix is what earlier requests of the
|
||||
// same turn already produced and is not part of the returned text.
|
||||
func (b *Bot) streamMessages(ctx context.Context, params anthropic.BetaMessageNewParams, onSegment func(string) error, onProgress func(string), prefix string) (string, anthropic.BetaMessage, error) {
|
||||
stream := b.anthropicClient.Beta.Messages.NewStreaming(ctx, params)
|
||||
defer func() {
|
||||
if err := stream.Close(); err != nil {
|
||||
ErrorLogger.Printf("[stream] close failed: %v", err)
|
||||
}
|
||||
}()
|
||||
|
||||
var (
|
||||
message anthropic.BetaMessage
|
||||
done string // this request's completed text blocks, joined
|
||||
doneAll = prefix // same, with the turn's earlier requests in front
|
||||
currentKind string
|
||||
currentText strings.Builder
|
||||
currentThinking strings.Builder
|
||||
currentInputJSON strings.Builder
|
||||
currentTUseName, currentTUseServer, currentTUseID string
|
||||
currentTResultUseID, currentTResultServer string
|
||||
currentTResultIsError bool
|
||||
currentTResultContent string
|
||||
currentServerToolName, currentServerToolID string
|
||||
currentServerResult string
|
||||
mcpCalls = map[string]mcpCall{}
|
||||
)
|
||||
|
||||
for stream.Next() {
|
||||
e := stream.Current()
|
||||
if accErr := message.Accumulate(e); accErr != nil {
|
||||
ErrorLogger.Printf("[stream] accumulate failed: %v", accErr)
|
||||
}
|
||||
switch e.Type {
|
||||
case "content_block_start":
|
||||
cbs := e.AsContentBlockStart()
|
||||
currentKind = cbs.ContentBlock.Type
|
||||
currentText.Reset()
|
||||
currentThinking.Reset()
|
||||
currentInputJSON.Reset()
|
||||
currentServerResult = ""
|
||||
switch currentKind {
|
||||
case "mcp_tool_use":
|
||||
currentTUseName = cbs.ContentBlock.Name
|
||||
currentTUseServer = cbs.ContentBlock.ServerName
|
||||
currentTUseID = cbs.ContentBlock.ID
|
||||
case "mcp_tool_result":
|
||||
currentTResultUseID = cbs.ContentBlock.ToolUseID
|
||||
currentTResultServer = cbs.ContentBlock.ServerName
|
||||
currentTResultIsError = cbs.ContentBlock.IsError
|
||||
currentTResultContent = cbs.ContentBlock.JSON.Content.Raw()
|
||||
case "server_tool_use":
|
||||
currentServerToolName = cbs.ContentBlock.Name
|
||||
currentServerToolID = cbs.ContentBlock.ID
|
||||
case "web_search_tool_result", "web_fetch_tool_result":
|
||||
currentServerResult = cbs.ContentBlock.JSON.Content.Raw()
|
||||
}
|
||||
|
||||
case "content_block_delta":
|
||||
cbd := e.AsContentBlockDelta()
|
||||
switch cbd.Delta.Type {
|
||||
case "text_delta":
|
||||
if currentKind == "text" {
|
||||
currentText.WriteString(cbd.Delta.Text)
|
||||
if onProgress != nil {
|
||||
onProgress(joinNonEmpty(doneAll, currentText.String()))
|
||||
}
|
||||
}
|
||||
case "thinking_delta":
|
||||
if currentKind == "thinking" {
|
||||
currentThinking.WriteString(cbd.Delta.Thinking)
|
||||
}
|
||||
case "input_json_delta":
|
||||
if currentKind == "mcp_tool_use" || currentKind == "server_tool_use" {
|
||||
currentInputJSON.WriteString(cbd.Delta.PartialJSON)
|
||||
}
|
||||
}
|
||||
|
||||
case "content_block_stop":
|
||||
switch currentKind {
|
||||
case "text":
|
||||
seg := strings.TrimSpace(currentText.String())
|
||||
if seg != "" {
|
||||
done = joinNonEmpty(done, seg)
|
||||
doneAll = joinNonEmpty(doneAll, seg)
|
||||
if onSegment != nil {
|
||||
if cbErr := onSegment(seg); cbErr != nil {
|
||||
ErrorLogger.Printf("[stream] onSegment failed: %v", cbErr)
|
||||
}
|
||||
}
|
||||
}
|
||||
case "mcp_tool_use":
|
||||
mcpCalls[currentTUseID] = mcpCall{
|
||||
server: currentTUseServer,
|
||||
name: currentTUseName,
|
||||
input: currentInputJSON.String(),
|
||||
}
|
||||
InfoLogger.Printf("[mcp] tool_use server=%q name=%q id=%q input=%s",
|
||||
currentTUseServer, currentTUseName, currentTUseID, currentInputJSON.String())
|
||||
case "mcp_tool_result":
|
||||
preview := currentTResultContent
|
||||
if len(preview) > 500 {
|
||||
preview = preview[:500] + "...(truncated)"
|
||||
}
|
||||
InfoLogger.Printf("[mcp] tool_result tool_use_id=%q server=%q is_error=%v content=%s",
|
||||
currentTResultUseID, currentTResultServer, currentTResultIsError, preview)
|
||||
if strings.Contains(currentTResultContent, mcpUnsupportedSentinel) {
|
||||
total := mcpUnsupportedCount.Add(1)
|
||||
call := mcpCalls[currentTResultUseID]
|
||||
ErrorLogger.Printf("[%s][mcp][unsupported] connector could not serialize result "+
|
||||
"(total=%d): server=%q tool=%q input=%s tool_use_id=%q",
|
||||
b.config.ID, total, call.server, call.name, call.input, currentTResultUseID)
|
||||
}
|
||||
case "server_tool_use":
|
||||
InfoLogger.Printf("[web] %s id=%q input=%s",
|
||||
currentServerToolName, currentServerToolID, currentInputJSON.String())
|
||||
case "web_search_tool_result", "web_fetch_tool_result":
|
||||
preview := currentServerResult
|
||||
if len(preview) > 500 {
|
||||
preview = preview[:500] + "...(truncated)"
|
||||
}
|
||||
InfoLogger.Printf("[web] %s content=%s", currentKind, preview)
|
||||
case "thinking", "redacted_thinking":
|
||||
if summary := strings.TrimSpace(currentThinking.String()); summary != "" {
|
||||
if len(summary) > 500 {
|
||||
summary = summary[:500] + "...(truncated)"
|
||||
}
|
||||
InfoLogger.Printf("[thinking] block complete: %s", summary)
|
||||
func (b *Bot) getAnthropicResponse(ctx context.Context, messages []anthropic.Message, isNewChat, isAdminOrOwner, isEmojiOnly bool, username string, firstName string, lastName string, isPremium bool, languageCode string, messageTime int) (string, error) {
|
||||
// Use prompts from config
|
||||
var systemMessage string
|
||||
if isNewChat {
|
||||
systemMessage = b.config.SystemPrompts["new_chat"]
|
||||
} else {
|
||||
InfoLogger.Printf("[thinking] block complete (content omitted)")
|
||||
systemMessage = b.config.SystemPrompts["continue_conversation"]
|
||||
}
|
||||
|
||||
// Combine default prompt with custom instructions
|
||||
systemMessage = b.config.SystemPrompts["default"] + " " + b.config.SystemPrompts["custom_instructions"] + " " + systemMessage
|
||||
|
||||
// Handle username placeholder
|
||||
usernameValue := username
|
||||
if username == "" {
|
||||
usernameValue = "unknown" // Use "unknown" when username is not available
|
||||
}
|
||||
systemMessage = strings.ReplaceAll(systemMessage, "{username}", usernameValue)
|
||||
|
||||
// Handle firstname placeholder
|
||||
firstnameValue := firstName
|
||||
if firstName == "" {
|
||||
firstnameValue = "unknown" // Use "unknown" when first name is not available
|
||||
}
|
||||
systemMessage = strings.ReplaceAll(systemMessage, "{firstname}", firstnameValue)
|
||||
|
||||
// Handle lastname placeholder
|
||||
lastnameValue := lastName
|
||||
if lastName == "" {
|
||||
lastnameValue = "" // Empty string when last name is not available
|
||||
}
|
||||
systemMessage = strings.ReplaceAll(systemMessage, "{lastname}", lastnameValue)
|
||||
|
||||
// Handle language code placeholder
|
||||
langValue := languageCode
|
||||
if languageCode == "" {
|
||||
langValue = "en" // Default to English when language code is not available
|
||||
}
|
||||
systemMessage = strings.ReplaceAll(systemMessage, "{language}", langValue)
|
||||
|
||||
// Handle premium status
|
||||
premiumStatus := "regular user"
|
||||
if isPremium {
|
||||
premiumStatus = "premium user"
|
||||
}
|
||||
systemMessage = strings.ReplaceAll(systemMessage, "{premium_status}", premiumStatus)
|
||||
|
||||
// Handle time awareness
|
||||
timeObj := time.Unix(int64(messageTime), 0)
|
||||
hour := timeObj.Hour()
|
||||
var timeContext string
|
||||
if hour >= 5 && hour < 12 {
|
||||
timeContext = "morning"
|
||||
} else if hour >= 12 && hour < 18 {
|
||||
timeContext = "afternoon"
|
||||
} else if hour >= 18 && hour < 22 {
|
||||
timeContext = "evening"
|
||||
} else {
|
||||
timeContext = "night"
|
||||
}
|
||||
systemMessage = strings.ReplaceAll(systemMessage, "{time_context}", timeContext)
|
||||
|
||||
if !isAdminOrOwner {
|
||||
systemMessage += " " + b.config.SystemPrompts["avoid_sensitive"]
|
||||
}
|
||||
|
||||
if isEmojiOnly {
|
||||
systemMessage += " " + b.config.SystemPrompts["respond_with_emojis"]
|
||||
}
|
||||
|
||||
// Debug logging
|
||||
InfoLogger.Printf("Sending %d messages to Anthropic", len(messages))
|
||||
for i, msg := range messages {
|
||||
for _, content := range msg.Content {
|
||||
if content.Type == anthropic.MessagesContentTypeText {
|
||||
InfoLogger.Printf("Message %d: Role=%v, Text=%v", i, msg.Role, content.Text)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Ensure the roles are correct
|
||||
for i := range messages {
|
||||
switch messages[i].Role {
|
||||
case anthropic.RoleUser:
|
||||
messages[i].Role = anthropic.RoleUser
|
||||
case anthropic.RoleAssistant:
|
||||
messages[i].Role = anthropic.RoleAssistant
|
||||
default:
|
||||
if currentKind != "" {
|
||||
InfoLogger.Printf("[stream] block type=%q (unhandled)", currentKind)
|
||||
}
|
||||
}
|
||||
currentKind = ""
|
||||
// Default to 'user' if role is unrecognized
|
||||
messages[i].Role = anthropic.RoleUser
|
||||
}
|
||||
}
|
||||
|
||||
if err := stream.Err(); err != nil {
|
||||
return "", message, err
|
||||
model := anthropic.Model(b.config.Model)
|
||||
|
||||
// Create the request
|
||||
request := anthropic.MessagesRequest{
|
||||
Model: model, // Now `model` is of type anthropic.Model
|
||||
Messages: messages,
|
||||
System: systemMessage,
|
||||
MaxTokens: 1000,
|
||||
}
|
||||
|
||||
stopReason := string(message.StopReason)
|
||||
if stopReason != "" || message.Usage.OutputTokens > 0 {
|
||||
// cache_read/cache_write make the caching configuration falsifiable: a
|
||||
// breakpoint below the model's minimum cacheable prefix fails silently,
|
||||
// reporting cache_write=0 rather than raising an error.
|
||||
InfoLogger.Printf("[usage] model=%s in=%d out=%d thinking=%d cache_read=%d cache_write=%d stop=%s",
|
||||
params.Model, message.Usage.InputTokens, message.Usage.OutputTokens,
|
||||
message.Usage.OutputTokensDetails.ThinkingTokens,
|
||||
message.Usage.CacheReadInputTokens, message.Usage.CacheCreationInputTokens,
|
||||
stopReason)
|
||||
if message.StopReason == anthropic.BetaStopReasonMaxTokens {
|
||||
ErrorLogger.Printf("[usage] response truncated at max_tokens=%d - raise max_tokens (thinking counts toward it)",
|
||||
params.MaxTokens)
|
||||
}
|
||||
// Apply temperature if set in config
|
||||
if b.config.Temperature != nil {
|
||||
request.Temperature = b.config.Temperature
|
||||
}
|
||||
|
||||
return done, message, nil
|
||||
resp, err := b.anthropicClient.CreateMessages(ctx, request)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("error creating Anthropic message: %w", err)
|
||||
}
|
||||
|
||||
// joinNonEmpty is the one rule for assembling reply text from parts: a blank
|
||||
// line between them, empty parts skipped. The live draft, the final message,
|
||||
// and the stored partial all go through it, so they can never disagree.
|
||||
func joinNonEmpty(a, b string) string {
|
||||
switch {
|
||||
case a == "":
|
||||
return b
|
||||
case b == "":
|
||||
return a
|
||||
default:
|
||||
return a + "\n\n" + b
|
||||
}
|
||||
if len(resp.Content) == 0 || resp.Content[0].Type != anthropic.MessagesContentTypeText {
|
||||
return "", fmt.Errorf("unexpected response format from Anthropic")
|
||||
}
|
||||
|
||||
func emptyStreamError(stopReason, refusalCategory string, thinkingTokens, maxTokens int64) error {
|
||||
switch stopReason {
|
||||
case "max_tokens":
|
||||
return fmt.Errorf("output budget exhausted before any text (thinking used %d of %d max_tokens) - raise max_tokens",
|
||||
thinkingTokens, maxTokens)
|
||||
case "refusal":
|
||||
if refusalCategory == "" {
|
||||
return fmt.Errorf("the model declined the request (refusal)")
|
||||
}
|
||||
return fmt.Errorf("the model declined the request (refusal, category %s)", refusalCategory)
|
||||
}
|
||||
return fmt.Errorf("unexpected response format from Anthropic")
|
||||
return resp.Content[0].GetText(), nil
|
||||
}
|
||||
|
||||
@@ -1,169 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/anthropics/anthropic-sdk-go/v2"
|
||||
)
|
||||
|
||||
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.Files.Upload(ctx, anthropic.FileUploadParams{
|
||||
File: anthropic.File(bytes.NewReader(data), filename, contentType),
|
||||
})
|
||||
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.Files.Delete(ctx, fileID, anthropic.FileDeleteParams{})
|
||||
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
|
||||
}
|
||||
@@ -1,191 +0,0 @@
|
||||
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)
|
||||
}
|
||||
+153
-314
@@ -1,358 +1,197 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/anthropics/anthropic-sdk-go/v2"
|
||||
)
|
||||
|
||||
func TestTimeContextFor(t *testing.T) {
|
||||
cases := []struct {
|
||||
// TestLanguageCodeReplacement tests that language code is properly handled and replaced
|
||||
func TestLanguageCodeReplacement(t *testing.T) {
|
||||
// Test with provided language code
|
||||
systemMessage := "User's language preference: '{language}'"
|
||||
|
||||
// Test with a specific language code
|
||||
langValue := "fr"
|
||||
result := strings.ReplaceAll(systemMessage, "{language}", langValue)
|
||||
|
||||
if !strings.Contains(result, "User's language preference: 'fr'") {
|
||||
t.Errorf("Expected language code 'fr' to be replaced, got: %s", result)
|
||||
}
|
||||
|
||||
// Test with empty language code (should default to "en")
|
||||
langValue = ""
|
||||
if langValue == "" {
|
||||
langValue = "en" // Default to English when language code is not available
|
||||
}
|
||||
result = strings.ReplaceAll(systemMessage, "{language}", langValue)
|
||||
|
||||
if !strings.Contains(result, "User's language preference: 'en'") {
|
||||
t.Errorf("Expected default language code 'en' to be used, got: %s", result)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPremiumStatusReplacement tests that premium status is properly handled and replaced
|
||||
func TestPremiumStatusReplacement(t *testing.T) {
|
||||
systemMessage := "User is a {premium_status}"
|
||||
|
||||
// Test with premium user
|
||||
isPremium := true
|
||||
premiumStatus := "regular user"
|
||||
if isPremium {
|
||||
premiumStatus = "premium user"
|
||||
}
|
||||
result := strings.ReplaceAll(systemMessage, "{premium_status}", premiumStatus)
|
||||
|
||||
if !strings.Contains(result, "User is a premium user") {
|
||||
t.Errorf("Expected premium status to be replaced with 'premium user', got: %s", result)
|
||||
}
|
||||
|
||||
// Test with regular user
|
||||
isPremium = false
|
||||
premiumStatus = "regular user"
|
||||
if isPremium {
|
||||
premiumStatus = "premium user"
|
||||
}
|
||||
result = strings.ReplaceAll(systemMessage, "{premium_status}", premiumStatus)
|
||||
|
||||
if !strings.Contains(result, "User is a regular user") {
|
||||
t.Errorf("Expected premium status to be replaced with 'regular user', got: %s", result)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTimeContextCalculation tests that time context is correctly calculated for different hours
|
||||
func TestTimeContextCalculation(t *testing.T) {
|
||||
// Test cases for different hours
|
||||
testCases := []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)
|
||||
}
|
||||
}
|
||||
{3, "night"}, // Night: hours < 5 or hours >= 22
|
||||
{5, "morning"}, // Morning: 5 <= hours < 12
|
||||
{12, "afternoon"}, // Afternoon: 12 <= hours < 18
|
||||
{17, "afternoon"}, // Afternoon: 12 <= hours < 18
|
||||
{18, "evening"}, // Evening: 18 <= hours < 22
|
||||
{21, "evening"}, // Evening: 18 <= hours < 22
|
||||
{22, "night"}, // Night: hours < 5 or hours >= 22
|
||||
{23, "night"}, // Night: hours < 5 or hours >= 22
|
||||
}
|
||||
|
||||
func TestBuildUserContext(t *testing.T) {
|
||||
noon := int(time.Date(2025, 5, 15, 12, 0, 0, 0, time.Local).Unix())
|
||||
for _, tc := range testCases {
|
||||
t.Run(fmt.Sprintf("Hour_%d", tc.hour), func(t *testing.T) {
|
||||
// Create a timestamp for the specified hour
|
||||
testTime := time.Date(2025, 5, 15, tc.hour, 0, 0, 0, time.UTC)
|
||||
|
||||
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)
|
||||
}
|
||||
// Get the hour directly from the test time to ensure it's what we expect
|
||||
actualHour := testTime.Hour()
|
||||
if actualHour != tc.hour {
|
||||
t.Fatalf("Test setup error: expected hour %d, got %d", tc.hour, actualHour)
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
// Calculate time context using the same logic as in anthropic.go
|
||||
var timeContext string
|
||||
if actualHour >= 5 && actualHour < 12 {
|
||||
timeContext = "morning"
|
||||
} else if actualHour >= 12 && actualHour < 18 {
|
||||
timeContext = "afternoon"
|
||||
} else if actualHour >= 18 && actualHour < 22 {
|
||||
timeContext = "evening"
|
||||
} else {
|
||||
timeContext = "night"
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
// Check if the calculated time context matches the expected value
|
||||
if timeContext != tc.expected {
|
||||
t.Errorf("For hour %d: expected time context '%s', got '%s'",
|
||||
actualHour, tc.expected, timeContext)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
// TestSystemMessagePlaceholderReplacement tests that all placeholders are correctly replaced
|
||||
func TestSystemMessagePlaceholderReplacement(t *testing.T) {
|
||||
systemMessage := "The user you're talking to has username '{username}' and display name '{firstname} {lastname}'.\n" +
|
||||
"User's language preference: '{language}'\n" +
|
||||
"User is a {premium_status}\n" +
|
||||
"It's currently {time_context} in your timezone"
|
||||
|
||||
func TestWebSearchTools(t *testing.T) {
|
||||
t.Run("nil config yields no tools", func(t *testing.T) {
|
||||
if tools := webSearchTools(nil); tools != nil {
|
||||
t.Errorf("webSearchTools(nil) = %v, want nil", tools)
|
||||
}
|
||||
})
|
||||
// Set up test data
|
||||
username := "testuser"
|
||||
firstName := "Test"
|
||||
lastName := "User"
|
||||
isPremium := true
|
||||
languageCode := "de"
|
||||
|
||||
t.Run("search only when fetch off", func(t *testing.T) {
|
||||
tools := webSearchTools(&WebSearchConfig{
|
||||
AllowedDomains: []string{"example.com/hc"},
|
||||
MaxUses: 3,
|
||||
})
|
||||
if len(tools) != 1 {
|
||||
t.Fatalf("got %d tools, want 1 (search only)", len(tools))
|
||||
}
|
||||
if tools[0].OfWebSearchTool20260318 == nil {
|
||||
t.Fatalf("tools[0] is not a web_search tool")
|
||||
}
|
||||
if tools[0].OfWebFetchTool20260318 != nil {
|
||||
t.Error("web_fetch tool present but fetch is off")
|
||||
}
|
||||
})
|
||||
// Create a timestamp for a specific hour (e.g., 14:00 = afternoon)
|
||||
testTime := time.Date(2025, 5, 15, 14, 0, 0, 0, time.UTC)
|
||||
messageTime := int(testTime.Unix())
|
||||
|
||||
t.Run("search + fetch with allowlist and citations", func(t *testing.T) {
|
||||
tools := webSearchTools(&WebSearchConfig{
|
||||
AllowedDomains: []string{"example.com/hc", "docs.example.com"},
|
||||
MaxUses: 3,
|
||||
Fetch: true,
|
||||
MaxContentTokens: 50000,
|
||||
})
|
||||
if len(tools) != 2 {
|
||||
t.Fatalf("got %d tools, want 2 (search + fetch)", len(tools))
|
||||
// Handle username placeholder
|
||||
usernameValue := username
|
||||
if username == "" {
|
||||
usernameValue = "unknown"
|
||||
}
|
||||
systemMessage = strings.ReplaceAll(systemMessage, "{username}", usernameValue)
|
||||
|
||||
search := tools[0].OfWebSearchTool20260318
|
||||
if search == nil {
|
||||
t.Fatalf("tools[0] is not a web_search tool")
|
||||
}
|
||||
if !sameStrings(search.AllowedDomains, []string{"example.com/hc", "docs.example.com"}) {
|
||||
t.Errorf("search AllowedDomains = %v, want the path-scoped list unchanged", search.AllowedDomains)
|
||||
}
|
||||
if search.MaxUses.Value != 3 {
|
||||
t.Errorf("search MaxUses = %d, want 3", search.MaxUses.Value)
|
||||
// Handle firstname placeholder
|
||||
firstnameValue := firstName
|
||||
if firstName == "" {
|
||||
firstnameValue = "unknown"
|
||||
}
|
||||
systemMessage = strings.ReplaceAll(systemMessage, "{firstname}", firstnameValue)
|
||||
|
||||
fetch := tools[1].OfWebFetchTool20260318
|
||||
if fetch == nil {
|
||||
t.Fatalf("tools[1] is not a web_fetch tool")
|
||||
// Handle lastname placeholder
|
||||
lastnameValue := lastName
|
||||
if lastName == "" {
|
||||
lastnameValue = ""
|
||||
}
|
||||
if !sameStrings(fetch.AllowedDomains, []string{"example.com", "docs.example.com"}) {
|
||||
t.Errorf("fetch AllowedDomains = %v, want host-only [example.com docs.example.com]", fetch.AllowedDomains)
|
||||
}
|
||||
if fetch.MaxContentTokens.Value != 50000 {
|
||||
t.Errorf("fetch MaxContentTokens = %d, want 50000", fetch.MaxContentTokens.Value)
|
||||
}
|
||||
if fetch.Citations.Enabled.Value != true {
|
||||
t.Error("fetch citations not enabled")
|
||||
}
|
||||
})
|
||||
systemMessage = strings.ReplaceAll(systemMessage, "{lastname}", lastnameValue)
|
||||
|
||||
t.Run("fetch hosts are deduped", func(t *testing.T) {
|
||||
tools := webSearchTools(&WebSearchConfig{
|
||||
AllowedDomains: []string{"a.com/x", "a.com/y", "b.com"},
|
||||
Fetch: true,
|
||||
})
|
||||
fetch := tools[1].OfWebFetchTool20260318
|
||||
if fetch == nil {
|
||||
t.Fatalf("tools[1] is not a web_fetch tool")
|
||||
// Handle language code placeholder
|
||||
langValue := languageCode
|
||||
if languageCode == "" {
|
||||
langValue = "en"
|
||||
}
|
||||
if !sameStrings(fetch.AllowedDomains, []string{"a.com", "b.com"}) {
|
||||
t.Errorf("fetch AllowedDomains = %v, want deduped [a.com b.com]", fetch.AllowedDomains)
|
||||
}
|
||||
})
|
||||
systemMessage = strings.ReplaceAll(systemMessage, "{language}", langValue)
|
||||
|
||||
t.Run("social host is search-only via fetch_allowed_domains", func(t *testing.T) {
|
||||
tools := webSearchTools(&WebSearchConfig{
|
||||
AllowedDomains: []string{"helpshift.example/hc", "x.com/thatskygame"},
|
||||
FetchAllowedDomains: []string{"helpshift.example"},
|
||||
Fetch: true,
|
||||
})
|
||||
search := tools[0].OfWebSearchTool20260318
|
||||
if search == nil {
|
||||
t.Fatalf("tools[0] is not a web_search tool")
|
||||
}
|
||||
var searchHasSocial bool
|
||||
for _, d := range search.AllowedDomains {
|
||||
if d == "x.com/thatskygame" {
|
||||
searchHasSocial = true
|
||||
}
|
||||
}
|
||||
if !searchHasSocial {
|
||||
t.Errorf("search AllowedDomains = %v, want it to include x.com/thatskygame", search.AllowedDomains)
|
||||
// Handle premium status
|
||||
premiumStatus := "regular user"
|
||||
if isPremium {
|
||||
premiumStatus = "premium user"
|
||||
}
|
||||
systemMessage = strings.ReplaceAll(systemMessage, "{premium_status}", premiumStatus)
|
||||
|
||||
fetch := tools[1].OfWebFetchTool20260318
|
||||
if fetch == nil {
|
||||
t.Fatalf("tools[1] is not a web_fetch tool")
|
||||
// Handle time awareness
|
||||
timeObj := time.Unix(int64(messageTime), 0)
|
||||
hour := timeObj.Hour()
|
||||
var timeContext string
|
||||
if hour >= 5 && hour < 12 {
|
||||
timeContext = "morning"
|
||||
} else if hour >= 12 && hour < 18 {
|
||||
timeContext = "afternoon"
|
||||
} else if hour >= 18 && hour < 22 {
|
||||
timeContext = "evening"
|
||||
} else {
|
||||
timeContext = "night"
|
||||
}
|
||||
if !sameStrings(fetch.AllowedDomains, []string{"helpshift.example"}) {
|
||||
t.Errorf("fetch AllowedDomains = %v, want only [helpshift.example]", fetch.AllowedDomains)
|
||||
}
|
||||
for _, d := range fetch.AllowedDomains {
|
||||
if strings.HasPrefix(d, "x.com") {
|
||||
t.Errorf("fetch AllowedDomains leaked the social host: %v", fetch.AllowedDomains)
|
||||
}
|
||||
}
|
||||
})
|
||||
systemMessage = strings.ReplaceAll(systemMessage, "{time_context}", timeContext)
|
||||
|
||||
t.Run("wire shape carries allowed_domains", func(t *testing.T) {
|
||||
tools := webSearchTools(&WebSearchConfig{
|
||||
AllowedDomains: []string{"thatgamecompany.helpshift.com/hc"},
|
||||
MaxUses: 2,
|
||||
Fetch: true,
|
||||
})
|
||||
raw, err := json.Marshal(tools)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal: %v", err)
|
||||
// Check that all placeholders were replaced correctly
|
||||
if !strings.Contains(systemMessage, "username 'testuser'") {
|
||||
t.Errorf("Username not replaced correctly, got: %s", systemMessage)
|
||||
}
|
||||
body := string(raw)
|
||||
for _, want := range []string{
|
||||
"web_search_20260318",
|
||||
"web_fetch_20260318",
|
||||
"thatgamecompany.helpshift.com/hc",
|
||||
"allowed_domains",
|
||||
} {
|
||||
if !strings.Contains(body, want) {
|
||||
t.Errorf("wire body missing %q:\n%s", want, body)
|
||||
if !strings.Contains(systemMessage, "display name 'Test User'") {
|
||||
t.Errorf("Display name not replaced correctly, got: %s", systemMessage)
|
||||
}
|
||||
if !strings.Contains(systemMessage, "language preference: 'de'") {
|
||||
t.Errorf("Language preference not replaced correctly, got: %s", systemMessage)
|
||||
}
|
||||
})
|
||||
if !strings.Contains(systemMessage, "User is a premium user") {
|
||||
t.Errorf("Premium status not replaced correctly, got: %s", systemMessage)
|
||||
}
|
||||
|
||||
func sameStrings(got, want []string) bool {
|
||||
if len(got) != len(want) {
|
||||
return false
|
||||
}
|
||||
for i := range got {
|
||||
if got[i] != want[i] {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func TestFetchHosts(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
in []string
|
||||
want []string
|
||||
}{
|
||||
{"nil in nil out", nil, nil},
|
||||
{"empty in nil out", []string{}, nil},
|
||||
{"host passthrough", []string{"example.com"}, []string{"example.com"}},
|
||||
{"strip path", []string{"example.com/hc/en"}, []string{"example.com"}},
|
||||
{"dedup after strip", []string{"a.com/x", "a.com/y"}, []string{"a.com"}},
|
||||
{"preserve order and subdomains", []string{"docs.example.com/a", "example.com"}, []string{"docs.example.com", "example.com"}},
|
||||
{"drop empty leading slash", []string{"/oops", "ok.com"}, []string{"ok.com"}},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
if got := fetchHosts(tc.in); !sameStrings(got, tc.want) {
|
||||
t.Errorf("fetchHosts(%v) = %v, want %v", tc.in, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestEmptyStreamError(t *testing.T) {
|
||||
err := emptyStreamError("max_tokens", "", 3900, 4000)
|
||||
for _, want := range []string{"output budget exhausted", "3900", "4000"} {
|
||||
if !strings.Contains(err.Error(), want) {
|
||||
t.Errorf("max_tokens case: %q missing %q", err.Error(), want)
|
||||
}
|
||||
}
|
||||
if got := emptyStreamError("refusal", "cyber", 0, 1000).Error(); got != "the model declined the request (refusal, category cyber)" {
|
||||
t.Errorf("refusal case = %q", got)
|
||||
}
|
||||
if got := emptyStreamError("refusal", "", 0, 1000).Error(); got != "the model declined the request (refusal)" {
|
||||
t.Errorf("uncategorized refusal case = %q", got)
|
||||
}
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDynamicFilteringEnabled(t *testing.T) {
|
||||
off := false
|
||||
on := true
|
||||
cases := []struct {
|
||||
name string
|
||||
cfg *WebSearchConfig
|
||||
want bool
|
||||
}{
|
||||
{"nil config defaults on", nil, true},
|
||||
{"unset defaults on", &WebSearchConfig{}, true},
|
||||
{"explicit true", &WebSearchConfig{DynamicFiltering: &on}, true},
|
||||
{"explicit false", &WebSearchConfig{DynamicFiltering: &off}, false},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
if got := tc.cfg.DynamicFilteringEnabled(); got != tc.want {
|
||||
t.Errorf("%s: got %v, want %v", tc.name, got, tc.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Models older than Claude 4.6 reject the 2026 web tools unless the tool is
|
||||
// restricted to direct calls, so the opt-out must reach both tools.
|
||||
func TestWebSearchTools_AllowedCallers(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
off := false
|
||||
|
||||
tools := webSearchTools(&WebSearchConfig{Fetch: true})
|
||||
if got := tools[0].OfWebSearchTool20260318.AllowedCallers; got != nil {
|
||||
t.Errorf("default should omit allowed_callers, got %v", got)
|
||||
}
|
||||
if got := tools[1].OfWebFetchTool20260318.AllowedCallers; got != nil {
|
||||
t.Errorf("default should omit allowed_callers on fetch, got %v", got)
|
||||
}
|
||||
|
||||
tools = webSearchTools(&WebSearchConfig{Fetch: true, DynamicFiltering: &off})
|
||||
if got := tools[0].OfWebSearchTool20260318.AllowedCallers; len(got) != 1 || got[0] != "direct" {
|
||||
t.Errorf("search allowed_callers = %v, want [direct]", got)
|
||||
}
|
||||
if got := tools[1].OfWebFetchTool20260318.AllowedCallers; len(got) != 1 || got[0] != "direct" {
|
||||
t.Errorf("fetch allowed_callers = %v, want [direct]", got)
|
||||
if !strings.Contains(systemMessage, "It's currently afternoon in your timezone") {
|
||||
t.Errorf("Time context not replaced correctly, got: %s", systemMessage)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -8,17 +8,16 @@ import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/anthropics/anthropic-sdk-go/v2"
|
||||
"github.com/anthropics/anthropic-sdk-go/v2/option"
|
||||
"github.com/go-telegram/bot"
|
||||
"github.com/go-telegram/bot/models"
|
||||
"github.com/liushuangls/go-anthropic/v2"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type Bot struct {
|
||||
tgBot TelegramClient
|
||||
db *gorm.DB
|
||||
anthropicClient anthropic.Client
|
||||
anthropicClient *anthropic.Client
|
||||
chatMemories map[int64]*ChatMemory
|
||||
memorySize int
|
||||
chatMemoriesMu sync.RWMutex
|
||||
@@ -26,17 +25,10 @@ type Bot struct {
|
||||
userLimiters map[int64]*userLimiter
|
||||
userLimitersMu sync.RWMutex
|
||||
clock Clock
|
||||
botID uint
|
||||
albumBuffers map[string]*pendingAlbum
|
||||
albumBuffersMu sync.Mutex
|
||||
intakeBuffers map[int64]*pendingIntake
|
||||
intakeBuffersMu sync.Mutex
|
||||
intakeSeq uint64
|
||||
turns map[int64]*chatTurn
|
||||
turnsMu sync.Mutex
|
||||
turnSeq uint64
|
||||
botID uint // Reference to BotModel.ID
|
||||
}
|
||||
|
||||
// Helper function to determine message type
|
||||
func messageType(msg *models.Message) string {
|
||||
if msg.Sticker != nil {
|
||||
return "sticker"
|
||||
@@ -44,11 +36,13 @@ func messageType(msg *models.Message) string {
|
||||
return "text"
|
||||
}
|
||||
|
||||
// NewBot initializes and returns a new Bot instance.
|
||||
func NewBot(db *gorm.DB, config BotConfig, clock Clock, tgClient TelegramClient) (*Bot, error) {
|
||||
// Retrieve or create Bot entry in the database
|
||||
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}
|
||||
botEntry = BotModel{Identifier: config.ID, Name: config.ID} // Customize as needed
|
||||
if err := db.Create(&botEntry).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -56,9 +50,11 @@ func NewBot(db *gorm.DB, config BotConfig, clock Clock, tgClient TelegramClient)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Ensure the owner exists in the Users table
|
||||
var owner User
|
||||
err = db.Where("telegram_id = ? AND bot_id = ?", config.OwnerTelegramID, botEntry.ID).First(&owner).Error
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
// Assign the "owner" role
|
||||
var ownerRole Role
|
||||
err := db.Where("name = ?", "owner").First(&ownerRole).Error
|
||||
if err != nil {
|
||||
@@ -68,12 +64,13 @@ func NewBot(db *gorm.DB, config BotConfig, clock Clock, tgClient TelegramClient)
|
||||
owner = User{
|
||||
BotID: botEntry.ID,
|
||||
TelegramID: config.OwnerTelegramID,
|
||||
Username: "",
|
||||
Username: "", // Initialize as empty; will be updated upon interaction
|
||||
RoleID: ownerRole.ID,
|
||||
IsOwner: true,
|
||||
}
|
||||
|
||||
if err := db.Create(&owner).Error; err != nil {
|
||||
// If unique constraint is violated, another owner already exists
|
||||
if strings.Contains(err.Error(), "unique index") {
|
||||
return nil, fmt.Errorf("an owner already exists for this bot")
|
||||
}
|
||||
@@ -83,13 +80,8 @@ func NewBot(db *gorm.DB, config BotConfig, clock Clock, tgClient TelegramClient)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// An explicitly passed credential is used as is, so an empty key would go
|
||||
// out as an empty bearer token; leave it out to use the SDK's default chain.
|
||||
var anthropicOpts []option.RequestOption
|
||||
if config.AnthropicAPIKey != "" {
|
||||
anthropicOpts = append(anthropicOpts, option.WithAPIKey(config.AnthropicAPIKey))
|
||||
}
|
||||
anthropicClient := anthropic.NewClient(anthropicOpts...)
|
||||
// Use the per-bot Anthropic API key
|
||||
anthropicClient := anthropic.NewClient(config.AnthropicAPIKey)
|
||||
|
||||
b := &Bot{
|
||||
db: db,
|
||||
@@ -99,11 +91,8 @@ func NewBot(db *gorm.DB, config BotConfig, clock Clock, tgClient TelegramClient)
|
||||
config: config,
|
||||
userLimiters: make(map[int64]*userLimiter),
|
||||
clock: clock,
|
||||
botID: botEntry.ID,
|
||||
botID: botEntry.ID, // Ensure BotModel has ID field
|
||||
tgBot: tgClient,
|
||||
albumBuffers: make(map[string]*pendingAlbum),
|
||||
intakeBuffers: make(map[int64]*pendingIntake),
|
||||
turns: make(map[int64]*chatTurn),
|
||||
}
|
||||
|
||||
if tgClient == nil {
|
||||
@@ -118,6 +107,7 @@ func NewBot(db *gorm.DB, config BotConfig, clock Clock, tgClient TelegramClient)
|
||||
return b, nil
|
||||
}
|
||||
|
||||
// Start begins the bot's operation.
|
||||
func (b *Bot) Start(ctx context.Context) {
|
||||
b.tgBot.Start(ctx)
|
||||
}
|
||||
@@ -127,6 +117,7 @@ func (b *Bot) getOrCreateUser(userID int64, username string, isOwner bool) (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) {
|
||||
// Check if an owner already exists for this bot
|
||||
if isOwner {
|
||||
var existingOwner User
|
||||
err := b.db.Where("bot_id = ? AND is_owner = ?", b.botID, true).First(&existingOwner).Error
|
||||
@@ -142,7 +133,7 @@ func (b *Bot) getOrCreateUser(userID int64, username string, isOwner bool) (User
|
||||
if isOwner {
|
||||
roleName = "owner"
|
||||
} else {
|
||||
roleName = "user"
|
||||
roleName = "user" // Assign "user" role to non-owner users
|
||||
}
|
||||
|
||||
err := b.db.Where("name = ?", roleName).First(&role).Error
|
||||
@@ -160,6 +151,7 @@ func (b *Bot) getOrCreateUser(userID int64, username string, isOwner bool) (User
|
||||
}
|
||||
|
||||
if err := b.db.Create(&user).Error; err != nil {
|
||||
// If unique constraint is violated, another owner already exists
|
||||
if strings.Contains(err.Error(), "unique index") {
|
||||
return User{}, fmt.Errorf("an owner already exists for this bot")
|
||||
}
|
||||
@@ -203,9 +195,10 @@ func (b *Bot) createMessage(chatID, userID int64, username, userRole, text strin
|
||||
return message
|
||||
}
|
||||
|
||||
// storeMessage stores a message in the database and updates its ID
|
||||
func (b *Bot) storeMessage(message *Message) error {
|
||||
message.BotID = b.botID
|
||||
return b.db.Create(message).Error
|
||||
message.BotID = b.botID // Associate the message with the correct bot
|
||||
return b.db.Create(message).Error // This will update the message with its new ID
|
||||
}
|
||||
|
||||
func (b *Bot) getOrCreateChatMemory(chatID int64) *ChatMemory {
|
||||
@@ -219,27 +212,25 @@ func (b *Bot) getOrCreateChatMemory(chatID int64) *ChatMemory {
|
||||
|
||||
chatMemory, exists = b.chatMemories[chatID]
|
||||
if !exists {
|
||||
// Check if this is a new chat by querying the database
|
||||
var count int64
|
||||
b.db.Model(&Message{}).Where("chat_id = ? AND bot_id = ?", chatID, b.botID).Count(&count)
|
||||
isNewChat := count == 0
|
||||
isNewChat := count == 0 // Truly new chat if no messages exist
|
||||
|
||||
var messages []Message
|
||||
if !isNewChat {
|
||||
// Fetch existing messages only if it's not a new chat
|
||||
err := b.db.Where("chat_id = ? AND bot_id = ?", chatID, b.botID).
|
||||
Order("timestamp desc").
|
||||
Order("timestamp asc").
|
||||
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]
|
||||
}
|
||||
messages = []Message{} // Initialize an empty slice on error
|
||||
}
|
||||
} else {
|
||||
messages = []Message{}
|
||||
messages = []Message{} // Ensure messages is initialized for new chats
|
||||
}
|
||||
|
||||
chatMemory = &ChatMemory{
|
||||
@@ -254,224 +245,71 @@ func (b *Bot) getOrCreateChatMemory(chatID int64) *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
|
||||
}
|
||||
}
|
||||
|
||||
// addMessageToChatMemory adds a new message to the chat memory, ensuring the memory size is maintained.
|
||||
func (b *Bot) addMessageToChatMemory(chatMemory *ChatMemory, message Message) {
|
||||
b.chatMemoriesMu.Lock()
|
||||
defer b.chatMemoriesMu.Unlock()
|
||||
|
||||
// Add the new message
|
||||
chatMemory.Messages = append(chatMemory.Messages, message)
|
||||
|
||||
// Maintain the memory size
|
||||
if len(chatMemory.Messages) > chatMemory.Size {
|
||||
chatMemory.Messages = chatMemory.Messages[len(chatMemory.Messages)-chatMemory.Size:]
|
||||
}
|
||||
}
|
||||
|
||||
// prepareContextMessages converts chat memory into API messages.
|
||||
func (b *Bot) prepareContextMessages(chatMemory *ChatMemory) []anthropic.BetaMessageParam {
|
||||
func (b *Bot) prepareContextMessages(chatMemory *ChatMemory) []anthropic.Message {
|
||||
b.chatMemoriesMu.RLock()
|
||||
defer b.chatMemoriesMu.RUnlock()
|
||||
return b.buildContextMessages(chatMemory)
|
||||
}
|
||||
|
||||
// snapshotForTurn builds the model context for turn and records on it the
|
||||
// newest user message that context includes, both read under one lock so they
|
||||
// cannot disagree. The recorded id is what lets a later flush tell whether the
|
||||
// running turn already answers it.
|
||||
func (b *Bot) snapshotForTurn(turn *chatTurn, chatMemory *ChatMemory) []anthropic.BetaMessageParam {
|
||||
b.chatMemoriesMu.RLock()
|
||||
contextMessages := b.buildContextMessages(chatMemory)
|
||||
covered := newestUserMessageID(chatMemory)
|
||||
b.chatMemoriesMu.RUnlock()
|
||||
|
||||
b.markCovered(turn, covered)
|
||||
return contextMessages
|
||||
}
|
||||
|
||||
// newestUserMessageID returns the highest database id among the user messages
|
||||
// in memory. Ids are monotonic, so "newer" and "higher" agree.
|
||||
func newestUserMessageID(chatMemory *ChatMemory) uint {
|
||||
var newest uint
|
||||
for _, msg := range chatMemory.Messages {
|
||||
if msg.IsUser && msg.ID > newest {
|
||||
newest = msg.ID
|
||||
}
|
||||
}
|
||||
return newest
|
||||
}
|
||||
|
||||
func (b *Bot) newestUserMessageIDInChat(chatID int64) uint {
|
||||
chatMemory := b.getOrCreateChatMemory(chatID)
|
||||
b.chatMemoriesMu.RLock()
|
||||
defer b.chatMemoriesMu.RUnlock()
|
||||
return newestUserMessageID(chatMemory)
|
||||
}
|
||||
|
||||
// buildContextMessages is prepareContextMessages without the lock.
|
||||
func (b *Bot) buildContextMessages(chatMemory *ChatMemory) []anthropic.BetaMessageParam {
|
||||
// Debug logging
|
||||
InfoLogger.Printf("Chat memory contains %d messages", len(chatMemory.Messages))
|
||||
if b.config.DebugScreening {
|
||||
for i, msg := range chatMemory.Messages {
|
||||
InfoLogger.Printf("Message %d: IsUser=%v, Text=%q Images=%d", i, msg.IsUser, msg.Text, len(msg.ImageFileIDs))
|
||||
}
|
||||
InfoLogger.Printf("Message %d: IsUser=%v, Text=%q", i, msg.IsUser, msg.Text)
|
||||
}
|
||||
|
||||
var contextMessages []anthropic.BetaMessageParam
|
||||
// Note: consecutive messages with the same role are permitted.
|
||||
// The Anthropic API automatically merges them into a single turn rather than
|
||||
// returning an error. This can happen after a /clear (which only deletes user
|
||||
// messages, leaving assistant messages in the DB) followed by a restart.
|
||||
// See: https://platform.claude.com/docs/en/api/messages
|
||||
var contextMessages []anthropic.Message
|
||||
for _, msg := range chatMemory.Messages {
|
||||
blocks := contentBlocksForMessage(msg)
|
||||
if len(blocks) == 0 {
|
||||
role := anthropic.RoleUser
|
||||
if !msg.IsUser {
|
||||
role = anthropic.RoleAssistant
|
||||
}
|
||||
|
||||
textContent := strings.TrimSpace(msg.Text)
|
||||
if textContent == "" {
|
||||
// Skip empty messages
|
||||
continue
|
||||
}
|
||||
var param anthropic.BetaMessageParam
|
||||
if msg.IsUser {
|
||||
param = anthropic.NewBetaUserMessage(blocks...)
|
||||
} else {
|
||||
param = anthropic.BetaMessageParam{OfAssistant: &anthropic.BetaAssistantMessageParam{
|
||||
Content: blocks,
|
||||
}}
|
||||
}
|
||||
contextMessages = append(contextMessages, param)
|
||||
}
|
||||
|
||||
if b.config.CacheHistoryEnabled() {
|
||||
markTrailingCacheBreakpoint(contextMessages)
|
||||
contextMessages = append(contextMessages, anthropic.Message{
|
||||
Role: role,
|
||||
Content: []anthropic.MessageContent{
|
||||
anthropic.NewTextMessageContent(textContent),
|
||||
},
|
||||
})
|
||||
}
|
||||
return contextMessages
|
||||
}
|
||||
|
||||
// markTrailingCacheBreakpoint puts a cache_control breakpoint on the final
|
||||
// content block of the conversation, so the next turn reads the whole prefix
|
||||
// from cache instead of reprocessing it. The system prompt keeps its own
|
||||
// breakpoint; tools and system render ahead of messages, so the two compose.
|
||||
//
|
||||
// Caveat worth knowing when reading [usage] lines: chat memory is a sliding
|
||||
// window. Once it is full, each new turn evicts the oldest message, which
|
||||
// changes the prefix and forces a miss. Until then, and for any chat shorter
|
||||
// than the window, this converts a full-price reread into a cache read.
|
||||
func markTrailingCacheBreakpoint(messages []anthropic.BetaMessageParam) {
|
||||
if len(messages) == 0 {
|
||||
return
|
||||
}
|
||||
var blocks []anthropic.BetaContentBlockParamUnion
|
||||
switch last := messages[len(messages)-1]; {
|
||||
case last.OfUser != nil:
|
||||
blocks = last.OfUser.Content
|
||||
case last.OfAssistant != nil:
|
||||
blocks = last.OfAssistant.Content
|
||||
}
|
||||
if len(blocks) == 0 {
|
||||
return
|
||||
func (b *Bot) isNewChat(chatID int64) bool {
|
||||
var count int64
|
||||
b.db.Model(&Message{}).Where("chat_id = ? AND bot_id = ?", chatID, b.botID).Count(&count)
|
||||
return count == 0 // Only consider a chat new if it has 0 messages
|
||||
}
|
||||
|
||||
switch last := &blocks[len(blocks)-1]; {
|
||||
case last.OfText != nil:
|
||||
last.OfText.CacheControl = anthropic.NewBetaCacheControlEphemeralParam()
|
||||
case last.OfImage != nil:
|
||||
last.OfImage.CacheControl = anthropic.NewBetaCacheControlEphemeralParam()
|
||||
}
|
||||
}
|
||||
|
||||
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 {
|
||||
func (b *Bot) isAdminOrOwner(userID int64) 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 {
|
||||
err := b.db.Preload("Role").Where("telegram_id = ? AND bot_id = ?", userID, b.botID).First(&user).Error
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
if user.IsOwner {
|
||||
return true
|
||||
}
|
||||
return roleHasScope(user.Role, scope)
|
||||
}
|
||||
|
||||
var publicBotCommands = []models.BotCommand{
|
||||
{Command: "stats", Description: "Get bot statistics. Usage: /stats or /stats user [user_id]"},
|
||||
{Command: "whoami", Description: "Get your user information"},
|
||||
{Command: "clear", Description: "Clear chat history (soft delete). Admins: /clear [user_id]"},
|
||||
}
|
||||
|
||||
var adminBotCommands = []models.BotCommand{
|
||||
{Command: "clear_hard", Description: "Clear chat history (permanently delete). Admins: /clear_hard [user_id]"},
|
||||
{Command: "set_model", Description: "Switch the AI model (admin/owner only). Usage: /set_model <model-id>"},
|
||||
}
|
||||
|
||||
func (b *Bot) registerAdminCommandsForUser(ctx context.Context, telegramID int64) {
|
||||
allCommands := make([]models.BotCommand, 0, len(publicBotCommands)+len(adminBotCommands))
|
||||
allCommands = append(allCommands, publicBotCommands...)
|
||||
allCommands = append(allCommands, adminBotCommands...)
|
||||
_, err := b.tgBot.SetMyCommands(ctx, &bot.SetMyCommandsParams{
|
||||
Commands: allCommands,
|
||||
Scope: &models.BotCommandScopeChat{ChatID: telegramID},
|
||||
})
|
||||
if err != nil {
|
||||
ErrorLogger.Printf("Failed to register admin commands for user %d: %v", telegramID, err)
|
||||
}
|
||||
}
|
||||
|
||||
func setElevatedCommands(tgBot TelegramClient, users []User) {
|
||||
allCommands := make([]models.BotCommand, 0, len(publicBotCommands)+len(adminBotCommands))
|
||||
allCommands = append(allCommands, publicBotCommands...)
|
||||
allCommands = append(allCommands, adminBotCommands...)
|
||||
for _, u := range users {
|
||||
if u.TelegramID == 0 {
|
||||
continue
|
||||
}
|
||||
if !u.IsOwner && !roleHasScope(u.Role, ScopeModelSet) {
|
||||
continue
|
||||
}
|
||||
_, err := tgBot.SetMyCommands(context.Background(), &bot.SetMyCommandsParams{
|
||||
Commands: allCommands,
|
||||
Scope: &models.BotCommandScopeChat{ChatID: u.TelegramID},
|
||||
})
|
||||
if err != nil {
|
||||
ErrorLogger.Printf("Warning: could not set admin commands for user %d: %v", u.TelegramID, err)
|
||||
}
|
||||
}
|
||||
return user.Role.Name == "admin" || user.Role.Name == "owner"
|
||||
}
|
||||
|
||||
func initTelegramBot(token string, b *Bot) (TelegramClient, error) {
|
||||
@@ -484,32 +322,47 @@ func initTelegramBot(token string, b *Bot) (TelegramClient, error) {
|
||||
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
|
||||
// Define bot commands
|
||||
commands := []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]",
|
||||
},
|
||||
{
|
||||
Command: "clear_hard",
|
||||
Description: "Clear chat history (permanently delete). Admins: /clear_hard [user_id]",
|
||||
},
|
||||
}
|
||||
|
||||
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)
|
||||
// Set bot commands
|
||||
_, err = tgBot.SetMyCommands(context.Background(), &bot.SetMyCommandsParams{
|
||||
Commands: commands,
|
||||
})
|
||||
if err != nil {
|
||||
ErrorLogger.Printf("Error setting bot commands: %v", err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return tgBot, nil
|
||||
}
|
||||
|
||||
func (b *Bot) sendResponse(ctx context.Context, chatID int64, text string, businessConnectionID string) error {
|
||||
// Pass the outgoing message through the centralized screen for storage and chat memory update
|
||||
_, err := b.screenOutgoingMessage(chatID, text)
|
||||
if err != nil {
|
||||
ErrorLogger.Printf("Error storing assistant message: %v", err)
|
||||
return err
|
||||
}
|
||||
|
||||
// Prepare message parameters
|
||||
params := &bot.SendMessageParams{
|
||||
ChatID: chatID,
|
||||
Text: text,
|
||||
@@ -519,6 +372,7 @@ func (b *Bot) sendResponse(ctx context.Context, chatID int64, text string, busin
|
||||
params.BusinessConnectionID = businessConnectionID
|
||||
}
|
||||
|
||||
// Send the message via Telegram client
|
||||
_, err = b.tgBot.SendMessage(ctx, params)
|
||||
if err != nil {
|
||||
ErrorLogger.Printf("[%s] Error sending message to chat %d with BusinessConnectionID %s: %v",
|
||||
@@ -528,23 +382,9 @@ func (b *Bot) sendResponse(ctx context.Context, chatID int64, text string, busin
|
||||
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
|
||||
}
|
||||
|
||||
// sendStats sends the bot statistics to the specified chat.
|
||||
func (b *Bot) sendStats(ctx context.Context, chatID int64, userID int64, targetUserID int64, businessConnectionID string) {
|
||||
// If targetUserID is 0, show global stats
|
||||
if targetUserID == 0 {
|
||||
totalUsers, totalMessages, err := b.getStats()
|
||||
if err != nil {
|
||||
@@ -555,6 +395,7 @@ func (b *Bot) sendStats(ctx context.Context, chatID int64, userID int64, targetU
|
||||
return
|
||||
}
|
||||
|
||||
// Do NOT manually escape hyphens here
|
||||
statsMessage := fmt.Sprintf(
|
||||
"📊 Bot Statistics:\n\n"+
|
||||
"- Total Users: %d\n"+
|
||||
@@ -563,44 +404,17 @@ func (b *Bot) sendStats(ctx context.Context, chatID int64, userID int64, targetU
|
||||
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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Send the response through the centralized screen
|
||||
if err := b.sendResponse(ctx, chatID, statsMessage, businessConnectionID); err != nil {
|
||||
ErrorLogger.Printf("Error sending stats message: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// If targetUserID is not 0, show user-specific stats
|
||||
// Check permissions if the user is trying to view someone else's stats
|
||||
if targetUserID != userID {
|
||||
if !b.hasScope(userID, ScopeStatsViewAny) {
|
||||
if !b.isAdminOrOwner(userID) {
|
||||
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)
|
||||
@@ -609,6 +423,7 @@ func (b *Bot) sendStats(ctx context.Context, chatID int64, userID int64, targetU
|
||||
}
|
||||
}
|
||||
|
||||
// Get user stats
|
||||
username, messagesIn, messagesOut, totalMessages, err := b.getUserStats(targetUserID)
|
||||
if err != nil {
|
||||
ErrorLogger.Printf("Error fetching user stats: %v\n", err)
|
||||
@@ -618,6 +433,7 @@ func (b *Bot) sendStats(ctx context.Context, chatID int64, userID int64, targetU
|
||||
return
|
||||
}
|
||||
|
||||
// Build the user stats message
|
||||
userInfo := fmt.Sprintf("@%s (ID: %d)", username, targetUserID)
|
||||
if username == "" {
|
||||
userInfo = fmt.Sprintf("User ID: %d", targetUserID)
|
||||
@@ -639,6 +455,7 @@ func (b *Bot) sendStats(ctx context.Context, chatID int64, userID int64, targetU
|
||||
}
|
||||
}
|
||||
|
||||
// getStats retrieves the total number of users and messages from the database.
|
||||
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 {
|
||||
@@ -653,30 +470,36 @@ func (b *Bot) getStats() (int64, int64, error) {
|
||||
return totalUsers, totalMessages, nil
|
||||
}
|
||||
|
||||
// getUserStats retrieves statistics for a specific user
|
||||
func (b *Bot) getUserStats(userID int64) (string, int64, int64, int64, error) {
|
||||
// Get user information from database
|
||||
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)
|
||||
}
|
||||
|
||||
// Count messages sent by the user (IN)
|
||||
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
|
||||
}
|
||||
|
||||
// Count responses to the user (OUT)
|
||||
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 = ?",
|
||||
if err := b.db.Model(&Message{}).Where("chat_id IN (SELECT DISTINCT chat_id FROM messages WHERE user_id = ? AND bot_id = ?) AND bot_id = ? AND is_user = ?",
|
||||
userID, b.botID, b.botID, false).Count(&messagesOut).Error; err != nil {
|
||||
return "", 0, 0, 0, err
|
||||
}
|
||||
|
||||
// Total messages is the sum
|
||||
totalMessages := messagesIn + messagesOut
|
||||
|
||||
return user.Username, messagesIn, messagesOut, totalMessages, nil
|
||||
}
|
||||
|
||||
// isOnlyEmojis checks if the string consists solely of emojis.
|
||||
func isOnlyEmojis(s string) bool {
|
||||
for _, r := range s {
|
||||
if !isEmoji(r) {
|
||||
@@ -686,12 +509,14 @@ func isOnlyEmojis(s string) bool {
|
||||
return true
|
||||
}
|
||||
|
||||
// isEmoji determines if a rune is an emoji.
|
||||
// This is a simplistic check and can be expanded based on requirements.
|
||||
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)
|
||||
return (r >= 0x1F600 && r <= 0x1F64F) || // Emoticons
|
||||
(r >= 0x1F300 && r <= 0x1F5FF) || // Misc Symbols and Pictographs
|
||||
(r >= 0x1F680 && r <= 0x1F6FF) || // Transport and Map
|
||||
(r >= 0x2600 && r <= 0x26FF) || // Misc symbols
|
||||
(r >= 0x2700 && r <= 0x27BF) // Dingbats
|
||||
}
|
||||
|
||||
func (b *Bot) sendWhoAmI(ctx context.Context, chatID int64, userID int64, username string, businessConnectionID string) {
|
||||
@@ -721,11 +546,13 @@ func (b *Bot) sendWhoAmI(ctx context.Context, chatID int64, userID int64, userna
|
||||
role.Name,
|
||||
)
|
||||
|
||||
// Send the response through the centralized screen
|
||||
if err := b.sendResponse(ctx, chatID, whoAmIMessage, businessConnectionID); err != nil {
|
||||
ErrorLogger.Printf("Error sending /whoami message: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// screenIncomingMessage centralizes all incoming message processing: storing messages and updating chat memory.
|
||||
func (b *Bot) screenIncomingMessage(message *models.Message) (Message, error) {
|
||||
if b.config.DebugScreening {
|
||||
start := time.Now()
|
||||
@@ -741,8 +568,9 @@ func (b *Bot) screenIncomingMessage(message *models.Message) (Message, error) {
|
||||
}()
|
||||
}
|
||||
|
||||
userRole := "user"
|
||||
userRole := string(anthropic.RoleUser)
|
||||
|
||||
// Determine message text based on message type
|
||||
messageText := message.Text
|
||||
if message.Sticker != nil {
|
||||
if message.Sticker.Emoji != "" {
|
||||
@@ -751,32 +579,34 @@ func (b *Bot) screenIncomingMessage(message *models.Message) (Message, error) {
|
||||
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)
|
||||
userMessage.TelegramMessageID = message.ID
|
||||
|
||||
// Handle sticker-specific details if present
|
||||
if message.Sticker != nil {
|
||||
userMessage.StickerFileID = message.Sticker.FileID
|
||||
userMessage.StickerEmoji = message.Sticker.Emoji
|
||||
userMessage.StickerEmoji = message.Sticker.Emoji // Store the sticker emoji
|
||||
if message.Sticker.Thumbnail != nil {
|
||||
userMessage.StickerPNGFile = message.Sticker.Thumbnail.FileID
|
||||
}
|
||||
}
|
||||
|
||||
// Get the chat memory before storing the message
|
||||
chatMemory := b.getOrCreateChatMemory(message.Chat.ID)
|
||||
|
||||
// Store the message and get its ID
|
||||
if err := b.storeMessage(&userMessage); err != nil {
|
||||
return Message{}, err
|
||||
}
|
||||
|
||||
// Add the message to the chat memory
|
||||
b.addMessageToChatMemory(chatMemory, userMessage)
|
||||
|
||||
return userMessage, nil
|
||||
}
|
||||
|
||||
// screenOutgoingMessage handles storing of outgoing messages and updating chat memory.
|
||||
// It also marks the most recent unanswered user message as answered.
|
||||
func (b *Bot) screenOutgoingMessage(chatID int64, response string) (Message, error) {
|
||||
if b.config.DebugScreening {
|
||||
start := time.Now()
|
||||
@@ -791,25 +621,27 @@ func (b *Bot) screenOutgoingMessage(chatID int64, response string) (Message, err
|
||||
}()
|
||||
}
|
||||
|
||||
assistantMessage := b.createMessage(chatID, 0, "", "assistant", response, false)
|
||||
// Create and store the assistant message
|
||||
assistantMessage := b.createMessage(chatID, 0, "", string(anthropic.RoleAssistant), response, false)
|
||||
if err := b.storeMessage(&assistantMessage); err != nil {
|
||||
return Message{}, err
|
||||
}
|
||||
|
||||
// Mark every outstanding user message in the chat, not just the newest one.
|
||||
// A coalesced turn answers the whole batch, so a single-row update would
|
||||
// leave the earlier messages permanently unanswered. This also drops an
|
||||
// UPDATE ... ORDER BY ... LIMIT, which stock SQLite builds do not support.
|
||||
// Find and mark the most recent unanswered user message as answered
|
||||
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 messages as answered: %v", err)
|
||||
ErrorLogger.Printf("Error marking user message as answered: %v", err)
|
||||
// Continue even if there's an error updating the user message
|
||||
}
|
||||
|
||||
// Update chat memory with the message that now has an ID
|
||||
chatMemory := b.getOrCreateChatMemory(chatID)
|
||||
b.addMessageToChatMemory(chatMemory, assistantMessage)
|
||||
|
||||
@@ -817,26 +649,25 @@ func (b *Bot) screenOutgoingMessage(chatID int64, response string) (Message, err
|
||||
}
|
||||
|
||||
func (b *Bot) promoteUserToAdmin(promoterID, userToPromoteID int64) error {
|
||||
if !b.hasScope(promoterID, ScopeUserPromote) {
|
||||
// Check if the promoter is an owner or admin
|
||||
if !b.isAdminOrOwner(promoterID) {
|
||||
return errors.New("only admins or owners can promote users to admin")
|
||||
}
|
||||
|
||||
// Get the user to promote
|
||||
userToPromote, err := b.getOrCreateUser(userToPromoteID, "", false)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Get the admin role
|
||||
var adminRole Role
|
||||
if err := b.db.Where("name = ?", "admin").First(&adminRole).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Update the user's role
|
||||
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
|
||||
return b.db.Save(&userToPromote).Error
|
||||
}
|
||||
|
||||
-105
@@ -1,105 +0,0 @@
|
||||
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")
|
||||
}
|
||||
@@ -1,25 +1,32 @@
|
||||
// clock.go
|
||||
package main
|
||||
|
||||
import "time"
|
||||
|
||||
// Clock is an interface to abstract time-related functions.
|
||||
type Clock interface {
|
||||
Now() time.Time
|
||||
}
|
||||
|
||||
// RealClock implements Clock using the actual time.
|
||||
type RealClock struct{}
|
||||
|
||||
// Now returns the current local time.
|
||||
func (RealClock) Now() time.Time {
|
||||
return time.Now()
|
||||
}
|
||||
|
||||
// MockClock implements Clock for testing purposes.
|
||||
type MockClock struct {
|
||||
currentTime time.Time
|
||||
}
|
||||
|
||||
// Now returns the mocked current time.
|
||||
func (mc *MockClock) Now() time.Time {
|
||||
return mc.currentTime
|
||||
}
|
||||
|
||||
// Advance moves the current time forward by the specified duration.
|
||||
func (mc *MockClock) Advance(d time.Duration) {
|
||||
mc.currentTime = mc.currentTime.Add(d)
|
||||
}
|
||||
|
||||
@@ -5,69 +5,11 @@ import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/liushuangls/go-anthropic/v2"
|
||||
)
|
||||
|
||||
type MCPServer struct {
|
||||
Name string `json:"name"`
|
||||
URL string `json:"url"`
|
||||
AuthorizationToken string `json:"authorization_token,omitempty"`
|
||||
AllowedTools []string `json:"allowed_tools,omitempty"`
|
||||
}
|
||||
|
||||
type WebSearchConfig struct {
|
||||
AllowedDomains []string `json:"allowed_domains,omitempty"`
|
||||
BlockedDomains []string `json:"blocked_domains,omitempty"`
|
||||
FetchAllowedDomains []string `json:"fetch_allowed_domains,omitempty"`
|
||||
DynamicFiltering *bool `json:"dynamic_filtering,omitempty"`
|
||||
MaxUses int `json:"max_uses,omitempty"`
|
||||
Fetch bool `json:"fetch,omitempty"`
|
||||
MaxContentTokens int `json:"max_content_tokens,omitempty"`
|
||||
}
|
||||
|
||||
// validEfforts are the output_config.effort levels the API accepts. Models
|
||||
// older than Claude 4.6 (Haiku 4.5, Sonnet 4.5) reject the parameter entirely.
|
||||
var validEfforts = []string{"low", "medium", "high", "xhigh", "max"}
|
||||
|
||||
const (
|
||||
ThinkingModeAdaptive = "adaptive"
|
||||
ThinkingModeDisabled = "disabled"
|
||||
ThinkingDisplaySummarized = "summarized"
|
||||
ThinkingDisplayOmitted = "omitted"
|
||||
)
|
||||
|
||||
// maxDebounceMs bounds debounce_ms. Beyond this the bot reads as unresponsive
|
||||
// rather than deliberate, and the coalesced turn drifts far enough from the
|
||||
// user's last message that the reply feels stale.
|
||||
const maxDebounceMs = 30000
|
||||
|
||||
// DebounceWindow is the quiet period an intake buffer waits before dispatching a
|
||||
// coalesced turn. Zero disables debouncing entirely (the default), matching the
|
||||
// opt-in behavior of comparable gateways.
|
||||
func (c *BotConfig) DebounceWindow() time.Duration {
|
||||
if c.DebounceMs <= 0 {
|
||||
return 0
|
||||
}
|
||||
return time.Duration(c.DebounceMs) * time.Millisecond
|
||||
}
|
||||
|
||||
// DynamicFilteringEnabled reports whether the web tools may run inside code
|
||||
// execution to filter results before they reach the context window. Defaults to
|
||||
// true; set "dynamic_filtering": false for models older than Claude 4.6, which
|
||||
// reject the tool unless it is called directly.
|
||||
func (c *WebSearchConfig) DynamicFilteringEnabled() bool {
|
||||
return c == nil || c.DynamicFiltering == nil || *c.DynamicFiltering
|
||||
}
|
||||
|
||||
// CacheHistoryEnabled reports whether a cache_control breakpoint should be placed
|
||||
// on the trailing conversation block in addition to the system prompt. Defaults
|
||||
// to true; set "cache_history": false to opt out.
|
||||
func (c *BotConfig) CacheHistoryEnabled() bool {
|
||||
return c.CacheHistory == nil || *c.CacheHistory
|
||||
}
|
||||
|
||||
type BotConfig struct {
|
||||
ID string `json:"id"`
|
||||
TelegramToken string `json:"telegram_token"`
|
||||
@@ -75,32 +17,38 @@ type BotConfig struct {
|
||||
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"` // never sent (rejected since API 2026-10-21); parsed only to flag stale configs
|
||||
MaxTokens int `json:"max_tokens,omitempty"`
|
||||
Thinking string `json:"thinking,omitempty"`
|
||||
ThinkingDisplay string `json:"thinking_display,omitempty"`
|
||||
Effort string `json:"effort,omitempty"`
|
||||
DebounceMs int `json:"debounce_ms,omitempty"`
|
||||
StreamDrafts bool `json:"stream_drafts,omitempty"`
|
||||
CacheHistory *bool `json:"cache_history,omitempty"`
|
||||
Model anthropic.Model `json:"model"`
|
||||
Temperature *float32 `json:"temperature,omitempty"` // Controls creativity vs determinism (0.0-1.0)
|
||||
SystemPrompts map[string]string `json:"system_prompts"`
|
||||
Active bool `json:"active"`
|
||||
OwnerTelegramID int64 `json:"owner_telegram_id"`
|
||||
AnthropicAPIKey string `json:"anthropic_api_key"`
|
||||
ElevenLabsAPIKey string `json:"elevenlabs_api_key"`
|
||||
ElevenLabsVoiceID string `json:"elevenlabs_voice_id"`
|
||||
ElevenLabsModel string `json:"elevenlabs_model"`
|
||||
DebugScreening bool `json:"debug_screening"`
|
||||
MCPServers []MCPServer `json:"mcp_servers,omitempty"`
|
||||
WebSearch *WebSearchConfig `json:"web_search,omitempty"`
|
||||
ConfigFilePath string `json:"-"`
|
||||
DebugScreening bool `json:"debug_screening"` // Enable detailed screening logs
|
||||
}
|
||||
|
||||
// Custom unmarshalling to handle anthropic.Model
|
||||
func (c *BotConfig) UnmarshalJSON(data []byte) error {
|
||||
type Alias BotConfig
|
||||
aux := &struct {
|
||||
Model string `json:"model"`
|
||||
*Alias
|
||||
}{
|
||||
Alias: (*Alias)(c),
|
||||
}
|
||||
if err := json.Unmarshal(data, &aux); err != nil {
|
||||
return err
|
||||
}
|
||||
c.Model = anthropic.Model(aux.Model)
|
||||
return nil
|
||||
}
|
||||
|
||||
// validateConfigPath ensures the file path is within the allowed directory
|
||||
func validateConfigPath(configDir, filename string) (string, error) {
|
||||
// Clean the paths to remove any . or .. components
|
||||
configDir = filepath.Clean(configDir)
|
||||
filename = filepath.Clean(filename)
|
||||
|
||||
// Get absolute paths
|
||||
absConfigDir, err := filepath.Abs(configDir)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to get absolute path for config directory: %w", err)
|
||||
@@ -112,11 +60,13 @@ func validateConfigPath(configDir, filename string) (string, error) {
|
||||
return "", fmt.Errorf("failed to get absolute path for config file: %w", err)
|
||||
}
|
||||
|
||||
// Use filepath.Rel to check if the path is within the config directory
|
||||
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")
|
||||
}
|
||||
|
||||
// Verify file extension
|
||||
if filepath.Ext(absPath) != ".json" {
|
||||
return "", fmt.Errorf("invalid file extension: must be .json")
|
||||
}
|
||||
@@ -158,9 +108,6 @@ func loadAllConfigs(dir string) ([]BotConfig, error) {
|
||||
continue
|
||||
}
|
||||
|
||||
logConfigAdvisories(&config)
|
||||
|
||||
config.ConfigFilePath = validPath
|
||||
configs = append(configs, config)
|
||||
}
|
||||
}
|
||||
@@ -172,61 +119,6 @@ func loadAllConfigs(dir string) ([]BotConfig, error) {
|
||||
return configs, nil
|
||||
}
|
||||
|
||||
// logConfigAdvisories emits non-fatal boot-time notes about settings that are
|
||||
// valid but likely to surprise: silently ineffective, or costlier than intended.
|
||||
func logConfigAdvisories(config *BotConfig) {
|
||||
if config.Temperature != nil {
|
||||
InfoLogger.Printf("[%s] temperature=%g is ignored: the Claude API no longer accepts sampling parameters; remove the key",
|
||||
config.ID, *config.Temperature)
|
||||
}
|
||||
|
||||
if config.Thinking == ThinkingModeAdaptive && config.MaxTokens > 0 && config.MaxTokens < 4000 {
|
||||
InfoLogger.Printf("[%s] thinking=adaptive with max_tokens=%d: thinking tokens count toward max_tokens; consider >= 4000",
|
||||
config.ID, config.MaxTokens)
|
||||
}
|
||||
|
||||
if config.DebounceMs > 0 {
|
||||
InfoLogger.Printf("[%s] intake debounce enabled: coalescing rapid text messages over a %dms quiet window",
|
||||
config.ID, config.DebounceMs)
|
||||
} else {
|
||||
InfoLogger.Printf("[%s] intake debounce disabled: every message dispatches its own turn (set debounce_ms to coalesce rapid follow-ups)",
|
||||
config.ID)
|
||||
}
|
||||
|
||||
if config.Effort != "" {
|
||||
InfoLogger.Printf("[%s] effort=%s: caps reasoning depth and token spend; the API rejects it on models older than Claude 4.6",
|
||||
config.ID, config.Effort)
|
||||
}
|
||||
|
||||
if config.StreamDrafts {
|
||||
InfoLogger.Printf("[%s] draft streaming enabled: private chats get one growing draft with a Stop button; business and group chats keep per-block messages",
|
||||
config.ID)
|
||||
}
|
||||
|
||||
if ws := config.WebSearch; ws != nil && len(ws.AllowedDomains) == 0 && len(ws.BlockedDomains) == 0 {
|
||||
InfoLogger.Printf("[%s] web_search enabled with no allowed_domains/blocked_domains: the model may search the open web",
|
||||
config.ID)
|
||||
}
|
||||
|
||||
if ws := config.WebSearch; ws != nil && !ws.DynamicFilteringEnabled() {
|
||||
InfoLogger.Printf("[%s] web_search dynamic_filtering disabled: results enter context unfiltered (set it true on Claude 4.6+ models to cut tokens)",
|
||||
config.ID)
|
||||
}
|
||||
|
||||
if ws := config.WebSearch; ws != nil && len(ws.FetchAllowedDomains) > 0 {
|
||||
if !ws.Fetch {
|
||||
InfoLogger.Printf("[%s] web_search.fetch_allowed_domains is set but fetch is disabled: it has no effect",
|
||||
config.ID)
|
||||
}
|
||||
for _, d := range ws.FetchAllowedDomains {
|
||||
if strings.Contains(d, "/") {
|
||||
InfoLogger.Printf("[%s] web_search.fetch_allowed_domains entry %q includes a path: web_fetch matches host-only, so the whole host is fetchable",
|
||||
config.ID, d)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func validateConfig(config *BotConfig, ids, tokens map[string]bool) error {
|
||||
if config.ID == "" {
|
||||
return fmt.Errorf("missing 'id' field")
|
||||
@@ -248,54 +140,6 @@ func validateConfig(config *BotConfig, ids, tokens map[string]bool) error {
|
||||
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.Effort != "" && !slices.Contains(validEfforts, config.Effort) {
|
||||
return fmt.Errorf("invalid 'effort' value %q: must be one of %s (or omitted)",
|
||||
config.Effort, strings.Join(validEfforts, ", "))
|
||||
}
|
||||
|
||||
if config.MaxTokens < 0 {
|
||||
return fmt.Errorf("'max_tokens' must be greater than 0 when set")
|
||||
}
|
||||
|
||||
if config.DebounceMs < 0 {
|
||||
return fmt.Errorf("'debounce_ms' must be greater than 0 when set")
|
||||
}
|
||||
if config.DebounceMs > maxDebounceMs {
|
||||
return fmt.Errorf("'debounce_ms' of %d exceeds the maximum of %d (Telegram drops long-idle updates and users read silence as failure)",
|
||||
config.DebounceMs, maxDebounceMs)
|
||||
}
|
||||
|
||||
if ws := config.WebSearch; ws != nil {
|
||||
if len(ws.AllowedDomains) > 0 && len(ws.BlockedDomains) > 0 {
|
||||
return fmt.Errorf("'web_search' cannot set both allowed_domains and blocked_domains (the API rejects that)")
|
||||
}
|
||||
if ws.MaxUses < 0 {
|
||||
return fmt.Errorf("'web_search.max_uses' must be greater than 0 when set")
|
||||
}
|
||||
if ws.MaxContentTokens < 0 {
|
||||
return fmt.Errorf("'web_search.max_content_tokens' must be greater than 0 when set")
|
||||
}
|
||||
}
|
||||
|
||||
if config.MessagePerHour <= 0 {
|
||||
return fmt.Errorf("'messages_per_hour' must be greater than 0")
|
||||
}
|
||||
@@ -309,6 +153,7 @@ func validateConfig(config *BotConfig, ids, tokens map[string]bool) error {
|
||||
|
||||
func loadConfig(filename string) (BotConfig, error) {
|
||||
var config BotConfig
|
||||
// Use filepath.Clean before opening the file
|
||||
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)
|
||||
@@ -327,12 +172,15 @@ func loadConfig(filename string) (BotConfig, error) {
|
||||
return config, nil
|
||||
}
|
||||
|
||||
// Reload reloads the BotConfig from the specified filename within the given config directory
|
||||
func (c *BotConfig) Reload(configDir, filename string) error {
|
||||
// Validate the config path
|
||||
validPath, err := validateConfigPath(configDir, filename)
|
||||
if err != nil {
|
||||
return fmt.Errorf("invalid config path: %w", err)
|
||||
}
|
||||
|
||||
// Use filepath.Clean before opening the file
|
||||
cleanPath := filepath.Clean(validPath)
|
||||
file, err := os.OpenFile(cleanPath, os.O_RDONLY, 0)
|
||||
if err != nil {
|
||||
@@ -349,35 +197,6 @@ func (c *BotConfig) Reload(configDir, filename string) error {
|
||||
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
|
||||
c.Model = anthropic.Model(c.Model)
|
||||
return nil
|
||||
}
|
||||
|
||||
+7
-9
@@ -4,20 +4,18 @@
|
||||
"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-5-5",
|
||||
"max_tokens": 1000,
|
||||
"effort": "low",
|
||||
"debounce_ms": 2500,
|
||||
"model": "claude-3-5-haiku-latest",
|
||||
"temperature": 0.7,
|
||||
"debug_screening": false,
|
||||
"system_prompts": {
|
||||
"custom_instructions": "You are Atom, a helpful assistant texting through a limited Telegram interface with a 15-word maximum. Write like texting a friend - use shorthand, skip grammar, use slang/abbreviations. The system cuts off anything longer than 15 words.\n\n- Address the user by their first name, and reply in their preferred language (both are in the conversation context).\n- Use time-appropriate greetings based on the user's local time of day.\n- If a user asks about buying apples, inform them that we don't sell apples.\n- When asked for a joke, tell a clean, family-friendly joke about programming or technology.\n- If someone inquires about our services, explain that we offer AI-powered chatbot solutions.\n- For any questions about pricing, direct users to contact our sales team at [email protected].\n- If asked about your capabilities, be honest about what you can and cannot do.\nAlways maintain a friendly and professional tone.",
|
||||
"respond_with_emojis": "The user's message contains only emoji. Reply using only emoji."
|
||||
"default": "You are a helpful assistant.",
|
||||
"custom_instructions": "You are texting through a limited Telegram interface with 15-word maximum. Write like texting a friend - use shorthand, skip grammar, use slang/abbreviations. System cuts off anything longer than 15 words.\n\n- Your name is Atom.\n- The user you're talking to has username '{username}' and display name '{firstname} {lastname}'.\n- User's language preference: '{language}'\n- User is a {premium_status}\n- It's currently {time_context} in your timezone. Use appropriate time-based greetings and address the user by name.\n- If a user asks about buying apples, inform them that we don't sell apples.\n- When asked for a joke, tell a clean, family-friendly joke about programming or technology.\n- If someone inquires about our services, explain that we offer AI-powered chatbot solutions.\n- For any questions about pricing, direct users to contact our sales team at [email protected].\n- If asked about your capabilities, be honest about what you can and cannot do.\nAlways maintain a friendly and professional tone.",
|
||||
"continue_conversation": "Continuing our conversation. Remember previous context if relevant.",
|
||||
"avoid_sensitive": "Avoid discussing sensitive topics or providing harmful information.",
|
||||
"respond_with_emojis": "Since the user sent only emojis, respond using emojis only."
|
||||
}
|
||||
}
|
||||
|
||||
+40
-244
@@ -2,18 +2,21 @@ package main
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/liushuangls/go-anthropic/v2"
|
||||
)
|
||||
|
||||
// Set up loggers
|
||||
func TestMain(m *testing.M) {
|
||||
initLoggers()
|
||||
os.Exit(m.Run())
|
||||
}
|
||||
|
||||
// TestBotConfig_UnmarshalJSON tests the custom unmarshalling of BotConfig
|
||||
func TestBotConfig_UnmarshalJSON(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
jsonData := `{
|
||||
"id": "bot123",
|
||||
@@ -35,7 +38,7 @@ func TestBotConfig_UnmarshalJSON(t *testing.T) { //NOSONAR go:S100 -- underscore
|
||||
t.Fatalf("Failed to unmarshal JSON: %v", err)
|
||||
}
|
||||
|
||||
expectedModel := "claude-v1"
|
||||
expectedModel := anthropic.Model("claude-v1")
|
||||
if config.Model != expectedModel {
|
||||
t.Errorf("Expected model %s, got %s", expectedModel, config.Model)
|
||||
}
|
||||
@@ -45,8 +48,10 @@ func TestBotConfig_UnmarshalJSON(t *testing.T) { //NOSONAR go:S100 -- underscore
|
||||
t.Errorf("Expected ID %s, got %s", expectedID, config.ID)
|
||||
}
|
||||
|
||||
// Add more field checks as necessary
|
||||
}
|
||||
|
||||
// TestValidateConfigPath tests the validateConfigPath function
|
||||
func TestValidateConfigPath(t *testing.T) {
|
||||
execDir, err := os.Getwd()
|
||||
if err != nil {
|
||||
@@ -91,6 +96,7 @@ func TestValidateConfigPath(t *testing.T) {
|
||||
},
|
||||
}
|
||||
|
||||
// Create a subdirectory for testing
|
||||
subDir := filepath.Join(execDir, "subdir")
|
||||
if err := os.MkdirAll(subDir, 0755); err != nil {
|
||||
t.Fatalf("Failed to create subdir: %v", err)
|
||||
@@ -116,7 +122,9 @@ func TestValidateConfigPath(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestLoadConfig tests the loadConfig function
|
||||
func TestLoadConfig(t *testing.T) {
|
||||
// Create a temporary directory
|
||||
tempDir, err := os.MkdirTemp("", "config_test")
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create temp dir: %v", err)
|
||||
@@ -127,6 +135,7 @@ func TestLoadConfig(t *testing.T) {
|
||||
}
|
||||
}()
|
||||
|
||||
// Valid config JSON
|
||||
validConfig := `{
|
||||
"id": "bot123",
|
||||
"telegram_token": "token123",
|
||||
@@ -142,6 +151,7 @@ func TestLoadConfig(t *testing.T) {
|
||||
"anthropic_api_key": "api_key_123"
|
||||
}`
|
||||
|
||||
// Invalid config JSON
|
||||
invalidConfig := `{
|
||||
"id": "bot123",
|
||||
"telegram_token": "token123",
|
||||
@@ -149,11 +159,13 @@ func TestLoadConfig(t *testing.T) {
|
||||
"model": "claude-v1"
|
||||
}`
|
||||
|
||||
// Write valid config file
|
||||
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)
|
||||
}
|
||||
|
||||
// Write invalid config file
|
||||
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)
|
||||
@@ -206,6 +218,7 @@ func TestLoadConfig(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestValidateConfig tests the validateConfig function
|
||||
func TestValidateConfig(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
@@ -338,7 +351,9 @@ func TestValidateConfig(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestLoadAllConfigs tests the loadAllConfigs function
|
||||
func TestLoadAllConfigs(t *testing.T) {
|
||||
// Create a temporary directory
|
||||
tempDir, err := os.MkdirTemp("", "load_all_configs_test")
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create temp dir: %v", err)
|
||||
@@ -351,7 +366,7 @@ func TestLoadAllConfigs(t *testing.T) {
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
setupFiles map[string]string
|
||||
setupFiles map[string]string // filename -> content
|
||||
expectConfigs int
|
||||
expectError bool
|
||||
expectErrorMsg string
|
||||
@@ -509,6 +524,7 @@ func TestLoadAllConfigs(t *testing.T) {
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
// Clear the tempDir before each test
|
||||
if err := os.RemoveAll(tempDir); err != nil {
|
||||
t.Fatalf("Failed to remove temp dir: %v", err)
|
||||
}
|
||||
@@ -516,6 +532,7 @@ func TestLoadAllConfigs(t *testing.T) {
|
||||
t.Fatalf("Failed to create temp dir: %v", err)
|
||||
}
|
||||
|
||||
// Write the test files directly
|
||||
for filename, content := range tt.setupFiles {
|
||||
err := os.WriteFile(filepath.Join(tempDir, filename), []byte(content), 0644)
|
||||
if err != nil {
|
||||
@@ -535,7 +552,9 @@ func TestLoadAllConfigs(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestBotConfig_Reload tests the Reload method of BotConfig
|
||||
func TestBotConfig_Reload(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
// Create a temporary directory
|
||||
tempDir, err := os.MkdirTemp("", "reload_test")
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create temp dir: %v", err)
|
||||
@@ -546,6 +565,7 @@ func TestBotConfig_Reload(t *testing.T) { //NOSONAR go:S100 -- underscore separa
|
||||
}
|
||||
}()
|
||||
|
||||
// Create initial config file
|
||||
config1 := `{
|
||||
"id": "bot123",
|
||||
"telegram_token": "token123",
|
||||
@@ -565,11 +585,13 @@ func TestBotConfig_Reload(t *testing.T) { //NOSONAR go:S100 -- underscore separa
|
||||
t.Fatalf("Failed to write initial config: %v", err)
|
||||
}
|
||||
|
||||
// Initialize BotConfig
|
||||
var config BotConfig
|
||||
if err := config.Reload(tempDir, "config.json"); err != nil {
|
||||
t.Fatalf("Failed to reload config: %v", err)
|
||||
}
|
||||
|
||||
// Verify initial load
|
||||
if config.ID != "bot123" {
|
||||
t.Errorf("Expected ID 'bot123', got '%s'", config.ID)
|
||||
}
|
||||
@@ -577,6 +599,7 @@ func TestBotConfig_Reload(t *testing.T) { //NOSONAR go:S100 -- underscore separa
|
||||
t.Errorf("Expected Model 'claude-v1', got '%s'", config.Model)
|
||||
}
|
||||
|
||||
// Update config file
|
||||
config2 := `{
|
||||
"id": "bot123",
|
||||
"telegram_token": "token123_updated",
|
||||
@@ -595,10 +618,12 @@ func TestBotConfig_Reload(t *testing.T) { //NOSONAR go:S100 -- underscore separa
|
||||
t.Fatalf("Failed to write updated config: %v", err)
|
||||
}
|
||||
|
||||
// Reload config
|
||||
if err := config.Reload(tempDir, "config.json"); err != nil {
|
||||
t.Fatalf("Failed to reload updated config: %v", err)
|
||||
}
|
||||
|
||||
// Verify updated config
|
||||
if config.TelegramToken != "token123_updated" {
|
||||
t.Errorf("Expected TelegramToken 'token123_updated', got '%s'", config.TelegramToken)
|
||||
}
|
||||
@@ -613,6 +638,7 @@ func TestBotConfig_Reload(t *testing.T) { //NOSONAR go:S100 -- underscore separa
|
||||
}
|
||||
}
|
||||
|
||||
// TestBotConfig_UnmarshalJSON_Invalid tests unmarshalling with invalid model
|
||||
func TestBotConfig_UnmarshalJSON_Invalid(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
jsonData := `{
|
||||
"id": "bot123",
|
||||
@@ -640,11 +666,14 @@ func TestBotConfig_UnmarshalJSON_Invalid(t *testing.T) { //NOSONAR go:S100 -- un
|
||||
}
|
||||
}
|
||||
|
||||
// Helper function to check substring
|
||||
func contains(s, substr string) bool {
|
||||
return strings.Contains(s, substr)
|
||||
}
|
||||
|
||||
// TestTemperatureConfig tests that the temperature value is correctly loaded
|
||||
func TestTemperatureConfig(t *testing.T) {
|
||||
// Create a temporary directory
|
||||
tempDir, err := os.MkdirTemp("", "temperature_test")
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create temp dir: %v", err)
|
||||
@@ -655,6 +684,7 @@ func TestTemperatureConfig(t *testing.T) {
|
||||
}
|
||||
}()
|
||||
|
||||
// Create config with temperature
|
||||
configWithTemp := `{
|
||||
"id": "bot123",
|
||||
"telegram_token": "token123",
|
||||
@@ -670,6 +700,7 @@ func TestTemperatureConfig(t *testing.T) {
|
||||
"anthropic_api_key": "api_key_123"
|
||||
}`
|
||||
|
||||
// Create config without temperature
|
||||
configWithoutTemp := `{
|
||||
"id": "bot124",
|
||||
"telegram_token": "token124",
|
||||
@@ -684,6 +715,7 @@ func TestTemperatureConfig(t *testing.T) {
|
||||
"anthropic_api_key": "api_key_123"
|
||||
}`
|
||||
|
||||
// Write config files
|
||||
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)
|
||||
@@ -694,265 +726,29 @@ func TestTemperatureConfig(t *testing.T) {
|
||||
t.Fatalf("Failed to write config without temperature: %v", err)
|
||||
}
|
||||
|
||||
// Test loading config with temperature
|
||||
configWithTempObj, err := loadConfig(withTempPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to load config with temperature: %v", err)
|
||||
}
|
||||
|
||||
// Verify temperature is set correctly
|
||||
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)
|
||||
}
|
||||
|
||||
// Test loading config without temperature
|
||||
configWithoutTempObj, err := loadConfig(withoutTempPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to load config without temperature: %v", err)
|
||||
}
|
||||
|
||||
// Verify temperature is nil when not specified
|
||||
if configWithoutTempObj.Temperature != nil {
|
||||
t.Errorf("Expected Temperature to be nil, got %f", *configWithoutTempObj.Temperature)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBotConfig_PersistModel(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
tempDir, err := os.MkdirTemp("", "persist_model_test")
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create temp dir: %v", err)
|
||||
}
|
||||
defer func() {
|
||||
if err := os.RemoveAll(tempDir); err != nil {
|
||||
t.Errorf("Failed to remove temp directory: %v", err)
|
||||
}
|
||||
}()
|
||||
|
||||
initialJSON := `{
|
||||
"id": "bot1",
|
||||
"telegram_token": "token1",
|
||||
"model": "claude-v1",
|
||||
"messages_per_hour": 10,
|
||||
"messages_per_day": 100
|
||||
}`
|
||||
configPath := filepath.Join(tempDir, "config.json")
|
||||
if err := os.WriteFile(configPath, []byte(initialJSON), 0600); err != nil {
|
||||
t.Fatalf("Failed to write config file: %v", err)
|
||||
}
|
||||
|
||||
config := BotConfig{
|
||||
ID: "bot1",
|
||||
Model: "claude-v1",
|
||||
ConfigFilePath: configPath,
|
||||
}
|
||||
|
||||
if err := config.PersistModel("claude-sonnet-4-6"); err != nil {
|
||||
t.Fatalf("PersistModel() unexpected error: %v", err)
|
||||
}
|
||||
|
||||
if string(config.Model) != "claude-sonnet-4-6" {
|
||||
t.Errorf("in-memory model: got %q, want %q", config.Model, "claude-sonnet-4-6")
|
||||
}
|
||||
|
||||
data, err := os.ReadFile(configPath)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to read updated config file: %v", err)
|
||||
}
|
||||
var raw map[string]any
|
||||
if err := json.Unmarshal(data, &raw); err != nil {
|
||||
t.Fatalf("Failed to unmarshal updated config: %v", err)
|
||||
}
|
||||
if raw["model"] != "claude-sonnet-4-6" {
|
||||
t.Errorf("on-disk model: got %v, want %q", raw["model"], "claude-sonnet-4-6")
|
||||
}
|
||||
if raw["id"] != "bot1" {
|
||||
t.Errorf("on-disk id should be preserved: got %v, want %q", raw["id"], "bot1")
|
||||
}
|
||||
|
||||
noPath := BotConfig{Model: "claude-v1"}
|
||||
if err := noPath.PersistModel("claude-sonnet-4-6"); err == nil {
|
||||
t.Error("PersistModel with empty ConfigFilePath: expected error, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func thinkingTestConfig(id string) BotConfig {
|
||||
return BotConfig{
|
||||
ID: id,
|
||||
TelegramToken: "token-" + id,
|
||||
MemorySize: 10,
|
||||
MessagePerHour: 10,
|
||||
MessagePerDay: 100,
|
||||
TempBanDuration: "1h",
|
||||
Model: "claude-test",
|
||||
}
|
||||
}
|
||||
|
||||
func TestThinkingConfig(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
thinking string
|
||||
display string
|
||||
wantErr string
|
||||
}{
|
||||
{"absent", "", "", ""},
|
||||
{"adaptive", ThinkingModeAdaptive, "", ""},
|
||||
{"disabled", ThinkingModeDisabled, "", ""},
|
||||
{"adaptive summarized", ThinkingModeAdaptive, ThinkingDisplaySummarized, ""},
|
||||
{"adaptive omitted", ThinkingModeAdaptive, ThinkingDisplayOmitted, ""},
|
||||
{"legacy enabled rejected", "enabled", "", "invalid 'thinking'"},
|
||||
{"case sensitive", "Adaptive", "", "invalid 'thinking'"},
|
||||
{"unknown display", ThinkingModeAdaptive, "verbose", "invalid 'thinking_display'"},
|
||||
{"display without thinking", "", ThinkingDisplaySummarized, "'thinking_display' requires"},
|
||||
{"display with disabled", ThinkingModeDisabled, ThinkingDisplayOmitted, "'thinking_display' requires"},
|
||||
}
|
||||
for i, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
cfg := thinkingTestConfig(fmt.Sprintf("bot-think-%d", i))
|
||||
cfg.Thinking = tc.thinking
|
||||
cfg.ThinkingDisplay = tc.display
|
||||
err := validateConfig(&cfg, map[string]bool{}, map[string]bool{})
|
||||
if tc.wantErr == "" {
|
||||
if err != nil {
|
||||
t.Fatalf("validateConfig(thinking=%q display=%q) = %v, want nil", tc.thinking, tc.display, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
if err == nil {
|
||||
t.Fatalf("validateConfig(thinking=%q display=%q) = nil, want error containing %q", tc.thinking, tc.display, tc.wantErr)
|
||||
}
|
||||
if !contains(err.Error(), tc.wantErr) {
|
||||
t.Errorf("error %q does not contain %q", err.Error(), tc.wantErr)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestMaxTokensConfig(t *testing.T) {
|
||||
var withValue BotConfig
|
||||
if err := json.Unmarshal([]byte(`{"max_tokens": 4000}`), &withValue); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
if withValue.MaxTokens != 4000 {
|
||||
t.Errorf("MaxTokens = %d, want 4000", withValue.MaxTokens)
|
||||
}
|
||||
|
||||
var withoutValue BotConfig
|
||||
if err := json.Unmarshal([]byte(`{}`), &withoutValue); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
if withoutValue.MaxTokens != 0 {
|
||||
t.Errorf("MaxTokens = %d, want 0 when absent", withoutValue.MaxTokens)
|
||||
}
|
||||
|
||||
neg := thinkingTestConfig("bot-maxtok-neg")
|
||||
neg.MaxTokens = -1
|
||||
if err := validateConfig(&neg, map[string]bool{}, map[string]bool{}); err == nil {
|
||||
t.Error("validateConfig(max_tokens=-1) = nil, want error")
|
||||
}
|
||||
zero := thinkingTestConfig("bot-maxtok-zero")
|
||||
zero.MaxTokens = 0
|
||||
if err := validateConfig(&zero, map[string]bool{}, map[string]bool{}); err != nil {
|
||||
t.Errorf("validateConfig(max_tokens=0) = %v, want nil", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestThinkingConfigLoad(t *testing.T) {
|
||||
jsonData := `{
|
||||
"id": "bot-think-load",
|
||||
"thinking": "adaptive",
|
||||
"thinking_display": "omitted",
|
||||
"max_tokens": 4096
|
||||
}`
|
||||
var cfg BotConfig
|
||||
if err := json.Unmarshal([]byte(jsonData), &cfg); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
if cfg.Thinking != ThinkingModeAdaptive {
|
||||
t.Errorf("Thinking = %q, want %q", cfg.Thinking, ThinkingModeAdaptive)
|
||||
}
|
||||
if cfg.ThinkingDisplay != ThinkingDisplayOmitted {
|
||||
t.Errorf("ThinkingDisplay = %q, want %q", cfg.ThinkingDisplay, ThinkingDisplayOmitted)
|
||||
}
|
||||
if cfg.MaxTokens != 4096 {
|
||||
t.Errorf("MaxTokens = %d, want 4096", cfg.MaxTokens)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWebSearchConfig(t *testing.T) {
|
||||
t.Run("absent leaves nil", func(t *testing.T) {
|
||||
var cfg BotConfig
|
||||
if err := json.Unmarshal([]byte(`{}`), &cfg); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
if cfg.WebSearch != nil {
|
||||
t.Errorf("WebSearch = %+v, want nil when absent", cfg.WebSearch)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("loads allowlist", func(t *testing.T) {
|
||||
jsonData := `{
|
||||
"web_search": {
|
||||
"allowed_domains": ["example.com/hc", "docs.example.com"],
|
||||
"max_uses": 3,
|
||||
"fetch": true,
|
||||
"max_content_tokens": 50000
|
||||
}
|
||||
}`
|
||||
var cfg BotConfig
|
||||
if err := json.Unmarshal([]byte(jsonData), &cfg); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
if cfg.WebSearch == nil {
|
||||
t.Fatalf("WebSearch = nil, want populated")
|
||||
}
|
||||
if len(cfg.WebSearch.AllowedDomains) != 2 {
|
||||
t.Errorf("AllowedDomains = %v, want 2 entries", cfg.WebSearch.AllowedDomains)
|
||||
}
|
||||
if cfg.WebSearch.MaxUses != 3 {
|
||||
t.Errorf("MaxUses = %d, want 3", cfg.WebSearch.MaxUses)
|
||||
}
|
||||
if !cfg.WebSearch.Fetch {
|
||||
t.Error("Fetch = false, want true")
|
||||
}
|
||||
if cfg.WebSearch.MaxContentTokens != 50000 {
|
||||
t.Errorf("MaxContentTokens = %d, want 50000", cfg.WebSearch.MaxContentTokens)
|
||||
}
|
||||
})
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
ws *WebSearchConfig
|
||||
wantErr string
|
||||
}{
|
||||
{"nil ok", nil, ""},
|
||||
{"allowlist ok", &WebSearchConfig{AllowedDomains: []string{"example.com"}, MaxUses: 3}, ""},
|
||||
{"blocklist ok", &WebSearchConfig{BlockedDomains: []string{"evil.com"}}, ""},
|
||||
{"empty ok", &WebSearchConfig{}, ""},
|
||||
{"both allow and block rejected",
|
||||
&WebSearchConfig{AllowedDomains: []string{"a.com"}, BlockedDomains: []string{"b.com"}},
|
||||
"cannot set both allowed_domains and blocked_domains"},
|
||||
{"negative max_uses rejected",
|
||||
&WebSearchConfig{AllowedDomains: []string{"a.com"}, MaxUses: -1},
|
||||
"'web_search.max_uses' must be greater than 0"},
|
||||
{"negative max_content_tokens rejected",
|
||||
&WebSearchConfig{AllowedDomains: []string{"a.com"}, MaxContentTokens: -1},
|
||||
"'web_search.max_content_tokens' must be greater than 0"},
|
||||
}
|
||||
for i, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
cfg := thinkingTestConfig(fmt.Sprintf("bot-web-%d", i))
|
||||
cfg.WebSearch = tc.ws
|
||||
err := validateConfig(&cfg, map[string]bool{}, map[string]bool{})
|
||||
if tc.wantErr == "" {
|
||||
if err != nil {
|
||||
t.Fatalf("validateConfig(web_search=%+v) = %v, want nil", tc.ws, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
if err == nil {
|
||||
t.Fatalf("validateConfig(web_search=%+v) = nil, want error containing %q", tc.ws, tc.wantErr)
|
||||
}
|
||||
if !contains(err.Error(), tc.wantErr) {
|
||||
t.Errorf("error %q does not contain %q", err.Error(), tc.wantErr)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
// Additional tests can be added here to cover more scenarios
|
||||
|
||||
+9
-81
@@ -4,7 +4,6 @@ import (
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"gorm.io/driver/sqlite"
|
||||
@@ -12,27 +11,6 @@ import (
|
||||
"gorm.io/gorm/logger"
|
||||
)
|
||||
|
||||
// dbLogLevel reads DB_LOG_LEVEL. The default is warn (GORM's own default):
|
||||
// errors and slow queries only. "info" logs every statement, which is useful
|
||||
// alone at a terminal and, in a deployed bot, writes each user message and
|
||||
// reply into the container log twice over.
|
||||
func dbLogLevel() logger.LogLevel {
|
||||
raw := strings.ToLower(strings.TrimSpace(os.Getenv("DB_LOG_LEVEL")))
|
||||
switch raw {
|
||||
case "", "warn", "warning":
|
||||
return logger.Warn
|
||||
case "silent":
|
||||
return logger.Silent
|
||||
case "error":
|
||||
return logger.Error
|
||||
case "info":
|
||||
return logger.Info
|
||||
default:
|
||||
InfoLogger.Printf("unknown DB_LOG_LEVEL %q; falling back to warn", raw)
|
||||
return logger.Warn
|
||||
}
|
||||
}
|
||||
|
||||
func initDB() (*gorm.DB, error) {
|
||||
if err := os.MkdirAll("data", 0750); err != nil {
|
||||
return nil, fmt.Errorf("failed to create data directory: %w", err)
|
||||
@@ -42,19 +20,12 @@ func initDB() (*gorm.DB, error) {
|
||||
log.New(log.Writer(), "\r\n", log.LstdFlags),
|
||||
logger.Config{
|
||||
SlowThreshold: time.Second,
|
||||
LogLevel: dbLogLevel(),
|
||||
LogLevel: logger.Info,
|
||||
Colorful: false,
|
||||
// A missing row is routine here (the owner probe, a first-time
|
||||
// user), so it must not surface as an error once the level drops
|
||||
// below info.
|
||||
IgnoreRecordNotFoundError: true,
|
||||
// Log placeholders rather than inline values: a slow-query line
|
||||
// would otherwise carry the message text it was writing.
|
||||
ParameterizedQueries: true,
|
||||
},
|
||||
)
|
||||
|
||||
db, err := gorm.Open(sqlite.Open("data/bot.db?_journal_mode=WAL&_busy_timeout=5000&_foreign_keys=on"), &gorm.Config{
|
||||
db, err := gorm.Open(sqlite.Open("data/bot.db?_journal_mode=WAL&_busy_timeout=5000"), &gorm.Config{
|
||||
Logger: newLogger,
|
||||
})
|
||||
if err != nil {
|
||||
@@ -67,11 +38,17 @@ func initDB() (*gorm.DB, error) {
|
||||
}
|
||||
sqlDB.SetMaxOpenConns(1)
|
||||
|
||||
err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{}, &Scope{})
|
||||
// AutoMigrate the models
|
||||
err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to migrate database schema: %w", err)
|
||||
}
|
||||
|
||||
// Enforce unique owner per bot using raw SQL
|
||||
// Note: SQLite doesn't support partial indexes, but we can simulate it by making a unique index on (BotID, IsOwner)
|
||||
// and ensuring that IsOwner can only be true for one user per BotID.
|
||||
// This approach allows multiple users with IsOwner=false for the same BotID,
|
||||
// but only one user can have IsOwner=true per BotID.
|
||||
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)
|
||||
@@ -82,58 +59,9 @@ func initDB() (*gorm.DB, error) {
|
||||
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 {
|
||||
|
||||
@@ -13,9 +13,6 @@ services:
|
||||
# Optional: Environment variables (can be overridden with .env file)
|
||||
# environment:
|
||||
# - BOT_LOG_LEVEL=info
|
||||
# GORM SQL logging: silent | error | warn (default) | info.
|
||||
# "info" logs every statement, including message text.
|
||||
# - DB_LOG_LEVEL=warn
|
||||
|
||||
# Volume mounts
|
||||
volumes:
|
||||
|
||||
@@ -1,288 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/go-telegram/bot"
|
||||
)
|
||||
|
||||
const (
|
||||
// draftUpdateInterval throttles sendMessageDraft. Telegram documents no
|
||||
// limit for drafts specifically, so this matches its general guidance of
|
||||
// about one send per second per chat before 429s start.
|
||||
draftUpdateInterval = 1 * time.Second
|
||||
// draftKeepaliveInterval re-sends the current draft during silence, such
|
||||
// as a server-side web search. The preview lives roughly 30 seconds when
|
||||
// untouched; whether updates extend it is undocumented, so this keeps
|
||||
// well inside that either way.
|
||||
draftKeepaliveInterval = 10 * time.Second
|
||||
// telegramMessageLimit is the text cap, in characters, for both drafts
|
||||
// and sent messages.
|
||||
telegramMessageLimit = 4096
|
||||
)
|
||||
|
||||
// draftEligible reports whether a turn may stream into a Telegram draft.
|
||||
// Telegram documents drafts for private chats only, and its reference lists
|
||||
// no business_connection_id parameter for sendMessageDraft (the Go library's
|
||||
// struct carries the field, but nothing says the server honours it), so
|
||||
// business traffic keeps the segment path until that is verified live.
|
||||
func (b *Bot) draftEligible(chatID int64, businessConnectionID string) bool {
|
||||
return b.config.StreamDrafts && businessConnectionID == "" && chatID > 0
|
||||
}
|
||||
|
||||
// draftStream is one growing Telegram draft for one turn.
|
||||
//
|
||||
// A single sender goroutine owns every sendMessageDraft call, so sends are
|
||||
// strictly sequential, never run on the model's stream goroutine, and cannot
|
||||
// land at Telegram out of order. The stream goroutine only records the newest
|
||||
// text and signals; the sender throttles, keeps the draft alive through silent
|
||||
// stretches, and goes quiet after the first failure so a broken draft degrades
|
||||
// to a plain final message.
|
||||
type draftStream struct {
|
||||
b *Bot
|
||||
chatID int64
|
||||
draftID string
|
||||
interval time.Duration
|
||||
keepalive time.Duration
|
||||
onFirst func()
|
||||
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
kick chan struct{}
|
||||
done chan struct{}
|
||||
closeOnce sync.Once
|
||||
|
||||
mu sync.Mutex
|
||||
latest string // newest text from the model
|
||||
shown string // last text Telegram accepted
|
||||
failed bool
|
||||
}
|
||||
|
||||
// newDraftStream builds the stream; start launches the sender. The two are
|
||||
// separate so tests can shorten the intervals before anything runs.
|
||||
func (b *Bot) newDraftStream(parent context.Context, chatID int64, draftID string, onFirst func()) *draftStream {
|
||||
ctx, cancel := context.WithCancel(parent)
|
||||
return &draftStream{
|
||||
b: b,
|
||||
chatID: chatID,
|
||||
draftID: draftID,
|
||||
interval: draftUpdateInterval,
|
||||
keepalive: draftKeepaliveInterval,
|
||||
onFirst: onFirst,
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
kick: make(chan struct{}, 1),
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
func (d *draftStream) start() { go d.run() }
|
||||
|
||||
// update records the full reply text so far and wakes the sender. It never
|
||||
// blocks on Telegram.
|
||||
func (d *draftStream) update(text string) {
|
||||
d.mu.Lock()
|
||||
d.latest = text
|
||||
d.mu.Unlock()
|
||||
select {
|
||||
case d.kick <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
// run is the sender loop: wake on new text or on the keepalive deadline,
|
||||
// respect the throttle, send the newest text, repeat until closed or failed.
|
||||
func (d *draftStream) run() {
|
||||
defer close(d.done)
|
||||
|
||||
var lastSend time.Time
|
||||
for {
|
||||
var keepalive <-chan time.Time
|
||||
if !lastSend.IsZero() {
|
||||
keepalive = time.After(d.keepalive - time.Since(lastSend))
|
||||
}
|
||||
woken := false
|
||||
select {
|
||||
case <-d.ctx.Done():
|
||||
return
|
||||
case <-d.kick:
|
||||
case <-keepalive:
|
||||
woken = true
|
||||
}
|
||||
|
||||
if wait := d.interval - time.Since(lastSend); !lastSend.IsZero() && wait > 0 {
|
||||
select {
|
||||
case <-d.ctx.Done():
|
||||
return
|
||||
case <-time.After(wait):
|
||||
}
|
||||
}
|
||||
|
||||
d.mu.Lock()
|
||||
text, shown := d.latest, d.shown
|
||||
d.mu.Unlock()
|
||||
if text == "" || (text == shown && !woken) {
|
||||
continue
|
||||
}
|
||||
if !d.send(text) {
|
||||
return
|
||||
}
|
||||
lastSend = time.Now()
|
||||
}
|
||||
}
|
||||
|
||||
// send performs one sendMessageDraft call and reports whether the draft is
|
||||
// still usable afterwards.
|
||||
func (d *draftStream) send(text string) bool {
|
||||
text = clampRunes(text, telegramMessageLimit)
|
||||
_, err := d.b.tgBot.SendMessageDraft(d.ctx, &bot.SendMessageDraftParams{
|
||||
ChatID: d.chatID,
|
||||
DraftID: d.draftID,
|
||||
Text: text,
|
||||
CanStop: true,
|
||||
KeepOnStop: true,
|
||||
})
|
||||
if err != nil {
|
||||
d.mu.Lock()
|
||||
d.failed = true
|
||||
d.mu.Unlock()
|
||||
if d.ctx.Err() == nil {
|
||||
// Cosmetic from here on: the final message is still sent.
|
||||
InfoLogger.Printf("[%s] draft %s for chat %d failed, falling back to final message only: %v",
|
||||
d.b.config.ID, d.draftID, d.chatID, err)
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
d.mu.Lock()
|
||||
first := d.shown == ""
|
||||
d.shown = text
|
||||
d.mu.Unlock()
|
||||
if first && d.onFirst != nil {
|
||||
d.onFirst()
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// current returns the text the user has actually seen.
|
||||
func (d *draftStream) current() string {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
return d.shown
|
||||
}
|
||||
|
||||
// close stops the sender and waits for any send in flight, so nothing lands
|
||||
// on the chat after the final message. Telegram removes the draft itself once
|
||||
// the bot sends a message.
|
||||
func (d *draftStream) close() {
|
||||
d.closeOnce.Do(func() {
|
||||
d.cancel()
|
||||
<-d.done
|
||||
})
|
||||
}
|
||||
|
||||
// replyOutput is the sink a turn writes into. Two implementations: a draft
|
||||
// that grows in place and becomes one final message, or the original
|
||||
// one-message-per-text-block segments.
|
||||
type replyOutput struct {
|
||||
onSegment func(string) error
|
||||
onProgress func(string)
|
||||
finish func(ctx context.Context, joined string) error
|
||||
partial func() string
|
||||
close func()
|
||||
}
|
||||
|
||||
func (b *Bot) newReplyOutput(turnCtx context.Context, turn *chatTurn, chatID int64, businessConnectionID string) replyOutput {
|
||||
if b.draftEligible(chatID, businessConnectionID) {
|
||||
draftID := turn.draftIDString()
|
||||
d := b.newDraftStream(turnCtx, chatID, draftID, func() { b.markCommitted(turn) })
|
||||
b.setTurnDraft(turn, draftID)
|
||||
d.start()
|
||||
return replyOutput{
|
||||
onProgress: d.update,
|
||||
finish: func(ctx context.Context, joined string) error {
|
||||
d.close()
|
||||
return b.sendLongMessage(ctx, chatID, joined, businessConnectionID)
|
||||
},
|
||||
partial: d.current,
|
||||
close: d.close,
|
||||
}
|
||||
}
|
||||
|
||||
return replyOutput{
|
||||
onSegment: func(seg string) error {
|
||||
b.markCommitted(turn)
|
||||
return b.sendOneSegment(turnCtx, chatID, seg, businessConnectionID)
|
||||
},
|
||||
finish: func(context.Context, string) error { return nil },
|
||||
partial: func() string { return "" },
|
||||
close: func() {},
|
||||
}
|
||||
}
|
||||
|
||||
// sendLongMessage delivers text as one message, or as few as the character
|
||||
// cap allows, splitting on paragraph and line boundaries before cutting words.
|
||||
func (b *Bot) sendLongMessage(ctx context.Context, chatID int64, text, businessConnectionID string) error {
|
||||
for _, chunk := range splitTelegramMessage(text, telegramMessageLimit) {
|
||||
if err := b.sendOneSegment(ctx, chatID, chunk, businessConnectionID); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// splitTelegramMessage splits text into pieces of at most limit runes,
|
||||
// preferring paragraph breaks, then line breaks, then spaces.
|
||||
func splitTelegramMessage(text string, limit int) []string {
|
||||
text = strings.TrimSpace(text)
|
||||
if text == "" {
|
||||
return nil
|
||||
}
|
||||
var out []string
|
||||
for {
|
||||
off := runeOffset(text, limit)
|
||||
if off < 0 {
|
||||
return append(out, text)
|
||||
}
|
||||
window := text[:off]
|
||||
cut := -1
|
||||
for _, sep := range []string{"\n\n", "\n", " "} {
|
||||
if i := strings.LastIndex(window, sep); i > 0 {
|
||||
cut = i
|
||||
break
|
||||
}
|
||||
}
|
||||
if cut < 0 {
|
||||
cut = off
|
||||
}
|
||||
out = append(out, strings.TrimSpace(text[:cut]))
|
||||
text = strings.TrimSpace(text[cut:])
|
||||
if text == "" {
|
||||
return out
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// runeOffset returns the byte offset just past the n-th rune of s, or -1 when
|
||||
// s holds at most n runes.
|
||||
func runeOffset(s string, n int) int {
|
||||
count := 0
|
||||
for i := range s {
|
||||
if count == n {
|
||||
return i
|
||||
}
|
||||
count++
|
||||
}
|
||||
return -1
|
||||
}
|
||||
|
||||
// clampRunes truncates s to at most n runes.
|
||||
func clampRunes(s string, n int) string {
|
||||
if off := runeOffset(s, n); off >= 0 {
|
||||
return s[:off]
|
||||
}
|
||||
return s
|
||||
}
|
||||
-217
@@ -1,217 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/go-telegram/bot"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type draftRecorder struct {
|
||||
mu sync.Mutex
|
||||
calls []bot.SendMessageDraftParams
|
||||
err error
|
||||
delay time.Duration
|
||||
}
|
||||
|
||||
func (r *draftRecorder) send(_ context.Context, p *bot.SendMessageDraftParams) (bool, error) {
|
||||
time.Sleep(r.delay)
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.calls = append(r.calls, *p)
|
||||
return r.err == nil, r.err
|
||||
}
|
||||
|
||||
func (r *draftRecorder) texts() []string {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
out := make([]string, len(r.calls))
|
||||
for i, c := range r.calls {
|
||||
out[i] = c.Text
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (r *draftRecorder) count() int { return len(r.texts()) }
|
||||
|
||||
const testDraftInterval = 40 * time.Millisecond
|
||||
|
||||
func newTestDraft(t *testing.T, rec *draftRecorder, onFirst func()) *draftStream {
|
||||
t.Helper()
|
||||
b, mockTg := setupBotForTest(t, 123)
|
||||
mockTg.SendMessageDraftFunc = rec.send
|
||||
d := b.newDraftStream(context.Background(), 123, "7", onFirst)
|
||||
d.interval = testDraftInterval
|
||||
d.keepalive = time.Hour
|
||||
d.start()
|
||||
t.Cleanup(d.close)
|
||||
return d
|
||||
}
|
||||
|
||||
func eventuallyCalls(t *testing.T, rec *draftRecorder, n int) {
|
||||
t.Helper()
|
||||
require.Eventually(t, func() bool { return rec.count() == n }, time.Second, 2*time.Millisecond)
|
||||
}
|
||||
|
||||
func TestDraftStream_FirstUpdateSendsPromptlyThenThrottles(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
rec := &draftRecorder{}
|
||||
d := newTestDraft(t, rec, nil)
|
||||
|
||||
d.update("Hel")
|
||||
eventuallyCalls(t, rec, 1)
|
||||
assert.Equal(t, "Hel", rec.texts()[0], "the first text goes out without waiting for the window")
|
||||
|
||||
d.update("Hello")
|
||||
d.update("Hello, wor")
|
||||
time.Sleep(testDraftInterval / 4)
|
||||
assert.Equal(t, 1, rec.count(), "inside the window nothing is sent")
|
||||
|
||||
eventuallyCalls(t, rec, 2)
|
||||
assert.Equal(t, "Hello, wor", rec.texts()[1], "only the newest text is sent when the window elapses")
|
||||
}
|
||||
|
||||
func TestDraftStream_SendsAreSequential(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
rec := &draftRecorder{delay: 3 * testDraftInterval}
|
||||
d := newTestDraft(t, rec, nil)
|
||||
|
||||
d.update("a")
|
||||
time.Sleep(testDraftInterval)
|
||||
d.update("ab")
|
||||
time.Sleep(testDraftInterval)
|
||||
d.update("abc")
|
||||
|
||||
eventuallyCalls(t, rec, 2)
|
||||
time.Sleep(2 * testDraftInterval)
|
||||
assert.Equal(t, []string{"a", "abc"}, rec.texts(), "a slow send never overlaps the next; the newest text wins")
|
||||
}
|
||||
|
||||
func TestDraftStream_ParamsCarryStopButton(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
rec := &draftRecorder{}
|
||||
d := newTestDraft(t, rec, nil)
|
||||
|
||||
d.update("x")
|
||||
eventuallyCalls(t, rec, 1)
|
||||
|
||||
p := rec.calls[0]
|
||||
assert.Equal(t, "7", p.DraftID)
|
||||
assert.Equal(t, int64(123), p.ChatID)
|
||||
assert.True(t, p.CanStop)
|
||||
assert.True(t, p.KeepOnStop)
|
||||
}
|
||||
|
||||
func TestDraftStream_OnFirstFiresOnce(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
rec := &draftRecorder{}
|
||||
var mu sync.Mutex
|
||||
fired := 0
|
||||
d := newTestDraft(t, rec, func() { mu.Lock(); fired++; mu.Unlock() })
|
||||
|
||||
d.update("a")
|
||||
eventuallyCalls(t, rec, 1)
|
||||
d.update("ab")
|
||||
eventuallyCalls(t, rec, 2)
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
assert.Equal(t, 1, fired)
|
||||
}
|
||||
|
||||
func TestDraftStream_FailureStopsFurtherSends(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
rec := &draftRecorder{err: errors.New("TEXTDRAFT_PEER_INVALID")}
|
||||
fired := false
|
||||
d := newTestDraft(t, rec, func() { fired = true })
|
||||
|
||||
d.update("a")
|
||||
eventuallyCalls(t, rec, 1)
|
||||
d.update("ab")
|
||||
time.Sleep(3 * testDraftInterval)
|
||||
|
||||
assert.Equal(t, 1, rec.count(), "after a failure the draft goes quiet")
|
||||
assert.False(t, fired, "a rejected draft never commits the turn")
|
||||
assert.Equal(t, "", d.current(), "nothing was shown to the user")
|
||||
}
|
||||
|
||||
func TestDraftStream_ClampsToTelegramLimit(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
rec := &draftRecorder{}
|
||||
d := newTestDraft(t, rec, nil)
|
||||
|
||||
d.update(strings.Repeat("é", telegramMessageLimit+500))
|
||||
eventuallyCalls(t, rec, 1)
|
||||
|
||||
assert.Equal(t, telegramMessageLimit, len([]rune(rec.texts()[0])), "a draft over the cap is cut, not rejected")
|
||||
assert.Eventually(t, func() bool { return d.current() != "" }, time.Second, 2*time.Millisecond)
|
||||
}
|
||||
|
||||
func TestDraftStream_KeepaliveResendsDuringSilence(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
rec := &draftRecorder{}
|
||||
d := newTestDraft(t, rec, nil)
|
||||
d.keepalive = 2 * testDraftInterval
|
||||
|
||||
d.update("thinking")
|
||||
eventuallyCalls(t, rec, 1)
|
||||
|
||||
require.Eventually(t, func() bool { return rec.count() >= 2 }, time.Second, 2*time.Millisecond)
|
||||
assert.Equal(t, "thinking", rec.texts()[1], "silence re-sends the same text to keep the draft alive")
|
||||
}
|
||||
|
||||
func TestDraftStream_CloseWaitsAndSilences(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
rec := &draftRecorder{delay: 2 * testDraftInterval}
|
||||
d := newTestDraft(t, rec, nil)
|
||||
|
||||
d.update("a")
|
||||
time.Sleep(testDraftInterval / 4) // the send is now in flight
|
||||
d.close()
|
||||
|
||||
n := rec.count()
|
||||
assert.LessOrEqual(t, n, 1)
|
||||
d.update("after close")
|
||||
time.Sleep(3 * testDraftInterval)
|
||||
assert.Equal(t, n, rec.count(), "nothing is sent after close returns")
|
||||
}
|
||||
|
||||
func TestDraftEligible(t *testing.T) {
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
|
||||
assert.False(t, b.draftEligible(123, ""), "off by default")
|
||||
b.config.StreamDrafts = true
|
||||
assert.True(t, b.draftEligible(123, ""))
|
||||
assert.False(t, b.draftEligible(123, "biz-1"), "business connections keep segments")
|
||||
assert.False(t, b.draftEligible(-100123, ""), "groups keep segments")
|
||||
}
|
||||
|
||||
func TestSplitTelegramMessage(t *testing.T) {
|
||||
assert.Nil(t, splitTelegramMessage(" ", 10))
|
||||
assert.Equal(t, []string{"short"}, splitTelegramMessage("short", 10))
|
||||
|
||||
parts := splitTelegramMessage("para one\n\npara two\n\npara three", 20)
|
||||
assert.Equal(t, []string{"para one\n\npara two", "para three"}, parts, "prefers paragraph breaks")
|
||||
|
||||
words := splitTelegramMessage("one two three four", 9)
|
||||
assert.Equal(t, []string{"one two", "three", "four"}, words, "falls back to spaces")
|
||||
|
||||
for _, p := range splitTelegramMessage(strings.Repeat("a", 25), 10) {
|
||||
assert.LessOrEqual(t, len([]rune(p)), 10, "hard cut when there is no separator")
|
||||
}
|
||||
|
||||
multibyte := splitTelegramMessage(strings.Repeat("é", 12), 5)
|
||||
assert.Len(t, multibyte, 3, "limits count runes, not bytes")
|
||||
}
|
||||
|
||||
func TestClampRunes(t *testing.T) {
|
||||
assert.Equal(t, "abc", clampRunes("abc", 5))
|
||||
assert.Equal(t, "abc", clampRunes("abcde", 3))
|
||||
assert.Equal(t, "éé", clampRunes("ééé", 2))
|
||||
assert.Equal(t, "", clampRunes("x", 0))
|
||||
}
|
||||
|
||||
func TestJoinNonEmpty(t *testing.T) {
|
||||
assert.Equal(t, "", joinNonEmpty("", ""))
|
||||
assert.Equal(t, "a", joinNonEmpty("a", ""))
|
||||
assert.Equal(t, "b", joinNonEmpty("", "b"))
|
||||
assert.Equal(t, "a\n\nb", joinNonEmpty("a", "b"))
|
||||
}
|
||||
@@ -1,181 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/go-telegram/bot"
|
||||
"github.com/go-telegram/bot/models"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// storeUserMessage puts a user message with a Telegram id into DB and memory,
|
||||
// the way screenIncomingMessage does, without going through handleUpdate.
|
||||
func storeUserMessage(t *testing.T, b *Bot, chatID int64, telegramID int, text string) Message {
|
||||
t.Helper()
|
||||
msg := b.createMessage(chatID, chatID, "owner", "user", text, true)
|
||||
msg.TelegramMessageID = telegramID
|
||||
require.NoError(t, b.storeMessage(&msg))
|
||||
b.addMessageToChatMemory(b.getOrCreateChatMemory(chatID), msg)
|
||||
return msg
|
||||
}
|
||||
|
||||
func editUpdate(chatID int64, telegramID int, text string) *models.Update {
|
||||
return &models.Update{EditedMessage: &models.Message{
|
||||
ID: telegramID, Chat: models.Chat{ID: chatID}, From: &models.User{ID: chatID, Username: "owner"}, Text: text,
|
||||
}}
|
||||
}
|
||||
|
||||
func TestHandleUpdate_EditedMessageUpdatesTextAndRearmsWindow(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
b.config.DebounceMs = int(bufferOnly / time.Millisecond)
|
||||
const chatID int64 = 123
|
||||
from := &models.User{ID: chatID, Username: "owner"}
|
||||
|
||||
b.handleUpdate(context.Background(), nil, &models.Update{Message: &models.Message{
|
||||
ID: 77, Chat: models.Chat{ID: chatID}, From: from, Text: "How to use libtik",
|
||||
}})
|
||||
require.True(t, b.hasPendingIntake(chatID))
|
||||
b.intakeBuffersMu.Lock()
|
||||
seqBefore := b.intakeBuffers[chatID].seq
|
||||
timerBefore := b.intakeBuffers[chatID].timer
|
||||
b.intakeBuffersMu.Unlock()
|
||||
|
||||
b.handleUpdate(context.Background(), nil, editUpdate(chatID, 77, "How to use libtibik"))
|
||||
|
||||
var stored Message
|
||||
require.NoError(t, b.db.Where("chat_id = ? AND telegram_message_id = ?", chatID, 77).First(&stored).Error)
|
||||
assert.Equal(t, "How to use libtibik", stored.Text, "database copy follows the edit")
|
||||
|
||||
mem := b.getOrCreateChatMemory(chatID)
|
||||
assert.Equal(t, "How to use libtibik", mem.Messages[len(mem.Messages)-1].Text, "memory copy follows the edit")
|
||||
|
||||
require.True(t, b.hasPendingIntake(chatID), "an edit must not flush or drop the batch")
|
||||
b.intakeBuffersMu.Lock()
|
||||
defer b.intakeBuffersMu.Unlock()
|
||||
assert.Equal(t, seqBefore, b.intakeBuffers[chatID].seq, "same batch")
|
||||
assert.NotSame(t, timerBefore, b.intakeBuffers[chatID].timer, "the window restarted")
|
||||
assert.Equal(t, 1, b.intakeBuffers[chatID].count, "an edit is not a new message")
|
||||
}
|
||||
|
||||
func TestHandleUpdate_EditFromEmojiToTextClearsEmojiFlag(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
b.config.DebounceMs = int(bufferOnly / time.Millisecond)
|
||||
const chatID int64 = 123
|
||||
|
||||
b.handleUpdate(context.Background(), nil, &models.Update{Message: &models.Message{
|
||||
ID: 78, Chat: models.Chat{ID: chatID}, From: &models.User{ID: chatID}, Text: "👍",
|
||||
}})
|
||||
b.intakeBuffersMu.Lock()
|
||||
require.True(t, b.intakeBuffers[chatID].allEmojiOnly)
|
||||
b.intakeBuffersMu.Unlock()
|
||||
|
||||
b.handleUpdate(context.Background(), nil, editUpdate(chatID, 78, "can you explain that again?"))
|
||||
|
||||
b.intakeBuffersMu.Lock()
|
||||
defer b.intakeBuffersMu.Unlock()
|
||||
assert.False(t, b.intakeBuffers[chatID].allEmojiOnly, "a real question must not get an emoji-only reply")
|
||||
}
|
||||
|
||||
func TestHandleUpdate_CaptionEditApplies(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
const chatID int64 = 123
|
||||
storeUserMessage(t, b, chatID, 79, "old caption")
|
||||
|
||||
b.handleUpdate(context.Background(), nil, &models.Update{EditedMessage: &models.Message{
|
||||
ID: 79, Chat: models.Chat{ID: chatID}, From: &models.User{ID: chatID}, Caption: "new caption",
|
||||
}})
|
||||
|
||||
var stored Message
|
||||
require.NoError(t, b.db.Where("chat_id = ? AND telegram_message_id = ?", chatID, 79).First(&stored).Error)
|
||||
assert.Equal(t, "new caption", stored.Text)
|
||||
}
|
||||
|
||||
func TestHandleUpdate_EditOfAnsweredMessageOnlyCorrectsRecord(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
b.config.DebounceMs = int(bufferOnly / time.Millisecond)
|
||||
const chatID int64 = 123
|
||||
storeUserMessage(t, b, chatID, 80, "first question")
|
||||
_, err := b.screenOutgoingMessage(chatID, "first answer")
|
||||
require.NoError(t, err)
|
||||
|
||||
// A newer, unrelated message is waiting in its window.
|
||||
b.handleUpdate(context.Background(), nil, &models.Update{Message: &models.Message{
|
||||
ID: 81, Chat: models.Chat{ID: chatID}, From: &models.User{ID: chatID}, Text: "second question",
|
||||
}})
|
||||
b.intakeBuffersMu.Lock()
|
||||
timerBefore := b.intakeBuffers[chatID].timer
|
||||
b.intakeBuffersMu.Unlock()
|
||||
|
||||
b.handleUpdate(context.Background(), nil, editUpdate(chatID, 80, "first question, corrected"))
|
||||
|
||||
var stored Message
|
||||
require.NoError(t, b.db.Where("telegram_message_id = ?", 80).First(&stored).Error)
|
||||
assert.Equal(t, "first question, corrected", stored.Text, "the record is corrected")
|
||||
b.intakeBuffersMu.Lock()
|
||||
defer b.intakeBuffersMu.Unlock()
|
||||
assert.Same(t, timerBefore, b.intakeBuffers[chatID].timer, "an answered message's edit must not touch another batch's window")
|
||||
}
|
||||
|
||||
func TestHandleUpdate_EditRestartsUncommittedTurn(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
const chatID int64 = 123
|
||||
msg := storeUserMessage(t, b, chatID, 82, "how do I use libtik")
|
||||
turn, ctx, ok := b.beginTextTurn(textReq(chatID), msg.ID)
|
||||
require.True(t, ok)
|
||||
|
||||
b.handleUpdate(context.Background(), nil, editUpdate(chatID, 82, "how do I use libtibik"))
|
||||
|
||||
assert.ErrorIs(t, ctx.Err(), context.Canceled, "the model was reading the typo; rerun with the fix")
|
||||
assert.Equal(t, cancelRestart, b.turnReason(turn))
|
||||
next := b.finishTurn(turn)
|
||||
require.NotNil(t, next)
|
||||
assert.Equal(t, chatID, next.chatID)
|
||||
}
|
||||
|
||||
func TestHandleUpdate_EditOfUnknownMessageIsIgnored(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
|
||||
b.handleUpdate(context.Background(), nil, editUpdate(123, 999, "never stored"))
|
||||
|
||||
assert.False(t, b.hasPendingIntake(123))
|
||||
var count int64
|
||||
b.db.Model(&Message{}).Where("chat_id = ?", 123).Count(&count)
|
||||
assert.Zero(t, count, "an edit never creates a message")
|
||||
}
|
||||
|
||||
func TestHandleUpdate_StopGenerationCancelsMatchingTurn(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
turn, ctx, ok := b.beginTextTurn(textReq(123), 1)
|
||||
require.True(t, ok)
|
||||
b.setTurnDraft(turn, turn.draftIDString())
|
||||
|
||||
b.handleUpdate(context.Background(), nil, &models.Update{
|
||||
StoppedMessageGeneration: &models.MessageGenerationStopped{
|
||||
Chat: models.Chat{ID: 123}, DraftID: int(turn.seq),
|
||||
},
|
||||
})
|
||||
|
||||
assert.ErrorIs(t, ctx.Err(), context.Canceled)
|
||||
assert.Equal(t, cancelStop, b.turnReason(turn))
|
||||
b.finishTurn(turn)
|
||||
}
|
||||
|
||||
func TestClearChatHistory_DiscardsCollectedFlush(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, mockTg := setupBotForTest(t, 123)
|
||||
mockTg.SendMessageFunc = func(_ context.Context, _ *bot.SendMessageParams) (*models.Message, error) {
|
||||
return &models.Message{}, nil
|
||||
}
|
||||
const chatID int64 = 123
|
||||
turn, _, ok := b.beginTextTurn(textReq(chatID), 5)
|
||||
require.True(t, ok)
|
||||
b.markCommitted(turn)
|
||||
_, _, ok = b.beginTextTurn(textReq(chatID), 7)
|
||||
require.False(t, ok, "collected behind the running turn")
|
||||
|
||||
b.clearChatHistory(context.Background(), chatID, chatID, 0, 0, "", false)
|
||||
|
||||
assert.Nil(t, b.finishTurn(turn), "the collected flush must not rerun against the emptied chat")
|
||||
}
|
||||
@@ -1,99 +0,0 @@
|
||||
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_v2"); 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
|
||||
}
|
||||
Binary file not shown.
@@ -1,40 +1,23 @@
|
||||
module github.com/HugeFrog24/go-telegram-bot
|
||||
|
||||
go 1.27.0
|
||||
go 1.26.0
|
||||
|
||||
require (
|
||||
github.com/anthropics/anthropic-sdk-go/v2 v2.3.0
|
||||
github.com/go-telegram/bot v1.27.0
|
||||
github.com/stretchr/testify v1.12.1
|
||||
golang.org/x/sync v0.24.0
|
||||
golang.org/x/time v0.16.0
|
||||
github.com/go-telegram/bot v1.19.0
|
||||
github.com/liushuangls/go-anthropic/v2 v2.17.1
|
||||
github.com/stretchr/testify v1.11.1
|
||||
golang.org/x/time v0.14.0
|
||||
gorm.io/driver/sqlite v1.6.0
|
||||
gorm.io/gorm v1.31.2
|
||||
gorm.io/gorm v1.31.1
|
||||
)
|
||||
|
||||
require (
|
||||
github.com/bahlo/generic-list-go v0.2.0 // indirect
|
||||
github.com/buger/jsonparser v1.6.1 // indirect
|
||||
github.com/cespare/xxhash/v2 v2.3.0 // indirect
|
||||
github.com/go-logr/logr v1.4.4 // indirect
|
||||
github.com/go-logr/stdr v1.2.2 // indirect
|
||||
github.com/invopop/jsonschema v0.14.0 // indirect
|
||||
github.com/davecgh/go-spew v1.1.1 // indirect
|
||||
github.com/jinzhu/inflection v1.0.0 // indirect
|
||||
github.com/jinzhu/now v1.1.5 // indirect
|
||||
github.com/mattn/go-sqlite3 v1.14.52 // indirect
|
||||
github.com/pb33f/go-yaml v0.1.1 // indirect
|
||||
github.com/pb33f/ordered-map/v2 v2.3.2 // indirect
|
||||
github.com/standard-webhooks/standard-webhooks/libraries v0.0.1 // indirect
|
||||
github.com/stretchr/objx v0.8.0 // indirect
|
||||
github.com/tidwall/gjson v1.20.0 // indirect
|
||||
github.com/tidwall/match v1.2.0 // indirect
|
||||
github.com/tidwall/pretty v1.2.2 // indirect
|
||||
github.com/tidwall/sjson v1.2.5 // indirect
|
||||
go.opentelemetry.io/auto/sdk v1.2.1 // indirect
|
||||
go.opentelemetry.io/otel v1.47.0 // indirect
|
||||
go.opentelemetry.io/otel/log v1.47.0 // indirect
|
||||
go.opentelemetry.io/otel/metric v1.47.0 // indirect
|
||||
go.opentelemetry.io/otel/trace v1.47.0 // indirect
|
||||
go.yaml.in/yaml/v3 v3.0.5 // indirect
|
||||
golang.org/x/text v0.43.0 // indirect
|
||||
github.com/mattn/go-sqlite3 v1.14.34 // indirect
|
||||
github.com/pmezard/go-difflib v1.0.0 // indirect
|
||||
github.com/stretchr/objx v0.5.3 // indirect
|
||||
golang.org/x/text v0.34.0 // indirect
|
||||
gopkg.in/yaml.v3 v3.0.1 // indirect
|
||||
)
|
||||
|
||||
@@ -1,152 +1,30 @@
|
||||
cloud.google.com/go/auth v0.7.2 h1:uiha352VrCDMXg+yoBtaD0tUF4Kv9vrtrWPYXwutnDE=
|
||||
cloud.google.com/go/auth v0.7.2/go.mod h1:VEc4p5NNxycWQTMQEDQF0bd6aTMb6VgYDXEwiJJQAbs=
|
||||
cloud.google.com/go/auth/oauth2adapt v0.2.3 h1:MlxF+Pd3OmSudg/b1yZ5lJwoXCEaeedAguodky1PcKI=
|
||||
cloud.google.com/go/auth/oauth2adapt v0.2.3/go.mod h1:tMQXOfZzFuNuUxOypHlQEXgdfX5cuhwU+ffUuXRJE8I=
|
||||
cloud.google.com/go/compute/metadata v0.5.0 h1:Zr0eK8JbFv6+Wi4ilXAR8FJ3wyNdpxHKJNPos6LTZOY=
|
||||
cloud.google.com/go/compute/metadata v0.5.0/go.mod h1:aHnloV2TPI38yx4s9+wAZhHykWvVCfu7hQbF+9CWoiY=
|
||||
github.com/anthropics/anthropic-sdk-go/v2 v2.3.0 h1:wEaN13ToNCcjdqP+47ALC78Jt+DkGMmwqL5zzPQTGkQ=
|
||||
github.com/anthropics/anthropic-sdk-go/v2 v2.3.0/go.mod h1:LAOInMKgn6reTIrPvKXa5weyqzxsJkxSBuzhnxzukpQ=
|
||||
github.com/aws/aws-sdk-go-v2 v1.38.0 h1:UCRQ5mlqcFk9HJDIqENSLR3wiG1VTWlyUfLDEvY7RxU=
|
||||
github.com/aws/aws-sdk-go-v2 v1.38.0/go.mod h1:9Q0OoGQoboYIAJyslFyF1f5K1Ryddop8gqMhWx/n4Wg=
|
||||
github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.3 h1:tW1/Rkad38LA15X4UQtjXZXNKsCgkshC3EbmcUmghTg=
|
||||
github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.3/go.mod h1:UbnqO+zjqk3uIt9yCACHJ9IVNhyhOCnYk8yA19SAWrM=
|
||||
github.com/aws/aws-sdk-go-v2/config v1.27.27 h1:HdqgGt1OAP0HkEDDShEl0oSYa9ZZBSOmKpdpsDMdO90=
|
||||
github.com/aws/aws-sdk-go-v2/config v1.27.27/go.mod h1:MVYamCg76dFNINkZFu4n4RjDixhVr51HLj4ErWzrVwg=
|
||||
github.com/aws/aws-sdk-go-v2/credentials v1.17.27 h1:2raNba6gr2IfA0eqqiP2XiQ0UVOpGPgDSi0I9iAP+UI=
|
||||
github.com/aws/aws-sdk-go-v2/credentials v1.17.27/go.mod h1:gniiwbGahQByxan6YjQUMcW4Aov6bLC3m+evgcoN4r4=
|
||||
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.11 h1:KreluoV8FZDEtI6Co2xuNk/UqI9iwMrOx/87PBNIKqw=
|
||||
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.11/go.mod h1:SeSUYBLsMYFoRvHE0Tjvn7kbxaUhl75CJi1sbfhMxkU=
|
||||
github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.15 h1:SoNJ4RlFEQEbtDcCEt+QG56MY4fm4W8rYirAmq+/DdU=
|
||||
github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.15/go.mod h1:U9ke74k1n2bf+RIgoX1SXFed1HLs51OgUSs+Ph0KJP8=
|
||||
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.15 h1:C6WHdGnTDIYETAm5iErQUiVNsclNx9qbJVPIt03B6bI=
|
||||
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.15/go.mod h1:ZQLZqhcu+JhSrA9/NXRm8SkDvsycE+JkV3WGY41e+IM=
|
||||
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0 h1:hT8rVHwugYE2lEfdFE0QWVo81lF7jMrYJVDWI+f+VxU=
|
||||
github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0/go.mod h1:8tu/lYfQfFe6IGnaOdrpVgEL2IrrDOf6/m9RQum4NkY=
|
||||
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.3 h1:dT3MqvGhSoaIhRseqw2I0yH81l7wiR2vjs57O51EAm8=
|
||||
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.3/go.mod h1:GlAeCkHwugxdHaueRr4nhPuY+WW+gR8UjlcqzPr1SPI=
|
||||
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.17 h1:HGErhhrxZlQ044RiM+WdoZxp0p+EGM62y3L6pwA4olE=
|
||||
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.17/go.mod h1:RkZEx4l0EHYDJpWppMJ3nD9wZJAa8/0lq9aVC+r2UII=
|
||||
github.com/aws/aws-sdk-go-v2/service/sso v1.22.4 h1:BXx0ZIxvrJdSgSvKTZ+yRBeSqqgPM89VPlulEcl37tM=
|
||||
github.com/aws/aws-sdk-go-v2/service/sso v1.22.4/go.mod h1:ooyCOXjvJEsUw7x+ZDHeISPMhtwI3ZCB7ggFMcFfWLU=
|
||||
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.26.4 h1:yiwVzJW2ZxZTurVbYWA7QOrAaCYQR72t0wrSBfoesUE=
|
||||
github.com/aws/aws-sdk-go-v2/service/ssooidc v1.26.4/go.mod h1:0oxfLkpz3rQ/CHlx5hB7H69YUpFiI1tql6Q6Ne+1bCw=
|
||||
github.com/aws/aws-sdk-go-v2/service/sts v1.30.3 h1:ZsDKRLXGWHk8WdtyYMoGNO7bTudrvuKpDKgMVRlepGE=
|
||||
github.com/aws/aws-sdk-go-v2/service/sts v1.30.3/go.mod h1:zwySh8fpFyXp9yOr/KVzxOl8SRqgf/IDw5aUt9UKFcQ=
|
||||
github.com/aws/smithy-go v1.22.5 h1:P9ATCXPMb2mPjYBgueqJNCA5S9UfktsW0tTxi+a7eqw=
|
||||
github.com/aws/smithy-go v1.22.5/go.mod h1:t1ufH5HMublsJYulve2RKmHDC15xu1f26kHCp/HgceI=
|
||||
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.6.1 h1:I0phFv0PlbLHnM7TZAVjZ2MJ2/eWRTDyuO7GLR98IEs=
|
||||
github.com/buger/jsonparser v1.6.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0=
|
||||
github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
|
||||
github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
|
||||
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/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg=
|
||||
github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
|
||||
github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
|
||||
github.com/go-logr/logr v1.4.4 h1:tG4xh9yMsRCAiodLVTxyrkzSZ9+o0L1Kg/+cPVcbP/8=
|
||||
github.com/go-logr/logr v1.4.4/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
|
||||
github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
|
||||
github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
|
||||
github.com/go-telegram/bot v1.27.0 h1:PG9yBDzM8O26oUerE9SZwWlHon4RpH8WcNslZ7CEZmE=
|
||||
github.com/go-telegram/bot v1.27.0/go.mod h1:i2TRs7fXWIeaceF3z7KzsMt/he0TwkVC680mvdTFYeM=
|
||||
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da h1:oI5xCqsCo564l8iNU+DwB5epxmsaqB+rhGL0m5jtYqE=
|
||||
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc=
|
||||
github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
|
||||
github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
|
||||
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
|
||||
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
|
||||
github.com/google/jsonschema-go v0.4.2 h1:tmrUohrwoLZZS/P3x7ex0WAVknEkBZM46iALbcqoRA8=
|
||||
github.com/google/jsonschema-go v0.4.2/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE=
|
||||
github.com/google/s2a-go v0.1.7 h1:60BLSyTrOV4/haCDW4zb1guZItoSq8foHCXrAnjBo/o=
|
||||
github.com/google/s2a-go v0.1.7/go.mod h1:50CgR4k1jNlWBu4UfS4AcfhVe1r6pdZPygJ3R8F0Qdw=
|
||||
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
|
||||
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
github.com/googleapis/enterprise-certificate-proxy v0.3.2 h1:Vie5ybvEvT75RniqhfFxPRy3Bf7vr3h0cechB90XaQs=
|
||||
github.com/googleapis/enterprise-certificate-proxy v0.3.2/go.mod h1:VLSiSSBs/ksPL8kq3OBOQ6WRI2QnaFynd1DCjZ62+V0=
|
||||
github.com/invopop/jsonschema v0.14.0 h1:MHQqLhvpNUZfw+hM3AZDYK7jxO8FZoQeQM77g8iyZjg=
|
||||
github.com/invopop/jsonschema v0.14.0/go.mod h1:ygm6C2EaVNMBDPpaPlnOA2pFAxBnxGjFlMZABxm9n2I=
|
||||
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/go-telegram/bot v1.19.0 h1:tuvTQhgNietHFRN0HUDhuXsgfgkGSaO8WWwZQW3DMQg=
|
||||
github.com/go-telegram/bot v1.19.0/go.mod h1:i2TRs7fXWIeaceF3z7KzsMt/he0TwkVC680mvdTFYeM=
|
||||
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/mattn/go-sqlite3 v1.14.52 h1:wVbm2Qnf4OXkqhBTSPuCRZDRnxfbVrrmiCEroVdog8U=
|
||||
github.com/mattn/go-sqlite3 v1.14.52/go.mod h1:6JTjA44L93a0QCyJef5YvlPoKXntQPjzWv5gtm9sB6w=
|
||||
github.com/modelcontextprotocol/go-sdk v1.3.1 h1:TfqtNKOIWN4Z1oqmPAiWDC2Jq7K9OdJaooe0teoXASI=
|
||||
github.com/modelcontextprotocol/go-sdk v1.3.1/go.mod h1:DgVX498dMD8UJlseK1S5i1T4tFz2fkBk4xogC3D15nw=
|
||||
github.com/pb33f/go-yaml v0.1.1 h1:yWGmhVdwzionRnSc/hEJwcfq6PaZL2QdYzyReXBbX7Q=
|
||||
github.com/pb33f/go-yaml v0.1.1/go.mod h1:QvWdkHP5VjFYAeHMQUkaPQqREaQnBNOXmWzAjNrnSdo=
|
||||
github.com/pb33f/ordered-map/v2 v2.3.2 h1:wDyaZ2Pv9QLh64X4utCeD5Zoi9nIv2aW3lwv172O5PQ=
|
||||
github.com/pb33f/ordered-map/v2 v2.3.2/go.mod h1:1OhFrXu3OYw3kM+FXF+Ug3ugmmZ9LE8z0JfQMLvtV0w=
|
||||
github.com/segmentio/asm v1.1.3 h1:WM03sfUOENvvKexOLp+pCqgb/WDjsi7EK8gIsICtzhc=
|
||||
github.com/segmentio/asm v1.1.3/go.mod h1:Ld3L4ZXGNcSLRg4JBsZ3//1+f/TjYl0Mzen/DQy1EJg=
|
||||
github.com/segmentio/encoding v0.5.4 h1:OW1VRern8Nw6ITAtwSZ7Idrl3MXCFwXHPgqESYfvNt0=
|
||||
github.com/segmentio/encoding v0.5.4/go.mod h1:HS1ZKa3kSN32ZHVZ7ZLPLXWvOVIiZtyJnO1gPH1sKt0=
|
||||
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.8.0 h1:rm9oMkeprG/wHifkZP4EuDww9S0mvEaQ/p9cXTNT4cw=
|
||||
github.com/stretchr/objx v0.8.0/go.mod h1:xSzHnPHkdEv7vtgdQej8atlCuU1mdX+ELge8bn0hCsE=
|
||||
github.com/stretchr/testify v1.12.1 h1:EuwCh5fleGS7H32xRwO3wRGT7DxrDhLAT6FF8MpWDWE=
|
||||
github.com/stretchr/testify v1.12.1/go.mod h1:MDEgiDPPsNp5cuIrHPPCyornHKgEVbtFUmoNlxoYthg=
|
||||
github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
|
||||
github.com/tidwall/gjson v1.20.0 h1:+agJ3rEzKcCXKDKo0ml26UROcpcAlnZyjq+TjbY5Fto=
|
||||
github.com/tidwall/gjson v1.20.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.2 h1:dz1jrRuE7or/74V490B4/GP1pZm5WKlt2bgCP5A83w8=
|
||||
github.com/tidwall/pretty v1.2.2/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/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=
|
||||
github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4=
|
||||
go.opencensus.io v0.24.0 h1:y73uSU6J157QMP2kn2r30vwW1A2W2WFwSCGnAVxeaD0=
|
||||
go.opencensus.io v0.24.0/go.mod h1:vNK8G9p7aAivkbmorf4v+7Hgx+Zs0yY+0fOtgBfjQKo=
|
||||
go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64=
|
||||
go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y=
|
||||
go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.49.0 h1:4Pp6oUg3+e/6M4C0A/3kJ2VYa++dsWVTtGgLVj5xtHg=
|
||||
go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.49.0/go.mod h1:Mjt1i1INqiaoZOMGR1RIUJN+i3ChKoFRqzrRQhlkbs0=
|
||||
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.49.0 h1:jq9TW8u3so/bN+JPT166wjOI6/vQPF6Xe7nMNIltagk=
|
||||
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.49.0/go.mod h1:p8pYQP+m5XfbZm9fxtSKAbM6oIllS7s2AfxrChvc7iw=
|
||||
go.opentelemetry.io/otel v1.47.0 h1:j7ALJ/zgkS7Z6aeJW09p8VC9804bC+PpeTfCD4XPnOM=
|
||||
go.opentelemetry.io/otel v1.47.0/go.mod h1:8wS9O2qfXrYrzp6hIF/HOYJJf/wIhFPhR2xLuP+iXQU=
|
||||
go.opentelemetry.io/otel/log v1.47.0 h1:cOTS1CcLbSQeZKanGJ+0JpF/+t4PELi3O3bbl2lqCcI=
|
||||
go.opentelemetry.io/otel/log v1.47.0/go.mod h1:9byitSQ5pLC6PpqwGXjqdMKya6ZTswHRZh2vvXT33nw=
|
||||
go.opentelemetry.io/otel/metric v1.47.0 h1:4PptaldXx3Eat1XjMZ68pPJEs5wrhlemctZE9a3UdWY=
|
||||
go.opentelemetry.io/otel/metric v1.47.0/go.mod h1:ADGSXxRrXM6bjbvLo535EstVFlPpPYZm4LBKixjDHwU=
|
||||
go.opentelemetry.io/otel/sdk v1.33.0 h1:iax7M131HuAm9QkZotNHEfstof92xM+N8sr3uHXc2IM=
|
||||
go.opentelemetry.io/otel/sdk v1.33.0/go.mod h1:A1Q5oi7/9XaMlIWzPSxLRWOI8nG3FnzHJNbiENQuihM=
|
||||
go.opentelemetry.io/otel/trace v1.47.0 h1:JOjX/Oci8K94QHddo+bbfya/Ai/nf6/dt9ZfrFNWSrM=
|
||||
go.opentelemetry.io/otel/trace v1.47.0/go.mod h1:jNaSLa2PZEYFG6fRjJABAu+bw4FS08uDmPg28lTghu0=
|
||||
go.yaml.in/yaml/v3 v3.0.5 h1:N6y/pJk8buWs9NY5ERU2HSMfm+IuD/OtfdAnq6kESPw=
|
||||
go.yaml.in/yaml/v3 v3.0.5/go.mod h1:HVTZu1O7/Vkt2N+BFy8Zza+lnLsABggaTM2ZpNIGuKg=
|
||||
golang.org/x/crypto v0.40.0 h1:r4x+VvoG5Fm+eJcxMaY8CQM7Lb0l1lsmjGBQ6s8BfKM=
|
||||
golang.org/x/crypto v0.40.0/go.mod h1:Qr1vMER5WyS2dfPHAlsOj01wgLbsyWtFn/aY+5+ZdxY=
|
||||
golang.org/x/net v0.41.0 h1:vBTly1HeNPEn3wtREYfy4GZ/NECgw2Cnl+nK6Nz3uvw=
|
||||
golang.org/x/net v0.41.0/go.mod h1:B/K4NNqkfmg07DQYrbwvSluqCJOOXwUjeb/5lOisjbA=
|
||||
golang.org/x/oauth2 v0.30.0 h1:dnDm7JmhM45NNpd8FDDeLhK6FwqbOf4MLCM9zb1BOHI=
|
||||
golang.org/x/oauth2 v0.30.0/go.mod h1:B++QgG3ZKulg6sRPGD/mqlHQs5rB3Ml9erfeDY7xKlU=
|
||||
golang.org/x/sync v0.24.0 h1:iqiDrGoDgyuvAKrhLQthUEzU/R378N78gXbBjxed7oA=
|
||||
golang.org/x/sync v0.24.0/go.mod h1:sUUOizhqBxiL6pEWpqNLUiaJn1ShEbZ6BBqskPbjZm0=
|
||||
golang.org/x/sys v0.35.0 h1:vz1N37gP5bs89s7He8XuIYXpyY0+QlsKmzipCbUtyxI=
|
||||
golang.org/x/sys v0.35.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
|
||||
golang.org/x/text v0.43.0 h1:1QivrlhwAsnMOcqgOdl9my7s9OyZl8O5/UGUTJVXoWc=
|
||||
golang.org/x/text v0.43.0/go.mod h1:hbSIYA/amXcRXTFSZqh8tWHEoqRvvbKPJ2uz/KU+a+s=
|
||||
golang.org/x/time v0.16.0 h1:vMb6ptszcQMkcwiRTAuNNU50gom6++Q/6gY2hDM6VDE=
|
||||
golang.org/x/time v0.16.0/go.mod h1:rVKOqvZeKvrDKTQiAHJ7wmwP0RzleSphoEA9RcdLA0s=
|
||||
google.golang.org/api v0.189.0 h1:equMo30LypAkdkLMBqfeIqtyAnlyig1JSZArl4XPwdI=
|
||||
google.golang.org/api v0.189.0/go.mod h1:FLWGJKb0hb+pU2j+rJqwbnsF+ym+fQs73rbJ+KAUgy8=
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20240722135656-d784300faade h1:oCRSWfwGXQsqlVdErcyTt4A93Y8fo0/9D4b1gnI++qo=
|
||||
google.golang.org/genproto/googleapis/rpc v0.0.0-20240722135656-d784300faade/go.mod h1:Ue6ibwXGpU+dqIcODieyLOcgj7z8+IcskoNIgZxtrFY=
|
||||
google.golang.org/grpc v1.64.1 h1:LKtvyfbX3UGVPFcGqJ9ItpVWW6oN/2XqTxfAnwRRXiA=
|
||||
google.golang.org/grpc v1.64.1/go.mod h1:hiQF4LFZelK2WKaP6W0L92zGHtiQdZxk8CrSdvyjeP0=
|
||||
google.golang.org/protobuf v1.34.2 h1:6xV6lTsCfpGD21XK49h7MhtcApnLqkfYgPcdHftf6hg=
|
||||
google.golang.org/protobuf v1.34.2/go.mod h1:qYOHts0dSfpeUzUFpOMr/WGzszTmLH+DiWniOlNbLDw=
|
||||
gopkg.in/yaml.v2 v2.2.8 h1:obN1ZagJSUGI0Ek/LBmuj4SNLPfIny3KsKFopxRdj10=
|
||||
gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
|
||||
github.com/liushuangls/go-anthropic/v2 v2.17.1 h1:ca3oFzgQHs9/mJr+xx2XFQIYcQLM2rDCqieUx0g+8p4=
|
||||
github.com/liushuangls/go-anthropic/v2 v2.17.1/go.mod h1:a550cJXPoTG2FL3DvfKG2zzD5O2vjgvo4tHtoGPzFLU=
|
||||
github.com/mattn/go-sqlite3 v1.14.34 h1:3NtcvcUnFBPsuRcno8pUtupspG/GM+9nZ88zgJcp6Zk=
|
||||
github.com/mattn/go-sqlite3 v1.14.34/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y=
|
||||
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/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=
|
||||
golang.org/x/text v0.34.0 h1:oL/Qq0Kdaqxa1KbNeMKwQq0reLCCaFtqu2eNuSeNHbk=
|
||||
golang.org/x/text v0.34.0/go.mod h1:homfLqTYRFyVYemLBFl5GgL/DWEiH5wcsQ5gSh1yziA=
|
||||
golang.org/x/time v0.14.0 h1:MRx4UaLrDotUKUdCIqzPC48t1Y9hANFKIRpNx+Te8PI=
|
||||
golang.org/x/time v0.14.0/go.mod h1:eL/Oa2bBBK0TkX57Fyni+NgnyQQN4LitPmob2Hjnqw4=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
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.2 h1:3o8FXNo9v9S858gil+3LlZA1LkCOzgb4g5BL64FgaCo=
|
||||
gorm.io/gorm v1.31.2/go.mod h1:XyQVbO2k6YkOis7C2437jSit3SsDK72s7n7rsSHd+Gs=
|
||||
gorm.io/gorm v1.31.1 h1:7CA8FTFz/gRfgqgpeKIBcervUn3xSyPUmr6B2WXJ7kg=
|
||||
gorm.io/gorm v1.31.1/go.mod h1:XyQVbO2k6YkOis7C2437jSit3SsDK72s7n7rsSHd+Gs=
|
||||
|
||||
+131
-477
@@ -2,294 +2,16 @@ package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/anthropics/anthropic-sdk-go/v2"
|
||||
"github.com/go-telegram/bot"
|
||||
"github.com/go-telegram/bot/models"
|
||||
"golang.org/x/sync/errgroup"
|
||||
"github.com/liushuangls/go-anthropic/v2"
|
||||
)
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
// Voice is a committed turn: it waits for any running turn and is never
|
||||
// restarted, so a text sent right after a voice note is answered afterwards.
|
||||
turn, turnCtx := b.beginTurn(ctx, chatID)
|
||||
defer b.endTurn(turn)
|
||||
ctx = turnCtx
|
||||
|
||||
stopTyping := b.startChatAction(ctx, chatID, businessConnectionID, models.ChatActionTyping)
|
||||
defer stopTyping()
|
||||
|
||||
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()
|
||||
|
||||
contextMessages := b.snapshotForTurn(turn, b.getOrCreateChatMemory(chatID))
|
||||
response, err := b.getAnthropicResponse(ctx, chatID, contextMessages, false, username, firstName, lastName, isPremium, languageCode, messageTime, nil, 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
|
||||
}
|
||||
|
||||
// Switch the indicator once the model is done and synthesis begins, so the
|
||||
// client shows "recording audio" rather than "typing" for a voice reply.
|
||||
stopTyping()
|
||||
stopRecording := b.startChatAction(ctx, chatID, businessConnectionID, models.ChatActionUploadVoice)
|
||||
defer stopRecording()
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
// The Files API uploads depend on nothing the chat's running turn is doing,
|
||||
// so they overlap it instead of queueing behind it. On an album this is the
|
||||
// longest wait in the bot, and the typing indicator covers it.
|
||||
stopTyping := b.startChatAction(ctx, chatID, businessConnectionID, models.ChatActionTyping)
|
||||
finalUploaded, caption, err := b.uploadPhotoItems(ctx, items, chatID, businessConnectionID)
|
||||
stopTyping()
|
||||
if err != nil || len(finalUploaded) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
req := turnRequest{
|
||||
ctx: ctx, chatID: chatID, userID: userID,
|
||||
username: username, firstName: firstName, lastName: lastName,
|
||||
languageCode: languageCode, isPremium: isPremium, messageTime: messageTime,
|
||||
businessConnectionID: businessConnectionID,
|
||||
}
|
||||
// Photos wait for any running turn and are never restarted.
|
||||
turn, turnCtx := b.beginTurn(ctx, chatID)
|
||||
|
||||
b.runTurn(turn, turnCtx, req, func(ctx context.Context, out replyOutput) (string, error) {
|
||||
chatMemory := b.getOrCreateChatMemory(chatID)
|
||||
userMessage := b.createMessage(chatID, userID, username, "user", caption, true)
|
||||
userMessage.ImageFileIDs = finalUploaded
|
||||
userMessage.TelegramMessageID = items[0].ID
|
||||
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 "", errTurnSilent
|
||||
}
|
||||
b.addMessageToChatMemory(chatMemory, userMessage)
|
||||
|
||||
contextMessages := b.snapshotForTurn(turn, chatMemory)
|
||||
return b.getAnthropicResponse(
|
||||
ctx, chatID, contextMessages, false,
|
||||
username, firstName, lastName, isPremium, languageCode, messageTime,
|
||||
out.onSegment, out.onProgress,
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
// uploadPhotoItems uploads every photo in items to the Files API in parallel
|
||||
// and returns their file ids with the album caption. On any failure it deletes
|
||||
// the uploads that did succeed, tells the user, and returns errTurnSilent.
|
||||
func (b *Bot) uploadPhotoItems(ctx context.Context, items []*models.Message, chatID int64, businessConnectionID string) ([]string, string, error) {
|
||||
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 nil, "", errTurnSilent
|
||||
}
|
||||
finalUploaded := make([]string, 0, len(uploaded))
|
||||
for _, fid := range uploaded {
|
||||
if fid != "" {
|
||||
finalUploaded = append(finalUploaded, fid)
|
||||
}
|
||||
}
|
||||
return finalUploaded, caption, nil
|
||||
}
|
||||
|
||||
// respondToChat runs one assistant turn against the chat's current memory and
|
||||
// streams the reply back. Both the immediate path and the debounced flush go
|
||||
// through here, so a coalesced turn is byte-for-byte the same request as a
|
||||
// single-message one: the messages were already written to memory at intake, and
|
||||
// the model simply sees more of them.
|
||||
//
|
||||
// If a turn is already running for the chat, beginTextTurn decides whether this
|
||||
// flush is dropped, restarts it, or runs after it; see turn.go.
|
||||
func (b *Bot) respondToChat(req turnRequest) {
|
||||
wantID := b.newestUserMessageIDInChat(req.chatID)
|
||||
if wantID == 0 {
|
||||
// Nothing of the user's is in memory, for example because /clear
|
||||
// landed after this flush was collected. There is nothing to answer.
|
||||
InfoLogger.Printf("[%s] turn: chat %d has no user message to answer; skipping", b.config.ID, req.chatID)
|
||||
return
|
||||
}
|
||||
turn, turnCtx, ok := b.beginTextTurn(req, wantID)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
b.runTurn(turn, turnCtx, req, func(ctx context.Context, out replyOutput) (string, error) {
|
||||
contextMessages := b.snapshotForTurn(turn, b.getOrCreateChatMemory(req.chatID))
|
||||
return b.getAnthropicResponse(
|
||||
ctx, req.chatID, contextMessages, req.isEmojiOnly,
|
||||
req.username, req.firstName, req.lastName, req.isPremium, req.languageCode, req.messageTime,
|
||||
out.onSegment, out.onProgress,
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
func (b *Bot) anthropicErrorResponse(err error, userID int64) string {
|
||||
isElevated := b.hasScope(userID, ScopeModelSet)
|
||||
|
||||
if errors.Is(err, ErrModelNotFound) && isElevated {
|
||||
return fmt.Sprintf(
|
||||
"⚠️ Model `%s` is no longer available (deprecated or removed by Anthropic).\n"+
|
||||
"Use /set_model <model-id> to switch. Current models: https://platform.claude.com/docs/en/about-claude/models/overview",
|
||||
b.config.Model,
|
||||
)
|
||||
}
|
||||
|
||||
if isElevated {
|
||||
var apiErr *anthropic.Error
|
||||
if errors.As(err, &apiErr) {
|
||||
body := apiErr.RawJSON()
|
||||
if len(body) > 800 {
|
||||
body = body[:800] + "...(truncated)"
|
||||
}
|
||||
out := fmt.Sprintf("⚠️ Anthropic API error %d:\n%s", apiErr.StatusCode, body)
|
||||
if apiErr.RequestID != "" {
|
||||
out += fmt.Sprintf("\nRequest-ID: %s", apiErr.RequestID)
|
||||
}
|
||||
return out
|
||||
}
|
||||
return fmt.Sprintf("⚠️ Anthropic call failed: %v", err)
|
||||
}
|
||||
|
||||
return "I'm sorry, I'm having trouble processing your request right now."
|
||||
}
|
||||
|
||||
func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.Update) {
|
||||
if stopped := update.StoppedMessageGeneration; stopped != nil {
|
||||
b.handleStopGeneration(stopped.Chat.ID, stopped.DraftID)
|
||||
return
|
||||
}
|
||||
if edited := update.EditedMessage; edited != nil {
|
||||
b.handleEditedMessage(ctx, edited)
|
||||
return
|
||||
}
|
||||
if edited := update.EditedBusinessMessage; edited != nil {
|
||||
b.handleEditedMessage(ctx, edited)
|
||||
return
|
||||
}
|
||||
|
||||
var message *models.Message
|
||||
|
||||
if update.Message != nil {
|
||||
@@ -297,9 +19,11 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U
|
||||
} else if update.BusinessMessage != nil {
|
||||
message = update.BusinessMessage
|
||||
} else {
|
||||
// No message to process
|
||||
return
|
||||
}
|
||||
|
||||
// Extract businessConnectionID if available
|
||||
var businessConnectionID string
|
||||
if update.BusinessConnection != nil {
|
||||
businessConnectionID = update.BusinessConnection.ID
|
||||
@@ -308,6 +32,8 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U
|
||||
}
|
||||
|
||||
if message.From == nil {
|
||||
// Channel posts and some automated messages have no sender — ignore them.
|
||||
// see: https://core.telegram.org/bots/api#message
|
||||
return
|
||||
}
|
||||
|
||||
@@ -321,17 +47,50 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U
|
||||
messageTime := message.Date
|
||||
text := message.Text
|
||||
|
||||
// Check if it's a new chat
|
||||
isNewChatFlag := b.isNewChat(chatID)
|
||||
|
||||
// Screen incoming message (store to DB + add to chat memory)
|
||||
userMsg, err := b.screenIncomingMessage(message)
|
||||
if err != nil {
|
||||
ErrorLogger.Printf("Error storing user message: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
// Determine if the user is the owner
|
||||
var isOwner bool
|
||||
if b.db.Where("telegram_id = ? AND bot_id = ? AND is_owner = ?", userID, b.botID, true).First(&User{}).Error == nil {
|
||||
err = b.db.Where("telegram_id = ? AND bot_id = ? AND is_owner = ?", userID, b.botID, true).First(&User{}).Error
|
||||
if err == nil {
|
||||
isOwner = true
|
||||
}
|
||||
|
||||
// Get the chat memory which now contains the user's message
|
||||
chatMemory := b.getOrCreateChatMemory(chatID)
|
||||
contextMessages := b.prepareContextMessages(chatMemory)
|
||||
|
||||
if isNewChatFlag {
|
||||
|
||||
// Get response from Anthropic using the context messages
|
||||
response, err := b.getAnthropicResponse(ctx, contextMessages, true, isOwner, false, username, firstName, lastName, isPremium, languageCode, messageTime)
|
||||
if err != nil {
|
||||
ErrorLogger.Printf("Error getting Anthropic response: %v", err)
|
||||
// Use the same error message as in the non-new chat case
|
||||
response = "I'm sorry, I'm having trouble processing your request right now."
|
||||
}
|
||||
|
||||
// Send the AI-generated response or error message
|
||||
if err := b.sendResponse(ctx, chatID, response, businessConnectionID); err != nil {
|
||||
ErrorLogger.Printf("Error sending response: %v", err)
|
||||
return
|
||||
}
|
||||
} else {
|
||||
user, err := b.getOrCreateUser(userID, username, isOwner)
|
||||
if err != nil {
|
||||
ErrorLogger.Printf("Error getting or creating user: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
// Update the username if it's empty or has changed
|
||||
if user.Username != username {
|
||||
user.Username = username
|
||||
if err := b.db.Save(&user).Error; err != nil {
|
||||
@@ -339,50 +98,27 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U
|
||||
}
|
||||
}
|
||||
|
||||
// Media never waits on the text debounce window. Cancelling here does not
|
||||
// discard the buffered text: those messages are already in chat memory, so
|
||||
// the turn this media triggers answers them too.
|
||||
if message.MediaGroupID != "" && len(message.Photo) > 0 {
|
||||
b.cancelIntake(chatID)
|
||||
b.bufferAlbumItem(ctx, message, chatID, userID, username, firstName, lastName,
|
||||
isPremium, languageCode, messageTime, businessConnectionID)
|
||||
return
|
||||
}
|
||||
if len(message.Photo) > 0 {
|
||||
b.cancelIntake(chatID)
|
||||
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
|
||||
}
|
||||
|
||||
// Check if the message is a command
|
||||
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":
|
||||
// Parse command parameters
|
||||
parts := strings.Fields(message.Text)
|
||||
|
||||
// Default: show global stats
|
||||
if len(parts) == 1 {
|
||||
b.sendStats(ctx, chatID, userID, 0, businessConnectionID)
|
||||
return
|
||||
}
|
||||
|
||||
// Check for "user" parameter
|
||||
if len(parts) >= 2 && parts[1] == "user" {
|
||||
targetUserID := userID
|
||||
targetUserID := userID // Default to current user
|
||||
|
||||
// If a user ID is provided, parse it
|
||||
if len(parts) >= 3 {
|
||||
var parseErr error
|
||||
targetUserID, parseErr = strconv.ParseInt(parts[2], 10, 64)
|
||||
@@ -399,6 +135,7 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U
|
||||
return
|
||||
}
|
||||
|
||||
// Invalid parameter
|
||||
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)
|
||||
}
|
||||
@@ -433,33 +170,6 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U
|
||||
}
|
||||
b.clearChatHistory(ctx, chatID, userID, targetUserID, targetChatID, businessConnectionID, false)
|
||||
return
|
||||
case "/set_model":
|
||||
if !b.hasScope(userID, ScopeModelSet) {
|
||||
if err := b.sendResponse(ctx, chatID, "Permission denied. Only admins and owners can change the model.", businessConnectionID); err != nil {
|
||||
ErrorLogger.Printf("Error sending response: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
parts := strings.Fields(message.Text)
|
||||
if len(parts) < 2 || strings.TrimSpace(parts[1]) == "" {
|
||||
if err := b.sendResponse(ctx, chatID, "Usage: /set_model <model-id>", businessConnectionID); err != nil {
|
||||
ErrorLogger.Printf("Error sending response: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
newModel := strings.TrimSpace(parts[1])
|
||||
if err := b.config.PersistModel(newModel); err != nil {
|
||||
ErrorLogger.Printf("Failed to persist model change: %v", err)
|
||||
if err := b.sendResponse(ctx, chatID, fmt.Sprintf("Model updated in memory to `%s`, but failed to save to config file: %v", newModel, err), businessConnectionID); err != nil {
|
||||
ErrorLogger.Printf("Error sending response: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
InfoLogger.Printf("Model changed to %s by user %d", newModel, userID)
|
||||
if err := b.sendResponse(ctx, chatID, fmt.Sprintf("✅ Model updated to `%s` and saved to config.", newModel), businessConnectionID); err != nil {
|
||||
ErrorLogger.Printf("Error sending response: %v", err)
|
||||
}
|
||||
return
|
||||
case "/clear_hard":
|
||||
parts := strings.Fields(message.Text)
|
||||
var targetUserID, targetChatID int64
|
||||
@@ -492,125 +202,44 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U
|
||||
}
|
||||
}
|
||||
|
||||
// Rate limit check applies to all message types including stickers
|
||||
if !b.checkRateLimits(userID) {
|
||||
b.sendRateLimitExceededMessage(ctx, chatID, businessConnectionID)
|
||||
return
|
||||
}
|
||||
|
||||
if message.Voice != nil {
|
||||
b.cancelIntake(chatID)
|
||||
b.handleVoiceMessage(ctx, message, userMsg, chatID, userID, username, firstName, lastName, isPremium, languageCode, messageTime, businessConnectionID)
|
||||
return
|
||||
}
|
||||
|
||||
// Check if the message contains a sticker
|
||||
if message.Sticker != nil {
|
||||
b.cancelIntake(chatID)
|
||||
turn, turnCtx := b.beginTurn(ctx, chatID)
|
||||
defer b.endTurn(turn)
|
||||
contextMessages := b.snapshotForTurn(turn, b.getOrCreateChatMemory(chatID))
|
||||
b.handleStickerMessage(turnCtx, chatID, userMsg, message, contextMessages, businessConnectionID)
|
||||
b.handleStickerMessage(ctx, chatID, userMsg, message, businessConnectionID)
|
||||
return
|
||||
}
|
||||
|
||||
// Proceed only if the message contains text
|
||||
if text == "" {
|
||||
InfoLogger.Printf("Received a non-text message from user %d in chat %d", userID, chatID)
|
||||
return
|
||||
}
|
||||
|
||||
// Determine if the text contains only emojis
|
||||
isEmojiOnly := isOnlyEmojis(text)
|
||||
|
||||
// Plain text is the only thing that debounces: it is what users fragment
|
||||
// across several sends, and it is the only kind whose meaning survives being
|
||||
// read as one turn.
|
||||
if b.config.DebounceWindow() > 0 {
|
||||
b.bufferIntake(ctx, chatID, userID, username, firstName, lastName,
|
||||
isPremium, languageCode, messageTime, businessConnectionID, isEmojiOnly)
|
||||
// Prepare context messages for Anthropic
|
||||
chatMemory := b.getOrCreateChatMemory(chatID)
|
||||
contextMessages := b.prepareContextMessages(chatMemory)
|
||||
|
||||
// Get response from Anthropic
|
||||
response, err := b.getAnthropicResponse(ctx, contextMessages, false, isOwner, isEmojiOnly, username, firstName, lastName, isPremium, languageCode, messageTime) // isNewChat is false here
|
||||
if err != nil {
|
||||
ErrorLogger.Printf("Error getting Anthropic response: %v", err)
|
||||
response = "I'm sorry, I'm having trouble processing your request right now."
|
||||
}
|
||||
|
||||
// Send the response
|
||||
if err := b.sendResponse(ctx, chatID, response, businessConnectionID); err != nil {
|
||||
ErrorLogger.Printf("Error sending response: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
b.respondToChat(turnRequest{
|
||||
ctx: ctx, chatID: chatID, userID: userID, isEmojiOnly: isEmojiOnly,
|
||||
username: username, firstName: firstName, lastName: lastName,
|
||||
isPremium: isPremium, languageCode: languageCode, messageTime: messageTime,
|
||||
businessConnectionID: businessConnectionID,
|
||||
})
|
||||
}
|
||||
|
||||
// handleEditedMessage applies a Telegram edit to the stored copy of a user
|
||||
// message. Mobile clients deliver autocorrect fixes as edits within a few
|
||||
// seconds of the send, so the model should see the corrected text. What
|
||||
// happens next depends on where the message is:
|
||||
//
|
||||
// - still in its quiet window: the window restarts as if it had just arrived;
|
||||
// - being read by a turn that has shown nothing yet: that turn restarts;
|
||||
// - already answered, or shown: only the record is corrected.
|
||||
func (b *Bot) handleEditedMessage(ctx context.Context, edited *models.Message) {
|
||||
text := edited.Text
|
||||
if text == "" {
|
||||
text = edited.Caption
|
||||
}
|
||||
if edited.From == nil || text == "" {
|
||||
return
|
||||
}
|
||||
chatID := edited.Chat.ID
|
||||
|
||||
// Memory is the gate: only a message still in memory can influence a turn,
|
||||
// and the check is cheap. It also keeps edits from chats this bot never
|
||||
// stored off the database; a business connection delivers every edit on
|
||||
// the whole account.
|
||||
var msgID uint
|
||||
answered := false
|
||||
b.chatMemoriesMu.Lock()
|
||||
if mem, exists := b.chatMemories[chatID]; exists {
|
||||
for i := len(mem.Messages) - 1; i >= 0; i-- {
|
||||
m := &mem.Messages[i]
|
||||
if m.IsUser && m.TelegramMessageID == edited.ID {
|
||||
m.Text = text
|
||||
msgID = m.ID
|
||||
break
|
||||
}
|
||||
if !m.IsUser {
|
||||
answered = true
|
||||
}
|
||||
}
|
||||
}
|
||||
b.chatMemoriesMu.Unlock()
|
||||
if msgID == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
if err := b.db.Model(&Message{}).Where("id = ?", msgID).Update("text", text).Error; err != nil {
|
||||
ErrorLogger.Printf("[%s] applying edit to message %d in chat %d: %v", b.config.ID, edited.ID, chatID, err)
|
||||
}
|
||||
if answered {
|
||||
return
|
||||
}
|
||||
|
||||
isEmojiOnly := isOnlyEmojis(text)
|
||||
if b.touchIntake(chatID, isEmojiOnly) {
|
||||
InfoLogger.Printf("[%s] edit to message %d re-armed the quiet window for chat %d", b.config.ID, edited.ID, chatID)
|
||||
return
|
||||
}
|
||||
|
||||
req := turnRequest{
|
||||
ctx: ctx, chatID: chatID, userID: edited.From.ID,
|
||||
username: edited.From.Username, firstName: edited.From.FirstName, lastName: edited.From.LastName,
|
||||
languageCode: edited.From.LanguageCode, isPremium: edited.From.IsPremium, messageTime: edited.Date,
|
||||
businessConnectionID: edited.BusinessConnectionID, isEmojiOnly: isEmojiOnly,
|
||||
}
|
||||
if b.restartForEdit(req, msgID) {
|
||||
InfoLogger.Printf("[%s] edit to message %d restarted the running turn for chat %d", b.config.ID, edited.ID, chatID)
|
||||
}
|
||||
}
|
||||
|
||||
// handleStopGeneration routes Telegram's stopped_message_generation update,
|
||||
// sent when the user taps Stop under a streaming draft.
|
||||
func (b *Bot) handleStopGeneration(chatID int64, draftID int) {
|
||||
if b.stopTurn(chatID, draftID) {
|
||||
InfoLogger.Printf("[%s] user stopped generation in chat %d (draft %d)", b.config.ID, chatID, draftID)
|
||||
return
|
||||
}
|
||||
InfoLogger.Printf("[%s] stop for chat %d ignored: draft %d is not the running turn", b.config.ID, chatID, draftID)
|
||||
}
|
||||
|
||||
func (b *Bot) sendRateLimitExceededMessage(ctx context.Context, chatID int64, businessConnectionID string) {
|
||||
@@ -619,11 +248,14 @@ func (b *Bot) sendRateLimitExceededMessage(ctx context.Context, chatID int64, bu
|
||||
}
|
||||
}
|
||||
|
||||
func (b *Bot) handleStickerMessage(ctx context.Context, chatID int64, userMessage Message, message *models.Message, contextMessages []anthropic.BetaMessageParam, businessConnectionID string) {
|
||||
func (b *Bot) handleStickerMessage(ctx context.Context, chatID int64, userMessage Message, message *models.Message, businessConnectionID string) {
|
||||
// userMessage was already screened (stored + added to memory) by handleUpdate — do not call screenIncomingMessage again.
|
||||
|
||||
response, err := b.generateStickerResponse(ctx, userMessage, contextMessages, businessConnectionID)
|
||||
// Generate AI response about the sticker
|
||||
response, err := b.generateStickerResponse(ctx, userMessage)
|
||||
if err != nil {
|
||||
ErrorLogger.Printf("Error generating sticker response: %v", err)
|
||||
// Provide a fallback dynamic response based on sticker type
|
||||
if message.Sticker.IsAnimated {
|
||||
response = "Wow, that's a cool animated sticker!"
|
||||
} else if message.Sticker.IsVideo {
|
||||
@@ -633,22 +265,42 @@ func (b *Bot) handleStickerMessage(ctx context.Context, chatID int64, userMessag
|
||||
}
|
||||
}
|
||||
|
||||
// Send the response
|
||||
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, businessConnectionID string) (string, error) {
|
||||
stopTyping := b.startChatAction(ctx, message.ChatID, businessConnectionID, models.ChatActionTyping)
|
||||
defer stopTyping()
|
||||
|
||||
func (b *Bot) generateStickerResponse(ctx context.Context, message Message) (string, error) {
|
||||
// Example: Use the sticker type to generate a response
|
||||
if message.StickerFileID != "" {
|
||||
// Create message content with emoji information if available
|
||||
var messageContent string
|
||||
if message.StickerEmoji != "" {
|
||||
messageContent = fmt.Sprintf("User sent a sticker: %s", message.StickerEmoji)
|
||||
} else {
|
||||
messageContent = "User sent a sticker."
|
||||
}
|
||||
|
||||
// Prepare context with information about the sticker
|
||||
contextMessages := []anthropic.Message{
|
||||
{
|
||||
Role: anthropic.RoleUser,
|
||||
Content: []anthropic.MessageContent{
|
||||
anthropic.NewTextMessageContent(messageContent),
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
// Treat sticker messages like emoji messages to get emoji responses
|
||||
// Convert the timestamp to Unix time for the messageTime parameter
|
||||
messageTime := int(message.Timestamp.Unix())
|
||||
response, err := b.getAnthropicResponse(ctx, message.ChatID, contextMessages, true, message.Username, "", "", false, "", messageTime, nil, nil)
|
||||
response, err := b.getAnthropicResponse(ctx, contextMessages, false, false, true, message.Username, "", "", false, "", messageTime)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
return response, nil
|
||||
}
|
||||
|
||||
@@ -656,12 +308,10 @@ func (b *Bot) generateStickerResponse(ctx context.Context, message Message, cont
|
||||
}
|
||||
|
||||
func (b *Bot) clearChatHistory(ctx context.Context, chatID int64, currentUserID int64, targetUserID int64, targetChatID int64, businessConnectionID string, hardDelete bool) {
|
||||
// If targetUserID is provided and different from currentUserID, check permissions
|
||||
if targetUserID != 0 && targetUserID != currentUserID {
|
||||
requiredScope := ScopeHistoryClearAny
|
||||
if hardDelete {
|
||||
requiredScope = ScopeHistoryClearHardAny
|
||||
}
|
||||
if !b.hasScope(currentUserID, requiredScope) {
|
||||
// Check if the current user is an admin or owner
|
||||
if !b.isAdminOrOwner(currentUserID) {
|
||||
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)
|
||||
@@ -669,6 +319,7 @@ func (b *Bot) clearChatHistory(ctx context.Context, chatID int64, currentUserID
|
||||
return
|
||||
}
|
||||
|
||||
// Check if the target user exists
|
||||
var targetUser User
|
||||
err := b.db.Where("telegram_id = ? AND bot_id = ?", targetUserID, b.botID).First(&targetUser).Error
|
||||
if err != nil {
|
||||
@@ -679,38 +330,52 @@ func (b *Bot) clearChatHistory(ctx context.Context, chatID int64, currentUserID
|
||||
return
|
||||
}
|
||||
} else {
|
||||
// If no targetUserID is provided, set it to currentUserID
|
||||
targetUserID = currentUserID
|
||||
}
|
||||
|
||||
// Delete messages from the database
|
||||
//
|
||||
// Assumption: this bot is primarily used in private DMs, where each user's messages
|
||||
// are stored with chat_id == their own user_id — not the caller's chat_id. Scoping
|
||||
// a cross-user delete by the caller's chatID would therefore match 0 rows.
|
||||
//
|
||||
// When clearing another user's history the default (targetChatID == 0) deletes all
|
||||
// of that user's messages across every chat for this bot — the natural meaning of
|
||||
// "/clear <userID>" (wipe their entire history with the bot).
|
||||
//
|
||||
// When targetChatID != 0 the deletion is scoped to that specific chat, which is
|
||||
// useful for group moderation ("/clear <userID> <chatID>").
|
||||
var err error
|
||||
if hardDelete {
|
||||
// Permanently delete messages
|
||||
if targetUserID == currentUserID {
|
||||
err = b.hardDeleteScope(ctx, "chat_id = ? AND bot_id = ?", chatID, b.botID)
|
||||
// Deleting own messages — scope to the current chat only.
|
||||
err = b.db.Unscoped().Where("chat_id = ? AND bot_id = ? AND user_id = ?", chatID, b.botID, targetUserID).Delete(&Message{}).Error
|
||||
InfoLogger.Printf("User %d permanently deleted their own chat history in chat %d", currentUserID, chatID)
|
||||
} else {
|
||||
// Deleting another user's messages — scope bot-wide by default; chat-scoped if targetChatID given (see above).
|
||||
if targetChatID != 0 {
|
||||
err = b.hardDeleteScope(ctx, "chat_id = ? AND bot_id = ?", targetChatID, b.botID)
|
||||
err = b.db.Unscoped().Where("chat_id = ? AND bot_id = ? AND user_id = ?", targetChatID, b.botID, targetUserID).Delete(&Message{}).Error
|
||||
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)
|
||||
err = b.db.Unscoped().Where("bot_id = ? AND user_id = ?", b.botID, targetUserID).Delete(&Message{}).Error
|
||||
InfoLogger.Printf("Admin/owner %d permanently deleted all chat history for user %d", currentUserID, targetUserID)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Soft delete messages
|
||||
if targetUserID == currentUserID {
|
||||
err = b.db.Where("chat_id = ? AND bot_id = ?", chatID, b.botID).Delete(&Message{}).Error
|
||||
// Deleting own messages — scope to the current chat only.
|
||||
err = b.db.Where("chat_id = ? AND bot_id = ? AND user_id = ?", chatID, b.botID, targetUserID).Delete(&Message{}).Error
|
||||
InfoLogger.Printf("User %d soft deleted their own chat history in chat %d", currentUserID, chatID)
|
||||
} else {
|
||||
// Deleting another user's messages — scope bot-wide by default; chat-scoped if targetChatID given (see above).
|
||||
if targetChatID != 0 {
|
||||
err = b.db.Where("chat_id = ? AND bot_id = ?", targetChatID, b.botID).Delete(&Message{}).Error
|
||||
err = b.db.Where("chat_id = ? AND bot_id = ? AND user_id = ?", targetChatID, b.botID, targetUserID).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)
|
||||
}
|
||||
}
|
||||
@@ -724,39 +389,28 @@ func (b *Bot) clearChatHistory(ctx context.Context, chatID int64, currentUserID
|
||||
return
|
||||
}
|
||||
|
||||
// Drop any armed intake buffer for the same chat before clearing memory.
|
||||
// Otherwise the debounce timer fires moments later and repopulates the chat
|
||||
// with the very messages that were just deleted — the openclaw/openclaw#51046
|
||||
// failure mode, but with a privacy consequence rather than a stray reply.
|
||||
clearedChatID := chatID
|
||||
if targetUserID != currentUserID {
|
||||
clearedChatID = targetChatID
|
||||
if clearedChatID == 0 {
|
||||
clearedChatID = targetUserID
|
||||
}
|
||||
}
|
||||
if discarded := b.cancelIntake(clearedChatID); discarded > 0 {
|
||||
InfoLogger.Printf("[%s] discarded %d buffered message(s) for chat %d on history clear",
|
||||
b.config.ID, discarded, clearedChatID)
|
||||
}
|
||||
// Likewise a flush collected behind a running turn: rerun against the
|
||||
// emptied chat, it would call the model with no user message at all.
|
||||
b.discardRerun(clearedChatID)
|
||||
|
||||
// Evict the relevant in-memory cache entry so the next access rebuilds from
|
||||
// the now-clean DB. Applies to all cases: own history, cross-user
|
||||
// scoped to a specific chat, and bot-wide cross-user clear.
|
||||
b.chatMemoriesMu.Lock()
|
||||
if targetUserID == currentUserID {
|
||||
// Own history is always scoped to the current chat.
|
||||
delete(b.chatMemories, chatID)
|
||||
} else if targetChatID != 0 {
|
||||
// Admin cleared a specific chat — evict that chat's cache.
|
||||
delete(b.chatMemories, targetChatID)
|
||||
} else {
|
||||
// Bot-wide clear: primary use-case is DMs where chatID == userID.
|
||||
delete(b.chatMemories, targetUserID)
|
||||
}
|
||||
b.chatMemoriesMu.Unlock()
|
||||
|
||||
// Send a confirmation message
|
||||
var confirmationMessage string
|
||||
if targetUserID == currentUserID {
|
||||
confirmationMessage = "Your chat history has been cleared."
|
||||
} else {
|
||||
// Get the username of the target user if available
|
||||
var targetUser User
|
||||
err := b.db.Where("telegram_id = ? AND bot_id = ?", targetUserID, b.botID).First(&targetUser).Error
|
||||
if err == nil && targetUser.Username != "" {
|
||||
|
||||
+70
-274
@@ -2,10 +2,6 @@ package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -17,6 +13,7 @@ import (
|
||||
)
|
||||
|
||||
func TestHandleUpdate_NewChat(t *testing.T) {
|
||||
// Setup
|
||||
db := setupTestDB(t)
|
||||
mockClock := &MockClock{
|
||||
currentTime: time.Now(),
|
||||
@@ -24,7 +21,7 @@ func TestHandleUpdate_NewChat(t *testing.T) {
|
||||
|
||||
config := BotConfig{
|
||||
ID: "test_bot",
|
||||
OwnerTelegramID: 123,
|
||||
OwnerTelegramID: 123, // owner's ID
|
||||
TelegramToken: "test_token",
|
||||
MemorySize: 10,
|
||||
MessagePerHour: 5,
|
||||
@@ -36,6 +33,7 @@ func TestHandleUpdate_NewChat(t *testing.T) {
|
||||
|
||||
mockTgClient := &MockTelegramClient{}
|
||||
|
||||
// Create bot model first
|
||||
botModel := &BotModel{
|
||||
Identifier: config.ID,
|
||||
Name: config.ID,
|
||||
@@ -43,6 +41,7 @@ func TestHandleUpdate_NewChat(t *testing.T) {
|
||||
err := db.Create(botModel).Error
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Create bot config
|
||||
configModel := &ConfigModel{
|
||||
BotID: botModel.ID,
|
||||
MemorySize: config.MemorySize,
|
||||
@@ -56,34 +55,40 @@ func TestHandleUpdate_NewChat(t *testing.T) {
|
||||
err = db.Create(configModel).Error
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Create bot instance
|
||||
b, err := NewBot(db, config, mockClock, mockTgClient)
|
||||
assert.NoError(t, err)
|
||||
|
||||
testCases := []struct {
|
||||
name string
|
||||
userID int64
|
||||
wantSubstr string
|
||||
isOwner bool
|
||||
wantResp string
|
||||
}{
|
||||
{
|
||||
name: "Owner First Message",
|
||||
userID: 123,
|
||||
wantSubstr: "Anthropic call failed:",
|
||||
userID: 123, // owner's ID
|
||||
isOwner: true,
|
||||
wantResp: "I'm sorry, I'm having trouble processing your request right now.",
|
||||
},
|
||||
{
|
||||
name: "Regular User First Message",
|
||||
userID: 456,
|
||||
wantSubstr: "I'm sorry, I'm having trouble processing your request right now.",
|
||||
isOwner: false,
|
||||
wantResp: "I'm sorry, I'm having trouble processing your request right now.",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
// Setup mock response expectations for error case to test fallback messages
|
||||
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)
|
||||
assert.Equal(t, tc.wantResp, params.Text)
|
||||
return &models.Message{}, nil
|
||||
}
|
||||
|
||||
// Create update with new message
|
||||
update := &models.Update{
|
||||
Message: &models.Message{
|
||||
Chat: models.Chat{ID: tc.userID},
|
||||
@@ -95,23 +100,24 @@ func TestHandleUpdate_NewChat(t *testing.T) {
|
||||
},
|
||||
}
|
||||
|
||||
// Handle the update
|
||||
b.handleUpdate(context.Background(), nil, update)
|
||||
|
||||
// Verify message was stored
|
||||
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)
|
||||
|
||||
// Verify response was stored
|
||||
var respMsg Message
|
||||
err = db.Where("chat_id = ? AND is_user = ?", tc.userID, false).
|
||||
Order("timestamp DESC").
|
||||
First(&respMsg).Error
|
||||
err = db.Where("chat_id = ? AND is_user = ? AND text = ?", tc.userID, false, tc.wantResp).First(&respMsg).Error
|
||||
assert.NoError(t, err)
|
||||
assert.Contains(t, respMsg.Text, tc.wantSubstr)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestClearChatHistory(t *testing.T) {
|
||||
// Setup
|
||||
db := setupTestDB(t)
|
||||
mockClock := &MockClock{
|
||||
currentTime: time.Now(),
|
||||
@@ -119,7 +125,7 @@ func TestClearChatHistory(t *testing.T) {
|
||||
|
||||
config := BotConfig{
|
||||
ID: "test_bot",
|
||||
OwnerTelegramID: 123,
|
||||
OwnerTelegramID: 123, // owner's ID
|
||||
TelegramToken: "test_token",
|
||||
MemorySize: 10,
|
||||
MessagePerHour: 5,
|
||||
@@ -131,6 +137,7 @@ func TestClearChatHistory(t *testing.T) {
|
||||
|
||||
mockTgClient := &MockTelegramClient{}
|
||||
|
||||
// Create bot model first
|
||||
botModel := &BotModel{
|
||||
Identifier: config.ID,
|
||||
Name: config.ID,
|
||||
@@ -138,6 +145,7 @@ func TestClearChatHistory(t *testing.T) {
|
||||
err := db.Create(botModel).Error
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Create bot config
|
||||
configModel := &ConfigModel{
|
||||
BotID: botModel.ID,
|
||||
MemorySize: config.MemorySize,
|
||||
@@ -151,18 +159,22 @@ func TestClearChatHistory(t *testing.T) {
|
||||
err = db.Create(configModel).Error
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Create bot instance
|
||||
b, err := NewBot(db, config, mockClock, mockTgClient)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Create test users
|
||||
ownerID := int64(123)
|
||||
adminID := int64(456)
|
||||
regularUserID := int64(789)
|
||||
nonExistentUserID := int64(999)
|
||||
chatID := int64(1000)
|
||||
|
||||
// Create admin role
|
||||
adminRole, err := b.getRoleByName("admin")
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Create admin user
|
||||
adminUser := User{
|
||||
BotID: b.botID,
|
||||
TelegramID: adminID,
|
||||
@@ -174,6 +186,7 @@ func TestClearChatHistory(t *testing.T) {
|
||||
err = db.Create(&adminUser).Error
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Create regular user
|
||||
regularRole, err := b.getRoleByName("user")
|
||||
assert.NoError(t, err)
|
||||
regularUser := User{
|
||||
@@ -187,11 +200,15 @@ func TestClearChatHistory(t *testing.T) {
|
||||
err = db.Create(®ularUser).Error
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Create test messages for each user.
|
||||
// Each user's messages are stored with chat_id == their own user_id, mirroring
|
||||
// how Telegram private DMs work (chat_id == user_id in 1-on-1 bot conversations).
|
||||
// Using a shared artificial chatID here would mask the cross-user delete bug.
|
||||
for _, userID := range []int64{ownerID, adminID, regularUserID} {
|
||||
for i := 0; i < 5; i++ {
|
||||
message := Message{
|
||||
BotID: b.botID,
|
||||
ChatID: userID,
|
||||
ChatID: userID, // per-user chat, not a shared chatID
|
||||
UserID: userID,
|
||||
Username: "test",
|
||||
UserRole: "user",
|
||||
@@ -204,6 +221,7 @@ func TestClearChatHistory(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// Test cases
|
||||
testCases := []struct {
|
||||
name string
|
||||
currentUserID int64
|
||||
@@ -266,7 +284,7 @@ func TestClearChatHistory(t *testing.T) {
|
||||
targetUserID: adminID,
|
||||
hardDelete: false,
|
||||
expectedError: true,
|
||||
expectedCount: 5,
|
||||
expectedCount: 5, // Messages should remain
|
||||
expectedMsg: "Permission denied. Only admins and owners can clear other users' histories.",
|
||||
},
|
||||
{
|
||||
@@ -275,7 +293,7 @@ func TestClearChatHistory(t *testing.T) {
|
||||
targetUserID: nonExistentUserID,
|
||||
hardDelete: false,
|
||||
expectedError: true,
|
||||
expectedCount: 5,
|
||||
expectedCount: 5, // Messages should remain for admin
|
||||
expectedMsg: "User with ID 999 not found.",
|
||||
},
|
||||
{
|
||||
@@ -288,23 +306,29 @@ func TestClearChatHistory(t *testing.T) {
|
||||
expectedMsg: "Chat history for user @regular (ID: 789) has been cleared.",
|
||||
},
|
||||
{
|
||||
// targetChatID scopes the delete to a specific chat; messages in other chats survive.
|
||||
// We seed messages with ChatID == userID (per-user DM), so targeting a different chatID
|
||||
// should leave the user's messages untouched (expectedCount == 5).
|
||||
name: "Admin clears regular user's history scoped to non-matching chat",
|
||||
currentUserID: adminID,
|
||||
targetUserID: regularUserID,
|
||||
targetChatID: int64(9999),
|
||||
targetChatID: int64(9999), // a chat the user has no messages in
|
||||
hardDelete: false,
|
||||
expectedError: false,
|
||||
expectedCount: 5,
|
||||
expectedCount: 5, // messages in chat 789 are unaffected
|
||||
expectedMsg: "Chat history for user @regular (ID: 789) has been cleared.",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
// Reset messages for the test case
|
||||
if tc.name != "Owner hard deletes regular user's history" {
|
||||
// Delete all messages for the target user
|
||||
err = db.Where("user_id = ?", tc.targetUserID).Delete(&Message{}).Error
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Recreate messages for the target user
|
||||
for i := 0; i < 5; i++ {
|
||||
message := Message{
|
||||
BotID: b.botID,
|
||||
@@ -321,16 +345,20 @@ func TestClearChatHistory(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// Setup mock response expectations
|
||||
var sentMessage string
|
||||
mockTgClient.SendMessageFunc = func(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error) {
|
||||
sentMessage = params.Text
|
||||
return &models.Message{}, nil
|
||||
}
|
||||
|
||||
// Call the clearChatHistory method
|
||||
b.clearChatHistory(context.Background(), chatID, tc.currentUserID, tc.targetUserID, tc.targetChatID, tc.businessConnID, tc.hardDelete)
|
||||
|
||||
// Verify the response message
|
||||
assert.Equal(t, tc.expectedMsg, sentMessage)
|
||||
|
||||
// Count remaining messages for the target user
|
||||
var count int64
|
||||
if tc.hardDelete {
|
||||
db.Unscoped().Model(&Message{}).Where("user_id = ? AND chat_id = ?", tc.targetUserID, chatID).Count(&count)
|
||||
@@ -343,6 +371,7 @@ func TestClearChatHistory(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestStatsCommand(t *testing.T) {
|
||||
// Setup
|
||||
db := setupTestDB(t)
|
||||
mockClock := &MockClock{
|
||||
currentTime: time.Now(),
|
||||
@@ -350,7 +379,7 @@ func TestStatsCommand(t *testing.T) {
|
||||
|
||||
config := BotConfig{
|
||||
ID: "test_bot",
|
||||
OwnerTelegramID: 123,
|
||||
OwnerTelegramID: 123, // owner's ID
|
||||
TelegramToken: "test_token",
|
||||
MemorySize: 10,
|
||||
MessagePerHour: 5,
|
||||
@@ -362,6 +391,7 @@ func TestStatsCommand(t *testing.T) {
|
||||
|
||||
mockTgClient := &MockTelegramClient{}
|
||||
|
||||
// Create bot model first
|
||||
botModel := &BotModel{
|
||||
Identifier: config.ID,
|
||||
Name: config.ID,
|
||||
@@ -369,6 +399,7 @@ func TestStatsCommand(t *testing.T) {
|
||||
err := db.Create(botModel).Error
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Create bot config
|
||||
configModel := &ConfigModel{
|
||||
BotID: botModel.ID,
|
||||
MemorySize: config.MemorySize,
|
||||
@@ -382,17 +413,21 @@ func TestStatsCommand(t *testing.T) {
|
||||
err = db.Create(configModel).Error
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Create bot instance
|
||||
b, err := NewBot(db, config, mockClock, mockTgClient)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Create test users
|
||||
ownerID := int64(123)
|
||||
adminID := int64(456)
|
||||
regularUserID := int64(789)
|
||||
chatID := int64(1000)
|
||||
|
||||
// Create admin role
|
||||
adminRole, err := b.getRoleByName("admin")
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Create admin user
|
||||
adminUser := User{
|
||||
BotID: b.botID,
|
||||
TelegramID: adminID,
|
||||
@@ -404,6 +439,7 @@ func TestStatsCommand(t *testing.T) {
|
||||
err = db.Create(&adminUser).Error
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Create regular user
|
||||
regularRole, err := b.getRoleByName("user")
|
||||
assert.NoError(t, err)
|
||||
regularUser := User{
|
||||
@@ -417,8 +453,10 @@ func TestStatsCommand(t *testing.T) {
|
||||
err = db.Create(®ularUser).Error
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Create test messages for each user
|
||||
for _, userID := range []int64{ownerID, adminID, regularUserID} {
|
||||
for i := 0; i < 5; i++ {
|
||||
// User message
|
||||
userMessage := Message{
|
||||
BotID: b.botID,
|
||||
ChatID: chatID,
|
||||
@@ -432,6 +470,7 @@ func TestStatsCommand(t *testing.T) {
|
||||
err = db.Create(&userMessage).Error
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Bot response
|
||||
botMessage := Message{
|
||||
BotID: b.botID,
|
||||
ChatID: chatID,
|
||||
@@ -447,6 +486,7 @@ func TestStatsCommand(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// Test cases
|
||||
testCases := []struct {
|
||||
name string
|
||||
command string
|
||||
@@ -515,12 +555,14 @@ func TestStatsCommand(t *testing.T) {
|
||||
|
||||
for _, tc := range testCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
// Setup mock response expectations
|
||||
var sentMessage string
|
||||
mockTgClient.SendMessageFunc = func(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error) {
|
||||
sentMessage = params.Text
|
||||
return &models.Message{}, nil
|
||||
}
|
||||
|
||||
// Create update with command
|
||||
update := &models.Update{
|
||||
Message: &models.Message{
|
||||
Chat: models.Chat{ID: chatID},
|
||||
@@ -533,19 +575,22 @@ func TestStatsCommand(t *testing.T) {
|
||||
{
|
||||
Type: "bot_command",
|
||||
Offset: 0,
|
||||
Length: 6,
|
||||
Length: 6, // Length of "/stats"
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
// Handle the update
|
||||
b.handleUpdate(context.Background(), nil, update)
|
||||
|
||||
// Verify the response message contains the expected text
|
||||
assert.Contains(t, sentMessage, tc.expectedMsg)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Helper function to get username by ID for test
|
||||
func getUsernameByID(id int64) string {
|
||||
switch id {
|
||||
case 123:
|
||||
@@ -565,266 +610,17 @@ func setupTestDB(t *testing.T) *gorm.DB {
|
||||
t.Fatalf("Failed to open test database: %v", err)
|
||||
}
|
||||
|
||||
err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{}, &Scope{})
|
||||
// AutoMigrate the models
|
||||
err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{})
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to migrate database schema: %v", err)
|
||||
}
|
||||
|
||||
// Create default roles
|
||||
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-haiku-4-5",
|
||||
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-haiku-4-5", 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-haiku-4-5","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))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,143 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
)
|
||||
|
||||
// pendingIntake holds a chat's coalescing window. The buffered message bodies
|
||||
// are deliberately absent: screenIncomingMessage has already persisted each
|
||||
// message to the database and to chat memory by the time it is buffered, so the
|
||||
// flushed turn picks them all up from memory. What is kept here is the request
|
||||
// the turn needs, always refreshed to the most recent message in the batch.
|
||||
type pendingIntake struct {
|
||||
turnRequest
|
||||
allEmojiOnly bool // every message in the batch was emoji only
|
||||
count int
|
||||
seq uint64
|
||||
timer *time.Timer
|
||||
}
|
||||
|
||||
// bufferIntake holds a text message for the configured quiet window instead of
|
||||
// dispatching a turn immediately, resetting the window on each new message.
|
||||
// Rapid follow-ups therefore produce one reply rather than one per message.
|
||||
func (b *Bot) bufferIntake(
|
||||
ctx context.Context,
|
||||
chatID, userID int64,
|
||||
username, firstName, lastName string,
|
||||
isPremium bool,
|
||||
languageCode string,
|
||||
messageTime int,
|
||||
businessConnectionID string,
|
||||
isEmojiOnly bool,
|
||||
) {
|
||||
window := b.config.DebounceWindow()
|
||||
|
||||
b.intakeBuffersMu.Lock()
|
||||
defer b.intakeBuffersMu.Unlock()
|
||||
|
||||
pending, exists := b.intakeBuffers[chatID]
|
||||
if !exists {
|
||||
b.intakeSeq++
|
||||
pending = &pendingIntake{seq: b.intakeSeq, allEmojiOnly: true}
|
||||
b.intakeBuffers[chatID] = pending
|
||||
}
|
||||
|
||||
// Reply metadata tracks the most recent message in the batch.
|
||||
pending.turnRequest = turnRequest{
|
||||
ctx: ctx, chatID: chatID, userID: userID,
|
||||
username: username, firstName: firstName, lastName: lastName,
|
||||
isPremium: isPremium, languageCode: languageCode, messageTime: messageTime,
|
||||
businessConnectionID: businessConnectionID,
|
||||
}
|
||||
pending.allEmojiOnly = pending.allEmojiOnly && isEmojiOnly
|
||||
pending.count++
|
||||
|
||||
b.armIntakeLocked(pending, window)
|
||||
}
|
||||
|
||||
// armIntakeLocked starts or restarts the quiet-window timer. The caller holds
|
||||
// intakeBuffersMu.
|
||||
func (b *Bot) armIntakeLocked(pending *pendingIntake, window time.Duration) {
|
||||
if pending.timer != nil {
|
||||
pending.timer.Stop()
|
||||
}
|
||||
ctx, chatID, seq := pending.ctx, pending.chatID, pending.seq
|
||||
pending.timer = time.AfterFunc(window, func() {
|
||||
b.flushIntake(ctx, chatID, seq)
|
||||
})
|
||||
}
|
||||
|
||||
// touchIntake restarts the quiet window for a chat with a pending batch, as an
|
||||
// edit to one of its messages does, and reports whether there was one.
|
||||
// isEmojiOnly is the edited text's classification; an edit can only make the
|
||||
// batch less emoji-only, never more, since other messages are not re-read.
|
||||
func (b *Bot) touchIntake(chatID int64, isEmojiOnly bool) bool {
|
||||
b.intakeBuffersMu.Lock()
|
||||
defer b.intakeBuffersMu.Unlock()
|
||||
|
||||
pending, exists := b.intakeBuffers[chatID]
|
||||
if !exists {
|
||||
return false
|
||||
}
|
||||
pending.allEmojiOnly = pending.allEmojiOnly && isEmojiOnly
|
||||
b.armIntakeLocked(pending, b.config.DebounceWindow())
|
||||
return true
|
||||
}
|
||||
|
||||
// flushIntake dispatches the coalesced turn for a chat. seq guards against a
|
||||
// timer that had already fired before its Stop call landed: a stale goroutine
|
||||
// would otherwise flush a buffer belonging to a later batch.
|
||||
func (b *Bot) flushIntake(ctx context.Context, chatID int64, seq uint64) {
|
||||
b.intakeBuffersMu.Lock()
|
||||
pending, exists := b.intakeBuffers[chatID]
|
||||
if !exists || pending.seq != seq {
|
||||
b.intakeBuffersMu.Unlock()
|
||||
return
|
||||
}
|
||||
delete(b.intakeBuffers, chatID)
|
||||
captured := *pending
|
||||
b.intakeBuffersMu.Unlock()
|
||||
|
||||
if captured.count > 1 {
|
||||
InfoLogger.Printf("[%s] intake flush: coalesced %d messages into one turn for chat %d",
|
||||
b.config.ID, captured.count, chatID)
|
||||
}
|
||||
|
||||
req := captured.turnRequest
|
||||
req.ctx = ctx
|
||||
req.isEmojiOnly = captured.allEmojiOnly
|
||||
b.respondToChat(req)
|
||||
}
|
||||
|
||||
// cancelIntake drops a chat's pending buffer without dispatching, returning how
|
||||
// many messages were discarded.
|
||||
//
|
||||
// This is the fix for the class of bug in openclaw/openclaw#51046, where a stop
|
||||
// command aborted the running turn but left the debounce buffer armed, so the
|
||||
// timer fired afterwards and started the very turn the user had just cancelled.
|
||||
// Here the stakes are higher than a stray turn: /clear and /clear_hard delete
|
||||
// chat memory, and a surviving buffer would repopulate it moments later with
|
||||
// content the user asked to have removed.
|
||||
func (b *Bot) cancelIntake(chatID int64) int {
|
||||
b.intakeBuffersMu.Lock()
|
||||
defer b.intakeBuffersMu.Unlock()
|
||||
|
||||
pending, exists := b.intakeBuffers[chatID]
|
||||
if !exists {
|
||||
return 0
|
||||
}
|
||||
if pending.timer != nil {
|
||||
pending.timer.Stop()
|
||||
}
|
||||
delete(b.intakeBuffers, chatID)
|
||||
return pending.count
|
||||
}
|
||||
|
||||
// hasPendingIntake reports whether a chat currently holds a buffered batch.
|
||||
func (b *Bot) hasPendingIntake(chatID int64) bool {
|
||||
b.intakeBuffersMu.Lock()
|
||||
defer b.intakeBuffersMu.Unlock()
|
||||
_, exists := b.intakeBuffers[chatID]
|
||||
return exists
|
||||
}
|
||||
@@ -1,176 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/go-telegram/bot"
|
||||
"github.com/go-telegram/bot/models"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// bufferOnly parks a message in the intake buffer without letting the flush run,
|
||||
// by using a window long enough that no test waits it out.
|
||||
const bufferOnly = 10 * time.Second
|
||||
|
||||
func bufferText(b *Bot, chatID int64, isEmojiOnly bool) {
|
||||
b.bufferIntake(context.Background(), chatID, 555,
|
||||
"tester", "Test", "User", false, "en", int(time.Now().Unix()), "", isEmojiOnly)
|
||||
}
|
||||
|
||||
func TestBufferIntake_CoalescesIntoSingleTurn(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
b.config.DebounceMs = int(bufferOnly / time.Millisecond)
|
||||
|
||||
// The burst this whole feature exists for: one real question, then filler.
|
||||
for i := 0; i < 5; i++ {
|
||||
bufferText(b, 900, false)
|
||||
}
|
||||
|
||||
b.intakeBuffersMu.Lock()
|
||||
defer b.intakeBuffersMu.Unlock()
|
||||
assert.Len(t, b.intakeBuffers, 1, "rapid messages must share one buffer entry")
|
||||
assert.Equal(t, 5, b.intakeBuffers[900].count, "all five messages land in the same batch")
|
||||
}
|
||||
|
||||
// Separate chats must not share a window; one user's burst cannot delay another's.
|
||||
func TestBufferIntake_IsolatesChats(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
b.config.DebounceMs = int(bufferOnly / time.Millisecond)
|
||||
|
||||
bufferText(b, 910, false)
|
||||
bufferText(b, 911, false)
|
||||
bufferText(b, 911, false)
|
||||
|
||||
b.intakeBuffersMu.Lock()
|
||||
defer b.intakeBuffersMu.Unlock()
|
||||
assert.Equal(t, 1, b.intakeBuffers[910].count)
|
||||
assert.Equal(t, 2, b.intakeBuffers[911].count)
|
||||
}
|
||||
|
||||
func TestBufferIntake_ResetsWindowAndKeepsLatestMetadata(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
b.config.DebounceMs = int(bufferOnly / time.Millisecond)
|
||||
|
||||
b.bufferIntake(context.Background(), 901, 1, "first", "First", "", false, "en", 1000, "", true)
|
||||
b.bufferIntake(context.Background(), 901, 2, "second", "Second", "", true, "de", 2000, "biz-42", false)
|
||||
|
||||
b.intakeBuffersMu.Lock()
|
||||
defer b.intakeBuffersMu.Unlock()
|
||||
pending := b.intakeBuffers[901]
|
||||
require.NotNil(t, pending)
|
||||
|
||||
assert.Equal(t, 2, pending.count)
|
||||
// OpenClaw semantics: reply metadata follows the most recent message.
|
||||
assert.Equal(t, "second", pending.username)
|
||||
assert.Equal(t, int64(2), pending.userID)
|
||||
assert.Equal(t, "de", pending.languageCode)
|
||||
assert.Equal(t, 2000, pending.messageTime)
|
||||
assert.Equal(t, "biz-42", pending.businessConnectionID)
|
||||
assert.True(t, pending.isPremium)
|
||||
// One non-emoji message makes the whole coalesced turn non-emoji.
|
||||
assert.False(t, pending.allEmojiOnly)
|
||||
}
|
||||
|
||||
func TestBufferIntake_AllEmojiOnlySurvivesWhenEveryMessageIsEmoji(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
b.config.DebounceMs = int(bufferOnly / time.Millisecond)
|
||||
|
||||
bufferText(b, 902, true)
|
||||
bufferText(b, 902, true)
|
||||
|
||||
b.intakeBuffersMu.Lock()
|
||||
defer b.intakeBuffersMu.Unlock()
|
||||
assert.True(t, b.intakeBuffers[902].allEmojiOnly)
|
||||
}
|
||||
|
||||
func TestCancelIntake_DiscardsPendingBatch(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
b.config.DebounceMs = int(bufferOnly / time.Millisecond)
|
||||
|
||||
bufferText(b, 903, false)
|
||||
bufferText(b, 903, false)
|
||||
require.True(t, b.hasPendingIntake(903))
|
||||
|
||||
assert.Equal(t, 2, b.cancelIntake(903), "cancel reports how many messages it discarded")
|
||||
assert.False(t, b.hasPendingIntake(903))
|
||||
assert.Equal(t, 0, b.cancelIntake(903), "cancelling an empty chat is a no-op")
|
||||
}
|
||||
|
||||
// A cancelled buffer must never dispatch afterwards. This is the openclaw#51046
|
||||
// shape: the timer had already been armed when the cancel landed.
|
||||
func TestCancelIntake_TimerDoesNotFireAfterCancel(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
b.config.DebounceMs = 20
|
||||
|
||||
bufferText(b, 904, false)
|
||||
b.cancelIntake(904)
|
||||
|
||||
time.Sleep(80 * time.Millisecond)
|
||||
assert.False(t, b.hasPendingIntake(904), "cancelled buffer must not resurrect")
|
||||
}
|
||||
|
||||
// A stale timer from a cancelled batch must not flush a newer batch early.
|
||||
func TestFlushIntake_IgnoresSupersededSequence(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
b.config.DebounceMs = int(bufferOnly / time.Millisecond)
|
||||
|
||||
bufferText(b, 905, false)
|
||||
b.intakeBuffersMu.Lock()
|
||||
staleSeq := b.intakeBuffers[905].seq
|
||||
b.intakeBuffersMu.Unlock()
|
||||
|
||||
b.cancelIntake(905)
|
||||
bufferText(b, 905, false) // new batch, new sequence
|
||||
|
||||
// The old timer firing late must not consume the new batch.
|
||||
b.flushIntake(context.Background(), 905, staleSeq)
|
||||
assert.True(t, b.hasPendingIntake(905), "superseded flush must leave the newer batch armed")
|
||||
}
|
||||
|
||||
func TestClearChatHistory_CancelsPendingIntake(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, mockTg := setupBotForTest(t, 123)
|
||||
b.config.DebounceMs = int(bufferOnly / time.Millisecond)
|
||||
mockTg.SendMessageFunc = func(_ context.Context, _ *bot.SendMessageParams) (*models.Message, error) {
|
||||
return &models.Message{}, nil
|
||||
}
|
||||
|
||||
const chatID int64 = 906
|
||||
bufferText(b, chatID, false)
|
||||
require.True(t, b.hasPendingIntake(chatID))
|
||||
|
||||
b.clearChatHistory(context.Background(), chatID, 123, 0, 0, "", false)
|
||||
|
||||
assert.False(t, b.hasPendingIntake(chatID),
|
||||
"clearing history must disarm the buffer, or deleted messages get replayed into memory")
|
||||
}
|
||||
|
||||
func TestDebounceWindow(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
ms int
|
||||
want time.Duration
|
||||
}{
|
||||
{"unset disables debouncing", 0, 0},
|
||||
{"negative disables debouncing", -1, 0},
|
||||
{"positive converts to duration", 2500, 2500 * time.Millisecond},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
c := BotConfig{DebounceMs: tc.ms}
|
||||
assert.Equal(t, tc.want, c.DebounceWindow())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestCacheHistoryEnabled_DefaultsOn(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
assert.True(t, (&BotConfig{}).CacheHistoryEnabled(), "cache_history defaults to enabled")
|
||||
|
||||
off := false
|
||||
assert.False(t, (&BotConfig{CacheHistory: &off}).CacheHistoryEnabled())
|
||||
|
||||
on := true
|
||||
assert.True(t, (&BotConfig{CacheHistory: &on}).CacheHistoryEnabled())
|
||||
}
|
||||
@@ -5,13 +5,23 @@ import (
|
||||
"os"
|
||||
)
|
||||
|
||||
// For log management, use journalctl commands:
|
||||
// - View logs: journalctl -u telegram-bot
|
||||
// - Follow logs: journalctl -u telegram-bot -f
|
||||
// - View errors: journalctl -u telegram-bot -p err
|
||||
// Refer to the documentation for details on systemd unit setup.
|
||||
|
||||
// Initialize loggers for informational and error messages.
|
||||
var (
|
||||
InfoLogger *log.Logger
|
||||
ErrorLogger *log.Logger
|
||||
)
|
||||
|
||||
// initLoggers sets up separate loggers for stdout and stderr.
|
||||
func initLoggers() {
|
||||
// InfoLogger writes to stdout with specific flags.
|
||||
InfoLogger = log.New(os.Stdout, "INFO: ", log.Ldate|log.Ltime|log.Lshortfile)
|
||||
|
||||
// ErrorLogger writes to stderr with specific flags.
|
||||
ErrorLogger = log.New(os.Stderr, "ERROR: ", log.Ldate|log.Ltime|log.Lshortfile)
|
||||
}
|
||||
|
||||
@@ -8,30 +8,38 @@ import (
|
||||
)
|
||||
|
||||
func main() {
|
||||
// Initialize custom loggers
|
||||
initLoggers()
|
||||
|
||||
// Log the start of the application
|
||||
InfoLogger.Println("Starting Telegram Bot Application")
|
||||
|
||||
// Initialize database
|
||||
db, err := initDB()
|
||||
if err != nil {
|
||||
ErrorLogger.Fatalf("Error initializing database: %v", err)
|
||||
}
|
||||
|
||||
// Load all bot configurations
|
||||
configs, err := loadAllConfigs("config")
|
||||
if err != nil {
|
||||
ErrorLogger.Fatalf("Error loading configurations: %v", err)
|
||||
}
|
||||
|
||||
// Create a WaitGroup to manage goroutines
|
||||
var wg sync.WaitGroup
|
||||
|
||||
// Set up context with cancellation
|
||||
ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
|
||||
defer cancel()
|
||||
|
||||
// Initialize and start each bot
|
||||
for _, config := range configs {
|
||||
wg.Add(1)
|
||||
go func(cfg BotConfig) {
|
||||
defer wg.Done()
|
||||
|
||||
// Create Bot instance without TelegramClient initially
|
||||
realClock := RealClock{}
|
||||
bot, err := NewBot(db, cfg, realClock, nil)
|
||||
if err != nil {
|
||||
@@ -39,14 +47,17 @@ func main() {
|
||||
return
|
||||
}
|
||||
|
||||
// Start the bot in a separate goroutine
|
||||
go bot.Start(ctx)
|
||||
|
||||
// Keep the bot running until the context is cancelled
|
||||
<-ctx.Done()
|
||||
|
||||
InfoLogger.Printf("Bot %s stopped", cfg.ID)
|
||||
}(config)
|
||||
}
|
||||
|
||||
// Wait for all bots to finish
|
||||
wg.Wait()
|
||||
|
||||
InfoLogger.Println("All bots have stopped. Exiting application.")
|
||||
|
||||
@@ -8,10 +8,10 @@ import (
|
||||
|
||||
type BotModel struct {
|
||||
gorm.Model
|
||||
Identifier string `gorm:"uniqueIndex"`
|
||||
Identifier string `gorm:"uniqueIndex"` // Renamed from ID to Identifier
|
||||
Name string
|
||||
Configs []ConfigModel `gorm:"foreignKey:BotID;constraint:OnDelete:CASCADE"`
|
||||
Users []User `gorm:"foreignKey:BotID;constraint:OnDelete:CASCADE"`
|
||||
Users []User `gorm:"foreignKey:BotID;constraint:OnDelete:CASCADE"` // Associated users
|
||||
Messages []Message `gorm:"foreignKey:BotID;constraint:OnDelete:CASCADE"`
|
||||
}
|
||||
|
||||
@@ -22,7 +22,7 @@ type ConfigModel struct {
|
||||
MessagePerHour int `json:"messages_per_hour"`
|
||||
MessagePerDay int `json:"messages_per_day"`
|
||||
TempBanDuration string `json:"temp_ban_duration"`
|
||||
SystemPrompts string `json:"system_prompts"`
|
||||
SystemPrompts string `json:"system_prompts"` // Consider JSON string or separate table
|
||||
TelegramToken string `json:"telegram_token"`
|
||||
Active bool `json:"active"`
|
||||
}
|
||||
@@ -33,62 +33,40 @@ type Message struct {
|
||||
ChatID int64 `gorm:"index"`
|
||||
UserID int64 `gorm:"index"`
|
||||
Username string `gorm:"index"`
|
||||
UserRole string
|
||||
UserRole string // Store the role as a 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"`
|
||||
// TelegramMessageID is Telegram's message_id, kept so an edited_message
|
||||
// update can find and correct the stored copy. Zero on rows stored before
|
||||
// edits were tracked.
|
||||
TelegramMessageID int `gorm:"index"`
|
||||
StickerEmoji string // Store the emoji associated with the sticker
|
||||
DeletedAt gorm.DeletedAt `gorm:"index"` // Add soft delete field
|
||||
AnsweredOn *time.Time `gorm:"index"` // Tracks when a user message was answered (NULL for assistant messages and unanswered user messages)
|
||||
}
|
||||
|
||||
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"`
|
||||
BusinessConnectionID string // New field to store the business connection ID
|
||||
}
|
||||
|
||||
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"`
|
||||
BotID uint `gorm:"uniqueIndex:idx_user_bot;index"` // Foreign key to BotModel
|
||||
TelegramID int64 `gorm:"uniqueIndex:idx_user_bot;not null"` // Unique per (telegram_id, bot_id) pair
|
||||
Username string
|
||||
RoleID uint
|
||||
Role Role `gorm:"foreignKey:RoleID"`
|
||||
IsOwner bool `gorm:"default:false"`
|
||||
IsOwner bool `gorm:"default:false"` // Indicates if the user is the owner
|
||||
}
|
||||
|
||||
// idx_user_bot is a composite unique index on (bot_id, telegram_id),
|
||||
// allowing the same Telegram user to be registered independently on each bot.
|
||||
func (User) TableName() string {
|
||||
return "users"
|
||||
}
|
||||
|
||||
@@ -1,157 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/anthropics/anthropic-sdk-go/v2"
|
||||
"github.com/anthropics/anthropic-sdk-go/v2/packages/param"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestMarkTrailingCacheBreakpoint(t *testing.T) {
|
||||
t.Run("marks the final text block of the final turn", func(t *testing.T) {
|
||||
msgs := []anthropic.BetaMessageParam{
|
||||
anthropic.NewBetaUserMessage(anthropic.NewBetaTextBlock("older")),
|
||||
anthropic.NewBetaUserMessage(
|
||||
anthropic.NewBetaTextBlock("first"),
|
||||
anthropic.NewBetaTextBlock("last"),
|
||||
),
|
||||
}
|
||||
|
||||
markTrailingCacheBreakpoint(msgs)
|
||||
|
||||
last := msgs[1].OfUser.Content[1].OfText
|
||||
require.NotNil(t, last)
|
||||
assert.False(t, param.IsOmitted(last.CacheControl),
|
||||
"the trailing block carries the breakpoint")
|
||||
|
||||
// Everything earlier stays unmarked: one breakpoint, not one per block.
|
||||
assert.True(t, param.IsOmitted(msgs[1].OfUser.Content[0].OfText.CacheControl))
|
||||
assert.True(t, param.IsOmitted(msgs[0].OfUser.Content[0].OfText.CacheControl))
|
||||
})
|
||||
|
||||
t.Run("marks a trailing image block", func(t *testing.T) {
|
||||
msgs := []anthropic.BetaMessageParam{
|
||||
anthropic.NewBetaUserMessage(
|
||||
anthropic.NewBetaImageBlock(anthropic.BetaFileImageSourceParam{FileID: "file_1"}),
|
||||
),
|
||||
}
|
||||
|
||||
markTrailingCacheBreakpoint(msgs)
|
||||
|
||||
require.NotNil(t, msgs[0].OfUser.Content[0].OfImage)
|
||||
assert.False(t, param.IsOmitted(msgs[0].OfUser.Content[0].OfImage.CacheControl))
|
||||
})
|
||||
|
||||
t.Run("tolerates empty input", func(t *testing.T) {
|
||||
assert.NotPanics(t, func() { markTrailingCacheBreakpoint(nil) })
|
||||
assert.NotPanics(t, func() {
|
||||
markTrailingCacheBreakpoint([]anthropic.BetaMessageParam{{}})
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func TestPrepareContextMessages_CacheHistoryToggle(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
memory := &ChatMemory{
|
||||
Messages: []Message{
|
||||
{IsUser: true, Text: "hello"},
|
||||
{IsUser: false, Text: "hi there"},
|
||||
},
|
||||
Size: 10,
|
||||
}
|
||||
|
||||
t.Run("enabled by default", func(t *testing.T) {
|
||||
msgs := b.prepareContextMessages(memory)
|
||||
require.Len(t, msgs, 2)
|
||||
require.NotNil(t, msgs[0].OfUser, "user memory becomes a user turn")
|
||||
require.NotNil(t, msgs[1].OfAssistant, "bot memory becomes an assistant turn")
|
||||
assert.False(t, param.IsOmitted(msgs[1].OfAssistant.Content[0].OfText.CacheControl))
|
||||
})
|
||||
|
||||
t.Run("opt-out leaves history unmarked", func(t *testing.T) {
|
||||
off := false
|
||||
b.config.CacheHistory = &off
|
||||
defer func() { b.config.CacheHistory = nil }()
|
||||
|
||||
msgs := b.prepareContextMessages(memory)
|
||||
require.Len(t, msgs, 2)
|
||||
assert.True(t, param.IsOmitted(msgs[1].OfAssistant.Content[0].OfText.CacheControl))
|
||||
})
|
||||
}
|
||||
|
||||
func TestValidateConfig_DebounceMs(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
base := func(ms int) *BotConfig {
|
||||
return &BotConfig{
|
||||
ID: "b",
|
||||
TelegramToken: "t",
|
||||
Model: "claude-sonnet-4-6",
|
||||
MessagePerHour: 1,
|
||||
MessagePerDay: 1,
|
||||
DebounceMs: ms,
|
||||
}
|
||||
}
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
ms int
|
||||
wantErr bool
|
||||
}{
|
||||
{"omitted is valid", 0, false},
|
||||
{"typical chat window is valid", 2500, false},
|
||||
{"at the ceiling is valid", maxDebounceMs, false},
|
||||
{"negative is rejected", -1, true},
|
||||
{"above the ceiling is rejected", maxDebounceMs + 1, true},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
err := validateConfig(base(tc.ms), map[string]bool{}, map[string]bool{})
|
||||
if tc.wantErr {
|
||||
assert.Error(t, err)
|
||||
return
|
||||
}
|
||||
assert.NoError(t, err)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateConfig_Effort(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
base := func(effort string) *BotConfig {
|
||||
return &BotConfig{
|
||||
ID: "b",
|
||||
TelegramToken: "t",
|
||||
Model: "claude-sonnet-5",
|
||||
MessagePerHour: 1,
|
||||
MessagePerDay: 1,
|
||||
Effort: effort,
|
||||
}
|
||||
}
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
effort string
|
||||
wantErr bool
|
||||
}{
|
||||
{"omitted is valid", "", false},
|
||||
{"low is valid", "low", false},
|
||||
{"medium is valid", "medium", false},
|
||||
{"high is valid", "high", false},
|
||||
{"xhigh is valid", "xhigh", false},
|
||||
{"max is valid", "max", false},
|
||||
{"unknown level is rejected", "turbo", true},
|
||||
{"wrong case is rejected", "Low", true},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
err := validateConfig(base(tc.effort), map[string]bool{}, map[string]bool{})
|
||||
if tc.wantErr {
|
||||
assert.Error(t, err)
|
||||
return
|
||||
}
|
||||
assert.NoError(t, err)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -33,20 +33,26 @@ func (b *Bot) checkRateLimits(userID int64) bool {
|
||||
|
||||
now := limiter.clock.Now()
|
||||
|
||||
// Check if the user is currently banned
|
||||
if now.Before(limiter.banUntil) {
|
||||
return false
|
||||
}
|
||||
|
||||
// Reset hourly limiter if an hour has passed since the last reset
|
||||
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
|
||||
}
|
||||
|
||||
// Reset daily limiter if 24 hours have passed since the last reset
|
||||
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
|
||||
}
|
||||
|
||||
// Check if the message exceeds rate limits.
|
||||
// Reserve from both limiters first, then cancel both if either is over budget.
|
||||
// This prevents consuming a token from one limiter when the other rejects.
|
||||
dailyRes := limiter.dailyLimiter.ReserveN(now, 1)
|
||||
hourlyRes := limiter.hourlyLimiter.ReserveN(now, 1)
|
||||
if dailyRes.DelayFrom(now) > 0 || hourlyRes.DelayFrom(now) > 0 {
|
||||
@@ -54,6 +60,7 @@ func (b *Bot) checkRateLimits(userID int64) bool {
|
||||
hourlyRes.CancelAt(now)
|
||||
banDuration, err := time.ParseDuration(b.config.TempBanDuration)
|
||||
if err != nil {
|
||||
// If parsing fails, default to a 24-hour ban
|
||||
banDuration = 24 * time.Hour
|
||||
}
|
||||
limiter.banUntil = now.Add(banDuration)
|
||||
|
||||
+23
-5
@@ -5,22 +5,27 @@ import (
|
||||
"time"
|
||||
)
|
||||
|
||||
// TestCheckRateLimits tests the checkRateLimits method of the Bot.
|
||||
// It verifies that users are allowed or denied based on their message rates.
|
||||
func TestCheckRateLimits(t *testing.T) {
|
||||
// Create a mock clock starting at a fixed time
|
||||
mockClock := &MockClock{
|
||||
currentTime: time.Date(2023, 10, 1, 0, 0, 0, 0, time.UTC),
|
||||
}
|
||||
|
||||
// Create a mock configuration with reduced timeframes for testing
|
||||
config := BotConfig{
|
||||
ID: "bot1",
|
||||
MemorySize: 10,
|
||||
MessagePerHour: 5,
|
||||
MessagePerDay: 10,
|
||||
TempBanDuration: "1m",
|
||||
MessagePerHour: 5, // Allow 5 messages per hour
|
||||
MessagePerDay: 10, // Allow 10 messages per day
|
||||
TempBanDuration: "1m", // Temporary ban duration of 1 minute for testing
|
||||
SystemPrompts: make(map[string]string),
|
||||
TelegramToken: "YOUR_TELEGRAM_BOT_TOKEN",
|
||||
OwnerTelegramID: 123456789,
|
||||
}
|
||||
|
||||
// Initialize the Bot with mock data and MockClock
|
||||
bot := &Bot{
|
||||
config: config,
|
||||
userLimiters: make(map[int64]*userLimiter),
|
||||
@@ -29,39 +34,52 @@ func TestCheckRateLimits(t *testing.T) {
|
||||
|
||||
userID := int64(12345)
|
||||
|
||||
// Helper function to simulate message sending
|
||||
sendMessage := func() bool {
|
||||
return bot.checkRateLimits(userID)
|
||||
}
|
||||
|
||||
// Send 5 messages within the hourly limit
|
||||
for i := 0; i < config.MessagePerHour; i++ {
|
||||
if !sendMessage() {
|
||||
t.Errorf("Expected message %d to be allowed", i+1)
|
||||
}
|
||||
}
|
||||
|
||||
// 6th message should exceed the hourly limit and trigger a ban
|
||||
if sendMessage() {
|
||||
t.Errorf("Expected message to be denied due to hourly limit exceeded")
|
||||
}
|
||||
|
||||
// Attempt to send another message immediately, should still be banned
|
||||
if sendMessage() {
|
||||
t.Errorf("Expected message to be denied while user is banned")
|
||||
}
|
||||
|
||||
mockClock.Advance(time.Minute)
|
||||
// Fast-forward time by TempBanDuration to lift the ban
|
||||
mockClock.Advance(time.Minute) // Banned for 1 minute
|
||||
|
||||
mockClock.Advance(time.Hour)
|
||||
// Advance time to allow hourly limiter to replenish
|
||||
mockClock.Advance(time.Hour) // Advance by 1 hour
|
||||
|
||||
// Send another message, should be allowed now
|
||||
if !sendMessage() {
|
||||
t.Errorf("Expected message to be allowed after ban duration")
|
||||
}
|
||||
|
||||
// Send additional messages to reach the daily limit
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
// Attempt to exceed the daily limit
|
||||
if sendMessage() {
|
||||
t.Errorf("Expected message to be denied due to daily limit exceeded")
|
||||
}
|
||||
}
|
||||
|
||||
// To ensure thread safety and avoid race conditions during testing,
|
||||
// you can run the tests with the `-race` flag:
|
||||
// go test -race -v
|
||||
|
||||
+3
-6
@@ -1,3 +1,4 @@
|
||||
// telegram_client.go
|
||||
package main
|
||||
|
||||
import (
|
||||
@@ -7,13 +8,9 @@ import (
|
||||
"github.com/go-telegram/bot/models"
|
||||
)
|
||||
|
||||
// TelegramClient defines the methods required from the Telegram bot.
|
||||
type TelegramClient interface {
|
||||
SendMessage(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error)
|
||||
SendMessageDraft(ctx context.Context, params *bot.SendMessageDraftParams) (bool, error)
|
||||
SendAudio(ctx context.Context, params *bot.SendAudioParams) (*models.Message, error)
|
||||
SendChatAction(ctx context.Context, params *bot.SendChatActionParams) (bool, 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)
|
||||
// Add other methods if needed.
|
||||
}
|
||||
|
||||
+4
-48
@@ -1,3 +1,4 @@
|
||||
// telegram_client_mock.go
|
||||
package main
|
||||
|
||||
import (
|
||||
@@ -8,18 +9,14 @@ import (
|
||||
"github.com/stretchr/testify/mock"
|
||||
)
|
||||
|
||||
// MockTelegramClient is a mock implementation of TelegramClient for testing.
|
||||
type MockTelegramClient struct {
|
||||
mock.Mock
|
||||
SendMessageFunc func(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error)
|
||||
SendMessageDraftFunc func(ctx context.Context, params *bot.SendMessageDraftParams) (bool, error)
|
||||
SendAudioFunc func(ctx context.Context, params *bot.SendAudioParams) (*models.Message, error)
|
||||
SendChatActionFunc func(ctx context.Context, params *bot.SendChatActionParams) (bool, 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)
|
||||
}
|
||||
|
||||
// SendMessage mocks sending a message.
|
||||
func (m *MockTelegramClient) SendMessage(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error) {
|
||||
if m.SendMessageFunc != nil {
|
||||
return m.SendMessageFunc(ctx, params)
|
||||
@@ -31,48 +28,7 @@ func (m *MockTelegramClient) SendMessage(ctx context.Context, params *bot.SendMe
|
||||
return nil, args.Error(1)
|
||||
}
|
||||
|
||||
func (m *MockTelegramClient) SendMessageDraft(ctx context.Context, params *bot.SendMessageDraftParams) (bool, error) {
|
||||
if m.SendMessageDraftFunc != nil {
|
||||
return m.SendMessageDraftFunc(ctx, params)
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
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) SendChatAction(ctx context.Context, params *bot.SendChatActionParams) (bool, error) {
|
||||
if m.SendChatActionFunc != nil {
|
||||
return m.SendChatActionFunc(ctx, params)
|
||||
}
|
||||
return true, 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 ""
|
||||
}
|
||||
|
||||
// Start mocks starting the Telegram client.
|
||||
func (m *MockTelegramClient) Start(ctx context.Context) {
|
||||
if m.StartFunc != nil {
|
||||
m.StartFunc(ctx)
|
||||
|
||||
@@ -1,52 +0,0 @@
|
||||
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
|
||||
}
|
||||
@@ -1,53 +0,0 @@
|
||||
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)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1,321 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/go-telegram/bot/models"
|
||||
)
|
||||
|
||||
// cancelReason records why a running turn's context was cancelled, so the
|
||||
// turn can tell a deliberate interruption from a real failure.
|
||||
type cancelReason int
|
||||
|
||||
const (
|
||||
cancelNone cancelReason = iota
|
||||
// cancelRestart: a newer message or an edit arrived before anything was
|
||||
// shown to the user. The turn is abandoned and rerun with it included.
|
||||
cancelRestart
|
||||
// cancelStop: the user pressed Telegram's Stop button under the draft.
|
||||
cancelStop
|
||||
)
|
||||
|
||||
// errTurnSilent marks a failure runTurn must not report: the turn body has
|
||||
// either already told the user, or deliberately has nothing to say.
|
||||
var errTurnSilent = errors.New("turn: failure handled by the turn body")
|
||||
|
||||
// turnRequest is everything one assistant turn needs to run and reply.
|
||||
type turnRequest struct {
|
||||
ctx context.Context
|
||||
chatID, userID int64
|
||||
username, firstName, lastName, languageCode string
|
||||
isPremium bool
|
||||
messageTime int
|
||||
businessConnectionID string
|
||||
isEmojiOnly bool
|
||||
}
|
||||
|
||||
// chatTurn is one running assistant turn. At most one exists per chat.
|
||||
//
|
||||
// This is the second half of the debounce design. The intake buffer decides
|
||||
// when a burst is quiet enough to answer; this decides what happens to a
|
||||
// message that lands after that, while the model is already working. Without
|
||||
// it every late fragment started another concurrent turn, and a slow typist
|
||||
// got one reply per fragment, each split into several messages.
|
||||
type chatTurn struct {
|
||||
seq uint64
|
||||
chatID int64
|
||||
cancel context.CancelFunc
|
||||
done chan struct{}
|
||||
|
||||
// Guarded by Bot.turnsMu.
|
||||
draftID string // Telegram draft id shown for this turn, "" if none
|
||||
coveredID uint // highest user Message.ID this turn's context includes
|
||||
committed bool // must not be restarted: output shown, or work not cheaply redone
|
||||
reason cancelReason // why cancel was called, if it was
|
||||
rerun *turnRequest // a flush that arrived mid-turn and may still need a turn
|
||||
rerunWant uint // the newest user message that flush wanted answered
|
||||
}
|
||||
|
||||
func (t *chatTurn) draftIDString() string { return strconv.FormatUint(t.seq, 10) }
|
||||
|
||||
func (b *Bot) registerTurnLocked(parent context.Context, chatID int64, coveredID uint, committed bool) (*chatTurn, context.Context) {
|
||||
ctx, cancel := context.WithCancel(parent)
|
||||
b.turnSeq++
|
||||
t := &chatTurn{
|
||||
seq: b.turnSeq,
|
||||
chatID: chatID,
|
||||
cancel: cancel,
|
||||
done: make(chan struct{}),
|
||||
coveredID: coveredID,
|
||||
committed: committed,
|
||||
}
|
||||
b.turns[chatID] = t
|
||||
return t, ctx
|
||||
}
|
||||
|
||||
// beginTextTurn applies the late-message policy for a text flush. wantID is the
|
||||
// newest user message the flush wants answered. ok=false means the caller must
|
||||
// not run a turn, because one of three things happened instead:
|
||||
//
|
||||
// - dropped: the running turn's context already includes wantID, so its reply
|
||||
// answers this flush too;
|
||||
// - restart: nothing has been shown yet, so the running turn is cancelled and
|
||||
// recorded to rerun with the new message in context. This is the only case
|
||||
// where tokens are wasted, and it is bounded by the time to first output;
|
||||
// - collect: output is already on screen, so the flush is recorded to run as
|
||||
// its own turn once this one finishes, unless the finished turn turns out to
|
||||
// have covered it after all (see finishTurn).
|
||||
func (b *Bot) beginTextTurn(req turnRequest, wantID uint) (*chatTurn, context.Context, bool) {
|
||||
b.turnsMu.Lock()
|
||||
defer b.turnsMu.Unlock()
|
||||
|
||||
if running, exists := b.turns[req.chatID]; exists {
|
||||
switch {
|
||||
case wantID <= running.coveredID:
|
||||
InfoLogger.Printf("[%s] turn: chat %d flush dropped, message %d already covered by running turn %d",
|
||||
b.config.ID, req.chatID, wantID, running.seq)
|
||||
case !running.committed:
|
||||
running.reason = cancelRestart
|
||||
running.rerun, running.rerunWant = &req, wantID
|
||||
running.cancel()
|
||||
InfoLogger.Printf("[%s] turn: chat %d restarting turn %d, message %d arrived before any output",
|
||||
b.config.ID, req.chatID, running.seq, wantID)
|
||||
default:
|
||||
running.rerun, running.rerunWant = &req, wantID
|
||||
InfoLogger.Printf("[%s] turn: chat %d collected message %d, runs after turn %d",
|
||||
b.config.ID, req.chatID, wantID, running.seq)
|
||||
}
|
||||
return nil, nil, false
|
||||
}
|
||||
|
||||
turn, ctx := b.registerTurnLocked(req.ctx, req.chatID, wantID, false)
|
||||
return turn, ctx, true
|
||||
}
|
||||
|
||||
// beginTurn registers a turn for a media handler, waiting for any running turn
|
||||
// on the chat to finish first. Media is never dropped or coalesced, and the
|
||||
// turn is committed from the start: cancelling half-finished uploads or a voice
|
||||
// transcription to save a few tokens is not worth the mess.
|
||||
//
|
||||
// The wait does not watch parent: the running turn derives from the same
|
||||
// long-lived context, so if parent is cancelled it finishes promptly anyway,
|
||||
// and registering over it would put two turns on one chat.
|
||||
func (b *Bot) beginTurn(parent context.Context, chatID int64) (*chatTurn, context.Context) {
|
||||
for {
|
||||
b.turnsMu.Lock()
|
||||
running, exists := b.turns[chatID]
|
||||
if !exists {
|
||||
turn, ctx := b.registerTurnLocked(parent, chatID, 0, true)
|
||||
b.turnsMu.Unlock()
|
||||
return turn, ctx
|
||||
}
|
||||
b.turnsMu.Unlock()
|
||||
<-running.done
|
||||
}
|
||||
}
|
||||
|
||||
// markCovered records the newest user message the turn's context includes.
|
||||
func (b *Bot) markCovered(turn *chatTurn, coveredID uint) {
|
||||
b.turnsMu.Lock()
|
||||
defer b.turnsMu.Unlock()
|
||||
if coveredID > turn.coveredID {
|
||||
turn.coveredID = coveredID
|
||||
}
|
||||
}
|
||||
|
||||
// markCommitted flags that the user has seen output, so the turn can no longer
|
||||
// be restarted, only collected behind.
|
||||
func (b *Bot) markCommitted(turn *chatTurn) {
|
||||
b.turnsMu.Lock()
|
||||
defer b.turnsMu.Unlock()
|
||||
turn.committed = true
|
||||
}
|
||||
|
||||
func (b *Bot) setTurnDraft(turn *chatTurn, draftID string) {
|
||||
b.turnsMu.Lock()
|
||||
defer b.turnsMu.Unlock()
|
||||
turn.draftID = draftID
|
||||
}
|
||||
|
||||
func (b *Bot) turnReason(turn *chatTurn) cancelReason {
|
||||
b.turnsMu.Lock()
|
||||
defer b.turnsMu.Unlock()
|
||||
return turn.reason
|
||||
}
|
||||
|
||||
// stopTurn handles Telegram's stopped_message_generation update. The draft id
|
||||
// must match the running turn's, so a stale tap on an old draft cannot kill a
|
||||
// newer turn. Reports whether a turn was cancelled.
|
||||
func (b *Bot) stopTurn(chatID int64, draftID int) bool {
|
||||
b.turnsMu.Lock()
|
||||
defer b.turnsMu.Unlock()
|
||||
|
||||
running, exists := b.turns[chatID]
|
||||
if !exists || running.draftID == "" || running.draftID != strconv.Itoa(draftID) {
|
||||
return false
|
||||
}
|
||||
running.reason = cancelStop
|
||||
running.cancel()
|
||||
return true
|
||||
}
|
||||
|
||||
// restartForEdit reruns a turn that is reading the pre-edit text of message
|
||||
// msgID and has shown nothing yet. Reports whether it did. A committed turn
|
||||
// keeps going: the user has already seen a reply to the old wording.
|
||||
func (b *Bot) restartForEdit(req turnRequest, msgID uint) bool {
|
||||
b.turnsMu.Lock()
|
||||
defer b.turnsMu.Unlock()
|
||||
|
||||
running, exists := b.turns[req.chatID]
|
||||
if !exists || running.committed || running.coveredID < msgID {
|
||||
return false
|
||||
}
|
||||
running.reason = cancelRestart
|
||||
running.rerun, running.rerunWant = &req, msgID
|
||||
running.cancel()
|
||||
return true
|
||||
}
|
||||
|
||||
// discardRerun drops a flush collected behind the chat's running turn. Used
|
||||
// by /clear, where rerunning against an emptied chat would call the model
|
||||
// with no user message at all.
|
||||
func (b *Bot) discardRerun(chatID int64) {
|
||||
b.turnsMu.Lock()
|
||||
defer b.turnsMu.Unlock()
|
||||
if running, exists := b.turns[chatID]; exists {
|
||||
running.rerun, running.rerunWant = nil, 0
|
||||
}
|
||||
}
|
||||
|
||||
// finishTurn unregisters the turn and returns the flush that still needs a turn
|
||||
// of its own, if any.
|
||||
//
|
||||
// A restart always reruns: the cancelled turn delivered nothing. A collected
|
||||
// flush reruns only if this turn's final context did not include its message.
|
||||
// Media turns snapshot after their uploads or transcription, so a text that
|
||||
// arrived during that work was collected and then answered anyway; rerunning
|
||||
// it would produce a second, unprompted reply. Stop does not discard the
|
||||
// flush: the message predates the tap and still deserves an answer.
|
||||
func (b *Bot) finishTurn(turn *chatTurn) *turnRequest {
|
||||
b.turnsMu.Lock()
|
||||
if b.turns[turn.chatID] == turn {
|
||||
delete(b.turns, turn.chatID)
|
||||
}
|
||||
rerun, want, reason, covered := turn.rerun, turn.rerunWant, turn.reason, turn.coveredID
|
||||
b.turnsMu.Unlock()
|
||||
|
||||
turn.cancel()
|
||||
close(turn.done)
|
||||
|
||||
if rerun == nil {
|
||||
return nil
|
||||
}
|
||||
if reason != cancelRestart && want <= covered {
|
||||
InfoLogger.Printf("[%s] turn %d for chat %d covered collected message %d; no rerun",
|
||||
b.config.ID, turn.seq, turn.chatID, want)
|
||||
return nil
|
||||
}
|
||||
return rerun
|
||||
}
|
||||
|
||||
// endTurn finishes the turn and runs whatever was collected behind it on its
|
||||
// own goroutine, so a chain of collected turns does not nest on the stack of
|
||||
// whichever update handler happened to run the first one.
|
||||
func (b *Bot) endTurn(turn *chatTurn) {
|
||||
if next := b.finishTurn(turn); next != nil {
|
||||
go b.respondToChat(*next)
|
||||
}
|
||||
}
|
||||
|
||||
// runTurn drives a registered turn to completion: keeps the typing indicator
|
||||
// alive, picks the reply sink (draft or segments), runs call, delivers and
|
||||
// records the reply, and afterwards runs anything collected while it was busy.
|
||||
//
|
||||
// call receives the turn's cancellable context and the sink, and returns the
|
||||
// joined reply text. Both the debounced text path and the photo path go
|
||||
// through here, so the two behave identically once the model is involved.
|
||||
//
|
||||
// Delivery of a finished reply uses the request's context, not the turn's: a
|
||||
// Stop or restart that lands after the stream ended must not turn a complete
|
||||
// reply into one that is stored but never sent.
|
||||
func (b *Bot) runTurn(turn *chatTurn, turnCtx context.Context, req turnRequest,
|
||||
call func(ctx context.Context, out replyOutput) (string, error)) {
|
||||
defer b.endTurn(turn)
|
||||
|
||||
stopTyping := b.startChatAction(turnCtx, req.chatID, req.businessConnectionID, models.ChatActionTyping)
|
||||
defer stopTyping()
|
||||
|
||||
out := b.newReplyOutput(turnCtx, turn, req.chatID, req.businessConnectionID)
|
||||
defer out.close()
|
||||
|
||||
joined, err := call(turnCtx, out)
|
||||
if err != nil {
|
||||
if errors.Is(err, errTurnSilent) {
|
||||
return
|
||||
}
|
||||
switch b.turnReason(turn) {
|
||||
case cancelRestart:
|
||||
InfoLogger.Printf("[%s] turn %d for chat %d abandoned for restart", b.config.ID, turn.seq, req.chatID)
|
||||
return
|
||||
case cancelStop:
|
||||
b.deliverStoppedTurn(turn, req, out)
|
||||
return
|
||||
}
|
||||
ErrorLogger.Printf("Error getting Anthropic response: %v", err)
|
||||
if sendErr := b.sendResponse(req.ctx, req.chatID, b.anthropicErrorResponse(err, req.userID), req.businessConnectionID); sendErr != nil {
|
||||
ErrorLogger.Printf("Error sending response: %v", sendErr)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if finishErr := out.finish(req.ctx, joined); finishErr != nil {
|
||||
ErrorLogger.Printf("[%s] delivering reply to chat %d: %v", b.config.ID, req.chatID, finishErr)
|
||||
return
|
||||
}
|
||||
if _, storeErr := b.screenOutgoingMessage(req.chatID, joined); storeErr != nil {
|
||||
ErrorLogger.Printf("Error recording assistant turn: %v", storeErr)
|
||||
}
|
||||
}
|
||||
|
||||
// deliverStoppedTurn keeps what the user watched being written. Telegram drops
|
||||
// the draft shortly after a stop unless the bot re-sends it as a real message,
|
||||
// and memory must hold the same text the user saw, or the next turn continues
|
||||
// from a reply that was never delivered.
|
||||
func (b *Bot) deliverStoppedTurn(turn *chatTurn, req turnRequest, out replyOutput) {
|
||||
partial := strings.TrimSpace(out.partial())
|
||||
InfoLogger.Printf("[%s] turn %d for chat %d stopped by user with %d chars shown",
|
||||
b.config.ID, turn.seq, req.chatID, len(partial))
|
||||
if partial == "" {
|
||||
return
|
||||
}
|
||||
if err := b.sendLongMessage(req.ctx, req.chatID, partial, req.businessConnectionID); err != nil {
|
||||
ErrorLogger.Printf("[%s] sending stopped partial to chat %d: %v", b.config.ID, req.chatID, err)
|
||||
return
|
||||
}
|
||||
if _, err := b.screenOutgoingMessage(req.chatID, partial); err != nil {
|
||||
ErrorLogger.Printf("Error recording stopped assistant turn: %v", err)
|
||||
}
|
||||
}
|
||||
-215
@@ -1,215 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func textReq(chatID int64) turnRequest {
|
||||
return turnRequest{ctx: context.Background(), chatID: chatID, userID: 555, username: "tester"}
|
||||
}
|
||||
|
||||
func TestBeginTextTurn_RunsWhenIdle(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
|
||||
turn, ctx, ok := b.beginTextTurn(textReq(1), 5)
|
||||
require.True(t, ok)
|
||||
require.NotNil(t, turn)
|
||||
assert.NoError(t, ctx.Err())
|
||||
|
||||
assert.Nil(t, b.finishTurn(turn))
|
||||
_, _, ok = b.beginTextTurn(textReq(1), 6)
|
||||
assert.True(t, ok, "chat is free again after finish")
|
||||
}
|
||||
|
||||
func TestBeginTextTurn_DropsFlushAlreadyCovered(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
turn, ctx, _ := b.beginTextTurn(textReq(2), 5)
|
||||
|
||||
_, _, ok := b.beginTextTurn(textReq(2), 5)
|
||||
|
||||
assert.False(t, ok)
|
||||
assert.NoError(t, ctx.Err(), "a covered flush must not cancel the running turn")
|
||||
assert.Nil(t, b.finishTurn(turn), "nothing to rerun")
|
||||
}
|
||||
|
||||
func TestBeginTextTurn_RestartsUncommittedTurn(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
turn, ctx, _ := b.beginTextTurn(textReq(3), 5)
|
||||
|
||||
_, _, ok := b.beginTextTurn(textReq(3), 7)
|
||||
|
||||
assert.False(t, ok)
|
||||
assert.ErrorIs(t, ctx.Err(), context.Canceled, "a newer message before any output cancels the turn")
|
||||
assert.Equal(t, cancelRestart, b.turnReason(turn))
|
||||
next := b.finishTurn(turn)
|
||||
require.NotNil(t, next, "the cancelled turn must be rerun")
|
||||
assert.Equal(t, int64(3), next.chatID)
|
||||
}
|
||||
|
||||
func TestFinishTurn_RestartRerunsEvenIfSnapshotRacedIn(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
turn, _, _ := b.beginTextTurn(textReq(30), 5)
|
||||
_, _, ok := b.beginTextTurn(textReq(30), 7)
|
||||
require.False(t, ok)
|
||||
|
||||
// The cancelled turn's snapshot happened to include message 7, but it
|
||||
// delivered nothing, so the rerun must still happen.
|
||||
b.markCovered(turn, 7)
|
||||
|
||||
assert.NotNil(t, b.finishTurn(turn), "a restart always reruns")
|
||||
}
|
||||
|
||||
func TestBeginTextTurn_CollectsBehindCommittedTurn(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
turn, ctx, _ := b.beginTextTurn(textReq(4), 5)
|
||||
b.markCommitted(turn)
|
||||
|
||||
_, _, ok := b.beginTextTurn(textReq(4), 7)
|
||||
|
||||
assert.False(t, ok)
|
||||
assert.NoError(t, ctx.Err(), "output already shown: never cancel")
|
||||
assert.Equal(t, cancelNone, b.turnReason(turn))
|
||||
require.NotNil(t, b.finishTurn(turn), "the collected flush runs afterwards")
|
||||
}
|
||||
|
||||
func TestFinishTurn_DropsCollectedFlushTheTurnEndedUpCovering(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
// A media turn: committed from the start, coverage known only later.
|
||||
turn, _ := b.beginTurn(context.Background(), 40)
|
||||
_, _, ok := b.beginTextTurn(textReq(40), 7)
|
||||
require.False(t, ok, "collected behind the media turn")
|
||||
|
||||
// The media turn's late snapshot includes message 7 and answers it.
|
||||
b.markCovered(turn, 7)
|
||||
|
||||
assert.Nil(t, b.finishTurn(turn), "rerunning would answer message 7 twice")
|
||||
}
|
||||
|
||||
func TestStopTurn_RequiresMatchingDraft(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
turn, ctx, _ := b.beginTextTurn(textReq(5), 5)
|
||||
|
||||
assert.False(t, b.stopTurn(5, 1), "no draft shown yet: nothing to stop")
|
||||
b.setTurnDraft(turn, "42")
|
||||
assert.False(t, b.stopTurn(5, 41), "a stale draft id is ignored")
|
||||
assert.NoError(t, ctx.Err())
|
||||
|
||||
assert.True(t, b.stopTurn(5, 42))
|
||||
assert.ErrorIs(t, ctx.Err(), context.Canceled)
|
||||
assert.Equal(t, cancelStop, b.turnReason(turn))
|
||||
b.finishTurn(turn)
|
||||
}
|
||||
|
||||
func TestFinishTurn_StopKeepsCollectedFlush(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
turn, _, _ := b.beginTextTurn(textReq(6), 5)
|
||||
b.markCommitted(turn)
|
||||
b.setTurnDraft(turn, turn.draftIDString())
|
||||
_, _, ok := b.beginTextTurn(textReq(6), 7)
|
||||
require.False(t, ok)
|
||||
|
||||
require.True(t, b.stopTurn(6, int(turn.seq)))
|
||||
|
||||
assert.NotNil(t, b.finishTurn(turn), "a message sent before the Stop tap still deserves an answer")
|
||||
}
|
||||
|
||||
func TestDiscardRerun_DropsCollectedFlush(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
turn, _, _ := b.beginTextTurn(textReq(8), 5)
|
||||
b.markCommitted(turn)
|
||||
_, _, ok := b.beginTextTurn(textReq(8), 7)
|
||||
require.False(t, ok)
|
||||
|
||||
b.discardRerun(8)
|
||||
|
||||
assert.Nil(t, b.finishTurn(turn), "/clear must not let the collected flush rerun on an empty chat")
|
||||
}
|
||||
|
||||
func TestRestartForEdit(t *testing.T) {
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
|
||||
t.Run("uncommitted turn reading the edited message restarts", func(t *testing.T) {
|
||||
turn, ctx, _ := b.beginTextTurn(textReq(9), 5)
|
||||
assert.True(t, b.restartForEdit(textReq(9), 5))
|
||||
assert.ErrorIs(t, ctx.Err(), context.Canceled)
|
||||
assert.Equal(t, cancelRestart, b.turnReason(turn))
|
||||
assert.NotNil(t, b.finishTurn(turn))
|
||||
})
|
||||
|
||||
t.Run("committed turn keeps going", func(t *testing.T) {
|
||||
turn, ctx, _ := b.beginTextTurn(textReq(10), 5)
|
||||
b.markCommitted(turn)
|
||||
assert.False(t, b.restartForEdit(textReq(10), 5))
|
||||
assert.NoError(t, ctx.Err())
|
||||
b.finishTurn(turn)
|
||||
})
|
||||
|
||||
t.Run("turn that never read the message is left alone", func(t *testing.T) {
|
||||
turn, ctx, _ := b.beginTextTurn(textReq(11), 3)
|
||||
assert.False(t, b.restartForEdit(textReq(11), 5))
|
||||
assert.NoError(t, ctx.Err())
|
||||
b.finishTurn(turn)
|
||||
})
|
||||
|
||||
t.Run("no running turn", func(t *testing.T) {
|
||||
assert.False(t, b.restartForEdit(textReq(12), 5))
|
||||
})
|
||||
}
|
||||
|
||||
func TestBeginTurn_WaitsForRunningTurn(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
running, _, _ := b.beginTextTurn(textReq(7), 5)
|
||||
|
||||
acquired := make(chan *chatTurn, 1)
|
||||
go func() {
|
||||
turn, _ := b.beginTurn(context.Background(), 7)
|
||||
acquired <- turn
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-acquired:
|
||||
t.Fatal("a media turn must wait for the running turn")
|
||||
case <-time.After(50 * time.Millisecond):
|
||||
}
|
||||
|
||||
b.finishTurn(running)
|
||||
|
||||
select {
|
||||
case turn := <-acquired:
|
||||
assert.True(t, turn.committed, "media turns are committed from the start")
|
||||
b.finishTurn(turn)
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("the media turn never acquired the chat")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBeginTurn_CancelledParentStillWaits(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, _ := setupBotForTest(t, 123)
|
||||
running, _, _ := b.beginTextTurn(textReq(13), 5)
|
||||
parent, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
acquired := make(chan struct{})
|
||||
go func() {
|
||||
turn, _ := b.beginTurn(parent, 13)
|
||||
b.finishTurn(turn)
|
||||
close(acquired)
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-acquired:
|
||||
t.Fatal("a cancelled parent must not register over the running turn")
|
||||
case <-time.After(50 * time.Millisecond):
|
||||
}
|
||||
b.finishTurn(running)
|
||||
select {
|
||||
case <-acquired:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("never acquired after the running turn finished")
|
||||
}
|
||||
}
|
||||
@@ -1,65 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"github.com/go-telegram/bot"
|
||||
"github.com/go-telegram/bot/models"
|
||||
)
|
||||
|
||||
// typingRefreshInterval re-arms the chat action before Telegram expires it.
|
||||
// The Bot API sets the status "for 5 seconds or less", so anything at or above
|
||||
// 5s leaves visible gaps. Telegram also clears the status as soon as the bot
|
||||
// sends a message, so streamed segments naturally interrupt it until the next
|
||||
// tick; there is no API call to clear it early.
|
||||
const typingRefreshInterval = 4 * time.Second
|
||||
|
||||
// startChatAction shows a chat action (typing, uploading a photo, recording a
|
||||
// voice note) and keeps it alive until the returned stop function runs.
|
||||
//
|
||||
// The returned function is idempotent and MUST be deferred by the caller. A
|
||||
// keepalive loop that can outlive its turn is the failure mode behind
|
||||
// openclaw/openclaw#27177, where the indicator stuck on until the process was
|
||||
// restarted, so the loop here owns a derived context and exits on the first of:
|
||||
// stop being called, or the parent context ending.
|
||||
func (b *Bot) startChatAction(
|
||||
ctx context.Context,
|
||||
chatID int64,
|
||||
businessConnectionID string,
|
||||
action models.ChatAction,
|
||||
) (stop func()) {
|
||||
actionCtx, cancel := context.WithCancel(ctx)
|
||||
|
||||
send := func() {
|
||||
params := &bot.SendChatActionParams{
|
||||
ChatID: chatID,
|
||||
Action: action,
|
||||
}
|
||||
if businessConnectionID != "" {
|
||||
params.BusinessConnectionID = businessConnectionID
|
||||
}
|
||||
if _, err := b.tgBot.SendChatAction(actionCtx, params); err != nil {
|
||||
// Cosmetic only: a failed indicator must never abort the turn.
|
||||
InfoLogger.Printf("[%s] chat action %q failed for chat %d: %v",
|
||||
b.config.ID, action, chatID, err)
|
||||
}
|
||||
}
|
||||
|
||||
send()
|
||||
|
||||
go func() {
|
||||
ticker := time.NewTicker(typingRefreshInterval)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-actionCtx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
send()
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
return cancel
|
||||
}
|
||||
-133
@@ -1,133 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/go-telegram/bot"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestStartChatAction_SendsImmediately(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, mockTg := setupBotForTest(t, 123)
|
||||
|
||||
var mu sync.Mutex
|
||||
var got []*bot.SendChatActionParams
|
||||
mockTg.SendChatActionFunc = func(_ context.Context, p *bot.SendChatActionParams) (bool, error) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
got = append(got, p)
|
||||
return true, nil
|
||||
}
|
||||
|
||||
stop := b.startChatAction(context.Background(), 42, "biz-7", "typing")
|
||||
stop()
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
require.Len(t, got, 1, "the indicator must show before the slow work starts, not after")
|
||||
assert.Equal(t, int64(42), got[0].ChatID)
|
||||
assert.EqualValues(t, "typing", got[0].Action)
|
||||
assert.Equal(t, "biz-7", got[0].BusinessConnectionID,
|
||||
"business chats need the connection id or the indicator never renders")
|
||||
}
|
||||
|
||||
func TestStartChatAction_OmitsEmptyBusinessConnectionID(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, mockTg := setupBotForTest(t, 123)
|
||||
|
||||
var mu sync.Mutex
|
||||
var captured *bot.SendChatActionParams
|
||||
mockTg.SendChatActionFunc = func(_ context.Context, p *bot.SendChatActionParams) (bool, error) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
captured = p
|
||||
return true, nil
|
||||
}
|
||||
|
||||
stop := b.startChatAction(context.Background(), 42, "", "typing")
|
||||
stop()
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
require.NotNil(t, captured)
|
||||
assert.Empty(t, captured.BusinessConnectionID)
|
||||
}
|
||||
|
||||
// The openclaw#27177 guard: once the turn ends, the keepalive must stop. A loop
|
||||
// that can re-arm after completion left the indicator stuck on until restart.
|
||||
func TestStartChatAction_StopHaltsKeepalive(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, mockTg := setupBotForTest(t, 123)
|
||||
|
||||
var calls atomic.Int32
|
||||
mockTg.SendChatActionFunc = func(_ context.Context, _ *bot.SendChatActionParams) (bool, error) {
|
||||
calls.Add(1)
|
||||
return true, nil
|
||||
}
|
||||
|
||||
stop := b.startChatAction(context.Background(), 42, "", "typing")
|
||||
stop()
|
||||
after := calls.Load()
|
||||
|
||||
// Well past a refresh tick had the loop survived.
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
assert.Equal(t, after, calls.Load(), "no chat action may be sent after stop")
|
||||
}
|
||||
|
||||
// Cancelling the parent context must also tear the loop down, so a turn aborted
|
||||
// upstream cannot leak a goroutine that keeps calling Telegram.
|
||||
func TestStartChatAction_ParentCancelHaltsKeepalive(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, mockTg := setupBotForTest(t, 123)
|
||||
|
||||
var calls atomic.Int32
|
||||
mockTg.SendChatActionFunc = func(_ context.Context, _ *bot.SendChatActionParams) (bool, error) {
|
||||
calls.Add(1)
|
||||
return true, nil
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
stop := b.startChatAction(ctx, 42, "", "typing")
|
||||
defer stop()
|
||||
|
||||
cancel()
|
||||
after := calls.Load()
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
assert.Equal(t, after, calls.Load(), "parent cancellation must stop the keepalive")
|
||||
}
|
||||
|
||||
func TestStartChatAction_StopIsIdempotent(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, mockTg := setupBotForTest(t, 123)
|
||||
mockTg.SendChatActionFunc = func(_ context.Context, _ *bot.SendChatActionParams) (bool, error) {
|
||||
return true, nil
|
||||
}
|
||||
|
||||
stop := b.startChatAction(context.Background(), 42, "", "typing")
|
||||
// The voice path calls stop early and again via defer.
|
||||
assert.NotPanics(t, func() {
|
||||
stop()
|
||||
stop()
|
||||
})
|
||||
}
|
||||
|
||||
// A failed indicator is cosmetic and must never surface as a turn failure.
|
||||
func TestStartChatAction_SendErrorIsNonFatal(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
b, mockTg := setupBotForTest(t, 123)
|
||||
mockTg.SendChatActionFunc = func(_ context.Context, _ *bot.SendChatActionParams) (bool, error) {
|
||||
return false, assert.AnError
|
||||
}
|
||||
|
||||
assert.NotPanics(t, func() {
|
||||
stop := b.startChatAction(context.Background(), 42, "", "typing")
|
||||
stop()
|
||||
})
|
||||
}
|
||||
|
||||
// Telegram clears the status after "5 seconds or less", so the refresh must be
|
||||
// strictly under that or the indicator visibly drops out mid-turn.
|
||||
func TestTypingRefreshInterval_UnderTelegramExpiry(t *testing.T) { //NOSONAR go:S100 -- underscore separation is idiomatic in Go test names
|
||||
assert.Less(t, typingRefreshInterval, 5*time.Second,
|
||||
"Telegram expires a chat action after at most 5s")
|
||||
}
|
||||
+63
-36
@@ -12,36 +12,29 @@ import (
|
||||
"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) {
|
||||
// Initialize loggers
|
||||
initLoggers()
|
||||
|
||||
db, err := gorm.Open(sqlite.Open(memoryDSN), &gorm.Config{})
|
||||
// Initialize in-memory database for testing
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
if err != nil {
|
||||
t.Fatalf(errOpenDB, err)
|
||||
t.Fatalf("Failed to open in-memory database: %v", err)
|
||||
}
|
||||
|
||||
err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{}, &Scope{})
|
||||
// Migrate the schema
|
||||
err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{})
|
||||
if err != nil {
|
||||
t.Fatalf(errMigrateSchema, err)
|
||||
t.Fatalf("Failed to migrate database schema: %v", err)
|
||||
}
|
||||
|
||||
// Create default roles
|
||||
err = createDefaultRoles(db)
|
||||
if err != nil {
|
||||
t.Fatalf(errCreateRoles, err)
|
||||
}
|
||||
if err := createDefaultScopes(db); err != nil {
|
||||
t.Fatalf(errCreateScopes, err)
|
||||
t.Fatalf("Failed to create default roles: %v", err)
|
||||
}
|
||||
|
||||
// Create a bot configuration
|
||||
config := BotConfig{
|
||||
ID: "test_bot",
|
||||
TelegramToken: "TEST_TELEGRAM_TOKEN",
|
||||
@@ -54,41 +47,49 @@ func TestOwnerAssignment(t *testing.T) {
|
||||
OwnerTelegramID: 111111111,
|
||||
}
|
||||
|
||||
// Initialize MockClock
|
||||
mockClock := &MockClock{
|
||||
currentTime: time.Now(),
|
||||
}
|
||||
|
||||
// Initialize MockTelegramClient
|
||||
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")
|
||||
}
|
||||
// Simulate successful message sending
|
||||
return &models.Message{ID: 1, Chat: models.Chat{ID: chatID}}, nil
|
||||
},
|
||||
}
|
||||
|
||||
// Create the bot with the mock Telegram client
|
||||
bot, err := NewBot(db, config, mockClock, mockTGClient)
|
||||
if err != nil {
|
||||
t.Fatalf(errCreateBot, err)
|
||||
t.Fatalf("Failed to create bot: %v", err)
|
||||
}
|
||||
|
||||
// Verify that the owner exists
|
||||
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)
|
||||
}
|
||||
|
||||
// Attempt to create another owner for the same bot
|
||||
_, err = bot.getOrCreateUser(222222222, "AnotherOwner", true)
|
||||
if err == nil {
|
||||
t.Fatalf("Expected error when creating a second owner, but got none")
|
||||
}
|
||||
|
||||
// Verify that the error message is appropriate
|
||||
expectedErrorMsg := "an owner already exists for this bot"
|
||||
if err.Error() != expectedErrorMsg {
|
||||
t.Fatalf("Unexpected error message: %v", err)
|
||||
}
|
||||
|
||||
// Assign admin role to a new user
|
||||
regularUser, err := bot.getOrCreateUser(333333333, "RegularUser", false)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create regular user: %v", err)
|
||||
@@ -98,6 +99,7 @@ func TestOwnerAssignment(t *testing.T) {
|
||||
t.Fatalf("Expected role 'user', got '%s'", regularUser.Role.Name)
|
||||
}
|
||||
|
||||
// Attempt to change an existing user to owner
|
||||
_, err = bot.getOrCreateUser(333333333, "AdminUser", true)
|
||||
if err == nil {
|
||||
t.Fatalf("Expected error when changing existing user to owner, but got none")
|
||||
@@ -108,27 +110,30 @@ func TestOwnerAssignment(t *testing.T) {
|
||||
t.Fatalf("Unexpected error message: %v", err)
|
||||
}
|
||||
|
||||
// If you need to test admin creation, you should do it through a separate admin creation function
|
||||
// or by updating an existing user's role with proper authorization checks
|
||||
}
|
||||
|
||||
func TestPromoteUserToAdmin(t *testing.T) {
|
||||
// Initialize loggers
|
||||
initLoggers()
|
||||
|
||||
db, err := gorm.Open(sqlite.Open(memoryDSN), &gorm.Config{})
|
||||
// Initialize in-memory database for testing
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
if err != nil {
|
||||
t.Fatalf(errOpenDB, err)
|
||||
t.Fatalf("Failed to open in-memory database: %v", err)
|
||||
}
|
||||
|
||||
err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{}, &Scope{})
|
||||
// Migrate the schema
|
||||
err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{})
|
||||
if err != nil {
|
||||
t.Fatalf(errMigrateSchema, err)
|
||||
t.Fatalf("Failed to migrate database schema: %v", err)
|
||||
}
|
||||
|
||||
// Create default roles
|
||||
err = createDefaultRoles(db)
|
||||
if err != nil {
|
||||
t.Fatalf(errCreateRoles, err)
|
||||
}
|
||||
if err := createDefaultScopes(db); err != nil {
|
||||
t.Fatalf(errCreateScopes, err)
|
||||
t.Fatalf("Failed to create default roles: %v", err)
|
||||
}
|
||||
|
||||
config := BotConfig{
|
||||
@@ -148,14 +153,16 @@ func TestPromoteUserToAdmin(t *testing.T) {
|
||||
|
||||
bot, err := NewBot(db, config, mockClock, mockTGClient)
|
||||
if err != nil {
|
||||
t.Fatalf(errCreateBot, err)
|
||||
t.Fatalf("Failed to create bot: %v", err)
|
||||
}
|
||||
|
||||
// Create an owner
|
||||
owner, err := bot.getOrCreateUser(config.OwnerTelegramID, "OwnerUser", true)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create owner: %v", err)
|
||||
}
|
||||
|
||||
// Test promoting a user to admin
|
||||
regularUser, err := bot.getOrCreateUser(444444444, "RegularUser", false)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create regular user: %v", err)
|
||||
@@ -166,6 +173,7 @@ func TestPromoteUserToAdmin(t *testing.T) {
|
||||
t.Fatalf("Failed to promote user to admin: %v", err)
|
||||
}
|
||||
|
||||
// Refresh user data
|
||||
promotedUser, err := bot.getOrCreateUser(444444444, "RegularUser", false)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to get promoted user: %v", err)
|
||||
@@ -176,31 +184,37 @@ func TestPromoteUserToAdmin(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestGetOrCreateUser tests the getOrCreateUser method of the Bot.
|
||||
// It verifies that a new user is created when one does not exist,
|
||||
// and an existing user is returned when one does exist.
|
||||
func TestGetOrCreateUser(t *testing.T) {
|
||||
// Initialize loggers
|
||||
initLoggers()
|
||||
|
||||
db, err := gorm.Open(sqlite.Open(memoryDSN), &gorm.Config{})
|
||||
// Initialize in-memory database for testing
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
if err != nil {
|
||||
t.Fatalf(errOpenDB, err)
|
||||
t.Fatalf("Failed to open in-memory database: %v", err)
|
||||
}
|
||||
|
||||
err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{}, &Scope{})
|
||||
// Migrate the schema
|
||||
err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{})
|
||||
if err != nil {
|
||||
t.Fatalf(errMigrateSchema, err)
|
||||
t.Fatalf("Failed to migrate database schema: %v", err)
|
||||
}
|
||||
|
||||
// Create default roles
|
||||
err = createDefaultRoles(db)
|
||||
if err != nil {
|
||||
t.Fatalf(errCreateRoles, err)
|
||||
}
|
||||
if err := createDefaultScopes(db); err != nil {
|
||||
t.Fatalf(errCreateScopes, err)
|
||||
t.Fatalf("Failed to create default roles: %v", err)
|
||||
}
|
||||
|
||||
// Create a mock clock starting at a fixed time
|
||||
mockClock := &MockClock{
|
||||
currentTime: time.Date(2023, 10, 1, 0, 0, 0, 0, time.UTC),
|
||||
}
|
||||
|
||||
// Create a mock configuration
|
||||
config := BotConfig{
|
||||
ID: "bot1",
|
||||
MemorySize: 10,
|
||||
@@ -212,49 +226,62 @@ func TestGetOrCreateUser(t *testing.T) {
|
||||
OwnerTelegramID: 123456789,
|
||||
}
|
||||
|
||||
// Initialize MockTelegramClient
|
||||
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")
|
||||
}
|
||||
// Simulate successful message sending
|
||||
return &models.Message{ID: 1, Chat: models.Chat{ID: chatID}}, nil
|
||||
},
|
||||
}
|
||||
|
||||
// Create the bot with the mock Telegram client
|
||||
bot, err := NewBot(db, config, mockClock, mockTGClient)
|
||||
if err != nil {
|
||||
t.Fatalf(errCreateBot, err)
|
||||
t.Fatalf("Failed to create bot: %v", err)
|
||||
}
|
||||
|
||||
// Verify that the owner exists
|
||||
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)
|
||||
}
|
||||
|
||||
// Attempt to create another owner for the same bot
|
||||
_, err = bot.getOrCreateUser(222222222, "AnotherOwner", true)
|
||||
if err == nil {
|
||||
t.Fatalf("Expected error when creating a second owner, but got none")
|
||||
}
|
||||
|
||||
// Create a new user
|
||||
newUser, err := bot.getOrCreateUser(987654321, "TestUser", false)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create a new user: %v", err)
|
||||
}
|
||||
|
||||
// Verify that the new user was created
|
||||
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)
|
||||
}
|
||||
|
||||
// Get the existing user
|
||||
existingUser, err := bot.getOrCreateUser(987654321, "TestUser", false)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to get existing user: %v", err)
|
||||
}
|
||||
|
||||
// Verify that the existing user is the same as the new user
|
||||
if existingUser.ID != userInDB.ID {
|
||||
t.Fatalf("Expected to get the existing user, but got a different user")
|
||||
}
|
||||
}
|
||||
|
||||
// To ensure thread safety and avoid race conditions during testing,
|
||||
// you can run the tests with the `-race` flag:
|
||||
// go test -race -v
|
||||
|
||||
Reference in New Issue
Block a user