notarius/internal/framework/pipeline/debug.go

package pipeline

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"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
	WriteBytes(name string, data []byte) error
}

type noopDebugRecorder struct{}

func NoopDebugRecorder() DebugRecorder { return noopDebugRecorder{} }

func (noopDebugRecorder) Enabled() bool                   { return false }
func (noopDebugRecorder) WriteJSON(string, any) error     { return nil }
func (noopDebugRecorder) WriteBytes(string, []byte) 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 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"`
	Ref      source.SourceRef    `json:"ref"`
	Content  debugBinaryEnvelope `json:"content"`
	Units    []source.SourceUnit `json:"units,omitempty"`
	Metadata map[string]any      `json:"metadata,omitempty"`
}

type debugSerializedOutput struct {
	LaneID        string                   `json:"lane_id"`
	NormalizerKey string                   `json:"normalizer_key"`
	SourceID      string                   `json:"source_id,omitempty"`
	Kind          contracts.ArtifactKind   `json:"artifact_kind,omitempty"`
	Schema        contracts.ArtifactSchema `json:"schema"`
	SchemaDigest  string                   `json:"schema_digest"`
	Content       debugBinaryEnvelope      `json:"content"`
}

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 debugLLMPromptArtifact struct {
	CallID string                    `json:"call_id"`
	Prompt *contracts.LLMDebugPrompt `json:"prompt,omitempty"`
}

type debugLLMResponseArtifact struct {
	CallID      string                             `json:"call_id"`
	Response    *contracts.LLMDebugResponse        `json:"response,omitempty"`
	Fallback    *debugStructuredCompletionResponse `json:"fallback,omitempty"`
	ContentPath string                             `json:"content_path,omitempty"`
	Error       string                             `json:"error,omitempty"`
}

type debugLLMCallReference struct {
	CallID              string `json:"call_id"`
	PromptPath          string `json:"prompt_path,omitempty"`
	ResponsePath        string `json:"response_path"`
	ResponseContentPath string `json:"response_content_path,omitempty"`
	PromptID            string `json:"prompt_id,omitempty"`
	ProfileID           string `json:"profile_id,omitempty"`
	Model               string `json:"model,omitempty"`
	Error               bool   `json:"error,omitempty"`
}

