First Commit

llmrec/capture.go

@@ -0,0 +1,162 @@
+package llmrec
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"strings"
+	"sync"
+)
+
+type captureContextKey struct{}
+
+// Capture collects the untouched wire bodies of one logical LLM call. The
+// Recorder creates it and puts it on the context; the HTTP transport that norma
+// dials through (agent.quotaAwareTransport) finds it there and fills it in.
+//
+// This exists because everything the Recorder itself sees is already normalized:
+// llm.CompletionRequest is re-serialized rather than the body buildBody() sent,
+// and the response arrives as decoded StreamEvents, not the SSE frames. For
+// debugging a live provider, the bytes on the wire are the only ground truth.
+//
+// norma's doStream retries the request-establishment phase, so one Stream can
+// issue several HTTP attempts. Each attempt is kept: the discarded ones (see
+// norma/llm/retry.go, which closes non-final bodies unread) are exactly what
+// makes rate-limit and gateway failures diagnosable.
+//
+// The transport writes from norma's stream-reading goroutine while the Recorder
+// snapshots at stream end, so all state is mutex-guarded.
+type Capture struct {
+	mu       sync.Mutex
+	request  string
+	attempts []*attempt
+}
+
+type attempt struct {
+	status int
+	body   strings.Builder
+}
+
+// NewCapture returns a context carrying a fresh Capture, plus the Capture itself.
+func NewCapture(ctx context.Context) (context.Context, *Capture) {
+	c := &Capture{}
+	return context.WithValue(ctx, captureContextKey{}, c), c
+}
+
+// CaptureFrom returns the Capture attached to ctx, or nil when raw capture is
+// off. Callers must tolerate nil — recording is a toggle, and non-recorded
+// providers dial through the same transport.
+func CaptureFrom(ctx context.Context) *Capture {
+	if ctx == nil {
+		return nil
+	}
+	c, _ := ctx.Value(captureContextKey{}).(*Capture)
+	return c
+}
+
+// SetRequest stores the outgoing request body. Retries re-send identical bytes,
+// so only the first attempt's body is kept.
+func (c *Capture) SetRequest(body string) {
+	if c == nil || body == "" {
+		return
+	}
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.request == "" {
+		c.request = body
+	}
+}
+
+// TeeResponse opens a new attempt and wraps rc so everything read from it is
+// mirrored into that attempt. It tees rather than reads because a successful
+// response is an SSE stream that must keep streaming to the caller.
+func (c *Capture) TeeResponse(status int, rc io.ReadCloser) io.ReadCloser {
+	if c == nil || rc == nil {
+		return rc
+	}
+	a := &attempt{status: status}
+	c.mu.Lock()
+	c.attempts = append(c.attempts, a)
+	c.mu.Unlock()
+	return &teeBody{rc: rc, c: c, a: a}
+}
+
+// RawRequest returns the request body as sent, or "" if nothing was captured.
+func (c *Capture) RawRequest() string {
+	if c == nil {
+		return ""
+	}
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.request
+}
+
+// RawResponse returns the response bytes as received. A single attempt yields
+// the untouched original (copy-pasteable straight into a replay); multiple
+// attempts are concatenated behind per-attempt header lines so a retry sequence
+// stays readable.
+func (c *Capture) RawResponse() string {
+	if c == nil {
+		return ""
+	}
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	switch len(c.attempts) {
+	case 0:
+		return ""
+	case 1:
+		return c.attempts[0].body.String()
+	}
+	var b strings.Builder
+	for i, a := range c.attempts {
+		fmt.Fprintf(&b, "===== attempt %d/%d — HTTP %d =====\n", i+1, len(c.attempts), a.status)
+		body := a.body.String()
+		b.WriteString(body)
+		if !strings.HasSuffix(body, "\n") {
+			b.WriteByte('\n')
+		}
+	}
+	return b.String()
+}
+
+// Attempt is one HTTP round trip's status code and response bytes.
+type Attempt struct {
+	Status int
+	Body   string
+}
+
+// Attempts returns every round trip in order. Unlike RawResponse — which drops
+// the header line for a lone attempt so the bytes stay replayable — this always
+// carries the status code, for callers that must report "HTTP 401 + body".
+func (c *Capture) Attempts() []Attempt {
+	if c == nil {
+		return nil
+	}
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	out := make([]Attempt, 0, len(c.attempts))
+	for _, a := range c.attempts {
+		out = append(out, Attempt{Status: a.status, Body: a.body.String()})
+	}
+	return out
+}
+
+// teeBody mirrors reads into a Capture attempt, guarded by the Capture's mutex
+// so a snapshot taken mid-stream never races the writer.
+type teeBody struct {
+	rc io.ReadCloser
+	c  *Capture
+	a  *attempt
+}
+
+func (t *teeBody) Read(p []byte) (int, error) {
+	n, err := t.rc.Read(p)
+	if n > 0 {
+		t.c.mu.Lock()
+		t.a.body.Write(p[:n])
+		t.c.mu.Unlock()
+	}
+	return n, err
+}
+
+func (t *teeBody) Close() error { return t.rc.Close() }

