Assistant/cmd/cognitive-worker/main.go

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"git.cin.su/m/Assistant/adapters/models"
	"git.cin.su/m/Assistant/contracts/capability"
	"git.cin.su/m/Assistant/contracts/cognitive"
	"git.cin.su/m/Assistant/contracts/identity"
	secretcontracts "git.cin.su/m/Assistant/contracts/secrets"
	platformsecrets "git.cin.su/m/Assistant/platform/secrets"
	"git.cin.su/m/Assistant/runtime/health"
	"io"
	"log"
	"net/http"
	"os"
	"strings"
	"time"
)

const maxOpenAISecretBytes = 8 << 10

type advisoryProvider interface {
	Advise(context.Context, cognitive.AdvisoryRequest) (cognitive.AdvisoryResponse, error)
}

func main() {
	addr := os.Getenv("LISTEN_ADDR")
	if addr == "" {
		addr = ":8083"
	}
	openAIProvider, err := loadOpenAIAdvisoryProvider(os.Getenv("OPENAI_API_KEY_FILE"))
	if err != nil {
		log.Fatal(err)
	}
	mux := http.NewServeMux()
	resolver, err := newHTTPBindingResolver(envOr("CORE_RUNTIME_URL", "http://core-runtime:8080"), &http.Client{Timeout: 12 * time.Second})
	if err != nil {
		log.Fatal(err)
	}
	providers := map[capability.ProviderKind][]providerRegistration{}
	discoveries := map[string]providerRegistration{}
	var builtinLocal *providerRegistration
	externalRecovery := []providerRegistration{}
	cpuAPIKey, err := loadProviderAPIKey("LOCAL_CPU_PROVIDER_API_KEY", "LOCAL_CPU_PROVIDER_API_KEY_FILE", "LOCAL_CPU_PROVIDER_CREDENTIAL_REF", "credential://domovoy/local-cpu/provider-api-key/v1")
	if err != nil {
		log.Fatal(err)
	}
	cpuAdapter, err := models.NewOpenAICompatibleLocalAdapterWithAPIKey(envOr("LOCAL_CPU_PROVIDER_URL", "http://docker.cin.su:18082/v1"), cpuAPIKey, nil)
	if err != nil {
		log.Fatal(err)
	}
	cpuAuthMode := capability.AuthNone
	cpuCredentialRef := capability.CredentialRef("")
	if cpuAPIKey != "" {
		cpuAuthMode = capability.AuthKey
		cpuCredentialRef = capability.CredentialRef(envOr("LOCAL_CPU_PROVIDER_CREDENTIAL_REF", "credential://domovoy/local-cpu/provider-api-key/v1"))
	}
	cpu := providerRegistration{ProviderInstanceID: "local-cpu-primary", Kind: "openai-compatible-chat-v1", EndpointRef: "local-cpu-primary-v1", CredentialRef: cpuCredentialRef, AuthMode: cpuAuthMode, BillingMode: capability.BillingFree, DataPolicy: capability.DataPolicyLocalOnly, Provider: cpuAdapter, Discovery: cpuAdapter}
	providers[cpu.Kind] = append(providers[cpu.Kind], cpu)
	discoveries["cpu"] = cpu
	mediumURL := strings.TrimSpace(os.Getenv("LOCAL_MEDIUM_PROVIDER_URL"))
	mediumAPIKey := ""
	if mediumURL != "" {
		mediumAPIKey, err = loadProviderAPIKey("LOCAL_MEDIUM_PROVIDER_API_KEY", "LOCAL_MEDIUM_PROVIDER_API_KEY_FILE", "LOCAL_MEDIUM_PROVIDER_CREDENTIAL_REF", "credential://domovoy/local-medium/provider-api-key/v1")
		if err != nil {
			log.Fatal(err)
		}
		mediumTimeout := durationEnv("LOCAL_MEDIUM_PROVIDER_TIMEOUT", 60*time.Second)
		mediumAdapter, adapterErr := models.NewOpenAICompatibleLocalAdapterWithAPIKeyAndTimeout(strings.TrimRight(mediumURL, "/"), mediumAPIKey, nil, mediumTimeout)
		if adapterErr != nil {
			log.Fatal(adapterErr)
		}
		mediumAuthMode := capability.AuthNone
		mediumCredentialRef := capability.CredentialRef("")
		if mediumAPIKey != "" {
			mediumAuthMode = capability.AuthKey
			mediumCredentialRef = capability.CredentialRef(envOr("LOCAL_MEDIUM_PROVIDER_CREDENTIAL_REF", "credential://domovoy/local-medium/provider-api-key/v1"))
		}
		medium := providerRegistration{ProviderInstanceID: "local-medium-primary", Kind: "openai-compatible-chat-v1", EndpointRef: "local-medium-primary-v1", CredentialRef: mediumCredentialRef, AuthMode: mediumAuthMode, BillingMode: capability.BillingFree, DataPolicy: capability.DataPolicyLocalOnly, Provider: mediumAdapter, Discovery: mediumAdapter}
		providers[medium.Kind] = append(providers[medium.Kind], medium)
		discoveries["medium"] = medium
	}
	if lifeboat, enabled, lifeboatErr := newLifeboatProvider(os.Getenv("LOCAL_LIFEBOAT_PROVIDER_URL")); lifeboatErr != nil {
		log.Fatal(lifeboatErr)
	} else if enabled {
		lifeboat.DataPolicy = capability.DataPolicyLocalOnly
		providers[lifeboat.Kind] = append(providers[lifeboat.Kind], lifeboat)
		discoveries["lifeboat"] = lifeboat
		builtinLocal = &lifeboat
	}
	gpuAdapter, err := models.NewOllamaAdapter(envOr("LOCAL_GPU_PROVIDER_URL", "http://192.168.220.91:11434"), nil)
	if err != nil {
		log.Fatal(err)
	}
	gpu := providerRegistration{ProviderInstanceID: "local-gpu-opportunistic", Kind: "ollama-v1", EndpointRef: "local-gpu-opportunistic-v1", AuthMode: capability.AuthNone, BillingMode: capability.BillingFree, DataPolicy: capability.DataPolicyExternalMessageOnly, Provider: gpuAdapter, Discovery: gpuAdapter}
	providers[gpu.Kind] = append(providers[gpu.Kind], gpu)
	discoveries["gpu"] = gpu
	externalRecovery = append(externalRecovery, gpu)
	if openAIProvider != nil {
		kind := capability.ProviderKind(envOr("OPENAI_PROVIDER_KIND", "openai-responses-v1"))
		providers[kind] = append(providers[kind], providerRegistration{ProviderInstanceID: "openai-bootstrap", Kind: kind, EndpointRef: capability.EndpointBindingRef(envOr("OPENAI_ENDPOINT_REF", "openai-responses-default")), CredentialRef: capability.CredentialRef(envOr("OPENAI_CREDENTIAL_REF", "openai-advisory-v1")), AuthMode: capability.AuthKey, BillingMode: capability.BillingPerInput, DataPolicy: capability.DataPolicyExternalMessageOnly, Provider: openAIProvider})
	}
	gateway := cognitiveGateway{resolver: resolver, providers: providers, builtinLocal: builtinLocal, externalRecovery: externalRecovery, now: time.Now, allowCanary: os.Getenv("COGNITIVE_ALLOW_CANARY") == "true", canaryMaxOutputBytes: 4096, admission: newHTTPAdmissionClient(envOr("CORE_RUNTIME_URL", "http://core-runtime:8080"), &http.Client{Timeout: 5 * time.Second})}
	mux.Handle("/healthz", health.Handler{Runtime: "cognitive-worker", SemanticGeneration: 1, ArchitectureGeneration: 1})
	readinessClient := &http.Client{Timeout: 2 * time.Second}
	cpuReadinessClient := &http.Client{Timeout: 10 * time.Second}
	mediumReadinessClient := &http.Client{Timeout: 10 * time.Second}
	lifeboatURL := strings.TrimSuffix(strings.TrimRight(envOr("LOCAL_LIFEBOAT_PROVIDER_URL", "http://cognitive-lifeboat:8086/v1"), "/"), "/v1") + "/health"
	fleetClient := &http.Client{Timeout: 2 * time.Second}
	preferredFleet := capability.FleetTier(envOr("COGNITIVE_PREFERRED_FLEET_TIER", string(capability.FleetTierHeavy)))
	cpuFleetHeaders := map[string]string{}
	if cpuAPIKey != "" {
		cpuFleetHeaders["Authorization"] = "Bearer " + cpuAPIKey
	}
	cpuEndpoint := strings.TrimRight(envOr("LOCAL_CPU_PROVIDER_URL", "http://docker.cin.su:18082/v1"), "/") + "/models"
	cpuHealthProbe := health.OptionalHTTPProbeWithHeaders(cpuReadinessClient, cpuEndpoint, cpuFleetHeaders)
	cpuHealthProbe.Timeout = 10 * time.Second
	mediumEndpoint := ""
	mediumHeaders := map[string]string{}
	mediumHealthProbe := health.Probe{}
	if mediumURL != "" {
		mediumEndpoint = strings.TrimRight(mediumURL, "/") + "/models"
		if mediumAPIKey != "" {
			mediumHeaders["Authorization"] = "Bearer " + mediumAPIKey
		}
		mediumHealthProbe = health.OptionalHTTPProbeWithHeaders(mediumReadinessClient, mediumEndpoint, mediumHeaders)
		mediumHealthProbe.Timeout = 10 * time.Second
	}
	readinessProbes := []health.Probe{
		health.HTTPProbe(readinessClient, strings.TrimRight(envOr("CORE_RUNTIME_URL", "http://core-runtime:8080"), "/")+"/healthz"),
		cpuHealthProbe,
	}
	if mediumURL != "" {
		readinessProbes = append(readinessProbes, mediumHealthProbe)
	}
	readinessProbes = append(readinessProbes,
		health.HTTPProbe(readinessClient, lifeboatURL),
		health.OptionalHTTPProbe(readinessClient, strings.TrimRight(envOr("LOCAL_GPU_PROVIDER_URL", "http://192.168.220.91:11434"), "/")+"/api/tags"),
	)
	mux.Handle("/readyz", health.ReadinessHandler{Runtime: "cognitive-worker", Probes: readinessProbes})
	mux.Handle("/v1/fleet/readiness", fleetReadinessHandler(fleetClient, preferredFleet, []fleetEndpoint{
		{ProviderID: "local-cpu-primary", ModelID: "local-cpu-model", Tier: capability.FleetTierFast, URL: cpuEndpoint, Headers: cpuFleetHeaders, Client: &http.Client{Timeout: 10 * time.Second}},
		{ProviderID: "local-medium-primary", ModelID: "local-medium-model", Tier: capability.FleetTierMedium, URL: mediumEndpoint, Headers: mediumHeaders, Client: mediumReadinessClient},
		{ProviderID: "local-gpu-opportunistic", ModelID: "gpu-opportunistic", Tier: capability.FleetTierHeavy, URL: strings.TrimRight(envOr("LOCAL_GPU_PROVIDER_URL", "http://192.168.220.91:11434"), "/") + "/api/tags"},
		{ProviderID: "cognitive-lifeboat", ModelID: "qwen3-8b-lifeboat", Tier: capability.FleetTierLifeboat, URL: lifeboatURL},
	}))
	mux.HandleFunc("POST /v1/proposals/deterministic", deterministicProposal)
	mux.HandleFunc("POST /v1/advisory-policy", gateway.policyHandler)
	mux.HandleFunc("POST /v1/advisories", gateway.advisoryHandler)
	mux.HandleFunc("POST /v1/interpretations", gateway.interpretationHandler)
	discovery := discoveryService{resolver: resolver, providers: discoveries, owner: identity.PrincipalID(envOr("COGNITIVE_OWNER_PRINCIPAL_ID", "domovoy-creator-principal")), now: time.Now}
	mux.HandleFunc("POST /v1/provider-discovery/{provider}", discovery.handler)
	refreshProviderDiscoveries(resolver, discovery)
	log.Printf("cognitive-worker listening on %s", addr)
	if err := http.ListenAndServe(addr, mux); err != nil {
		log.Fatal(err)
	}
}

