notarius/internal/framework/pipeline/debug.go

package pipeline

import (
	"context"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"fmt"
	"path"
	"regexp"
	"strings"
	"sync"
	"time"
	"unicode/utf8"

	"gitea.maximumdirect.net/eric/notarius/internal/core/artifacts"
	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
)

type DebugRecorder interface {
	Enabled() bool
	WriteJSON(name string, payload any) error
}

type noopDebugRecorder struct{}

func NoopDebugRecorder() DebugRecorder { return noopDebugRecorder{} }

func (noopDebugRecorder) Enabled() bool               { return false }
func (noopDebugRecorder) WriteJSON(string, any) error { return nil }
func debugPathComponent(value string) string {
	value = strings.TrimSpace(value)
	if value == "" {
		return "_"
	}
	var b strings.Builder
	for _, r := range value {
		switch {
		case r >= 'a' && r <= 'z':
			b.WriteRune(r)
		case r >= 'A' && r <= 'Z':
			b.WriteRune(r)
		case r >= '0' && r <= '9':
			b.WriteRune(r)
		case r == '-' || r == '_' || r == '.':
			b.WriteRune(r)
		default:
			b.WriteString(fmt.Sprintf("~%x", r))
		}
	}
	out := b.String()
	if out == "." || out == ".." || strings.Contains(out, "..") {
		return "_"
	}
	return out
}

type debugTimedEnvelope struct {
	Stage       string                  `json:"stage,omitempty"`
	LaneID      string                  `json:"lane_id,omitempty"`
	ModuleKey   string                  `json:"module_key,omitempty"`
	Attempt     int                     `json:"attempt,omitempty"`
	StartedAt   time.Time               `json:"started_at"`
	CompletedAt time.Time               `json:"completed_at"`
	DurationMS  int64                   `json:"duration_ms"`
	Payload     any                     `json:"payload,omitempty"`
	LLMCalls    []debugLLMCallReference `json:"llm_calls,omitempty"`
	Error       string                  `json:"error,omitempty"`
}

type debugBinaryEnvelope struct {
	ContentBase64 string              `json:"content_base64,omitempty"`
	ContentDigest string              `json:"content_digest,omitempty"`
	MediaType     string              `json:"media_type,omitempty"`
	Metadata      map[string]any      `json:"metadata,omitempty"`
	Warnings      []contracts.Warning `json:"warnings,omitempty"`
}

type debugRawPayload struct {
	Content debugBinaryEnvelope `json:"content"`
}

type debugSourceInput struct {
	SourceID string              `json:"source_id,omitempty"`
	Path     string              `json:"path,omitempty"`
	Raw      debugBinaryEnvelope `json:"raw,omitempty"`
	Options  map[string]any      `json:"options,omitempty"`
	Metadata map[string]any      `json:"metadata,omitempty"`
}

type debugSourceDocument struct {
	ID       string              `json:"id"`
	Kind     string              `json:"kind"`
	Format   string              `json:"format,omitempty"`
	Digest   string              `json:"digest,omitempty"`
	Units    []source.SourceUnit `json:"units,omitempty"`
	Metadata map[string]any      `json:"metadata,omitempty"`
}

type debugSourceChunk struct {
	ID          string              `json:"id"`
	SourceID    string              `json:"source_id"`
	Index       int                 `json:"index"`
	StartUnitID int                 `json:"start_unit_id"`
	EndUnitID   int                 `json:"end_unit_id"`
	Content     debugBinaryEnvelope `json:"content"`
	Units       []source.SourceUnit `json:"units,omitempty"`
	Metadata    map[string]any      `json:"metadata,omitempty"`
}

type debugExtractOutput struct {
	LaneID       string                   `json:"lane_id"`
	ExtractorKey string                   `json:"extractor_key"`
	SourceID     string                   `json:"source_id"`
	ChunkID      string                   `json:"chunk_id"`
	ChunkIndex   int                      `json:"chunk_index"`
	Schema       contracts.ResponseSchema `json:"schema,omitempty"`
	Payload      debugBinaryEnvelope      `json:"payload"`
}

