LocalAI/tests/e2e/mock-backend/main.go

package main

import (
	"context"
	"crypto/sha256"
	"encoding/base64"
	"encoding/binary"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"math"
	"net"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/mudler/LocalAI/pkg/grpc/grpcerrors"
	pb "github.com/mudler/LocalAI/pkg/grpc/proto"
	"github.com/mudler/xlog"
	"google.golang.org/grpc"
)

var (
	addr = flag.String("addr", "localhost:50051", "the address to connect to")
)

// MockBackend implements the Backend gRPC service with mocked responses.
// When tools are present but the prompt already contains MCP tool results
// (indicated by the marker from the mock MCP server), it returns a plain
// text response instead of another tool call, letting the MCP loop complete.
type MockBackend struct {
	pb.UnimplementedBackendServer
	quantizationMu      sync.RWMutex
	quantizationOutputs map[string]string
}

var (
	pngFixture = []byte{
		0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, 0x00, 0x00, 0x00, 0x0d,
		0x49, 0x48, 0x44, 0x52, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01,
		0x08, 0x06, 0x00, 0x00, 0x00, 0x1f, 0x15, 0xc4, 0x89, 0x00, 0x00, 0x00,
		0x0d, 0x49, 0x44, 0x41, 0x54, 0x08, 0xd7, 0x63, 0xf8, 0xcf, 0xc0, 0xf0,
		0x1f, 0x00, 0x05, 0x00, 0x01, 0xff, 0x89, 0x99, 0x3d, 0x1d, 0x00, 0x00,
		0x00, 0x00, 0x49, 0x45, 0x4e, 0x44, 0xae, 0x42, 0x60, 0x82,
	}
	videoFixture = []byte("\x00\x00\x00\x18ftypisomMOCK-VIDEO")
	glbFixture   = []byte{'g', 'l', 'T', 'F', 2, 0, 0, 0, 12, 0, 0, 0}
)

type namedFixtureInput struct {
	name  string
	value string
}

type ttsFixtureReference struct {
	Audio string `json:"audio"`
}

// safeLocalFixturePath is the only gate through which fixture inputs may
// reach filesystem APIs. Distributed staging produces short absolute paths;
// inline base64, data URIs, URLs, and long opaque values must remain literals.
func safeLocalFixturePath(value string) (string, bool) {
	if value == "" || len(value) > 4096 || strings.IndexByte(value, 0) >= 0 || isInlineFixtureValue(value) || !filepath.IsAbs(value) {
		return "", false
	}
	clean := filepath.Clean(value)
	if clean != value {
		return "", false
	}
	for _, component := range strings.Split(value, string(filepath.Separator)) {
		if len(component) > 255 {
			return "", false
		}
	}

	current := string(filepath.Separator)
	for _, component := range strings.Split(strings.TrimPrefix(clean, string(filepath.Separator)), string(filepath.Separator)) {
		current = filepath.Join(current, component)
		info, err := os.Lstat(current)
		if err != nil || info.Mode()&os.ModeSymlink != 0 {
			return "", false
		}
	}
	info, err := os.Lstat(clean)
	if err != nil || !info.Mode().IsRegular() {
		return "", false
	}
	return clean, true
}

func isInlineFixtureValue(value string) bool {
	lower := strings.ToLower(value)
	if strings.HasPrefix(lower, "data:") || strings.Contains(value, "://") {
		return true
	}
	// Standard base64 may legitimately begin with '/'. Require a meaningful
	// payload size so short Unix paths such as /tmp/foo are not ambiguous.
	if len(value) < 64 {
		return false
	}
	if _, err := base64.StdEncoding.DecodeString(value); err == nil {
		return true
	}
	_, err := base64.RawStdEncoding.DecodeString(value)
	return err == nil
}

func writeFixture(path string, data []byte) error {
	if path == "" {
		return nil
	}
	if err := os.MkdirAll(filepath.Dir(path), 0750); err != nil {
		return err
	}
	// #nosec G703 -- this test backend writes only the destination allocated by
	// the worker staging layer; the binary conformance suite verifies its root.
	return os.WriteFile(path, data, 0600)
}

func fixtureArtifact(base []byte, markers string) []byte {
	if markers == "" {
		return base
	}
	out := append([]byte(nil), base...)
	out = append(out, []byte("\nMOCK-INPUTS:"+markers+"\n")...)
	return out
}

func fixtureDigest(path string) (string, error) {
	// #nosec G304 -- callers pass only paths accepted by safeLocalFixturePath,
	// including its symlink and optional expected-worker-root checks.
	data, err := os.ReadFile(path)
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("sha256:%x", sha256.Sum256(data)), nil
}

func fixtureInputMarker(input namedFixtureInput) (string, error) {
	if input.value == "" {
		return "", nil
	}
	if path, ok := safeLocalFixturePath(input.value); ok {
		if expectedRoot := os.Getenv("LOCALAI_MOCK_EXPECT_STAGING_ROOT"); expectedRoot != "" {
			realRoot, rootErr := filepath.EvalSymlinks(expectedRoot)
			realPath, pathErr := filepath.EvalSymlinks(path)
			if rootErr != nil || pathErr != nil {
				return "", fmt.Errorf("validating staged %s root", input.name)
			}
			rel, relErr := filepath.Rel(realRoot, realPath)
			if relErr != nil || rel == "." || rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) {
				return "", fmt.Errorf("staged %s path is outside expected worker root", input.name)
			}
		}
		digest, err := fixtureDigest(path)
		if err != nil {
			return "", fmt.Errorf("reading staged %s: %w", input.name, err)
		}
		return input.name + "=" + digest, nil
	}
	return fmt.Sprintf("%s=inline-sha256:%x", input.name, sha256.Sum256([]byte(input.value))), nil
}

func fixtureInputMarkers(inputs ...namedFixtureInput) (string, error) {
	markers := make([]string, 0, len(inputs))
	for _, input := range inputs {
		marker, err := fixtureInputMarker(input)
		if err != nil {
			return "", err
		}
		if marker != "" {
			markers = append(markers, marker)
		}
	}
	return strings.Join(markers, "; "), nil
}

func predictFixtureMarkers(in *pb.PredictOptions) (string, error) {
	inputs := make([]namedFixtureInput, 0, len(in.Images)+len(in.Videos)+len(in.Audios))
	for i, value := range in.Images {
		inputs = append(inputs, namedFixtureInput{name: fmt.Sprintf("image[%d]", i), value: value})
	}
	for i, value := range in.Videos {
		inputs = append(inputs, namedFixtureInput{name: fmt.Sprintf("video[%d]", i), value: value})
	}
	for i, value := range in.Audios {
		inputs = append(inputs, namedFixtureInput{name: fmt.Sprintf("audio[%d]", i), value: value})
	}
	return fixtureInputMarkers(inputs...)
}

func ttsFixtureInputs(in *pb.TTSRequest) ([]namedFixtureInput, error) {
	inputs := []namedFixtureInput{
		{name: "model", value: in.Model},
		{name: "voice", value: in.Voice},
	}
	raw := in.Params["multi_reference_cond"]
	if raw == "" {
		return inputs, nil
	}
	var references []ttsFixtureReference
	if err := json.Unmarshal([]byte(raw), &references); err != nil {
		return nil, fmt.Errorf("decoding TTS fixture references: %w", err)
	}
	for i, reference := range references {
		inputs = append(inputs, namedFixtureInput{
			name:  fmt.Sprintf("reference[%d]", i),
			value: reference.Audio,
		})
	}
	return inputs, nil
}

func fixtureResult(message, markers string, err error) *pb.Result {
	if err != nil {
		return &pb.Result{Message: err.Error(), Success: false}
	}
	if markers != "" {
		message += "; " + markers
	}
	return &pb.Result{Message: message, Success: true}
}

func fixtureInt32(value int) int32 {
	if value > math.MaxInt32 {
		return math.MaxInt32
	}
	if value < math.MinInt32 {
		return math.MinInt32
	}
	// #nosec G115 -- the two guards above prove the conversion is in range.
	return int32(value)
}

