xinmi/internal/workbench/diagnostics_test.go

package workbench

import (
	"context"
	"errors"
	"fmt"
	"io"
	"strings"
	"testing"
	"time"

	"github.com/Tencent/WeKnora/internal/types"
)

func TestJobErrorDetailsNeverExposeProviderBody(t *testing.T) {
	for _, code := range []int{400, 401, 403, 404, 429, 500, 524} {
		err := fmt.Errorf("create chat completion: error, status code: %d, secret-api-key and private prompt", code)
		class, message := jobErrorDetails(err)
		if class != fmt.Sprintf("provider_http_%d", code) || !strings.Contains(message, fmt.Sprintf("HTTP %d", code)) {
			t.Fatalf("unexpected diagnostic %s %s", class, message)
		}
		if strings.Contains(message, "secret") || strings.Contains(message, "private") {
			t.Fatal("provider response leaked")
		}
	}
	for _, tc := range []struct {
		err   error
		class string
	}{
		{context.DeadlineExceeded, "timeout"}, {context.Canceled, "cancelled"}, {io.ErrUnexpectedEOF, "provider_disconnected"}, {errors.New("raw private-prompt"), "internal_error"},
	} {
		class, message := jobErrorDetails(fmt.Errorf("call failed: %w", tc.err))
		if class != tc.class || strings.Contains(message, "private-prompt") {
			t.Fatalf("unexpected diagnostic %s %s", class, message)
		}
	}
}

func TestFailedJobKeepsProgressAndDoesNotCallDBFaultForbidden(t *testing.T) {
	for _, dbFault := range []bool{false, true} {
		s := sqliteWorkbench(t)
		if !dbFault {
			if err := s.db.AutoMigrate(&types.TenantMember{}); err != nil {
				t.Fatal(err)
			}
		}
		until := time.Now().Add(time.Minute)
		job := Job{ID: "diagnostic-job", TenantID: 1, OwnerID: "owner", ResourceID: "resource", Kind: "training-maps", Status: "running", Progress: 5, LeaseToken: "lease", LeaseUntil: &until}
		if err := s.db.Create(&job).Error; err != nil {
			t.Fatal(err)
		}
		if err := s.db.Model(&Job{}).Where("id = ?", job.ID).Update("progress", 45).Error; err != nil {
			t.Fatal(err)
		}
		s.run(context.Background(), &job)
		var stored Job
		if err := s.db.First(&stored, "id = ?", job.ID).Error; err != nil {
			t.Fatal(err)
		}
		if stored.Status != "failed" || stored.Progress != 45 {
			t.Fatalf("lost failure state/progress: %+v", stored)
		}
		if dbFault && stored.Error == ErrForbidden.Error() {
			t.Fatal("database failure reported as missing permission")
		}
		if !dbFault && stored.Error != ErrForbidden.Error() {
			t.Fatal("missing member was allowed")
		}
		var logged OperationLog
		if err := s.db.Where("kind = ? AND trace_id = ?", "system", job.ID).Take(&logged).Error; err != nil {
			t.Fatal("background failure was not persisted:", err)
		}
		if logged.TenantID != 1 || logged.OwnerID != "owner" || logged.ResourceID != "resource" || logged.Metadata["source"] != "background_job" {
			t.Fatalf("failure attribution lost: %+v", logged)
		}
		if dbFault && logged.Level != "error" || !dbFault && logged.Level != "warn" {
			t.Fatalf("wrong failure classification: %+v", logged)
		}
	}
}

func TestJobFailurePersistsSafeStructuredDiagnostic(t *testing.T) {
	s := sqliteWorkbench(t)
	job := Job{ID: "format-job", TenantID: 7, OwnerID: "author", ResourceID: "map", Action: "generate", Kind: "training-maps", Status: "failed", Progress: 70, Generation: JSON{"model_id": "model", "agent_id": "agent", "attempt": 2, "untrusted": "private prompt"}}
	if err := s.db.Create(&job).Error; err != nil {
		t.Fatal(err)
	}
	_, err := parseJSON[mapDaysOutput](`{"days":"private-body-secret"}`)
	s.recordJobFailure(context.Background(), &job, "execute", err)
	var row OperationLog
	if err := s.db.First(&row).Error; err != nil {
		t.Fatal(err)
	}
	if row.Metadata["error_class"] != "structured_output_invalid" || row.Metadata["structured_output"] == nil || row.Metadata["model_id"] != "model" || row.Level != "error" || row.TenantID != 7 {
		t.Fatalf("diagnostic was not retained: %+v", row)
	}
	for _, secret := range []string{"private-body-secret", "private prompt", "untrusted"} {
		if strings.Contains(compact(row), secret) {
			t.Fatalf("private data leaked: %s", secret)
		}
	}
}