type debugMergeOutput struct {
	LaneID    string                   `json:"lane_id"`
	MergerKey string                   `json:"merger_key"`
	SourceID  string                   `json:"source_id,omitempty"`
	Schema    contracts.ResponseSchema `json:"schema,omitempty"`
	Payload   debugBinaryEnvelope      `json:"payload"`
}

type debugNormalizeOutput struct {
	LaneID        string                   `json:"lane_id"`
	NormalizerKey string                   `json:"normalizer_key"`
	SourceID      string                   `json:"source_id,omitempty"`
	Schema        contracts.ResponseSchema `json:"schema,omitempty"`
	Payload       debugBinaryEnvelope      `json:"payload"`
}

type debugLLMInputMaterial struct {
	Name      string `json:"name"`
	MediaType string `json:"media_type,omitempty"`
	Content   string `json:"content_base64,omitempty"`
	Digest    string `json:"digest,omitempty"`
	OriginURI string `json:"origin_uri,omitempty"`
	SizeBytes int64  `json:"size_bytes,omitempty"`
}

type debugStructuredCompletionRequest struct {
	StageName     string                           `json:"stage_name"`
	PromptID      string                           `json:"prompt_id,omitempty"`
	PromptVersion string                           `json:"prompt_version,omitempty"`
	ProfileID     string                           `json:"profile_id,omitempty"`
	SessionID     string                           `json:"session_id,omitempty"`
	Inputs        map[string]debugLLMInputMaterial `json:"inputs,omitempty"`
	Vars          map[string]any                   `json:"vars,omitempty"`
}

type debugStructuredCompletionResponse struct {
	Content          string `json:"content,omitempty"`
	Provider         string `json:"provider,omitempty"`
	Model            string `json:"model,omitempty"`
	ProfileID        string `json:"profile_id,omitempty"`
	PromptTokens     int    `json:"prompt_tokens,omitempty"`
	CompletionTokens int    `json:"completion_tokens,omitempty"`
	TotalTokens      int    `json:"total_tokens,omitempty"`
}

type debugStructuredLLMCall struct {
	Request  debugStructuredCompletionRequest  `json:"request"`
	Response debugStructuredCompletionResponse `json:"response,omitempty"`
	Error    string                            `json:"error,omitempty"`
}

type debugLLMCallReference struct {
	CallID        string `json:"call_id"`
	CanonicalPath string `json:"canonical_path"`
	ScopedPath    string `json:"scoped_path,omitempty"`
	PromptID      string `json:"prompt_id,omitempty"`
	ProfileID     string `json:"profile_id,omitempty"`
	Error         bool   `json:"error,omitempty"`
}

type debugValidationRequest struct {
	Stage          string                   `json:"stage"`
	LaneID         string                   `json:"lane_id,omitempty"`
	ModuleKey      string                   `json:"module_key"`
	SourceID       string                   `json:"source_id,omitempty"`
	SessionID      string                   `json:"session_id,omitempty"`
	LLMProfile     string                   `json:"llm_profile,omitempty"`
	Options        map[string]any           `json:"options,omitempty"`
	Metadata       map[string]any           `json:"metadata,omitempty"`
	Schema         contracts.ResponseSchema `json:"schema,omitempty"`
	Payload        *debugBinaryEnvelope     `json:"payload,omitempty"`
	ChunkID        string                   `json:"chunk_id,omitempty"`
	ChunkIndex     int                      `json:"chunk_index,omitempty"`
	Chunk          *debugSourceChunk        `json:"chunk,omitempty"`
	Chunks         []debugSourceChunk       `json:"chunks,omitempty"`
	ExtractOutputs []debugExtractOutput     `json:"extract_outputs,omitempty"`
	MergeOutput    *debugMergeOutput        `json:"merge_output,omitempty"`
}

type debugValidationCall struct {
	ValidatorName string                     `json:"validator_name"`
	Request       debugValidationRequest     `json:"request"`
	Result        contracts.ValidationResult `json:"result,omitempty"`
	Error         string                     `json:"error,omitempty"`
}

