spark-mcp/internal/logging/tool_call.go

package logging

import (
	"log/slog"
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"time"
)

const toolCallResultMaxBytes = 10 * 1024

var toolNameSanitizer = regexp.MustCompile(`[^a-zA-Z0-9_]+`)

func sanitizeToolName(name string) string {
	s := toolNameSanitizer.ReplaceAllString(name, "_")
	if len(s) > 32 {
		s = s[:32]
	}
	s = strings.Trim(s, "_")
	if s == "" {
		s = "tool"
	}
	return s
}

func shortHexID(n int) (string, error) {
	b := make([]byte, n/2+n%2)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return hex.EncodeToString(b)[:n], nil
}

type toolCallLogger struct {
	file      *os.File
	toolName  string
	startedAt time.Time
	resultRaw json.RawMessage
	err       error
	endOnce   sync.Once
}

func (t *toolCallLogger) WithResult(result any) {
	t.resultRaw = marshalResult(result)
}

func (t *toolCallLogger) WithError(err error) {
	t.err = err
}

func (t *toolCallLogger) End() {
	t.endOnce.Do(func() {
		defer func() { _ = t.file.Close() }()

		durationMs := time.Since(t.startedAt).Milliseconds()
		errStr := ""
		if t.err != nil {
			errStr = t.err.Error()
		}

		end := map[string]any{
			"event":       "end",
			"duration_ms": durationMs,
			"result":      nil,
			"error":       errStr,
		}
		if t.resultRaw != nil {
			end["result"] = json.RawMessage(t.resultRaw)
		}

		b, err := json.Marshal(end)
		if err != nil {
			b = []byte(fmt.Sprintf(`{"event":"end","duration_ms":%d,"result":null,"error":%q}`, durationMs, errStr))
		}
		_, _ = t.file.Write(b)
		_, _ = t.file.WriteString("\n")
	})
}

// StartToolCall creates a per-tool-call log file and returns a logger that
// records the call lifecycle.
func StartToolCall(ctx context.Context, base *slog.Logger, toolName string, params any) (ToolCallLogger, error) {
	logDirMu.RLock()
	dir := logDir
	logDirMu.RUnlock()
	if dir == "" {
		return nil, fmt.Errorf("logging: log directory not initialized; call New first")
	}

	toolsDir := filepath.Join(dir, "tools")
	if err := os.MkdirAll(toolsDir, 0o700); err != nil {
		return nil, fmt.Errorf("logging: create tools log dir: %w", err)
	}

	safeName := sanitizeToolName(toolName)
	ts := time.Now().Format("20060102_150405.000")
	id, err := shortHexID(6)
	if err != nil {
		return nil, fmt.Errorf("logging: generate tool call id: %w", err)
	}
	filename := fmt.Sprintf("%s_%s_%s.log", ts, safeName, id)
	path := filepath.Join(toolsDir, filename)

	f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600)
	if err != nil {
		return nil, fmt.Errorf("logging: create tool call log file: %w", err)
	}

	startedAt := time.Now()
	start := map[string]any{
		"event":      "start",
		"tool":       toolName,
		"params":     marshalValue(params),
		"started_at": startedAt.Format(time.RFC3339Nano),
	}
	if d, ok := ctx.Deadline(); ok {
		start["deadline"] = d.Format(time.RFC3339Nano)
	}

	b, err := json.Marshal(start)
	if err != nil {
		_ = f.Close()
		return nil, fmt.Errorf("logging: marshal start record: %w", err)
	}
	if _, err := f.Write(b); err != nil {
		_ = f.Close()
		return nil, fmt.Errorf("logging: write start record: %w", err)
	}
	if _, err := f.WriteString("\n"); err != nil {
		_ = f.Close()
		return nil, fmt.Errorf("logging: write start record: %w", err)
	}

	if base != nil {
		base.Info("tool.call.start", "tool", toolName, "file", filename)
	}

	return &toolCallLogger{
		file:      f,
		toolName:  toolName,
		startedAt: startedAt,
	}, nil
}

func marshalValue(v any) json.RawMessage {
	if v == nil {
		return []byte("null")
	}
	b, err := json.MarshalIndent(v, "", "  ")
	if err != nil {
		b, err = json.Marshal(v)
		if err != nil {
			return mustMarshalString("<unserializable>")
		}
	}
	return b
}

func marshalResult(v any) json.RawMessage {
	if v == nil {
		return []byte("null")
	}
	b, err := json.MarshalIndent(v, "", "  ")
	if err != nil {
		b, err = json.Marshal(v)
		if err != nil {
			return mustMarshalString("<unserializable>")
		}
	}
	if len(b) > toolCallResultMaxBytes {
		trunc := string(b[:toolCallResultMaxBytes]) + "...[truncated]"
		return mustMarshalString(trunc)
	}
	return b
}

func mustMarshalString(s string) json.RawMessage {
	b, err := json.Marshal(s)
	if err != nil {
		return []byte(`"` + strings.ReplaceAll(s, `"`, `\"`) + `"`)
	}
	return b
}