tutti/services/tuttid/data/workspace/sqlite_model_plans_test.go

package workspace

import (
	"context"
	"errors"
	"testing"
	"time"

	agenttargetbiz "github.com/tutti-os/tutti/services/tuttid/biz/agenttarget"
	modelbindingbiz "github.com/tutti-os/tutti/services/tuttid/biz/modelbinding"
	modelplanbiz "github.com/tutti-os/tutti/services/tuttid/biz/modelplan"
	modelpolicybiz "github.com/tutti-os/tutti/services/tuttid/biz/modelpolicy"
	workspacebiz "github.com/tutti-os/tutti/services/tuttid/biz/workspace"
)

func createModelPlanTestWorkspace(t *testing.T, store *SQLiteStore, workspaceID string) {
	t.Helper()
	if err := store.Create(context.Background(), workspacebiz.Summary{
		ID:   workspaceID,
		Name: "Model Plan Workspace",
	}); err != nil {
		t.Fatalf("Create workspace error = %v", err)
	}
}

func TestSQLiteStoreModelPlanRoundTrip(t *testing.T) {
	t.Parallel()

	ctx := context.Background()
	store := openTestSQLiteStore(t)
	createModelPlanTestWorkspace(t, store, "ws-plans")

	now := time.UnixMilli(1700000000000).UTC()
	plan := modelplanbiz.Plan{
		ID:           "mp-test",
		WorkspaceID:  "ws-plans",
		Name:         "Volc Coding Plan",
		TemplateKind: modelplanbiz.TemplateCodingPlan,
		Protocol:     modelplanbiz.ProtocolAnthropic,
		APIKey:       "sk-secret-value",
		BaseURL:      "https://open.volcengineapi.example/v1",
		Models: []modelplanbiz.Model{
			{
				ID:   "doubao-seed-code",
				Name: "Doubao Seed Code",
				Pricing: &modelplanbiz.ModelPricing{
					Currency:                   "USD",
					InputMicrosPerMillion:      100,
					OutputMicrosPerMillion:     200,
					CacheReadMicrosPerMillion:  30,
					CacheWriteMicrosPerMillion: 40,
				},
			},
		},
		DefaultModel: "doubao-seed-code",
		Enabled:      true,
		Detection: modelplanbiz.DetectionSnapshot{
			CheckedAt: now,
			Model:     "doubao-seed-code",
			Stages: []modelplanbiz.StageResult{
				{Stage: modelplanbiz.StageNetwork, Status: modelplanbiz.StagePassed, CheckedAt: now},
			},
		},
		CreatedAt: now,
		UpdatedAt: now,
	}
	if err := store.PutModelPlan(ctx, plan); err != nil {
		t.Fatalf("PutModelPlan() error = %v", err)
	}

	loaded, err := store.GetModelPlan(ctx, "ws-plans", "mp-test")
	if err != nil {
		t.Fatalf("GetModelPlan() error = %v", err)
	}
	if loaded.APIKey != "sk-secret-value" {
		t.Fatalf("GetModelPlan() api key = %q, want decrypted secret", loaded.APIKey)
	}
	if loaded.Revision != 1 {
		t.Fatalf("GetModelPlan() revision = %d, want 1", loaded.Revision)
	}
	if loaded.Protocol != modelplanbiz.ProtocolAnthropic || loaded.TemplateKind != modelplanbiz.TemplateCodingPlan {
		t.Fatalf("GetModelPlan() protocol/template = %q/%q", loaded.Protocol, loaded.TemplateKind)
	}
	if len(loaded.Models) != 1 || loaded.Models[0].ID != "doubao-seed-code" {
		t.Fatalf("GetModelPlan() models = %#v", loaded.Models)
	}
	if loaded.Models[0].Pricing == nil || loaded.Models[0].Pricing.OutputMicrosPerMillion != 200 {
		t.Fatalf("GetModelPlan() pricing = %#v, want persisted pricing", loaded.Models[0].Pricing)
	}
	if outcome, ok := loaded.Detection.StageOutcome(modelplanbiz.StageNetwork); !ok || outcome.Status != modelplanbiz.StagePassed {
		t.Fatalf("GetModelPlan() detection = %#v", loaded.Detection)
	}

	// The ciphertext column must never contain the raw key.
	var ciphertext string
	if err := store.readDB.QueryRowContext(ctx, `SELECT api_key_ciphertext FROM model_plans WHERE plan_id = 'mp-test'`).Scan(&ciphertext); err != nil {
		t.Fatalf("read ciphertext error = %v", err)
	}
	if ciphertext == "" || ciphertext == "sk-secret-value" {
		t.Fatalf("api_key_ciphertext = %q, want encrypted value", ciphertext)
	}

	plans, err := store.ListModelPlans(ctx, "ws-plans")
	if err != nil {
		t.Fatalf("ListModelPlans() error = %v", err)
	}
	if len(plans) != 1 {
		t.Fatalf("ListModelPlans() len = %d, want 1", len(plans))
	}

	loaded.Revision = 2
	loaded.APIKey = "sk-rotated-value"
	loaded.BaseURL = "https://rotated.example/v1"
	loaded.UpdatedAt = now.Add(time.Minute)
	if err := store.PutModelPlan(ctx, loaded); err != nil {
		t.Fatalf("PutModelPlan(revision 2) error = %v", err)
	}
	current, err := store.GetModelPlan(ctx, "ws-plans", "mp-test")
	if err != nil {
		t.Fatalf("GetModelPlan(revision 2) error = %v", err)
	}
	if current.Revision != 2 || current.APIKey != "sk-rotated-value" || current.BaseURL != "https://rotated.example/v1" {
		t.Fatalf("current plan = revision %d key %q url %q", current.Revision, current.APIKey, current.BaseURL)
	}
	historical, err := store.GetModelPlanRevision(ctx, "ws-plans", "mp-test", 1)
	if err != nil {
		t.Fatalf("GetModelPlanRevision(1) error = %v", err)
	}
	if historical.Revision != 1 || historical.APIKey != "sk-secret-value" || historical.BaseURL != "https://open.volcengineapi.example/v1" {
		t.Fatalf("historical plan = revision %d key %q url %q", historical.Revision, historical.APIKey, historical.BaseURL)
	}
	var historicalCiphertext string
	if err := store.readDB.QueryRowContext(ctx, `SELECT api_key_ciphertext FROM model_plan_revisions WHERE plan_id = 'mp-test' AND revision = 1`).Scan(&historicalCiphertext); err != nil {
		t.Fatalf("read historical ciphertext error = %v", err)
	}
	if historicalCiphertext == "" || historicalCiphertext == "sk-secret-value" {
		t.Fatalf("historical ciphertext = %q, want encrypted value", historicalCiphertext)
	}

	if err := store.DeleteModelPlan(ctx, "ws-plans", "mp-test"); err != nil {
		t.Fatalf("DeleteModelPlan() error = %v", err)
	}
	if _, err := store.GetModelPlan(ctx, "ws-plans", "mp-test"); !errors.Is(err, ErrModelPlanNotFound) {
		t.Fatalf("GetModelPlan() after delete error = %v, want ErrModelPlanNotFound", err)
	}
	if preserved, err := store.GetModelPlanRevision(ctx, "ws-plans", "mp-test", 1); err != nil || preserved.APIKey != "sk-secret-value" {
		t.Fatalf("GetModelPlanRevision() after current delete = %#v, %v", preserved, err)
	}
	if err := store.DeleteModelPlan(ctx, "ws-plans", "mp-test"); !errors.Is(err, ErrModelPlanNotFound) {
		t.Fatalf("DeleteModelPlan() second delete error = %v, want ErrModelPlanNotFound", err)
	}
}