// lastLoadParams records the most recent LoadModel parameters so a Predict
// call can echo them back. Used by the path-resolution e2e test, which needs
// to verify that relative draft_model / mmproj / modelfile paths in the YAML
// config arrive at the backend already resolved against the models directory.
// Each backend binary serves a single model, so a single value is enough.
var (
	lastLoadParamsMu sync.RWMutex
	lastLoadParams   *pb.ModelOptions
)

func recordLoadParams(opts *pb.ModelOptions) {
	lastLoadParamsMu.Lock()
	defer lastLoadParamsMu.Unlock()
	lastLoadParams = opts
}

func snapshotLoadParams() *pb.ModelOptions {
	lastLoadParamsMu.RLock()
	defer lastLoadParamsMu.RUnlock()
	return lastLoadParams
}

// checkModelIdentity mirrors the guard the real backends apply (pkg/grpc,
// backend/python/common/model_identity.py, backend/cpp/*/grpc-server.cpp) so
// the distributed e2e suite exercises the #10952 fix rather than only the
// unit tests. Empty on either side means skip, which is why every existing
// spec that sends a bare request struct keeps working.
//
// It takes an interface rather than *pb.PredictOptions because every modality
// request message now carries the field, and the rule must not drift per RPC.
func checkModelIdentity(in interface{ GetModelIdentity() string }) error {
	if in == nil || in.GetModelIdentity() == "" {
		return nil
	}
	opts := snapshotLoadParams()
	if opts == nil || opts.Model == "" || opts.Model == in.GetModelIdentity() {
		return nil
	}
	return grpcerrors.ModelMismatch("mock-backend", opts.Model, in.GetModelIdentity())
}

// promptHasToolResults checks if the prompt contains evidence of prior tool
// execution — specifically the output from the mock MCP server's get_weather tool.
func promptHasToolResults(prompt string) bool {
	return strings.Contains(prompt, "Weather in")
}

func (m *MockBackend) Health(ctx context.Context, in *pb.HealthMessage) (*pb.Reply, error) {
	xlog.Debug("Health check called")
	return &pb.Reply{Message: []byte("OK")}, nil
}

func (m *MockBackend) LoadModel(ctx context.Context, in *pb.ModelOptions) (*pb.Result, error) {
	xlog.Debug("LoadModel called",
		"model", in.Model,
		"modelfile", in.ModelFile,
		"draft_model", in.DraftModel,
		"mmproj", in.MMProj)
	recordLoadParams(in)
	inputs := []namedFixtureInput{
		{name: "model_file", value: in.ModelFile},
		{name: "draft_model", value: in.DraftModel},
		{name: "mmproj", value: in.MMProj},
		{name: "original_config_file", value: in.OriginalConfigFile},
	}
	if _, err := os.Stat(in.ModelFile + ".json"); err == nil {
		inputs = append(inputs, namedFixtureInput{name: "model_companion", value: in.ModelFile + ".json"})
	}
	markers, err := fixtureInputMarkers(inputs...)
	if err != nil {
		return nil, err
	}
	return &pb.Result{
		Message: "Model loaded successfully (mocked) " + markers,
		Success: true,
	}, nil
}

func (m *MockBackend) Predict(ctx context.Context, in *pb.PredictOptions) (*pb.Reply, error) {
	if err := checkModelIdentity(in); err != nil {
		return nil, err
	}
	xlog.Debug("Predict called", "prompt", in.Prompt)
	if strings.Contains(in.Prompt, "MOCK_ERROR") {
		return nil, fmt.Errorf("mock backend predict error: simulated failure")
	}
	if strings.Contains(in.Prompt, "ECHO_FIXTURE_INPUTS") {
		markers, err := predictFixtureMarkers(in)
		if err != nil {
			return nil, err
		}
		return &pb.Reply{Message: []byte(markers), Tokens: 1, PromptTokens: 1}, nil
	}

	// ECHO_LOAD_PARAMS lets path-resolution tests inspect what LoadModel
	// received without adding a new RPC. The reply carries a JSON snapshot
	// of the relevant ModelOptions fields so the test can assert that
	// relative paths from the YAML have been resolved before reaching the
	// backend.
	if strings.Contains(in.Prompt, "ECHO_LOAD_PARAMS") {
		opts := snapshotLoadParams()
		snapshot := map[string]string{}
		if opts != nil {
			snapshot["model"] = opts.Model
			snapshot["model_file"] = opts.ModelFile
			snapshot["draft_model"] = opts.DraftModel
			snapshot["mmproj"] = opts.MMProj
			snapshot["engine_args"] = opts.EngineArgs
			snapshot["original_config_file"] = opts.OriginalConfigFile
			snapshot["fixture_env"] = opts.EnvVars["FIXTURE_ENV"]
		}
		payload, err := json.Marshal(snapshot)
		if err != nil {
			return nil, fmt.Errorf("mock backend echo error: %w", err)
		}
		return &pb.Reply{
			Message:      payload,
			Tokens:       int32(len(snapshot)),
			PromptTokens: 1,
		}, nil
	}

	// ECHO_PREDICT_METADATA lets tests assert exactly what the REST layer
	// forwarded to the backend as gRPC PredictOptions.Metadata (e.g. the
	// chat_template_kwargs blob and the standalone enable_thinking/reasoning_effort
	// keys). The reply carries a JSON snapshot of in.Metadata so an HTTP-level
	// test can pin the request -> gRPC mapping without a new RPC.
	if strings.Contains(in.Prompt, "ECHO_PREDICT_METADATA") {
		payload, err := json.Marshal(in.Metadata)
		if err != nil {
			return nil, fmt.Errorf("mock backend echo metadata error: %w", err)
		}
		return &pb.Reply{
			Message:      payload,
			Tokens:       int32(len(in.Metadata)),
			PromptTokens: 1,
		}, nil
	}

	// ECHO_SERVED_MODEL returns the loaded model file path so router e2e
	// tests can verify which candidate actually served the request without
	// adding a new RPC. The router fans out to a single backend process per
	// candidate, so lastLoadParams.Model is unique per candidate.
	if strings.Contains(in.Prompt, "ECHO_SERVED_MODEL") {
		opts := snapshotLoadParams()
		modelID := ""
		if opts != nil {
			modelID = opts.Model
		}
		return &pb.Reply{
			Message:      []byte("SERVED_MODEL=" + modelID),
			Tokens:       2,
			PromptTokens: 1,
		}, nil
	}

	// Simulate C++ autoparser: tool call via ChatDeltas, empty message
	if strings.Contains(in.Prompt, "AUTOPARSER_TOOL_CALL") {
		toolName := mockToolNameFromRequest(in)
		if toolName == "" {
			toolName = "search_collections"
		}
		return &pb.Reply{
			Message:      []byte{},
			Tokens:       10,
			PromptTokens: 5,
			ChatDeltas: []*pb.ChatDelta{
				{ReasoningContent: "I need to search for information."},
				{
					ToolCalls: []*pb.ToolCallDelta{
						{
							Index:     0,
							Id:        "call_mock_123",
							Name:      toolName,
							Arguments: `{"query":"localai"}`,
						},
					},
				},
			},
		}, nil
	}

	// Simulate C++ autoparser: content via ChatDeltas, empty message
	if strings.Contains(in.Prompt, "AUTOPARSER_CONTENT") {
		return &pb.Reply{
			Message:      []byte{},
			Tokens:       10,
			PromptTokens: 5,
			ChatDeltas: []*pb.ChatDelta{
				{ReasoningContent: "Let me compose a response."},
				{Content: "LocalAI is an open-source AI platform."},
			},
		}, nil
	}

	// Simulate Gemma 4 / thinking model with C++ autoparser:
	// - Message contains the clean content (autoparser extracts it from OAI choices[0].message.content)
	// - ChatDeltas contain both reasoning and content separately
	// This reproduces the bug where Go-side PrependThinkingTokenIfNeeded
	// incorrectly prepends a thinking start token to the clean content,
	// causing the entire response to be classified as unclosed reasoning.
	if strings.Contains(in.Prompt, "AUTOPARSER_THINKING_CONTENT") {
		return &pb.Reply{
			Message:      []byte("I am a helpful AI assistant designed to assist you with a wide range of tasks."),
			Tokens:       20,
			PromptTokens: 50,
			ChatDeltas: []*pb.ChatDelta{
				{
					ReasoningContent: "The user is asking a simple introductory question. I should respond directly.",
					Content:          "I am a helpful AI assistant designed to assist you with a wide range of tasks.",
				},
			},
		}, nil
	}

	// Simulate multiple tool calls in a single response (Go-side JSON parser path).
	if strings.Contains(in.Prompt, "MULTI_TOOL_CALL") {
		return &pb.Reply{
			Message: []byte(`{"name": "get_weather", "arguments": {"location": "Rome"}}
{"name": "get_weather", "arguments": {"location": "Paris"}}`),
			Tokens:       30,
			PromptTokens: 10,
		}, nil
	}
	var response string
	toolName := mockToolNameFromRequest(in)
	if toolName != "" && !promptHasToolResults(in.Prompt) {
		// First call with tools: return a tool call so the MCP loop executes it.
		response = fmt.Sprintf(`{"name": "%s", "arguments": {"location": "San Francisco"}}`, toolName)
	} else if toolName != "" {
		// Subsequent call: tool results already in prompt, return final text.
		response = "Based on the tool results, the weather in San Francisco is sunny, 72°F."
	} else {
		response = "This is a mocked response."
	}
	return &pb.Reply{
		Message:                []byte(response),
		Tokens:                 10,
		PromptTokens:           5,
		TimingPromptProcessing: 0.1,
		TimingTokenGeneration:  0.2,
	}, nil
}