func durationEnv(name string, fallback time.Duration) time.Duration {
	value := strings.TrimSpace(os.Getenv(name))
	if value == "" {
		return fallback
	}
	parsed, err := time.ParseDuration(value)
	if err != nil || parsed <= 0 {
		log.Printf("invalid %s=%q; using %s", name, value, fallback)
		return fallback
	}
	return parsed
}

func refreshProviderDiscoveries(resolver bindingResolver, service discoveryService) {
	for name := range service.providers {
		ctx, cancel := context.WithTimeout(context.Background(), 8*time.Second)
		registration, ok := service.providers[name]
		if !ok || registration.Discovery == nil {
			cancel()
			continue
		}
		discovered, probe := registration.Discovery.Discover(ctx)
		cancel()
		snapshot := capability.CognitiveDiscoverySnapshot{
			ProviderInstanceID:   registration.ProviderInstanceID,
			ProviderKind:         registration.Kind,
			EndpointRef:          registration.EndpointRef,
			CredentialRef:        registration.CredentialRef,
			AuthMode:             registration.AuthMode,
			BillingMode:          registration.BillingMode,
			Enabled:              true,
			OwnerPrincipalID:     service.owner,
			CapabilityID:         creatorAdvisoryCapabilityID,
			DataPolicy:           registration.DataPolicy,
			LimitPolicyRef:       "creator-advisory-local-v1",
			Health:               probe.Health,
			FailureClass:         capability.CognitiveFailureType(probe.Failure),
			ObservedAtUnixMillis: time.Now().UnixMilli(),
		}
		for _, model := range discovered {
			snapshot.Models = append(snapshot.Models, capability.DiscoveredModelSnapshot{
				ProviderModelID: model.ProviderModelID,
				Fingerprint:     model.Fingerprint,
				MaxContextBytes: model.MaxContextBytes,
				Modalities:      model.Modalities,
				Capabilities:    model.Capabilities,
				Installed:       capability.OperationalYes,
				Loaded:          capability.OperationalNotObserved,
				Invokable:       capability.OperationalNotObserved,
			})
		}
		if _, err := resolver.RegisterDiscovery(context.Background(), snapshot); err != nil {
			log.Printf("startup cognitive discovery refresh failed for %s: %v", name, err)
		}
	}
}