type debugValidationCall struct {
	ValidatorName string                     `json:"validator_name"`
	Request       any                        `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
	}
	if _, ok := client.(*debugLLMClient); ok {
		return client
	}
	return &debugLLMClient{inner: client, recorder: recorder}
}

// WithDebugLLMRecording decorates a shared LLM client so calls made by
// construction-injected modules participate in the run's debug recording.
func WithDebugLLMRecording(client contracts.StructuredLLMClient, recorder DebugRecorder) contracts.StructuredLLMClient {
	return wrapDebugLLMClient(client, recorder)
}

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

	started := time.Now().UTC()
	response, err := client.inner.CompleteStructured(ctx, req, out)
	completed := time.Now().UTC()
	errorText := ""
	if err != nil {
		errorText = err.Error()
	}

	scopePrefix := cleanDebugPath(req.StageName)
	if scopePrefix == "_" {
		scopePrefix = "llm"
	}
	if scope := debugLLMScopeFromContext(ctx); scope != nil {
		scopePrefix = scope.prefix
	}
	promptPath := ""
	var writeErr error
	if response.Debug != nil && response.Debug.Prompt != nil {
		promptPath = path.Join(scopePrefix, "prompt-"+callID+".json")
		writeErr = errors.Join(writeErr, writeDebugTimed(client.recorder, promptPath, debugTimedEnvelope{
			Stage:       req.StageName,
			ModuleKey:   req.StageName,
			StartedAt:   started,
			CompletedAt: completed,
			DurationMS:  completed.Sub(started).Milliseconds(),
			Payload: debugLLMPromptArtifact{
				CallID: callID,
				Prompt: response.Debug.Prompt,
			},
		}))
	}
	responsePath := path.Join(scopePrefix, "response-"+callID+".json")
	responseMaterial := debugResponseMaterial(response)
	fallbackMaterial := debugCompletionFallback(response)
	responseContentPath, responseForArtifact, fallbackForArtifact, contentErr := writeDebugResponseContent(client.recorder, scopePrefix, callID, responseMaterial, fallbackMaterial)
	writeErr = errors.Join(writeErr, contentErr)
	writeErr = errors.Join(writeErr, writeDebugTimed(client.recorder, responsePath, debugTimedEnvelope{
		Stage:       req.StageName,
		ModuleKey:   req.StageName,
		StartedAt:   started,
		CompletedAt: completed,
		DurationMS:  completed.Sub(started).Milliseconds(),
		Payload: debugLLMResponseArtifact{
			CallID:      callID,
			Response:    responseForArtifact,
			Fallback:    fallbackForArtifact,
			ContentPath: responseContentPath,
			Error:       errorText,
		},
		Error: errorText,
	}))
	callRef := debugLLMCallReference{
		CallID:              callID,
		PromptPath:          promptPath,
		ResponsePath:        responsePath,
		ResponseContentPath: responseContentPath,
		PromptID:            req.PromptID,
		ProfileID:           debugFirstNonEmptyString(response.ProfileID, req.ProfileID),
		Model:               debugFirstNonEmptyString(response.Model, debugResponseModel(response)),
		Error:               err != nil,
	}
	if scope := debugLLMScopeFromContext(ctx); scope != nil {
		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 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 source.Chunk) debugSourceChunk {
	return debugSourceChunk{
		ID:       chunk.ID,
		SourceID: chunk.SourceID,
		Index:    chunk.Index,
		Ref:      chunk.Ref,
		Content:  debugContentEnvelope(chunk.Content, chunk.MediaType, chunk.Metadata, nil),
		Units:    cloneSourceUnits(chunk.Units),
		Metadata: redactSensitiveMap(chunk.Metadata),
	}
}

func debugSourceChunkEnvelopes(chunks []source.Chunk) []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 debugSerializedOutputEnvelope(output contracts.SerializedOutput) debugSerializedOutput {
	schema := contracts.CloneArtifactSchema(output.Artifact.Schema)
	digest := contracts.DigestArtifactSchema(schema)
	schema.JSONSchema = nil
	content := debugContentEnvelope(output.Artifact.Content, output.Artifact.MediaType, output.Artifact.Metadata, nil)
	content.ContentDigest = debugContentDigest(output.Artifact.Content)
	return debugSerializedOutput{
		LaneID: output.LaneID, NormalizerKey: output.NormalizerKey, SourceID: output.SourceID,
		Kind: output.Artifact.Kind, Schema: schema, SchemaDigest: digest,
		Content: content,
	}
}

func debugSerializedOutputEnvelopes(outputs []contracts.SerializedOutput) []debugSerializedOutput {
	if len(outputs) == 0 {
		return nil
	}
	out := make([]debugSerializedOutput, 0, len(outputs))
	for _, output := range outputs {
		out = append(out, debugSerializedOutputEnvelope(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:          string(response.Content),
		Provider:         response.Provider,
		Model:            response.Model,
		ProfileID:        response.ProfileID,
		PromptTokens:     response.PromptTokens,
		CompletionTokens: response.CompletionTokens,
		TotalTokens:      response.TotalTokens,
	}
}

func debugCompletionFallback(response contracts.StructuredCompletionResponse) *debugStructuredCompletionResponse {
	if response.Debug != nil && response.Debug.Response != nil {
		return nil
	}
	fallback := debugCompletionResponse(response)
	if fallback.Content == "" &&
		fallback.Provider == "" &&
		fallback.Model == "" &&
		fallback.ProfileID == "" &&
		fallback.PromptTokens == 0 &&
		fallback.CompletionTokens == 0 &&
		fallback.TotalTokens == 0 {
		return nil
	}
	return &fallback
}

func writeDebugResponseContent(recorder DebugRecorder, scopePrefix string, callID string, response *contracts.LLMDebugResponse, fallback *debugStructuredCompletionResponse) (string, *contracts.LLMDebugResponse, *debugStructuredCompletionResponse, error) {
	responseCopy := cloneDebugResponseWithoutContent(response)
	fallbackCopy := cloneDebugFallbackWithoutContent(fallback)
	content := ""
	if response != nil {
		content = response.Content
	}
	if content == "" && fallback != nil {
		content = fallback.Content
	}
	if content == "" {
		return "", responseCopy, fallbackCopy, nil
	}

	contentPath, data := debugResponseContentFile(scopePrefix, callID, content)
	if recorder == nil || !recorder.Enabled() {
		return contentPath, responseCopy, fallbackCopy, nil
	}
	if err := recorder.WriteBytes(contentPath, data); err != nil {
		return contentPath, responseCopy, fallbackCopy, err
	}
	return contentPath, responseCopy, fallbackCopy, nil
}

func cloneDebugResponseWithoutContent(response *contracts.LLMDebugResponse) *contracts.LLMDebugResponse {
	if response == nil {
		return nil
	}
	clone := *response
	clone.Content = ""
	return &clone
}

func cloneDebugFallbackWithoutContent(fallback *debugStructuredCompletionResponse) *debugStructuredCompletionResponse {
	if fallback == nil {
		return nil
	}
	clone := *fallback
	clone.Content = ""
	return &clone
}

func debugResponseContentFile(scopePrefix string, callID string, content string) (string, []byte) {
	raw := []byte(content)
	if json.Valid(raw) {
		var formatted bytes.Buffer
		if err := json.Indent(&formatted, raw, "", "  "); err == nil {
			formatted.WriteByte('\n')
			return path.Join(scopePrefix, "response-content-"+callID+".json"), formatted.Bytes()
		}
	}
	return path.Join(scopePrefix, "response-content-"+callID+".txt"), raw
}

func debugResponseMaterial(response contracts.StructuredCompletionResponse) *contracts.LLMDebugResponse {
	if response.Debug == nil {
		return nil
	}
	return response.Debug.Response
}

func debugResponseModel(response contracts.StructuredCompletionResponse) string {
	if response.Debug == nil || response.Debug.Response == nil {
		return ""
	}
	return response.Debug.Response.ModelName
}

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 debugWarningEnvelopes(warnings []contracts.Warning) []contracts.Warning {
	out := cloneWarnings(warnings)
	for i := range out {
		out[i].Message = string(redactSecretBytes([]byte(out[i].Message)))
	}
	return out
}

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