func (m *MockBackend) PredictStream(in *pb.PredictOptions, stream pb.Backend_PredictStreamServer) error {
	if err := checkModelIdentity(in); err != nil {
		return err
	}
	xlog.Debug("PredictStream called", "prompt", in.Prompt)
	if strings.Contains(in.Prompt, "MOCK_ERROR_IMMEDIATE") {
		return fmt.Errorf("mock backend stream error: simulated failure")
	}
	if strings.Contains(in.Prompt, "MOCK_ERROR_MIDSTREAM") {
		for _, r := range "Partial resp" {
			if err := stream.Send(&pb.Reply{Message: []byte(string(r))}); err != nil {
				return err
			}
		}
		return fmt.Errorf("mock backend stream error: simulated mid-stream failure")
	}
	if strings.Contains(in.Prompt, "MOCK_SLOW_STREAM") {
		select {
		case <-time.After(500 * time.Millisecond):
		case <-stream.Context().Done():
			return stream.Context().Err()
		}
	}
	if strings.Contains(in.Prompt, "ECHO_FIXTURE_INPUTS") {
		markers, err := predictFixtureMarkers(in)
		if err != nil {
			return err
		}
		return stream.Send(&pb.Reply{Message: []byte(markers), Tokens: 1, PromptTokens: 1})
	}

	// Simulate C++ autoparser behavior: tool calls delivered via ChatDeltas
	// with empty message (autoparser clears raw message during parsing).
	if strings.Contains(in.Prompt, "AUTOPARSER_TOOL_CALL") {
		toolName := mockToolNameFromRequest(in)
		if toolName == "" {
			toolName = "search_collections"
		}
		// Phase 1: Stream reasoning tokens with empty message (autoparser active)
		reasoning := "I need to search for information."
		for _, r := range reasoning {
			if err := stream.Send(&pb.Reply{
				Message: []byte{}, // autoparser clears raw message
				ChatDeltas: []*pb.ChatDelta{
					{ReasoningContent: string(r)},
				},
			}); err != nil {
				return err
			}
		}
		// Phase 2: Emit tool call via ChatDeltas (no raw message)
		if err := stream.Send(&pb.Reply{
			Message: []byte{}, // autoparser clears raw message
			ChatDeltas: []*pb.ChatDelta{
				{
					ToolCalls: []*pb.ToolCallDelta{
						{
							Index:     0,
							Id:        "call_mock_123",
							Name:      toolName,
							Arguments: `{"query":"localai"}`,
						},
					},
				},
			},
		}); err != nil {
			return err
		}
		return nil
	}

	// Simulate C++ autoparser behavior: content delivered via ChatDeltas
	// with empty message (autoparser clears raw message during parsing).
	if strings.Contains(in.Prompt, "AUTOPARSER_CONTENT") {
		// Phase 1: Stream reasoning via ChatDeltas
		reasoning := "Let me compose a response."
		for _, r := range reasoning {
			if err := stream.Send(&pb.Reply{
				Message: []byte{},
				ChatDeltas: []*pb.ChatDelta{
					{ReasoningContent: string(r)},
				},
			}); err != nil {
				return err
			}
		}
		// Phase 2: Stream content via ChatDeltas (no raw message)
		content := "LocalAI is an open-source AI platform."
		for _, r := range content {
			if err := stream.Send(&pb.Reply{
				Message: []byte{},
				ChatDeltas: []*pb.ChatDelta{
					{Content: string(r)},
				},
			}); err != nil {
				return err
			}
		}
		return nil
	}

	// Simulate tool calls streamed as whole JSON objects (Go-side parser path).
	// Each object is sent as a complete chunk so the incremental parser can
	// detect tool calls mid-stream (unlike char-by-char which only parses after
	// streaming completes).
	if strings.Contains(in.Prompt, "MULTI_TOOL_CALL") {
		chunks := []string{
			`{"name": "get_weather", "arguments": {"location": "Rome"}}`,
			"\n",
			`{"name": "get_weather", "arguments": {"location": "Paris"}}`,
		}
		for i, chunk := range chunks {
			if err := stream.Send(&pb.Reply{
				Message: []byte(chunk),
				Tokens:  int32(i + 1),
			}); err != nil {
				return err
			}
		}
		return nil
	}

	// Simulate single tool call streamed as whole JSON (Go-side parser path).
	if strings.Contains(in.Prompt, "SINGLE_TOOL_CALL") {
		if err := stream.Send(&pb.Reply{
			Message: []byte(`{"name": "get_weather", "arguments": {"location": "San Francisco"}}`),
			Tokens:  1,
		}); err != nil {
			return err
		}
		return nil
	}

	var toStream string
	toolName := mockToolNameFromRequest(in)
	switch {
	case toolName != "" && !promptHasToolResults(in.Prompt):
		toStream = fmt.Sprintf(`{"name": "%s", "arguments": {"location": "San Francisco"}}`, toolName)
	case toolName != "":
		toStream = "Based on the tool results, the weather in San Francisco is sunny, 72°F."
	case strings.Contains(in.Prompt, "MOCK_LEAK_EMAIL"):
		// PII streaming test fixture: emit a response containing an email
		// address so the streaming PII filter has something to mask. The
		// content is split character-by-character below, so the mask
		// must hold across chunk boundaries.
		toStream = "Sure — here it is: alice@example.com is the address."
	default:
		toStream = "This is a mocked streaming response."
	}
	for i, r := range toStream {
		if err := stream.Send(&pb.Reply{
			Message: []byte(string(r)),
			Tokens:  int32(i + 1),
		}); err != nil {
			return err
		}
	}
	return nil
}

// mockToolNameFromRequest returns the first tool name from the request's Tools JSON (same as other endpoints).
func mockToolNameFromRequest(in *pb.PredictOptions) string {
	if in.Tools == "" {
		return ""
	}
	var tools []struct {
		Function struct {
			Name string `json:"name"`
		} `json:"function"`
	}
	if err := json.Unmarshal([]byte(in.Tools), &tools); err != nil || len(tools) == 0 || tools[0].Function.Name == "" {
		return ""
	}
	return tools[0].Function.Name
}

