ai-resume/internal/server/stream.go

package server

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"

	"gitea.lab.audasmedia.com.au/sam/ai-resume/internal/stage"
	"gitea.lab.audasmedia.com.au/sam/ai-resume/internal/web/templates"
)

// StreamEvent is the SSE event the panel listens for. The binding rendered by
// stage.SSEDriver uses the same name, so a client needs no configuration.
const StreamEvent = "stage"

// maxStreamKeys caps how many fragments one connection may ask for, so a
// crafted query cannot turn the demonstration into a long-lived loop.
const maxStreamKeys = 24

// handleStream points the stage panel at an event stream.
//
// This is the Phase 2 proof of the §2.4 contract: the panel markup does not
// change, only the binding does. The frames come from the fragments the build
// already rendered, so nothing calls a model and the site still works with
// every AI service stopped. Phase 5 replaces the source behind the same
// interface (stage.Source) and this handler does not change.
func (s *Server) handleStream(w http.ResponseWriter, r *http.Request) {
	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "streaming is not supported by this server", http.StatusInternalServerError)
		return
	}
	if s.manifest == nil || len(s.manifest.Entries) == 0 {
		msg := "the stage manifest is not part of this build"
		if s.manifestErr != nil {
			s.log.Warn("stream unavailable", "error", s.manifestErr)
		}
		http.Error(w, msg, http.StatusServiceUnavailable)
		return
	}

	keys := s.streamKeys(r.URL.Query())
	if len(keys) == 0 {
		keys = s.defaultStreamKeys()
	}
	if len(keys) == 0 {
		http.Error(w, "no fragments to stream", http.StatusNotFound)
		return
	}

	h := w.Header()
	h.Set("Content-Type", "text/event-stream; charset=utf-8")
	h.Set("Cache-Control", "no-store")
	h.Set("Connection", "keep-alive")
	h.Set("X-Accel-Buffering", "no")
	w.WriteHeader(http.StatusOK)

	// A client that reconnects gets a short pause rather than an immediate retry.
	if _, err := io.WriteString(w, "retry: 2000\n\n"); err != nil {
		return
	}
	flusher.Flush()

	ctx := r.Context()
	source := s.newStreamSource(ctx)
	err := source.Stream(ctx, keys, func(f stage.Fragment) error {
		var buf bytes.Buffer
		if f.Rendered {
			// The build already rendered the whole article; send those bytes.
			if err := f.Body.Render(ctx, &buf); err != nil {
				return err
			}
		} else {
			// A structured source (Phase 5) hands over parts, so the panel
			// renders them here with the same component the route driver uses.
			if err := templates.FragmentView(f).Render(ctx, &buf); err != nil {
				return err
			}
		}
		if err := writeSSE(w, StreamEvent, buf.String()); err != nil {
			return err
		}
		flusher.Flush()
		return nil
	})
	if err != nil && ctx.Err() == nil {
		// Headers are already sent, so there is no status to change. Log it and
		// let the panel keep the last fragment it received.
		s.log.Warn("stage stream ended early", "error", err)
	}
}

// streamKeys reads the fragment ids a caller asked for. Both `?k=a&k=b` and
// `?k=a,b` work, because both are natural to write by hand.
func (s *Server) streamKeys(q url.Values) []string {
	var keys []string
	for _, raw := range q["k"] {
		for _, part := range strings.Split(raw, ",") {
			id := strings.TrimSpace(part)
			if id != "" && validFragmentID(id) {
				keys = append(keys, id)
			}
			if len(keys) >= maxStreamKeys {
				return keys
			}
		}
	}
	return keys
}

// defaultStreamKeys is the demonstration stream: the six area landing
// fragments, in navigation order. The build writes those first, so the manifest
// alone is enough — the server never reads the content tree.
func (s *Server) defaultStreamKeys() []string {
	if s.manifest == nil {
		return nil
	}
	keys := make([]string, 0, 6)
	for _, e := range s.manifest.Entries {
		if !strings.HasPrefix(e.ID, "doc-") {
			continue
		}
		keys = append(keys, e.ID)
		if len(keys) == 6 {
			break
		}
	}
	if len(keys) > 0 {
		return keys
	}
	for _, e := range s.manifest.Entries {
		keys = append(keys, e.ID)
		if len(keys) == 6 {
			break
		}
	}
	return keys
}

// newStreamSource wires the manifest to the files on disk and paces the frames.
func (s *Server) newStreamSource(ctx context.Context) *stage.ManifestSource {
	read := func(e stage.Entry) (string, error) {
		file := filepath.Join(s.dist, filepath.FromSlash(e.File))
		if !within(s.dist, file) {
			return "", fmt.Errorf("fragment %s is outside the build directory", e.ID)
		}
		raw, err := os.ReadFile(file)
		if err != nil {
			return "", fmt.Errorf("fragment %s: %w", e.ID, err)
		}
		return string(raw), nil
	}
	delay := s.cfg.StreamDelay
	frame := func(i int) {
		if i == 0 || delay <= 0 {
			return
		}
		select {
		case <-time.After(delay):
		case <-ctx.Done():
		}
	}
	return stage.NewManifestSource(s.manifest, read, frame)
}

// writeSSE writes one event. A payload with newlines is split across data
// lines, which is what the server-sent events format expects.
func writeSSE(w io.Writer, event, data string) error {
	var b strings.Builder
	b.WriteString("event: ")
	b.WriteString(event)
	b.WriteByte('\n')
	for _, line := range strings.Split(data, "\n") {
		b.WriteString("data: ")
		b.WriteString(line)
		b.WriteByte('\n')
	}
	b.WriteByte('\n')
	_, err := io.WriteString(w, b.String())
	return err
}