type debugLLMClient struct {
	inner    contracts.StructuredLLMClient
	recorder DebugRecorder
	mu       sync.Mutex
	counter  int
}

type debugLLMScope struct {
	prefix string
	parent *debugLLMScope
	mu     sync.Mutex
	calls  []debugLLMCallReference
}

type debugLLMScopeContextKey struct{}

func wrapDebugLLMClient(client contracts.StructuredLLMClient, recorder DebugRecorder) contracts.StructuredLLMClient {
	if client == nil || recorder == nil || !recorder.Enabled() {
		return client
	}
	return &debugLLMClient{inner: client, recorder: recorder}
}

func (client *debugLLMClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
	client.mu.Lock()
	client.counter++
	callID := fmt.Sprintf("call-%04d", client.counter)
	client.mu.Unlock()

	started := time.Now().UTC()
	response, err := client.inner.CompleteStructured(ctx, req, out)
	completed := time.Now().UTC()
	payload := debugStructuredLLMCall{
		Request:  debugCompletionRequest(req),
		Response: debugCompletionResponse(response),
	}
	if err != nil {
		payload.Error = err.Error()
	}
	envelope := debugTimedEnvelope{
		Stage:       req.StageName,
		ModuleKey:   req.StageName,
		StartedAt:   started,
		CompletedAt: completed,
		DurationMS:  completed.Sub(started).Milliseconds(),
		Payload:     payload,
		Error:       payload.Error,
	}
	canonicalPath := path.Join("llm", callID+".json")
	writeErr := writeDebugTimed(client.recorder, canonicalPath, envelope)
	callRef := debugLLMCallReference{
		CallID:        callID,
		CanonicalPath: canonicalPath,
		PromptID:      req.PromptID,
		ProfileID:     debugFirstNonEmptyString(response.ProfileID, req.ProfileID),
		Error:         err != nil,
	}
	if scope := debugLLMScopeFromContext(ctx); scope != nil {
		scopedPath := path.Join(scope.prefix, "llm-"+callID+".json")
		callRef.ScopedPath = scopedPath
		scopedWriteErr := writeDebugTimed(client.recorder, scopedPath, envelope)
		if scopedWriteErr != nil {
			writeErr = errors.Join(writeErr, scopedWriteErr)
		}
		scope.record(callRef)
	}
	if err != nil {
		return response, err
	}
	if writeErr != nil {
		return response, fmt.Errorf("write LLM debug artifact: %w", writeErr)
	}
	return response, err
}

func withDebugLLMScope(ctx context.Context, prefix string) (context.Context, *debugLLMScope) {
	if ctx == nil {
		ctx = context.Background()
	}
	prefix = cleanDebugPath(prefix)
	scope := &debugLLMScope{
		prefix: prefix,
		parent: debugLLMScopeFromContext(ctx),
	}
	return context.WithValue(ctx, debugLLMScopeContextKey{}, scope), scope
}

func debugLLMScopeFromContext(ctx context.Context) *debugLLMScope {
	if ctx == nil {
		return nil
	}
	scope, _ := ctx.Value(debugLLMScopeContextKey{}).(*debugLLMScope)
	return scope
}

func (scope *debugLLMScope) record(ref debugLLMCallReference) {
	if scope == nil {
		return
	}
	scope.mu.Lock()
	scope.calls = append(scope.calls, ref)
	scope.mu.Unlock()
	if scope.parent != nil {
		scope.parent.record(ref)
	}
}

func (scope *debugLLMScope) references() []debugLLMCallReference {
	if scope == nil {
		return nil
	}
	scope.mu.Lock()
	defer scope.mu.Unlock()
	if len(scope.calls) == 0 {
		return nil
	}
	out := make([]debugLLMCallReference, len(scope.calls))
	copy(out, scope.calls)
	return out
}

func cleanDebugPath(value string) string {
	parts := strings.Split(path.Clean(strings.TrimSpace(value)), "/")
	out := make([]string, 0, len(parts))
	for _, part := range parts {
		out = append(out, debugPathComponent(part))
	}
	if len(out) == 0 {
		return "_"
	}
	return path.Join(out...)
}