func (m *MockBackend) Embedding(ctx context.Context, in *pb.PredictOptions) (*pb.EmbeddingResult, error) {
	if err := checkModelIdentity(in); err != nil {
		return nil, err
	}
	// The embeddings path ships the text in PredictOptions.Embeddings
	// (see core/backend/embeddings.go), not Prompt; check both so the
	// markers below work however the caller packed the request.
	text := in.Embeddings
	if text == "" {
		text = in.Prompt
	}
	xlog.Debug("Embedding called", "text", text)
	// Deterministic per-token mode for Go-side pooling tests: a prompt
	// carrying the "per-token:" marker yields len(fields) vectors of dim 8
	// with vec[i][j] = (i+1)/(j+2), so endpoint tests can assert exact
	// pooled goldens. The marker may sit mid-prompt (the embeddings
	// messages[] path renders conversations as "<role>: <content>" lines),
	// so match anywhere and tokenize what follows the first occurrence.
	if idx := strings.Index(text, "per-token:"); idx >= 0 {
		fields := strings.Fields(text[idx+len("per-token:"):])
		tokens := len(fields)
		const dim = 8
		flat := make([]float32, 0, tokens*dim)
		for i := 0; i < tokens; i++ {
			for j := 0; j < dim; j++ {
				flat = append(flat, float32(i+1)/float32(j+2))
			}
		}
		return &pb.EmbeddingResult{
			Embeddings:   flat,
			Tokens:       int32(tokens),
			Dim:          dim,
			PromptTokens: int32(tokens),
			Layout:       pb.EmbeddingLayout_EMBEDDING_LAYOUT_PER_TOKEN,
		}, nil
	}
	// Legacy mode: a prompt carrying "no-shape:" simulates a backend built
	// before EmbeddingResult carried shape and layout fields, so tests can
	// assert the fail-closed error when Go-side pooling is requested.
	legacyShape := strings.Contains(text, "no-shape:")
	// Return a mock embedding vector of 768 dimensions
	embeddings := make([]float32, 768)
	for i := range embeddings {
		embeddings[i] = float32(i%100) / 100.0 // Pattern: 0.0, 0.01, 0.02, ..., 0.99, 0.0, ...
	}
	if legacyShape {
		return &pb.EmbeddingResult{Embeddings: embeddings}, nil
	}
	return &pb.EmbeddingResult{
		Embeddings:   embeddings,
		Tokens:       1,
		Dim:          768,
		PromptTokens: 1,
		Layout:       pb.EmbeddingLayout_EMBEDDING_LAYOUT_FINAL,
	}, nil
}

func (m *MockBackend) GenerateImage(ctx context.Context, in *pb.GenerateImageRequest) (*pb.Result, error) {
	if err := checkModelIdentity(in); err != nil {
		return nil, err
	}
	xlog.Debug("GenerateImage called", "prompt", in.PositivePrompt)
	inputs := []namedFixtureInput{
		{name: "negative_prompt", value: in.NegativePrompt},
		{name: "src", value: in.Src},
	}
	for i, ref := range in.RefImages {
		inputs = append(inputs, namedFixtureInput{name: fmt.Sprintf("ref_image[%d]", i), value: ref})
	}
	markers, err := fixtureInputMarkers(inputs...)
	if err == nil {
		err = writeFixture(in.Dst, fixtureArtifact(pngFixture, markers))
	}
	return fixtureResult("Image generated successfully (mocked)", markers, err), nil
}

func (m *MockBackend) GenerateVideo(ctx context.Context, in *pb.GenerateVideoRequest) (*pb.Result, error) {
	if err := checkModelIdentity(in); err != nil {
		return nil, err
	}
	xlog.Debug("GenerateVideo called", "prompt", in.Prompt)
	markers, err := fixtureInputMarkers(
		namedFixtureInput{name: "start_image", value: in.StartImage},
		namedFixtureInput{name: "end_image", value: in.EndImage},
		namedFixtureInput{name: "audio", value: in.Audio},
	)
	if err == nil {
		err = writeFixture(in.Dst, fixtureArtifact(videoFixture, markers))
	}
	return fixtureResult("Video generated successfully (mocked)", markers, err), nil
}

func (m *MockBackend) Generate3D(ctx context.Context, in *pb.Generate3DRequest) (*pb.Result, error) {
	markers, err := fixtureInputMarkers(namedFixtureInput{name: "src", value: in.Src})
	if err == nil {
		err = writeFixture(in.Dst, fixtureArtifact(glbFixture, markers))
	}
	return fixtureResult("3D asset generated successfully (mocked)", markers, err), nil
}

func (m *MockBackend) Animate3D(ctx context.Context, in *pb.Animate3DRequest) (*pb.Result, error) {
	if err := checkModelIdentity(in); err != nil {
		return nil, err
	}
	names := make([]string, 0, len(in.Inputs))
	for name := range in.Inputs {
		names = append(names, name)
	}
	sort.Strings(names)
	inputs := make([]namedFixtureInput, 0, len(names))
	for _, name := range names {
		input := in.Inputs[name]
		if input != nil {
			inputs = append(inputs, namedFixtureInput{name: name, value: input.Data})
		}
	}
	markers, err := fixtureInputMarkers(inputs...)
	if err == nil {
		err = writeFixture(in.Dst, fixtureArtifact(glbFixture, markers))
	}
	return fixtureResult("3D animation generated successfully (mocked)", markers, err), nil
}

func (m *MockBackend) UpscaleImage(ctx context.Context, in *pb.UpscaleImageRequest) (*pb.Result, error) {
	markers, err := fixtureInputMarkers(namedFixtureInput{name: "src", value: in.Src})
	if err == nil {
		err = writeFixture(in.Dst, fixtureArtifact(pngFixture, markers))
	}
	return fixtureResult("Image upscaled successfully (mocked)", markers, err), nil
}

func (m *MockBackend) TTS(ctx context.Context, in *pb.TTSRequest) (*pb.Result, error) {
	if err := checkModelIdentity(in); err != nil {
		return nil, err
	}
	xlog.Debug("TTS called", "text", in.Text)
	inputs, err := ttsFixtureInputs(in)
	markers := ""
	if err == nil {
		markers, err = fixtureInputMarkers(inputs...)
	}
	if err == nil {
		err = writeMinimalWAV(in.Dst)
	}
	return fixtureResult("TTS audio generated successfully (mocked)", markers, err), nil
}

func (m *MockBackend) TTSStream(in *pb.TTSRequest, stream pb.Backend_TTSStreamServer) error {
	if err := checkModelIdentity(in); err != nil {
		return err
	}
	xlog.Debug("TTSStream called", "text", in.Text)
	inputs, err := ttsFixtureInputs(in)
	if err != nil {
		return err
	}
	markers, err := fixtureInputMarkers(inputs...)
	if err != nil {
		return err
	}
	metadata := map[string]any{"sample_rate": ttsSampleRate()}
	if markers != "" {
		metadata["fixture_inputs"] = markers
	}
	message, err := json.Marshal(metadata)
	if err != nil {
		return err
	}
	if err := stream.Send(&pb.Reply{Message: message}); err != nil {
		return err
	}
	pcm := minimalPCM(ttsSampleRate())
	const chunks = 3
	chunkSize := (len(pcm) + chunks - 1) / chunks
	for start := 0; start < len(pcm); start += chunkSize {
		end := min(start+chunkSize, len(pcm))
		chunk := pcm[start:end]
		if err := stream.Send(&pb.Reply{Audio: chunk}); err != nil {
			return err
		}
	}
	return nil
}

func (m *MockBackend) SoundGeneration(ctx context.Context, in *pb.SoundGenerationRequest) (*pb.Result, error) {
	if err := checkModelIdentity(in); err != nil {
		return nil, err
	}
	xlog.Debug("SoundGeneration called",
		"text", in.Text,
		"caption", in.GetCaption(),
		"lyrics", in.GetLyrics(),
		"think", in.GetThink(),
		"bpm", in.GetBpm(),
		"keyscale", in.GetKeyscale(),
		"language", in.GetLanguage(),
		"timesignature", in.GetTimesignature(),
		"instrumental", in.GetInstrumental())
	markers, err := fixtureInputMarkers(
		namedFixtureInput{name: "model", value: in.Model},
		namedFixtureInput{name: "src", value: in.GetSrc()},
	)
	if err == nil {
		err = writeMinimalWAV(in.Dst)
	}
	return fixtureResult("Sound generated successfully (mocked)", markers, err), nil
}

