339 lines
9.5 KiB
Go
339 lines
9.5 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"crypto/subtle"
|
|
"embed"
|
|
"encoding/json"
|
|
"errors"
|
|
"flag"
|
|
"io/fs"
|
|
"log"
|
|
"net/http"
|
|
"net/url"
|
|
"os"
|
|
"os/signal"
|
|
"strings"
|
|
"sync"
|
|
"syscall"
|
|
"time"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
//go:embed web/*
|
|
var assets embed.FS
|
|
|
|
type server struct {
|
|
cfg config
|
|
token, origin string
|
|
slots chan struct{}
|
|
newProvider func(string) (asrProvider, error)
|
|
sessionLimit time.Duration
|
|
}
|
|
|
|
type modelOption struct {
|
|
ID string `json:"id"`
|
|
Enabled bool `json:"enabled"`
|
|
}
|
|
|
|
func (s *server) bailianCredentials(model string) (string, string) {
|
|
if strings.HasPrefix(model, "fun-") {
|
|
return s.cfg.funASRKey, s.cfg.funASRBase
|
|
}
|
|
return s.cfg.bailianKey, s.cfg.bailianBase
|
|
}
|
|
|
|
func (s *server) models() []modelOption {
|
|
out := []modelOption{}
|
|
for _, id := range strings.Split(env("BAILIAN_ASR_MODELS", "fun-asr-realtime"), ",") {
|
|
id = strings.TrimSpace(id)
|
|
if id != "" && id != "volc-bigmodel" {
|
|
key, _ := s.bailianCredentials(id)
|
|
out = append(out, modelOption{id, key != ""})
|
|
}
|
|
}
|
|
return append(out, modelOption{"volc-bigmodel", s.cfg.volcAppKey != "" || s.cfg.volcAppID != "" && s.cfg.volcToken != ""})
|
|
}
|
|
|
|
func (s *server) provider(id string) (asrProvider, error) {
|
|
for _, model := range s.models() {
|
|
if model.ID != id {
|
|
continue
|
|
}
|
|
if !model.Enabled {
|
|
return nil, errors.New("provider credentials are not configured")
|
|
}
|
|
if id == "volc-bigmodel" {
|
|
return newVolcASR(agentCfg{VolcAppKey: s.cfg.volcAppKey, VolcAppID: s.cfg.volcAppID, VolcAccessToken: s.cfg.volcToken, VolcResourceID: s.cfg.volcResourceID, VolcWS: s.cfg.volcWS}), nil
|
|
}
|
|
key, endpoint := s.bailianCredentials(id)
|
|
return newBailianASR(id, agentCfg{BailianKey: key, BailianWssBaseURL: endpoint}), nil
|
|
}
|
|
return nil, errors.New("unknown ASR model")
|
|
}
|
|
|
|
func (s *server) authorized(r *http.Request) bool {
|
|
token := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ")
|
|
for _, protocol := range websocket.Subprotocols(r) {
|
|
if strings.HasPrefix(protocol, "auth.") {
|
|
token = strings.TrimPrefix(protocol, "auth.")
|
|
}
|
|
}
|
|
return len(token) == len(s.token) && subtle.ConstantTimeCompare([]byte(token), []byte(s.token)) == 1
|
|
}
|
|
|
|
func (s *server) sameOrigin(r *http.Request) bool {
|
|
origin := r.Header.Get("Origin")
|
|
if origin == "" {
|
|
return true
|
|
} // CLI clients still need the access token.
|
|
if s.origin != "" {
|
|
return origin == s.origin
|
|
}
|
|
u, err := url.Parse(origin)
|
|
return err == nil && u.User == nil && u.Path == "" && u.RawQuery == "" && u.Fragment == "" && (u.Scheme == "http" || u.Scheme == "https") && u.Host == r.Host
|
|
}
|
|
|
|
func (s *server) handler(ctx context.Context) http.Handler {
|
|
mux := http.NewServeMux()
|
|
mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write([]byte(`{"status":"ok","scope":"asr-only","llm":false,"tts":false}`))
|
|
})
|
|
mux.HandleFunc("GET /api/models", func(w http.ResponseWriter, r *http.Request) {
|
|
if !s.authorized(r) {
|
|
http.Error(w, "unauthorized", http.StatusUnauthorized)
|
|
return
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_ = json.NewEncoder(w).Encode(s.models())
|
|
})
|
|
mux.HandleFunc("GET /ws", func(w http.ResponseWriter, r *http.Request) { s.serveWS(ctx, w, r) })
|
|
web, _ := fs.Sub(assets, "web")
|
|
mux.Handle("GET /", http.FileServerFS(web))
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("X-Content-Type-Options", "nosniff")
|
|
w.Header().Set("Referrer-Policy", "no-referrer")
|
|
w.Header().Set("Cache-Control", "no-store")
|
|
w.Header().Set("Content-Security-Policy", "default-src 'self'; script-src 'self'; style-src 'self'; connect-src 'self'; frame-ancestors 'none'; base-uri 'none'; form-action 'self'")
|
|
mux.ServeHTTP(w, r)
|
|
})
|
|
}
|
|
|
|
func (s *server) serveWS(parent context.Context, w http.ResponseWriter, r *http.Request) {
|
|
if !s.authorized(r) {
|
|
http.Error(w, "unauthorized", http.StatusUnauthorized)
|
|
return
|
|
}
|
|
if !s.sameOrigin(r) {
|
|
http.Error(w, "origin rejected", http.StatusForbidden)
|
|
return
|
|
}
|
|
select {
|
|
case s.slots <- struct{}{}:
|
|
defer func() { <-s.slots }()
|
|
default:
|
|
http.Error(w, "session limit reached", http.StatusTooManyRequests)
|
|
return
|
|
}
|
|
up := websocket.Upgrader{Subprotocols: []string{"asr.v1"}, CheckOrigin: s.sameOrigin, HandshakeTimeout: 5 * time.Second}
|
|
conn, err := up.Upgrade(w, r, nil)
|
|
if err != nil {
|
|
return
|
|
}
|
|
limit := s.sessionLimit
|
|
if limit == 0 {
|
|
limit = 5 * time.Minute
|
|
}
|
|
ctx, cancel := context.WithTimeout(parent, limit)
|
|
defer cancel()
|
|
defer conn.Close()
|
|
conn.SetReadLimit(64 << 10)
|
|
stopClose := context.AfterFunc(ctx, func() { _ = conn.Close() })
|
|
defer stopClose()
|
|
var writeMu sync.Mutex
|
|
var stateMu sync.Mutex
|
|
var active *asrRun
|
|
send := func(expected *asrRun, value any) {
|
|
writeMu.Lock()
|
|
defer writeMu.Unlock()
|
|
stateMu.Lock()
|
|
valid := expected == nil || active == expected
|
|
stateMu.Unlock()
|
|
if !valid || ctx.Err() != nil {
|
|
return
|
|
}
|
|
_ = conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
|
|
if err := conn.WriteJSON(value); err != nil {
|
|
cancel()
|
|
}
|
|
}
|
|
stop := func() {
|
|
stateMu.Lock()
|
|
old := active
|
|
active = nil
|
|
stateMu.Unlock()
|
|
if old != nil {
|
|
old.cancel()
|
|
old.provider.close()
|
|
}
|
|
}
|
|
defer stop()
|
|
send(nil, map[string]any{"type": "ready", "sample_rate": 16000, "channels": 1, "format": "pcm_s16le", "max_session_seconds": int(limit.Seconds()), "llm": false, "tts": false})
|
|
for {
|
|
kind, data, err := conn.ReadMessage()
|
|
if err != nil {
|
|
return
|
|
}
|
|
if kind == websocket.BinaryMessage {
|
|
stateMu.Lock()
|
|
run := active
|
|
stateMu.Unlock()
|
|
if run == nil {
|
|
send(nil, map[string]string{"type": "error", "error": "start ASR before sending audio"})
|
|
continue
|
|
}
|
|
if len(data)%2 != 0 {
|
|
send(run, map[string]string{"type": "error", "error": "PCM16 must have an even byte length"})
|
|
continue
|
|
}
|
|
if err := run.provider.sendAudio(data); err != nil {
|
|
send(run, map[string]string{"type": "error", "error": "upstream audio write failed"})
|
|
stop()
|
|
}
|
|
continue
|
|
}
|
|
var msg struct {
|
|
Type string `json:"type"`
|
|
Model string `json:"model"`
|
|
}
|
|
if kind != websocket.TextMessage || json.Unmarshal(data, &msg) != nil {
|
|
send(nil, map[string]string{"type": "error", "error": "invalid control message"})
|
|
continue
|
|
}
|
|
switch msg.Type {
|
|
case "start":
|
|
stop()
|
|
factory := s.newProvider
|
|
if factory == nil {
|
|
factory = s.provider
|
|
}
|
|
p, err := factory(msg.Model)
|
|
if err != nil {
|
|
send(nil, map[string]string{"type": "error", "error": err.Error()})
|
|
continue
|
|
}
|
|
runCtx, runCancel := context.WithCancel(ctx)
|
|
run := &asrRun{provider: p, cancel: runCancel, ctx: runCtx}
|
|
if err := p.start(runCtx); err != nil {
|
|
runCancel()
|
|
p.close()
|
|
send(nil, map[string]string{"type": "error", "error": "ASR start failed; check provider credentials, model and endpoint"})
|
|
continue
|
|
}
|
|
stateMu.Lock()
|
|
active = run
|
|
stateMu.Unlock()
|
|
send(run, map[string]string{"type": "asr-started", "model": msg.Model})
|
|
go func() {
|
|
defer runCancel()
|
|
defer p.close()
|
|
for {
|
|
select {
|
|
case <-runCtx.Done():
|
|
return
|
|
case ev, ok := <-p.events():
|
|
if !ok {
|
|
send(run, map[string]string{"type": "asr-stopped"})
|
|
stateMu.Lock()
|
|
if active == run {
|
|
active = nil
|
|
}
|
|
stateMu.Unlock()
|
|
return
|
|
}
|
|
send(run, map[string]string{"type": ev.Typ, "text": ev.Text, "code": ev.Code, "error": ev.Error})
|
|
}
|
|
}
|
|
}()
|
|
case "finish":
|
|
stateMu.Lock()
|
|
run := active
|
|
stateMu.Unlock()
|
|
if run != nil {
|
|
run.finishOnce.Do(func() {
|
|
if err := run.provider.finish(); err != nil {
|
|
send(run, map[string]string{"type": "error", "error": "upstream finish failed"})
|
|
stop()
|
|
return
|
|
}
|
|
go func() {
|
|
select {
|
|
case <-run.ctx.Done():
|
|
case <-time.After(10 * time.Second):
|
|
stateMu.Lock()
|
|
same := active == run
|
|
stateMu.Unlock()
|
|
if same {
|
|
run.cancel()
|
|
run.provider.close()
|
|
send(run, map[string]string{"type": "asr-stopped"})
|
|
}
|
|
}
|
|
}()
|
|
})
|
|
}
|
|
case "cancel":
|
|
stop()
|
|
send(nil, map[string]string{"type": "asr-stopped"})
|
|
default:
|
|
send(nil, map[string]string{"type": "error", "error": "unsupported command; LLM/TTS are disabled"})
|
|
}
|
|
}
|
|
}
|
|
|
|
type asrRun struct {
|
|
provider asrProvider
|
|
cancel context.CancelFunc
|
|
ctx context.Context
|
|
finishOnce sync.Once
|
|
}
|
|
|
|
func main() {
|
|
health := flag.Bool("healthcheck", false, "check local process liveness")
|
|
flag.Parse()
|
|
if *health {
|
|
client := http.Client{Timeout: 2 * time.Second}
|
|
resp, err := client.Get("http://127.0.0.1:8080/healthz")
|
|
if err != nil {
|
|
os.Exit(1)
|
|
}
|
|
defer resp.Body.Close()
|
|
if resp.StatusCode != 200 {
|
|
os.Exit(1)
|
|
}
|
|
return
|
|
}
|
|
cfg := loadConfig()
|
|
token := os.Getenv("ASR_WEB_TOKEN")
|
|
origin := os.Getenv("PUBLIC_ORIGIN")
|
|
if err := validateConfig(token, origin, cfg); err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
s := &server{cfg: cfg, token: token, origin: origin, slots: make(chan struct{}, 4)}
|
|
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
|
|
defer stop()
|
|
srv := &http.Server{Addr: env("ASR_LISTEN", "127.0.0.1:8080"), Handler: s.handler(ctx), ReadHeaderTimeout: 5 * time.Second, IdleTimeout: 60 * time.Second, MaxHeaderBytes: 16 << 10}
|
|
go func() {
|
|
<-ctx.Done()
|
|
deadline, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
_ = srv.Shutdown(deadline)
|
|
}()
|
|
log.Printf("ASR-only service listening on %s; LLM/TTS disabled", srv.Addr)
|
|
if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
|
|
log.Fatal(err)
|
|
}
|
|
}
|