fix(localai-proxy): bound live transcription setup and track in-flight turns precisely

A caller that gives up before the ready ack now ends the call with Canceled
and closes the upstream socket, and setup has a 3 minute default bound when
request_timeout_seconds is unset, so a hung upstream cannot hold the call
or block failover. A closing session now waits only for a turn the upstream
committed (or is still speaking), not for turns it discarded. Audio held
before the ready ack is capped at 5 s, and NaN samples become silence.

Assisted-by: Claude:claude-opus-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

backend/go/localai-proxy/live.go

@@ -28,27 +28,53 @@ const (
 	defaultLiveSampleRate = 16000
 
 	// finalWait bounds how long a closing session waits for an utterance the
-	// upstream VAD already started. The upstream transcribes only after its
+	// upstream still has in flight. The upstream transcribes only after its
 	// VAD sees the speech end, so a slow model can finish after the client
 	// stops sending; waiting forever would pin the gRPC call on a hung
 	// upstream.
 	finalWait = 5 * time.Second
 
+	// commitGrace is how long a speech_stopped waits for its
+	// input_audio_buffer.committed. The upstream sends the two back to back,
+	// so a stop with no commit inside this window is a discarded turn and
+	// must not hold a closing session for the full finalWait.
+	commitGrace = 500 * time.Millisecond
+
+	// maxBacklogSeconds caps the audio held before the ready ack. Callers
+	// wait for the ack before streaming, so more than this means a client
+	// that ignores the contract, and holding it unbounded while a cold
+	// upstream loads models would grow memory without limit.
+	maxBacklogSeconds = 5
+
 	// liveWriteTimeout turns an upstream that stops reading into an error
 	// instead of a call blocked on a full socket buffer.
 	liveWriteTimeout = 10 * time.Second
 
 	// liveHandshakeTimeout bounds the WebSocket upgrade only. The upstream
 	// warms the pipeline models before it sends session.created, which can
-	// take minutes on a cold box, so the setup phase after the upgrade is
-	// bounded by request_timeout_seconds instead.
+	// take minutes on a cold box, so the setup phase after the upgrade has
+	// its own, longer bound.
 	liveHandshakeTimeout = 30 * time.Second
+
+	// defaultLiveSetupTimeout bounds the session setup (upgrade to
+	// session.updated) when request_timeout_seconds is unset. Core waits for
+	// the ready ack with a plain Recv, so without a bound a hung upstream
+	// would hold the call, and the failover that should move the stage to
+	// the next target, forever. It is generous because a cold upstream loads
+	// the pipeline's VAD and transcription models before it answers.
+	defaultLiveSetupTimeout = 3 * time.Minute
 )
 