func TestModelPlansMigrationBackfillsLegacyManagedCredentials(t *testing.T) {
	t.Parallel()

	ctx := context.Background()
	store := openTestSQLiteStore(t)
	createModelPlanTestWorkspace(t, store, "ws-legacy")

	ciphertext, err := encryptManagedCredential("legacy-key")
	if err != nil {
		t.Fatalf("encryptManagedCredential() error = %v", err)
	}
	if _, err := store.writeDB.ExecContext(ctx, `
INSERT INTO managed_model_provider_credentials (
  workspace_id, provider_id, enabled, api_key_ciphertext, base_url, models_json, updated_at_unix_ms
) VALUES ('ws-legacy', 'anthropic', 1, ?, 'https://relay.example/api/anthropic', '[{"id":"claude-x","name":"Claude X","provider":"anthropic"}]', 1700000000000)
`, ciphertext); err != nil {
		t.Fatalf("insert legacy credential error = %v", err)
	}

	resetModelPlanMigrations(t, store)
	if err := store.applyModelPlansV1(ctx); err != nil {
		t.Fatalf("applyModelPlansV1() error = %v", err)
	}
	if err := store.applyAgentModelBindingsV1(ctx); err != nil {
		t.Fatalf("applyAgentModelBindingsV1() error = %v", err)
	}
	if err := store.applyModelPlanRevisionsV1(ctx); err != nil {
		t.Fatalf("applyModelPlanRevisionsV1() error = %v", err)
	}

	plan, err := store.GetModelPlan(ctx, "ws-legacy", "mp-migrated-anthropic")
	if err != nil {
		t.Fatalf("GetModelPlan(migrated) error = %v", err)
	}
	if plan.APIKey != "legacy-key" {
		t.Fatalf("migrated plan api key = %q, want legacy-key", plan.APIKey)
	}
	if plan.Protocol != modelplanbiz.ProtocolAnthropic {
		t.Fatalf("migrated plan protocol = %q, want anthropic", plan.Protocol)
	}
	if !plan.Enabled {
		t.Fatalf("migrated plan enabled = false, want true")
	}
	if plan.BaseURL != "https://relay.example/api/anthropic" {
		t.Fatalf("migrated plan base url = %q", plan.BaseURL)
	}
	if len(plan.Models) != 1 || plan.Models[0].ID != "claude-x" || plan.Models[0].Name != "Claude X" {
		t.Fatalf("migrated plan models = %#v", plan.Models)
	}
	if plan.Revision != 1 {
		t.Fatalf("migrated plan revision = %d, want 1", plan.Revision)
	}
	historical, err := store.GetModelPlanRevision(ctx, "ws-legacy", "mp-migrated-anthropic", 1)
	if err != nil || historical.APIKey != "legacy-key" {
		t.Fatalf("migrated plan revision = %#v, %v", historical, err)
	}
	if plan.Status() != modelplanbiz.StatusUndetected {
		t.Fatalf("migrated plan status = %q, want undetected", plan.Status())
	}
}

