amcs/internal/tools/capture.go

package tools

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

	"github.com/google/uuid"
	"github.com/modelcontextprotocol/go-sdk/mcp"

	"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)
}

// MetadataQueuer queues a thought for background metadata retry. Both
// MetadataRetryer and EnrichmentRetryer satisfy this.
type MetadataQueuer interface {
	QueueThought(id uuid.UUID)
}

type CaptureTool struct {
	store           *store.DB
	embeddings      *ai.EmbeddingRunner
	metadata        *ai.MetadataRunner
	capture         config.CaptureConfig
	sessions        *session.ActiveProjects
	metadataTimeout time.Duration
	retryer         MetadataQueuer
	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, embeddings *ai.EmbeddingRunner, metadata *ai.MetadataRunner, capture config.CaptureConfig, metadataTimeout time.Duration, sessions *session.ActiveProjects, retryer MetadataQueuer, embedRetryer EmbeddingQueuer, log *slog.Logger) *CaptureTool {
	return &CaptureTool{store: db, embeddings: embeddings, metadata: metadata, 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
	}

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

	created, err := t.store.InsertThought(ctx, thought, t.embeddings.PrimaryModel())
	if err != nil {
		return nil, CaptureOutput{}, err
	}
	if project != nil {
		_ = t.store.TouchProject(ctx, project.ID)
	}

	if t.retryer != nil || t.embedRetryer != nil {
		t.launchEnrichment(created.ID, content)
	}

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

func (t *CaptureTool) launchEnrichment(id uuid.UUID, content string) {
	go func() {
		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
		defer cancel()

		if t.retryer != nil {
			attemptedAt := time.Now().UTC()
			rawMetadata := metadata.Fallback(t.capture)
			extracted, err := t.metadata.ExtractMetadata(ctx, content)
			if err != nil {
				failed := metadata.MarkMetadataFailed(rawMetadata, t.capture, attemptedAt, err)
				if _, updateErr := t.store.UpdateThoughtMetadata(ctx, id, failed); updateErr != nil {
					t.log.Warn("deferred metadata failure could not be persisted",
						slog.String("thought_id", id.String()),
						slog.String("error", updateErr.Error()),
					)
				}
				t.log.Warn("deferred metadata extraction failed",
					slog.String("thought_id", id.String()),
					slog.String("provider", t.metadata.PrimaryProvider()),
					slog.String("error", err.Error()),
				)
				t.retryer.QueueThought(id)
			} else {
				completed := metadata.MarkMetadataComplete(extracted, t.capture, attemptedAt)
				if _, updateErr := t.store.UpdateThoughtMetadata(ctx, id, completed); updateErr != nil {
					t.log.Warn("deferred metadata completion could not be persisted",
						slog.String("thought_id", id.String()),
						slog.String("error", updateErr.Error()),
					)
				}
			}
		}

		if t.embedRetryer != nil {
			if _, err := t.embeddings.Embed(ctx, content); err != nil {
				t.log.Warn("deferred embedding failed",
					slog.String("thought_id", id.String()),
					slog.String("provider", t.embeddings.PrimaryProvider()),
					slog.String("error", err.Error()),
				)
			}
			t.embedRetryer.QueueThought(ctx, id, content)
		}
	}()
}