func (m *MockBackend) SoundDetection(ctx context.Context, in *pb.SoundDetectionRequest) (*pb.SoundDetectionResponse, error) {
	if err := checkModelIdentity(in); err != nil {
		return nil, err
	}
	markers, err := fixtureInputMarkers(namedFixtureInput{name: "src", value: in.Src})
	if err != nil {
		return nil, err
	}
	label := "mocked_sound"
	if markers != "" {
		label += "; " + markers
	}
	return &pb.SoundDetectionResponse{Detections: []*pb.SoundClass{{
		Label: label,
		Score: 0.99,
		Index: 1,
	}}}, nil
}

// ttsSampleRate returns the sample rate to use for TTS output, configurable
// via the MOCK_TTS_SAMPLE_RATE environment variable (default 16000).
func ttsSampleRate() int {
	if s := os.Getenv("MOCK_TTS_SAMPLE_RATE"); s != "" {
		if v, err := strconv.Atoi(s); err == nil && v > 0 {
			return v
		}
	}
	return 16000
}

// writeMinimalWAV writes a WAV file containing a 440Hz sine wave (0.5s)
// so that tests can verify audio integrity end-to-end. The sample rate
// is configurable via MOCK_TTS_SAMPLE_RATE to test rate mismatch bugs.
func writeMinimalWAV(path string) error {
	if path == "" {
		return nil
	}
	sampleRate := ttsSampleRate()
	const numChannels = 1
	const bitsPerSample = 16
	pcm := minimalPCM(sampleRate)
	dataSize := len(pcm)
	const headerLen = 44
	if err := os.MkdirAll(filepath.Dir(path), 0750); err != nil {
		return err
	}
	// #nosec G304 -- path is the output allocated by the worker staging layer;
	// the binary conformance suite verifies it is under the worker root.
	f, err := os.OpenFile(path, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0600)
	if err != nil {
		return err
	}
	defer f.Close()
	// RIFF header
	_, _ = f.Write([]byte("RIFF"))
	_ = binary.Write(f, binary.LittleEndian, uint32(headerLen-8+dataSize))
	_, _ = f.Write([]byte("WAVE"))
	// fmt chunk
	_, _ = f.Write([]byte("fmt "))
	_ = binary.Write(f, binary.LittleEndian, uint32(16))
	_ = binary.Write(f, binary.LittleEndian, uint16(1))
	_ = binary.Write(f, binary.LittleEndian, uint16(numChannels))
	_ = binary.Write(f, binary.LittleEndian, uint32(sampleRate))
	_ = binary.Write(f, binary.LittleEndian, uint32(sampleRate*numChannels*(bitsPerSample/8)))
	_ = binary.Write(f, binary.LittleEndian, uint16(numChannels*(bitsPerSample/8)))
	_ = binary.Write(f, binary.LittleEndian, uint16(bitsPerSample))
	// data chunk — 440Hz sine wave
	_, _ = f.Write([]byte("data"))
	// #nosec G115 -- minimalPCM is half a second at the bounded fixture sample
	// rate, so its byte length is well below the WAV uint32 limit.
	_ = binary.Write(f, binary.LittleEndian, uint32(dataSize))
	_, err = f.Write(pcm)
	return err
}

func minimalPCM(sampleRate int) []byte {
	const freq = 440.0
	const durationSec = 0.5
	numSamples := int(float64(sampleRate) * durationSec)
	pcm := make([]byte, numSamples*2)
	for i := range numSamples {
		t := float64(i) / float64(sampleRate)
		sample := int16(math.MaxInt16 / 2 * math.Sin(2*math.Pi*freq*t))
		// #nosec G115 -- WAV PCM stores the signed int16 bit pattern as two
		// little-endian bytes; this conversion intentionally preserves the bits.
		binary.LittleEndian.PutUint16(pcm[i*2:], uint16(sample))
	}
	return pcm
}

func (m *MockBackend) AudioTranscription(ctx context.Context, in *pb.TranscriptRequest) (*pb.TranscriptResult, error) {
	if err := checkModelIdentity(in); err != nil {
		return nil, err
	}
	dst := in.GetDst()
	wavSR := 0
	dataLen := 0
	rms := 0.0

	marker, err := fixtureInputMarker(namedFixtureInput{name: "audio", value: dst})
	if err != nil {
		return nil, err
	}
	if path, ok := safeLocalFixturePath(dst); ok {
		// #nosec G304 -- safeLocalFixturePath rejects non-absolute, unclean,
		// non-regular and symlinked paths before this fixture read.
		if data, readErr := os.ReadFile(path); readErr == nil {
			if len(data) >= 44 {
				wavSR = int(binary.LittleEndian.Uint32(data[24:28]))
				dataLen = int(binary.LittleEndian.Uint32(data[40:44]))

				// Compute RMS of the PCM payload (16-bit LE samples)
				pcm := data[44:]
				var sumSq float64
				nSamples := len(pcm) / 2
				for i := range nSamples {
					s := int16(pcm[2*i]) | int16(pcm[2*i+1])<<8
					v := float64(s)
					sumSq += v * v
				}
				if nSamples > 0 {
					rms = math.Sqrt(sumSq / float64(nSamples))
				}
			}
		}
	}

	xlog.Debug("AudioTranscription called", "input", marker, "wav_sample_rate", wavSR, "data_len", dataLen, "rms", rms)

	text := fmt.Sprintf("transcribed: rms=%.1f samples=%d sr=%d", rms, dataLen/2, wavSR)
	if marker != "" {
		text += "; " + marker
	}
	return &pb.TranscriptResult{
		Text: text,
		Segments: []*pb.TranscriptSegment{
			{
				Id:     0,
				Start:  0,
				End:    3000,
				Text:   text,
				Tokens: []int32{1, 2, 3, 4, 5, 6},
			},
		},
	}, nil
}

func (m *MockBackend) AudioTranscriptionStream(in *pb.TranscriptRequest, stream pb.Backend_AudioTranscriptionStreamServer) error {
	result, err := m.AudioTranscription(stream.Context(), in)
	if err != nil {
		return err
	}
	if err := stream.Send(&pb.TranscriptStreamResponse{Delta: result.Text}); err != nil {
		return err
	}
	return stream.Send(&pb.TranscriptStreamResponse{FinalResult: result})
}

func (m *MockBackend) TokenizeString(ctx context.Context, in *pb.PredictOptions) (*pb.TokenizationResponse, error) {
	if err := checkModelIdentity(in); err != nil {
		return nil, err
	}
	xlog.Debug("TokenizeString called", "prompt_len", len(in.Prompt))
	// Approximate BPE: ~4 chars/token, minimum 1. Realistic enough for the
	// router's fitMessages to exercise the budget/rune-pretrim path with
	// recognisable counts that scale with input size.
	n := max((len(in.Prompt)+3)/4, 1)
	tokens := make([]int32, n)
	for i := range tokens {
		tokens[i] = int32(i + 1)
	}
	return &pb.TokenizationResponse{
		Length: int32(n),
		Tokens: tokens,
	}, nil
}

// Score implements deterministic marker-driven ranking for router e2e
// tests. The Score RPC receives the full rendered routing prompt (system
// prompt + chat envelope + user turn), and the system prompt by construction
// lists every policy label — so any keyword-against-prompt heuristic would
// match every candidate. Instead we look for an explicit `ROUTE_HINT=<label>`
// marker, which only appears when a test deliberately places one in a user
// message. The candidate whose extracted label equals the hint gets a large
// log-prob boost; all others stay at the base. With no hint, every candidate
// scores equally, softmax is uniform, and (with a sensible activation
// threshold) the router falls back.
func (m *MockBackend) Score(ctx context.Context, in *pb.ScoreRequest) (*pb.ScoreResponse, error) {
	if err := checkModelIdentity(in); err != nil {
		return nil, err
	}
	xlog.Debug("Score called", "candidates", len(in.Candidates))
	hint := extractRouteHint(in.Prompt)
	out := &pb.ScoreResponse{Candidates: make([]*pb.CandidateScore, len(in.Candidates))}
	for i, c := range in.Candidates {
		label := extractRouteLabel(c)
		// Base -5 (softmax ≈ 0.003), hint match +5 → 0 (softmax ≈ 0.99).
		logProb := -5.0
		if hint != "" && label == hint {
			logProb = 0.0
		}
		// num_tokens matches TokenizeString's heuristic so per-token mean
		// log-prob consumers see consistent values.
		nTok := max((len(c)+3)/4, 1)
		out.Candidates[i] = &pb.CandidateScore{
			LogProb:                 logProb,
			NumTokens:               int32(nTok),
			LengthNormalizedLogProb: logProb / float64(nTok),
		}
	}
	return out, nil
}

