diff --git a/backend/go/localai-proxy/live.go b/backend/go/localai-proxy/live.go index 622358729..0193b973a 100644 --- a/backend/go/localai-proxy/live.go +++ b/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) } diff --git a/backend/go/localai-proxy/live_test.go b/backend/go/localai-proxy/live_test.go index 22e30e6d1..c35b4463b 100644 --- a/backend/go/localai-proxy/live_test.go +++ b/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) + }) })