llmrec/capture_test.go

@@ -0,0 +1,103 @@
+package llmrec
+
+import (
+	"context"
+	"io"
+	"strings"
+	"testing"
+)
+
+func TestCaptureFromNilAndMissing(t *testing.T) {
+	var nilCtx context.Context // the transport reaches here on any unrecorded call
+	if CaptureFrom(nilCtx) != nil {
+		t.Fatal("nil context should yield no capture")
+	}
+	if CaptureFrom(context.Background()) != nil {
+		t.Fatal("bare context should yield no capture")
+	}
+}
+
+// A nil *Capture is the "recording off" path: the transport calls these on every
+// request, so they must all be no-ops rather than panics.
+func TestNilCaptureIsInert(t *testing.T) {
+	var c *Capture
+	c.SetRequest("body")
+	if got := c.RawRequest(); got != "" {
+		t.Fatalf("RawRequest()=%q want empty", got)
+	}
+	if got := c.RawResponse(); got != "" {
+		t.Fatalf("RawResponse()=%q want empty", got)
+	}
+	rc := io.NopCloser(strings.NewReader("x"))
+	if c.TeeResponse(200, rc) != rc {
+		t.Fatal("nil capture must pass the body through untouched")
+	}
+}
+
+func TestCaptureSingleAttemptKeepsBytesVerbatim(t *testing.T) {
+	ctx, c := NewCapture(context.Background())
+	if CaptureFrom(ctx) != c {
+		t.Fatal("capture not retrievable from its own context")
+	}
+	c.SetRequest(`{"model":"x"}`)
+	// Retries re-send identical bytes; only the first is kept.
+	c.SetRequest(`{"model":"ignored"}`)
+
+	const sse = "event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: message_stop\ndata: {}\n\n"
+	body := c.TeeResponse(200, io.NopCloser(strings.NewReader(sse)))
+	got, err := io.ReadAll(body)
+	if err != nil {
+		t.Fatalf("read: %v", err)
+	}
+	if string(got) != sse {
+		t.Fatal("tee altered the stream delivered to the caller")
+	}
+	if err := body.Close(); err != nil {
+		t.Fatalf("close: %v", err)
+	}
+
+	if c.RawRequest() != `{"model":"x"}` {
+		t.Fatalf("RawRequest()=%q", c.RawRequest())
+	}
+	// A lone attempt must stay byte-identical — no headers, no framing added, so
+	// it can be replayed as-is.
+	if c.RawResponse() != sse {
+		t.Fatalf("RawResponse()=%q want verbatim SSE", c.RawResponse())
+	}
+}
+
+// norma's doStream discards the bodies of retried attempts (retry.go closes them
+// unread), so the capture is the only place a 429 body survives.
+func TestCaptureRetriedAttemptsAreAllKept(t *testing.T) {
+	_, c := NewCapture(context.Background())
+
+	first := c.TeeResponse(429, io.NopCloser(strings.NewReader(`{"error":"rate_limit"}`)))
+	if _, err := io.ReadAll(first); err != nil {
+		t.Fatalf("read first: %v", err)
+	}
+	second := c.TeeResponse(200, io.NopCloser(strings.NewReader("data: ok\n")))
+	if _, err := io.ReadAll(second); err != nil {
+		t.Fatalf("read second: %v", err)
+	}
+
+	raw := c.RawResponse()
+	if !strings.Contains(raw, `{"error":"rate_limit"}`) {
+		t.Fatalf("dropped the retried attempt body: %q", raw)
+	}
+	if !strings.Contains(raw, "data: ok") {
+		t.Fatalf("dropped the final attempt body: %q", raw)
+	}
+	if !strings.Contains(raw, "attempt 1/2 — HTTP 429") || !strings.Contains(raw, "attempt 2/2 — HTTP 200") {
+		t.Fatalf("attempts not delimited: %q", raw)
+	}
+	if strings.Index(raw, "rate_limit") > strings.Index(raw, "data: ok") {
+		t.Fatal("attempts stored out of order")
+	}
+}
+
+func TestCaptureEmptyWhenNoHTTPHappened(t *testing.T) {
+	_, c := NewCapture(context.Background())
+	if c.RawRequest() != "" || c.RawResponse() != "" {
+		t.Fatal("unused capture should be empty")
+	}
+}

