claude-code-proxy/proxy/internal/service/storage_payload.go

package service

import (
	"encoding/json"
	"log"
	"strings"

	"github.com/seifghazi/claude-code-monitor/internal/config"
	"github.com/seifghazi/claude-code-monitor/internal/model"
)

const redactionPlaceholder = "[REDACTED]"

// maxStoredBodyBytes is the maximum serialized size of a request body stored in the DB.
// Bodies larger than this (e.g. 1M context payloads) are replaced with a metadata summary.
const maxStoredBodyBytes = 512 * 1024 // 512 KB

func prepareRequestBodyForStorage(cfg *config.StorageConfig, body interface{}) (interface{}, error) {
	if shouldSuppressBodies(cfg) {
		return storageBodyPlaceholder("metadata_only"), nil
	}
	if cfg != nil && !cfg.CaptureRequestBody {
		return storageBodyPlaceholder("request_body_disabled"), nil
	}

	normalized, err := normalizeJSONValue(body)
	if err != nil {
		return nil, err
	}

	redacted := redactJSONValue(normalized, redactedFieldSet(redactedFields(cfg)))

	// Check serialized size; if too large, store a lightweight summary instead
	data, err := json.Marshal(redacted)
	if err != nil {
		return redacted, nil
	}
	if len(data) > maxStoredBodyBytes {
		return truncatedBodySummary(redacted, len(data)), nil
	}

	return redacted, nil
}

// truncatedBodySummary extracts key metadata from an oversized request body
// so the DB row stays small while retaining useful diagnostic info.
func truncatedBodySummary(body interface{}, originalBytes int) map[string]interface{} {
	summary := map[string]interface{}{
		"_truncated":     true,
		"_original_bytes": originalBytes,
	}
	if m, ok := body.(map[string]interface{}); ok {
		if v, ok := m["model"]; ok {
			summary["model"] = v
		}
		if v, ok := m["stream"]; ok {
			summary["stream"] = v
		}
		if v, ok := m["max_tokens"]; ok {
			summary["max_tokens"] = v
		}
		if msgs, ok := m["messages"].([]interface{}); ok {
			summary["message_count"] = len(msgs)
		}
		if sys, ok := m["system"].([]interface{}); ok {
			summary["system_count"] = len(sys)
		}
		if tools, ok := m["tools"].([]interface{}); ok {
			summary["tool_count"] = len(tools)
		}
	}
	return summary
}

func prepareResponseForStorage(cfg *config.StorageConfig, logger *log.Logger, response *model.ResponseLog) (*model.ResponseLog, error) {
	if response == nil {
		return nil, nil
	}

	clone := *response
	if shouldSuppressBodies(cfg) || (cfg != nil && !cfg.CaptureResponseBody) {
		clone.Body = nil
		clone.BodyText = ""
		clone.StreamingChunks = nil
		clone.ChunkTimings = nil
		return &clone, nil
	}

	if len(clone.Body) > 0 {
		sanitizedBody, err := sanitizeRawJSON(clone.Body, redactedFieldSet(redactedFields(cfg)))
		if err != nil {
			if logger != nil {
				logger.Printf("Warning: failed to redact response body: %v", err)
			}
		} else {
			clone.Body = sanitizedBody
		}
		// Cap stored response body size
		if len(clone.Body) > maxStoredBodyBytes {
			clone.Body = json.RawMessage(`{"_truncated":true}`)
		}
	}

	// Cap stored streaming chunks to avoid huge DB rows on long streams
	const maxStoredChunks = 500
	if len(clone.StreamingChunks) > maxStoredChunks {
		clone.StreamingChunks = clone.StreamingChunks[:maxStoredChunks]
	}
	if len(clone.ChunkTimings) > maxStoredChunks {
		clone.ChunkTimings = clone.ChunkTimings[:maxStoredChunks]
	}

	return &clone, nil
}

func shouldSuppressBodies(cfg *config.StorageConfig) bool {
	return cfg != nil && cfg.MetadataOnly
}

func redactedFields(cfg *config.StorageConfig) []string {
	if cfg == nil {
		return nil
	}
	return cfg.RedactedFields
}

func normalizeJSONValue(value interface{}) (interface{}, error) {
	if value == nil {
		return nil, nil
	}

	data, err := json.Marshal(value)
	if err != nil {
		return nil, err
	}

	var normalized interface{}
	if err := json.Unmarshal(data, &normalized); err != nil {
		return nil, err
	}

	return normalized, nil
}

func sanitizeRawJSON(raw json.RawMessage, redacted map[string]struct{}) (json.RawMessage, error) {
	if len(raw) == 0 {
		return raw, nil
	}

	var value interface{}
	if err := json.Unmarshal(raw, &value); err != nil {
		return raw, err
	}

	sanitized := redactJSONValue(value, redacted)
	data, err := json.Marshal(sanitized)
	if err != nil {
		return raw, err
	}

	return json.RawMessage(data), nil
}

func redactJSONValue(value interface{}, redacted map[string]struct{}) interface{} {
	switch typed := value.(type) {
	case map[string]interface{}:
		result := make(map[string]interface{}, len(typed))
		for key, child := range typed {
			if _, ok := redacted[strings.ToLower(key)]; ok {
				result[key] = redactionPlaceholder
				continue
			}
			result[key] = redactJSONValue(child, redacted)
		}
		return result
	case []interface{}:
		result := make([]interface{}, len(typed))
		for i, child := range typed {
			result[i] = redactJSONValue(child, redacted)
		}
		return result
	default:
		return value
	}
}

func storageBodyPlaceholder(mode string) map[string]interface{} {
	return map[string]interface{}{
		"_storage_mode": mode,
	}
}

func redactedFieldSet(fields []string) map[string]struct{} {
	set := make(map[string]struct{}, len(fields))
	for _, field := range fields {
		field = strings.TrimSpace(strings.ToLower(field))
		if field == "" {
			continue
		}
		set[field] = struct{}{}
	}
	return set
}