zero/internal/provideronboarding/localruntime.go

package provideronboarding

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/Gitlawb/zero/internal/providercatalog"
)

// LocalRuntime describes a local, OpenAI-compatible model server that ZERO can
// adopt with no API key. These are the runtimes the launch audience is most
// likely to already be running (Ollama, LM Studio), so the first-run wizard
// auto-detects them on their default ports and offers them with no key step.
type LocalRuntime struct {
	CatalogID    string
	Name         string
	BaseURL      string
	DefaultModel string
	RequiresKey  bool
}

// DetectedLocalRuntime is a LocalRuntime that has been probed. Reachable reports
// whether its default endpoint answered the probe; Models is the list of model
// ids the runtime advertised (best-effort, empty when the probe could not parse
// a model list).
type DetectedLocalRuntime struct {
	LocalRuntime
	Reachable bool
	Models    []string
}

// LocalDetectOptions configures DetectLocalRuntimes. All fields are optional;
// the zero value probes the built-in candidates with a short default timeout
// using http.DefaultClient. Tests inject HTTPClient and Candidates.
type LocalDetectOptions struct {
	HTTPClient *http.Client
	Timeout    time.Duration
	Candidates []LocalRuntime
}

const defaultLocalProbeTimeout = 400 * time.Millisecond

// LocalRuntimeCandidates returns the built-in local-runtime candidates derived
// from the provider catalog. Only catalog entries flagged Local are returned, so
// the default ports stay the single source of truth (catalog.go).
func LocalRuntimeCandidates() []LocalRuntime {
	candidates := make([]LocalRuntime, 0, 2)
	for _, descriptor := range providercatalog.All() {
		if !descriptor.Local {
			continue
		}
		candidates = append(candidates, LocalRuntime{
			CatalogID:    descriptor.ID,
			Name:         descriptor.Name,
			BaseURL:      descriptor.DefaultBaseURL,
			DefaultModel: descriptor.DefaultModel,
			RequiresKey:  descriptor.RequiresAuth,
		})
	}
	return candidates
}

// DetectLocalRuntimes probes each candidate's default endpoint and returns only
// the runtimes that answered successfully. A runtime is "reachable" when its
// OpenAI-compatible models endpoint returns a 2xx; transport errors (closed
// port), timeouts, and non-2xx responses are treated as "not running" and the
// candidate is dropped. Detection never errors — a machine with nothing running
// locally simply yields an empty slice.
func DetectLocalRuntimes(ctx context.Context, options LocalDetectOptions) []DetectedLocalRuntime {
	candidates := options.Candidates
	if len(candidates) == 0 {
		candidates = LocalRuntimeCandidates()
	}
	timeout := options.Timeout
	if timeout <= 0 {
		timeout = defaultLocalProbeTimeout
	}
	client := options.HTTPClient
	if client == nil {
		client = &http.Client{Timeout: timeout}
	}

	detected := make([]DetectedLocalRuntime, 0, len(candidates))
	for _, candidate := range candidates {
		models, reachable := probeLocalRuntime(ctx, client, timeout, candidate)
		if !reachable {
			continue
		}
		detected = append(detected, DetectedLocalRuntime{
			LocalRuntime: candidate,
			Reachable:    true,
			Models:       models,
		})
	}
	return detected
}

// SetupAction returns the no-key onboarding action for a detected local runtime.
func (runtime DetectedLocalRuntime) SetupAction() Action {
	descriptor := providercatalog.Descriptor{ID: runtime.CatalogID, RequiresAuth: false}
	command := SetupCommand(descriptor, runtime.Name, true)
	name := strings.TrimSpace(runtime.Name)
	if name == "" {
		name = runtime.CatalogID
	}
	return Action{
		Label:   "Use local runtime",
		Command: command,
		Detail:  "Detected " + name + " on " + runtime.BaseURL + " — no API key required.",
	}
}

func probeLocalRuntime(ctx context.Context, client *http.Client, timeout time.Duration, candidate LocalRuntime) ([]string, bool) {
	endpoint := localModelsEndpoint(candidate.BaseURL)
	if endpoint == "" {
		return nil, false
	}
	probeCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()
	request, err := http.NewRequestWithContext(probeCtx, http.MethodGet, endpoint, nil)
	if err != nil {
		return nil, false
	}
	response, err := client.Do(request)
	if err != nil {
		return nil, false
	}
	defer func() {
		_ = response.Body.Close()
	}()
	if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices {
		return nil, false
	}
	return decodeLocalModelIDs(response.Body), true
}

// decodeLocalModelIDs best-effort parses the OpenAI-compatible models list
// ({"data":[{"id":"..."}]}). The body is bounded so a misbehaving local server
// cannot stream unbounded data into the wizard. A parse failure is not fatal:
// the runtime is still reachable, we just return no model ids.
func decodeLocalModelIDs(body io.Reader) []string {
	data, err := io.ReadAll(io.LimitReader(body, 256*1024))
	if err != nil {
		return nil
	}
	var parsed struct {
		Data []struct {
			ID string `json:"id"`
		} `json:"data"`
	}
	if err := json.Unmarshal(data, &parsed); err != nil {
		return nil
	}
	ids := make([]string, 0, len(parsed.Data))
	seen := map[string]bool{}
	for _, entry := range parsed.Data {
		id := strings.TrimSpace(entry.ID)
		if id == "" || seen[id] {
			continue
		}
		seen[id] = true
		ids = append(ids, id)
	}
	if len(ids) == 0 {
		return nil
	}
	return ids
}

// localModelsEndpoint resolves the OpenAI-compatible "list models" path for a
// base URL. Local base URLs in the catalog already end in /v1, so the endpoint
// is baseURL + "/models"; a trailing slash is trimmed so we never emit "//".
func localModelsEndpoint(baseURL string) string {
	trimmed := strings.TrimRight(strings.TrimSpace(baseURL), "/")
	if trimmed == "" {
		return ""
	}
	return trimmed + "/models"
}