Sync from internal 2026-07-07

Author: Nexus-ABC

src/pkg/spec/openai/chat.go

@@ -1,20 +1,38 @@
 package openai
 
 import (
+	"strconv"
+	"sync/atomic"
+	"time"
+
 	"nexus-mock-provider/pkg/core"
 	"nexus-mock-provider/pkg/sse"
 
 	"github.com/gin-gonic/gin"
 )
 
 const (
-	completionID        = "chatcmpl-mock"
 	created       int64 = 1700000000
 	objChat             = "chat.completion"
 	objChunk            = "chat.completion.chunk"
 	roleAssistant       = "assistant"
 )
 
+// Real providers (and the rig's audit test) key per-request records by the UPSTREAM response id, so a
+// constant id makes clients that upsert on it (e.g. LiteLLM's SpendLogs) collapse every request into
+// one row. Emit a UNIQUE id per completion. Hot path is lock-free: a per-process nonce computed once
+// at init (keeps ids unique across restarts) + a single atomic increment — no lock, one small alloc,
+// negligible vs marshaling the response. Keeps the "chatcmpl-mock" prefix so the rig's health/smoke
+// greps (substring match) still pass.
+var (
+	completionSeq    uint64
+	completionIDBase = "chatcmpl-mock-" + strconv.FormatInt(time.Now().UnixNano(), 36) + "-"
+)
+
+func nextCompletionID() string {
+	return completionIDBase + strconv.FormatUint(atomic.AddUint64(&completionSeq, 1), 36)
+}
+
 // ---- request ----
 
 type chatRequest struct {
@@ -104,7 +122,7 @@ func handleChat(c *gin.Context) {
 func writeNonStream(c *gin.Context, r core.ChatResult) {
 	reason := string(r.FinishReason)
 	c.JSON(200, completion{
-		ID:      completionID,
+		ID:      nextCompletionID(),
 		Object:  objChat,
 		Created: created,
 		Model:   r.Model,
@@ -128,7 +146,11 @@ func writeNonStream(c *gin.Context, r core.ChatResult) {
 func streamFrames(r core.ChatResult, includeUsage bool) []sse.Frame {
 	frames := make([]sse.Frame, 0, len(r.Tokens)+4)
 
-	frames = append(frames, sse.DataFrame(jsonStr(chunk(r.Model, []choice{{
+	// One id for the WHOLE completion — every chunk of a streamed response shares it (OpenAI
+	// semantics), generated once so it's stable across this response's frames but unique per response.
+	id := nextCompletionID()
+
+	frames = append(frames, sse.DataFrame(jsonStr(chunk(id, r.Model, []choice{{
 		Index: 0, Delta: &delta{Role: roleAssistant},
 	}}, nil))))
 
@@ -137,18 +159,18 @@ func streamFrames(r core.ChatResult, includeUsage bool) []sse.Frame {
 		if i > 0 {
 			piece = " " + tok
 		}
-		frames = append(frames, sse.DataFrame(jsonStr(chunk(r.Model, []choice{{
+		frames = append(frames, sse.DataFrame(jsonStr(chunk(id, r.Model, []choice{{
 			Index: 0, Delta: &delta{Content: piece},
 		}}, nil))))
 	}
 
 	reason := string(r.FinishReason)
-	frames = append(frames, sse.DataFrame(jsonStr(chunk(r.Model, []choice{{
+	frames = append(frames, sse.DataFrame(jsonStr(chunk(id, r.Model, []choice{{
 		Index: 0, Delta: &delta{}, FinishReason: &reason,
 	}}, nil))))
 
 	if includeUsage {
-		frames = append(frames, sse.DataFrame(jsonStr(chunk(r.Model, []choice{}, &usage{
+		frames = append(frames, sse.DataFrame(jsonStr(chunk(id, r.Model, []choice{}, &usage{
 			PromptTokens:     r.PromptTokens,
 			CompletionTokens: r.CompletionTokens,
 			TotalTokens:      r.PromptTokens + r.CompletionTokens,
@@ -159,9 +181,9 @@ func streamFrames(r core.ChatResult, includeUsage bool) []sse.Frame {
 	return frames
 }
 
-func chunk(model string, choices []choice, u *usage) completion {
+func chunk(id, model string, choices []choice, u *usage) completion {
 	return completion{
-		ID:      completionID,
+		ID:      id,
 		Object:  objChunk,
 		Created: created,
 		Model:   model,

src/pkg/spec/openai/completion_id_test.go

@@ -0,0 +1,64 @@
+package openai
+
+import (
+	"strings"
+	"sync"
+	"testing"
+)
+
+// Every id must be unique AND carry the "chatcmpl-mock" prefix the rig's health greps rely on.
+func TestNextCompletionID_UniqueAndPrefixed(t *testing.T) {
+	const n = 100_000
+	seen := make(map[string]struct{}, n)
+	for i := 0; i < n; i++ {
+		id := nextCompletionID()
+		if !strings.HasPrefix(id, "chatcmpl-mock-") {
+			t.Fatalf("id %q missing chatcmpl-mock prefix", id)
+		}
+		if _, dup := seen[id]; dup {
+			t.Fatalf("duplicate id %q at %d", id, i)
+		}
+		seen[id] = struct{}{}
+	}
+}
+
+// Lock-free counter must stay unique under heavy concurrency (run with -race).
+func TestNextCompletionID_ConcurrentUnique(t *testing.T) {
+	const goroutines, per = 64, 10_000
+	var mu sync.Mutex
+	seen := make(map[string]struct{}, goroutines*per)
+	var wg sync.WaitGroup
+	for g := 0; g < goroutines; g++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			local := make([]string, per)
+			for i := range local {
+				local[i] = nextCompletionID()
+			}
+			mu.Lock()
+			for _, id := range local {
+				if _, dup := seen[id]; dup {
+					t.Errorf("duplicate id %q under concurrency", id)
+				}
+				seen[id] = struct{}{}
+			}
+			mu.Unlock()
+		}()
+	}
+	wg.Wait()
+	if len(seen) != goroutines*per {
+		t.Fatalf("expected %d unique ids, got %d", goroutines*per, len(seen))
+	}
+}
+
+// Hot path is a single atomic increment + one format — must stay in the low-ns range so the mock's
+// throughput is unaffected. (Reported as ns/op; typically well under ~50 ns on the c6i target.)
+func BenchmarkNextCompletionID(b *testing.B) {
+	b.ReportAllocs()
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			_ = nextCompletionID()
+		}
+	})
+}