loco-operator/internal/reconciler/agent.go

package reconciler

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"maps"
	"sort"
	"strings"

	"github.com/LocoStack/loco-operator/api/v1alpha1"
	"github.com/LocoStack/loco-operator/pkg/templates/litellm"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
)

type AgentReconciler struct {
	client    client.Client
	scheme    *runtime.Scheme
	stack     *v1alpha1.Stack
	component *v1alpha1.Component
}

func NewAgentReconciler(client client.Client, scheme *runtime.Scheme, stack *v1alpha1.Stack, component *v1alpha1.Component) *AgentReconciler {
	return &AgentReconciler{
		client:    client,
		scheme:    scheme,
		stack:     stack,
		component: component,
	}
}

type agentComponentEndpoint struct {
	Name string
	URL  string
}

type AgentComponents struct {
	Gateway        agentComponentEndpoint
	GatewayKey     string
	Model          agentComponentEndpoint
	Skills         []string
	Tools          []agentComponentEndpoint
	KnowledgeBases []agentComponentEndpoint
	Agents         []agentComponentEndpoint
}

func (r *AgentReconciler) getDependencies(ctx context.Context) ([]client.Object, *AgentComponents, error) {
	comps := &AgentComponents{
		Tools:          []agentComponentEndpoint{},
		KnowledgeBases: []agentComponentEndpoint{},
	}
	endpointResolver := NewEndpointResolver(r.client, r.scheme, r.stack.Namespace)
	deps := make([]client.Object, 0)
	for _, dep := range r.component.Spec.Dependencies {
		obj, endpoint, err := endpointResolver.ResolveEndpoint(ctx, r.stack, dep)
		if err != nil {
			return nil, nil, fmt.Errorf("Failed to resolve endpoint for %s %s: %w", dep.Kind, dep.Name, err)
		}
		switch dep.Kind {
		case "ExternalModel":
			fallthrough
		case "ManagedModel":
			var category string
			if extModel, ok := obj.(*v1alpha1.ExternalModel); ok {
				category = extModel.Spec.Category
				comps.Model.Name = extModel.Spec.ModelName
			} else if managedModel, ok := obj.(*v1alpha1.ManagedModel); ok {
				category = managedModel.Spec.Category
				comps.Model.Name = managedModel.Spec.ModelName
			} else {
				return nil, nil, fmt.Errorf("Unsupported dependency kind for agent: %s", dep.Kind)
			}
			if category != "language" {
				return nil, nil, fmt.Errorf("Unsupported model category for agent: %s", category)
			}
			comps.Model.URL = endpoint
		case "Skill":
			comps.Skills = append(comps.Skills, dep.Name)
		case "ExternalTool":
			fallthrough
		case "ManagedTool":
			var toolName string
			if extTool, ok := obj.(*v1alpha1.ExternalTool); ok {
				toolName = extTool.Spec.ToolName
			} else if managedTool, ok := obj.(*v1alpha1.ManagedTool); ok {
				toolName = managedTool.Spec.ToolName
			}
			comps.Tools = append(comps.Tools, agentComponentEndpoint{Name: toolName, URL: endpoint})
		case "ExternalKnowledgeBase":
			fallthrough
		case "ManagedKnowledgeBase":
			comps.KnowledgeBases = append(comps.KnowledgeBases, agentComponentEndpoint{Name: dep.Name, URL: endpoint})
		case "Agent":
			comps.Agents = append(comps.Agents, agentComponentEndpoint{Name: dep.Name, URL: endpoint})
		default:
			return nil, nil, fmt.Errorf("Unsupported dependency kind for agent: %s", dep.Kind)
		}
		deps = append(deps, obj)
	}
	sort.Strings(comps.Skills)
	sort.Slice(comps.Tools, func(i, j int) bool {
		return comps.Tools[i].Name < comps.Tools[j].Name
	})
	sort.Slice(comps.KnowledgeBases, func(i, j int) bool {
		return comps.KnowledgeBases[i].Name < comps.KnowledgeBases[j].Name
	})
	sort.Slice(comps.Agents, func(i, j int) bool {
		return comps.Agents[i].Name < comps.Agents[j].Name
	})

	gatewayComp, err := endpointResolver.resolveGatewayEndpoint(ctx, r.stack)
	if err != nil {
		return nil, nil, fmt.Errorf("Failed to resolve gateway for %s: %w", r.component.Name, err)
	}
	comps.Gateway = agentComponentEndpoint{Name: gatewayComp.Name, URL: gatewayComp.Status.Endpoint}
	masterKey, err := r.gatewayMasterKey(ctx, gatewayComp)
	if err != nil {
		return nil, nil, fmt.Errorf("resolving gateway master key: %w", err)
	}
	comps.GatewayKey = masterKey
	return deps, comps, nil
}

func (r *AgentReconciler) gatewayMasterKey(ctx context.Context, gatewayComp *v1alpha1.Component) (string, error) {
	if gatewayComp == nil {
		return "", nil
	}
	secret := &corev1.Secret{}
	secretName := litellm.MasterKeySecretName(gatewayComp.Name)
	if err := r.client.Get(ctx, client.ObjectKey{Name: secretName, Namespace: r.component.GetNamespace()}, secret); err != nil {
		return "", fmt.Errorf("Failed to get gateway master key secret %s: %w", secretName, err)
	}
	key, ok := secret.Data[litellm.LITELLM_AUTH_SECRET_KEY]
	if !ok {
		return "", fmt.Errorf("Secret %s does not contain key %q", secretName, litellm.LITELLM_AUTH_SECRET_KEY)
	}
	return string(key), nil
}

