From 91ead361b7907df149c176888fe913654f1bcf45 Mon Sep 17 00:00:00 2001 From: erio Date: Sun, 19 Apr 2026 14:00:01 +0800 Subject: [PATCH] test(signature): add unit tests + fix gjson ForEach indexing bug MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Covers: - ReplaceThinkingSignaturesInBody / InClaudeRequest: M>N cycling, empty pool, no thinking blocks, empty signature preservation, string-content messages, nil pool guard - Strip / Pool rectifier strategies: Strip stage-1 unconditional, stage-2 gated on tool error; Pool empty→proceed=false (rule A), no replacements→abort, one-shot semantics (stage-2 always declines), pool error treated as empty - BucketFor: oauth/setup-token share, apikey per-id, unknown empty - Harvester SSE: content_block_start extraction, per-response dedupe, Skip callback, bucket/capacity guards - Harvester non-streaming: buffers until Close then parses once Along the way caught and fixed a path-construction bug in ReplaceThinkingSignaturesInBody: it used gjson.ForEach's key.Raw which is empty for array indices, producing malformed sjson paths like "messages..content..signature". Switched to manual index counters so paths are always well-formed integers. --- .../service/signature/harvester_test.go | 127 +++++++++++ .../service/signature/rectifier_test.go | 199 ++++++++++++++++++ backend/internal/service/signature/replace.go | 12 +- .../service/signature/replace_test.go | 148 +++++++++++++ 4 files changed, 483 insertions(+), 3 deletions(-) create mode 100644 backend/internal/service/signature/harvester_test.go create mode 100644 backend/internal/service/signature/rectifier_test.go create mode 100644 backend/internal/service/signature/replace_test.go diff --git a/backend/internal/service/signature/harvester_test.go b/backend/internal/service/signature/harvester_test.go new file mode 100644 index 0000000000..8230ad90f4 --- /dev/null +++ b/backend/internal/service/signature/harvester_test.go @@ -0,0 +1,127 @@ +//go:build unit + +package signature + +import ( + "context" + "io" + "strings" + "testing" +) + +func TestHarvester_SSE_ExtractsContentBlockStartSignatures(t *testing.T) { + pool := newFakePool() + h := NewHarvester(pool, 10) + + // Two thinking content_block_start events + one unrelated event. + body := strings.Join([]string{ + `event: message_start`, + `data: {"type":"message_start","message":{}}`, + ``, + `event: content_block_start`, + `data: {"type":"content_block_start","index":0,"content_block":{"type":"thinking","thinking":"","signature":"SIG-A"}}`, + ``, + `event: content_block_start`, + `data: {"type":"content_block_start","index":1,"content_block":{"type":"thinking","thinking":"","signature":"SIG-B"}}`, + ``, + `event: content_block_delta`, + `data: {"type":"content_block_delta","delta":{"type":"thinking_delta","thinking":"hi"}}`, + ``, + }, "\n") + "\n" + + rc := h.Wrap(context.Background(), io.NopCloser(strings.NewReader(body)), HarvestOptions{ + Bucket: "oauth", + Streaming: true, + }) + if _, err := io.ReadAll(rc); err != nil { + t.Fatalf("read: %v", err) + } + _ = rc.Close() + + got := pool.entries["oauth"] + if len(got) != 2 { + t.Fatalf("expected 2 signatures harvested, got %d: %v", len(got), got) + } + // Newest-first semantics: fakePool prepends on Add, so the last-emitted is index 0. + if got[0] != "SIG-B" || got[1] != "SIG-A" { + t.Errorf("ordering: got %v, want [SIG-B, SIG-A]", got) + } +} + +func TestHarvester_SSE_DeduplicatesWithinSameResponse(t *testing.T) { + pool := newFakePool() + h := NewHarvester(pool, 10) + body := strings.Repeat( + "event: content_block_start\ndata: {\"type\":\"content_block_start\",\"content_block\":{\"type\":\"thinking\",\"signature\":\"DUP\"}}\n\n", + 5, + ) + rc := h.Wrap(context.Background(), io.NopCloser(strings.NewReader(body)), HarvestOptions{ + Bucket: "oauth", + Streaming: true, + }) + _, _ = io.ReadAll(rc) + _ = rc.Close() + + if got := pool.entries["oauth"]; len(got) != 1 { + t.Errorf("expected dedup to 1 entry within single response, got %d: %v", len(got), got) + } +} + +func TestHarvester_SkipCallbackBlocksIngestion(t *testing.T) { + pool := newFakePool() + h := NewHarvester(pool, 10) + body := `data: {"type":"content_block_start","content_block":{"type":"thinking","signature":"NO"}}` + "\n\n" + rc := h.Wrap(context.Background(), io.NopCloser(strings.NewReader(body)), HarvestOptions{ + Bucket: "oauth", + Streaming: true, + Skip: func() bool { return true }, + }) + _, _ = io.ReadAll(rc) + _ = rc.Close() + if len(pool.entries["oauth"]) != 0 { + t.Errorf("skip callback must prevent harvesting, got %v", pool.entries["oauth"]) + } +} + +func TestHarvester_NonStreaming_ParsesOnClose(t *testing.T) { + pool := newFakePool() + h := NewHarvester(pool, 10) + body := `{"id":"msg_1","content":[ + {"type":"thinking","thinking":"x","signature":"NS-A"}, + {"type":"text","text":"hi"}, + {"type":"thinking","thinking":"y","signature":"NS-B"} + ]}` + rc := h.Wrap(context.Background(), io.NopCloser(strings.NewReader(body)), HarvestOptions{ + Bucket: "oauth", + Streaming: false, + }) + _, _ = io.ReadAll(rc) + // Nothing should be harvested before Close on non-streaming path. + if len(pool.entries["oauth"]) != 0 { + t.Errorf("non-streaming must not ingest before Close; got %v", pool.entries["oauth"]) + } + _ = rc.Close() + if got := pool.entries["oauth"]; len(got) != 2 { + t.Fatalf("expected 2 entries after Close, got %d: %v", len(got), got) + } +} + +func TestHarvester_EmptyBucketDisablesWrap(t *testing.T) { + pool := newFakePool() + h := NewHarvester(pool, 10) + inner := io.NopCloser(strings.NewReader("irrelevant")) + out := h.Wrap(context.Background(), inner, HarvestOptions{Bucket: ""}) + if out != inner { + t.Errorf("empty bucket must return the original body unmodified") + } +} + +func TestHarvester_ZeroCapacityDisablesWrap(t *testing.T) { + pool := newFakePool() + h := NewHarvester(pool, 0) + inner := io.NopCloser(strings.NewReader("irrelevant")) + out := h.Wrap(context.Background(), inner, HarvestOptions{Bucket: "oauth"}) + if out != inner { + t.Errorf("zero capacity must return the original body unmodified") + } +} diff --git a/backend/internal/service/signature/rectifier_test.go b/backend/internal/service/signature/rectifier_test.go new file mode 100644 index 0000000000..8e01a06b4e --- /dev/null +++ b/backend/internal/service/signature/rectifier_test.go @@ -0,0 +1,199 @@ +//go:build unit + +package signature + +import ( + "context" + "encoding/json" + "errors" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" +) + +// fakePool is an in-memory SignaturePool used by tests; implements the +// ordering contract (TopN returns most-recently-added first). +type fakePool struct { + entries map[string][]string // bucket -> signatures, index 0 = newest +} + +func newFakePool() *fakePool { return &fakePool{entries: map[string][]string{}} } + +func (p *fakePool) Add(_ context.Context, bucket, sig string, _ time.Time, _ int) error { + // Prepend so index 0 is the newest. + p.entries[bucket] = append([]string{sig}, p.entries[bucket]...) + return nil +} +func (p *fakePool) TopN(_ context.Context, bucket string, n int) ([]string, error) { + list := p.entries[bucket] + if n > len(list) { + n = len(list) + } + out := make([]string, n) + copy(out, list[:n]) + return out, nil +} +func (p *fakePool) Size(_ context.Context, bucket string) (int64, error) { + return int64(len(p.entries[bucket])), nil +} + +func TestBucketFor(t *testing.T) { + if got := BucketFor("oauth", 1); got != "oauth" { + t.Errorf("oauth bucket: got %q, want oauth", got) + } + if got := BucketFor("setup-token", 1); got != "oauth" { + t.Errorf("setup-token must share oauth bucket, got %q", got) + } + if got := BucketFor("apikey", 42); got != "apikey:42" { + t.Errorf("apikey bucket: got %q, want apikey:42", got) + } + if got := BucketFor("bedrock", 1); got != "" { + t.Errorf("bedrock has no pool: got %q, want empty", got) + } +} + +func TestStripClaudeRectifier_Stage1AlwaysApplied(t *testing.T) { + r := &StripClaudeRectifier{ + FilterStage1: func(b []byte) []byte { return append([]byte("stripped1:"), b...) }, + FilterStage2: func(b []byte) []byte { return append([]byte("stripped2:"), b...) }, + } + out, proceed := r.Apply(context.Background(), ClaudeInput{Body: []byte("orig")}, StageThinkingOnly) + if !proceed || string(out) != "stripped1:orig" { + t.Fatalf("stage 1: proceed=%v out=%q", proceed, out) + } +} + +func TestStripClaudeRectifier_Stage2GatedOnToolError(t *testing.T) { + r := &StripClaudeRectifier{ + FilterStage1: func(b []byte) []byte { return b }, + FilterStage2: func(b []byte) []byte { return append([]byte("stripped2:"), b...) }, + } + // Non-tool-related error: stage 2 must decline. + out, proceed := r.Apply(context.Background(), ClaudeInput{Body: []byte("orig"), LastErrMsg: "invalid signature in thinking block"}, StageThinkingAndTools) + if proceed { + t.Fatalf("stage 2 should decline for non-tool-related errors") + } + if out != nil { + t.Errorf("expected nil body when declining, got %q", out) + } + // Tool-related error: stage 2 must proceed. + out, proceed = r.Apply(context.Background(), ClaudeInput{Body: []byte("orig"), LastErrMsg: "tool_use block signature invalid"}, StageThinkingAndTools) + if !proceed || string(out) != "stripped2:orig" { + t.Fatalf("stage 2 on tool error: proceed=%v out=%q", proceed, out) + } +} + +func TestPoolClaudeRectifier_EmptyPoolSignalsAbort(t *testing.T) { + pool := newFakePool() + r := &PoolClaudeRectifier{Pool: pool, Capacity: 10} + in := ClaudeInput{AccountType: "oauth", AccountID: 1, Body: []byte(`{"messages":[]}`)} + out, proceed := r.Apply(context.Background(), in, StageThinkingOnly) + if proceed { + t.Fatalf("empty pool must signal proceed=false (rule A)") + } + if out != nil { + t.Errorf("expected nil body when pool empty, got %q", out) + } +} + +func TestPoolClaudeRectifier_NoThinkingBlocksSignalsAbort(t *testing.T) { + pool := newFakePool() + _ = pool.Add(context.Background(), "oauth", "good", time.Now(), 10) + r := &PoolClaudeRectifier{Pool: pool, Capacity: 10} + // Pool has a sig, but the request has no thinking blocks to replace. + in := ClaudeInput{AccountType: "oauth", AccountID: 1, Body: []byte(`{"messages":[{"role":"user","content":"hi"}]}`)} + out, proceed := r.Apply(context.Background(), in, StageThinkingOnly) + if proceed { + t.Fatalf("replaced=0 must signal proceed=false") + } + if out != nil { + t.Errorf("expected nil body when nothing replaced, got %q", out) + } +} + +func TestPoolClaudeRectifier_ReplacesAndStopsAtStage2(t *testing.T) { + pool := newFakePool() + _ = pool.Add(context.Background(), "oauth", "good", time.Now(), 10) + r := &PoolClaudeRectifier{Pool: pool, Capacity: 10} + body := []byte(`{"messages":[{"role":"assistant","content":[{"type":"thinking","thinking":"x","signature":"bad"}]}]}`) + in := ClaudeInput{AccountType: "oauth", AccountID: 1, Body: body} + + out, proceed := r.Apply(context.Background(), in, StageThinkingOnly) + if !proceed { + t.Fatalf("pool with sigs + thinking blocks must proceed") + } + if string(out) == string(body) { + t.Errorf("body should be mutated") + } + + // Stage 2: pool is a one-shot strategy, must decline. + out2, proceed2 := r.Apply(context.Background(), in, StageThinkingAndTools) + if proceed2 || out2 != nil { + t.Fatalf("pool must decline stage 2 (one-shot)") + } +} + +func TestPoolAntigravityRectifier_Stages(t *testing.T) { + r := &PoolAntigravityRectifier{} + st := r.Stages() + if len(st) != 1 || st[0] != StageThinkingOnly { + t.Errorf("pool antigravity stages: got %v, want [StageThinkingOnly]", st) + } +} + +func TestPoolAntigravityRectifier_EmptyPool(t *testing.T) { + pool := newFakePool() + r := &PoolAntigravityRectifier{Pool: pool, Capacity: 10} + req := &antigravity.ClaudeRequest{ + Messages: []antigravity.ClaudeMessage{ + {Role: "assistant", Content: json.RawMessage(`[{"type":"thinking","thinking":"x","signature":"bad"}]`)}, + }, + } + applied, proceed, err := r.Apply(context.Background(), AntigravityInput{AccountType: "oauth", AccountID: 1, Request: req}, StageThinkingOnly) + if applied || proceed || err != nil { + t.Errorf("empty pool: applied=%v proceed=%v err=%v; want all false/nil", applied, proceed, err) + } +} + +func TestPoolAntigravityRectifier_AppliedAndOneShot(t *testing.T) { + pool := newFakePool() + _ = pool.Add(context.Background(), "oauth", "g0", time.Now(), 10) + r := &PoolAntigravityRectifier{Pool: pool, Capacity: 10} + req := &antigravity.ClaudeRequest{ + Messages: []antigravity.ClaudeMessage{ + {Role: "assistant", Content: json.RawMessage(`[{"type":"thinking","thinking":"x","signature":"bad"}]`)}, + }, + } + applied, proceed, err := r.Apply(context.Background(), AntigravityInput{AccountType: "oauth", AccountID: 1, Request: req}, StageThinkingOnly) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !applied { + t.Errorf("expected applied=true when pool produces replacements") + } + if proceed { + t.Errorf("expected proceed=false (pool is one-shot)") + } +} + +// errPool is a SignaturePool that always returns errors — used to confirm +// rectifiers handle pool failures gracefully. +type errPool struct{} + +func (errPool) Add(context.Context, string, string, time.Time, int) error { + return errors.New("add failed") +} +func (errPool) TopN(context.Context, string, int) ([]string, error) { + return nil, errors.New("topn failed") +} +func (errPool) Size(context.Context, string) (int64, error) { return 0, errors.New("size failed") } + +func TestPoolClaudeRectifier_PoolErrorIsTreatedAsEmpty(t *testing.T) { + r := &PoolClaudeRectifier{Pool: errPool{}, Capacity: 10} + in := ClaudeInput{AccountType: "oauth", AccountID: 1, Body: []byte(`{"messages":[]}`)} + _, proceed := r.Apply(context.Background(), in, StageThinkingOnly) + if proceed { + t.Errorf("pool error must behave like empty pool (proceed=false)") + } +} diff --git a/backend/internal/service/signature/replace.go b/backend/internal/service/signature/replace.go index 714a3bd7d5..822de9dc4b 100644 --- a/backend/internal/service/signature/replace.go +++ b/backend/internal/service/signature/replace.go @@ -35,12 +35,18 @@ func ReplaceThinkingSignaturesInBody(body []byte, pool []string) ([]byte, int) { out := body replaced := 0 - msgs.ForEach(func(mKey, msg gjson.Result) bool { + mIdx := 0 + msgs.ForEach(func(_, msg gjson.Result) bool { + thisMsg := mIdx + mIdx++ content := msg.Get("content") if !content.IsArray() { return true } - content.ForEach(func(bKey, blk gjson.Result) bool { + bIdx := 0 + content.ForEach(func(_, blk gjson.Result) bool { + thisBlk := bIdx + bIdx++ if blk.Get("type").String() != "thinking" { return true } @@ -49,7 +55,7 @@ func ReplaceThinkingSignaturesInBody(body []byte, pool []string) ([]byte, int) { return true } newSig := pool[replaced%len(pool)] - path := fmt.Sprintf("messages.%s.content.%s.signature", mKey.Raw, bKey.Raw) + path := fmt.Sprintf("messages.%d.content.%d.signature", thisMsg, thisBlk) if next, err := sjson.SetBytes(out, path, newSig); err == nil { out = next replaced++ diff --git a/backend/internal/service/signature/replace_test.go b/backend/internal/service/signature/replace_test.go new file mode 100644 index 0000000000..4b111f73a0 --- /dev/null +++ b/backend/internal/service/signature/replace_test.go @@ -0,0 +1,148 @@ +//go:build unit + +package signature + +import ( + "encoding/json" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" + "github.com/tidwall/gjson" +) + +func TestReplaceThinkingSignaturesInBody_CyclesPoolForMGreaterThanN(t *testing.T) { + body := []byte(`{ + "model": "claude-opus-4-7", + "messages": [ + {"role": "assistant", "content": [ + {"type": "thinking", "thinking": "a", "signature": "bad1"}, + {"type": "thinking", "thinking": "b", "signature": "bad2"}, + {"type": "thinking", "thinking": "c", "signature": "bad3"}, + {"type": "thinking", "thinking": "d", "signature": "bad4"}, + {"type": "thinking", "thinking": "e", "signature": "bad5"} + ]} + ] + }`) + pool := []string{"good0", "good1"} + + out, replaced := ReplaceThinkingSignaturesInBody(body, pool) + if replaced != 5 { + t.Fatalf("expected 5 replacements, got %d", replaced) + } + + sigs := gjson.GetBytes(out, "messages.0.content.#.signature").Array() + if len(sigs) != 5 { + t.Fatalf("expected 5 signatures in result, got %d", len(sigs)) + } + // Expected cycle: good0, good1, good0, good1, good0 + want := []string{"good0", "good1", "good0", "good1", "good0"} + for i, s := range sigs { + if s.String() != want[i] { + t.Errorf("sig[%d]: got %q, want %q", i, s.String(), want[i]) + } + } +} + +func TestReplaceThinkingSignaturesInBody_EmptyPool(t *testing.T) { + body := []byte(`{"messages":[{"role":"assistant","content":[{"type":"thinking","thinking":"x","signature":"bad"}]}]}`) + out, n := ReplaceThinkingSignaturesInBody(body, nil) + if n != 0 { + t.Fatalf("expected 0 replacements on empty pool, got %d", n) + } + if string(out) != string(body) { + t.Fatalf("empty pool must not modify body") + } +} + +func TestReplaceThinkingSignaturesInBody_NoThinkingBlocks(t *testing.T) { + body := []byte(`{"messages":[{"role":"user","content":"hi"}]}`) + out, n := ReplaceThinkingSignaturesInBody(body, []string{"g"}) + if n != 0 { + t.Fatalf("expected 0 replacements when no thinking blocks, got %d", n) + } + if string(out) != string(body) { + t.Fatalf("body must be unchanged") + } +} + +func TestReplaceThinkingSignaturesInBody_SkipsEmptySignatureAndOtherTypes(t *testing.T) { + body := []byte(`{ + "messages": [ + {"role":"assistant","content":[ + {"type":"text","text":"hello"}, + {"type":"thinking","thinking":"keep-empty-sig","signature":""}, + {"type":"redacted_thinking","data":"xyz"}, + {"type":"thinking","thinking":"replace-me","signature":"bad"} + ]} + ] + }`) + pool := []string{"good"} + out, n := ReplaceThinkingSignaturesInBody(body, pool) + if n != 1 { + t.Fatalf("expected 1 replacement, got %d", n) + } + // Empty signature stays empty (not replaced) + if gjson.GetBytes(out, "messages.0.content.1.signature").String() != "" { + t.Errorf("empty signature must be preserved unchanged") + } + // Real bad signature gets replaced + if gjson.GetBytes(out, "messages.0.content.3.signature").String() != "good" { + t.Errorf("thinking signature should be replaced with pool entry") + } +} + +func TestReplaceThinkingSignaturesInClaudeRequest_Cycles(t *testing.T) { + req := &antigravity.ClaudeRequest{ + Messages: []antigravity.ClaudeMessage{ + { + Role: "assistant", + Content: json.RawMessage(`[ + {"type":"thinking","thinking":"a","signature":"bad1"}, + {"type":"thinking","thinking":"b","signature":"bad2"}, + {"type":"thinking","thinking":"c","signature":"bad3"} + ]`), + }, + }, + } + pool := []string{"G0", "G1"} + replaced, err := ReplaceThinkingSignaturesInClaudeRequest(req, pool) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if replaced != 3 { + t.Fatalf("expected 3 replacements, got %d", replaced) + } + sigs := gjson.GetBytes(req.Messages[0].Content, "#.signature").Array() + want := []string{"G0", "G1", "G0"} + for i, s := range sigs { + if s.String() != want[i] { + t.Errorf("sig[%d]: got %q, want %q", i, s.String(), want[i]) + } + } +} + +func TestReplaceThinkingSignaturesInClaudeRequest_StringContentUntouched(t *testing.T) { + req := &antigravity.ClaudeRequest{ + Messages: []antigravity.ClaudeMessage{ + {Role: "user", Content: json.RawMessage(`"plain text"`)}, + }, + } + replaced, err := ReplaceThinkingSignaturesInClaudeRequest(req, []string{"x"}) + if err != nil || replaced != 0 { + t.Fatalf("expected 0 replacements for string content, got replaced=%d err=%v", replaced, err) + } +} + +func TestReplaceThinkingSignaturesInClaudeRequest_NilPoolNoop(t *testing.T) { + orig := json.RawMessage(`[{"type":"thinking","thinking":"a","signature":"bad"}]`) + req := &antigravity.ClaudeRequest{ + Messages: []antigravity.ClaudeMessage{{Role: "assistant", Content: orig}}, + } + replaced, err := ReplaceThinkingSignaturesInClaudeRequest(req, nil) + if err != nil || replaced != 0 { + t.Fatalf("expected 0 replacements for nil pool, got replaced=%d err=%v", replaced, err) + } + if string(req.Messages[0].Content) != string(orig) { + t.Errorf("content must be unchanged for nil pool") + } +}