teleport/lib/cache/summarizer.go

// Teleport
// Copyright (C) 2026 Gravitational, Inc.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package cache

import (
	"context"
	"iter"

	"github.com/gravitational/trace"
	"google.golang.org/protobuf/proto"

	headerv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/header/v1"
	summarizerv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/summarizer/v1"
	"github.com/gravitational/teleport/api/types"
	"github.com/gravitational/teleport/api/utils/clientutils"
	"github.com/gravitational/teleport/lib/itertools/stream"
	"github.com/gravitational/teleport/lib/services"
)

type inferenceModelIndex struct{}

var inferenceModelNameIndex = inferenceModelIndex{}

func newInferenceModelCollection(upstream services.Summarizer, w types.WatchKind) (*collection[*summarizerv1.InferenceModel, inferenceModelIndex], error) {
	if upstream == nil {
		return nil, trace.BadParameter("missing parameter Summarizer")
	}

	return &collection[*summarizerv1.InferenceModel, inferenceModelIndex]{
		store: newStore(
			types.KindInferenceModel,
			proto.CloneOf[*summarizerv1.InferenceModel],
			map[inferenceModelIndex]func(*summarizerv1.InferenceModel) string{
				inferenceModelNameIndex: func(r *summarizerv1.InferenceModel) string {
					return r.GetMetadata().GetName()
				},
			}),
		fetcher: func(ctx context.Context, loadSecrets bool) ([]*summarizerv1.InferenceModel, error) {
			out, err := stream.Collect(clientutils.Resources(ctx, upstream.ListInferenceModels))
			return out, trace.Wrap(err)
		},
		headerTransform: func(hdr *types.ResourceHeader) *summarizerv1.InferenceModel {
			return summarizerv1.InferenceModel_builder{
				Kind:    hdr.Kind,
				Version: hdr.Version,
				Metadata: headerv1.Metadata_builder{
					Name: hdr.Metadata.Name,
				}.Build(),
			}.Build()
		},
		watch: w,
	}, nil
}

func (c *Cache) GetInferenceModel(ctx context.Context, name string) (*summarizerv1.InferenceModel, error) {
	ctx, span := c.Tracer.Start(ctx, "cache/GetInferenceModel")
	defer span.End()

	getter := genericGetter[*summarizerv1.InferenceModel, inferenceModelIndex]{
		cache:       c,
		collection:  c.collections.inferenceModels,
		index:       inferenceModelNameIndex,
		upstreamGet: c.Config.Summarizer.GetInferenceModel,
	}
	out, err := getter.get(ctx, name)
	return out, trace.Wrap(err)
}

func (c *Cache) ListInferenceModels(ctx context.Context, pageSize int, pageToken string) ([]*summarizerv1.InferenceModel, string, error) {
	ctx, span := c.Tracer.Start(ctx, "cache/ListInferenceModels")
	defer span.End()

	lister := genericLister[*summarizerv1.InferenceModel, inferenceModelIndex]{
		cache:        c,
		collection:   c.collections.inferenceModels,
		index:        inferenceModelNameIndex,
		upstreamList: c.Config.Summarizer.ListInferenceModels,
		nextToken: func(t *summarizerv1.InferenceModel) string {
			return t.GetMetadata().GetName()
		},
	}
	out, next, err := lister.list(ctx, pageSize, pageToken)
	return out, next, trace.Wrap(err)
}

type inferenceSecretIndex struct{}

var inferenceSecretNameIndex = inferenceSecretIndex{}