// AddSkillsToTmpl adds skill mounts to the template,
// mounting each dependency Skill's ConfigMap under <skillsRoot>/skills/<skillName>,
// and returns a hash of the mounted files. Agent runtimes read their skills once at
// startup, and a ConfigMap edited in place never restarts the Pod, so callers fold that
// hash into the Pod template to roll the agent out whenever a Skill's content changes.
// A Skill's ConfigMap stores one entry per file (SKILL.md plus its scripts/references/
// assets), keyed by v1alpha1.SkillConfigMapKey(path) rather than the raw relative path,
// since ConfigMap keys cannot contain "/". The ConfigMap's v1alpha1.SkillFilePathsAnnotation
// records the key -> original relative path mapping, so it can be read directly rather than
// re-fetched and recomputed from the Skill spec; a ConfigMapVolumeSource with per-key Items
// then lays the files back out as a real directory tree using that mapping.
func (r *AgentReconciler) AddSkillsToTmpl(ctx context.Context, comps *AgentComponents, tmpl *v1alpha1.Template, skillsRoot string) (string, error) {
	volumes := make([]corev1.Volume, 0, len(comps.Skills))
	volumeMounts := make([]corev1.VolumeMount, 0, len(comps.Skills))
	skillsHash := sha256.New()
	for _, skillName := range comps.Skills {
		cm := &corev1.ConfigMap{}
		if err := r.client.Get(ctx, client.ObjectKey{Namespace: r.stack.Namespace, Name: skillName}, cm); err != nil {
			return "", fmt.Errorf("Failed to get ConfigMap for Skill %s: %w", skillName, err)
		}
		filePathsJSON, ok := cm.Annotations[v1alpha1.SkillFilePathsAnnotation]
		if !ok {
			return "", fmt.Errorf("ConfigMap for Skill %s is missing the %q annotation", skillName, v1alpha1.SkillFilePathsAnnotation)
		}
		var filePaths map[string]string
		if err := json.Unmarshal([]byte(filePathsJSON), &filePaths); err != nil {
			return "", fmt.Errorf("unmarshaling file paths for Skill %s: %w", skillName, err)
		}

		keys := sortedKeys(filePaths)

		items := make([]corev1.KeyToPath, 0, len(keys))
		for _, key := range keys {
			items = append(items, corev1.KeyToPath{
				Key:  key,
				Path: filePaths[key],
			})
			_, _ = fmt.Fprintf(skillsHash, "%s/%s=%s\n", skillName, filePaths[key], cm.Data[key])
		}
		volumes = append(volumes, corev1.Volume{
			Name: skillName,
			VolumeSource: corev1.VolumeSource{
				ConfigMap: &corev1.ConfigMapVolumeSource{
					LocalObjectReference: corev1.LocalObjectReference{Name: skillName},
					Items:                items,
				},
			},
		})
		volumeMounts = append(volumeMounts, corev1.VolumeMount{
			Name:      skillName,
			MountPath: fmt.Sprintf("%s/skills/%s", skillsRoot, skillName),
			ReadOnly:  true,
		})
	}
	tmpl.Spec.Volumes = append(tmpl.Spec.Volumes, volumes...)
	tmpl.Spec.Runtime.VolumeMounts = append(tmpl.Spec.Runtime.VolumeMounts, volumeMounts...)
	return hex.EncodeToString(skillsHash.Sum(nil)), nil
}

type mcpConfig struct {
	MCPServers map[string]mcpServer `json:"mcpServers"`
}

type mcpServer struct {
	URL     string            `json:"url"`
	Headers map[string]string `json:"headers,omitempty"`
}

// reconcileMcpConfig writes the agent's mcp.json as a Secret (not a ConfigMap): when a
// gateway is present, the generated config embeds the resolved litellm master key value
// in the auth header, so the credential must not land in a ConfigMap.
func (r *AgentReconciler) reconcileMcpConfig(ctx context.Context, comps *AgentComponents) (string, error) {
	cfg := mcpConfig{MCPServers: make(map[string]mcpServer)}
	if comps.Gateway.URL != "" && len(comps.Tools) > 0 {
		toolNames := make([]string, 0, len(comps.Tools))
		for _, tool := range comps.Tools {
			toolNames = append(toolNames, tool.Name)
		}
		sort.Strings(toolNames)
		cfg.MCPServers["litellm"] = mcpServer{
			URL: fmt.Sprintf("%s/%s/mcp", comps.Gateway.URL, strings.Join(toolNames, ",")),
			Headers: map[string]string{
				litellm.MCPAuthHeaderName: "Bearer " + comps.GatewayKey,
			},
		}
	} else {
		for _, tool := range comps.Tools {
			cfg.MCPServers[tool.Name] = mcpServer{URL: tool.URL}
		}
	}
	jsonBytes, err := json.Marshal(cfg)
	if err != nil {
		return "", fmt.Errorf("marshaling agent mcp config: %w", err)
	}
	h := sha256.Sum256(jsonBytes)
	configHash := hex.EncodeToString(h[:])
	secret := &corev1.Secret{ObjectMeta: metav1.ObjectMeta{Name: fmt.Sprintf("%s-mcp-config", r.component.GetName()), Namespace: r.component.GetNamespace()}}
	_, err = controllerutil.CreateOrUpdate(ctx, r.client, secret, func() error {
		if secret.Labels == nil {
			secret.Labels = map[string]string{}
		}
		maps.Copy(secret.Labels, ResourceLabels(r.stack.Name, r.component.Spec.Category, r.component.GetName(), ""))
		if secret.Annotations == nil {
			secret.Annotations = make(map[string]string)
		}
		secret.Annotations["hash"] = configHash
		secret.Data = map[string][]byte{
			"mcp.json": jsonBytes,
		}
		return controllerutil.SetControllerReference(r.component, secret, r.scheme)
	})
	return configHash, err
}