internal/router/loading.go
Rama main · réplica de lectura del repositorio original.
package router
import (
"context"
"fmt"
"math/rand"
"net/http"
"strings"
"sync"
"time"
"github.com/quartermaster-labs/quartermaster/internal/logmon"
)
var loadingPaths = []string{
"/v1/chat/completions",
}
func isLoadingPath(path string) bool {
for _, p := range loadingPaths {
if strings.HasPrefix(path, p) {
return true
}
}
return false
}
// loadingWriter keeps a streaming client's connection alive while its model
// loads, and narrates the wait on the way.
//
// It narrates in SSE *comment* frames (`: text`), never in `data:` frames. The
// connection is held open either way, but a comment is discarded by every
// conforming SSE parser, so nothing this writer emits can reach a client as
// model output. It used to send synthetic `reasoning_content` deltas: those
// prepended fabricated reasoning to the assistant message and carried none of
// the chunk fields (id/object/model/created) a strict client expects. Whatever
// is added here stays a comment -- progress belongs beside the stream, not
// inside it.
type loadingWriter struct {
hasWritten bool
writer http.ResponseWriter
req *http.Request
ctx context.Context
logger *logmon.Monitor
modelName string
startTime time.Time
pendingMu sync.Mutex
pendingUpdate string
// statusMu guards the last status sent, so an unchanged one is not resent
// on every queue broadcast.
statusMu sync.Mutex
lastStatus string
// writeMu serializes writes to the underlying writer and guards released.
// Once released is set, the streaming goroutine must not touch the writer
// again — ServeHTTP has reclaimed it (to run the real handler or to return)
// and writing/flushing a finalized response panics.
writeMu sync.Mutex
released bool
// closed by start when the goroutine finishes (after cleanup messages)
done chan struct{}
// test-only: closed when start enters its loop
loopStarted chan struct{}
// test-only: override the 1s tick interval
tickDuration time.Duration
}
func newLoadingWriter(logger *logmon.Monitor, modelName string, w http.ResponseWriter, req *http.Request) *loadingWriter {
s := &loadingWriter{
writer: w,
req: req,
ctx: req.Context(),
logger: logger,
modelName: modelName,
startTime: time.Now(),
tickDuration: 750 * time.Millisecond,
}
s.Header().Set("Content-Type", "text/event-stream")
s.Header().Set("Cache-Control", "no-cache")
s.Header().Set("Connection", "keep-alive")
s.WriteHeader(http.StatusOK)
s.sendComment(fmt.Sprintf("quartermaster loading model: %s", modelName))
// Loading until told otherwise: this writer only exists because the model
// was not ready, and a request that turns out to be queued learns so from
// its first position broadcast, a moment later.
s.setStatus("loading")
return s
}
// statusPrefix marks the one comment frame in this stream that is meant to be
// PARSED rather than read. Everything else here is narration for a human
// reading a curl -N; this tells a client which of the two waits it is in --
// queued behind another model's turn, or watching its own model load -- so it
// can say so instead of guessing. Still a comment, so it stays invisible to
// every client that does not opt in.
const statusPrefix = "qm-status: "
// setStatus publishes a machine-readable wait state ("waiting 2", "loading").
// Sent immediately rather than through pendingUpdate: the tick loop paces
// narration, and a status the client renders should not wait its own turn
// behind a whimsical remark. Deduped, because queue broadcasts repeat.
func (s *loadingWriter) setStatus(status string) {
s.statusMu.Lock()
unchanged := s.lastStatus == status
s.lastStatus = status
s.statusMu.Unlock()
if unchanged {
return
}
s.sendComment(statusPrefix + status)
}
func (s *loadingWriter) setUpdate(msg string) {
s.pendingMu.Lock()
s.pendingUpdate = msg
s.pendingMu.Unlock()
}
func (s *loadingWriter) start(ctx context.Context) {
s.done = make(chan struct{})
defer close(s.done)
defer func() {
// Skip cleanup writes if the client disconnected — the connection
// is being torn down and flushing against it will panic.
if s.ctx.Err() != nil {
return
}
duration := time.Since(s.startTime)
s.sendComment(fmt.Sprintf("model ready (%.2fs)", duration.Seconds()))
}()
remarks := make([]string, len(loadingRemarks))
copy(remarks, loadingRemarks)
rand.Shuffle(len(remarks), func(i, j int) {
remarks[i], remarks[j] = remarks[j], remarks[i]
})
ri := 0
nextRemarkIn := time.Duration(2+rand.Intn(4)) * time.Second
lastRemarkTime := time.Time{}
ticker := time.NewTicker(s.tickDuration)
defer ticker.Stop()
if s.loopStarted != nil {
close(s.loopStarted)
}
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
s.pendingMu.Lock()
update := s.pendingUpdate
s.pendingUpdate = ""
s.pendingMu.Unlock()
if update != "" {
s.sendComment(update)
lastRemarkTime = time.Now()
nextRemarkIn = time.Duration(5+rand.Intn(5)) * time.Second
} else if time.Since(lastRemarkTime) >= nextRemarkIn {
remark := remarks[ri%len(remarks)]
ri++
s.sendComment(remark)
lastRemarkTime = time.Now()
nextRemarkIn = time.Duration(5+rand.Intn(5)) * time.Second
} else {
// bare keepalive comment: holds the connection without saying
// anything
s.sendComment("")
}
}
}
}
func (s *loadingWriter) waitForCompletion(timeout time.Duration) bool {
if s.done == nil {
return true
}
select {
case <-s.done:
return true
case <-time.After(timeout):
return false
}
}
// sendComment emits one SSE comment frame. Comments are the protocol's own
// keepalive: they hold the connection open and every conforming parser drops
// them, so this narration can never be mistaken for model output. An empty
// text is a bare keepalive. Newlines would end the comment and turn the rest
// into an unintended field, so they are flattened.
func (s *loadingWriter) sendComment(text string) {
text = strings.NewReplacer("\r", " ", "\n", " ").Replace(text)
s.writeMu.Lock()
defer s.writeMu.Unlock()
// Once ServeHTTP has reclaimed the writer (release), writing/flushing it
// races the real handler or panics on a finalized response. Stop here.
if s.released {
return
}
if _, err := fmt.Fprintf(s.writer, ": %s\n\n", text); err != nil {
s.logger.Debugf("<%s> Failed to write SSE comment (client likely disconnected): %v", s.modelName, err)
return
}
if flusher, ok := s.writer.(http.Flusher); ok {
flusher.Flush()
}
}
// release fences the loadingWriter off from the underlying ResponseWriter.
// After it returns, the streaming goroutine will not write to or flush the
// writer again: any in-flight write completes under writeMu first, and later
// writes short-circuit on released. The caller can then safely hand the writer
// to the real handler or let ServeHTTP return without racing a finalized
// response (a use-after-return Flush panics on the recycled *bufio.Writer).
func (s *loadingWriter) release() {
s.writeMu.Lock()
s.released = true
s.writeMu.Unlock()
}
func (s *loadingWriter) Header() http.Header {
return s.writer.Header()
}
func (s *loadingWriter) Write(data []byte) (int, error) {
return s.writer.Write(data)
}
func (s *loadingWriter) WriteHeader(statusCode int) {
if s.hasWritten {
return
}
s.hasWritten = true
s.writer.WriteHeader(statusCode)
s.Flush()
}
func (s *loadingWriter) Flush() {
if flusher, ok := s.writer.(http.Flusher); ok {
flusher.Flush()
}
}