From 37810d918a871c2b4ea20f75aa55e91ab1319322 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Fri, 8 May 2026 14:52:35 +0200 Subject: [PATCH 1/6] fix(collections): rehydrate in-process collections lazily after init failure (#468) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When NewInProcessBackend boots, it iterates every on-disk collection and calls newVectorEngine to construct its engine wrapper. For the postgres engine that constructor performs a "test embedding" probe that requires the embedding model to be reachable. If the embedding service is briefly unavailable at boot — e.g. the node hosting the embedding model has temporary NATS connectivity issues — the construction fails and the collection is silently dropped from the in-memory map. Subsequent operations against that collection (Upload / Search / ListEntries / …) all return "collection not found" indefinitely, even after the embedding service comes back, because nothing ever retries. Worse, the collection still exists on disk (its JSON sidecar) and its data still exists in the vector DB (e.g. the per-collection documents_col_ table in PostgreSQL). From the user's perspective their collection has silently disappeared. Two changes: 1. NewInProcessBackend always registers every on-disk collection in state.Collections — even when newVectorEngine returns nil. A nil entry is a placeholder meaning "known on disk, not yet loaded". 2. backendInProcess.lookup centralises the cache read for every operation. If the cache holds a placeholder, it retries newVectorEngine now, under the write lock. So as soon as the embedding service is reachable again, the next request to the collection will rehydrate it transparently. If init still fails or the collection isn't on disk at all, lookup returns (nil, false) and the operation surfaces "collection not found" as before. The state.EnsureCollection callback used by the internal RAG provider already handled the placeholder case (it re-inits whenever the cache entry is missing or nil), so it needs no change. This does not address the underlying probe-and-cache pattern in LocalRecall's NewPersistentPostgresCollection — read-only operations on existing collections still require an embedding probe at engine construction. That is a separate, deeper fix worth pursuing in LocalRecall directly. --- webui/collections/inprocess.go | 87 +++++++++++++++++++++------------- 1 file changed, 53 insertions(+), 34 deletions(-) diff --git a/webui/collections/inprocess.go b/webui/collections/inprocess.go index f508a73..c9b8e42 100644 --- a/webui/collections/inprocess.go +++ b/webui/collections/inprocess.go @@ -63,6 +63,39 @@ type backendInProcess struct { var _ Backend = (*backendInProcess)(nil) +// lookup returns the cached collection KB for name. If the cache holds a +// placeholder (nil entry — the engine init failed at startup, e.g. because +// the embedding service was momentarily unreachable when iterating over +// existing collections in NewInProcessBackend) it attempts to re-initialise +// the engine now so a transient outage doesn't permanently 404 a collection +// that still has data on disk / in the vector DB. Returns (nil, false) only +// when the collection isn't known at all, or when re-init still fails. +func (b *backendInProcess) lookup(name string) (*rag.PersistentKB, bool) { + b.state.Mu.RLock() + kb, exists := b.state.Collections[name] + b.state.Mu.RUnlock() + if !exists { + return nil, false + } + if kb != nil { + return kb, true + } + // Placeholder: collection is known on disk but its engine wrapper failed + // to construct earlier. Retry under the write lock. + b.state.Mu.Lock() + defer b.state.Mu.Unlock() + if kb, ok := b.state.Collections[name]; ok && kb != nil { + return kb, true + } + kb = newVectorEngine(b.cfg.VectorEngine, b.openAIClient, b.cfg.LLMAPIURL, b.cfg.LLMAPIKey, name, b.cfg.CollectionDBPath, b.cfg.FileAssets, b.cfg.EmbeddingModel, b.cfg.DatabaseURL, b.cfg.MaxChunkingSize, b.cfg.ChunkOverlap) + if kb == nil { + return nil, false + } + b.state.Collections[name] = kb + b.state.SourceManager.RegisterCollection(name, kb) + return kb, true +} + func (b *backendInProcess) ListCollections() ([]string, error) { return rag.ListAllCollections(b.cfg.CollectionDBPath), nil } @@ -80,9 +113,7 @@ func (b *backendInProcess) CreateCollection(name string) error { } func (b *backendInProcess) Upload(collection, filename string, fileBody io.Reader) (string, error) { - b.state.Mu.RLock() - kb, exists := b.state.Collections[collection] - b.state.Mu.RUnlock() + kb, exists := b.lookup(collection) if !exists { return "", fmt.Errorf("collection not found: %s", collection) } @@ -108,9 +139,7 @@ func (b *backendInProcess) Upload(collection, filename string, fileBody io.Reade } func (b *backendInProcess) ListEntries(collection string) ([]string, error) { - b.state.Mu.RLock() - kb, exists := b.state.Collections[collection] - b.state.Mu.RUnlock() + kb, exists := b.lookup(collection) if !exists { return nil, fmt.Errorf("collection not found: %s", collection) } @@ -118,9 +147,7 @@ func (b *backendInProcess) ListEntries(collection string) ([]string, error) { } func (b *backendInProcess) GetEntryContent(collection, entry string) (string, int, error) { - b.state.Mu.RLock() - kb, exists := b.state.Collections[collection] - b.state.Mu.RUnlock() + kb, exists := b.lookup(collection) if !exists { return "", 0, fmt.Errorf("collection not found: %s", collection) } @@ -128,9 +155,7 @@ func (b *backendInProcess) GetEntryContent(collection, entry string) (string, in } func (b *backendInProcess) Search(collection, query string, maxResults int) ([]SearchResult, error) { - b.state.Mu.RLock() - kb, exists := b.state.Collections[collection] - b.state.Mu.RUnlock() + kb, exists := b.lookup(collection) if !exists { return nil, fmt.Errorf("collection not found: %s", collection) } @@ -159,22 +184,18 @@ func (b *backendInProcess) Search(collection, query string, maxResults int) ([]S } func (b *backendInProcess) Reset(collection string) error { - b.state.Mu.Lock() - kb, exists := b.state.Collections[collection] - if exists { - delete(b.state.Collections, collection) - } - b.state.Mu.Unlock() + kb, exists := b.lookup(collection) if !exists { return fmt.Errorf("collection not found: %s", collection) } + b.state.Mu.Lock() + delete(b.state.Collections, collection) + b.state.Mu.Unlock() return kb.Reset() } func (b *backendInProcess) DeleteEntry(collection, entry string) ([]string, error) { - b.state.Mu.RLock() - kb, exists := b.state.Collections[collection] - b.state.Mu.RUnlock() + kb, exists := b.lookup(collection) if !exists { return nil, fmt.Errorf("collection not found: %s", collection) } @@ -186,9 +207,7 @@ func (b *backendInProcess) DeleteEntry(collection, entry string) ([]string, erro } func (b *backendInProcess) AddSource(collection, url string, intervalMin int) error { - b.state.Mu.RLock() - kb, exists := b.state.Collections[collection] - b.state.Mu.RUnlock() + kb, exists := b.lookup(collection) if !exists { return fmt.Errorf("collection not found: %s", collection) } @@ -201,9 +220,7 @@ func (b *backendInProcess) RemoveSource(collection, url string) error { } func (b *backendInProcess) ListSources(collection string) ([]SourceInfo, error) { - b.state.Mu.RLock() - kb, exists := b.state.Collections[collection] - b.state.Mu.RUnlock() + kb, exists := b.lookup(collection) if !exists { return nil, fmt.Errorf("collection not found: %s", collection) } @@ -220,9 +237,7 @@ func (b *backendInProcess) ListSources(collection string) ([]SourceInfo, error) } func (b *backendInProcess) GetEntryFilePath(collection, entry string) (string, error) { - b.state.Mu.RLock() - kb, exists := b.state.Collections[collection] - b.state.Mu.RUnlock() + kb, exists := b.lookup(collection) if !exists { return "", fmt.Errorf("collection not found: %s", collection) } @@ -230,9 +245,7 @@ func (b *backendInProcess) GetEntryFilePath(collection, entry string) (string, e } func (b *backendInProcess) EntryExists(collection, entry string) bool { - b.state.Mu.RLock() - kb, exists := b.state.Collections[collection] - b.state.Mu.RUnlock() + kb, exists := b.lookup(collection) if !exists { return false } @@ -257,8 +270,14 @@ func NewInProcessBackend(cfg *Config) (Backend, *State) { colls := rag.ListAllCollections(cfg.CollectionDBPath) for _, c := range colls { collection := newVectorEngine(cfg.VectorEngine, openAIClient, cfg.LLMAPIURL, cfg.LLMAPIKey, c, cfg.CollectionDBPath, cfg.FileAssets, cfg.EmbeddingModel, cfg.DatabaseURL, cfg.MaxChunkingSize, cfg.ChunkOverlap) + // Register every on-disk collection — even when the engine wrapper + // failed to construct (e.g. the embedding service was momentarily + // unreachable). A nil entry marks "known on disk but not yet loaded"; + // backendInProcess.lookup will rehydrate lazily on first access so a + // transient outage at boot doesn't permanently 404 collections whose + // data is still on disk / in the vector DB. + st.Collections[c] = collection if collection != nil { - st.Collections[c] = collection st.SourceManager.RegisterCollection(c, collection) } } From 22280d4c88083a9e282faf2bd9a2411f9832992f Mon Sep 17 00:00:00 2001 From: Bas Hulsken Date: Tue, 2 Jun 2026 16:12:18 +0200 Subject: [PATCH 2/6] allow configuring LOCALAGI to use a custom base url (#469) allow configuring LOCALAGI to use a custom base url other than the default :3000 --- README.md | 6 ++++-- cmd/env.go | 2 ++ cmd/serve.go | 2 +- webui/react-ui/vite.config.js | 2 +- 4 files changed, 8 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 4ee5fe4..cb3b881 100644 --- a/README.md +++ b/README.md @@ -239,7 +239,8 @@ LocalAGI supports environment configurations. Note that these environment variab | `LOCALAGI_LLM_API_KEY` | API authentication | | `LOCALAGI_TIMEOUT` | Request timeout settings | | `LOCALAGI_STATE_DIR` | Where state gets stored | -| `LOCALAGI_BASE_URL` | Optional base URL for the app (only relevant when using an external LocalRAG URL; not used for built-in knowledge base) | +| `LOCALAGI_LOCALRAG_URL` | Optional URL when using an external LocalRAG URL; not used for built-in knowledge base | +| `LOCALAGI_BASE_URL` | Optional base URL for the app (defaults to ":3000") | | `LOCALAGI_ENABLE_CONVERSATIONS_LOGGING` | Toggle conversation logs | | `LOCALAGI_API_KEYS` | A comma separated list of api keys used for authentication | | `LOCALAGI_CUSTOM_ACTIONS_DIR` | Directory containing custom Go action files to be automatically loaded | @@ -1049,7 +1050,8 @@ LocalAGI supports environment configurations. Note that these environment variab | `LOCALAGI_LLM_API_KEY` | API authentication | | `LOCALAGI_TIMEOUT` | Request timeout settings | | `LOCALAGI_STATE_DIR` | Where state gets stored | -| `LOCALAGI_BASE_URL` | Optional base URL for built-in knowledge base (default `http://localhost:3000`) | +| `LOCALAGI_LOCALRAG_URL` | Optional URL when using an external LocalRAG URL; not used for built-in knowledge base | +| `LOCALAGI_BASE_URL` | Optional base URL for the app (defaults to ":3000") | | `LOCALAGI_SSHBOX_URL` | LocalAGI SSHBox URL, e.g. user:pass@ip:port | | `LOCALAGI_ENABLE_CONVERSATIONS_LOGGING` | Toggle conversation logs | | `LOCALAGI_API_KEYS` | A comma separated list of api keys used for authentication | diff --git a/cmd/env.go b/cmd/env.go index 683ad1f..c06bc52 100644 --- a/cmd/env.go +++ b/cmd/env.go @@ -20,6 +20,7 @@ type Env struct { // Directories and paths StateDir string + LocalAGIURL string LocalRAGURL string CustomActionsDir string SSHBoxURL string @@ -51,6 +52,7 @@ func LoadEnv() Env { TTSModel: envOrDefault("LOCALAGI_TTS_MODEL", ""), Timeout: envOrDefault("LOCALAGI_TIMEOUT", "5m"), StateDir: envOrDefault("LOCALAGI_STATE_DIR", ""), + LocalAGIURL: envOrDefault("LOCALAGI_BASE_URL", ":3000"), LocalRAGURL: os.Getenv("LOCALAGI_LOCALRAG_URL"), CustomActionsDir: os.Getenv("LOCALAGI_CUSTOM_ACTIONS_DIR"), SSHBoxURL: os.Getenv("LOCALAGI_SSHBOX_URL"), diff --git a/cmd/serve.go b/cmd/serve.go index 0e68980..c85b74f 100644 --- a/cmd/serve.go +++ b/cmd/serve.go @@ -123,6 +123,6 @@ func runServe(cmd *cobra.Command, args []string) error { return err } - log.Fatal(app.Listen(":3000")) + log.Fatal(app.Listen(env.LocalAGIURL)) return nil } diff --git a/webui/react-ui/vite.config.js b/webui/react-ui/vite.config.js index 324cc4a..c7dcd7a 100644 --- a/webui/react-ui/vite.config.js +++ b/webui/react-ui/vite.config.js @@ -6,7 +6,7 @@ export default defineConfig(({ mode }) => { const env = loadEnv(mode, process.cwd(), '') // Define backend URL with port from environment variable or default to 8080 - const backendUrl = `http://${env.BACKEND_HOST || 'localhost'}:${env.BACKEND_PORT || '3000'}` + const backendUrl = `http://${env.LOCALAGI_BASE_URL || 'localhost:3000'}` return { plugins: [react()], From 14aed1ae433606a5ffa094308709086047644d64 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Sat, 6 Jun 2026 09:12:51 +0200 Subject: [PATCH 3/6] chore: bump localrecall to index-backed hybrid search (#477) chore: bump localrecall to index-backed RRF hybrid search Pulls in mudler/LocalRecall#46 (merged), which rewrites the PostgreSQL hybrid search using the canonical Reciprocal Rank Fusion pattern (index-backed candidate retrieval + FULL OUTER JOIN + weighted RRF). Fixes the full sequential scan that blew past the statement timeout on large collections (mudler/LocalAI#10186). Signed-off-by: Ettore Di Giacinto --- go.mod | 2 +- go.sum | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/go.mod b/go.mod index 14bcf99..6c0f060 100644 --- a/go.mod +++ b/go.mod @@ -18,7 +18,7 @@ require ( github.com/jung-kurt/gofpdf v1.16.2 github.com/modelcontextprotocol/go-sdk v1.2.0 github.com/mudler/cogito v0.9.5-0.20260315222927-63abdec7189b - github.com/mudler/localrecall v0.6.1-0.20260507074622-a7724fef6f81 + github.com/mudler/localrecall v0.6.3-0.20260606070048-9a3b3321a9cd github.com/mudler/skillserver v0.0.5-0.20260221145827-0639a82c8f49 github.com/mudler/xlog v0.0.5 github.com/onsi/ginkgo/v2 v2.28.1 diff --git a/go.sum b/go.sum index 52fe00f..cad3d56 100644 --- a/go.sum +++ b/go.sum @@ -303,8 +303,8 @@ github.com/mschoch/smat v0.2.0 h1:8imxQsjDm8yFEAVBe7azKmKSgzSkZXDuKkSq9374khM= github.com/mschoch/smat v0.2.0/go.mod h1:kc9mz7DoBKqDyiRL7VZN8KvXQMWeTaVnttLRXOlotKw= github.com/mudler/cogito v0.9.5-0.20260315222927-63abdec7189b h1:A74T2Lauvg61KodYqsjTYDY05kPLcW+efVZjd23dghU= github.com/mudler/cogito v0.9.5-0.20260315222927-63abdec7189b/go.mod h1:6sfja3lcu2nWRzEc0wwqGNu/eCG3EWgij+8s7xyUeQ4= -github.com/mudler/localrecall v0.6.1-0.20260507074622-a7724fef6f81 h1:8D9NJ/ikhsJCxUwbdzIzadw6RqDrW+L0FPqpQQSeux8= -github.com/mudler/localrecall v0.6.1-0.20260507074622-a7724fef6f81/go.mod h1:28k5n19raUrkuwXkacdNsBlj8yuSnGhpT16tu+2+4dU= +github.com/mudler/localrecall v0.6.3-0.20260606070048-9a3b3321a9cd h1:trn9D5UHAE6zdRyD2uX04W1tLSslAwozVwcyNTd72Ak= +github.com/mudler/localrecall v0.6.3-0.20260606070048-9a3b3321a9cd/go.mod h1:28k5n19raUrkuwXkacdNsBlj8yuSnGhpT16tu+2+4dU= github.com/mudler/skillserver v0.0.5-0.20260221145827-0639a82c8f49 h1:dAF1ALXqqapRZo80x56BIBBcPrPbRNerbd66rdyO8J4= github.com/mudler/skillserver v0.0.5-0.20260221145827-0639a82c8f49/go.mod h1:z3yFhcL9bSykmmh6xgGu0hyoItd4CnxgtWMEWw8uFJU= github.com/mudler/xlog v0.0.5 h1:2unBuVC5rNGhCC86UaA94TElWFml80NL5XLK+kAmNuU= From 9a6b32f9bcc7067f24629ed76c3eb36061afec02 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 18 Jun 2026 16:48:34 +0200 Subject: [PATCH 4/6] =?UTF-8?q?chore:=20bump=20localrecall=20to=20include?= =?UTF-8?q?=20PostgreSQL=20table-name=20sanitization=E2=80=A6=20(#478)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit chore: bump localrecall to include PostgreSQL table-name sanitization fix Pulls mudler/LocalRecall#48, which makes sanitizeTableName allowlist valid identifier characters so collection names containing ':' (e.g. the per-user "legacy-api-key:" namespace used by LocalAI) no longer break PostgreSQL CREATE TABLE with "syntax error at or near ':'". Ref: mudler/LocalAI#10375 Signed-off-by: Ettore Di Giacinto --- go.mod | 2 +- go.sum | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/go.mod b/go.mod index 6c0f060..8747469 100644 --- a/go.mod +++ b/go.mod @@ -18,7 +18,7 @@ require ( github.com/jung-kurt/gofpdf v1.16.2 github.com/modelcontextprotocol/go-sdk v1.2.0 github.com/mudler/cogito v0.9.5-0.20260315222927-63abdec7189b - github.com/mudler/localrecall v0.6.3-0.20260606070048-9a3b3321a9cd + github.com/mudler/localrecall v0.6.3-0.20260618142827-d0073dd5dc32 github.com/mudler/skillserver v0.0.5-0.20260221145827-0639a82c8f49 github.com/mudler/xlog v0.0.5 github.com/onsi/ginkgo/v2 v2.28.1 diff --git a/go.sum b/go.sum index cad3d56..54aab36 100644 --- a/go.sum +++ b/go.sum @@ -303,8 +303,8 @@ github.com/mschoch/smat v0.2.0 h1:8imxQsjDm8yFEAVBe7azKmKSgzSkZXDuKkSq9374khM= github.com/mschoch/smat v0.2.0/go.mod h1:kc9mz7DoBKqDyiRL7VZN8KvXQMWeTaVnttLRXOlotKw= github.com/mudler/cogito v0.9.5-0.20260315222927-63abdec7189b h1:A74T2Lauvg61KodYqsjTYDY05kPLcW+efVZjd23dghU= github.com/mudler/cogito v0.9.5-0.20260315222927-63abdec7189b/go.mod h1:6sfja3lcu2nWRzEc0wwqGNu/eCG3EWgij+8s7xyUeQ4= -github.com/mudler/localrecall v0.6.3-0.20260606070048-9a3b3321a9cd h1:trn9D5UHAE6zdRyD2uX04W1tLSslAwozVwcyNTd72Ak= -github.com/mudler/localrecall v0.6.3-0.20260606070048-9a3b3321a9cd/go.mod h1:28k5n19raUrkuwXkacdNsBlj8yuSnGhpT16tu+2+4dU= +github.com/mudler/localrecall v0.6.3-0.20260618142827-d0073dd5dc32 h1:RP4BVGTHHpJIrGAwqRD3Wq1wmURmc1SxhwacnIWgI+g= +github.com/mudler/localrecall v0.6.3-0.20260618142827-d0073dd5dc32/go.mod h1:28k5n19raUrkuwXkacdNsBlj8yuSnGhpT16tu+2+4dU= github.com/mudler/skillserver v0.0.5-0.20260221145827-0639a82c8f49 h1:dAF1ALXqqapRZo80x56BIBBcPrPbRNerbd66rdyO8J4= github.com/mudler/skillserver v0.0.5-0.20260221145827-0639a82c8f49/go.mod h1:z3yFhcL9bSykmmh6xgGu0hyoItd4CnxgtWMEWw8uFJU= github.com/mudler/xlog v0.0.5 h1:2unBuVC5rNGhCC86UaA94TElWFml80NL5XLK+kAmNuU= From f32543e82ee1144f74a47b32b506de2c28d0c5de Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Sat, 22 Aug 2026 08:33:30 +0200 Subject: [PATCH 5/6] telegram rich streaming (#488) * docs: design Telegram rich response streaming * feat(telegram): add rich message API client * feat(telegram): preserve rich response formatting * feat(agent): add request stream callbacks * feat(telegram): add rich stream sessions * fix(telegram): preserve initial stream preview * fix(telegram): preserve streamed delivery fallbacks * fix(telegram): complete stream worker delivery * feat(telegram): integrate rich response streaming * fix(telegram): complete native streaming integration * fix(telegram): finish rich streaming lifecycle * fix(telegram): preserve final delivery after job completion * fix(telegram): honor delivery lifecycle contexts * fix(telegram): serialize streamed status delivery * fix(telegram): return preview flush errors * fix(telegram): preserve final-only delivery context * docs(telegram): explain streaming fallbacks --- README.md | 8 + core/agent/agent.go | 28 +- core/agent/stream_callback_test.go | 54 ++ core/types/job.go | 7 + ...26-08-21-telegram-rich-streaming-design.md | 176 ++++++ services/connectors/telegram.go | 493 ++++++++++------- services/connectors/telegram_api.go | 145 +++++ services/connectors/telegram_api_test.go | 166 ++++++ services/connectors/telegram_format.go | 194 +++++++ services/connectors/telegram_format_test.go | 92 ++++ .../connectors/telegram_integration_test.go | 461 ++++++++++++++++ services/connectors/telegram_stream.go | 414 ++++++++++++++ services/connectors/telegram_stream_test.go | 507 ++++++++++++++++++ 13 files changed, 2544 insertions(+), 201 deletions(-) create mode 100644 core/agent/stream_callback_test.go create mode 100644 docs/superpowers/specs/2026-08-21-telegram-rich-streaming-design.md create mode 100644 services/connectors/telegram_api.go create mode 100644 services/connectors/telegram_api_test.go create mode 100644 services/connectors/telegram_format.go create mode 100644 services/connectors/telegram_format_test.go create mode 100644 services/connectors/telegram_integration_test.go create mode 100644 services/connectors/telegram_stream.go create mode 100644 services/connectors/telegram_stream_test.go diff --git a/README.md b/README.md index cb3b881..22016d7 100644 --- a/README.md +++ b/README.md @@ -14,6 +14,8 @@ Try on [![Telegram](https://img.shields.io/badge/Telegram-2CA5E0?style=for-the-badge&logo=telegram&logoColor=white)](https://t.me/LocalAGI_bot) +Telegram response streaming is enabled by default (`"streaming": "true"`). Private chats use native rich drafts, while groups progressively edit a placeholder. Set `"streaming": "false"` to suppress previews; final responses still use rich Markdown with MarkdownV2 and plain-text fallbacks. + Create customizable AI assistants, automations, chat bots and agents that run 100% locally. No need for agentic Python libraries or cloud service keys, just bring your GPU (or even just CPU) and a web browser. @@ -826,6 +828,12 @@ Configuration options: - `mention_only`: When enabled, bot only responds when mentioned in groups - `admins`: Comma-separated list of Telegram usernames allowed to use the bot in private chats - `channel_id`: Optional channel ID for the bot to send messages to +- `streaming`: Show progressive responses. Defaults to `true`; set it to `false` for final-only output. + +Private chats use native rich drafts when the configured Telegram Bot API +supports the current rich-message methods. If those methods are unavailable, +the connector automatically falls back to progressive message edits. Final +responses fall back from rich Markdown to MarkdownV2 and then plain text. > **Important**: For group functionality to work properly: > 1. Go to @BotFather diff --git a/core/agent/agent.go b/core/agent/agent.go index 115c7c4..0074bea 100644 --- a/core/agent/agent.go +++ b/core/agent/agent.go @@ -225,6 +225,21 @@ func (a *Agent) Context() context.Context { return a.context.Context } +func (a *Agent) streamCallbackForJob(job *types.Job) func(cogito.StreamEvent) { + agentCallback := a.options.streamCallback + requestCallback := job.StreamCallback + if agentCallback == nil { + return requestCallback + } + if requestCallback == nil { + return agentCallback + } + return func(event cogito.StreamEvent) { + agentCallback(event) + requestCallback(event) + } +} + // Ask is a blocking call that returns the response as soon as it's ready. // It discards any other computation. func (a *Agent) Ask(opts ...types.JobOption) *types.JobResult { @@ -896,6 +911,7 @@ func (a *Agent) addFunctionResultToConversation(ctx context.Context, chosenActio } func (a *Agent) consumeJob(job *types.Job, role string) { + streamCallback := a.streamCallbackForJob(job) if err := job.GetContext().Err(); err != nil { job.Result.Finish(fmt.Errorf("expired")) return @@ -1063,8 +1079,8 @@ func (a *Agent) consumeJob(job *types.Job, role string) { return } // Forward reasoning to stream callback - if a.options.streamCallback != nil { - a.options.streamCallback(cogito.StreamEvent{ + if streamCallback != nil { + streamCallback(cogito.StreamEvent{ Type: cogito.StreamEventReasoning, Content: s, }) @@ -1161,12 +1177,12 @@ func (a *Agent) consumeJob(job *types.Job, role string) { } // Forward tool selection to stream callback - if a.options.streamCallback != nil { + if streamCallback != nil { toolName := tc.Name if chosenAction != nil { toolName = chosenAction.Definition().Name.String() } - a.options.streamCallback(cogito.StreamEvent{ + streamCallback(cogito.StreamEvent{ Type: cogito.StreamEventToolCall, ToolName: toolName, ToolArgs: fmt.Sprintf("%v", tc.Arguments), @@ -1361,8 +1377,8 @@ func (a *Agent) consumeJob(job *types.Job, role string) { cogitoOpts = append(cogitoOpts, cogito.WithMaxRetries(a.options.maxAttempts)) } - if a.options.streamCallback != nil { - cogitoOpts = append(cogitoOpts, cogito.WithStreamCallback(a.options.streamCallback)) + if streamCallback != nil { + cogitoOpts = append(cogitoOpts, cogito.WithStreamCallback(streamCallback)) } fragment, err = cogito.ExecuteTools( diff --git a/core/agent/stream_callback_test.go b/core/agent/stream_callback_test.go new file mode 100644 index 0000000..212d68e --- /dev/null +++ b/core/agent/stream_callback_test.go @@ -0,0 +1,54 @@ +package agent + +import ( + "testing" + + "github.com/mudler/LocalAGI/core/types" + "github.com/mudler/cogito" +) + +func TestStreamCallbackForJobCombinesAgentAndRequestCallbacks(t *testing.T) { + var agentEvents, firstRequestEvents, secondRequestEvents []cogito.StreamEvent + a := &Agent{options: &options{ + streamCallback: func(event cogito.StreamEvent) { + agentEvents = append(agentEvents, event) + }, + }} + first := types.NewJob(types.WithStreamCallback(func(event cogito.StreamEvent) { + firstRequestEvents = append(firstRequestEvents, event) + })) + second := types.NewJob(types.WithStreamCallback(func(event cogito.StreamEvent) { + secondRequestEvents = append(secondRequestEvents, event) + })) + + firstEvent := cogito.StreamEvent{Content: "first"} + secondEvent := cogito.StreamEvent{Content: "second"} + a.streamCallbackForJob(first)(firstEvent) + a.streamCallbackForJob(second)(secondEvent) + + if len(agentEvents) != 2 || agentEvents[0].Content != "first" || agentEvents[1].Content != "second" { + t.Fatalf("agent callback events = %#v, want first and second events", agentEvents) + } + if len(firstRequestEvents) != 1 || firstRequestEvents[0].Content != "first" { + t.Fatalf("first request callback events = %#v, want only first event", firstRequestEvents) + } + if len(secondRequestEvents) != 1 || secondRequestEvents[0].Content != "second" { + t.Fatalf("second request callback events = %#v, want only second event", secondRequestEvents) + } +} + +func TestStreamCallbackForJobNilRequestCallbackPreservesAgentCallback(t *testing.T) { + var events []cogito.StreamEvent + a := &Agent{options: &options{ + streamCallback: func(event cogito.StreamEvent) { + events = append(events, event) + }, + }} + + callback := a.streamCallbackForJob(types.NewJob(types.WithStreamCallback(nil))) + callback(cogito.StreamEvent{Content: "agent"}) + + if len(events) != 1 || events[0].Content != "agent" { + t.Fatalf("agent callback events = %#v, want agent event", events) + } +} diff --git a/core/types/job.go b/core/types/job.go index 3183903..a482454 100644 --- a/core/types/job.go +++ b/core/types/job.go @@ -22,6 +22,7 @@ type Job struct { Result *JobResult ReasoningCallback func(ActionCurrentState) bool ResultCallback func(ActionState) + StreamCallback func(cogito.StreamEvent) ConversationHistory []openai.ChatCompletionMessage UUID string Metadata map[string]interface{} @@ -82,6 +83,12 @@ func WithResultCallback(f func(ActionState)) JobOption { } } +func WithStreamCallback(f func(cogito.StreamEvent)) JobOption { + return func(j *Job) { + j.StreamCallback = f + } +} + func WithMetadata(metadata map[string]any) JobOption { return func(j *Job) { j.Metadata = metadata diff --git a/docs/superpowers/specs/2026-08-21-telegram-rich-streaming-design.md b/docs/superpowers/specs/2026-08-21-telegram-rich-streaming-design.md new file mode 100644 index 0000000..6242255 --- /dev/null +++ b/docs/superpowers/specs/2026-08-21-telegram-rich-streaming-design.md @@ -0,0 +1,176 @@ +# Telegram rich streaming design + +## Summary + +The Telegram connector streams agent output as Telegram renders it. Private chats use native rich-message drafts. Other chats use progressive message edits. The connector sends the final answer as a persistent rich message. + +Streaming is enabled by default. Operators can disable it with the connector setting `streaming=false`. + +## Goals + +- Show answer text while the model generates it. +- Render the model's Markdown in Telegram. +- Match the native streaming behavior of Hermes Agent. +- Keep final delivery reliable when a streaming or formatting API is unavailable. +- Isolate simultaneous responses so that one job cannot receive another job's output. +- Preserve the existing reasoning and tool-status feedback. + +## Non-goals + +- Stream generated audio, images, or songs. +- Add rich-message media blocks. +- Change the format of conversation history. +- Change connectors other than Telegram. +- Guarantee native drafts in groups. Telegram limits draft methods to private chats. + +## User experience + +### Private chats + +The connector starts a native rich draft with a Telegram thinking block. It appends answer deltas to the draft as the model produces them. Telegram animates updates that use the same nonzero draft ID. + +The connector sends the completed answer with `sendRichMessage`. The persistent message uses the model's raw Markdown. The ephemeral draft then expires according to Telegram behavior. + +### Groups and supergroups + +The connector sends the existing thinking placeholder as a reply. It edits that message as answer deltas arrive. The final edit contains formatted content when Telegram accepts it. + +### Reasoning and tool activity + +Before answer content starts, the preview can show reasoning and tool activity through the existing status accumulator. Answer content takes priority after the first content delta. The connector does not expose hidden reasoning that the agent does not already publish. + +### Disabled streaming + +If `streaming=false`, the connector keeps the thinking placeholder and sends only the completed answer. Rich Markdown formatting remains enabled. + +## Architecture + +### Per-request stream callback + +The agent API gains a per-request stream callback option. `Ask` passes this callback to the Cogito request without replacing the agent's existing callback. + +The agent invokes both callbacks for each event: + +- The existing agent callback continues to publish SSE events. +- The request callback receives events for one Telegram job. + +This boundary prevents concurrent jobs from sharing an unscoped global callback. It also prevents Telegram integration from disabling web UI streaming. + +### Telegram stream session + +Each incoming Telegram request creates one stream session. The session owns these values: + +- Chat ID and chat type. +- Draft ID or placeholder message ID. +- Accumulated answer content. +- The most recent delivered content. +- Delivery mode: rich draft, edited message, or final-only. +- Throttle timer and cancellation context. + +The request callback only appends event data and signals the session worker. One worker performs Telegram API calls in order. This structure avoids concurrent edits and protects the model callback from network latency. + +The handler closes and flushes the session before final delivery. Cleanup stops its timer and releases its goroutine on success, cancellation, or error. + +### Telegram Bot API client + +The pinned `go-telegram/bot` version does not expose the Bot API 10.2 rich-message methods. A small internal client calls these methods through the Bot API HTTP endpoint: + +- `sendRichMessageDraft` +- `sendRichMessage` + +The client uses typed request and response structures. It uses the request context and returns Telegram's error description. Existing operations continue to use `go-telegram/bot`. + +The client interface stays private to the connector package. Tests replace it with a fake implementation. + +## Streaming policy + +The session coalesces character and token deltas. It sends at most one preview update every 400 milliseconds. This interval gives visible progress without one API request per token. + +The worker skips an update when its content equals the last delivered content. It flushes pending content when the agent reports completion or `Ask` returns. + +Telegram drafts expire after 30 seconds. The worker sends an unchanged heartbeat before expiry when generation remains active. The heartbeat uses the same draft ID. + +The connector respects Telegram flood-control responses. If Telegram returns a retry delay, the worker delays the next preview update. Preview retries never delay final delivery beyond the request context. + +## Markdown and delivery fallback + +Rich delivery passes the unescaped response in `InputRichMessage.markdown`. This mode supports GitHub-style Markdown where Telegram supports it. It also supports headings, lists, tables, code blocks, quotations, details, footnotes, and formulas. + +The connector does not call `bot.EscapeMarkdown` on rich Markdown. The current escape step removes formatting and must not remain on this path. + +Delivery uses this fallback order: + +1. Send a rich draft in a private chat. +2. If the draft fails, switch that response to progressive edits. +3. Finalize with a rich persistent message or rich edit. +4. If rich formatting fails, convert the content to Telegram MarkdownV2. +5. If MarkdownV2 fails, send plain text without a parse mode. + +A preview failure changes only the active response. The connector can try native drafts again for the next response. + +The URL appendix becomes ordinary Markdown and stays part of the raw response. Link previews remain disabled for streamed and final output where the API supports that option. + +## Message limits + +Rich messages support more text than legacy messages, but the connector still handles every Telegram limit. A splitter operates on UTF-8 text and preserves complete code fences when possible. + +For a response that requires multiple persistent messages: + +- The first message finalizes the active preview. +- The connector sends later chunks in order. +- Every chunk uses the same formatting fallback order. +- Group chunks retain the reply relationship where Telegram permits it. + +The streaming preview shows the current tail within the applicable draft or edit limit. The accumulated session retains the complete response for final delivery. + +## Error handling + +- An empty final response replaces the preview with the existing internal-error text. +- A cancelled job stops preview updates and keeps the content already shown. +- A preview API failure logs the method and Telegram description without logging the bot token. +- A final rich-delivery failure falls back to MarkdownV2, then plain text. +- A total final-delivery failure returns through the existing connector error path. +- Multimedia and text-to-speech behavior remains unchanged. + +## Configuration + +`TelegramConfigMeta` adds this field: + +- `streaming`: A boolean that enables progressive Telegram output. The default is `true` when the field is absent. + +No token or BotFather setting is required. Unsupported Bot API servers use the fallback path. + +## Testing + +Unit tests use named, table-driven subtests where cases share a contract. Tests cover these behaviors: + +- Private chats start and update one rich draft with a stable draft ID. +- Private chats send one persistent rich final response. +- Groups edit one placeholder instead of creating a draft. +- Character deltas coalesce into throttled updates. +- A final flush includes pending content. +- Concurrent jobs do not mix their content. +- The request callback does not replace the SSE callback. +- Draft failure switches only the current response to message edits. +- Rich final failure falls back to MarkdownV2 and then plain text. +- Raw Markdown reaches the rich-message client without escaping. +- Long UTF-8 responses split without data loss. +- Cancellation stops timers and workers. +- `streaming=false` suppresses preview updates but keeps rich final formatting. + +HTTP contract tests use `httptest.Server` to inspect Bot API payloads. They do not contact Telegram. Relevant package tests also run with the race detector. + +## Documentation + +The Telegram section in `README.md` documents the default streaming behavior and the `streaming` setting. It states that native drafts require a current Telegram Bot API. It also describes the automatic edit and formatting fallbacks. + +## Acceptance criteria + +- A private-chat user sees an animated rich preview while the model writes. +- A group-chat user sees the reply message update while the model writes. +- The final answer renders supported Markdown instead of showing Markdown source. +- Streaming is active when the configuration omits `streaming`. +- Setting `streaming=false` restores final-only output. +- An unsupported rich-message API does not prevent final delivery. +- Simultaneous Telegram jobs do not mix content, and existing SSE delivery continues unchanged. +- The test suite and race-enabled connector tests pass. diff --git a/services/connectors/telegram.go b/services/connectors/telegram.go index a4033d3..9a2b883 100644 --- a/services/connectors/telegram.go +++ b/services/connectors/telegram.go @@ -22,20 +22,24 @@ import ( "github.com/mudler/LocalAGI/core/agent" "github.com/mudler/LocalAGI/core/types" "github.com/mudler/LocalAGI/pkg/config" - "github.com/mudler/LocalAGI/services/connectors/common" "github.com/mudler/LocalAGI/pkg/xstrings" "github.com/mudler/LocalAGI/services/actions" + "github.com/mudler/LocalAGI/services/connectors/common" "github.com/mudler/xlog" "github.com/sashabaranov/go-openai" ) const telegramThinkingMessage = "🤔 thinking..." const telegramMaxMessageLength = 3000 +const telegramStreamingMetadataKey = "telegram_streaming" type Telegram struct { Token string bot *bot.Bot agent *agent.Agent + api telegramAPI + + streaming bool admins []string @@ -53,6 +57,155 @@ type Telegram struct { mentionOnly bool } +func telegramAskOptions(history []openai.ChatCompletionMessage, jobUUID string, metadata map[string]any, session *telegramStreamSession) []types.JobOption { + opts := []types.JobOption{ + types.WithConversationHistory(history), + types.WithUUID(jobUUID), + types.WithMetadata(metadata), + } + if session != nil { + opts = append(opts, types.WithStreamCallback(session.Accept)) + } + return opts +} + +func telegramNewJobWithStream(parent context.Context, api telegramAPI, chatID int64, private bool, delivery telegramStreamDelivery, history []openai.ChatCompletionMessage, jobUUID string, metadata map[string]any) (*types.Job, *telegramStreamSession) { + if metadata == nil { + metadata = make(map[string]any) + } + metadata[telegramStreamingMetadataKey] = true + opts := append(telegramAskOptions(history, jobUUID, metadata, nil), types.WithContext(parent)) + job := types.NewJob(opts...) + session := newTelegramStreamSessionWithContexts(parent, job.GetContext(), api, chatID, private, delivery, telegramDraftHeartbeatInterval) + job.StreamCallback = session.Accept + return job, session +} + +func telegramUseLegacyStatusDelivery(job *types.Job) bool { + if job == nil || job.Metadata == nil { + return true + } + streaming, _ := job.Metadata[telegramStreamingMetadataKey].(bool) + return !streaming +} + +func telegramDeliverLegacyStatus(job *types.Job, deliver func()) { + if telegramUseLegacyStatusDelivery(job) { + deliver() + } +} + +type telegramMessageBot interface { + SendMessage(context.Context, *bot.SendMessageParams) (*models.Message, error) + EditMessageText(context.Context, *bot.EditMessageTextParams) (*models.Message, error) + DeleteMessage(context.Context, *bot.DeleteMessageParams) (bool, error) +} + +type telegramJobExecutor interface { + Execute(*types.Job) *types.JobResult +} + +func telegramExecuteJob(executor telegramJobExecutor, job *types.Job) *types.JobResult { + return executor.Execute(job) +} + +func (t *Telegram) telegramDelivery(_ context.Context, b telegramMessageBot, chatID int64, replyTo int, jobUUID string, initialMessageID int) telegramStreamDelivery { + var mu sync.Mutex + messageID := initialMessageID + ensurePlaceholder := func(ctx context.Context, text string, mode models.ParseMode) (int, bool, error) { + mu.Lock() + defer mu.Unlock() + if messageID != 0 { + return messageID, false, nil + } + params := &bot.SendMessageParams{ChatID: chatID, Text: text, ParseMode: mode} + disabled := true + params.LinkPreviewOptions = &models.LinkPreviewOptions{IsDisabled: &disabled} + if replyTo != 0 { + params.ReplyParameters = &models.ReplyParameters{MessageID: replyTo} + } + msg, err := b.SendMessage(ctx, params) + if err != nil { + return 0, false, err + } + messageID = msg.ID + t.placeholderMutex.Lock() + t.placeholders[jobUUID] = messageID + t.placeholderMutex.Unlock() + return messageID, true, nil + } + sendChunks := func(ctx context.Context, chunks []string, mode models.ParseMode) error { + for i, chunk := range chunks { + if i == 0 { + if id, created, err := ensurePlaceholder(ctx, chunk, mode); err != nil { + return err + } else if !created { + disabled := true + if _, err := b.EditMessageText(ctx, &bot.EditMessageTextParams{ChatID: chatID, MessageID: id, Text: chunk, ParseMode: mode, LinkPreviewOptions: &models.LinkPreviewOptions{IsDisabled: &disabled}}); err != nil { + return err + } + } else if mode != "" { + // Creation already delivered identical text. Parse mode is relevant + // only to final fallback, whose placeholders normally pre-exist. + _ = id + } + continue + } + params := &bot.SendMessageParams{ChatID: chatID, Text: chunk, ParseMode: mode} + disabled := true + params.LinkPreviewOptions = &models.LinkPreviewOptions{IsDisabled: &disabled} + if replyTo != 0 { + params.ReplyParameters = &models.ReplyParameters{MessageID: replyTo} + } + if _, err := b.SendMessage(ctx, params); err != nil { + return err + } + } + return nil + } + return telegramStreamDelivery{ + editPreview: func(ctx context.Context, _ int64, text string) error { + id, created, err := ensurePlaceholder(ctx, text, "") + if err != nil { + return err + } + if created { + return nil + } + disabled := true + _, err = b.EditMessageText(ctx, &bot.EditMessageTextParams{ChatID: chatID, MessageID: id, Text: text, LinkPreviewOptions: &models.LinkPreviewOptions{IsDisabled: &disabled}}) + return err + }, + finalMarkdown: func(ctx context.Context, _ int64, chunks []string) error { + return sendChunks(ctx, chunks, models.ParseModeMarkdown) + }, + finalPlain: func(ctx context.Context, _ int64, chunks []string) error { + return sendChunks(ctx, chunks, "") + }, + clearPreview: func(ctx context.Context, _ int64) error { + mu.Lock() + id := messageID + messageID = 0 + mu.Unlock() + if id == 0 { + return nil + } + t.placeholderMutex.Lock() + delete(t.placeholders, jobUUID) + t.placeholderMutex.Unlock() + _, err := b.DeleteMessage(ctx, &bot.DeleteMessageParams{ChatID: chatID, MessageID: id}) + return err + }, + replyTo: replyTo, + } +} + +func telegramFinalSession(ctx context.Context, api telegramAPI, chatID int64, private bool, delivery telegramStreamDelivery) *telegramStreamSession { + finalCtx := ctx + ctx, cancel := context.WithCancel(ctx) + return &telegramStreamSession{ctx: ctx, finalCtx: finalCtx, cancel: cancel, api: api, chatID: chatID, private: private, delivery: delivery} +} + // isBotMentioned checks if the bot is mentioned in the message func (t *Telegram) isBotMentioned(message string, botUsername string) bool { return strings.Contains(message, "@"+botUsername) @@ -261,7 +414,7 @@ func (t *Telegram) handleGroupMessage(ctx context.Context, b *bot.Bot, a *agent. // Add chat ID and conversation_id for tracking and cancel-previous-on-new-message metadata := map[string]interface{}{ - "chatID": update.Message.Chat.ID, + "chatID": update.Message.Chat.ID, types.MetadataKeyConversationID: fmt.Sprintf("telegram:%d", update.Message.Chat.ID), } @@ -282,12 +435,15 @@ func (t *Telegram) handleGroupMessage(ctx context.Context, b *bot.Bot, a *agent. currentConv := a.SharedState().ConversationTracker.GetConversation(fmt.Sprintf("telegram:%d", update.Message.Chat.ID)) - // Create a new job with the conversation history and metadata - job := types.NewJob( - types.WithConversationHistory(currentConv), - types.WithUUID(jobUUID), - types.WithMetadata(metadata), - ) + delivery := t.telegramDelivery(ctx, b, update.Message.Chat.ID, update.Message.ID, jobUUID, msg.ID) + var streamSession *telegramStreamSession + var job *types.Job + if t.streaming { + job, streamSession = telegramNewJobWithStream(ctx, t.api, update.Message.Chat.ID, false, delivery, currentConv, jobUUID, metadata) + defer streamSession.Close() + } else { + job = types.NewJob(telegramAskOptions(currentConv, jobUUID, metadata, nil)...) + } // Mark this chat as having an active job t.activeJobsMutex.Lock() @@ -313,20 +469,16 @@ func (t *Telegram) handleGroupMessage(ctx context.Context, b *bot.Bot, a *agent. t.placeholderMutex.Unlock() }() - res := a.Ask( - types.WithConversationHistory(currentConv), - types.WithUUID(jobUUID), - types.WithMetadata(metadata), - ) + res := telegramExecuteJob(a, job) + if streamSession != nil { + if err := streamSession.Flush(); err != nil { + xlog.Error("Error flushing Telegram stream", "error", err) + } + } if res.Response == "" { xlog.Error("Empty response from agent") - _, err := b.EditMessageText(ctx, &bot.EditMessageTextParams{ - ChatID: update.Message.Chat.ID, - MessageID: msg.ID, - Text: "there was an internal error. try again!", - }) - if err != nil { + if err := delivery.finalPlain(ctx, update.Message.Chat.ID, []string{"there was an internal error. try again!"}); err != nil { xlog.Error("Error updating error message", "error", err) } return @@ -360,12 +512,9 @@ func (t *Telegram) handleGroupMessage(ctx context.Context, b *bot.Bot, a *agent. xlog.Error("Error sending audio response", "error", err) } else { xlog.Debug("Audio response sent successfully") - // Remove the thinking placeholder message before returning - _, err := t.bot.DeleteMessage(ctx, &bot.DeleteMessageParams{ - ChatID: update.Message.Chat.ID, - MessageID: msg.ID, - }) - if err != nil { + // Remove any legacy preview before returning. Native drafts are + // superseded by the audio message itself. + if err := delivery.clearPreview(ctx, update.Message.Chat.ID); err != nil { xlog.Error("Error deleting thinking placeholder", "error", err) } // Don't send text response if audio was sent successfully @@ -374,13 +523,7 @@ func (t *Telegram) handleGroupMessage(ctx context.Context, b *bot.Bot, a *agent. } } - // Update the message with the final response - formattedResponse := formatResponseWithURLs(res.Response, urls) - - // Split the message if it's too long - messages := xstrings.SplitParagraph(formattedResponse, telegramMaxMessageLength) - - if len(messages) == 0 { + if len(telegramFormatResponse(res.Response, urls, telegramMaxMessageLength)) == 0 { _, err := b.EditMessageText(ctx, &bot.EditMessageTextParams{ ChatID: update.Message.Chat.ID, MessageID: msg.ID, @@ -392,31 +535,14 @@ func (t *Telegram) handleGroupMessage(ctx context.Context, b *bot.Bot, a *agent. return } - // Update the first message - _, err = b.EditMessageText(ctx, &bot.EditMessageTextParams{ - ChatID: update.Message.Chat.ID, - MessageID: msg.ID, - Text: messages[0], - ParseMode: models.ParseModeMarkdown, - }) - if err != nil { - xlog.Error("Error updating message", "error", err) - return + var finalErr error + if streamSession != nil { + finalErr = streamSession.Finalize(res.Response, urls) + } else { + finalErr = telegramFinalSession(ctx, t.api, update.Message.Chat.ID, false, delivery).deliverFinal(res.Response, urls) } - - // Send additional chunks as new messages - for i := 1; i < len(messages); i++ { - _, err = b.SendMessage(ctx, &bot.SendMessageParams{ - ChatID: update.Message.Chat.ID, - Text: messages[i], - ParseMode: models.ParseModeMarkdown, - ReplyParameters: &models.ReplyParameters{ - MessageID: update.Message.ID, - }, - }) - if err != nil { - xlog.Error("Error sending additional message", "error", err) - } + if finalErr != nil { + xlog.Error("Error delivering final Telegram response", "error", finalErr) } } @@ -433,29 +559,31 @@ func (t *Telegram) AgentResultCallback() func(state types.ActionState) { return } - // Update placeholder with tool result if still in progress - t.placeholderMutex.Lock() - msgID, exists := t.placeholders[job.UUID] - if exists && msgID != 0 && t.bot != nil { - acc, ok := t.jobStatus[job.UUID] - if !ok { - acc = common.NewStatusAccumulator() - t.jobStatus[job.UUID] = acc - } - acc.AppendToolResult(common.ActionDisplayName(state.Action), state.Result) - thought := acc.BuildMessage(telegramThinkingMessage, telegramMaxMessageLength) - t.placeholderMutex.Unlock() - _, err := t.bot.EditMessageText(t.agent.Context(), &bot.EditMessageTextParams{ - ChatID: chatID, - MessageID: msgID, - Text: thought, - }) - if err != nil { - xlog.Error("Error updating tool result message", "error", err) - } + telegramDeliverLegacyStatus(job, func() { + // Update placeholder with tool result if still in progress. t.placeholderMutex.Lock() - } - t.placeholderMutex.Unlock() + msgID, exists := t.placeholders[job.UUID] + if exists && msgID != 0 && t.bot != nil { + acc, ok := t.jobStatus[job.UUID] + if !ok { + acc = common.NewStatusAccumulator() + t.jobStatus[job.UUID] = acc + } + acc.AppendToolResult(common.ActionDisplayName(state.Action), state.Result) + thought := acc.BuildMessage(telegramThinkingMessage, telegramMaxMessageLength) + t.placeholderMutex.Unlock() + _, err := t.bot.EditMessageText(t.agent.Context(), &bot.EditMessageTextParams{ + ChatID: chatID, + MessageID: msgID, + Text: thought, + }) + if err != nil { + xlog.Error("Error updating tool result message", "error", err) + } + t.placeholderMutex.Lock() + } + t.placeholderMutex.Unlock() + }) t.activeJobsMutex.Lock() delete(t.activeJobs, chatID) @@ -465,46 +593,48 @@ func (t *Telegram) AgentResultCallback() func(state types.ActionState) { func (t *Telegram) AgentReasoningCallback() func(state types.ActionCurrentState) bool { return func(state types.ActionCurrentState) bool { - t.placeholderMutex.Lock() - msgID, exists := t.placeholders[state.Job.UUID] - chatID := int64(0) - if state.Job.Metadata != nil { - if ch, ok := state.Job.Metadata["chatID"].(int64); ok { - chatID = ch + telegramDeliverLegacyStatus(state.Job, func() { + t.placeholderMutex.Lock() + msgID, exists := t.placeholders[state.Job.UUID] + chatID := int64(0) + if state.Job.Metadata != nil { + if ch, ok := state.Job.Metadata["chatID"].(int64); ok { + chatID = ch + } } - } - if !exists || msgID == 0 || chatID == 0 || t.bot == nil { + if !exists || msgID == 0 || chatID == 0 || t.bot == nil { + t.placeholderMutex.Unlock() + return + } + + if state.Reasoning == "" && state.Action == nil { + t.placeholderMutex.Unlock() + return + } + + acc, ok := t.jobStatus[state.Job.UUID] + if !ok { + acc = common.NewStatusAccumulator() + t.jobStatus[state.Job.UUID] = acc + } + if state.Reasoning != "" { + acc.AppendReasoning(state.Reasoning) + } + if state.Action != nil { + acc.AppendToolCall(common.ActionDisplayName(state.Action), state.Params.String()) + } + thought := acc.BuildMessage(telegramThinkingMessage, telegramMaxMessageLength) t.placeholderMutex.Unlock() - return true - } - if state.Reasoning == "" && state.Action == nil { - t.placeholderMutex.Unlock() - return true - } - - acc, ok := t.jobStatus[state.Job.UUID] - if !ok { - acc = common.NewStatusAccumulator() - t.jobStatus[state.Job.UUID] = acc - } - if state.Reasoning != "" { - acc.AppendReasoning(state.Reasoning) - } - if state.Action != nil { - acc.AppendToolCall(common.ActionDisplayName(state.Action), state.Params.String()) - } - thought := acc.BuildMessage(telegramThinkingMessage, telegramMaxMessageLength) - t.placeholderMutex.Unlock() - - _, err := t.bot.EditMessageText(t.agent.Context(), &bot.EditMessageTextParams{ - ChatID: chatID, - MessageID: msgID, - Text: thought, + _, err := t.bot.EditMessageText(t.agent.Context(), &bot.EditMessageTextParams{ + ChatID: chatID, + MessageID: msgID, + Text: thought, + }) + if err != nil { + xlog.Error("Error updating reasoning message", "error", err) + } }) - if err != nil { - xlog.Error("Error updating reasoning message", "error", err) - } return true } } @@ -672,19 +802,6 @@ func (t *Telegram) handleMultimediaContent(ctx context.Context, chatID int64, re return urls, nil } -// formatResponseWithURLs formats the response text and creates message entities for URLs -func formatResponseWithURLs(response string, urls []string) string { - finalResponse := response - if len(urls) > 0 { - finalResponse += "\n\nReferences:\n" - for i, url := range urls { - finalResponse += fmt.Sprintf("🔗 %d. %s\n", i+1, url) - } - } - - return bot.EscapeMarkdown(finalResponse) -} - func (t *Telegram) handleUpdate(ctx context.Context, b *bot.Bot, a *agent.Agent, update *models.Update) { if update.Message == nil || update.Message.From == nil { xlog.Debug("Message or user is nil", "update", update) @@ -738,27 +855,23 @@ func (t *Telegram) handleUpdate(ctx context.Context, b *bot.Bot, a *agent.Agent, message, ) - // Send initial placeholder message - msg, err := b.SendMessage(ctx, &bot.SendMessageParams{ - ChatID: update.Message.Chat.ID, - Text: bot.EscapeMarkdown(telegramThinkingMessage), - ParseMode: models.ParseModeMarkdown, - }) - if err != nil { - xlog.Error("Error sending initial message", "error", err) - return + msg := &models.Message{} + jobUUID := types.NewJob().UUID + if !t.streaming { + msg, err = b.SendMessage(ctx, &bot.SendMessageParams{ChatID: update.Message.Chat.ID, Text: bot.EscapeMarkdown(telegramThinkingMessage), ParseMode: models.ParseModeMarkdown}) + if err != nil { + xlog.Error("Error sending initial message", "error", err) + return + } + jobUUID = fmt.Sprintf("%d", msg.ID) + t.placeholderMutex.Lock() + t.placeholders[jobUUID] = msg.ID + t.placeholderMutex.Unlock() } - // Store the UUID->placeholder message mapping - jobUUID := fmt.Sprintf("%d", msg.ID) - - t.placeholderMutex.Lock() - t.placeholders[jobUUID] = msg.ID - t.placeholderMutex.Unlock() - // Add chat ID and conversation_id for tracking and cancel-previous-on-new-message metadata := map[string]interface{}{ - "chatID": update.Message.Chat.ID, + "chatID": update.Message.Chat.ID, types.MetadataKeyConversationID: fmt.Sprintf("telegram:%d", update.Message.Chat.ID), } @@ -767,12 +880,15 @@ func (t *Telegram) handleUpdate(ctx context.Context, b *bot.Bot, a *agent.Agent, metadata["originalMessageType"] = "audio" } - // Create a new job with the conversation history and metadata - job := types.NewJob( - types.WithConversationHistory(currentConv), - types.WithUUID(jobUUID), - types.WithMetadata(metadata), - ) + delivery := t.telegramDelivery(ctx, b, update.Message.Chat.ID, 0, jobUUID, msg.ID) + var streamSession *telegramStreamSession + var job *types.Job + if t.streaming { + job, streamSession = telegramNewJobWithStream(ctx, t.api, update.Message.Chat.ID, true, delivery, currentConv, jobUUID, metadata) + defer streamSession.Close() + } else { + job = types.NewJob(telegramAskOptions(currentConv, jobUUID, metadata, nil)...) + } // Mark this chat as having an active job t.activeJobsMutex.Lock() @@ -798,20 +914,16 @@ func (t *Telegram) handleUpdate(ctx context.Context, b *bot.Bot, a *agent.Agent, t.placeholderMutex.Unlock() }() - res := a.Ask( - types.WithConversationHistory(currentConv), - types.WithUUID(jobUUID), - types.WithMetadata(metadata), - ) + res := telegramExecuteJob(a, job) + if streamSession != nil { + if err := streamSession.Flush(); err != nil { + xlog.Error("Error flushing Telegram stream", "error", err) + } + } if res.Response == "" { xlog.Error("Empty response from agent") - _, err := b.EditMessageText(ctx, &bot.EditMessageTextParams{ - ChatID: update.Message.Chat.ID, - MessageID: msg.ID, - Text: "there was an internal error. try again!", - }) - if err != nil { + if err := delivery.finalPlain(ctx, update.Message.Chat.ID, []string{"there was an internal error. try again!"}); err != nil { xlog.Error("Error updating error message", "error", err) } return @@ -844,12 +956,9 @@ func (t *Telegram) handleUpdate(ctx context.Context, b *bot.Bot, a *agent.Agent, xlog.Error("Error sending audio response", "error", err) } else { xlog.Debug("Audio response sent successfully") - // Remove the thinking placeholder message before returning - _, err := t.bot.DeleteMessage(ctx, &bot.DeleteMessageParams{ - ChatID: update.Message.Chat.ID, - MessageID: msg.ID, - }) - if err != nil { + // Remove any legacy preview before returning. Native drafts are + // superseded by the audio message itself. + if err := delivery.clearPreview(ctx, update.Message.Chat.ID); err != nil { xlog.Error("Error deleting thinking placeholder", "error", err) } // Don't send text response if audio was sent successfully @@ -858,13 +967,7 @@ func (t *Telegram) handleUpdate(ctx context.Context, b *bot.Bot, a *agent.Agent, } } - // Update the message with the final response - formattedResponse := formatResponseWithURLs(res.Response, urls) - - // Split the message if it's too long - messages := xstrings.SplitParagraph(formattedResponse, telegramMaxMessageLength) - - if len(messages) == 0 { + if len(telegramFormatResponse(res.Response, urls, telegramMaxMessageLength)) == 0 { _, err := b.EditMessageText(ctx, &bot.EditMessageTextParams{ ChatID: update.Message.Chat.ID, MessageID: msg.ID, @@ -877,28 +980,14 @@ func (t *Telegram) handleUpdate(ctx context.Context, b *bot.Bot, a *agent.Agent, return } - // Update the first message - _, err = b.EditMessageText(ctx, &bot.EditMessageTextParams{ - ChatID: update.Message.Chat.ID, - MessageID: msg.ID, - Text: messages[0], - ParseMode: models.ParseModeMarkdown, - }) - if err != nil { - xlog.Error("Error updating message", "error", err) - return + var finalErr error + if streamSession != nil { + finalErr = streamSession.Finalize(res.Response, urls) + } else { + finalErr = telegramFinalSession(ctx, t.api, update.Message.Chat.ID, true, delivery).deliverFinal(res.Response, urls) } - - // Send additional chunks as new messages - for i := 1; i < len(messages); i++ { - _, err = b.SendMessage(ctx, &bot.SendMessageParams{ - ChatID: update.Message.Chat.ID, - Text: messages[i], - ParseMode: models.ParseModeMarkdown, - }) - if err != nil { - xlog.Error("Error sending additional message", "error", err) - } + if finalErr != nil { + xlog.Error("Error delivering final Telegram response", "error", finalErr) } } @@ -1032,8 +1121,15 @@ func NewTelegramConnector(config map[string]string) (*Telegram, error) { admins = append(admins, strings.Split(config["admins"], ",")...) } + streaming := true + if value, ok := config["streaming"]; ok { + streaming = value != "false" + } + return &Telegram{ Token: token, + api: newTelegramHTTPAPI(token, http.DefaultClient, ""), + streaming: streaming, admins: admins, placeholders: make(map[string]int), jobStatus: make(map[string]*common.StatusAccumulator), @@ -1077,5 +1173,12 @@ func TelegramConfigMeta() []config.Field { Type: config.FieldTypeCheckbox, HelpText: "Bot will only respond when mentioned in group chats", }, + { + Name: "streaming", + Label: "Streaming", + Type: config.FieldTypeCheckbox, + DefaultValue: true, + HelpText: "Show progressive response previews (native rich drafts in private chats and edited placeholders in groups)", + }, } } diff --git a/services/connectors/telegram_api.go b/services/connectors/telegram_api.go new file mode 100644 index 0000000..f804d70 --- /dev/null +++ b/services/connectors/telegram_api.go @@ -0,0 +1,145 @@ +package connectors + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" +) + +const telegramBotAPIBaseURL = "https://api.telegram.org" + +type telegramAPI interface { + sendRichMessageDraft(context.Context, telegramRichMessageDraft) error + sendRichMessage(context.Context, telegramRichMessage) error +} + +type telegramInputRichMessage struct { + Markdown string `json:"markdown"` +} + +type telegramRichMessageDraft struct { + ChatID int64 `json:"chat_id"` + DraftID int64 `json:"draft_id"` + RichMessage telegramInputRichMessage `json:"rich_message"` +} + +type telegramRichMessage struct { + ChatID int64 `json:"chat_id"` + RichMessage telegramInputRichMessage `json:"rich_message"` + ReplyParameters *telegramReplyParameters `json:"reply_parameters,omitempty"` +} + +type telegramReplyParameters struct { + MessageID int `json:"message_id"` +} + +type telegramAPIError struct { + Method string + ErrorCode int + Description string + RetryAfter int +} + +func (e *telegramAPIError) Error() string { + return fmt.Sprintf("telegram %s: %s", e.Method, e.Description) +} + +type telegramHTTPAPI struct { + token string + client *http.Client + baseURL string +} + +func newTelegramHTTPAPI(token string, client *http.Client, baseURL string) telegramAPI { + if client == nil { + client = http.DefaultClient + } + if baseURL == "" { + baseURL = telegramBotAPIBaseURL + } + return &telegramHTTPAPI{ + token: token, + client: client, + baseURL: strings.TrimRight(baseURL, "/"), + } +} + +func (a *telegramHTTPAPI) sendRichMessageDraft(ctx context.Context, input telegramRichMessageDraft) error { + if input.DraftID == 0 { + return fmt.Errorf("telegram sendRichMessageDraft: draft_id must be nonzero") + } + return a.call(ctx, "sendRichMessageDraft", input) +} + +func (a *telegramHTTPAPI) sendRichMessage(ctx context.Context, input telegramRichMessage) error { + return a.call(ctx, "sendRichMessage", input) +} + +type telegramAPIResponse struct { + OK bool `json:"ok"` + Result json.RawMessage `json:"result"` + ErrorCode int `json:"error_code"` + Description string `json:"description"` + Parameters struct { + RetryAfter int `json:"retry_after"` + } `json:"parameters"` +} + +func (a *telegramHTTPAPI) call(ctx context.Context, method string, input any) error { + body, err := json.Marshal(input) + if err != nil { + return a.safeError(method, "encode request", err) + } + + endpoint := a.baseURL + "/bot" + a.token + "/" + method + req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(body)) + if err != nil { + return a.safeError(method, "create request", err) + } + req.Header.Set("Content-Type", "application/json") + + resp, err := a.client.Do(req) + if err != nil { + return a.safeError(method, "send request", err) + } + defer resp.Body.Close() + + responseBody, err := io.ReadAll(resp.Body) + if err != nil { + return a.safeError(method, "read response", err) + } + + var result telegramAPIResponse + if err := json.Unmarshal(responseBody, &result); err != nil { + return a.safeError(method, "decode response", err) + } + if !result.OK { + description := a.redact(result.Description) + if description == "" { + description = http.StatusText(resp.StatusCode) + } + return &telegramAPIError{ + Method: method, + ErrorCode: result.ErrorCode, + Description: description, + RetryAfter: result.Parameters.RetryAfter, + } + } + + return nil +} + +func (a *telegramHTTPAPI) safeError(method, action string, err error) error { + return fmt.Errorf("telegram %s: %s: %s", method, action, a.redact(err.Error())) +} + +func (a *telegramHTTPAPI) redact(value string) string { + if a.token == "" { + return value + } + return strings.ReplaceAll(value, a.token, "[REDACTED]") +} diff --git a/services/connectors/telegram_api_test.go b/services/connectors/telegram_api_test.go new file mode 100644 index 0000000..9ddd8d2 --- /dev/null +++ b/services/connectors/telegram_api_test.go @@ -0,0 +1,166 @@ +package connectors + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func TestTelegramAPISendsRichMessageDraft(t *testing.T) { + t.Parallel() + + const token = "123456:test-token" + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got, want := r.URL.Path, "/bot"+token+"/sendRichMessageDraft"; got != want { + t.Errorf("path = %q, want %q", got, want) + } + if got, want := r.Header.Get("Content-Type"), "application/json"; got != want { + t.Errorf("Content-Type = %q, want %q", got, want) + } + + var payload map[string]any + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + t.Fatalf("decode request: %v", err) + } + want := map[string]any{ + "chat_id": float64(42), + "draft_id": float64(77), + "rich_message": map[string]any{ + "markdown": "**working**", + }, + } + if !equalJSON(payload, want) { + t.Errorf("payload = %#v, want %#v", payload, want) + } + + _, _ = w.Write([]byte(`{"ok":true,"result":true}`)) + })) + defer server.Close() + + api := newTelegramHTTPAPI(token, server.Client(), server.URL) + err := api.sendRichMessageDraft(context.Background(), telegramRichMessageDraft{ + ChatID: 42, + DraftID: 77, + RichMessage: telegramInputRichMessage{ + Markdown: "**working**", + }, + }) + if err != nil { + t.Fatalf("sendRichMessageDraft() error = %v", err) + } +} + +func TestTelegramAPISendsRichMessageWithoutLinkPreviewField(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got, want := r.URL.Path, "/bottoken/sendRichMessage"; got != want { + t.Errorf("path = %q, want %q", got, want) + } + var payload map[string]any + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + t.Fatalf("decode request: %v", err) + } + if _, exists := payload["link_preview_options"]; exists { + t.Errorf("payload unexpectedly contains link_preview_options: %#v", payload) + } + want := map[string]any{ + "chat_id": float64(-1001), + "reply_parameters": map[string]any{"message_id": float64(55)}, + "rich_message": map[string]any{ + "markdown": "[docs](https://example.com)", + }, + } + if !equalJSON(payload, want) { + t.Errorf("payload = %#v, want %#v", payload, want) + } + _, _ = w.Write([]byte(`{"ok":true,"result":{"message_id":12}}`)) + })) + defer server.Close() + + api := newTelegramHTTPAPI("token", server.Client(), server.URL) + err := api.sendRichMessage(context.Background(), telegramRichMessage{ + ChatID: -1001, + ReplyParameters: &telegramReplyParameters{MessageID: 55}, + RichMessage: telegramInputRichMessage{ + Markdown: "[docs](https://example.com)", + }, + }) + if err != nil { + t.Fatalf("sendRichMessage() error = %v", err) + } +} + +func TestTelegramAPIRichMessageOmitsReplyParametersWhenUnset(t *testing.T) { + t.Parallel() + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var payload map[string]any + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + t.Fatal(err) + } + if _, ok := payload["reply_parameters"]; ok { + t.Fatalf("private payload has reply_parameters: %#v", payload) + } + _, _ = w.Write([]byte(`{"ok":true,"result":true}`)) + })) + defer server.Close() + api := newTelegramHTTPAPI("token", server.Client(), server.URL) + if err := api.sendRichMessage(t.Context(), telegramRichMessage{ChatID: 1, RichMessage: telegramInputRichMessage{Markdown: "ok"}}); err != nil { + t.Fatal(err) + } +} + +func TestTelegramAPIReturnsRetryAfterAndRedactsEchoedToken(t *testing.T) { + t.Parallel() + + const token = "123456:exact-secret-token" + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(`{"ok":false,"error_code":429,"description":"send failed for 123456:exact-secret-token then 123456:exact-secret-token","parameters":{"retry_after":3}}`)) + })) + defer server.Close() + + api := newTelegramHTTPAPI(token, server.Client(), server.URL) + err := api.sendRichMessageDraft(context.Background(), telegramRichMessageDraft{ + ChatID: 1, DraftID: 9, RichMessage: telegramInputRichMessage{Markdown: "text"}, + }) + if err == nil { + t.Fatal("sendRichMessageDraft() error = nil, want Telegram API error") + } + if strings.Contains(err.Error(), token) { + t.Fatalf("error leaked bot token: %q", err) + } + if got := err.Error(); !strings.Contains(got, "sendRichMessageDraft") || !strings.Contains(got, "send failed for [REDACTED] then [REDACTED]") { + t.Errorf("error = %q, want method and redacted description", got) + } + + var apiErr *telegramAPIError + if !errors.As(err, &apiErr) { + t.Fatalf("error type = %T, want *telegramAPIError", err) + } + if apiErr.ErrorCode != 429 || apiErr.RetryAfter != 3 { + t.Errorf("API error = %#v, want error code 429 and retry_after 3", apiErr) + } +} + +func TestTelegramAPIRejectsZeroDraftID(t *testing.T) { + t.Parallel() + + api := newTelegramHTTPAPI("token", http.DefaultClient, "http://unused.invalid") + err := api.sendRichMessageDraft(context.Background(), telegramRichMessageDraft{ + ChatID: 1, RichMessage: telegramInputRichMessage{Markdown: "text"}, + }) + if err == nil || !strings.Contains(err.Error(), "draft_id must be nonzero") { + t.Fatalf("sendRichMessageDraft() error = %v, want nonzero draft ID error", err) + } +} + +func equalJSON(got, want map[string]any) bool { + gotJSON, _ := json.Marshal(got) + wantJSON, _ := json.Marshal(want) + return string(gotJSON) == string(wantJSON) +} diff --git a/services/connectors/telegram_format.go b/services/connectors/telegram_format.go new file mode 100644 index 0000000..40f09ed --- /dev/null +++ b/services/connectors/telegram_format.go @@ -0,0 +1,194 @@ +package connectors + +import ( + "fmt" + "regexp" + "strings" + "unicode/utf8" +) + +var ( + telegramLinkPattern = regexp.MustCompile(`(?m)\[([^\]]+)\]\(([^\n]+)\)`) + telegramBoldPattern = regexp.MustCompile(`\*\*([^*\n]+)\*\*`) + telegramItalicPattern = regexp.MustCompile(`\*([^*\n]+)\*`) + telegramHeadingPattern = regexp.MustCompile(`(?m)^#{1,6}[ \t]+(.+)$`) +) + +func telegramFormatResponse(response string, urls []string, limit int) []string { + return telegramSplitMarkdown(formatResponseWithURLs(response, urls), limit) +} + +// formatResponseWithURLs preserves ordinary Markdown for Telegram's rich API. +func formatResponseWithURLs(response string, urls []string) string { + if len(urls) == 0 { + return response + } + + var result strings.Builder + result.WriteString(response) + result.WriteString("\n\nReferences:\n") + for i, url := range urls { + fmt.Fprintf(&result, "🔗 %d. %s\n", i+1, url) + } + return result.String() +} + +func telegramMarkdownV2(markdown string) string { + parts := strings.Split(markdown, "```") + for i := range parts { + if i%2 == 1 { + parts[i] = escapeTelegramCode(parts[i]) + continue + } + parts[i] = telegramMarkdownV2Text(parts[i]) + } + return strings.Join(parts, "```") +} + +func telegramMarkdownV2Text(text string) string { + tokens := []string{} + protect := func(value string) string { + tokens = append(tokens, value) + return fmt.Sprintf("\x00%d\x00", len(tokens)-1) + } + inlineCodePattern := regexp.MustCompile("`([^`\\n]+)`") + text = inlineCodePattern.ReplaceAllStringFunc(text, func(match string) string { + return protect("`" + escapeTelegramCode(strings.TrimSuffix(strings.TrimPrefix(match, "`"), "`")) + "`") + }) + + text = telegramLinkPattern.ReplaceAllStringFunc(text, func(match string) string { + groups := telegramLinkPattern.FindStringSubmatch(match) + label := escapeTelegramMarkdownV2(groups[1]) + url := strings.NewReplacer(`\`, `\\`, `(`, `\(`, `)`, `\)`).Replace(groups[2]) + return protect("[" + label + "](" + url + ")") + }) + text = telegramHeadingPattern.ReplaceAllStringFunc(text, func(match string) string { + groups := telegramHeadingPattern.FindStringSubmatch(match) + return protect("*" + escapeTelegramMarkdownV2(groups[1]) + "*") + }) + text = telegramBoldPattern.ReplaceAllStringFunc(text, func(match string) string { + groups := telegramBoldPattern.FindStringSubmatch(match) + return protect("*" + escapeTelegramMarkdownV2(groups[1]) + "*") + }) + text = telegramItalicPattern.ReplaceAllStringFunc(text, func(match string) string { + groups := telegramItalicPattern.FindStringSubmatch(match) + return protect("_" + escapeTelegramMarkdownV2(groups[1]) + "_") + }) + text = escapeTelegramMarkdownV2(text) + for i, token := range tokens { + text = strings.ReplaceAll(text, fmt.Sprintf("\x00%d\x00", i), token) + } + return text +} + +func escapeTelegramMarkdownV2(text string) string { + var result strings.Builder + for _, r := range text { + if strings.ContainsRune(`_*[]()~`+"`"+`>#+-=|{}.!\\`, r) { + result.WriteByte('\\') + } + result.WriteRune(r) + } + return result.String() +} + +func escapeTelegramCode(text string) string { + return strings.NewReplacer(`\`, `\\`, "`", "\\`").Replace(text) +} + +func telegramPlainText(markdown string) string { + text := telegramLinkPattern.ReplaceAllString(markdown, "$1 ($2)") + text = telegramHeadingPattern.ReplaceAllString(text, "$1") + text = telegramBoldPattern.ReplaceAllString(text, "$1") + text = telegramItalicPattern.ReplaceAllString(text, "$1") + text = strings.ReplaceAll(text, "```", "") + text = strings.ReplaceAll(text, "`", "") + return text +} + +func telegramSplitMarkdown(text string, limit int) []string { + if text == "" || limit <= 0 { + return []string{} + } + if utf8.RuneCountInString(text) <= limit { + return []string{text} + } + + chunks := []string{} + remaining := text + openFence := "" + for remaining != "" { + prefix := "" + if openFence != "" { + prefix = "```" + openFence + "\n" + } + if utf8.RuneCountInString(prefix) >= limit { + piece, rest := splitTelegramText(remaining, limit) + chunks = append(chunks, piece) + remaining = rest + openFence = "" + continue + } + capacity := limit - utf8.RuneCountInString(prefix) + piece, rest := splitTelegramText(remaining, capacity) + fenceAfter := telegramFenceState(openFence, piece) + if fenceAfter != "" && rest != "" { + close := "```\n" + if !strings.HasSuffix(piece, "\n") { + close = "\n" + close + } + closeLen := utf8.RuneCountInString(close) + if capacity <= closeLen { + piece, rest = splitTelegramText(remaining, limit) + chunks = append(chunks, piece) + remaining = rest + openFence = "" + continue + } + piece, rest = splitTelegramText(remaining, capacity-closeLen) + fenceAfter = telegramFenceState(openFence, piece) + if fenceAfter != "" { + piece += close + } + } + chunks = append(chunks, prefix+piece) + remaining = rest + openFence = fenceAfter + } + return chunks +} + +func splitTelegramText(text string, limit int) (string, string) { + if limit <= 0 { + return "", text + } + runes := []rune(text) + if len(runes) <= limit { + return text, "" + } + cut := limit + for i := limit; i > 0; i-- { + if runes[i-1] == '\n' { + cut = i + break + } + } + return string(runes[:cut]), string(runes[cut:]) +} + +func telegramFenceState(current, text string) string { + state := current + lines := strings.Split(text, "\n") + for _, line := range lines { + trimmed := strings.TrimSpace(line) + if !strings.HasPrefix(trimmed, "```") { + continue + } + if state == "" { + state = strings.TrimSpace(strings.TrimPrefix(trimmed, "```")) + continue + } + state = "" + } + return state +} diff --git a/services/connectors/telegram_format_test.go b/services/connectors/telegram_format_test.go new file mode 100644 index 0000000..82e7045 --- /dev/null +++ b/services/connectors/telegram_format_test.go @@ -0,0 +1,92 @@ +package connectors + +import ( + "strings" + "testing" + "unicode/utf8" +) + +func TestTelegramFormatRawMarkdownAndURLs(t *testing.T) { + t.Parallel() + + chunks := telegramFormatResponse("# Title\n\n**bold**", []string{"https://example.com/a_(b)"}, 200) + want := "# Title\n\n**bold**\n\nReferences:\n🔗 1. https://example.com/a_(b)\n" + if len(chunks) != 1 || chunks[0] != want { + t.Fatalf("telegramFormatResponse() = %#v, want [%q]", chunks, want) + } +} + +func TestTelegramFormatMarkdownV2(t *testing.T) { + t.Parallel() + + input := "# Heading\n\n**bold** and *italic* with [link](https://example.com/a_(b)).\n\n```go\nfmt.Println(`ok`)\n```" + got := telegramMarkdownV2(input) + want := "*Heading*\n\n*bold* and _italic_ with [link](https://example.com/a_\\(b\\))\\.\n\n```go\nfmt.Println(\\`ok\\`)\n```" + if got != want { + t.Fatalf("telegramMarkdownV2() = %q, want %q", got, want) + } +} + +func TestTelegramFormatMarkdownV2PreservesInlineCode(t *testing.T) { + t.Parallel() + got := telegramMarkdownV2("Use `a\\b` and `x` now.") + want := "Use `a\\\\b` and `x` now\\." + if got != want { + t.Fatalf("telegramMarkdownV2() = %q, want %q", got, want) + } +} + +func TestTelegramFormatPlainText(t *testing.T) { + t.Parallel() + + got := telegramPlainText("# Heading\n\n**bold** and [link](https://example.com)") + want := "Heading\n\nbold and link (https://example.com)" + if got != want { + t.Fatalf("telegramPlainText() = %q, want %q", got, want) + } +} + +func TestTelegramSplitUTF8WithoutDataLoss(t *testing.T) { + t.Parallel() + + input := strings.Repeat("🙂", 11) + chunks := telegramSplitMarkdown(input, 4) + if strings.Join(chunks, "") != input { + t.Fatalf("joined chunks differ: %#v", chunks) + } + for _, chunk := range chunks { + if !utf8.ValidString(chunk) || utf8.RuneCountInString(chunk) > 4 { + t.Fatalf("invalid chunk %q", chunk) + } + } +} + +func TestTelegramSplitBalancesFencedCodeBlocks(t *testing.T) { + t.Parallel() + + input := "before\n```go\n" + strings.Repeat("line\n", 8) + "```\nafter" + chunks := telegramSplitMarkdown(input, 28) + if len(chunks) < 2 { + t.Fatalf("got %d chunks, want multiple", len(chunks)) + } + for _, chunk := range chunks { + if strings.Count(chunk, "```")%2 != 0 { + t.Fatalf("unbalanced code fence in %q", chunk) + } + } + joined := strings.Join(chunks, "") + joined = strings.ReplaceAll(joined, "```\n```go\n", "") + if joined != input { + t.Fatalf("split lost content:\n%q\nwant:\n%q", joined, input) + } +} + +func TestTelegramSplitTinyLimitMakesProgress(t *testing.T) { + t.Parallel() + + input := "```go\n🙂🙂\n```" + chunks := telegramSplitMarkdown(input, 5) + if strings.Join(chunks, "") != input { + t.Fatalf("joined chunks differ: %#v", chunks) + } +} diff --git a/services/connectors/telegram_integration_test.go b/services/connectors/telegram_integration_test.go new file mode 100644 index 0000000..81881b7 --- /dev/null +++ b/services/connectors/telegram_integration_test.go @@ -0,0 +1,461 @@ +package connectors + +import ( + "context" + "errors" + "fmt" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/go-telegram/bot" + "github.com/go-telegram/bot/models" + "github.com/mudler/LocalAGI/core/types" + "github.com/mudler/cogito" +) + +type recordingTelegramExecutor struct{ got *types.Job } + +func (r *recordingTelegramExecutor) Execute(j *types.Job) *types.JobResult { + r.got = j + return j.Result +} + +func TestTelegramExecutesTrackedJobIdentity(t *testing.T) { + tracked := types.NewJob() + executor := &recordingTelegramExecutor{} + telegramExecuteJob(executor, tracked) + if executor.got != tracked { + t.Fatal("executor received a different job") + } +} + +func TestTelegramStreamingJobMarksLegacyStatusDeliveryDisabled(t *testing.T) { + metadata := map[string]any{"chatID": int64(42)} + job, session := telegramNewJobWithStream(t.Context(), &telegramStreamAPI{}, 42, true, telegramStreamDelivery{}, nil, "job", metadata) + defer session.Close() + + if telegramUseLegacyStatusDelivery(job) { + t.Fatal("streaming job routed to legacy status delivery") + } + if got, ok := job.Metadata[telegramStreamingMetadataKey].(bool); !ok || !got { + t.Fatalf("streaming metadata = %#v, want true", job.Metadata[telegramStreamingMetadataKey]) + } +} + +func TestTelegramNonStreamingJobKeepsLegacyStatusDelivery(t *testing.T) { + job := types.NewJob(types.WithMetadata(map[string]any{"chatID": int64(42)})) + if !telegramUseLegacyStatusDelivery(job) { + t.Fatal("non-streaming job did not route to legacy status delivery") + } +} + +func TestTelegramStreamingCallbacksDoNotInvokeLegacyStatusPath(t *testing.T) { + job := types.NewJob(types.WithMetadata(map[string]any{telegramStreamingMetadataKey: true})) + calls := 0 + telegramDeliverLegacyStatus(job, func() { calls++ }) + if calls != 0 { + t.Fatalf("legacy status calls = %d, want 0", calls) + } +} + +func TestTelegramNonStreamingCallbacksInvokeLegacyStatusPath(t *testing.T) { + job := types.NewJob() + calls := 0 + telegramDeliverLegacyStatus(job, func() { calls++ }) + if calls != 1 { + t.Fatalf("legacy status calls = %d, want 1", calls) + } +} + +type recordingTelegramBot struct { + sends, edits []models.ParseMode + texts []string + sendParams []*bot.SendMessageParams + deleted []int + nextID int +} + +type contextBlockingTelegramBot struct { + recordingTelegramBot + started chan struct{} + returned chan struct{} +} + +func (b *contextBlockingTelegramBot) EditMessageText(ctx context.Context, _ *bot.EditMessageTextParams) (*models.Message, error) { + close(b.started) + <-ctx.Done() + close(b.returned) + return nil, ctx.Err() +} + +func (b *recordingTelegramBot) SendMessage(_ context.Context, p *bot.SendMessageParams) (*models.Message, error) { + b.sends = append(b.sends, p.ParseMode) + b.sendParams = append(b.sendParams, p) + b.texts = append(b.texts, p.Text) + b.nextID++ + return &models.Message{ID: b.nextID}, nil +} +func (b *recordingTelegramBot) EditMessageText(_ context.Context, p *bot.EditMessageTextParams) (*models.Message, error) { + b.edits = append(b.edits, p.ParseMode) + b.texts = append(b.texts, p.Text) + return &models.Message{ID: p.MessageID}, nil +} +func (b *recordingTelegramBot) DeleteMessage(_ context.Context, p *bot.DeleteMessageParams) (bool, error) { + b.deleted = append(b.deleted, p.MessageID) + return true, nil +} + +func TestTelegramDeliveryClearResetsPlaceholderAndFallbackSendsAfterRich(t *testing.T) { + tg := &Telegram{placeholders: map[string]int{}} + b := &recordingTelegramBot{nextID: 10} + d := tg.telegramDelivery(t.Context(), b, -1, 7, "job", 10) + if err := d.clearPreview(t.Context(), -1); err != nil { + t.Fatal(err) + } + if err := d.finalMarkdown(t.Context(), -1, []string{"later"}); err != nil { + t.Fatal(err) + } + if len(b.deleted) != 1 || len(b.edits) != 0 || len(b.sends) != 1 { + t.Fatalf("deleted/edits/sends = %v/%d/%d", b.deleted, len(b.edits), len(b.sends)) + } + if b.sendParams[0].ReplyParameters == nil || b.sendParams[0].ReplyParameters.MessageID != 7 { + t.Fatalf("fallback reply = %#v", b.sendParams[0].ReplyParameters) + } + if b.sendParams[0].LinkPreviewOptions == nil || b.sendParams[0].LinkPreviewOptions.IsDisabled == nil || !*b.sendParams[0].LinkPreviewOptions.IsDisabled { + t.Fatalf("link previews not disabled: %#v", b.sendParams[0]) + } +} + +func TestTelegramStreamingJobCancellationStillAllowsFinalDelivery(t *testing.T) { + api := &telegramStreamAPI{} + job, session := telegramNewJobWithStream(t.Context(), api, 1, true, telegramStreamDelivery{}, nil, "job", nil) + defer session.Close() + job.Cancel() + if err := session.Finalize("completed answer", nil); err != nil { + t.Fatalf("Finalize after normal job completion: %v", err) + } + _, finals := api.snapshot() + if len(finals) != 1 || finals[0].RichMessage.Markdown != "completed answer" { + t.Fatalf("finals = %#v, want persistent completed answer", finals) + } +} + +func TestTelegramDeliveryPreviewCancellationAbortsInflightLegacyEdit(t *testing.T) { + tg := &Telegram{placeholders: map[string]int{"job": 10}} + b := &contextBlockingTelegramBot{ + started: make(chan struct{}), + returned: make(chan struct{}), + } + delivery := tg.telegramDelivery(context.Background(), b, -1, 7, "job", 10) + job, session := telegramNewJobWithStream(t.Context(), &telegramStreamAPI{}, -1, false, delivery, nil, "job", nil) + defer session.Close() + + session.Accept(cogito.StreamEvent{Type: cogito.StreamEventContent, Content: "working"}) + select { + case <-b.started: + case <-time.After(time.Second): + t.Fatal("legacy preview edit did not start") + } + + job.Cancel() + select { + case <-b.returned: + case <-time.After(time.Second): + t.Fatal("preview cancellation did not unblock the legacy edit") + } +} + +type cancellingTelegramAPI struct { + started chan struct{} + returned chan struct{} + calls atomic.Int32 +} + +func (a *cancellingTelegramAPI) sendRichMessageDraft(ctx context.Context, _ telegramRichMessageDraft) error { + if a.calls.Add(1) == 1 { + close(a.started) + } + <-ctx.Done() + close(a.returned) + return ctx.Err() +} +func (*cancellingTelegramAPI) sendRichMessage(context.Context, telegramRichMessage) error { return nil } + +func TestTelegramTrackedJobCancellationAbortsInflightDraftRequest(t *testing.T) { + api := &cancellingTelegramAPI{started: make(chan struct{}), returned: make(chan struct{})} + job, session := telegramNewJobWithStream(t.Context(), api, 1, true, telegramStreamDelivery{}, nil, "job", nil) + defer session.Close() + select { + case <-api.started: + case <-time.After(time.Second): + t.Fatal("draft request did not start") + } + job.Cancel() + select { + case <-api.returned: + case <-time.After(time.Second): + t.Fatal("in-flight draft request was not cancelled") + } + session.Accept(cogito.StreamEvent{Type: cogito.StreamEventContent, Content: "must not restart previews"}) + time.Sleep(telegramStreamInterval + 50*time.Millisecond) + if got := api.calls.Load(); got != 1 { + t.Fatalf("draft calls after job cancellation = %d, want 1", got) + } + select { + case <-session.done: + t.Fatal("job cancellation stopped final-delivery orchestration") + default: + } + if err := session.Finalize("answer after cancellation", nil); err != nil { + t.Fatalf("Finalize after cancellation: %v", err) + } +} + +type orderedTelegramAPI struct { + mu sync.Mutex + events *[]string + calls int +} + +type contextRecordingTelegramAPI struct { + contexts []context.Context +} + +func (*contextRecordingTelegramAPI) sendRichMessageDraft(context.Context, telegramRichMessageDraft) error { + return nil +} + +func (a *contextRecordingTelegramAPI) sendRichMessage(ctx context.Context, _ telegramRichMessage) error { + a.contexts = append(a.contexts, ctx) + if ctx == nil { + return errors.New("rich message received nil context") + } + if err := ctx.Err(); err != nil { + return fmt.Errorf("rich message received canceled context: %w", err) + } + return errors.New("force fallback delivery") +} + +func (*orderedTelegramAPI) sendRichMessageDraft(context.Context, telegramRichMessageDraft) error { + return nil +} +func (a *orderedTelegramAPI) sendRichMessage(_ context.Context, m telegramRichMessage) error { + a.mu.Lock() + defer a.mu.Unlock() + a.calls++ + *a.events = append(*a.events, "rich:"+m.RichMessage.Markdown) + if a.calls == 2 { + return errors.New("rejected") + } + return nil +} + +type orderedTelegramBot struct { + recordingTelegramBot + events *[]string +} + +func (b *orderedTelegramBot) SendMessage(ctx context.Context, p *bot.SendMessageParams) (*models.Message, error) { + *b.events = append(*b.events, "send:"+p.Text) + return b.recordingTelegramBot.SendMessage(ctx, p) +} +func (b *orderedTelegramBot) DeleteMessage(ctx context.Context, p *bot.DeleteMessageParams) (bool, error) { + *b.events = append(*b.events, "delete") + return b.recordingTelegramBot.DeleteMessage(ctx, p) +} + +func TestTelegramMixedRichFallbackUsesRealDeliveryInOrder(t *testing.T) { + var events []string + tg := &Telegram{placeholders: map[string]int{"job": 10}} + b := &orderedTelegramBot{recordingTelegramBot: recordingTelegramBot{nextID: 10}, events: &events} + delivery := tg.telegramDelivery(t.Context(), b, -1, 7, "job", 10) + api := &orderedTelegramAPI{events: &events} + s := telegramFinalSession(t.Context(), api, -1, false, delivery) + answer := strings.Repeat("a", telegramMaxMessageLength) + "second" + if err := s.deliverFinal(answer, nil); err != nil { + t.Fatal(err) + } + want := []string{"delete", "rich:" + strings.Repeat("a", telegramMaxMessageLength), "rich:second", "send:second"} + if len(events) != len(want) { + t.Fatalf("events = %#v, want %#v", events, want) + } + for i := range want { + if events[i] != want[i] { + t.Fatalf("events = %#v, want %#v", events, want) + } + } + if len(b.edits) != 0 { + t.Fatalf("fallback edited old placeholder: %#v", b.edits) + } +} + +func TestTelegramDeliveryCreatesLazyPlaceholderWithoutIdenticalEditAndUsesMarkdownV2(t *testing.T) { + tg := &Telegram{placeholders: map[string]int{}} + b := &recordingTelegramBot{} + d := tg.telegramDelivery(t.Context(), b, 1, 0, "job", 0) + if err := d.editPreview(t.Context(), 1, "thinking"); err != nil { + t.Fatal(err) + } + if len(b.sends) != 1 || len(b.edits) != 0 { + t.Fatalf("sends/edits = %d/%d", len(b.sends), len(b.edits)) + } + if err := d.finalMarkdown(t.Context(), 1, []string{"final"}); err != nil { + t.Fatal(err) + } + if len(b.edits) != 1 || b.edits[0] != models.ParseModeMarkdown { + t.Fatalf("parse modes = %#v", b.edits) + } +} + +func TestTelegramDeliveryLazyFinalCreationCarriesMarkdownV2ParseMode(t *testing.T) { + tg := &Telegram{placeholders: map[string]int{}} + b := &recordingTelegramBot{} + d := tg.telegramDelivery(t.Context(), b, 1, 0, "job", 0) + if err := d.finalMarkdown(t.Context(), 1, []string{"*final*"}); err != nil { + t.Fatal(err) + } + if len(b.sends) != 1 || b.sends[0] != models.ParseModeMarkdown || len(b.edits) != 0 { + t.Fatalf("send modes/edits = %#v/%d", b.sends, len(b.edits)) + } +} + +func TestTelegramStreamingDefaultsEnabled(t *testing.T) { + tg, err := NewTelegramConnector(map[string]string{"token": "test"}) + if err != nil { + t.Fatal(err) + } + if !tg.streaming { + t.Fatal("streaming = false, want true when omitted") + } +} + +func TestTelegramStreamingCanBeDisabled(t *testing.T) { + tg, err := NewTelegramConnector(map[string]string{"token": "test", "streaming": "false"}) + if err != nil { + t.Fatal(err) + } + if tg.streaming { + t.Fatal("streaming = true, want false") + } +} + +func TestTelegramAskOptionsAttachMatchingSession(t *testing.T) { + api := &telegramStreamAPI{} + session := newTelegramStreamSession(t.Context(), api, 42, true, telegramStreamDelivery{}) + defer session.Close() + + job := types.NewJob(telegramAskOptions(nil, "job", map[string]any{"chatID": int64(42)}, session)...) + if job.StreamCallback == nil { + t.Fatal("request stream callback is nil") + } + job.StreamCallback(cogito.StreamEvent{Type: cogito.StreamEventContent, Content: "hello"}) + if err := session.Flush(); err != nil { + t.Fatal(err) + } + drafts, _ := api.snapshot() + if got := drafts[len(drafts)-1].RichMessage.Markdown; got != "hello" { + t.Fatalf("preview = %q, want hello", got) + } +} + +func TestTelegramAskOptionsWithoutSessionDoesNotStream(t *testing.T) { + job := types.NewJob(telegramAskOptions(nil, "job", nil, nil)...) + if job.StreamCallback != nil { + t.Fatal("request stream callback is set while streaming is disabled") + } +} + +func TestTelegramGroupFinalAttemptsRichBeforeFallback(t *testing.T) { + api := &telegramStreamAPI{} + session := telegramFinalSession(t.Context(), api, -42, false, telegramStreamDelivery{}) + if err := session.deliverFinal("**answer**", nil); err != nil { + t.Fatal(err) + } + _, finals := api.snapshot() + if len(finals) != 1 || finals[0].RichMessage.Markdown != "**answer**" { + t.Fatalf("rich finals = %#v, want raw Markdown attempted once", finals) + } +} + +func TestTelegramFinalDeliveryContinuesAfterPreviewCleanupFailure(t *testing.T) { + api := &telegramStreamAPI{} + session := telegramFinalSession(t.Context(), api, -42, false, telegramStreamDelivery{ + clearPreview: func(context.Context, int64) error { + return errors.New("delete failed") + }, + }) + if err := session.deliverFinal("persistent answer", nil); err != nil { + t.Fatalf("deliverFinal after cleanup failure: %v", err) + } + _, finals := api.snapshot() + if len(finals) != 1 || finals[0].RichMessage.Markdown != "persistent answer" { + t.Fatalf("finals = %#v, want persistent answer", finals) + } +} + +func TestTelegramFinalOnlyDeliveryUsesSuppliedLiveContext(t *testing.T) { + type contextKey struct{} + supplied := context.WithValue(t.Context(), contextKey{}, "final-only") + api := &contextRecordingTelegramAPI{} + var cleanupContexts, markdownContexts, plainContexts []context.Context + session := telegramFinalSession(supplied, api, -42, false, telegramStreamDelivery{ + clearPreview: func(ctx context.Context, _ int64) error { + cleanupContexts = append(cleanupContexts, ctx) + if ctx == nil { + return errors.New("cleanup received nil context") + } + return ctx.Err() + }, + finalMarkdown: func(ctx context.Context, _ int64, _ []string) error { + markdownContexts = append(markdownContexts, ctx) + if ctx == nil { + return errors.New("markdown fallback received nil context") + } + if err := ctx.Err(); err != nil { + return err + } + return errors.New("force plain fallback delivery") + }, + finalPlain: func(ctx context.Context, _ int64, _ []string) error { + plainContexts = append(plainContexts, ctx) + if ctx == nil { + return errors.New("plain fallback received nil context") + } + return ctx.Err() + }, + }) + defer session.cancel() + + if err := session.deliverFinal("final answer", nil); err != nil { + t.Fatalf("deliverFinal() error = %v", err) + } + + operations := []struct { + name string + contexts []context.Context + }{ + {name: "cleanup", contexts: cleanupContexts}, + {name: "rich API", contexts: api.contexts}, + {name: "Markdown fallback", contexts: markdownContexts}, + {name: "plain fallback", contexts: plainContexts}, + } + for _, operation := range operations { + t.Run(operation.name, func(t *testing.T) { + if len(operation.contexts) != 1 { + t.Fatalf("contexts = %#v, want one call", operation.contexts) + } + if operation.contexts[0] != supplied { + t.Fatalf("context = %#v, want supplied context %#v", operation.contexts[0], supplied) + } + if err := operation.contexts[0].Err(); err != nil { + t.Fatalf("context is not live: %v", err) + } + if got := operation.contexts[0].Value(contextKey{}); got != "final-only" { + t.Fatalf("context value = %#v, want final-only", got) + } + }) + } +} diff --git a/services/connectors/telegram_stream.go b/services/connectors/telegram_stream.go new file mode 100644 index 0000000..ac798c1 --- /dev/null +++ b/services/connectors/telegram_stream.go @@ -0,0 +1,414 @@ +package connectors + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "time" + + "github.com/mudler/cogito" +) + +const telegramStreamInterval = 400 * time.Millisecond +const telegramDraftHeartbeatInterval = 25 * time.Second + +type telegramStreamDelivery struct { + editPreview func(context.Context, int64, string) error + finalMarkdown func(context.Context, int64, []string) error + finalPlain func(context.Context, int64, []string) error + clearPreview func(context.Context, int64) error + replyTo int +} + +type telegramStreamCommand struct { + kind uint8 + markdown string + urls []string + done chan error +} + +type telegramStreamSession struct { + ctx context.Context + finalCtx context.Context + cancel context.CancelFunc + api telegramAPI + chatID int64 + private bool + parentCtx context.Context + draftID int64 + delivery telegramStreamDelivery + + mu sync.Mutex + content string + status string + version uint64 + dirty bool + thinkingPending bool + closed bool + wake chan struct{} + command chan telegramStreamCommand + done chan struct{} + heartbeat time.Duration + lastDraftAt time.Time +} + +var telegramDraftSequence atomic.Int64 + +func newTelegramStreamSession(parent context.Context, api telegramAPI, chatID int64, private bool, delivery telegramStreamDelivery) *telegramStreamSession { + return newTelegramStreamSessionWithHeartbeat(parent, api, chatID, private, delivery, telegramDraftHeartbeatInterval) +} + +func newTelegramStreamSessionWithHeartbeat(parent context.Context, api telegramAPI, chatID int64, private bool, delivery telegramStreamDelivery, heartbeat time.Duration) *telegramStreamSession { + return newTelegramStreamSessionWithContexts(parent, parent, api, chatID, private, delivery, heartbeat) +} + +func newTelegramStreamSessionWithContexts(finalParent, previewParent context.Context, api telegramAPI, chatID int64, private bool, delivery telegramStreamDelivery, heartbeat time.Duration) *telegramStreamSession { + finalCtx, cancelFinal := context.WithCancel(finalParent) + ctx, cancelPreview := context.WithCancel(previewParent) + stopPreviewLink := context.AfterFunc(finalCtx, cancelPreview) + draftID := telegramDraftSequence.Add(1) + if draftID == 0 { + draftID = telegramDraftSequence.Add(1) + } + s := &telegramStreamSession{ + ctx: ctx, finalCtx: finalCtx, parentCtx: previewParent, cancel: func() { + stopPreviewLink() + cancelPreview() + cancelFinal() + }, api: api, chatID: chatID, private: private, heartbeat: heartbeat, + draftID: draftID, delivery: delivery, wake: make(chan struct{}, 1), + command: make(chan telegramStreamCommand), done: make(chan struct{}), dirty: true, thinkingPending: true, + } + go s.run() + s.signal() + return s +} + +func (s *telegramStreamSession) Accept(event cogito.StreamEvent) { + if event.Type == cogito.StreamEventDone { + return + } + if event.Type != cogito.StreamEventContent && event.Type != cogito.StreamEventReasoning && event.Type != cogito.StreamEventStatus && event.Type != cogito.StreamEventToolCall && event.Type != cogito.StreamEventToolResult { + return + } + s.mu.Lock() + if s.closed || s.ctx.Err() != nil { + s.mu.Unlock() + return + } + if event.Type == cogito.StreamEventContent && event.Content != "" { + s.content += event.Content + } else if s.content == "" { + s.status = telegramStreamStatus(event) + } + s.version++ + s.dirty = true + s.mu.Unlock() + s.signal() +} + +func (s *telegramStreamSession) Flush() error { + return s.execute(telegramStreamCommand{kind: 1, done: make(chan error, 1)}) +} + +func (s *telegramStreamSession) Finalize(markdown string, urls []string) error { + return s.execute(telegramStreamCommand{kind: 2, markdown: markdown, urls: urls, done: make(chan error, 1)}) +} + +func (s *telegramStreamSession) Close() { + s.mu.Lock() + if s.closed { + s.mu.Unlock() + <-s.done + return + } + s.closed = true + s.mu.Unlock() + s.cancel() + <-s.done +} + +func (s *telegramStreamSession) execute(command telegramStreamCommand) error { + select { + case s.command <- command: + case <-s.done: + return s.ctx.Err() + } + select { + case err := <-command.done: + return err + case <-s.done: + return s.ctx.Err() + } +} + +func (s *telegramStreamSession) signal() { + select { + case s.wake <- struct{}{}: + default: + } +} + +func (s *telegramStreamSession) run() { + defer close(s.done) + var timer *time.Timer + var timerC <-chan time.Time + var nextPreview time.Time + var retryUntil time.Time + var flushWaiters []chan error + previewDone := s.ctx.Done() + previewStopped := false + stopTimer := func() { + if timer != nil && !timer.Stop() { + select { + case <-timer.C: + default: + } + } + timerC = nil + } + schedule := func(at time.Time) { + d := time.Until(at) + if d < 0 { + d = 0 + } + if timer == nil { + timer = time.NewTimer(d) + } else { + stopTimer() + timer.Reset(d) + } + timerC = timer.C + } + defer stopTimer() + finishFlushes := func() { + s.mu.Lock() + dirty := s.dirty + s.mu.Unlock() + if dirty { + return + } + for _, waiter := range flushWaiters { + waiter <- nil + } + flushWaiters = nil + } + schedulePending := func() { + s.mu.Lock() + dirty := s.dirty + lastDraftAt := s.lastDraftAt + s.mu.Unlock() + if !dirty { + finishFlushes() + if s.private && !lastDraftAt.IsZero() && s.heartbeat > 0 { + schedule(lastDraftAt.Add(s.heartbeat)) + } + return + } + when := nextPreview + if retryUntil.After(when) { + when = retryUntil + } + schedule(when) + } + + for { + select { + case <-s.finalCtx.Done(): + return + case <-previewDone: + previewDone = nil + previewStopped = true + stopTimer() + s.mu.Lock() + s.dirty = false + s.mu.Unlock() + for _, waiter := range flushWaiters { + waiter <- s.ctx.Err() + } + flushWaiters = nil + case command := <-s.command: + if command.kind == 1 { + if previewStopped { + command.done <- s.ctx.Err() + continue + } + if time.Now().Before(retryUntil) { + command.done <- errors.New("Telegram preview pending retry") + continue + } + attempted, retry, err := s.deliverPreview() + if attempted { + nextPreview = time.Now().Add(telegramStreamInterval) + } + if retry > 0 { + retryUntil = time.Now().Add(retry) + command.done <- errors.New("Telegram preview pending retry") + continue + } + if err != nil { + command.done <- err + schedulePending() + continue + } + flushWaiters = append(flushWaiters, command.done) + schedulePending() + } else { + stopTimer() + s.mu.Lock() + s.dirty = false + s.mu.Unlock() + command.done <- s.deliverFinal(command.markdown, command.urls) + } + case <-s.wake: + now := time.Now() + when := nextPreview + if retryUntil.After(when) { + when = retryUntil + } + if when.After(now) { + schedule(when) + continue + } + attempted, retry, _ := s.deliverPreview() + if attempted { + nextPreview = time.Now().Add(telegramStreamInterval) + } + if retry > 0 { + retryUntil = time.Now().Add(retry) + } + schedulePending() + case <-timerC: + timerC = nil + s.mu.Lock() + if s.private && !s.dirty && !s.lastDraftAt.IsZero() && s.heartbeat > 0 { + s.dirty = true + } + s.mu.Unlock() + s.signal() + } + } +} + +func (s *telegramStreamSession) previewSnapshot() (string, uint64, bool, bool) { + s.mu.Lock() + defer s.mu.Unlock() + if !s.dirty { + return "", 0, false, false + } + if s.thinkingPending { + return telegramThinkingMessage, s.version, true, true + } + text := s.content + if text == "" { + text = s.status + if text == "" { + text = telegramThinkingMessage + } + } + return telegramPreviewTail(text, telegramMaxMessageLength), s.version, true, false +} + +func telegramPreviewTail(text string, limit int) string { + runes := []rune(text) + if len(runes) <= limit { + return text + } + return string(runes[len(runes)-limit:]) +} + +func telegramStreamStatus(event cogito.StreamEvent) string { + if event.Content != "" { + return event.Content + } + if event.ToolName != "" { + return "Using " + event.ToolName + "…" + } + if event.Type == cogito.StreamEventToolResult { + return "Tool completed…" + } + return "" +} + +func (s *telegramStreamSession) deliverPreview() (bool, time.Duration, error) { + if s.ctx.Err() != nil { + return false, 0, nil + } + text, version, ok, thinking := s.previewSnapshot() + if !ok { + return false, 0, nil + } + var err error + if s.private { + err = s.api.sendRichMessageDraft(s.ctx, telegramRichMessageDraft{ChatID: s.chatID, DraftID: s.draftID, RichMessage: telegramInputRichMessage{Markdown: text}}) + var apiErr *telegramAPIError + if errors.As(err, &apiErr) && apiErr.RetryAfter > 0 { + return true, time.Duration(apiErr.RetryAfter) * time.Second, err + } + if err != nil { + s.private = false + if s.delivery.editPreview != nil { + err = s.delivery.editPreview(s.ctx, s.chatID, text) + } + } + } else if s.delivery.editPreview != nil { + err = s.delivery.editPreview(s.ctx, s.chatID, text) + } + var apiErr *telegramAPIError + if errors.As(err, &apiErr) && apiErr.RetryAfter > 0 { + return true, time.Duration(apiErr.RetryAfter) * time.Second, err + } + if err == nil { + s.mu.Lock() + if s.private { + s.lastDraftAt = time.Now() + } + if thinking { + s.thinkingPending = false + } + if s.version == version && (!thinking || version == 0) { + s.dirty = false + } + s.mu.Unlock() + } + return true, 0, err +} + +func (s *telegramStreamSession) deliverFinal(markdown string, urls []string) error { + formatted := telegramFormatResponse(markdown, urls, telegramMaxMessageLength) + if s.delivery.clearPreview != nil { + _ = s.delivery.clearPreview(s.finalCtx, s.chatID) + } + failedAt := -1 + for i, chunk := range formatted { + final := telegramRichMessage{ChatID: s.chatID, RichMessage: telegramInputRichMessage{Markdown: chunk}} + if !s.private && s.delivery.replyTo != 0 { + final.ReplyParameters = &telegramReplyParameters{MessageID: s.delivery.replyTo} + } + if err := s.api.sendRichMessage(s.finalCtx, final); err == nil { + continue + } + failedAt = i + break + } + if failedAt < 0 { + return nil + } + remaining := formatted[failedAt:] + markdownChunks := make([]string, len(remaining)) + for i, chunk := range remaining { + markdownChunks[i] = telegramMarkdownV2(chunk) + } + if s.delivery.finalMarkdown != nil && s.delivery.finalMarkdown(s.finalCtx, s.chatID, markdownChunks) == nil { + return nil + } + plainChunks := make([]string, len(remaining)) + for i, chunk := range remaining { + plainChunks[i] = telegramPlainText(chunk) + } + if s.delivery.finalPlain != nil { + return s.delivery.finalPlain(s.finalCtx, s.chatID, plainChunks) + } + return nil +} diff --git a/services/connectors/telegram_stream_test.go b/services/connectors/telegram_stream_test.go new file mode 100644 index 0000000..7a6fe63 --- /dev/null +++ b/services/connectors/telegram_stream_test.go @@ -0,0 +1,507 @@ +package connectors + +import ( + "context" + "errors" + "runtime" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/mudler/cogito" +) + +type telegramStreamAPI struct { + mu sync.Mutex + drafts []telegramRichMessageDraft + finals []telegramRichMessage + draftErr func(int) error + finalErr func(int) error + inCall atomic.Int32 + maxCalls atomic.Int32 + block time.Duration +} + +func (a *telegramStreamAPI) sendRichMessageDraft(_ context.Context, draft telegramRichMessageDraft) error { + n := a.inCall.Add(1) + defer a.inCall.Add(-1) + for old := a.maxCalls.Load(); n > old && !a.maxCalls.CompareAndSwap(old, n); old = a.maxCalls.Load() { + } + if a.block > 0 { + time.Sleep(a.block) + } + a.mu.Lock() + a.drafts = append(a.drafts, draft) + i := len(a.drafts) + a.mu.Unlock() + if a.draftErr != nil { + return a.draftErr(i) + } + return nil +} + +func (a *telegramStreamAPI) sendRichMessage(_ context.Context, final telegramRichMessage) error { + n := a.inCall.Add(1) + defer a.inCall.Add(-1) + for old := a.maxCalls.Load(); n > old && !a.maxCalls.CompareAndSwap(old, n); old = a.maxCalls.Load() { + } + a.mu.Lock() + a.finals = append(a.finals, final) + i := len(a.finals) + a.mu.Unlock() + if a.finalErr != nil { + return a.finalErr(i) + } + return nil +} + +func (a *telegramStreamAPI) snapshot() ([]telegramRichMessageDraft, []telegramRichMessage) { + a.mu.Lock() + defer a.mu.Unlock() + return append([]telegramRichMessageDraft(nil), a.drafts...), append([]telegramRichMessage(nil), a.finals...) +} + +func waitTelegramStream(t *testing.T, condition func() bool) { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + if condition() { + return + } + time.Sleep(5 * time.Millisecond) + } + t.Fatal("timed out waiting for stream delivery") +} + +func TestTelegramStreamPrivateUsesStableDraftAndRateLimitsSerializedCalls(t *testing.T) { + api := &telegramStreamAPI{block: 25 * time.Millisecond} + s := newTelegramStreamSession(context.Background(), api, 42, true, telegramStreamDelivery{}) + defer s.Close() + waitTelegramStream(t, func() bool { d, _ := api.snapshot(); return len(d) == 1 }) + + for _, delta := range []string{"one", " two", " three"} { + s.Accept(cogito.StreamEvent{Type: cogito.StreamEventContent, Content: delta}) + } + time.Sleep(200 * time.Millisecond) + if drafts, _ := api.snapshot(); len(drafts) != 1 { + t.Fatalf("calls before interval = %d, want 1", len(drafts)) + } + waitTelegramStream(t, func() bool { d, _ := api.snapshot(); return len(d) == 2 }) + drafts, _ := api.snapshot() + if drafts[0].DraftID == 0 || drafts[1].DraftID != drafts[0].DraftID { + t.Fatalf("draft IDs = %d, %d, want same nonzero ID", drafts[0].DraftID, drafts[1].DraftID) + } + if drafts[0].RichMessage.Markdown != telegramThinkingMessage || drafts[1].RichMessage.Markdown != "one two three" { + t.Fatalf("drafts = %#v", drafts) + } + if api.maxCalls.Load() != 1 { + t.Fatalf("maximum concurrent API calls = %d, want 1", api.maxCalls.Load()) + } +} + +func TestTelegramStreamAlwaysDeliversThinkingBeforeImmediateContent(t *testing.T) { + previous := runtime.GOMAXPROCS(1) + t.Cleanup(func() { runtime.GOMAXPROCS(previous) }) + + api := &telegramStreamAPI{} + s := newTelegramStreamSession(context.Background(), api, 42, true, telegramStreamDelivery{}) + defer s.Close() + s.Accept(cogito.StreamEvent{Type: cogito.StreamEventContent, Content: "immediate"}) + + waitTelegramStream(t, func() bool { drafts, _ := api.snapshot(); return len(drafts) == 2 }) + drafts, _ := api.snapshot() + if got := drafts[0].RichMessage.Markdown; got != telegramThinkingMessage { + t.Fatalf("initial draft = %q, want thinking draft", got) + } + if got := drafts[1].RichMessage.Markdown; got != "immediate" { + t.Fatalf("content draft = %q, want immediate content without another event", got) + } +} + +func TestTelegramStreamRetryAfterRetainsLatestPreview(t *testing.T) { + api := &telegramStreamAPI{draftErr: func(i int) error { + if i == 2 { + return &telegramAPIError{Method: "sendRichMessageDraft", ErrorCode: 429, RetryAfter: 1} + } + return nil + }} + s := newTelegramStreamSession(context.Background(), api, 9, true, telegramStreamDelivery{}) + defer s.Close() + waitTelegramStream(t, func() bool { d, _ := api.snapshot(); return len(d) == 1 }) + s.Accept(cogito.StreamEvent{Type: cogito.StreamEventContent, Content: "first"}) + waitTelegramStream(t, func() bool { d, _ := api.snapshot(); return len(d) == 2 }) + s.Accept(cogito.StreamEvent{Type: cogito.StreamEventContent, Content: " latest"}) + waitTelegramStream(t, func() bool { d, _ := api.snapshot(); return len(d) == 3 }) + drafts, _ := api.snapshot() + if got := drafts[2].RichMessage.Markdown; got != "first latest" { + t.Fatalf("retried preview = %q", got) + } +} + +func TestTelegramStreamEditRetryAfterReschedulesPendingPreviewWithoutNewContent(t *testing.T) { + api := &telegramStreamAPI{} + var calls atomic.Int32 + got := make(chan string, 2) + s := newTelegramStreamSession(context.Background(), api, -9, false, telegramStreamDelivery{editPreview: func(_ context.Context, _ int64, text string) error { + if calls.Add(1) == 1 { + return &telegramAPIError{Method: "editMessageText", ErrorCode: 429, RetryAfter: 1} + } + got <- text + return nil + }}) + defer s.Close() + + select { + case text := <-got: + if text != telegramThinkingMessage { + t.Fatalf("retried preview = %q, want thinking preview", text) + } + case <-time.After(2 * time.Second): + t.Fatal("pending preview was not retried after edit retry_after") + } + if got := calls.Load(); got != 2 { + t.Fatalf("edit calls = %d, want 2", got) + } +} + +func TestTelegramStreamNativeFailureFallsBackOnlyForThatSession(t *testing.T) { + failing := &telegramStreamAPI{draftErr: func(int) error { return errors.New("unsupported") }} + healthy := &telegramStreamAPI{} + var mu sync.Mutex + edits := map[int64][]string{} + hooks := telegramStreamDelivery{editPreview: func(_ context.Context, chatID int64, text string) error { + mu.Lock() + edits[chatID] = append(edits[chatID], text) + mu.Unlock() + return nil + }} + a := newTelegramStreamSession(context.Background(), failing, 1, true, hooks) + b := newTelegramStreamSession(context.Background(), healthy, 2, true, hooks) + defer a.Close() + defer b.Close() + waitTelegramStream(t, func() bool { d, _ := healthy.snapshot(); return len(d) == 1 }) + a.Accept(cogito.StreamEvent{Type: cogito.StreamEventContent, Content: "fallback"}) + b.Accept(cogito.StreamEvent{Type: cogito.StreamEventContent, Content: "native"}) + waitTelegramStream(t, func() bool { + mu.Lock() + defer mu.Unlock() + return len(edits[1]) > 0 && edits[1][len(edits[1])-1] == "fallback" + }) + waitTelegramStream(t, func() bool { d, _ := healthy.snapshot(); return len(d) >= 2 }) + mu.Lock() + otherEdits := len(edits[2]) + mu.Unlock() + if otherEdits != 0 { + t.Fatalf("healthy session used edit fallback %d times", otherEdits) + } +} + +func TestTelegramStreamFlushAndFinalizePromptlyBypassPreviewRetry(t *testing.T) { + api := &telegramStreamAPI{draftErr: func(i int) error { + if i == 2 { + return &telegramAPIError{ErrorCode: 429, RetryAfter: 10} + } + return nil + }} + s := newTelegramStreamSession(context.Background(), api, 5, true, telegramStreamDelivery{}) + defer s.Close() + waitTelegramStream(t, func() bool { d, _ := api.snapshot(); return len(d) == 1 }) + s.Accept(cogito.StreamEvent{Type: cogito.StreamEventContent, Content: "answer"}) + waitTelegramStream(t, func() bool { d, _ := api.snapshot(); return len(d) == 2 }) + start := time.Now() + if err := s.Finalize("answer", nil); err != nil { + t.Fatal(err) + } + if time.Since(start) > 500*time.Millisecond { + t.Fatal("final delivery waited for retry_after") + } + _, finals := api.snapshot() + if len(finals) != 1 || finals[0].RichMessage.Markdown != "answer" { + t.Fatalf("finals = %#v", finals) + } +} + +func TestTelegramStreamFlushDeliversPendingContentBeforeReturning(t *testing.T) { + api := &telegramStreamAPI{} + s := newTelegramStreamSession(context.Background(), api, 5, true, telegramStreamDelivery{}) + defer s.Close() + waitTelegramStream(t, func() bool { d, _ := api.snapshot(); return len(d) == 1 }) + + s.Accept(cogito.StreamEvent{Type: cogito.StreamEventContent, Content: "pending answer"}) + if err := s.Flush(); err != nil { + t.Fatal(err) + } + drafts, _ := api.snapshot() + if len(drafts) != 2 { + t.Fatalf("drafts when Flush returned = %d, want pending content delivered", len(drafts)) + } + if got := drafts[1].RichMessage.Markdown; got != "pending answer" { + t.Fatalf("flushed preview = %q, want pending answer", got) + } +} + +func TestTelegramStreamFlushReportsPendingRetryPromptly(t *testing.T) { + api := &telegramStreamAPI{draftErr: func(i int) error { + if i == 2 { + return &telegramAPIError{ErrorCode: 429, RetryAfter: 1} + } + return nil + }} + s := newTelegramStreamSession(context.Background(), api, 5, true, telegramStreamDelivery{}) + defer s.Close() + waitTelegramStream(t, func() bool { d, _ := api.snapshot(); return len(d) == 1 }) + + s.Accept(cogito.StreamEvent{Type: cogito.StreamEventContent, Content: "pending after retry"}) + waitTelegramStream(t, func() bool { d, _ := api.snapshot(); return len(d) == 2 }) + start := time.Now() + if err := s.Flush(); err == nil { + t.Fatal("Flush error = nil, want pending preview error") + } + if time.Since(start) > 500*time.Millisecond { + t.Fatal("Flush waited for retry_after") + } +} + +func TestTelegramStreamFlushReportsPersistentEditErrorAndFinalizeSucceeds(t *testing.T) { + previewErr := errors.New("persistent edit failure") + var editCalls atomic.Int32 + api := &telegramStreamAPI{} + s := newTelegramStreamSession(context.Background(), api, -5, false, telegramStreamDelivery{ + editPreview: func(context.Context, int64, string) error { + editCalls.Add(1) + return previewErr + }, + }) + defer s.Close() + waitTelegramStream(t, func() bool { return editCalls.Load() == 1 }) + + s.Accept(cogito.StreamEvent{Type: cogito.StreamEventContent, Content: "answer"}) + started := time.Now() + if err := s.Flush(); !errors.Is(err, previewErr) { + t.Fatalf("Flush error = %v, want %v", err, previewErr) + } + if elapsed := time.Since(started); elapsed > 500*time.Millisecond { + t.Fatalf("Flush took %v, want less than 500ms", elapsed) + } + if got := editCalls.Load(); got != 2 { + t.Fatalf("edit calls after Flush = %d, want 2", got) + } + s.mu.Lock() + dirty := s.dirty + s.mu.Unlock() + if !dirty { + t.Fatal("Flush cleared preview state after delivery error") + } + + if err := s.Finalize("answer", nil); err != nil { + t.Fatalf("Finalize after Flush error: %v", err) + } + _, finals := api.snapshot() + if len(finals) != 1 || finals[0].RichMessage.Markdown != "answer" { + t.Fatalf("finals = %#v", finals) + } + time.Sleep(50 * time.Millisecond) + if got := editCalls.Load(); got != 2 { + t.Fatalf("edit calls after Finalize = %d, want no busy retry loop", got) + } +} + +func TestTelegramStreamPreviewUsesUTF8SafeTail(t *testing.T) { + api := &telegramStreamAPI{} + s := newTelegramStreamSession(context.Background(), api, 5, true, telegramStreamDelivery{}) + defer s.Close() + waitTelegramStream(t, func() bool { d, _ := api.snapshot(); return len(d) == 1 }) + full := strings.Repeat("🙂", telegramMaxMessageLength+10) + s.Accept(cogito.StreamEvent{Type: cogito.StreamEventContent, Content: full}) + if err := s.Flush(); err != nil { + t.Fatal(err) + } + drafts, _ := api.snapshot() + got := drafts[len(drafts)-1].RichMessage.Markdown + if len([]rune(got)) != telegramMaxMessageLength || got != strings.Repeat("🙂", telegramMaxMessageLength) { + t.Fatalf("preview rune length = %d, want tail of %d", len([]rune(got)), telegramMaxMessageLength) + } + if s.content != full { + t.Fatal("preview truncation discarded final content") + } +} + +func TestTelegramStreamPrivateShowsPublishedStatusBeforeContent(t *testing.T) { + api := &telegramStreamAPI{} + s := newTelegramStreamSession(context.Background(), api, 5, true, telegramStreamDelivery{}) + defer s.Close() + waitTelegramStream(t, func() bool { d, _ := api.snapshot(); return len(d) == 1 }) + s.Accept(cogito.StreamEvent{Type: cogito.StreamEventReasoning, Content: "checking sources"}) + if err := s.Flush(); err != nil { + t.Fatal(err) + } + drafts, _ := api.snapshot() + if got := drafts[len(drafts)-1].RichMessage.Markdown; got != "checking sources" { + t.Fatalf("status = %q", got) + } + s.Accept(cogito.StreamEvent{Type: cogito.StreamEventContent, Content: "answer"}) + if err := s.Flush(); err != nil { + t.Fatal(err) + } + drafts, _ = api.snapshot() + if got := drafts[len(drafts)-1].RichMessage.Markdown; got != "answer" { + t.Fatalf("answer = %q", got) + } +} + +func TestTelegramStreamDoneEventIsNonblocking(t *testing.T) { + api := &telegramStreamAPI{block: 500 * time.Millisecond} + s := newTelegramStreamSession(context.Background(), api, 5, true, telegramStreamDelivery{}) + defer s.Close() + waitTelegramStream(t, func() bool { return api.inCall.Load() == 1 }) + + returned := make(chan struct{}) + go func() { + s.Accept(cogito.StreamEvent{Type: cogito.StreamEventDone}) + close(returned) + }() + select { + case <-returned: + case <-time.After(100 * time.Millisecond): + t.Fatal("Done event blocked on stream delivery") + } +} + +func TestTelegramStreamCancelAndCloseStopPendingThrottledDelivery(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + api := &telegramStreamAPI{} + s := newTelegramStreamSession(ctx, api, 7, true, telegramStreamDelivery{}) + waitTelegramStream(t, func() bool { d, _ := api.snapshot(); return len(d) == 1 }) + s.Accept(cogito.StreamEvent{Type: cogito.StreamEventContent, Content: "pending"}) + time.Sleep(100 * time.Millisecond) + if drafts, _ := api.snapshot(); len(drafts) != 1 { + t.Fatalf("drafts before throttle elapsed = %d, want 1", len(drafts)) + } + + cancel() + s.Close() + drafts, _ := api.snapshot() + time.Sleep(500 * time.Millisecond) + after, _ := api.snapshot() + if len(after) != len(drafts) { + t.Fatalf("calls after close = %d, before = %d", len(after), len(drafts)) + } + select { + case <-s.done: + default: + t.Fatal("worker did not terminate") + } +} + +func TestTelegramStreamGroupEditsPlaceholder(t *testing.T) { + api := &telegramStreamAPI{} + got := make(chan string, 2) + s := newTelegramStreamSession(context.Background(), api, -10, false, telegramStreamDelivery{editPreview: func(_ context.Context, _ int64, text string) error { got <- text; return nil }}) + defer s.Close() + select { + case text := <-got: + if text != telegramThinkingMessage { + t.Fatalf("initial edit = %q", text) + } + case <-time.After(time.Second): + t.Fatal("missing initial edit") + } + s.Accept(cogito.StreamEvent{Type: cogito.StreamEventContent, Content: "group answer"}) + select { + case text := <-got: + if text != "group answer" { + t.Fatalf("content edit = %q", text) + } + case <-time.After(time.Second): + t.Fatal("missing content edit") + } +} + +func TestTelegramStreamLongPrivateFinalUsesRichMarkdownForEveryChunkInOrder(t *testing.T) { + api := &telegramStreamAPI{} + s := newTelegramStreamSession(context.Background(), api, 42, true, telegramStreamDelivery{}) + defer s.Close() + + markdown := strings.Repeat("a", telegramMaxMessageLength) + strings.Repeat("b", 17) + if err := s.Finalize(markdown, nil); err != nil { + t.Fatal(err) + } + _, finals := api.snapshot() + if len(finals) != 2 { + t.Fatalf("rich final calls = %d, want 2", len(finals)) + } + if got := finals[0].RichMessage.Markdown; got != strings.Repeat("a", telegramMaxMessageLength) { + t.Fatalf("first rich chunk length/content = %d/%q", len(got), got[:min(len(got), 20)]) + } + if got := finals[1].RichMessage.Markdown; got != strings.Repeat("b", 17) { + t.Fatalf("second rich chunk = %q", got) + } +} + +func TestTelegramStreamLongPrivateFinalFallsBackWithoutLosingOrReorderingChunks(t *testing.T) { + api := &telegramStreamAPI{finalErr: func(i int) error { + if i == 1 { + return errors.New("rich markdown rejected") + } + return nil + }} + var markdownChunks, plainChunks []string + s := newTelegramStreamSession(context.Background(), api, 42, true, telegramStreamDelivery{ + finalMarkdown: func(_ context.Context, _ int64, chunks []string) error { + markdownChunks = append([]string(nil), chunks...) + return errors.New("MarkdownV2 rejected") + }, + finalPlain: func(_ context.Context, _ int64, chunks []string) error { + plainChunks = append([]string(nil), chunks...) + return nil + }, + }) + defer s.Close() + + markdown := strings.Repeat("a", telegramMaxMessageLength) + strings.Repeat("b", 17) + if err := s.Finalize(markdown, nil); err != nil { + t.Fatal(err) + } + _, finals := api.snapshot() + if len(finals) != 1 { + t.Fatalf("rich final calls = %d, want rich delivery to stop at first failure", len(finals)) + } + wantMarkdown := []string{strings.Repeat("a", telegramMaxMessageLength), strings.Repeat("b", 17)} + if len(markdownChunks) != 2 || markdownChunks[0] != wantMarkdown[0] || markdownChunks[1] != wantMarkdown[1] { + t.Fatalf("MarkdownV2 fallback chunks lost or reordered: lengths %d, %d", len(markdownChunks), len(plainChunks)) + } + wantPlain := []string{strings.Repeat("a", telegramMaxMessageLength), strings.Repeat("b", 17)} + if len(plainChunks) != 2 || plainChunks[0] != wantPlain[0] || plainChunks[1] != wantPlain[1] { + t.Fatalf("plain fallback chunks lost or reordered: %#v", plainChunks) + } +} + +func TestTelegramStreamNativeDraftHeartbeatAndClose(t *testing.T) { + api := &telegramStreamAPI{} + s := newTelegramStreamSessionWithHeartbeat(context.Background(), api, 42, true, telegramStreamDelivery{}, 20*time.Millisecond) + waitTelegramStream(t, func() bool { d, _ := api.snapshot(); return len(d) >= 2 }) + drafts, _ := api.snapshot() + if drafts[0].DraftID != drafts[1].DraftID || drafts[0].RichMessage.Markdown != drafts[1].RichMessage.Markdown { + t.Fatalf("heartbeats changed draft: %#v", drafts[:2]) + } + s.Close() + n := len(drafts) + time.Sleep(50 * time.Millisecond) + after, _ := api.snapshot() + if len(after) != n { + t.Fatalf("heartbeat continued after close: %d -> %d", n, len(after)) + } +} + +func TestTelegramStreamGroupDoesNotHeartbeat(t *testing.T) { + var calls atomic.Int32 + s := newTelegramStreamSessionWithHeartbeat(context.Background(), &telegramStreamAPI{}, -1, false, telegramStreamDelivery{editPreview: func(context.Context, int64, string) error { calls.Add(1); return nil }}, 20*time.Millisecond) + defer s.Close() + waitTelegramStream(t, func() bool { return calls.Load() == 1 }) + time.Sleep(60 * time.Millisecond) + if calls.Load() != 1 { + t.Fatalf("group heartbeat calls = %d", calls.Load()) + } +} From 5435cbd62bd81ed41c231ee1ca891158b6c90097 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Sat, 22 Aug 2026 09:06:24 +0200 Subject: [PATCH 6/6] Change runner from arc-runner-localagent to ubuntu-latest --- .github/workflows/image.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/image.yml b/.github/workflows/image.yml index 7f825b6..35c4d6e 100644 --- a/.github/workflows/image.yml +++ b/.github/workflows/image.yml @@ -11,8 +11,8 @@ concurrency: cancel-in-progress: true jobs: containerImages: - #runs-on: ubuntu-latest - runs-on: arc-runner-localagent + runs-on: ubuntu-latest + #runs-on: arc-runner-localagent steps: - name: Checkout uses: actions/checkout@v6