func debugFirstNonEmptyString(values ...string) string {
	for _, value := range values {
		if strings.TrimSpace(value) != "" {
			return value
		}
	}
	return ""
}

func (client *debugLLMClient) LLMProfileManifests() []artifacts.LLMProfileManifest {
	provider, ok := client.inner.(contracts.LLMProfileManifestProvider)
	if !ok {
		return nil
	}
	return provider.LLMProfileManifests()
}

func writeDebugTimed(recorder DebugRecorder, name string, envelope debugTimedEnvelope) error {
	if recorder == nil || !recorder.Enabled() {
		return nil
	}
	if envelope.CompletedAt.IsZero() {
		envelope.CompletedAt = time.Now().UTC()
	}
	if envelope.StartedAt.IsZero() {
		envelope.StartedAt = envelope.CompletedAt
	}
	if envelope.DurationMS == 0 {
		envelope.DurationMS = envelope.CompletedAt.Sub(envelope.StartedAt).Milliseconds()
	}
	return recorder.WriteJSON(name, envelope)
}

func debugEnvelopeWithLLMCalls(envelope debugTimedEnvelope, scope *debugLLMScope) debugTimedEnvelope {
	if scope != nil {
		envelope.LLMCalls = scope.references()
	}
	return envelope
}

func debugContentEnvelope(content []byte, mediaType string, metadata map[string]any, warnings []contracts.Warning) debugBinaryEnvelope {
	content = redactSecretBytes(content)
	return debugBinaryEnvelope{
		ContentBase64: base64.StdEncoding.EncodeToString(content),
		ContentDigest: debugContentDigest(content),
		MediaType:     mediaType,
		Metadata:      redactSensitiveMap(metadata),
		Warnings:      cloneWarnings(warnings),
	}
}

func debugPayloadEnvelope(payload contracts.RawPayload) debugBinaryEnvelope {
	return debugContentEnvelope(payload.Content, payload.MediaType, payload.Metadata, payload.Warnings)
}

func debugSourceDocumentEnvelope(doc *source.SourceDocument) *debugSourceDocument {
	if doc == nil {
		return nil
	}
	return &debugSourceDocument{
		ID:       doc.ID,
		Kind:     doc.Kind,
		Format:   doc.Format,
		Digest:   doc.Digest,
		Units:    cloneSourceUnits(doc.Units),
		Metadata: redactSensitiveMap(doc.Metadata),
	}
}

func debugSourceChunkEnvelope(chunk contracts.SourceChunk) debugSourceChunk {
	return debugSourceChunk{
		ID:          chunk.ID,
		SourceID:    chunk.SourceID,
		Index:       chunk.Index,
		StartUnitID: chunk.StartUnitID,
		EndUnitID:   chunk.EndUnitID,
		Content:     debugContentEnvelope(chunk.Content, chunk.MediaType, chunk.Metadata, nil),
		Units:       cloneSourceUnits(chunk.Units),
		Metadata:    redactSensitiveMap(chunk.Metadata),
	}
}

func debugSourceChunkEnvelopes(chunks []contracts.SourceChunk) []debugSourceChunk {
	if len(chunks) == 0 {
		return nil
	}
	out := make([]debugSourceChunk, 0, len(chunks))
	for _, chunk := range chunks {
		out = append(out, debugSourceChunkEnvelope(chunk))
	}
	return out
}

func debugExtractOutputEnvelope(output contracts.ExtractOutput) debugExtractOutput {
	output.Schema.JSONSchema = nil
	return debugExtractOutput{
		LaneID:       output.LaneID,
		ExtractorKey: output.ExtractorKey,
		SourceID:     output.SourceID,
		ChunkID:      output.ChunkID,
		ChunkIndex:   output.ChunkIndex,
		Schema:       output.Schema,
		Payload:      debugPayloadEnvelope(output.Payload),
	}
}