func loadProviderAPIKey(directEnvName, fileEnvName, refEnvName, defaultRef string) (string, error) {
	if _, set := os.LookupEnv(directEnvName); set {
		return "", fmt.Errorf("%s direct environment secret is forbidden; use %s", directEnvName, fileEnvName)
	}
	fileName := strings.TrimSpace(os.Getenv(fileEnvName))
	if fileName == "" {
		return "", nil
	}
	ref := strings.TrimSpace(os.Getenv(refEnvName))
	if ref == "" {
		ref = defaultRef
	}
	if ref == "" {
		return "", fmt.Errorf("%s is required when %s is configured", refEnvName, fileEnvName)
	}
	resolver, err := platformsecrets.NewFileResolver(envOr("SECRET_BACKEND_ROOT", "/run/secrets"))
	if err != nil {
		return "", err
	}
	return resolver.ResolveFile(secretcontracts.SecretFileRef{CredentialRef: secretcontracts.CredentialRef(ref), Purpose: secretcontracts.SecretPurposeProviderAPIKey, Consumer: secretcontracts.RuntimeConsumerCognitive, Path: fileName})
}

func loadLocalCPUProviderAPIKey() (string, error) {
	return loadProviderAPIKey("LOCAL_CPU_PROVIDER_API_KEY", "LOCAL_CPU_PROVIDER_API_KEY_FILE", "LOCAL_CPU_PROVIDER_CREDENTIAL_REF", "credential://domovoy/local-cpu/provider-api-key/v1")
}

