From edcf8e486a731cc9d7b260510257d7a84a9ffc56 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Wed, 2 Sep 2026 23:25:07 +1000 Subject: [PATCH] fix: Telegram interruption can hang forever while closing an uncooperative stream --- internal/serve/telegram.go | 139 +++++++++++++----- internal/serve/telegram_test.go | 244 ++++++++++++++++++++++++++++++++ 2 files changed, 351 insertions(+), 32 deletions(-) diff --git a/internal/serve/telegram.go b/internal/serve/telegram.go index ea4b14fb6..528d400cb 100644 --- a/internal/serve/telegram.go +++ b/internal/serve/telegram.go @@ -1843,20 +1843,45 @@ func (m *telegramSessionMgr) streamReplyWithAdmission(ctx context.Context, bot b var stream llm.Stream var runnerDone <-chan struct{} - waitForRunnerDone := func() { - if runnerDone == nil { + streamCloseDone := make(chan struct{}) + var streamCloseOnce sync.Once + startStreamClose := func() { + streamCloseOnce.Do(func() { + if stream == nil { + close(streamCloseDone) + return + } + go func() { + _ = stream.Close() + close(streamCloseDone) + }() + }) + } + cleanupTimeout := func() time.Duration { + timeout := telegramRunnerCleanupTimeout + if timeout <= 0 { + timeout = runpkg.DefaultRunnerCleanupTimeout + } + return timeout + } + cleanupDetached := false + markCleanupDetached := func() { + if cleanupDetached { return } - cleanupTimeout := telegramRunnerCleanupTimeout - if cleanupTimeout <= 0 { - cleanupTimeout = runpkg.DefaultRunnerCleanupTimeout + cleanupDetached = true + sess.runtimeStale.Store(true) + log.Printf("[telegram] stream for chat %d did not stop after cancellation; detaching and resetting runtime before reuse", chatID) + } + waitForRunnerDone := func() { + if runnerDone == nil || cleanupDetached { + return } - if !runpkg.WaitForRunnerDone(context.Background(), runnerDone, cleanupTimeout) { - sess.runtimeStale.Store(true) - log.Printf("[telegram] runner for chat %d did not stop within %s after stream cancellation; detaching and resetting runtime before reuse", chatID, cleanupTimeout) + timeout := cleanupTimeout() + if !runpkg.WaitForRunnerDone(context.Background(), runnerDone, timeout) { + markCleanupDetached() } } - defer waitForRunnerDone() if m.settings.Runner != nil { pipe := runpkg.NewEventPipe(streamCtx, ui.DefaultStreamBufferSize) stream = pipe @@ -1934,8 +1959,27 @@ func (m *telegramSessionMgr) streamReplyWithAdmission(ctx context.Context, bot b } return fmt.Errorf("stream: %w", err) } - defer stream.Close() + sess.cancelMu.Lock() + if sess.streamToken == streamToken { + sess.runnerDone = streamCloseDone + } + sess.cancelMu.Unlock() } + defer func() { + streamCancel() + startStreamClose() + waitForRunnerDone() + if cleanupDetached { + return + } + closeCtx, cancel := context.WithTimeout(context.Background(), cleanupTimeout()) + defer cancel() + select { + case <-streamCloseDone: + case <-closeCtx.Done(): + markCleanupDetached() + } + }() // Send placeholder message to obtain a message ID for live editing. placeholder, err := bot.Send(tgbotapi.NewMessage(chatID, "⏳")) @@ -1997,7 +2041,22 @@ func (m *telegramSessionMgr) streamReplyWithAdmission(ctx context.Context, bot b }() // Goroutine: consume stream events. + streamConsumerDone := make(chan struct{}) + if runnerDone == nil { + streamCleanupDone := make(chan struct{}) + go func() { + <-streamConsumerDone + <-streamCloseDone + close(streamCleanupDone) + }() + sess.cancelMu.Lock() + if sess.streamToken == streamToken { + sess.runnerDone = streamCleanupDone + } + sess.cancelMu.Unlock() + } go func() { + defer close(streamConsumerDone) for { ev, recvErr := stream.Recv() if recvErr == io.EOF { @@ -2328,23 +2387,35 @@ func (m *telegramSessionMgr) streamReplyWithAdmission(ctx context.Context, bot b userInterrupted := false streamDoneDrained := false stopStreamAndWait := func(waitCtx context.Context) bool { - if m.settings.Runner != nil { - streamCancel() - } else if stream != nil { - stream.Close() - } else { - streamCancel() - } - if streamDoneDrained { - return true + streamCancel() + startStreamClose() + + var pendingStreamDone <-chan error + if !streamDoneDrained { + pendingStreamDone = streamDone } - select { - case <-streamDone: - streamDoneDrained = true - return true - case <-waitCtx.Done(): - return false + pendingCloseDone := (<-chan struct{})(streamCloseDone) + pendingRunnerDone := runnerDone + for pendingStreamDone != nil || pendingCloseDone != nil || pendingRunnerDone != nil { + select { + case <-pendingStreamDone: + streamDoneDrained = true + pendingStreamDone = nil + case <-pendingCloseDone: + pendingCloseDone = nil + case <-pendingRunnerDone: + pendingRunnerDone = nil + case <-waitCtx.Done(): + markCleanupDetached() + return false + } } + return true + } + stopStreamWithCleanupTimeout := func() bool { + waitCtx, cancel := context.WithTimeout(context.Background(), cleanupTimeout()) + defer cancel() + return stopStreamAndWait(waitCtx) } loop: for { @@ -2441,8 +2512,9 @@ loop: if userInterrupted { // Stop and drain anything already in flight so history and persistence - // snapshots include the final callback-produced messages. - stopStreamAndWait(context.Background()) + // snapshots include the final callback-produced messages. If cleanup does + // not finish promptly, detach it and avoid racing the stale producer. + drained := stopStreamWithCleanupTimeout() textMu.Lock() partial := textBuf.String() @@ -2466,8 +2538,11 @@ loop: } sendEdit(currentMsgID, display, true) - // Preserve partial history so conversation context isn't lost. - salvagePartialHistory(streamCtx, "AddMessage(assistant_interrupt_fallback)") + // Preserve partial history so conversation context isn't lost, but only + // after every producer has stopped mutating the callback snapshot. + if drained { + salvagePartialHistory(streamCtx, "AddMessage(assistant_interrupt_fallback)") + } if m.store != nil && sess.meta != nil { m.runStoreOpWithTimeout(sess.meta.ID, "UpdateStatus(interrupted)", func(storeCtx context.Context) error { @@ -2478,10 +2553,10 @@ loop: } if streamErr != nil { - if !streamDoneDrained { - stopStreamAndWait(context.Background()) + drained := stopStreamWithCleanupTimeout() + if drained { + salvagePartialHistory(streamCtx, "AddMessage(assistant_error_fallback)") } - salvagePartialHistory(streamCtx, "AddMessage(assistant_error_fallback)") if strings.Contains(streamErr.Error(), "stream timed out") { _, _ = bot.Send(tgbotapi.NewMessage(chatID, "⌛ Response timed out — please try again.")) } diff --git a/internal/serve/telegram_test.go b/internal/serve/telegram_test.go index b67d0bbdd..cc17a92ba 100644 --- a/internal/serve/telegram_test.go +++ b/internal/serve/telegram_test.go @@ -20,6 +20,7 @@ import ( tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api/v5" "github.com/samsaffron/term-llm/internal/config" "github.com/samsaffron/term-llm/internal/llm" + runpkg "github.com/samsaffron/term-llm/internal/run" "github.com/samsaffron/term-llm/internal/session" "github.com/samsaffron/term-llm/internal/testutil" ) @@ -297,6 +298,76 @@ func (s *blockingTextStream) Close() error { return nil } +type uncooperativeCloseProvider struct { + started chan struct{} + release chan struct{} + stopped chan struct{} + runs atomic.Int32 +} + +func (p *uncooperativeCloseProvider) Name() string { return "uncooperative-close" } + +func (p *uncooperativeCloseProvider) Credential() string { return "mock" } + +func (p *uncooperativeCloseProvider) Capabilities() llm.Capabilities { return llm.Capabilities{} } + +func (p *uncooperativeCloseProvider) Stream(ctx context.Context, req llm.Request) (llm.Stream, error) { + if p.runs.Add(1) == 1 { + return &uncooperativeCloseStream{ + ctx: ctx, + started: p.started, + release: p.release, + stopped: p.stopped, + }, nil + } + return &oneShotTextStream{text: "replacement answer"}, nil +} + +type uncooperativeCloseStream struct { + ctx context.Context + started chan struct{} + release chan struct{} + stopped chan struct{} + startOnce sync.Once + stopOnce sync.Once + sent bool +} + +func (s *uncooperativeCloseStream) Recv() (llm.Event, error) { + if !s.sent { + s.sent = true + s.startOnce.Do(func() { close(s.started) }) + return llm.Event{Type: llm.EventTextDelta, Text: "partial answer"}, nil + } + <-s.ctx.Done() + return llm.Event{}, s.ctx.Err() +} + +func (s *uncooperativeCloseStream) Close() error { + <-s.release + s.stopOnce.Do(func() { close(s.stopped) }) + return nil +} + +type uncooperativeSequenceRunner struct { + started chan struct{} + release chan struct{} + stopped chan struct{} + runs atomic.Int32 +} + +func (r *uncooperativeSequenceRunner) Run(ctx context.Context, req runpkg.Request, sink runpkg.EventSink) (runpkg.Result, error) { + if r.runs.Add(1) == 1 { + sink.Event(llm.Event{Type: llm.EventTextDelta, Text: "partial answer"}) + close(r.started) + <-r.release + close(r.stopped) + return runpkg.Result{}, nil + } + sink.Event(llm.Event{Type: llm.EventTextDelta, Text: "replacement answer"}) + return runpkg.Result{}, nil +} + type errorAfterTextProvider struct { text string err error @@ -2685,6 +2756,179 @@ func TestStreamReply_WatchdogTimeoutIsNotTreatedAsUserInterrupt(t *testing.T) { } } +func TestStreamReply_UncooperativeCleanupIsBounded(t *testing.T) { + previousCleanupTimeout := telegramRunnerCleanupTimeout + telegramRunnerCleanupTimeout = 25 * time.Millisecond + defer func() { telegramRunnerCleanupTimeout = previousCleanupTimeout }() + + tests := []struct { + name string + watchdog bool + newMgr func(started, release, stopped chan struct{}) *telegramSessionMgr + }{ + { + name: "stream Close ignores cancellation", + newMgr: func(started, release, stopped chan struct{}) *telegramSessionMgr { + provider := &uncooperativeCloseProvider{started: started, release: release, stopped: stopped} + return &telegramSessionMgr{ + sessions: make(map[int64]*telegramSession), + allowedUserIDs: map[int64]struct{}{7: {}}, + idleTimeout: time.Hour, + tickerInterval: 5 * time.Millisecond, + settings: Settings{ + MaxTurns: 5, + NewSession: func(context.Context) (*SessionRuntime, error) { + return &SessionRuntime{ + Engine: llm.NewEngine(provider, llm.NewToolRegistry()), + ProviderName: "mock", + ModelName: "test", + }, nil + }, + }, + } + }, + }, + { + name: "Runner.Run ignores watchdog cancellation", + watchdog: true, + newMgr: func(started, release, stopped chan struct{}) *telegramSessionMgr { + runner := &uncooperativeSequenceRunner{started: started, release: release, stopped: stopped} + return &telegramSessionMgr{ + sessions: make(map[int64]*telegramSession), + allowedUserIDs: map[int64]struct{}{7: {}}, + idleTimeout: time.Hour, + tickerInterval: 5 * time.Millisecond, + streamEventTimeout: 15 * time.Millisecond, + settings: Settings{ + Runner: runner, + MaxTurns: 5, + NewSession: func(context.Context) (*SessionRuntime, error) { + provider := llm.NewMockProvider("mock") + return &SessionRuntime{ + Engine: llm.NewEngine(provider, llm.NewToolRegistry()), + Provider: provider, + ProviderName: "mock", + ModelName: "test", + }, nil + }, + }, + } + }, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + started := make(chan struct{}) + release := make(chan struct{}) + stopped := make(chan struct{}) + released := false + defer func() { + if !released { + close(release) + } + }() + + mgr := tc.newMgr(started, release, stopped) + sess, err := mgr.getOrCreate(context.Background(), 42) + if err != nil { + t.Fatalf("getOrCreate failed: %v", err) + } + bot := &fakeBotSender{} + replyResult := make(chan error, 1) + go func() { + replyResult <- mgr.streamReply(context.Background(), bot, sess, 42, llm.UserText("first request")) + }() + + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("first response did not start") + } + + sess.cancelMu.Lock() + cancelStream := sess.streamCancel + replyDone := sess.replyDone + sess.cancelMu.Unlock() + if cancelStream == nil || replyDone == nil { + t.Fatal("active response did not publish cancellation state") + } + if !tc.watchdog { + cancelStream() + } + + start := time.Now() + select { + case err := <-replyResult: + if tc.watchdog { + if err == nil || !strings.Contains(err.Error(), "stream timed out") { + t.Fatalf("streamReply error = %v, want watchdog timeout", err) + } + } else if err != nil { + t.Fatalf("streamReply returned error after interrupt: %v", err) + } + case <-time.After(time.Second): + t.Fatal("streamReply remained blocked on uncooperative cleanup") + } + if elapsed := time.Since(start); elapsed > 500*time.Millisecond { + t.Fatalf("streamReply cleanup took %s, want bounded cleanup", elapsed) + } + if !sess.runtimeStale.Load() { + t.Fatal("uncooperative cleanup did not mark runtime stale") + } + select { + case <-replyDone: + default: + t.Fatal("replyDone was not closed after detached cleanup") + } + + lockAcquired := make(chan struct{}) + go func() { + sess.mu.Lock() + sess.mu.Unlock() + close(lockAcquired) + }() + select { + case <-lockAcquired: + case <-time.After(time.Second): + t.Fatal("session mutex remained locked after detached cleanup") + } + + nextDone := make(chan struct{}) + go func() { + mgr.handleMessage(context.Background(), bot, &tgbotapi.Message{ + From: &tgbotapi.User{ID: 7, UserName: "sam"}, + Chat: &tgbotapi.Chat{ID: 42}, + Text: "next request", + }) + close(nextDone) + }() + select { + case <-nextDone: + case <-time.After(time.Second): + t.Fatal("replacement session did not process the next message") + } + mgr.mu.Lock() + replacement := mgr.sessions[42] + mgr.mu.Unlock() + if replacement == sess { + t.Fatal("stale runtime was reused for the next message") + } + if got := bot.lastText(); got != "replacement answer" { + t.Fatalf("replacement response = %q, want replacement answer", got) + } + + close(release) + released = true + select { + case <-stopped: + case <-time.After(time.Second): + t.Fatal("detached cleanup did not finish after test release") + } + }) + } +} + func TestStreamReply_InjectsCarryoverSystemNoteOnce(t *testing.T) { h := testutil.NewEngineHarness() h.Provider.AddTextResponse("first")