func TestModelPlansMigrationRollsBackFailedBackfillAndRetries(t *testing.T) {
	t.Parallel()

	ctx := context.Background()
	store := openTestSQLiteStore(t)
	createModelPlanTestWorkspace(t, store, "ws-retry")
	ciphertext, err := encryptManagedCredential("legacy-key")
	if err != nil {
		t.Fatalf("encryptManagedCredential() error = %v", err)
	}
	if _, err := store.writeDB.ExecContext(ctx, `
INSERT INTO managed_model_provider_credentials (
  workspace_id, provider_id, enabled, api_key_ciphertext, base_url, models_json, updated_at_unix_ms
) VALUES ('ws-retry', 'openai', 1, ?, 'https://relay.example/v1', '{invalid', 1700000000000)
`, ciphertext); err != nil {
		t.Fatalf("insert invalid legacy credential error = %v", err)
	}

	resetModelPlanMigrations(t, store)
	if err := store.applyModelPlansV1(ctx); err == nil {
		t.Fatal("applyModelPlansV1() error = nil, want invalid legacy models failure")
	}
	applied, err := store.hasMigration(ctx, schemaMigrationModelPlansV1)
	if err != nil {
		t.Fatalf("hasMigration() error = %v", err)
	}
	if applied {
		t.Fatal("failed model plan migration recorded its marker")
	}
	var tableCount int
	if err := store.writeDB.QueryRowContext(ctx, `SELECT COUNT(*) FROM sqlite_master WHERE type = 'table' AND name = 'model_plans'`).Scan(&tableCount); err != nil {
		t.Fatalf("inspect rolled back model_plans table error = %v", err)
	}
	if tableCount != 0 {
		t.Fatalf("model_plans table count after rollback = %d, want 0", tableCount)
	}

	if _, err := store.writeDB.ExecContext(ctx, `
UPDATE managed_model_provider_credentials
SET models_json = '[{"id":"gpt-retry","name":"GPT Retry"}]'
WHERE workspace_id = 'ws-retry' AND provider_id = 'openai'
`); err != nil {
		t.Fatalf("repair legacy credential error = %v", err)
	}
	if err := store.applyModelPlansV1(ctx); err != nil {
		t.Fatalf("retry applyModelPlansV1() error = %v", err)
	}
	if err := store.applyAgentModelBindingsV1(ctx); err != nil {
		t.Fatalf("retry applyAgentModelBindingsV1() error = %v", err)
	}
	if err := store.applyModelPlanRevisionsV1(ctx); err != nil {
		t.Fatalf("retry applyModelPlanRevisionsV1() error = %v", err)
	}
	if _, err := store.GetModelPlan(ctx, "ws-retry", "mp-migrated-openai"); err != nil {
		t.Fatalf("GetModelPlan(retried migration) error = %v", err)
	}
}