llmrec/llmrec.go

@@ -0,0 +1,359 @@
+// Package llmrec implements a recording decorator for llm.Provider. It intercepts
+// every Stream call, captures the full request and accumulated response, and
+// persists them to PostgreSQL for later inspection.
+package llmrec
+
+import (
+	"context"
+	"encoding/json"
+	"iter"
+	"log"
+	"strings"
+	"time"
+
+	"github.com/Autumn-27/artex/db"
+	"github.com/Autumn-27/norma/llm"
+	"github.com/Autumn-27/norma/transcript"
+)
+
+type taskIDContextKey struct{}
+
+// WithTaskID attaches the owning task registry id to an LLM call. Session ids
+// are based on exploration ids, which are not interchangeable with task ids.
+func WithTaskID(ctx context.Context, taskID string) context.Context {
+	taskID = strings.TrimSpace(taskID)
+	if taskID == "" {
+		return ctx
+	}
+	return context.WithValue(ctx, taskIDContextKey{}, taskID)
+}
+
+// TaskIDFrom returns the explicit task registry id attached by the task runtime.
+func TaskIDFrom(ctx context.Context) string {
+	if ctx == nil {
+		return ""
+	}
+	taskID, _ := ctx.Value(taskIDContextKey{}).(string)
+	return strings.TrimSpace(taskID)
+}
+
+// Recorder wraps an llm.Provider and records every completion call.
+type Recorder struct {
+	inner llm.Provider
+	pg    *db.DB
+	model string // model name (from config, not in CompletionRequest)
+	prof  string // LLM profile name (from llm_profiles)
+	// thinkingType / reasoningEffort 是配置级思考参数(思考开关 / 思考强度)。它们在
+	// norma 的 buildBody() 里从 provider 配置注入真正的 HTTP body,不出现在
+	// CompletionRequest 上,故 Recorder 需在此单独带一份,序列化时写进录制。
+	thinkingType    string
+	reasoningEffort string
+	enabled         func() bool // reports whether recording is currently on; nil = always record
+}
+
+// Wrap returns a Provider that records calls to pg when enabled() reports true.
+// profName is the LLM profile name (e.g. "default"); may be empty. thinkingType /
+// reasoningEffort are the config-level thinking params actually sent to the API
+// (empty = not sent). A nil enabled predicate records unconditionally.
+func Wrap(inner llm.Provider, pg *db.DB, model, profName, thinkingType, reasoningEffort string, enabled func() bool) *Recorder {
+	return &Recorder{
+		inner: inner, pg: pg, model: model, prof: profName,
+		thinkingType: thinkingType, reasoningEffort: reasoningEffort, enabled: enabled,
+	}
+}
+
+// parseSession extracts task id and worker role from session strings like
+// "exp1-worker-i3" → ("1", "worker") or "exp2-planner" → ("2", "planner").
+func parseSession(s string) (taskID, worker string) {
+	// format: exp<N>-<role>[-suffix]
+	if !strings.HasPrefix(s, "exp") {
+		return "", ""
+	}
+	rest := s[3:] // after "exp"
+	// split task number
+	i := strings.IndexByte(rest, '-')
+	if i < 0 {
+		return rest, ""
+	}
+	taskID = rest[:i]
+	rest = rest[i+1:]
+	// worker role is up to the next '-' (e.g. "worker" from "worker-i3")
+	if j := strings.IndexByte(rest, '-'); j >= 0 {
+		worker = rest[:j]
+	} else {
+		worker = rest
+	}
+	return taskID, worker
+}
+
+// Stream implements llm.Provider. It delegates to the inner provider, accumulates
+// the streamed events to reconstruct the response, and records the full exchange.
+//
+// The session identifier (e.g. "exp1-worker-i3") is carried on ctx by the norma
+// harness via transcript.WithSessionID; reading it per-call is race-free even when
+// planner and multiple workers share one Recorder instance.
+func (r *Recorder) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
+	// Body recording (the heavy debug trace: full request/response) is gated by the
+	// llm_record setting. Lightweight usage metering always runs — it powers token
+	// stats and must be complete even for interrupted/failed runs, so it is NOT
+	// gated. Only body serialization + response accumulation are skipped when off.
+	recordBodies := r.enabled == nil || r.enabled()
+	start := time.Now()
+	session := transcript.SessionIDFrom(ctx)
+	parsedID, worker := parseSession(session)
+	expID := db.ParseExpID(parsedID)
+	taskID := TaskIDFrom(ctx)
+	if taskID == "" {
+		// Backward compatibility for non-task callers. For task calls, the task
+		// runtime always supplies the registry id explicitly.
+		taskID = parsedID
+	}
+
+	// Serialize the request only when storing bodies (this is the expensive part).
+	// Alongside it, attach a Capture so the HTTP transport can hand back the
+	// untouched wire bodies — the normalized view below cannot reconstruct them
+	// (tool schemas are dropped, tool_use blocks never reach this layer, and the
+	// SSE framing is already decoded). See capture.go.
+	reqBody := ""
+	var capt *Capture
+	if recordBodies {
+		reqBody = r.serializeRequest(req)
+		ctx, capt = NewCapture(ctx)
+	}
+
+	return func(yield func(llm.StreamEvent, error) bool) {
+		var (
+			textBuf     strings.Builder
+			thinkingBuf strings.Builder
+			usage       llm.Usage
+			stopReason  string
+			streamErr   error
+		)
+
+		finished := false
+		finish := func(err error) {
+			if finished {
+				return
+			}
+			finished = true
+			status := "ok"
+			if err != nil {
+				status = "error"
+			}
+			// Lightweight metering row — always written.
+			r.recordUsage(taskID, expID, worker, usage, int(time.Since(start).Milliseconds()), status)
+			// Heavy trace row — only when body recording is on.
+			if recordBodies {
+				r.record(req, session, taskID, worker, reqBody, capt, start, textBuf.String(), thinkingBuf.String(), usage, stopReason, err)
+			}
+		}
+		defer func() { finish(ctx.Err()) }()
+
+		for ev, err := range r.inner.Stream(ctx, req) {
+			if err != nil {
+				streamErr = err
+				finish(streamErr)
+				if !yield(ev, err) {
+					return
+				}
+				return
+			}
+			// Always track usage (cheap); accumulate text/thinking only for bodies.
+			// Anthropic (and the other providers) split token usage across events:
+			// message_start carries ONLY the input side (input + cache), message_delta
+			// ONLY the output. They must be FOLDED with Add — overwriting on delta
+			// would zero out the input already counted at start (mirrors the SDK's own
+			// llm.Accumulator; see norma/llm/accumulate.go).
+			switch ev.Type {
+			case llm.SETextDelta:
+				if recordBodies {
+					textBuf.WriteString(ev.Text)
+				}
+			case llm.SEThinkingDelta:
+				if recordBodies {
+					thinkingBuf.WriteString(ev.Text)
+				}
+			case llm.SEMessageStart:
+				usage.Add(ev.Usage)
+			case llm.SEMessageDelta:
+				if ev.StopReason != "" {
+					stopReason = ev.StopReason
+				}
+				usage.Add(ev.Usage)
+			}
+			if !yield(ev, err) {
+				return
+			}
+		}
+
+		// Stream completed normally.
+		finish(streamErr)
+	}
+}
+
+// Complete implements the atomic completion path while preserving the same
+// usage metering, normalized body recording and raw wire capture as Stream.
+func (r *Recorder) Complete(ctx context.Context, req llm.CompletionRequest) (llm.Message, string, llm.Usage, error) {
+	recordBodies := r.enabled == nil || r.enabled()
+	start := time.Now()
+	session := transcript.SessionIDFrom(ctx)
+	parsedID, worker := parseSession(session)
+	expID := db.ParseExpID(parsedID)
+	taskID := TaskIDFrom(ctx)
+	if taskID == "" {
+		taskID = parsedID
+	}
+
+	reqBody := ""
+	var capt *Capture
+	if recordBodies {
+		reqBody = r.serializeRequest(req)
+		ctx, capt = NewCapture(ctx)
+	}
+
+	msg, stopReason, usage, err := r.inner.Complete(ctx, req)
+	status := "ok"
+	if err != nil {
+		status = "error"
+	}
+	r.recordUsage(taskID, expID, worker, usage, int(time.Since(start).Milliseconds()), status)
+	if recordBodies {
+		r.record(req, session, taskID, worker, reqBody, capt, start, msg.Text(), thinkingText(msg), usage, stopReason, err)
+	}
+	return msg, stopReason, usage, err
+}
+
+func thinkingText(msg llm.Message) string {
+	var b strings.Builder
+	for _, block := range msg.Content {
+		if block.Type == llm.BlockThinking && block.Thinking != "" {
+			b.WriteString(block.Thinking)
+		}
+	}
+	return b.String()
+}
+
+// recordUsage appends one lightweight metering row to llm_usage (no bodies). Skips
+// zero-token calls with no model, which carry nothing worth metering.
+func (r *Recorder) recordUsage(taskID string, expID int64, worker string, usage llm.Usage, latencyMs int, status string) {
+	if r.pg == nil {
+		return
+	}
+	if r.model == "" && usage.InputTokens == 0 && usage.OutputTokens == 0 &&
+		usage.CacheReadTokens == 0 && usage.CacheWriteTokens == 0 {
+		return
+	}
+	err := r.pg.InsertLLMUsage(&db.LLMUsage{
+		TaskID:        taskID,
+		ExplorationID: expID,
+		Worker:        worker,
+		Model:         r.model,
+		ProfileName:   r.prof,
+		LatencyMs:     latencyMs,
+		InputTokens:   usage.InputTokens,
+		OutputTokens:  usage.OutputTokens,
+		CacheRead:     usage.CacheReadTokens,
+		CacheWrite:    usage.CacheWriteTokens,
+		Status:        status,
+	})
+	if err != nil {
+		log.Printf("[llmusage] insert: %v", err)
+	}
+}
+
+// record persists one LLM call to PostgreSQL before the provider stream returns.
+// Keeping the write inside the owning task operation means task deletion can
+// drain calls and then remove records without a late async insert recreating one.
+func (r *Recorder) record(req llm.CompletionRequest, session, taskID, worker, reqBody string, capt *Capture, start time.Time, text, thinking string, usage llm.Usage, stopReason string, streamErr error) {
+	latency := int(time.Since(start).Milliseconds())
+	status := "ok"
+	errMsg := ""
+	if streamErr != nil {
+		status = "error"
+		errMsg = streamErr.Error()
+	}
+
+	// Build response body JSON.
+	resp := map[string]any{
+		"text":        text,
+		"stop_reason": stopReason,
+		"usage": map[string]int{
+			"input_tokens":       usage.InputTokens,
+			"output_tokens":      usage.OutputTokens,
+			"cache_read_tokens":  usage.CacheReadTokens,
+			"cache_write_tokens": usage.CacheWriteTokens,
+		},
+	}
+	if thinking != "" {
+		resp["thinking"] = thinking
+	}
+	respBody, _ := json.Marshal(resp)
+
+	model := ""
+	if len(req.System) > 0 {
+		// model is not in CompletionRequest; use the configured model name
+	}
+	model = r.model
+
+	rec := &db.LLMRecord{
+		SessionID:    session,
+		TaskID:       taskID,
+		Worker:       worker,
+		Model:        model,
+		ProfileName:  r.prof,
+		LatencyMs:    latency,
+		InputTokens:  usage.InputTokens,
+		OutputTokens: usage.OutputTokens,
+		CacheRead:    usage.CacheReadTokens,
+		CacheWrite:   usage.CacheWriteTokens,
+		Status:       status,
+		Error:        errMsg,
+		RequestBody:  reqBody,
+		ResponseBody: string(respBody),
+		// Raw wire bodies; empty when the transport did not fill the Capture
+		// (e.g. a provider dialing through a client without the capture hook, or
+		// a call that failed before any HTTP request went out).
+		RawRequest:  capt.RawRequest(),
+		RawResponse: capt.RawResponse(),
+	}
+
+	if err := r.pg.InsertLLMRecord(rec); err != nil {
+		log.Printf("[llmrec] insert: %v", err)
+	}
+}
+
+// serializeRequest builds a JSON representation of the completion request.
+func (r *Recorder) serializeRequest(req llm.CompletionRequest) string {
+	m := map[string]any{
+		"system":     req.System,
+		"messages":   req.Messages,
+		"max_tokens": req.MaxTokens,
+	}
+	// 记录本次调用实际发出的思考参数。type 采用「有效值」：每请求覆盖 req.Thinking
+	// 优先于配置级 thinkingType(与 norma buildBody 的判定一致，如 compaction 摘要会
+	// 强制 disabled)；effort 无每请求覆盖，直接取配置值。两者皆空则不写 thinking 字段。
+	effType := r.thinkingType
+	if req.Thinking != "" {
+		effType = req.Thinking
+	}
+	if effType != "" || r.reasoningEffort != "" {
+		m["thinking"] = map[string]string{"type": effType, "effort": r.reasoningEffort}
+	}
+	if len(req.Tools) > 0 {
+		// Store tool names only (full schemas are huge).
+		names := make([]string, len(req.Tools))
+		for i, t := range req.Tools {
+			names[i] = t.Name
+		}
+		m["tools"] = names
+		m["tools_count"] = len(req.Tools)
+	}
+	if req.Temperature != nil {
+		m["temperature"] = *req.Temperature
+	}
+	if len(req.Stop) > 0 {
+		m["stop"] = req.Stop
+	}
+	b, _ := json.Marshal(m)
+	return string(b)
+}