func deterministicAdvisory(w http.ResponseWriter, r *http.Request) {
	request, ok := decodeAdvisoryRequest(w, r)
	if !ok {
		return
	}
	response, err := models.DeterministicAdvisoryUnavailable(request)
	if err != nil {
		http.Error(w, "invalid advisory request", http.StatusBadRequest)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(response)
}

func openAIAdvisory(provider advisoryProvider) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if provider == nil {
			http.Error(w, "provider unavailable", http.StatusServiceUnavailable)
			return
		}
		request, ok := decodeAdvisoryRequest(w, r)
		if !ok {
			return
		}
		response, err := provider.Advise(r.Context(), request)
		if err != nil {
			http.Error(w, "invalid advisory request", http.StatusBadRequest)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(response)
	}
}

func decodeAdvisoryRequest(w http.ResponseWriter, r *http.Request) (cognitive.AdvisoryRequest, bool) {
	r.Body = http.MaxBytesReader(w, r.Body, 80*1024)
	decoder := json.NewDecoder(r.Body)
	decoder.DisallowUnknownFields()
	var request cognitive.AdvisoryRequest
	if err := decoder.Decode(&request); err != nil {
		http.Error(w, "invalid advisory request", http.StatusBadRequest)
		return cognitive.AdvisoryRequest{}, false
	}
	if err := decoder.Decode(&struct{}{}); err != io.EOF {
		http.Error(w, "invalid advisory request", http.StatusBadRequest)
		return cognitive.AdvisoryRequest{}, false
	}
	return request, true
}