+// liveSetupTimeout is defaultLiveSetupTimeout, as a variable so tests can
+// exercise the bound without waiting minutes.
+var liveSetupTimeout = defaultLiveSetupTimeout
+
 // realtimeEvent holds the fields the bridge reads from any upstream server
 // event; unrelated events decode into it harmlessly.
 type realtimeEvent struct {
-	Type       string `json:"type"`
+	Type string `json:"type"`
+	// ItemID is set on committed, delta, completed and failed events; the
+	// upstream uses the committed turn's id for its transcription events.
 	ItemID     string `json:"item_id"`
 	Delta      string `json:"delta"`
 	Transcript string `json:"transcript"`
@@ -99,21 +125,26 @@ func (p *LocalAIProxy) AudioTranscriptionLive(in <-chan *pb.TranscriptLiveReques
 	defer func() { _ = conn.Close() }()
 
 	s := &liveSession{
-		conn:     conn,
-		out:      out,
-		pipeline: cfg.realtimePipeline,
-		language: lc.GetLanguage(),
-		rate:     rate,
-		sent:     map[string]string{},
-		events:   make(chan realtimeEvent),
-		done:     make(chan struct{}),
+		conn:      conn,
+		out:       out,
+		pipeline:  cfg.realtimePipeline,
+		language:  lc.GetLanguage(),
+		rate:      rate,
+		sent:      map[string]string{},
+		committed: map[string]bool{},
+		events:    make(chan realtimeEvent),
+		done:      make(chan struct{}),
 	}
 	// Closing done releases the reader if it is blocked handing over an
 	// event; closing conn (deferred above, runs after this) unblocks its read.
 	defer close(s.done)
 	go s.readLoop()
 
-	return s.run(in, cfg.timeout)
+	setup := cfg.timeout
+	if setup <= 0 {
+		setup = liveSetupTimeout
+	}
+	return s.run(in, setup)
 }
 
 // dialRealtime opens the upstream WebSocket. A refused upgrade is mapped like
@@ -159,9 +190,16 @@ type liveSession struct {
 	language string
 	rate     int
 
-	sent    map[string]string // text already sent as deltas, per upstream item
-	final   []string          // completed transcripts, in order
-	pending int               // utterances started upstream and not yet transcribed
+	sent  map[string]string // text already sent as deltas, per upstream item
+	final []string          // completed transcripts, in order
+
+	// In-flight tracking, so a closing session waits only for an utterance
+	// the upstream will still transcribe. The upstream VAD emits
+	// speech_started, then either speech_stopped plus committed (a turn it
+	// will transcribe) or nothing at all (a turn it discarded as no speech).
+	speaking  bool            // between speech_started and speech_stopped
+	stopping  bool            // speech_stopped seen, its committed not yet
+	committed map[string]bool // committed items not yet completed
 
 	events  chan realtimeEvent
 	readErr error // set before events is closed
@@ -199,53 +237,47 @@ func (s *liveSession) run(in <-chan *pb.TranscriptLiveRequest, setupTimeout time
 		created, ready bool
 		inClosed       bool
 		backlog        []*pb.TranscriptLiveRequest
-		setupTimer     <-chan time.Time
+		backlogSamples int
 		drainTimer     <-chan time.Time
-		drain          *time.Timer
+		graceTimer     <-chan time.Time
 	)
-	defer func() {
-		if drain != nil {
-			drain.Stop()
-		}
-	}()
-	if setupTimeout > 0 {
-		t := time.NewTimer(setupTimeout)
-		defer t.Stop()
-		setupTimer = t.C
-	}
-
-	// finishOrDrain finalizes once nothing is in flight, else waits (bounded)
-	// for the upstream to finish the utterance it already started.
-	finishOrDrain := func() (bool, error) {
-		if s.pending <= 0 {
-			return true, s.finish()
-		}
-		if drain == nil {
-			drain = time.NewTimer(finalWait)
-			drainTimer = drain.C
-		}
-		return false, nil
-	}
+	setup := time.NewTimer(setupTimeout)
+	defer setup.Stop()
+	setupTimer := setup.C
+	drain := time.NewTimer(finalWait)
+	drain.Stop()
+	defer drain.Stop()
+	grace := time.NewTimer(commitGrace)
+	grace.Stop()
+	defer grace.Stop()
 
 	for {
 		select {
 		case req, ok := <-in:
 			if !ok {
 				in, inClosed = nil, true
-				if ready {
-					if done, err := finishOrDrain(); done {
-						return err
-					}
+				if !ready {
+					// The caller gave up waiting for the ready ack (core
+					// closes its send side on failure or cancel). Canceled
+					// rather than nil: there is no session to report as
+					// complete, and failover neither retries nor trips a
+					// target on a canceled call.
+					return status.Error(codes.Canceled, "localai-proxy: live transcription closed before the upstream session was ready")
+				}
+				if s.inFlight() {
+					drain.Reset(finalWait)
+					drainTimer = drain.C
+				}
+			} else if !ready {
+				// Hold audio sent before the ready ack rather than drop it,
+				// up to a bound.
+				backlogSamples += len(req.GetAudio().GetPcm())
+				if backlogSamples > maxBacklogSeconds*s.rate {
+					return status.Errorf(codes.InvalidArgument,
+						"localai-proxy: more than %d s of audio sent before the live transcription session was ready", maxBacklogSeconds)
 				}
-				continue
-			}
-			if !ready {
-				// Callers wait for the ready ack before streaming, but hold
-				// anything sent early rather than drop it.
 				backlog = append(backlog, req)
-				continue
-			}
-			if err := s.forward(req); err != nil {
+			} else if err := s.forward(req); err != nil {
 				return err
 			}
 
@@ -275,37 +307,54 @@ func (s *liveSession) run(in <-chan *pb.TranscriptLiveRequest, setupTimeout time
 					}
 				}
 				backlog = nil
-				if inClosed {
-					if done, err := finishOrDrain(); done {
-						return err
-					}
-				}
 			case "error":
 				return s.upstreamError("error", ev)
 			case "conversation.item.input_audio_transcription.failed":
 				return s.upstreamError("transcription failed", ev)
 			case "input_audio_buffer.speech_started":
-				s.pending++
+				s.speaking, s.stopping = true, false
+			case "input_audio_buffer.speech_stopped":
+				if s.speaking {
+					s.speaking, s.stopping = false, true
+					grace.Reset(commitGrace)
+					graceTimer = grace.C
+				}
+			case "input_audio_buffer.committed":
+				s.stopping, graceTimer = false, nil
+				if ev.ItemID != "" {
+					s.committed[ev.ItemID] = true
+				}
 			case "conversation.item.input_audio_transcription.delta":
 				s.delta(ev)
 			case "conversation.item.input_audio_transcription.completed":
 				s.completed(ev)
-				if inClosed && s.pending <= 0 {
-					return s.finish()
-				}
 			}
 
+		case <-graceTimer:
+			// A stop the upstream never committed: the turn was discarded.
+			s.stopping, graceTimer = false, nil
+
 		case <-setupTimer:
-			return status.Errorf(codes.Unavailable, "localai-proxy: upstream %s did not set up the transcription session in %s", realtimePath, setupTimeout)
+			return status.Errorf(codes.Unavailable, "localai-proxy: upstream %s did not set up the transcription session within %s", realtimePath, setupTimeout)
 
 		case <-drainTimer:
 			xlog.Warn("localai-proxy: upstream did not finish the last utterance in time; finalizing without it",
 				"pipeline", s.pipeline, "wait", finalWait)
 			return s.finish()
 		}
+
+		if inClosed && !s.inFlight() {
+			return s.finish()
+		}
 	}
 }
 