func TestRecoverExpiredJobsRecordsEachFailureOnce(t *testing.T) {
	s := sqliteWorkbench(t)
	expired := time.Now().Add(-time.Minute)
	for _, job := range []Job{
		{ID: "expired", TenantID: 1, OwnerID: "alice", Status: "running", LeaseUntil: &expired},
		{ID: "cancelled", TenantID: 1, OwnerID: "alice", Status: "cancelled", LeaseUntil: &expired},
	} {
		if err := s.db.Create(&job).Error; err != nil {
			t.Fatal(err)
		}
	}
	for i := 0; i < 2; i++ {
		if err := s.recoverExpiredJobs(context.Background(), time.Now()); err != nil {
			t.Fatal(err)
		}
	}
	var rows []OperationLog
	if err := s.db.Find(&rows).Error; err != nil {
		t.Fatal(err)
	}
	if len(rows) != 1 || rows[0].TraceID != "expired" || rows[0].Metadata["stage"] != "lease_recovery" {
		t.Fatalf("duplicate or missing lease recovery record: %+v", rows)
	}
}

func TestJobDiagnosticRedactsCredentials(t *testing.T) {
	err := fmt.Errorf("%w: 模型配置 apikey=private-value Authorization: Bearer private-token sk-private12345", ErrInvalid)
	_, message := jobErrorDetails(err)
	if strings.Contains(message, "private") || !strings.Contains(message, "模型配置") {
		t.Fatalf("credential redaction lost useful context or exposed a credential: %s", message)
	}
}

func TestTimeoutDiagnosticsIdentifyPhaseWithoutProviderBody(t *testing.T) {
	for _, phase := range []string{"model", "retrieval", "reading"} {
		err := phaseFailure(fmt.Errorf("private-provider-body: %w", context.DeadlineExceeded), phase, time.Now().Add(-3*time.Second))
		if !errors.Is(err, context.DeadlineExceeded) {
			t.Fatal("timeout no longer unwraps")
		}
		metadata := jobErrorMetadata(err)
		if metadata["error_class"] != phase+"_timeout" || metadata["phase"] != phase || metadata["elapsed_ms"].(int64) < 3000 || strings.Contains(compact(metadata), "private-provider-body") {
			t.Fatalf("invalid phase diagnostic: %+v", metadata)
		}
		if phaseFailure(err, "other", time.Now()) != err {
			t.Fatal("outer wrapper lost specific timeout phase")
		}
	}
	if phaseFailure(context.Canceled, "model", time.Now()) != context.Canceled {
		t.Fatal("cancellation was mislabeled timeout")
	}
}

func TestJobFailureKeepsTimeoutPhaseDuringStructuredRepair(t *testing.T) {
	for _, tc := range []struct {
		name, phase, class string
		err                error
	}{
		{"model timeout", "model", "model_timeout", phaseFailure(context.DeadlineExceeded, "model", time.Now().Add(-time.Second))},
		{"other failure", "structured_output_repair", "internal_error", errors.New("private provider response")},
	} {
		t.Run(tc.name, func(t *testing.T) {
			s := sqliteWorkbench(t)
			job := Job{ID: "repair-job", TenantID: 7, OwnerID: "author", ResourceID: "map", Action: "generate", Kind: "training-maps", Status: "failed", Generation: JSON{"phase": "structured_output_repair", "model_id": "model", "agent_id": "agent", "attempt": 2}}
			if err := s.db.Create(&job).Error; err != nil {
				t.Fatal(err)
			}
			s.recordJobFailure(context.Background(), &job, "execute", tc.err)
			var row OperationLog
			if err := s.db.First(&row).Error; err != nil {
				t.Fatal(err)
			}
			if row.Metadata["phase"] != tc.phase || row.Metadata["error_class"] != tc.class || row.Metadata["generation_phase"] != "structured_output_repair" {
				t.Fatalf("error phase or generation operation was lost: %+v", row.Metadata)
			}
			if row.Metadata["model_id"] != "model" || row.Metadata["agent_id"] != "agent" || row.Metadata["attempt"] != float64(2) {
				t.Fatalf("generation attribution was lost: %+v", row.Metadata)
			}
			if tc.phase == "model" && row.Metadata["elapsed_ms"].(float64) < 1000 {
				t.Fatalf("timeout elapsed time was lost: %+v", row.Metadata)
			}
		})
	}
}