Vectrum
Sign In
156 lines · 3.9 KB · Go
// Package metrics writes a JSONL event log: one JSON object per line, one line
// per thing worth counting later. It is wired in before there is anything to
// measure on purpose — retrofitting instrumentation is how projects end up with
// no numbers.
package metrics

import (
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sync"
	"time"

	"codey/internal/statedir"
)

// DefaultPath is where Codey logs unless told otherwise.
const DefaultPath = statedir.Dir + "/events.jsonl"

// Event is one line of the log. Extra is flattened into the same JSON object so
// the file stays greppable with jq without nested lookups.
type Event struct {
	Time    time.Time `json:"ts"`
	Session string    `json:"session,omitempty"`
	Type    string    `json:"type"`

	// Tool dispatch.
	Tool     string `json:"tool,omitempty"`
	Duration int64  `json:"dur_ms,omitempty"`
	OK       *bool  `json:"ok,omitempty"`
	Err      string `json:"err,omitempty"`

	// Turn accounting.
	Turn       int    `json:"turn,omitempty"`
	StopReason string `json:"stop_reason,omitempty"`

	// Token accounting, mirroring the Messages API usage block.
	InputTokens      int `json:"input_tokens,omitempty"`
	OutputTokens     int `json:"output_tokens,omitempty"`
	CacheReadTokens  int `json:"cache_read_input_tokens,omitempty"`
	CacheWriteTokens int `json:"cache_creation_input_tokens,omitempty"`

	Extra map[string]any `json:"-"`
}

// MarshalJSON flattens Extra alongside the declared fields. Keys in Extra never
// clobber a declared field.
func (e Event) MarshalJSON() ([]byte, error) {
	type alias Event
	b, err := json.Marshal(alias(e))
	if err != nil {
		return nil, err
	}
	if len(e.Extra) == 0 {
		return b, nil
	}
	var m map[string]json.RawMessage
	if err := json.Unmarshal(b, &m); err != nil {
		return nil, err
	}
	for k, v := range e.Extra {
		if _, taken := m[k]; taken {
			continue
		}
		ev, err := json.Marshal(v)
		if err != nil {
			return nil, fmt.Errorf("metrics: extra field %q: %w", k, err)
		}
		m[k] = ev
	}
	return json.Marshal(m)
}

// Logger appends events to a writer. The zero value and a nil *Logger are both
// usable and simply drop events, so callers never need a nil check.
type Logger struct {
	mu      sync.Mutex
	w       io.Writer
	closer  io.Closer
	session string
}

// Open appends to path, creating parent directories as needed.
func Open(path, session string) (*Logger, error) {
	if dir := filepath.Dir(path); dir != "." {
		if err := os.MkdirAll(dir, 0o755); err != nil {
			return nil, fmt.Errorf("metrics: %w", err)
		}
	}
	f, err := os.OpenFile(path, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o644)
	if err != nil {
		return nil, fmt.Errorf("metrics: %w", err)
	}
	return &Logger{w: f, closer: f, session: session}, nil
}

// NewWriter logs to an arbitrary writer. Useful in tests and in the eval runner,
// which keeps a separate log per run.
func NewWriter(w io.Writer, session string) *Logger {
	return &Logger{w: w, session: session}
}

// Session returns the session id stamped onto every event.
func (l *Logger) Session() string {
	if l == nil {
		return ""
	}
	return l.session
}

// Log writes one event. Errors are dropped: a broken metrics file must never
// take down a run.
func (l *Logger) Log(ev Event) {
	if l == nil || l.w == nil {
		return
	}
	if ev.Time.IsZero() {
		ev.Time = time.Now().UTC()
	}
	if ev.Session == "" {
		ev.Session = l.session
	}
	b, err := json.Marshal(ev)
	if err != nil {
		return
	}
	l.mu.Lock()
	defer l.mu.Unlock()
	l.w.Write(append(b, '\n'))
}

// ToolCall records one dispatch. err nil means success.
func (l *Logger) ToolCall(name string, d time.Duration, err error, extra map[string]any) {
	ok := err == nil
	ev := Event{
		Type:     "tool_call",
		Tool:     name,
		Duration: d.Milliseconds(),
		OK:       &ok,
		Extra:    extra,
	}
	if err != nil {
		ev.Err = err.Error()
	}
	l.Log(ev)
}

// Close closes the underlying file, if there is one.
func (l *Logger) Close() error {
	if l == nil || l.closer == nil {
		return nil
	}
	return l.closer.Close()
}