This commit is contained in:
HugeFrog24
2026-09-03 21:55:36 +02:00
parent c36e1846f5
commit 89458480f3
17 changed files with 1593 additions and 127 deletions
+32 -1
View File
@@ -77,6 +77,7 @@ Each bot is one JSON file in `config/` (see `config/default.json` for the templa
| `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. |
| `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. |
@@ -103,10 +104,40 @@ What deliberately does **not** wait:
- **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]
> Debouncing is the cheap fix and the reason there is no "cancel the in-flight request" mode. Anthropic bills input tokens plus any output already generated when a turn stops partway, and any web searches it already ran are billed and re-billed on the retry. A message that never dispatched costs nothing.
> 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
+36 -15
View File
@@ -27,7 +27,7 @@ var mcpUnsupportedCount atomic.Uint64
type mcpCall struct{ server, name, input string }
func (b *Bot) getAnthropicResponse(ctx context.Context, chatID int64, messages []anthropic.BetaMessageParam, isEmojiOnly bool, username string, firstName string, lastName string, isPremium bool, languageCode string, messageTime int, onSegment func(string) error) (string, error) {
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))
@@ -107,11 +107,13 @@ func (b *Bot) getAnthropicResponse(ctx context.Context, chatID int64, messages [
params.Tools = tools
}
var fullText strings.Builder
var full string
var lastMsg anthropic.BetaMessage
fileRetries, pauseContinuations := 0, 0
for {
joined, msg, streamErr := b.streamMessages(ctx, params, onSegment)
// 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 {
@@ -136,12 +138,7 @@ func (b *Bot) getAnthropicResponse(ctx context.Context, chatID int64, messages [
}
lastMsg = msg
if joined != "" {
if fullText.Len() > 0 {
fullText.WriteString("\n\n")
}
fullText.WriteString(joined)
}
full = joinNonEmpty(full, joined)
if msg.StopReason == anthropic.BetaStopReasonPauseTurn {
pauseContinuations++
@@ -156,11 +153,11 @@ func (b *Bot) getAnthropicResponse(ctx context.Context, chatID int64, messages [
break
}
if fullText.Len() == 0 {
if full == "" {
return "", emptyStreamError(string(lastMsg.StopReason),
lastMsg.Usage.OutputTokensDetails.ThinkingTokens, params.MaxTokens)
}
return fullText.String(), nil
return full, nil
}
func webSearchTools(cfg *WebSearchConfig) []anthropic.BetaToolUnionParam {
@@ -276,7 +273,12 @@ func thinkingParamFromConfig(mode, display string) (anthropic.BetaThinkingConfig
}
}
func (b *Bot) streamMessages(ctx context.Context, params anthropic.BetaMessageNewParams, onSegment func(string) error) (string, anthropic.BetaMessage, error) {
// 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 {
@@ -286,7 +288,8 @@ func (b *Bot) streamMessages(ctx context.Context, params anthropic.BetaMessageNe
var (
message anthropic.BetaMessage
allSegments []string
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
@@ -336,6 +339,9 @@ func (b *Bot) streamMessages(ctx context.Context, params anthropic.BetaMessageNe
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" {
@@ -352,7 +358,8 @@ func (b *Bot) streamMessages(ctx context.Context, params anthropic.BetaMessageNe
case "text":
seg := strings.TrimSpace(currentText.String())
if seg != "" {
allSegments = append(allSegments, 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)
@@ -428,7 +435,21 @@ func (b *Bot) streamMessages(ctx context.Context, params anthropic.BetaMessageNe
}
}
return strings.Join(allSegments, "\n\n"), message, nil
return done, message, nil
}
// 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
}
}
func emptyStreamError(stopReason string, thinkingTokens, maxTokens int64) error {
+43
View File
@@ -32,6 +32,9 @@ type Bot struct {
intakeBuffers map[int64]*pendingIntake
intakeBuffersMu sync.Mutex
intakeSeq uint64
turns map[int64]*chatTurn
turnsMu sync.Mutex
turnSeq uint64
}
func messageType(msg *models.Message) string {
@@ -94,6 +97,7 @@ func NewBot(db *gorm.DB, config BotConfig, clock Clock, tgClient TelegramClient)
tgBot: tgClient,
albumBuffers: make(map[string]*pendingAlbum),
intakeBuffers: make(map[int64]*pendingIntake),
turns: make(map[int64]*chatTurn),
}
if tgClient == nil {
@@ -276,10 +280,48 @@ func (b *Bot) addMessageToChatMemory(chatMemory *ChatMemory, message Message) {
}
}
// prepareContextMessages converts chat memory into API messages.
func (b *Bot) prepareContextMessages(chatMemory *ChatMemory) []anthropic.BetaMessageParam {
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 {
InfoLogger.Printf("Chat memory contains %d messages", len(chatMemory.Messages))
for i, msg := range chatMemory.Messages {
InfoLogger.Printf("Message %d: IsUser=%v, Text=%q Images=%d", i, msg.IsUser, msg.Text, len(msg.ImageFileIDs))
@@ -701,6 +743,7 @@ func (b *Bot) screenIncomingMessage(message *models.Message) (Message, error) {
}
userMessage := b.createMessage(message.Chat.ID, message.From.ID, message.From.Username, userRole, messageText, true)
userMessage.TelegramMessageID = message.ID
if message.Sticker != nil {
userMessage.StickerFileID = message.Sticker.FileID
+6
View File
@@ -67,6 +67,7 @@ type BotConfig struct {
Thinking string `json:"thinking,omitempty"`
ThinkingDisplay string `json:"thinking_display,omitempty"`
DebounceMs int `json:"debounce_ms,omitempty"`
StreamDrafts bool `json:"stream_drafts,omitempty"`
CacheHistory *bool `json:"cache_history,omitempty"`
SystemPrompts map[string]string `json:"system_prompts"`
Active bool `json:"active"`
@@ -172,6 +173,11 @@ func logConfigAdvisories(config *BotConfig) {
config.ID)
}
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)
+288
View File
@@ -0,0 +1,288 @@
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
View File
@@ -0,0 +1,217 @@
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"))
}
+181
View File
@@ -0,0 +1,181 @@
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")
}
Binary file not shown.
+7 -6
View File
@@ -3,9 +3,9 @@ module github.com/HugeFrog24/go-telegram-bot
go 1.26.0
require (
github.com/anthropics/anthropic-sdk-go v1.57.0
github.com/go-telegram/bot v1.22.0
github.com/stretchr/testify v1.11.1
github.com/anthropics/anthropic-sdk-go v1.69.0
github.com/go-telegram/bot v1.25.0
github.com/stretchr/testify v1.12.1
golang.org/x/sync v0.22.0
golang.org/x/time v0.15.0
gorm.io/driver/sqlite v1.6.0
@@ -14,14 +14,14 @@ require (
require (
github.com/bahlo/generic-list-go v0.2.0 // indirect
github.com/buger/jsonparser v1.2.0 // indirect
github.com/buger/jsonparser v1.6.1 // indirect
github.com/davecgh/go-spew v1.1.1 // indirect
github.com/invopop/jsonschema v0.14.0 // indirect
github.com/jinzhu/inflection v1.0.0 // indirect
github.com/jinzhu/now v1.1.5 // indirect
github.com/kr/pretty v0.3.1 // indirect
github.com/mailru/easyjson v0.9.2 // indirect
github.com/mattn/go-sqlite3 v1.14.48 // indirect
github.com/mattn/go-sqlite3 v1.14.50 // indirect
github.com/pb33f/ordered-map/v2 v2.3.1 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect
github.com/rogpeppe/go-internal v1.14.1 // indirect
@@ -32,8 +32,9 @@ require (
github.com/tidwall/pretty v1.2.1 // indirect
github.com/tidwall/sjson v1.2.5 // indirect
github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect
go.yaml.in/yaml/v3 v3.0.5 // indirect
go.yaml.in/yaml/v4 v4.0.0-rc.6 // indirect
golang.org/x/text v0.40.0 // indirect
golang.org/x/text v0.41.0 // indirect
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
gopkg.in/yaml.v3 v3.0.1 // indirect
)
+14
View File
@@ -4,10 +4,14 @@ github.com/anthropics/anthropic-sdk-go v1.52.0 h1:1TB9jt4DN87VMwS/hB1VK26tYzK0ip
github.com/anthropics/anthropic-sdk-go v1.52.0/go.mod h1:3EfIfmFqxH6rbiLcIP4tPFyXL/IHakx2wDG4OU+TIEI=
github.com/anthropics/anthropic-sdk-go v1.57.0 h1:iEAcPbUKfJ2Iqz9uN/jEndCNW2+x7OYLHDidXDhPjI0=
github.com/anthropics/anthropic-sdk-go v1.57.0/go.mod h1:3EfIfmFqxH6rbiLcIP4tPFyXL/IHakx2wDG4OU+TIEI=
github.com/anthropics/anthropic-sdk-go v1.69.0 h1:TC08TDog21aHW/pAj4wO0Xi7nMXt0n5Rybb8etkhL0M=
github.com/anthropics/anthropic-sdk-go v1.69.0/go.mod h1:3EfIfmFqxH6rbiLcIP4tPFyXL/IHakx2wDG4OU+TIEI=
github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk=
github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg=
github.com/buger/jsonparser v1.2.0 h1:4EFcvK1kD4jyj6YqNK6skK6w+y7FHHBR+XBCtxwu/6g=
github.com/buger/jsonparser v1.2.0/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0=
github.com/buger/jsonparser v1.6.1 h1:I0phFv0PlbLHnM7TZAVjZ2MJ2/eWRTDyuO7GLR98IEs=
github.com/buger/jsonparser v1.6.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0=
github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
@@ -19,6 +23,8 @@ github.com/go-telegram/bot v1.21.0 h1:Va/PbGc2vBDdv57GCUEEVV6ROlHWiC6SklJY9Hvhzp
github.com/go-telegram/bot v1.21.0/go.mod h1:i2TRs7fXWIeaceF3z7KzsMt/he0TwkVC680mvdTFYeM=
github.com/go-telegram/bot v1.22.0 h1:zK29OoTYMmR5emJrCtGa2SjaGleeZiUB/C1i7kc2lXE=
github.com/go-telegram/bot v1.22.0/go.mod h1:i2TRs7fXWIeaceF3z7KzsMt/he0TwkVC680mvdTFYeM=
github.com/go-telegram/bot v1.25.0 h1:C072msZtr+GjJiQsX/UJLirz+jzSNZVWkckxWufR7Ww=
github.com/go-telegram/bot v1.25.0/go.mod h1:i2TRs7fXWIeaceF3z7KzsMt/he0TwkVC680mvdTFYeM=
github.com/invopop/jsonschema v0.13.0 h1:KvpoAJWEjR3uD9Kbm2HWJmqsEaHt8lBUpd0qHcIi21E=
github.com/invopop/jsonschema v0.13.0/go.mod h1:ffZ5Km5SWWRAIN6wbDXItl95euhFz2uON45H2qjYt+0=
github.com/invopop/jsonschema v0.14.0 h1:MHQqLhvpNUZfw+hM3AZDYK7jxO8FZoQeQM77g8iyZjg=
@@ -42,6 +48,8 @@ github.com/mattn/go-sqlite3 v1.14.47 h1:jOBI62gS7nKeZv+as1oGEy0+1qISgXwH/QBlR6Kb
github.com/mattn/go-sqlite3 v1.14.47/go.mod h1:6JTjA44L93a0QCyJef5YvlPoKXntQPjzWv5gtm9sB6w=
github.com/mattn/go-sqlite3 v1.14.48 h1:7XHIgl0a8HwOaiK4E47ozLkST78rR9+OtNGx27D/TFs=
github.com/mattn/go-sqlite3 v1.14.48/go.mod h1:6JTjA44L93a0QCyJef5YvlPoKXntQPjzWv5gtm9sB6w=
github.com/mattn/go-sqlite3 v1.14.50 h1:dmdFvo1XG4MPzA4IkAmE9upVz/Nj31uRoM5+jC8hYbY=
github.com/mattn/go-sqlite3 v1.14.50/go.mod h1:6JTjA44L93a0QCyJef5YvlPoKXntQPjzWv5gtm9sB6w=
github.com/pb33f/ordered-map/v2 v2.3.1 h1:5319HDO0aw4DA4gzi+zv4FXU9UlSs3xGZ40wcP1nBjY=
github.com/pb33f/ordered-map/v2 v2.3.1/go.mod h1:qxFQgd0PkVUtOMCkTapqotNgzRhMPL7VvaHKbd1HnmQ=
github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
@@ -56,6 +64,8 @@ github.com/stretchr/objx v0.5.3 h1:jmXUvGomnU1o3W/V5h2VEradbpJDwGrzugQQvL0POH4=
github.com/stretchr/objx v0.5.3/go.mod h1:rDQraq+vQZU7Fde9LOZLr8Tax6zZvy4kuNKF+QYS+U0=
github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
github.com/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.19.0 h1:xwxm7n691Uf3u5OFjzngavjGTh55KX5q/9w9xHW88JU=
github.com/tidwall/gjson v1.19.0/go.mod h1:V37/opeE/JbLUOfH0QTXiNez2l0RUjYUhpT4szFQAfc=
@@ -69,6 +79,8 @@ github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY=
github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28=
github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc=
github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw=
go.yaml.in/yaml/v3 v3.0.5 h1:N6y/pJk8buWs9NY5ERU2HSMfm+IuD/OtfdAnq6kESPw=
go.yaml.in/yaml/v3 v3.0.5/go.mod h1:HVTZu1O7/Vkt2N+BFy8Zza+lnLsABggaTM2ZpNIGuKg=
go.yaml.in/yaml/v4 v4.0.0-rc.6 h1:1h7H1ohdUh93/FyE4YaDa1Zh64K6VVbjF4K6WUxMtH4=
go.yaml.in/yaml/v4 v4.0.0-rc.6/go.mod h1:aZqd9kCMsGL7AuUv/m/PvWLdg5sjJsZ4oHDEnfPPfY0=
golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4=
@@ -83,6 +95,8 @@ golang.org/x/text v0.38.0 h1:sXmwo9DwP3OK9EZ7PqAdaooSGozfl/3a6/xJcbzPRhE=
golang.org/x/text v0.38.0/go.mod h1:YXZt3QhHUKYT53r2lLKFIVi6Ao1jdzrTR/KQ09qyxF4=
golang.org/x/text v0.40.0 h1:Ub2Z6/xjgF1WrYQz2nuITOEegKFtiIy+rieRJ5lHZKs=
golang.org/x/text v0.40.0/go.mod h1:hpnzDAfGV753zIKo+wk3u1bVKCGPbrnF7+7LBF/UHVY=
golang.org/x/text v0.41.0 h1:vz/seA0lnX87Othu2f/0L24RcgrXD9/YFTSuGjj3rH8=
golang.org/x/text v0.41.0/go.mod h1:jvf1O8ajNzZqhSrQBPbutR/EB83Cc0CFrezNQIwbb5M=
golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U=
golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+180 -79
View File
@@ -28,6 +28,12 @@ func (b *Bot) handleVoiceMessage(ctx context.Context, message *models.Message, u
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()
@@ -54,9 +60,8 @@ func (b *Bot) handleVoiceMessage(ctx context.Context, message *models.Message, u
}
b.chatMemoriesMu.Unlock()
chatMemory := b.getOrCreateChatMemory(chatID)
contextMessages := b.prepareContextMessages(chatMemory)
response, err := b.getAnthropicResponse(ctx, chatID, contextMessages, false, username, firstName, lastName, isPremium, languageCode, messageTime, nil)
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 {
@@ -120,11 +125,53 @@ func (b *Bot) handlePhotoMessage(
return
}
// Covers the Files API uploads as well as the model turn; on an album this
// is the longest wait in the bot.
// 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)
defer stopTyping()
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)
@@ -157,7 +204,7 @@ func (b *Bot) handlePhotoMessage(
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
return nil, "", errTurnSilent
}
finalUploaded := make([]string, 0, len(uploaded))
for _, fid := range uploaded {
@@ -165,41 +212,7 @@ func (b *Bot) handlePhotoMessage(
finalUploaded = append(finalUploaded, fid)
}
}
if len(finalUploaded) == 0 {
return
}
chatMemory := b.getOrCreateChatMemory(chatID)
userMessage := b.createMessage(chatID, userID, username, "user", caption, true)
userMessage.ImageFileIDs = finalUploaded
if err := b.storeMessage(&userMessage); err != nil {
b.compensatingDelete(ctx, finalUploaded)
ErrorLogger.Printf("[%s] store photo message failed: %v", b.config.ID, err)
if sendErr := b.sendResponse(ctx, chatID, "Sorry, I had trouble saving your message.", businessConnectionID); sendErr != nil {
ErrorLogger.Printf("Error sending store failure message: %v", sendErr)
}
return
}
b.addMessageToChatMemory(chatMemory, userMessage)
contextMessages := b.prepareContextMessages(chatMemory)
joined, err := b.getAnthropicResponse(
ctx, chatID, contextMessages, false,
username, firstName, lastName, isPremium, languageCode, messageTime,
func(seg string) error {
return b.sendOneSegment(ctx, chatID, seg, businessConnectionID)
},
)
if err != nil {
ErrorLogger.Printf("Error getting Anthropic response for photo: %v", err)
if sendErr := b.sendResponse(ctx, chatID, b.anthropicErrorResponse(err, userID), businessConnectionID); sendErr != nil {
ErrorLogger.Printf("Error sending anthropic error response: %v", sendErr)
}
return
}
if _, storeErr := b.screenOutgoingMessage(chatID, joined); storeErr != nil {
ErrorLogger.Printf("Error recording assistant turn: %v", storeErr)
}
return finalUploaded, caption, nil
}
// respondToChat runs one assistant turn against the chat's current memory and
@@ -207,40 +220,30 @@ func (b *Bot) handlePhotoMessage(
// 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.
func (b *Bot) respondToChat(
ctx context.Context,
chatID, userID int64,
isEmojiOnly bool,
username, firstName, lastName string,
isPremium bool,
languageCode string,
messageTime int,
businessConnectionID string,
) {
stopTyping := b.startChatAction(ctx, chatID, businessConnectionID, models.ChatActionTyping)
defer stopTyping()
chatMemory := b.getOrCreateChatMemory(chatID)
contextMessages := b.prepareContextMessages(chatMemory)
joined, err := b.getAnthropicResponse(
ctx, chatID, contextMessages, isEmojiOnly,
username, firstName, lastName, isPremium, languageCode, messageTime,
func(seg string) error {
return b.sendOneSegment(ctx, chatID, seg, businessConnectionID)
},
)
if err != nil {
ErrorLogger.Printf("Error getting Anthropic response: %v", err)
if sendErr := b.sendResponse(ctx, chatID, b.anthropicErrorResponse(err, userID), businessConnectionID); sendErr != nil {
ErrorLogger.Printf("Error sending response: %v", sendErr)
}
//
// 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
}
if _, storeErr := b.screenOutgoingMessage(chatID, joined); storeErr != nil {
ErrorLogger.Printf("Error recording assistant turn: %v", storeErr)
}
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 {
@@ -274,6 +277,19 @@ func (b *Bot) anthropicErrorResponse(err error, userID int64) string {
}
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 {
@@ -489,8 +505,10 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U
if message.Sticker != nil {
b.cancelIntake(chatID)
contextMessages := b.prepareContextMessages(b.getOrCreateChatMemory(chatID))
b.handleStickerMessage(ctx, chatID, userMsg, message, contextMessages, businessConnectionID)
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)
return
}
@@ -510,9 +528,89 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U
return
}
b.respondToChat(ctx, chatID, userID, isEmojiOnly,
username, firstName, lastName, isPremium, languageCode, messageTime,
businessConnectionID)
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) {
@@ -547,7 +645,7 @@ func (b *Bot) generateStickerResponse(ctx context.Context, message Message, cont
if message.StickerFileID != "" {
messageTime := int(message.Timestamp.Unix())
response, err := b.getAnthropicResponse(ctx, message.ChatID, contextMessages, true, message.Username, "", "", false, "", messageTime, nil)
response, err := b.getAnthropicResponse(ctx, message.ChatID, contextMessages, true, message.Username, "", "", false, "", messageTime, nil, nil)
if err != nil {
return "", err
}
@@ -641,6 +739,9 @@ func (b *Bot) clearChatHistory(ctx context.Context, chatID int64, currentUserID
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)
b.chatMemoriesMu.Lock()
if targetUserID == currentUserID {
+40 -26
View File
@@ -8,18 +8,14 @@ import (
// 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 metadata
// 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 {
chatID, userID int64
username, firstName, lastName, languageCode string
isPremium bool
messageTime int
businessConnectionID string
allEmojiOnly bool
count int
seq uint64
timer *time.Timer
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
@@ -48,27 +44,47 @@ func (b *Bot) bufferIntake(
}
// Reply metadata tracks the most recent message in the batch.
pending.chatID = chatID
pending.userID = userID
pending.username = username
pending.firstName = firstName
pending.lastName = lastName
pending.isPremium = isPremium
pending.languageCode = languageCode
pending.messageTime = messageTime
pending.businessConnectionID = businessConnectionID
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()
}
seq := pending.seq
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.
@@ -88,12 +104,10 @@ func (b *Bot) flushIntake(ctx context.Context, chatID int64, seq uint64) {
b.config.ID, captured.count, chatID)
}
b.respondToChat(
ctx, chatID, captured.userID, captured.allEmojiOnly,
captured.username, captured.firstName, captured.lastName,
captured.isPremium, captured.languageCode, captured.messageTime,
captured.businessConnectionID,
)
req := captured.turnRequest
req.ctx = ctx
req.isEmojiOnly = captured.allEmojiOnly
b.respondToChat(req)
}
// cancelIntake drops a chat's pending buffer without dispatching, returning how
+4
View File
@@ -44,6 +44,10 @@ type Message struct {
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"`
}
type ChatMemory struct {
+1
View File
@@ -9,6 +9,7 @@ import (
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)
+8
View File
@@ -11,6 +11,7 @@ import (
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)
@@ -30,6 +31,13 @@ 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)
+321
View File
@@ -0,0 +1,321 @@
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
View File
@@ -0,0 +1,215 @@
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")
}
}