From 5139d152609b4fcb9921f8f8f95318dcbc1ad938 Mon Sep 17 00:00:00 2001
From: Rick Staa <rick.staa@outlook.com>
Date: Fri, 8 May 2026 15:09:33 +0200
Subject: [PATCH] byoc: drain data subscriber on /stream/stop before tearing
 down
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Fixes #3924. The gateway's data SSE proxy was tearing down its trickle
subscriber as soon as /stream/stop arrived, which races the runner's
on_stream_stop hook. Any final emit_data the runner publishes during
its stop hook (e.g. live_transcribe's last partial-window flush) is
written to the trickle channel after the gateway has already cancelled
its subscriber's context — so the segment is never read, never relayed
to the SSE caller.

Decouple the data subscriber's context from /stream/stop. After the
stream context is cancelled, give the subscriber up to drainTimeout
(10s) to keep reading. The trickle channel closes cleanly when the
runner SDK tears down its data publisher (DELETE → next Read returns
trickle.EOS), at which point the loop exits via the existing EOS path.
The timeout caps the drain so a hung runner can't keep the goroutine
alive forever.

Bug reproduced 3/3 against examples/runner/live_transcribe in the
livepeer-python-gateway SDK: the on_stream_stop final-flush transcript
consistently drops. Fix not yet validated end-to-end pending a test
image build.

Refs livepeer/livepeer-python-gateway#12 (companion SDK-side fix).
---
 byoc/trickle.go | 29 ++++++++++++++++++++++++++---
 1 file changed, 26 insertions(+), 3 deletions(-)

diff --git a/byoc/trickle.go b/byoc/trickle.go
index 3b90cf0825..38a3ea0ed6 100644
--- a/byoc/trickle.go
+++ b/byoc/trickle.go
@@ -794,12 +794,31 @@ func (bsg *BYOCGatewayServer) startDataSubscribe(ctx context.Context, url *url.U
 		return
 	}
 
+	// The data subscriber outlives /stream/stop's ctx cancellation by
+	// drainTimeout so we can read the runner's final emit_data segment(s)
+	// published from on_stream_stop before tearing down the SSE proxy.
+	// The trickle channel closes cleanly when the runner SDK tears down
+	// its data publisher (DELETE → trickle.EOS on the next Read), at which
+	// point we exit normally below. The timeout caps the drain so a hung
+	// runner can't keep this goroutine alive forever.
+	const drainTimeout = 10 * time.Second
+	drainCtx, drainCancel := context.WithCancel(context.Background())
+	go func() {
+		<-ctx.Done()
+		select {
+		case <-time.After(drainTimeout):
+			drainCancel()
+		case <-drainCtx.Done():
+		}
+	}()
+
 	// subscribe to the outputs
 	subscriber, err := trickle.NewTrickleSubscriber(trickle.TrickleSubscriberConfig{
 		URL: url.String(),
-		Ctx: ctx,
+		Ctx: drainCtx,
 	})
 	if err != nil {
+		drainCancel()
 		clog.Infof(ctx, "Failed to create data subscriber: %s", err)
 		return
 	}
@@ -809,6 +828,7 @@ func (bsg *BYOCGatewayServer) startDataSubscribe(ctx context.Context, url *url.U
 	// read segments from trickle subscription
 	go func() {
 		defer dataWriter.Close()
+		defer drainCancel()
 
 		var err error
 		firstSegment := true
@@ -819,12 +839,15 @@ func (bsg *BYOCGatewayServer) startDataSubscribe(ctx context.Context, url *url.U
 		const maxRetries = 5
 		for {
 			select {
-			case <-ctx.Done():
+			case <-drainCtx.Done():
 				clog.Info(ctx, "data subscribe done")
 				return
 			default:
 			}
-			if !params.inputStreamExists(params.liveParams.streamID) {
+			// Skip the input-stream-exists guard once /stream/stop has fired —
+			// during the drain window the input stream has already gone away,
+			// but we still want to read the runner's pending final segments.
+			if ctx.Err() == nil && !params.inputStreamExists(params.liveParams.streamID) {
 				clog.Infof(ctx, "data subscribe stopping, input stream does not exist.")
 				break
 			}