// extractRouteHint returns the label after the LAST occurrence of
// `ROUTE_HINT=` in the prompt, terminated by whitespace or end-of-string.
// Using the last occurrence makes the marker stable across long
// conversations: the *newest* user message's hint wins, mirroring how the
// router's fitMessages keeps the newest turn whole.
func extractRouteHint(prompt string) string {
	const key = "ROUTE_HINT="
	i := strings.LastIndex(prompt, key)
	if i < 0 {
		return ""
	}
	rest := prompt[i+len(key):]
	end := strings.IndexAny(rest, " \t\r\n<")
	if end < 0 {
		return rest
	}
	return rest[:end]
}

// extractRouteLabel returns the label inside `{"route": "<label>"}`. Returns
// "" on any shape it doesn't recognise — the caller treats that as a no-match.
func extractRouteLabel(candidate string) string {
	_, rest, ok := strings.Cut(candidate, `"route"`)
	if !ok {
		return ""
	}
	_, rest, ok = strings.Cut(rest, `"`)
	if !ok {
		return ""
	}
	label, _, ok := strings.Cut(rest, `"`)
	if !ok {
		return ""
	}
	return label
}

func (m *MockBackend) Detokenize(ctx context.Context, in *pb.DetokenizeRequest) (*pb.DetokenizeResponse, error) {
	xlog.Debug("Detokenize called", "tokens", in.Tokens)
	parts := make([]string, len(in.Tokens))
	for i, t := range in.Tokens {
		parts[i] = strconv.Itoa(int(t))
	}
	return &pb.DetokenizeResponse{
		Content: "detokenized: " + strings.Join(parts, " "),
	}, nil
}

func (m *MockBackend) Status(ctx context.Context, in *pb.HealthMessage) (*pb.StatusResponse, error) {
	xlog.Debug("Status called")
	return &pb.StatusResponse{
		State: pb.StatusResponse_READY,
		Memory: &pb.MemoryUsageData{
			Total: 1024 * 1024 * 100, // 100MB
			Breakdown: map[string]uint64{
				"mock": 1024 * 1024 * 50,
			},
		},
	}, nil
}

func (m *MockBackend) Detect(ctx context.Context, in *pb.DetectOptions) (*pb.DetectResponse, error) {
	if err := checkModelIdentity(in); err != nil {
		return nil, err
	}
	if _, err := fixtureInputMarkers(namedFixtureInput{name: "src", value: in.Src}); err != nil {
		return nil, err
	}
	xlog.Debug("Detect called", "src", in.Src)
	return &pb.DetectResponse{
		Detections: []*pb.Detection{
			{
				X:          10.0,
				Y:          20.0,
				Width:      100.0,
				Height:     200.0,
				Confidence: 0.95,
				ClassName:  "mocked_object",
			},
		},
	}, nil
}

func (m *MockBackend) Depth(ctx context.Context, in *pb.DepthRequest) (*pb.DepthResponse, error) {
	if err := checkModelIdentity(in); err != nil {
		return nil, err
	}
	markers, err := fixtureInputMarkers(namedFixtureInput{name: "src", value: in.Src})
	if err != nil {
		return nil, err
	}
	result := &pb.DepthResponse{
		Width: 2, Height: 1, Depth: []float32{1.25, 2.5},
		Confidence: []float32{0.9, 0.8}, Sky: []float32{0, 1},
		Extrinsics: []float32{1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0},
		Intrinsics: []float32{1, 0, 0, 0, 1, 0, 0, 0, 1}, IsMetric: true,
	}
	if in.Dst != "" {
		exportPath := filepath.Join(in.Dst, "nested", "depth.txt")
		if err := writeFixture(exportPath, []byte(markers)); err != nil {
			return nil, err
		}
		result.ExportPaths = []string{exportPath}
	}
	return result, nil
}

func (m *MockBackend) FaceAnalyze(ctx context.Context, in *pb.FaceAnalyzeRequest) (*pb.FaceAnalyzeResponse, error) {
	if err := checkModelIdentity(in); err != nil {
		return nil, err
	}
	return &pb.FaceAnalyzeResponse{Faces: []*pb.FaceAnalysis{{
		Region: &pb.FacialArea{X: 1, Y: 2, W: 3, H: 4}, FaceConfidence: 0.98,
		Age: 34, DominantGender: "Woman", Gender: map[string]float32{"Woman": 0.9},
		DominantEmotion: "happy", Emotion: map[string]float32{"happy": 0.8},
	}}}, nil
}

func (m *MockBackend) FaceVerify(ctx context.Context, in *pb.FaceVerifyRequest) (*pb.FaceVerifyResponse, error) {
	if err := checkModelIdentity(in); err != nil {
		return nil, err
	}
	return &pb.FaceVerifyResponse{
		Verified: true, Distance: 0.05, Threshold: 0.25, Confidence: 95,
		Model: "mock-face", Img1Area: &pb.FacialArea{X: 1, Y: 2, W: 3, H: 4},
		Img2Area: &pb.FacialArea{X: 5, Y: 6, W: 7, H: 8},
	}, nil
}

func (m *MockBackend) VoiceAnalyze(ctx context.Context, in *pb.VoiceAnalyzeRequest) (*pb.VoiceAnalyzeResponse, error) {
	if err := checkModelIdentity(in); err != nil {
		return nil, err
	}
	if _, err := fixtureInputMarkers(namedFixtureInput{name: "audio", value: in.Audio}); err != nil {
		return nil, err
	}
	return &pb.VoiceAnalyzeResponse{Segments: []*pb.VoiceAnalysis{{
		Start: 0, End: 1, Age: 42, DominantGender: "female",
		Gender: map[string]float32{"female": 0.95}, DominantEmotion: "neutral",
		Emotion: map[string]float32{"neutral": 0.9},
	}}}, nil
}

func (m *MockBackend) TokenClassify(ctx context.Context, in *pb.TokenClassifyRequest) (*pb.TokenClassifyResponse, error) {
	if err := checkModelIdentity(in); err != nil {
		return nil, err
	}
	const entity = "Alice"
	start := strings.Index(in.Text, entity)
	if start < 0 {
		return &pb.TokenClassifyResponse{}, nil
	}
	return &pb.TokenClassifyResponse{Entities: []*pb.TokenClassifyEntity{{
		EntityGroup: "PER", Start: fixtureInt32(start), End: fixtureInt32(start + len(entity)), Score: 0.99, Text: entity,
	}}}, nil
}

func (m *MockBackend) StoresSet(ctx context.Context, in *pb.StoresSetOptions) (*pb.Result, error) {
	xlog.Debug("StoresSet called", "keys", len(in.Keys))
	return &pb.Result{
		Message: "Keys set successfully (mocked)",
		Success: true,
	}, nil
}

func (m *MockBackend) StoresDelete(ctx context.Context, in *pb.StoresDeleteOptions) (*pb.Result, error) {
	xlog.Debug("StoresDelete called", "keys", len(in.Keys))
	return &pb.Result{
		Message: "Keys deleted successfully (mocked)",
		Success: true,
	}, nil
}

func (m *MockBackend) StoresGet(ctx context.Context, in *pb.StoresGetOptions) (*pb.StoresGetResult, error) {
	xlog.Debug("StoresGet called", "keys", len(in.Keys))
	// Return mock keys and values
	keys := make([]*pb.StoresKey, len(in.Keys))
	values := make([]*pb.StoresValue, len(in.Keys))
	for i := range in.Keys {
		keys[i] = in.Keys[i]
		values[i] = &pb.StoresValue{
			Bytes: []byte(fmt.Sprintf("mocked_value_%d", i)),
		}
	}
	return &pb.StoresGetResult{
		Keys:   keys,
		Values: values,
	}, nil
}