func debugExtractOutputEnvelopes(outputs []contracts.ExtractOutput) []debugExtractOutput {
	if len(outputs) == 0 {
		return nil
	}
	out := make([]debugExtractOutput, 0, len(outputs))
	for _, output := range outputs {
		out = append(out, debugExtractOutputEnvelope(output))
	}
	return out
}

func debugMergeOutputEnvelope(output contracts.MergeOutput) debugMergeOutput {
	output.Schema.JSONSchema = nil
	return debugMergeOutput{
		LaneID:    output.LaneID,
		MergerKey: output.MergerKey,
		SourceID:  output.SourceID,
		Schema:    output.Schema,
		Payload:   debugPayloadEnvelope(output.Payload),
	}
}

func debugNormalizeOutputEnvelope(output contracts.NormalizeOutput) debugNormalizeOutput {
	output.Schema.JSONSchema = nil
	return debugNormalizeOutput{
		LaneID:        output.LaneID,
		NormalizerKey: output.NormalizerKey,
		SourceID:      output.SourceID,
		Schema:        output.Schema,
		Payload:       debugPayloadEnvelope(output.Payload),
	}
}

func debugNormalizeOutputEnvelopes(outputs []contracts.NormalizeOutput) []debugNormalizeOutput {
	if len(outputs) == 0 {
		return nil
	}
	out := make([]debugNormalizeOutput, 0, len(outputs))
	for _, output := range outputs {
		out = append(out, debugNormalizeOutputEnvelope(output))
	}
	return out
}

type debugOutputFile struct {
	Name        string              `json:"name"`
	ContentType string              `json:"content_type,omitempty"`
	Content     debugBinaryEnvelope `json:"content"`
}

func debugOutputFiles(files []contracts.OutputFile) []debugOutputFile {
	if len(files) == 0 {
		return nil
	}
	out := make([]debugOutputFile, 0, len(files))
	for _, file := range files {
		out = append(out, debugOutputFile{
			Name:        file.Name,
			ContentType: file.ContentType,
			Content:     debugContentEnvelope(file.Bytes, file.ContentType, nil, nil),
		})
	}
	return out
}

func debugCompletionRequest(req contracts.StructuredCompletionRequest) debugStructuredCompletionRequest {
	inputs := make(map[string]debugLLMInputMaterial, len(req.Inputs))
	for key, material := range req.Inputs {
		inputs[key] = debugLLMInputMaterial{
			Name:      material.Name,
			MediaType: material.MediaType,
			Content:   base64.StdEncoding.EncodeToString(redactSecretBytes(material.Content)),
			Digest:    material.Digest,
			OriginURI: material.OriginURI,
			SizeBytes: material.SizeBytes,
		}
	}
	if len(inputs) == 0 {
		inputs = nil
	}
	return debugStructuredCompletionRequest{
		StageName:     req.StageName,
		PromptID:      req.PromptID,
		PromptVersion: req.PromptVersion,
		ProfileID:     req.ProfileID,
		SessionID:     req.SessionID,
		Inputs:        inputs,
		Vars:          redactSensitiveMap(req.Vars),
	}
}

func debugCompletionResponse(response contracts.StructuredCompletionResponse) debugStructuredCompletionResponse {
	return debugStructuredCompletionResponse{
		Content:          base64.StdEncoding.EncodeToString(redactSecretBytes(response.Content)),
		Provider:         response.Provider,
		Model:            response.Model,
		ProfileID:        response.ProfileID,
		PromptTokens:     response.PromptTokens,
		CompletionTokens: response.CompletionTokens,
		TotalTokens:      response.TotalTokens,
	}
}