func newInferenceSecretCollection(upstream services.Summarizer, w types.WatchKind) (*collection[*summarizerv1.InferenceSecret, inferenceSecretIndex], error) {
	if upstream == nil {
		return nil, trace.BadParameter("missing parameter Summarizer")
	}

	return &collection[*summarizerv1.InferenceSecret, inferenceSecretIndex]{
		store: newStore(
			types.KindInferenceSecret,
			proto.CloneOf[*summarizerv1.InferenceSecret],
			map[inferenceSecretIndex]func(*summarizerv1.InferenceSecret) string{
				inferenceSecretNameIndex: func(r *summarizerv1.InferenceSecret) string {
					return r.GetMetadata().GetName()
				},
			}),
		fetcher: func(ctx context.Context, loadSecrets bool) ([]*summarizerv1.InferenceSecret, error) {
			out, err := stream.Collect(clientutils.Resources(ctx, upstream.ListInferenceSecrets))
			return out, trace.Wrap(err)
		},
		headerTransform: func(hdr *types.ResourceHeader) *summarizerv1.InferenceSecret {
			return summarizerv1.InferenceSecret_builder{
				Kind:    hdr.Kind,
				Version: hdr.Version,
				Metadata: headerv1.Metadata_builder{
					Name: hdr.Metadata.Name,
				}.Build(),
			}.Build()
		},
		watch: w,
	}, nil
}

func (c *Cache) GetInferenceSecret(ctx context.Context, name string) (*summarizerv1.InferenceSecret, error) {
	ctx, span := c.Tracer.Start(ctx, "cache/GetInferenceSecret")
	defer span.End()

	getter := genericGetter[*summarizerv1.InferenceSecret, inferenceSecretIndex]{
		cache:       c,
		collection:  c.collections.inferenceSecrets,
		index:       inferenceSecretNameIndex,
		upstreamGet: c.Config.Summarizer.GetInferenceSecret,
	}
	out, err := getter.get(ctx, name)
	return out, trace.Wrap(err)
}

func (c *Cache) ListInferenceSecrets(ctx context.Context, pageSize int, pageToken string) ([]*summarizerv1.InferenceSecret, string, error) {
	ctx, span := c.Tracer.Start(ctx, "cache/ListInferenceSecrets")
	defer span.End()

	lister := genericLister[*summarizerv1.InferenceSecret, inferenceSecretIndex]{
		cache:        c,
		collection:   c.collections.inferenceSecrets,
		index:        inferenceSecretNameIndex,
		upstreamList: c.Config.Summarizer.ListInferenceSecrets,
		nextToken: func(t *summarizerv1.InferenceSecret) string {
			return t.GetMetadata().GetName()
		},
	}
	out, next, err := lister.list(ctx, pageSize, pageToken)
	return out, next, trace.Wrap(err)
}

type inferencePolicyIndex struct{}

var inferencePolicyNameIndex = inferencePolicyIndex{}

func newInferencePolicyCollection(upstream services.Summarizer, w types.WatchKind) (*collection[*summarizerv1.InferencePolicy, inferencePolicyIndex], error) {
	if upstream == nil {
		return nil, trace.BadParameter("missing parameter Summarizer")
	}

	return &collection[*summarizerv1.InferencePolicy, inferencePolicyIndex]{
		store: newStore(
			types.KindInferencePolicy,
			proto.CloneOf[*summarizerv1.InferencePolicy],
			map[inferencePolicyIndex]func(*summarizerv1.InferencePolicy) string{
				inferencePolicyNameIndex: func(r *summarizerv1.InferencePolicy) string {
					return r.GetMetadata().GetName()
				},
			}),
		fetcher: func(ctx context.Context, loadSecrets bool) ([]*summarizerv1.InferencePolicy, error) {
			out, err := stream.Collect(clientutils.Resources(ctx, upstream.ListInferencePolicies))
			return out, trace.Wrap(err)
		},
		headerTransform: func(hdr *types.ResourceHeader) *summarizerv1.InferencePolicy {
			return summarizerv1.InferencePolicy_builder{
				Kind:    hdr.Kind,
				Version: hdr.Version,
				Metadata: headerv1.Metadata_builder{
					Name: hdr.Metadata.Name,
				}.Build(),
			}.Build()
		},
		watch: w,
	}, nil
}

