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 diff --git a/README.md b/README.md index 4ee5fe4..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. @@ -239,7 +241,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 | @@ -825,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 @@ -1049,7 +1058,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/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/go.mod b/go.mod index 14bcf99..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.1-0.20260507074622-a7724fef6f81 + 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 52fe00f..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.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.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= 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()) + } +} 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) } } 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()],