func debugValidationRequestEnvelope(req contracts.ValidationRequest) debugValidationRequest {
	req.Schema.JSONSchema = nil
	out := debugValidationRequest{
		Stage:      req.Stage,
		LaneID:     req.LaneID,
		ModuleKey:  req.ModuleKey,
		SourceID:   req.SourceID,
		SessionID:  req.SessionID,
		LLMProfile: req.LLMProfile,
		Options:    redactSensitiveMap(req.Options),
		Metadata:   redactSensitiveMap(req.Metadata),
		Schema:     req.Schema,
		ChunkID:    req.ChunkID,
		ChunkIndex: req.ChunkIndex,
	}
	payload := debugPayloadEnvelope(req.Payload)
	out.Payload = &payload
	if req.Chunk != nil {
		chunk := debugSourceChunkEnvelope(*req.Chunk)
		out.Chunk = &chunk
	}
	out.Chunks = debugSourceChunkEnvelopes(req.Chunks)
	out.ExtractOutputs = debugExtractOutputEnvelopes(req.ExtractOutputs)
	if len(req.MergeOutput.Payload.Content) > 0 || req.MergeOutput.LaneID != "" {
		merge := debugMergeOutputEnvelope(req.MergeOutput)
		out.MergeOutput = &merge
	}
	return out
}

func debugValidationResultEnvelope(result contracts.ValidationResult) contracts.ValidationResult {
	result.Message = string(redactSecretBytes([]byte(result.Message)))
	result.DiagnosticArtifactPath = string(redactSecretBytes([]byte(result.DiagnosticArtifactPath)))
	for i := range result.Warnings {
		result.Warnings[i].Message = string(redactSecretBytes([]byte(result.Warnings[i].Message)))
	}
	return result
}

func debugRejectedOutputEnvelope(rejected contracts.RejectedOutput) contracts.RejectedOutput {
	rejected.Message = string(redactSecretBytes([]byte(rejected.Message)))
	rejected.DiagnosticArtifactPath = string(redactSecretBytes([]byte(rejected.DiagnosticArtifactPath)))
	return rejected
}

func debugRejectedOutputPtr(rejected *contracts.RejectedOutput) any {
	if rejected == nil {
		return nil
	}
	out := debugRejectedOutputEnvelope(*rejected)
	return out
}

func debugRejectedOutputEnvelopes(rejected []contracts.RejectedOutput) []contracts.RejectedOutput {
	if len(rejected) == 0 {
		return nil
	}
	out := make([]contracts.RejectedOutput, 0, len(rejected))
	for _, item := range rejected {
		out = append(out, debugRejectedOutputEnvelope(item))
	}
	return out
}

func debugContentDigest(content []byte) string {
	sum := sha256.Sum256(content)
	return "sha256:" + hex.EncodeToString(sum[:])
}

var secretPatterns = []*regexp.Regexp{
	regexp.MustCompile(`(?i)bearer\s+[a-z0-9._~+/=-]{8,}`),
	regexp.MustCompile(`(?i)sk-[a-z0-9_-]{8,}`),
}

func redactSecretBytes(content []byte) []byte {
	if len(content) == 0 || !utf8.Valid(content) {
		return append([]byte(nil), content...)
	}
	text := string(content)
	for _, pattern := range secretPatterns {
		text = pattern.ReplaceAllString(text, "[REDACTED]")
	}
	return []byte(text)
}

func redactSensitiveMap(values map[string]any) map[string]any {
	if len(values) == 0 {
		return nil
	}
	out := make(map[string]any, len(values))
	for key, value := range values {
		if sensitiveKey(key) {
			out[key] = "[REDACTED]"
			continue
		}
		out[key] = redactSensitiveValue(value)
	}
	return out
}

func redactSensitiveValue(value any) any {
	switch typed := value.(type) {
	case string:
		return string(redactSecretBytes([]byte(typed)))
	case map[string]any:
		return redactSensitiveMap(typed)
	case map[string]string:
		out := make(map[string]string, len(typed))
		for key, value := range typed {
			if sensitiveKey(key) {
				out[key] = "[REDACTED]"
			} else {
				out[key] = string(redactSecretBytes([]byte(value)))
			}
		}
		return out
	default:
		return value
	}
}

func sensitiveKey(key string) bool {
	key = strings.ToLower(key)
	return strings.Contains(key, "api_key") ||
		strings.Contains(key, "apikey") ||
		strings.Contains(key, "authorization") ||
		strings.Contains(key, "bearer") ||
		strings.Contains(key, "password") ||
		strings.Contains(key, "secret") ||
		strings.Contains(key, "token")
}