func (c *Cache) GetInferencePolicy(ctx context.Context, name string) (*summarizerv1.InferencePolicy, error) {
	ctx, span := c.Tracer.Start(ctx, "cache/GetInferencePolicy")
	defer span.End()

	getter := genericGetter[*summarizerv1.InferencePolicy, inferencePolicyIndex]{
		cache:       c,
		collection:  c.collections.inferencePolicies,
		index:       inferencePolicyNameIndex,
		upstreamGet: c.Config.Summarizer.GetInferencePolicy,
	}
	out, err := getter.get(ctx, name)
	return out, trace.Wrap(err)
}

func (c *Cache) ListInferencePolicies(ctx context.Context, pageSize int, pageToken string) ([]*summarizerv1.InferencePolicy, string, error) {
	ctx, span := c.Tracer.Start(ctx, "cache/ListInferencePolicies")
	defer span.End()

	lister := genericLister[*summarizerv1.InferencePolicy, inferencePolicyIndex]{
		cache:        c,
		collection:   c.collections.inferencePolicies,
		index:        inferencePolicyNameIndex,
		upstreamList: c.Config.Summarizer.ListInferencePolicies,
		nextToken: func(t *summarizerv1.InferencePolicy) string {
			return t.GetMetadata().GetName()
		},
	}
	out, next, err := lister.list(ctx, pageSize, pageToken)
	return out, next, trace.Wrap(err)
}

func (c *Cache) AllInferencePolicies(ctx context.Context) iter.Seq2[*summarizerv1.InferencePolicy, error] {
	return func(yield func(*summarizerv1.InferencePolicy, error) bool) {
		rg, err := acquireReadGuard(c, c.collections.inferencePolicies)
		if err != nil {
			yield(nil, trace.Wrap(err))
			return
		}
		defer rg.Release()

		if !rg.ReadCache() {
			for policy, err := range c.Config.Summarizer.AllInferencePolicies(ctx) {
				if !yield(policy, err) {
					return
				}
			}
			return
		}

		for policy := range rg.store.resources(inferencePolicyNameIndex, "", "") {
			if !yield(proto.CloneOf(policy), nil) {
				return
			}
		}
	}
}

type classifierIndex struct{}

var classifierNameIndex = classifierIndex{}

func newClassifierCollection(upstream services.Summarizer, w types.WatchKind) (*collection[*summarizerv1.Classifier, classifierIndex], error) {
	if upstream == nil {
		return nil, trace.BadParameter("missing parameter Summarizer")
	}

	return &collection[*summarizerv1.Classifier, classifierIndex]{
		store: newStore(
			types.KindClassifier,
			proto.CloneOf[*summarizerv1.Classifier],
			map[classifierIndex]func(*summarizerv1.Classifier) string{
				classifierNameIndex: func(r *summarizerv1.Classifier) string {
					return r.GetMetadata().GetName()
				},
			}),
		fetcher: func(ctx context.Context, loadSecrets bool) ([]*summarizerv1.Classifier, error) {
			out, err := stream.Collect(clientutils.Resources(ctx, upstream.ListClassifiers))
			return out, trace.Wrap(err)
		},
		headerTransform: func(hdr *types.ResourceHeader) *summarizerv1.Classifier {
			return summarizerv1.Classifier_builder{
				Kind:    hdr.Kind,
				Version: hdr.Version,
				Metadata: headerv1.Metadata_builder{
					Name: hdr.Metadata.Name,
				}.Build(),
			}.Build()
		},
		watch: w,
	}, nil
}

func (c *Cache) GetClassifier(ctx context.Context, name string) (*summarizerv1.Classifier, error) {
	ctx, span := c.Tracer.Start(ctx, "cache/GetClassifier")
	defer span.End()

	getter := genericGetter[*summarizerv1.Classifier, classifierIndex]{
		cache:       c,
		collection:  c.collections.classifiers,
		index:       classifierNameIndex,
		upstreamGet: c.Config.Summarizer.GetClassifier,
	}
	out, err := getter.get(ctx, name)
	return out, trace.Wrap(err)
}