llmrec/llmrec_test.go

@@ -0,0 +1,132 @@
+package llmrec
+
+import (
+	"context"
+	"iter"
+	"os"
+	"testing"
+
+	"github.com/Autumn-27/artex/db"
+	"github.com/Autumn-27/norma/llm"
+	"github.com/Autumn-27/norma/transcript"
+	"github.com/google/uuid"
+)
+
+type completeProvider struct{}
+
+func (completeProvider) Stream(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
+	return func(func(llm.StreamEvent, error) bool) {}
+}
+
+func (completeProvider) Complete(context.Context, llm.CompletionRequest) (llm.Message, string, llm.Usage, error) {
+	return llm.Message{
+		Role: llm.RoleAssistant,
+		Content: []llm.ContentBlock{
+			{Type: llm.BlockThinking, Thinking: "reasoning"},
+			llm.TextBlock("answer"),
+		},
+	}, "stop", llm.Usage{InputTokens: 7, OutputTokens: 3}, nil
+}
+
+func TestTaskIDContextUsesExplicitRegistryID(t *testing.T) {
+	ctx := WithTaskID(context.Background(), " 42 ")
+	if got := TaskIDFrom(ctx); got != "42" {
+		t.Fatalf("TaskIDFrom()=%q want 42", got)
+	}
+	if got := TaskIDFrom(WithTaskID(ctx, "   ")); got != "42" {
+		t.Fatalf("blank task id should preserve parent context, got %q", got)
+	}
+	if got := TaskIDFrom(nil); got != "" {
+		t.Fatalf("nil context returned %q", got)
+	}
+}
+
+func TestParseSessionFallbackIsExplorationScoped(t *testing.T) {
+	taskID, worker := parseSession("exp12-worker-i99")
+	if taskID != "12" || worker != "worker" {
+		t.Fatalf("parseSession()=(%q,%q)", taskID, worker)
+	}
+	if taskID, worker := parseSession("not-a-task"); taskID != "" || worker != "" {
+		t.Fatalf("unexpected non-session parse: (%q,%q)", taskID, worker)
+	}
+}
+
+func TestCompleteForwardsAtomicResponse(t *testing.T) {
+	recorder := Wrap(completeProvider{}, nil, "model", "profile", "", "", func() bool { return false })
+	msg, stopReason, usage, err := recorder.Complete(context.Background(), llm.CompletionRequest{})
+	if err != nil {
+		t.Fatalf("Complete() error: %v", err)
+	}
+	if got := msg.Text(); got != "answer" {
+		t.Fatalf("Complete() text=%q want answer", got)
+	}
+	if stopReason != "stop" {
+		t.Fatalf("Complete() stop reason=%q want stop", stopReason)
+	}
+	if usage.InputTokens != 7 || usage.OutputTokens != 3 {
+		t.Fatalf("Complete() usage=%+v", usage)
+	}
+}
+
+type meteredStreamProvider struct{ completeProvider }
+
+func (meteredStreamProvider) Stream(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
+	return func(yield func(llm.StreamEvent, error) bool) {
+		if !yield(llm.StreamEvent{Type: llm.SEMessageStart, Usage: llm.Usage{InputTokens: 23}}, nil) {
+			return
+		}
+		if !yield(llm.StreamEvent{Type: llm.SEMessageDelta, Usage: llm.Usage{OutputTokens: 5}}, nil) {
+			return
+		}
+		yield(llm.StreamEvent{Type: llm.SEMessageStop}, nil)
+	}
+}
+
+func TestSideUsageRecordedOnceOnConsumerCancellation(t *testing.T) {
+	dsn := os.Getenv("ARTEX_PG_DSN")
+	if dsn == "" {
+		t.Skip("requires isolated ARTEX_PG_DSN")
+	}
+	pg, err := db.Open(dsn)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer pg.Close()
+	// llm_usage 는 db.Open 의 기본 스키마에 들어 있지 않고 EnsureLLMUsageTable 로만
+	// 만들어진다(운영에서는 server/manager.go 가 호출한다). 격리된 빈 데이터베이스에서
+	// 이 테스트를 단독으로 돌릴 때 다른 패키지가 같은 데이터베이스에 테이블을 먼저
+	// 만들어 주기를 기대하지 않도록, 여기서 직접 테이블 존재를 보장한다.
+	if err := pg.EnsureLLMUsageTable(); err != nil {
+		t.Fatal(err)
+	}
+	for _, early := range []bool{false, true} {
+		profile := "btw-metering-" + uuid.NewString()
+		ctx, cancel := context.WithCancel(transcript.WithSessionID(t.Context(), "exp0-btw-test"))
+		recorder := Wrap(meteredStreamProvider{}, pg, "fixture", profile, "", "", func() bool { return false })
+		for event, err := range recorder.Stream(ctx, llm.CompletionRequest{}) {
+			if err != nil {
+				t.Fatal(err)
+			}
+			if early && event.Type == llm.SEMessageDelta {
+				cancel()
+				break // Consumer exits before the provider can return a cancellation event.
+			}
+		}
+		cancel()
+		var count, input, output int
+		var worker, status string
+		err = pg.QueryRow(`SELECT count(*),max(input_tokens),max(output_tokens),max(worker),max(status)
+FROM llm_usage WHERE profile_name=$1`, profile).Scan(&count, &input, &output, &worker, &status)
+		_, _ = pg.Exec(`DELETE FROM llm_usage WHERE profile_name=$1`, profile)
+		if err != nil {
+			t.Fatal(err)
+		}
+		wantStatus := "ok"
+		if early {
+			wantStatus = "error"
+		}
+		if count != 1 || input != 23 || output != 5 || worker != "btw" || status != wantStatus {
+			t.Fatalf("early=%v: count=%d usage=%d/%d worker=%q status=%q", early, count, input, output, worker, status)
+		}
+	}
+}