func loadOpenAIAdvisoryProvider(secretPath string) (advisoryProvider, error) {
	if os.Getenv("OPENAI_API_KEY") != "" {
		return nil, errors.New("direct OpenAI credential environment is forbidden; use OPENAI_API_KEY_FILE")
	}
	if secretPath == "" {
		return nil, nil
	}
	file, err := os.Open(secretPath)
	if err != nil {
		return nil, fmt.Errorf("open OpenAI credential file: %w", err)
	}
	defer file.Close()
	raw, err := io.ReadAll(io.LimitReader(file, maxOpenAISecretBytes+1))
	if err != nil {
		return nil, fmt.Errorf("read OpenAI credential file: %w", err)
	}
	defer clear(raw)
	if len(raw) > maxOpenAISecretBytes {
		return nil, errors.New("OpenAI credential file is oversized")
	}
	key := bytes.TrimSpace(raw)
	if len(key) == 0 {
		return nil, errors.New("OpenAI credential file is empty")
	}
	for _, character := range key {
		if character < 0x21 || character > 0x7e {
			return nil, errors.New("OpenAI credential file contains invalid characters")
		}
	}
	provider, err := models.NewOpenAIAdvisoryAdapter(key, nil)
	if err != nil {
		return nil, errors.New("OpenAI credential file is invalid")
	}
	return provider, nil
}

func deterministicProposal(w http.ResponseWriter, r *http.Request) {
	var request models.DeterministicRequest
	if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
		http.Error(w, "invalid deterministic proposal request", http.StatusBadRequest)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(models.Propose(request))
}

func envOr(name, fallback string) string {
	if value := os.Getenv(name); value != "" {
		return value
	}
	return fallback
}

func newLifeboatProvider(rawURL string) (providerRegistration, bool, error) {
	if rawURL == "" {
		return providerRegistration{}, false, nil
	}
	adapter, err := models.NewOpenAICompatibleLocalAdapter(rawURL, nil)
	if err != nil {
		return providerRegistration{}, false, err
	}
	return providerRegistration{
		ProviderInstanceID: "local-lifeboat",
		Kind:               "openai-compatible-chat-v1",
		EndpointRef:        "local-lifeboat-v1",
		AuthMode:           capability.AuthNone,
		BillingMode:        capability.BillingFree,
		Provider:           adapter,
		Discovery:          adapter,
	}, true, nil
}