+// inFlight reports an utterance the upstream is still expected to
+// transcribe.
+func (s *liveSession) inFlight() bool {
+	return s.speaking || s.stopping || len(s.committed) > 0
+}
+
 func (s *liveSession) sessionUpdate() map[string]any {
 	return map[string]any{
 		"type": "session.update",
@@ -343,7 +392,13 @@ func (s *liveSession) forward(req *pb.TranscriptLiveRequest) error {
 func pcm16LE(pcm []float32) []byte {
 	buf := make([]byte, len(pcm)*2)
 	for i, f := range pcm {
-		v := math.Max(-1, math.Min(1, float64(f)))
+		v := float64(f)
+		if math.IsNaN(v) {
+			// A NaN would convert to an arbitrary int16; silence is the
+			// only neutral value.
+			v = 0
+		}
+		v = math.Max(-1, math.Min(1, v))
 		binary.LittleEndian.PutUint16(buf[i*2:], uint16(int16(v*math.MaxInt16)))
 	}
 	return buf
@@ -368,9 +423,7 @@ func (s *liveSession) completed(ev realtimeEvent) {
 	if !ok {
 		rest = ""
 	}
-	if s.pending > 0 {
-		s.pending--
-	}
+	delete(s.committed, ev.ItemID)
 	if t := strings.TrimSpace(ev.Transcript); t != "" {
 		s.final = append(s.final, t)
 	}

backend/go/localai-proxy/live_test.go

@@ -4,6 +4,7 @@ import (
 	"encoding/base64"
 	"encoding/binary"
 	"encoding/json"
+	"math"
 	"net/http"
 	"runtime"
 	"strings"
@@ -147,6 +148,27 @@ func liveGoroutines() int {
 	return n
 }
 
+func ev(typ string) map[string]any { return map[string]any{"type": typ} }
+
+func item(typ, id string) map[string]any { return map[string]any{"type": typ, "item_id": id} }
+
+func completed(id, transcript string) map[string]any {
+	return map[string]any{"type": "conversation.item.input_audio_transcription.completed", "item_id": id, "transcript": transcript}
+}
+
+// stallAfterUpdate plays session.created, reads the session.update and then
+// never answers, as a hung upstream would, until the spec ends.
+func stallAfterUpdate(c *websocket.Conn, gone chan<- struct{}) {
+	wsSend(c, map[string]any{"type": "session.created", "session": map[string]any{}})
+	wsRecv(c)
+	for {
+		if _, ok := wsRecv(c); !ok {
+			close(gone)
+			return
+		}
+	}
+}
+
 func withPipeline(o *pb.ModelOptions) {
 	o.Options = append(o.Options, "realtime_pipeline:remote-pipe")
 }
@@ -249,25 +271,30 @@ var _ = Describe("AudioTranscriptionLive", func() {
 
 		lc.audio(0, 0.5, -1, 1, 2)
 		lc.audio(-0.25)
+		lc.audio(float32(math.NaN()))
 		close(lc.in)
 		Expect(lc.finish()).To(Succeed())
 
 		Eventually(frames, 2*time.Second).Should(Receive(Equal([][]int16{
 			{0, 16383, -32767, 32767, 32767},
 			{-8191},
+			{0},
 		})))
 	})
 
 	It("maps deltas and completions to Delta and Eou, and finishes with the full text", func() {
 		up := wsUpstream(func(c *websocket.Conn, _ *http.Request) {
 			wsHandshake(c)
-			wsSend(c, map[string]any{"type": "input_audio_buffer.speech_started"})
+			wsSend(c, ev("input_audio_buffer.speech_started"))
+			wsSend(c, ev("input_audio_buffer.speech_stopped"))
+			wsSend(c, item("input_audio_buffer.committed", "a"))
 			wsSend(c, map[string]any{"type": "conversation.item.input_audio_transcription.delta", "item_id": "a", "delta": "hel"})
 			wsSend(c, map[string]any{"type": "conversation.item.input_audio_transcription.delta", "item_id": "a", "delta": "lo"})
-			wsSend(c, map[string]any{"type": "input_audio_buffer.speech_stopped"})
-			wsSend(c, map[string]any{"type": "conversation.item.input_audio_transcription.completed", "item_id": "a", "transcript": "hello world"})
-			wsSend(c, map[string]any{"type": "input_audio_buffer.speech_started"})
-			wsSend(c, map[string]any{"type": "conversation.item.input_audio_transcription.completed", "item_id": "b", "transcript": "again"})
+			wsSend(c, completed("a", "hello world"))
+			wsSend(c, ev("input_audio_buffer.speech_started"))
+			wsSend(c, ev("input_audio_buffer.speech_stopped"))
+			wsSend(c, item("input_audio_buffer.committed", "b"))
+			wsSend(c, completed("b", "again"))
 			wsDrain(c)
 		})
 		DeferCleanup(up.Close)
@@ -291,29 +318,58 @@ var _ = Describe("AudioTranscriptionLive", func() {
 		Expect(lc.finish()).To(Succeed())
 	})
 
-	It("waits for an in-flight utterance before the final result", func() {
+	It("waits for a committed utterance before the final result", func() {
 		up := wsUpstream(func(c *websocket.Conn, _ *http.Request) {
 			wsHandshake(c)
-			wsSend(c, map[string]any{"type": "input_audio_buffer.speech_started"})
-			// The client closes its side now; the transcription lands later.
+			wsSend(c, ev("input_audio_buffer.speech_started"))
+			wsSend(c, ev("input_audio_buffer.speech_stopped"))
+			wsSend(c, item("input_audio_buffer.committed", "a"))
+			wsSend(c, map[string]any{"type": "conversation.item.input_audio_transcription.delta", "item_id": "a", "delta": "late"})
+			// The client closes its side once it sees the delta; the
+			// transcription completes later.
 			time.Sleep(300 * time.Millisecond)
-			wsSend(c, map[string]any{"type": "input_audio_buffer.speech_stopped"})
-			wsSend(c, map[string]any{"type": "conversation.item.input_audio_transcription.completed", "item_id": "a", "transcript": "late words"})
+			wsSend(c, completed("a", "late words"))
 			wsDrain(c)
 		})
 		DeferCleanup(up.Close)
 		lc := startLive(loadProxy(up, withPipeline))
 		lc.config("en", 16000)
 		Expect(lc.next().GetReady()).To(BeTrue())
+		Expect(lc.next().GetDelta()).To(Equal("late"))
 		close(lc.in)
 
 		r := lc.next()
-		Expect(r.GetDelta()).To(Equal("late words"))
+		Expect(r.GetDelta()).To(Equal(" words"))
 		Expect(r.GetEou()).To(BeTrue())
 		Expect(lc.next().GetFinalResult().GetText()).To(Equal("late words"))
 		Expect(lc.finish()).To(Succeed())
 	})
 
+	It("does not hold the close for a turn the upstream discarded", func() {
+		up := wsUpstream(func(c *websocket.Conn, _ *http.Request) {
+			wsHandshake(c)
+			wsSend(c, ev("input_audio_buffer.speech_started"))
+			wsSend(c, ev("input_audio_buffer.speech_stopped"))
+			wsSend(c, item("input_audio_buffer.committed", "a"))
+			wsSend(c, completed("a", "kept"))
+			// A stop that is never committed, then nothing more.
+			wsSend(c, ev("input_audio_buffer.speech_started"))
+			wsSend(c, ev("input_audio_buffer.speech_stopped"))
+			wsDrain(c)
+		})
+		DeferCleanup(up.Close)
+		lc := startLive(loadProxy(up, withPipeline))
+		lc.config("en", 16000)
+		Expect(lc.next().GetReady()).To(BeTrue())
+		Expect(lc.next().GetDelta()).To(Equal("kept"))
+		close(lc.in)
+
+		start := time.Now()
+		Expect(lc.next().GetFinalResult().GetText()).To(Equal("kept"))
+		Expect(lc.finish()).To(Succeed())
+		Expect(time.Since(start)).To(BeNumerically("<", finalWait/2))
+	})
+
 	It("ends with Unavailable on an upstream error event during setup", func() {
 		up := wsUpstream(func(c *websocket.Conn, _ *http.Request) {
 			wsSend(c, map[string]any{"type": "session.created", "session": map[string]any{}})
@@ -377,4 +433,101 @@ var _ = Describe("AudioTranscriptionLive", func() {
 		Eventually(liveGoroutines, time.Second).Should(Equal(before))
 		close(lc.in)
 	})
+
+	It("gives up with Canceled when the caller closes before the ready ack", func() {
+		before := liveGoroutines()
+		gone := make(chan struct{})
+		up := wsUpstream(func(c *websocket.Conn, _ *http.Request) { stallAfterUpdate(c, gone) })
+		DeferCleanup(up.Close)
+		lc := startLive(loadProxy(up, withPipeline))
+		lc.config("en", 16000)
+		lc.audio(0.1)
+		close(lc.in)
+
+		Expect(status.Code(lc.finish())).To(Equal(codes.Canceled))
+		Eventually(gone, 2*time.Second).Should(BeClosed(), "upstream socket left open")
+		Eventually(liveGoroutines, time.Second).Should(Equal(before))
+	})
+
+	It("bounds a hung setup by request_timeout_seconds", func() {
+		before := liveGoroutines()
+		gone := make(chan struct{})
+		up := wsUpstream(func(c *websocket.Conn, _ *http.Request) { stallAfterUpdate(c, gone) })
+		DeferCleanup(up.Close)
+		p := loadProxy(up, func(o *pb.ModelOptions) {
+			withPipeline(o)
+			o.Proxy.RequestTimeoutSeconds = 1
+		})
+		lc := startLive(p)
+		lc.config("en", 16000)
+
+		var err error
+		Eventually(lc.errc, 3*time.Second).Should(Receive(&err))
+		Expect(status.Code(err)).To(Equal(codes.Unavailable))
+		Expect(lc.out).To(BeClosed())
+		Eventually(gone, 2*time.Second).Should(BeClosed(), "upstream socket left open")
+		Eventually(liveGoroutines, time.Second).Should(Equal(before))
+		close(lc.in)
+	})
+
+	It("bounds a hung setup by default when request_timeout_seconds is unset", func() {
+		Expect(defaultLiveSetupTimeout).To(BeNumerically(">=", 2*time.Minute))
+		saved := liveSetupTimeout
+		liveSetupTimeout = 300 * time.Millisecond
+		DeferCleanup(func() { liveSetupTimeout = saved })
+
+		before := liveGoroutines()
+		gone := make(chan struct{})
+		up := wsUpstream(func(c *websocket.Conn, _ *http.Request) { stallAfterUpdate(c, gone) })
+		DeferCleanup(up.Close)
+		lc := startLive(loadProxy(up, withPipeline))
+		lc.config("en", 16000)
+
+		Expect(status.Code(lc.finish())).To(Equal(codes.Unavailable))
+		Eventually(gone, 2*time.Second).Should(BeClosed(), "upstream socket left open")
+		Eventually(liveGoroutines, time.Second).Should(Equal(before))
+		close(lc.in)
+	})
+
+	It("forwards audio sent before the ready ack once ready", func() {
+		release := make(chan struct{})
+		frames := make(chan [][]int16, 1)
+		up := wsUpstream(func(c *websocket.Conn, _ *http.Request) {
+			wsSend(c, map[string]any{"type": "session.created", "session": map[string]any{}})
+			wsRecv(c)
+			<-release
+			wsSend(c, map[string]any{"type": "session.updated", "session": map[string]any{}})
+			frames <- wsDrain(c)
+		})
+		DeferCleanup(up.Close)
+		lc := startLive(loadProxy(up, withPipeline))
+		lc.config("en", 16000)
+		lc.audio(0.5)
+		close(release)
+		Expect(lc.next().GetReady()).To(BeTrue())
+		close(lc.in)
+		Expect(lc.finish()).To(Succeed())
+		Eventually(frames, 2*time.Second).Should(Receive(Equal([][]int16{{16383}})))
+	})
+
+	It("refuses more than the backlog cap of audio before the ready ack", func() {
+		gone := make(chan struct{})
+		up := wsUpstream(func(c *websocket.Conn, _ *http.Request) { stallAfterUpdate(c, gone) })
+		DeferCleanup(up.Close)
+		lc := startLive(loadProxy(up, withPipeline))
+		lc.config("en", 16000)
+		second := make([]float32, 16000)
+		for range maxBacklogSeconds + 1 {
+			select {
+			case lc.in <- &pb.TranscriptLiveRequest{Payload: &pb.TranscriptLiveRequest_Audio{Audio: &pb.TranscriptLiveAudio{Pcm: second}}}:
+			case <-time.After(2 * time.Second):
+				Fail("bridge stopped reading audio before the cap")
+			}
+		}
+
+		err := lc.finish()
+		Expect(status.Code(err)).To(Equal(codes.InvalidArgument))
+		Expect(err.Error()).To(ContainSubstring("before the live transcription session was ready"))
+		close(lc.in)
+	})
 })