Maven/internal/worker/server.go

// worker/server.go — the worker-side listener + accept loop.
//
// Each call is dispatched to a single handler (Transcriber OR Synthesizer,
// depending on what the module process);
// the other verb returns ErrUnknownMethod — a stt process won't serve
// synthesize. robust to a misconfigured client (the daemon wiring chooses
// which module to dial; mixing the two is a config error caught cleanly by
// the wire, not a runtime goroutine panic). One Server per module process.
//
// The seam address decides the transport. On the default unix socket the
// perms mirror ipc.Server — dir 0700, socket 0600 ⇒ same unix user — and that
// is the whole auth floor, because the seam never leaves the box. A tcp
// address moves the module to another host and takes that floor away, so
// netaddr checks a shared token before the first frame. See internal/netaddr.
package worker

import (
	"context"
	"encoding/json"
	"fmt"
	"net"
	"sync"
	"sync/atomic"

	"github.com/kami/maven/internal/netaddr"
)

// Server — a worker module process's listener. Wires either a Transcriber,
// a Synthesizer, or both (the both case is unusual; the daemon typically
// runs two separate module processes). The unset verb returns
// ErrUnknownMethod.
type Server struct {
	t Transcriber
	s Synthesizer

	path string
	addr netaddr.Addr
	ln   net.Listener

	wg        sync.WaitGroup
	done      chan struct{}
	closeOnce sync.Once

	// connCount — assigned per accepted conn, used in logs to distinguish
	// concurrent connections. Monotonic; not load-bearing for correctness.
	connCount atomic.Uint64
}

// NewServer builds a Server with a Transcriber. The caller wires a
// Synthesizer via SetSynthesizer if this process serves tts. Use
// NewSynthesizerServer for the tts-only case (mirrors this constructor).
func NewServer(path string, t Transcriber) *Server {
	return &Server{t: t, path: path, done: make(chan struct{})}
}

// NewSynthesizerServer builds a Server with a Synthesizer (the tts module).
func NewSynthesizerServer(path string, s Synthesizer) *Server {
	return &Server{s: s, path: path, done: make(chan struct{})}
}

// SetSynthesizer wires the synthesize verb on a Transcriber-built Server.
// Used only when one process serves both (non-default; the daemon prefers
// two separate processes per the restart-free / fail-independent invariant).
func (srv *Server) SetSynthesizer(s Synthesizer) { srv.s = s }

// Listen binds the seam address the Server was built with.
//
// A bare path is a unix socket with 0700 dir + 0600 socket perms, the same
// floor as internal/ipc, and a stale socket is removed first so the worker
// process restarts cleanly after a crash. A "tcp://host:port?token=..."
// address binds a network listener instead, so this module can run on the
// workstation while core stays on homesrv; the token is mandatory there,
// because there is no filesystem to be the auth floor. See internal/netaddr.
func (srv *Server) Listen() error {
	addr, err := netaddr.Parse(srv.path)
	if err != nil {
		return err
	}
	ln, err := netaddr.Listen(addr)
	if err != nil {
		return err
	}
	srv.addr = addr
	srv.ln = ln
	return nil
}

// Path returns the bound socket path (after Listen; "" before).
func (srv *Server) Path() string { return srv.path }

// Serve accepts connections until the listener closes. Each connection is
// served in its own goroutine; a panicking handler tears down only that conn
// (the rest of the module keeps serving, restart-free per spec).
func (srv *Server) Serve() error {
	if srv.ln == nil {
		return fmt.Errorf("worker: serve before listen")
	}
	for {
		c, err := srv.ln.Accept()
		if err != nil {
			select {
			case <-srv.done:
				return nil
			default:
				return fmt.Errorf("worker: accept: %w", err)
			}
		}
		srv.wg.Add(1)
		go func(c net.Conn) {
			defer srv.wg.Done()
			defer c.Close()
			srv.serveConn(c)
		}(c)
	}
}

func (srv *Server) serveConn(c net.Conn) {
	id := srv.connCount.Add(1)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	for {
		var req Request
		if err := readFrame(c, &req); err != nil {
			return // EOF / malformed ⇒ end this conn
		}
		result, err := srv.safeDispatch(ctx, c.RemoteAddr(), id, req)
		resp := Response{}
		if err != nil {
			resp.Error = rpcErr(err)
		} else {
			resp.Result = result
		}
		if err := writeFrame(c, resp); err != nil {
			return
		}
	}
}

func (srv *Server) safeDispatch(ctx context.Context, addr net.Addr, id uint64, req Request) (result json.RawMessage, err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("worker: panic dispatching %s (conn %d): %v", req.Method, id, r)
		}
	}()
	return srv.dispatch(ctx, req)
}

func (srv *Server) dispatch(ctx context.Context, req Request) (json.RawMessage, error) {
	switch req.Method {
	case MethodTranscribe:
		if srv.t == nil {
			return nil, fmt.Errorf("%w: %s", ErrUnknownMethod, req.Method)
		}
		var p TranscribeReq
		if err := unmarshalParams(req.Params, &p); err != nil {
			return nil, err
		}
		resp, err := srv.t.Transcribe(ctx, p)
		if err != nil {
			return nil, err
		}
		return marshalResult(resp), nil

	case MethodSynthesize:
		if srv.s == nil {
			return nil, fmt.Errorf("%w: %s", ErrUnknownMethod, req.Method)
		}
		var p SynthesizeReq
		if err := unmarshalParams(req.Params, &p); err != nil {
			return nil, err
		}
		resp, err := srv.s.Synthesize(ctx, p)
		if err != nil {
			return nil, err
		}
		return marshalResult(resp), nil

	default:
		return nil, fmt.Errorf("%w: %s", ErrUnknownMethod, req.Method)
	}
}

// Close stops accepting and waits for in-flight connections to drain. The
// socket file is removed so a restart can rebind cleanly.
//
// Idempotent, and safe from two goroutines at once. The check-then-close it
// replaced let both callers see an open channel and the second close panicked,
// so a shutdown racing a signal handler took the process down the one way a
// clean shutdown is supposed to prevent.
func (srv *Server) Close() error {
	var err error
	srv.closeOnce.Do(func() {
		close(srv.done)
		if srv.ln == nil {
			return
		}
		err = srv.ln.Close()
		srv.wg.Wait()
		netaddr.Cleanup(srv.addr)
	})
	return err
}

func unmarshalParams(raw json.RawMessage, v any) error {
	if len(raw) == 0 {
		raw = []byte("null")
	}
	if err := json.Unmarshal(raw, v); err != nil {
		return fmt.Errorf("%w: %v", ErrBadParams, err)
	}
	return nil
}

func marshalResult(v any) json.RawMessage {
	if v == nil {
		return json.RawMessage("null")
	}
	b, _ := json.Marshal(v)
	return b
}