func TestModelPlanFirstUseCandidatesMigrationRepairsLegacyModelPlansV1Schema(t *testing.T) {
	t.Parallel()

	ctx := context.Background()
	store := openTestSQLiteStore(t)
	if _, err := store.writeDB.ExecContext(ctx, `
DROP TABLE model_plan_first_use_candidates;
DELETE FROM tuttid_schema_migrations WHERE id = ?;
`, schemaMigrationModelPlanFirstUseCandidatesV1); err != nil {
		t.Fatalf("simulate legacy model plans v1 schema: %v", err)
	}

	legacyMarkerApplied, err := store.hasMigration(ctx, schemaMigrationModelPlansV1)
	if err != nil {
		t.Fatalf("inspect legacy model plans marker: %v", err)
	}
	if !legacyMarkerApplied {
		t.Fatal("legacy model_plans_v1 marker missing")
	}
	if err := store.Migrate(ctx); err != nil {
		t.Fatalf("Migrate() legacy model plans schema error = %v", err)
	}

	for kind, name := range map[string]string{
		"table": "model_plan_first_use_candidates",
		"index": "idx_model_plan_first_use_candidates_plan",
	} {
		var count int
		if err := store.writeDB.QueryRowContext(ctx, `
SELECT COUNT(*) FROM sqlite_master WHERE type = ? AND name = ?
`, kind, name).Scan(&count); err != nil {
			t.Fatalf("inspect repaired %s %q: %v", kind, name, err)
		}
		if count != 1 {
			t.Fatalf("repaired %s %q count = %d, want 1", kind, name, count)
		}
	}

	if err := store.Migrate(ctx); err != nil {
		t.Fatalf("Migrate() idempotent retry error = %v", err)
	}
	var markerCount int
	if err := store.writeDB.QueryRowContext(ctx, `
SELECT COUNT(*) FROM tuttid_schema_migrations WHERE id = ?
`, schemaMigrationModelPlanFirstUseCandidatesV1).Scan(&markerCount); err != nil {
		t.Fatalf("count first-use candidate migration marker: %v", err)
	}
	if markerCount != 1 {
		t.Fatalf("first-use candidate migration marker count = %d, want 1", markerCount)
	}
}

func TestListAgentModelBindingsByModelPolicy(t *testing.T) {
	t.Parallel()

	ctx := context.Background()
	store := openTestSQLiteStore(t)
	createModelPlanTestWorkspace(t, store, "ws-pol")
	now := time.UnixMilli(1700000000000).UTC()
	plan := modelplanbiz.Plan{
		ID: "mp-1", WorkspaceID: "ws-pol", Name: "Plan", TemplateKind: modelplanbiz.TemplateCustom,
		Protocol: modelplanbiz.ProtocolOpenAI, Models: []modelplanbiz.Model{{ID: "m", Name: "M"}},
		CreatedAt: now, UpdatedAt: now,
	}
	if err := store.PutModelPlan(ctx, plan); err != nil {
		t.Fatalf("PutModelPlan() error = %v", err)
	}
	if err := store.PutModelPolicy(ctx, modelpolicybiz.Policy{
		ID: "pol-1", WorkspaceID: "ws-pol", Name: "Careful", CreatedAt: now, UpdatedAt: now,
	}); err != nil {
		t.Fatalf("PutModelPolicy() error = %v", err)
	}
	if err := store.PutAgentModelBinding(ctx, modelbindingbiz.Binding{
		WorkspaceID: "ws-pol", AgentTargetID: agenttargetbiz.IDLocalCodex,
		ModelPlanID: "mp-1", ModelPolicyID: "pol-1", UpdatedAt: now,
	}); err != nil {
		t.Fatalf("PutAgentModelBinding() error = %v", err)
	}

	refs, err := store.ListAgentModelBindingsByModelPolicy(ctx, "ws-pol", "pol-1")
	if err != nil {
		t.Fatalf("ListAgentModelBindingsByModelPolicy() error = %v", err)
	}
	if len(refs) != 1 || refs[0].AgentTargetID != agenttargetbiz.IDLocalCodex || refs[0].ModelPolicyID != "pol-1" {
		t.Fatalf("refs = %#v, want the codex binding referencing pol-1", refs)
	}

	other, err := store.ListAgentModelBindingsByModelPolicy(ctx, "ws-pol", "pol-none")
	if err != nil {
		t.Fatalf("ListAgentModelBindingsByModelPolicy(unreferenced) error = %v", err)
	}
	if len(other) != 0 {
		t.Fatalf("refs for unreferenced policy = %#v, want none", other)
	}
}

