crossbar/cmd/fakeupstream/main.go

// fakeupstream stands in for a llama-server router in tests and the smoke run. Do not edit.
//
//	fakeupstream -listen 127.0.0.1:18081 -name alpha -models a,b -down-file /tmp/alpha.down -slow 0
//
// /health answers 503 while the down file exists, 200 otherwise. /v1/models lists -models.
// /props answers a small JSON object. /v1/chat/completions echoes: a streamed answer of five
// SSE chunks 200 ms apart when the body has "stream": true, then a final chunk carrying
// "usage" and llama-server style "timings", then [DONE]; one JSON answer with usage and
// timings otherwise. -slow adds that many milliseconds before answering (for queue tests).
// Every response carries X-Upstream: <name>. /props reports -n-ctx and -slots. With -wol-listen,
// a valid wake-on-LAN magic packet for -wol-mac received on that UDP address removes the down
// file, so the fake "boots" when woken.
package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"strings"
	"time"
)

func main() {
	listen := flag.String("listen", "127.0.0.1:18081", "address to listen on")
	name := flag.String("name", "fake", "name reported in X-Upstream and answers")
	models := flag.String("models", "m", "comma-separated model ids for /v1/models")
	downFile := flag.String("down-file", "", "while this file exists, /health answers 503")
	slow := flag.Int("slow", 0, "milliseconds to wait before answering a completion")
	nCtx := flag.Int("n-ctx", 8192, "n_ctx reported by /props")
	slots := flag.Int("slots", 2, "total_slots reported by /props")
	wolListen := flag.String("wol-listen", "", "UDP address to listen on for a wake-on-LAN magic packet")
	wolMAC := flag.String("wol-mac", "aa:bb:cc:dd:ee:01", "MAC the magic packet must carry")
	flag.Parse()
	if *wolListen != "" && *downFile != "" {
		go wakeOnPacket(*wolListen, *wolMAC, *downFile)
	}

	ids := strings.Split(*models, ",")
	mux := http.NewServeMux()
	stamp := func(w http.ResponseWriter) { w.Header().Set("X-Upstream", *name) }
	usage := map[string]any{"prompt_tokens": 100, "completion_tokens": 10, "total_tokens": 110}
	timings := map[string]any{"prompt_n": 100, "cache_n": 90, "predicted_n": 10, "predicted_ms": 50.0}

	mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
		stamp(w)
		if *downFile != "" {
			if _, err := os.Stat(*downFile); err == nil {
				http.Error(w, `{"error":{"message":"Loading model"}}`, http.StatusServiceUnavailable)
				return
			}
		}
		writeJSON(w, map[string]string{"status": "ok"})
	})
	mux.HandleFunc("/v1/models", func(w http.ResponseWriter, r *http.Request) {
		stamp(w)
		data := []map[string]any{}
		for _, id := range ids {
			data = append(data, map[string]any{"id": id, "object": "model", "owned_by": *name})
		}
		writeJSON(w, map[string]any{"object": "list", "data": data})
	})
	mux.HandleFunc("/props", func(w http.ResponseWriter, r *http.Request) {
		stamp(w)
		writeJSON(w, map[string]any{"default_generation_settings": map[string]any{"n_ctx": *nCtx}, "total_slots": *slots, "model_path": *name})
	})
	mux.HandleFunc("/v1/chat/completions", func(w http.ResponseWriter, r *http.Request) {
		stamp(w)
		body, _ := io.ReadAll(io.LimitReader(r.Body, 1<<20))
		var req struct {
			Model  string `json:"model"`
			Stream bool   `json:"stream"`
		}
		_ = json.Unmarshal(body, &req)
		time.Sleep(time.Duration(*slow) * time.Millisecond)
		if !req.Stream {
			writeJSON(w, map[string]any{
				"id": "chatcmpl-fake", "object": "chat.completion", "model": req.Model,
				"choices": []map[string]any{{"index": 0, "message": map[string]string{"role": "assistant", "content": "hello from " + *name}, "finish_reason": "stop"}},
				"usage":   usage, "timings": timings,
			})
			return
		}
		w.Header().Set("Content-Type", "text/event-stream")
		w.Header().Set("Cache-Control", "no-cache")
		w.WriteHeader(http.StatusOK)
		fl, _ := w.(http.Flusher)
		flush := func() {
			if fl != nil {
				fl.Flush()
			}
		}
		for i := 1; i <= 5; i++ {
			chunk := map[string]any{"id": "chatcmpl-fake", "object": "chat.completion.chunk", "model": req.Model,
				"choices": []map[string]any{{"index": 0, "delta": map[string]string{"content": fmt.Sprintf("%s chunk %d ", *name, i)}}}}
			b, _ := json.Marshal(chunk)
			fmt.Fprintf(w, "data: %s\n\n", b)
			flush()
			time.Sleep(200 * time.Millisecond)
		}
		final := map[string]any{"id": "chatcmpl-fake", "object": "chat.completion.chunk", "model": req.Model,
			"choices": []map[string]any{}, "usage": usage, "timings": timings}
		b, _ := json.Marshal(final)
		fmt.Fprintf(w, "data: %s\n\n", b)
		flush()
		fmt.Fprint(w, "data: [DONE]\n\n")
	})
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		stamp(w)
		http.Error(w, `{"error":"not found"}`, http.StatusNotFound)
	})

	log.Printf("fakeupstream %s listening on %s models=%v slow=%dms", *name, *listen, ids, *slow)
	srv := &http.Server{Addr: *listen, Handler: mux, ReadHeaderTimeout: 5 * time.Second}
	log.Fatal(srv.ListenAndServe())
}

func writeJSON(w http.ResponseWriter, v any) {
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(v)
}

// wakeOnPacket removes downFile when a magic packet for mac arrives: 6×0xff then the MAC 16 times.
func wakeOnPacket(addr, mac, downFile string) {
	hw, err := net.ParseMAC(mac)
	if err != nil {
		log.Fatalf("wol-mac: %v", err)
	}
	pc, err := net.ListenPacket("udp4", addr)
	if err != nil {
		log.Fatalf("wol-listen: %v", err)
	}
	log.Printf("fakeupstream listening for wake-on-LAN on %s (mac %s)", addr, hw)
	buf := make([]byte, 256)
	for {
		n, _, err := pc.ReadFrom(buf)
		if err != nil {
			return
		}
		if n != 102 {
			continue
		}
		ok := true
		for i := 0; i < 6; i++ {
			ok = ok && buf[i] == 0xff
		}
		for i := 0; i < 16 && ok; i++ {
			for j := 0; j < 6; j++ {
				ok = ok && buf[6+6*i+j] == hw[j]
			}
		}
		if ok {
			log.Printf("magic packet received: waking (removing %s)", downFile)
			_ = os.Remove(downFile)
		}
	}
}