func (m *MockBackend) StoresFind(ctx context.Context, in *pb.StoresFindOptions) (*pb.StoresFindResult, error) {
	xlog.Debug("StoresFind called", "topK", in.TopK)
	// Return mock similar keys
	keys := []*pb.StoresKey{
		{Floats: []float32{0.1, 0.2, 0.3}},
		{Floats: []float32{0.4, 0.5, 0.6}},
	}
	values := []*pb.StoresValue{
		{Bytes: []byte("mocked_value_1")},
		{Bytes: []byte("mocked_value_2")},
	}
	similarities := []float32{0.95, 0.85}
	return &pb.StoresFindResult{
		Keys:         keys,
		Values:       values,
		Similarities: similarities,
	}, nil
}

func (m *MockBackend) Rerank(ctx context.Context, in *pb.RerankRequest) (*pb.RerankResult, error) {
	if err := checkModelIdentity(in); err != nil {
		return nil, err
	}
	xlog.Debug("Rerank called", "query", in.Query, "documents", len(in.Documents))
	// Return mock reranking results
	results := make([]*pb.DocumentResult, len(in.Documents))
	for i, doc := range in.Documents {
		results[i] = &pb.DocumentResult{
			Index:          int32(i),
			Text:           doc,
			RelevanceScore: 0.9 - float32(i)*0.1, // Decreasing scores
		}
	}
	return &pb.RerankResult{
		Usage: &pb.Usage{
			TotalTokens:  int32(len(in.Documents) * 10),
			PromptTokens: int32(len(in.Documents) * 10),
		},
		Results: results,
	}, nil
}

func (m *MockBackend) GetMetrics(ctx context.Context, in *pb.MetricsRequest) (*pb.MetricsResponse, error) {
	xlog.Debug("GetMetrics called")
	return &pb.MetricsResponse{
		SlotId:                0,
		PromptJsonForSlot:     `{"prompt":"mocked"}`,
		TokensPerSecond:       10.0,
		TokensGenerated:       100,
		PromptTokensProcessed: 50,
	}, nil
}

func (m *MockBackend) VAD(ctx context.Context, in *pb.VADRequest) (*pb.VADResponse, error) {
	if err := checkModelIdentity(in); err != nil {
		return nil, err
	}
	// Compute RMS of the received float32 audio to decide whether speech is present.
	var sumSq float64
	for _, s := range in.Audio {
		v := float64(s)
		sumSq += v * v
	}
	rms := 0.0
	if len(in.Audio) > 0 {
		rms = math.Sqrt(sumSq / float64(len(in.Audio)))
	}
	xlog.Debug("VAD called", "audio_length", len(in.Audio), "rms", rms)

	// If audio is near-silence, return no segments (no speech detected).
	if rms < 0.001 {
		return &pb.VADResponse{}, nil
	}

	// Audio has signal — return a single segment covering the duration.
	duration := float64(len(in.Audio)) / 16000.0
	return &pb.VADResponse{
		Segments: []*pb.VADSegment{
			{
				Start: 0.0,
				End:   float32(duration),
			},
		},
	}, nil
}

// Diarize returns a deterministic two-speaker layout that exercises the
// HTTP layer's normalisation: raw labels "5" and "2" should become
// SPEAKER_00 and SPEAKER_01 in first-seen order, the SPEAKER_00 totals
// should reflect two segments (1.0s + 1.5s = 2.5s), and IncludeText must
// gate the per-segment Text field.
func (m *MockBackend) Diarize(ctx context.Context, in *pb.DiarizeRequest) (*pb.DiarizeResponse, error) {
	if err := checkModelIdentity(in); err != nil {
		return nil, err
	}
	if _, err := fixtureInputMarkers(namedFixtureInput{name: "audio", value: in.Dst}); err != nil {
		return nil, err
	}
	xlog.Debug("Diarize called",
		"dst", in.Dst,
		"num_speakers", in.NumSpeakers,
		"include_text", in.IncludeText)

	seg := func(start, end float32, speaker, text string) *pb.DiarizeSegment {
		out := &pb.DiarizeSegment{Start: start, End: end, Speaker: speaker}
		if in.IncludeText {
			out.Text = text
		}
		return out
	}
	return &pb.DiarizeResponse{
		Segments: []*pb.DiarizeSegment{
			seg(0.0, 1.0, "5", "hello there"),
			seg(1.0, 2.0, "2", "general kenobi"),
			seg(2.0, 3.5, "5", "you are a bold one"),
		},
		NumSpeakers: 2,
		Duration:    3.5,
		Language:    in.Language,
	}, nil
}

func (m *MockBackend) AudioEncode(ctx context.Context, in *pb.AudioEncodeRequest) (*pb.AudioEncodeResult, error) {
	xlog.Debug("AudioEncode called", "pcm_len", len(in.PcmData), "sample_rate", in.SampleRate)
	// Return a single mock Opus frame per 960-sample chunk (20ms at 48kHz).
	numSamples := len(in.PcmData) / 2 // 16-bit samples
	frameSize := 960
	var frames [][]byte
	for offset := 0; offset+frameSize <= numSamples; offset += frameSize {
		// Minimal mock frame — just enough bytes to be non-empty.
		frames = append(frames, []byte{0xFC, 0xFF, 0xFE})
	}
	return &pb.AudioEncodeResult{
		Frames:          frames,
		SampleRate:      48000,
		SamplesPerFrame: int32(frameSize),
	}, nil
}

func (m *MockBackend) AudioDecode(ctx context.Context, in *pb.AudioDecodeRequest) (*pb.AudioDecodeResult, error) {
	xlog.Debug("AudioDecode called", "frames", len(in.Frames))
	// Return silent PCM (960 samples per frame at 48kHz, 16-bit LE).
	samplesPerFrame := 960
	totalSamples := len(in.Frames) * samplesPerFrame
	pcm := make([]byte, totalSamples*2)
	return &pb.AudioDecodeResult{
		PcmData:         pcm,
		SampleRate:      48000,
		SamplesPerFrame: int32(samplesPerFrame),
	}, nil
}

func (m *MockBackend) AudioTransform(ctx context.Context, in *pb.AudioTransformRequest) (*pb.AudioTransformResult, error) {
	if err := checkModelIdentity(in); err != nil {
		return nil, err
	}
	if _, err := fixtureInputMarker(namedFixtureInput{name: "audio", value: in.AudioPath}); err != nil {
		return nil, err
	}
	var input []byte
	if path, ok := safeLocalFixturePath(in.AudioPath); ok {
		var err error
		// #nosec G304 -- safeLocalFixturePath rejects non-absolute, unclean,
		// non-regular and symlinked paths before this fixture read.
		input, err = os.ReadFile(path)
		if err != nil {
			return nil, fmt.Errorf("reading staged audio transform input: %w", err)
		}
	} else {
		input = []byte(in.AudioPath)
	}
	referenceMarker, err := fixtureInputMarker(namedFixtureInput{name: "reference", value: in.ReferencePath})
	if err != nil {
		return nil, err
	}
	if err := writeFixture(in.Dst, fixtureArtifact(input, referenceMarker)); err != nil {
		return nil, fmt.Errorf("writing audio transform output: %w", err)
	}
	samples := 0
	if len(input) > 44 {
		samples = (len(input) - 44) / 2
	}
	return &pb.AudioTransformResult{
		Dst:               in.Dst,
		SampleRate:        fixtureInt32(ttsSampleRate()),
		Samples:           fixtureInt32(samples),
		ReferenceProvided: in.ReferencePath != "",
	}, nil
}

func (m *MockBackend) AudioTransformStream(stream pb.Backend_AudioTransformStreamServer) error {
	configured := false
	var frameIndex int64
	for {
		request, err := stream.Recv()
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return err
		}
		if request.GetConfig() != nil {
			configured = true
			continue
		}
		frame := request.GetFrame()
		if frame == nil {
			continue
		}
		if !configured {
			return fmt.Errorf("audio transform stream frame received before config")
		}
		if err := stream.Send(&pb.AudioTransformFrameResponse{
			Pcm:        append([]byte(nil), frame.AudioPcm...),
			FrameIndex: frameIndex,
		}); err != nil {
			return err
		}
		frameIndex++
	}
}

