Cyrene_v2/cmd/cyrene/observability_test.go

package main

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/auth"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/trace"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/services/gateway"
)

const obsTestSecret = "0123456789abcdef0123456789abcdef"

// recordingWriter captures the spans a fanout hands it.
type recordingWriter struct {
	spans []trace.SpanRecord
}

func (r *recordingWriter) WriteSpan(rec trace.SpanRecord) error {
	r.spans = append(r.spans, rec)
	return nil
}

// TestTraceFanout_LateSinkReceivesSpans proves the composition root can register
// a sink after the loop's tracer is already running: spans finished before the
// sink is added do not reach it, every span after does (the gateway is built
// after the loop, so this ordering is what A4 depends on).
func TestTraceFanout_LateSinkReceivesSpans(t *testing.T) {
	t.Parallel()
	early := &recordingWriter{}
	fanout := newTraceFanout(early)
	tracer := trace.New(fanout)

	_, first := tracer.StartSpan(context.Background(), "before", trace.SpanKindStateTransition)
	first.Finish(trace.SpanStatusOK)

	late := &recordingWriter{}
	fanout.Add(late)

	_, second := tracer.StartSpan(context.Background(), "after", trace.SpanKindLLM)
	second.Finish(trace.SpanStatusOK)

	if len(early.spans) != 2 {
		t.Fatalf("early sink spans = %d, want 2", len(early.spans))
	}
	if len(late.spans) != 1 || late.spans[0].ID != second.ID {
		t.Fatalf("late sink spans = %+v, want exactly the span %q", late.spans, second.ID)
	}
}

// TestWireGatewayTraceSink_Guards proves the boot helper tolerates a missing
// tracer or gateway, so a degraded boot cannot panic on the observability path.
func TestWireGatewayTraceSink_Guards(t *testing.T) {
	t.Parallel()
	wireGatewayTraceSink(&app{log: testLog()}, nil)
	a := &app{log: testLog(), traceWriter: newAppTraceWriter()}
	wireGatewayTraceSink(a, nil)
	if len(a.traceWriter.writers) != 1 {
		t.Fatalf("writers = %d, want 1 (the stderr recorder only)", len(a.traceWriter.writers))
	}
}

// TestGatewaySpanSink_FeedsTraceEndpointAndMetrics is the A4 acceptance: a span
// emitted by the tracer is written through the gateway sink (registered by the
// same helpers wire.go calls) and surfaces both in
// GET /api/v1/admin/observability/trace (a real span, not an empty buffer) and in
// the LLM metric counter.
func TestGatewaySpanSink_FeedsTraceEndpointAndMetrics(t *testing.T) {
	t.Parallel()
	jwt := auth.New(config.AuthConfig{JWTSecret: obsTestSecret})
	gw, err := gateway.New(&config.Config{}, gateway.Deps{Log: testLog(), Auth: jwt})
	if err != nil {
		t.Fatalf("gateway.New: %v", err)
	}
	t.Cleanup(gw.Close)
	ts := httptest.NewServer(gw.Handler())
	t.Cleanup(ts.Close)

	// Exactly the boot path: the app's fanout is built in wireRuntime and the
	// gateway sink is registered in wireFrontend.
	a := &app{log: testLog(), traceWriter: newAppTraceWriter()}
	wireGatewayTraceSink(a, gw)
	tracer := trace.New(a.traceWriter)

	ctx := context.Background()
	_, llmSpan := tracer.StartSpan(ctx, "llm.chat", trace.SpanKindLLM)
	llmSpan.Finish(trace.SpanStatusOK)
	_, stepSpan := tracer.StartSpan(ctx, "loop.step", trace.SpanKindStateTransition)
	_ = stepSpan.SetProp("session", "web")
	stepSpan.Finish(trace.SpanStatusOK)

	tok, err := jwt.Sign("acct", auth.RoleOwner, time.Hour)
	if err != nil {
		t.Fatalf("sign token: %v", err)
	}

	var traceBody struct {
		Data []struct {
			ID   string `json:"id"`
			Name string `json:"name"`
			Kind string `json:"kind"`
		} `json:"data"`
	}
	status, raw := obsGet(t, ts.URL+"/api/v1/admin/observability/trace", tok)
	if status != http.StatusOK {
		t.Fatalf("trace status = %d, want 200 (body=%s)", status, raw)
	}
	if err := json.Unmarshal(raw, &traceBody); err != nil {
		t.Fatalf("decode trace: %v (body=%s)", err, raw)
	}
	seen := map[string]string{}
	for _, sp := range traceBody.Data {
		seen[sp.ID] = sp.Kind
	}
	if seen[stepSpan.ID] != string(trace.SpanKindStateTransition) {
		t.Fatalf("trace endpoint = %+v, want the loop.step span %q", traceBody.Data, stepSpan.ID)
	}
	if seen[llmSpan.ID] != string(trace.SpanKindLLM) {
		t.Fatalf("trace endpoint = %+v, want the llm span %q", traceBody.Data, llmSpan.ID)
	}

	var metricsBody struct {
		Data struct {
			LLMCalls int64 `json:"llmCalls"`
		} `json:"data"`
	}
	status, raw = obsGet(t, ts.URL+"/api/v1/admin/observability/metrics", tok)
	if status != http.StatusOK {
		t.Fatalf("metrics status = %d, want 200 (body=%s)", status, raw)
	}
	if err := json.Unmarshal(raw, &metricsBody); err != nil {
		t.Fatalf("decode metrics: %v (body=%s)", err, raw)
	}
	if metricsBody.Data.LLMCalls != 1 {
		t.Fatalf("llmCalls = %d, want 1 (ObserveSpan counts the emitted llm span)", metricsBody.Data.LLMCalls)
	}
}

// obsGet performs an authenticated owner GET and returns the status plus the raw
// body.
func obsGet(t *testing.T, url, token string) (int, []byte) {
	t.Helper()
	req, err := http.NewRequest(http.MethodGet, url, nil)
	if err != nil {
		t.Fatalf("new request: %v", err)
	}
	req.Header.Set("Authorization", "Bearer "+token)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatalf("GET %s: %v", url, err)
	}
	defer resp.Body.Close()
	raw, err := io.ReadAll(resp.Body)
	if err != nil {
		t.Fatalf("read %s: %v", url, err)
	}
	return resp.StatusCode, raw
}