func (c *Cache) ListClassifiers(ctx context.Context, pageSize int, pageToken string) ([]*summarizerv1.Classifier, string, error) {
	ctx, span := c.Tracer.Start(ctx, "cache/ListClassifiers")
	defer span.End()

	lister := genericLister[*summarizerv1.Classifier, classifierIndex]{
		cache:        c,
		collection:   c.collections.classifiers,
		index:        classifierNameIndex,
		upstreamList: c.Config.Summarizer.ListClassifiers,
		nextToken: func(t *summarizerv1.Classifier) string {
			return t.GetMetadata().GetName()
		},
	}
	out, next, err := lister.list(ctx, pageSize, pageToken)
	return out, next, trace.Wrap(err)
}

func (c *Cache) RangeClassifiers(ctx context.Context, start, end string) iter.Seq2[*summarizerv1.Classifier, error] {
	lister := genericLister[*summarizerv1.Classifier, classifierIndex]{
		cache:        c,
		collection:   c.collections.classifiers,
		index:        classifierNameIndex,
		upstreamList: c.Config.Summarizer.ListClassifiers,
		nextToken: func(t *summarizerv1.Classifier) string {
			return t.GetMetadata().GetName()
		},
	}

	return func(yield func(*summarizerv1.Classifier, error) bool) {
		ctx, span := c.Tracer.Start(ctx, "cache/RangeClassifiers")
		defer span.End()

		for classifier, err := range lister.Range(ctx, start, end) {
			if !yield(classifier, err) {
				return
			}

			if err != nil {
				return
			}
		}
	}
}

type retrievalModelIndex struct{}

var retrievalModelNameIndex = retrievalModelIndex{}

func newRetrievalModelCollection(upstream services.Summarizer, w types.WatchKind) (*collection[*summarizerv1.RetrievalModel, retrievalModelIndex], error) {
	if upstream == nil {
		return nil, trace.BadParameter("missing parameter Summarizer")
	}

	return &collection[*summarizerv1.RetrievalModel, retrievalModelIndex]{
		store: newStore(
			types.KindRetrievalModel,
			proto.CloneOf[*summarizerv1.RetrievalModel],
			map[retrievalModelIndex]func(*summarizerv1.RetrievalModel) string{
				retrievalModelNameIndex: func(r *summarizerv1.RetrievalModel) string {
					return r.GetMetadata().GetName()
				},
			}),
		fetcher: func(ctx context.Context, loadSecrets bool) ([]*summarizerv1.RetrievalModel, error) {
			model, err := upstream.GetRetrievalModel(ctx)
			if err != nil {
				if trace.IsNotFound(err) {
					return nil, nil
				}
				return nil, trace.Wrap(err)
			}
			return []*summarizerv1.RetrievalModel{model}, nil
		},
		headerTransform: func(hdr *types.ResourceHeader) *summarizerv1.RetrievalModel {
			return summarizerv1.RetrievalModel_builder{
				Kind:    hdr.Kind,
				Version: hdr.Version,
				Metadata: headerv1.Metadata_builder{
					Name: hdr.Metadata.Name,
				}.Build(),
			}.Build()
		},
		watch: w,
	}, nil
}

func (c *Cache) GetRetrievalModel(ctx context.Context) (*summarizerv1.RetrievalModel, error) {
	ctx, span := c.Tracer.Start(ctx, "cache/GetRetrievalModel")
	defer span.End()

	getter := genericGetter[*summarizerv1.RetrievalModel, retrievalModelIndex]{
		cache:      c,
		collection: c.collections.retrievalModels,
		index:      retrievalModelNameIndex,
		upstreamGet: func(ctx context.Context, _ string) (*summarizerv1.RetrievalModel, error) {
			return c.Config.Summarizer.GetRetrievalModel(ctx)
		},
	}
	out, err := getter.get(ctx, types.MetaNameRetrievalModel)
	return out, trace.Wrap(err)
}