func TestModelPlanBindingForeignKeys(t *testing.T) {
	t.Parallel()

	ctx := context.Background()
	store := openTestSQLiteStore(t)
	createModelPlanTestWorkspace(t, store, "ws-bindings")
	now := time.UnixMilli(1700000000000).UTC()
	plan := modelplanbiz.Plan{
		ID: "mp-bound", WorkspaceID: "ws-bindings", Name: "Bound", TemplateKind: modelplanbiz.TemplateCustom,
		Protocol: modelplanbiz.ProtocolOpenAI, Models: []modelplanbiz.Model{{ID: "gpt-test", Name: "GPT Test"}},
		Enabled: true, CreatedAt: now, UpdatedAt: now,
	}
	if err := store.PutModelPlan(ctx, plan); err != nil {
		t.Fatalf("PutModelPlan() error = %v", err)
	}
	binding := modelbindingbiz.Binding{
		WorkspaceID: "ws-bindings", AgentTargetID: agenttargetbiz.IDLocalCodex,
		ModelPlanID: "mp-bound", DefaultModel: "gpt-test", UpdatedAt: now,
	}
	if err := store.PutAgentModelBinding(ctx, binding); err != nil {
		t.Fatalf("PutAgentModelBinding() error = %v", err)
	}
	if err := store.DeleteModelPlan(ctx, "ws-bindings", "mp-bound"); !errors.Is(err, ErrModelPlanReferenced) {
		t.Fatalf("DeleteModelPlan(referenced) error = %v, want ErrModelPlanReferenced", err)
	}

	if err := store.DeleteAgentTarget(ctx, agenttargetbiz.IDLocalCodex); err != nil {
		t.Fatalf("DeleteAgentTarget() error = %v", err)
	}
	if _, err := store.GetAgentModelBinding(ctx, "ws-bindings", agenttargetbiz.IDLocalCodex); !errors.Is(err, ErrAgentModelBindingNotFound) {
		t.Fatalf("GetAgentModelBinding() after target delete error = %v, want ErrAgentModelBindingNotFound", err)
	}
	if err := store.DeleteModelPlan(ctx, "ws-bindings", "mp-bound"); err != nil {
		t.Fatalf("DeleteModelPlan(after target cascade) error = %v", err)
	}
}

func resetModelPlanMigrations(t *testing.T, store *SQLiteStore) {
	t.Helper()
	ctx := context.Background()
	for _, statement := range []string{
		`DROP TABLE agent_target_model_bindings`,
		`DROP TABLE model_plan_first_use_candidates`,
		`DROP TABLE model_plan_revisions`,
		`DROP TABLE model_plans`,
	} {
		if _, err := store.writeDB.ExecContext(ctx, statement); err != nil {
			t.Fatalf("reset model plan migration with %q error = %v", statement, err)
		}
	}
	if _, err := store.writeDB.ExecContext(ctx, `
DELETE FROM tuttid_schema_migrations
WHERE id IN (?, ?, ?, ?)
`, schemaMigrationModelPlansV1, schemaMigrationModelPlanFirstUseCandidatesV1, schemaMigrationAgentModelBindingsV1, schemaMigrationModelPlanRevisionsV1); err != nil {
		t.Fatalf("reset model plan migration markers error = %v", err)
	}
}