amcs/internal/tools/capture.go

package tools

import (
	"context"
	"log/slog"
	"strings"
	"time"

	"github.com/google/uuid"
	"github.com/modelcontextprotocol/go-sdk/mcp"
	"golang.org/x/sync/errgroup"

	"git.warky.dev/wdevs/amcs/internal/ai"
	"git.warky.dev/wdevs/amcs/internal/config"
	"git.warky.dev/wdevs/amcs/internal/metadata"
	"git.warky.dev/wdevs/amcs/internal/session"
	"git.warky.dev/wdevs/amcs/internal/store"
	thoughttypes "git.warky.dev/wdevs/amcs/internal/types"
)

// EmbeddingQueuer queues a thought for background embedding generation.
type EmbeddingQueuer interface {
	QueueThought(ctx context.Context, id uuid.UUID, content string)
}

type CaptureTool struct {
	store           *store.DB
	provider        ai.Provider
	capture         config.CaptureConfig
	sessions        *session.ActiveProjects
	metadataTimeout time.Duration
	retryer         *MetadataRetryer
	embedRetryer    EmbeddingQueuer
	log             *slog.Logger
}

type CaptureInput struct {
	Content string `json:"content" jsonschema:"the thought or note to capture"`
	Project string `json:"project,omitempty" jsonschema:"optional project name or id to associate with the thought"`
}

type CaptureOutput struct {
	Thought thoughttypes.Thought `json:"thought"`
}

func NewCaptureTool(db *store.DB, provider ai.Provider, capture config.CaptureConfig, metadataTimeout time.Duration, sessions *session.ActiveProjects, retryer *MetadataRetryer, embedRetryer EmbeddingQueuer, log *slog.Logger) *CaptureTool {
	return &CaptureTool{store: db, provider: provider, capture: capture, sessions: sessions, metadataTimeout: metadataTimeout, retryer: retryer, embedRetryer: embedRetryer, log: log}
}

func (t *CaptureTool) Handle(ctx context.Context, req *mcp.CallToolRequest, in CaptureInput) (*mcp.CallToolResult, CaptureOutput, error) {
	content := strings.TrimSpace(in.Content)
	if content == "" {
		return nil, CaptureOutput{}, errRequiredField("content")
	}

	project, err := resolveProject(ctx, t.store, t.sessions, req, in.Project, false)
	if err != nil {
		return nil, CaptureOutput{}, err
	}

	var embedding []float32
	rawMetadata := metadata.Fallback(t.capture)
	metadataNeedsRetry := false
	embeddingNeedsRetry := false

	group, groupCtx := errgroup.WithContext(ctx)
	group.Go(func() error {
		vector, err := t.provider.Embed(groupCtx, content)
		if err != nil {
			t.log.Warn("embedding failed, thought will be saved without embedding",
				slog.String("provider", t.provider.Name()),
				slog.String("error", err.Error()),
			)
			embeddingNeedsRetry = true
			return nil
		}
		embedding = vector
		return nil
	})
	group.Go(func() error {
		metaCtx := groupCtx
		attemptedAt := time.Now().UTC()
		if t.metadataTimeout > 0 {
			var cancel context.CancelFunc
			metaCtx, cancel = context.WithTimeout(groupCtx, t.metadataTimeout)
			defer cancel()
		}
		extracted, err := t.provider.ExtractMetadata(metaCtx, content)
		if err != nil {
			t.log.Warn("metadata extraction failed, using fallback", slog.String("provider", t.provider.Name()), slog.String("error", err.Error()))
			rawMetadata = metadata.MarkMetadataPending(rawMetadata, t.capture, attemptedAt, err)
			metadataNeedsRetry = true
			return nil
		}
		rawMetadata = metadata.MarkMetadataComplete(extracted, t.capture, attemptedAt)
		return nil
	})

	if err := group.Wait(); err != nil {
		return nil, CaptureOutput{}, err
	}

	thought := thoughttypes.Thought{
		Content:   content,
		Embedding: embedding,
		Metadata:  metadata.Normalize(metadata.SanitizeExtracted(rawMetadata), t.capture),
	}
	if project != nil {
		thought.ProjectID = &project.ID
	}

	created, err := t.store.InsertThought(ctx, thought, t.provider.EmbeddingModel())
	if err != nil {
		return nil, CaptureOutput{}, err
	}
	if project != nil {
		_ = t.store.TouchProject(ctx, project.ID)
	}
	if metadataNeedsRetry && t.retryer != nil {
		t.retryer.QueueThought(created.ID)
	}
	if embeddingNeedsRetry && t.embedRetryer != nil {
		t.embedRetryer.QueueThought(ctx, created.ID, content)
	}

	return nil, CaptureOutput{Thought: created}, nil
}