func (m *MockBackend) ModelMetadata(ctx context.Context, in *pb.ModelOptions) (*pb.ModelMetadataResponse, error) {
	xlog.Debug("ModelMetadata called", "model", in.Model)
	return &pb.ModelMetadataResponse{
		SupportsThinking: false,
		RenderedTemplate: "",
	}, nil
}

// voiceEmbedFromWAV reads a 16-bit LE mono WAV and returns a 2-d speaker
// embedding derived from the signed DC offset of the samples. A positive DC
// bias maps to one orthogonal unit vector, a negative bias to the other, so
// e2e tests can deterministically simulate two distinct "speakers" that
// survive resampling (DC is sample-rate independent). Near-zero DC maps to a
// neutral vector equidistant from both. Returns nil for unreadable audio.
func voiceEmbedFromWAV(path string) []float32 {
	validated, ok := safeLocalFixturePath(path)
	if !ok {
		return nil
	}
	// #nosec G304 -- safeLocalFixturePath rejects non-absolute, unclean,
	// non-regular and symlinked paths before this fixture read.
	data, err := os.ReadFile(validated)
	if err != nil || len(data) < 44 {
		return nil
	}
	pcm := data[44:]
	n := len(pcm) / 2
	if n == 0 {
		return nil
	}
	var sum float64
	for i := 0; i < n; i++ {
		s := int16(pcm[2*i]) | int16(pcm[2*i+1])<<8
		sum += float64(s)
	}
	mean := sum / float64(n)
	switch {
	case mean > 500:
		return []float32{1, 0}
	case mean < -500:
		return []float32{0, 1}
	default:
		return []float32{0.7071, 0.7071}
	}
}

// VoiceEmbed returns a deterministic 2-d speaker embedding for the audio clip.
// See voiceEmbedFromWAV for the (test-only) DC-offset discrimination scheme.
func (m *MockBackend) VoiceEmbed(ctx context.Context, in *pb.VoiceEmbedRequest) (*pb.VoiceEmbedResponse, error) {
	if err := checkModelIdentity(in); err != nil {
		return nil, err
	}
	if _, err := fixtureInputMarkers(namedFixtureInput{name: "audio", value: in.Audio}); err != nil {
		return nil, err
	}
	emb := voiceEmbedFromWAV(in.GetAudio())
	xlog.Debug("VoiceEmbed called", "audio", in.GetAudio(), "embedding", emb)
	if len(emb) == 0 {
		return &pb.VoiceEmbedResponse{}, nil
	}
	return &pb.VoiceEmbedResponse{Embedding: emb, Model: "mock-speaker"}, nil
}

// VoiceVerify compares two clips by cosine distance over their mock embeddings.
func (m *MockBackend) VoiceVerify(ctx context.Context, in *pb.VoiceVerifyRequest) (*pb.VoiceVerifyResponse, error) {
	if err := checkModelIdentity(in); err != nil {
		return nil, err
	}
	if _, err := fixtureInputMarkers(
		namedFixtureInput{name: "audio1", value: in.Audio1},
		namedFixtureInput{name: "audio2", value: in.Audio2},
	); err != nil {
		return nil, err
	}
	a := voiceEmbedFromWAV(in.GetAudio1())
	b := voiceEmbedFromWAV(in.GetAudio2())
	dist := float32(1)
	if len(a) == 2 && len(b) == 2 {
		dist = 1 - (a[0]*b[0] + a[1]*b[1]) // both unit vectors
	}
	threshold := in.GetThreshold()
	if threshold == 0 {
		threshold = 0.25
	}
	xlog.Debug("VoiceVerify called", "distance", dist, "threshold", threshold)
	return &pb.VoiceVerifyResponse{
		Verified:  dist <= threshold,
		Distance:  dist,
		Threshold: threshold,
		Model:     "mock-speaker",
	}, nil
}

func (m *MockBackend) ExportModel(ctx context.Context, in *pb.ExportModelRequest) (*pb.Result, error) {
	markers, err := fixtureInputMarkers(
		namedFixtureInput{name: "checkpoint", value: in.CheckpointPath},
		namedFixtureInput{name: "model", value: in.Model},
	)
	if err == nil && in.OutputPath != "" {
		err = writeFixture(filepath.Join(in.OutputPath, "nested", "weights.bin"), []byte("MOCK-EXPORTED-WEIGHTS\n"))
	}
	if err == nil && in.OutputPath != "" {
		err = writeFixture(filepath.Join(in.OutputPath, "nested", "config.json"), []byte("{\"mock\":true}\n"))
	}
	return fixtureResult("Model exported successfully (mocked)", markers, err), nil
}

func (m *MockBackend) StartQuantization(ctx context.Context, in *pb.QuantizationRequest) (*pb.QuantizationJobResult, error) {
	marker, err := fixtureInputMarker(namedFixtureInput{name: "model", value: in.Model})
	if err != nil {
		return &pb.QuantizationJobResult{JobId: in.JobId, Success: false, Message: err.Error()}, nil
	}

	output := ""
	if in.OutputDir != "" {
		output = filepath.Join(in.OutputDir, "nested", in.JobId+".gguf")
		if err := writeFixture(output, []byte("MOCK-GGUF:"+in.QuantizationType+"\n")); err != nil {
			return &pb.QuantizationJobResult{JobId: in.JobId, Success: false, Message: err.Error()}, nil
		}
	}
	m.quantizationMu.Lock()
	if m.quantizationOutputs == nil {
		m.quantizationOutputs = map[string]string{}
	}
	m.quantizationOutputs[in.JobId] = output
	m.quantizationMu.Unlock()

	message := "Quantization started successfully (mocked)"
	if marker != "" {
		message += "; " + marker
	}
	return &pb.QuantizationJobResult{JobId: in.JobId, Success: true, Message: message}, nil
}

func (m *MockBackend) QuantizationProgress(in *pb.QuantizationProgressRequest, stream pb.Backend_QuantizationProgressServer) error {
	m.quantizationMu.RLock()
	output, ok := m.quantizationOutputs[in.JobId]
	m.quantizationMu.RUnlock()
	if !ok {
		return stream.Send(&pb.QuantizationProgressUpdate{
			JobId:   in.JobId,
			Status:  "failed",
			Message: "unknown mock quantization job",
		})
	}
	return stream.Send(&pb.QuantizationProgressUpdate{
		JobId:           in.JobId,
		ProgressPercent: 100,
		Status:          "completed",
		Message:         "Quantization completed successfully (mocked)",
		OutputFile:      output,
	})
}

func (m *MockBackend) StopQuantization(_ context.Context, in *pb.QuantizationStopRequest) (*pb.Result, error) {
	m.quantizationMu.Lock()
	defer m.quantizationMu.Unlock()
	if _, ok := m.quantizationOutputs[in.JobId]; !ok {
		return &pb.Result{Success: false, Message: "unknown mock quantization job"}, nil
	}
	delete(m.quantizationOutputs, in.JobId)
	return &pb.Result{Success: true, Message: "Quantization stopped successfully (mocked)"}, nil
}

func main() {
	xlog.SetLogger(xlog.NewLogger(xlog.LogLevel(os.Getenv("LOCALAI_LOG_LEVEL")), os.Getenv("LOCALAI_LOG_FORMAT")))

	flag.Parse()

	lis, err := net.Listen("tcp", *addr)
	if err != nil {
		log.Fatalf("failed to listen: %v", err)
	}

	s := grpc.NewServer(
		grpc.MaxRecvMsgSize(50*1024*1024), // 50MB
		grpc.MaxSendMsgSize(50*1024*1024), // 50MB
	)
	pb.RegisterBackendServer(s, &MockBackend{})

	xlog.Info("Mock gRPC Server listening", "address", lis.Addr())
	if err := s.Serve(lis); err != nil {
		log.Fatalf("failed to serve: %v", err)
	}
}