400 lines
16 KiB
Go
400 lines
16 KiB
Go
package openai
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"testing"
|
|
|
|
"iop/packages/go/streamgate"
|
|
)
|
|
|
|
func mustOpenAIRecoveryPlan(t *testing.T, id string, strategy streamgate.RecoveryStrategy, directive streamgate.RecoveryDirective) streamgate.RecoveryPlan {
|
|
t.Helper()
|
|
intent, err := streamgate.NewRecoveryIntent(strategy, directive, "openai.recovery", 10)
|
|
if err != nil {
|
|
t.Fatalf("NewRecoveryIntent: %v", err)
|
|
}
|
|
contributor, err := streamgate.NewRecoveryContributor("openai.edge", "filter.test", "rule.test")
|
|
if err != nil {
|
|
t.Fatalf("NewRecoveryContributor: %v", err)
|
|
}
|
|
policy, err := streamgate.NewRecoveryPolicySnapshot(3, map[streamgate.RecoveryStrategy]int{
|
|
streamgate.RecoveryStrategyExactReplay: 3,
|
|
streamgate.RecoveryStrategyContinuationRepair: 3,
|
|
streamgate.RecoveryStrategySchemaRepair: 3,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewRecoveryPolicySnapshot: %v", err)
|
|
}
|
|
usage, err := streamgate.NewRecoveryUsageSnapshot(0, nil)
|
|
if err != nil {
|
|
t.Fatalf("NewRecoveryUsageSnapshot: %v", err)
|
|
}
|
|
commit := streamgate.CommitStateTransportUncommitted
|
|
if strategy == streamgate.RecoveryStrategyContinuationRepair {
|
|
commit = streamgate.CommitStateStreamOpen
|
|
}
|
|
plan, err := streamgate.NewRecoveryPlan(streamgate.RecoveryEligibilityInput{
|
|
PlanID: id, IdempotencyKey: id + ":idempotency", Intent: intent,
|
|
Contributors: []streamgate.RecoveryContributor{contributor},
|
|
CommitState: commit, Policy: policy, Usage: usage,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewRecoveryPlan: %v", err)
|
|
}
|
|
return plan
|
|
}
|
|
|
|
func newOpenAIRebuilderFixture(t *testing.T, endpoint string, body []byte, maxBytes int64) (*openAIIngressSnapshot, *openAIRequestRebuilder, streamgate.RecoveryRequestSnapshotRef) {
|
|
t.Helper()
|
|
ingress, err := buildOpenAIIngressSnapshot(maxBytes, body, json.RawMessage(body))
|
|
if err != nil {
|
|
t.Fatalf("buildOpenAIIngressSnapshot: %v", err)
|
|
}
|
|
t.Cleanup(ingress.Close)
|
|
rebuilder, err := newOpenAIRequestRebuilder(ingress, endpoint)
|
|
if err != nil {
|
|
t.Fatalf("newOpenAIRequestRebuilder: %v", err)
|
|
}
|
|
t.Cleanup(rebuilder.Close)
|
|
ref, err := ingress.recoveryRef()
|
|
if err != nil {
|
|
t.Fatalf("recoveryRef: %v", err)
|
|
}
|
|
return ingress, rebuilder, ref
|
|
}
|
|
|
|
func TestOpenAIRequestRebuilderExactByteIdentity(t *testing.T) {
|
|
body := []byte("{\n \"unknown\" : [1, 2], \"model\" : \"alias\", \"messages\" : []\n}\n")
|
|
_, rebuilder, ref := newOpenAIRebuilderFixture(t, openAIRebuildEndpointChat, body, 4096)
|
|
directive, err := streamgate.NewRecoveryDirectiveExact(ref.SnapshotRef())
|
|
if err != nil {
|
|
t.Fatalf("NewRecoveryDirectiveExact: %v", err)
|
|
}
|
|
plan := mustOpenAIRecoveryPlan(t, "plan.exact", streamgate.RecoveryStrategyExactReplay, directive)
|
|
draft, err := rebuilder.RebuildRequest(context.Background(), ref, plan)
|
|
if err != nil {
|
|
t.Fatalf("RebuildRequest: %v", err)
|
|
}
|
|
lease, err := rebuilder.RebuiltStore().take(draft.RequestRef())
|
|
if err != nil {
|
|
t.Fatalf("take: %v", err)
|
|
}
|
|
defer lease.release()
|
|
got, err := lease.body()
|
|
if err != nil {
|
|
t.Fatalf("body: %v", err)
|
|
}
|
|
if !bytes.Equal(got, body) {
|
|
t.Fatalf("exact rebuild changed bytes:\n%s", got)
|
|
}
|
|
}
|
|
|
|
func TestOpenAIRequestRebuilderContinuationPreservesNonTargetBytes(t *testing.T) {
|
|
body := []byte("{\n \"unknown\" : { \"x\": 1 }, \"model\" : \"alias\", \"messages\" : [ {\"role\":\"user\",\"content\":\"old\"} ], \"tail\" : true\n}")
|
|
_, rebuilder, ref := newOpenAIRebuilderFixture(t, openAIRebuildEndpointChat, body, 8192)
|
|
patch := json.RawMessage(`[{"role":"assistant","content":"safe prefix"},{"role":"user","content":"continue"}]`)
|
|
if err := rebuilder.PatchStore().PutContinuation("snapshot.continue", 17, patch); err != nil {
|
|
t.Fatalf("PutContinuation: %v", err)
|
|
}
|
|
directive, _ := streamgate.NewRecoveryDirectiveContinuation(17, "snapshot.continue")
|
|
plan := mustOpenAIRecoveryPlan(t, "plan.continue", streamgate.RecoveryStrategyContinuationRepair, directive)
|
|
draft, err := rebuilder.RebuildRequest(context.Background(), ref, plan)
|
|
if err != nil {
|
|
t.Fatalf("RebuildRequest: %v", err)
|
|
}
|
|
lease, _ := rebuilder.RebuiltStore().take(draft.RequestRef())
|
|
defer lease.release()
|
|
got, _ := lease.body()
|
|
want := bytes.Replace(body, []byte(`[ {"role":"user","content":"old"} ]`), patch, 1)
|
|
if !bytes.Equal(got, want) {
|
|
t.Fatalf("continuation rebuild changed non-target bytes:\n got=%s\nwant=%s", got, want)
|
|
}
|
|
if draft.RetainedBytes() > draft.PeakBytes() || draft.PeakBytes() > draft.MaxBytes() {
|
|
t.Fatalf("invalid draft memory accounting: retained=%d peak=%d max=%d", draft.RetainedBytes(), draft.PeakBytes(), draft.MaxBytes())
|
|
}
|
|
}
|
|
|
|
func TestOpenAIRequestRebuilderResponsesSchemaPatch(t *testing.T) {
|
|
body := []byte(`{ "model" : "alias", "input" : { "old" : true }, "custom" : [3, 2, 1] }`)
|
|
_, rebuilder, ref := newOpenAIRebuilderFixture(t, openAIRebuildEndpointResponses, body, 4096)
|
|
patch := json.RawMessage(`[{"role":"user","content":"fixed"}]`)
|
|
if err := rebuilder.PatchStore().PutSchema("schema.responses", "patch.input", patch); err != nil {
|
|
t.Fatalf("PutSchema: %v", err)
|
|
}
|
|
directive, _ := streamgate.NewRecoveryDirectiveSchema("schema.responses", "patch.input")
|
|
plan := mustOpenAIRecoveryPlan(t, "plan.schema", streamgate.RecoveryStrategySchemaRepair, directive)
|
|
draft, err := rebuilder.RebuildRequest(context.Background(), ref, plan)
|
|
if err != nil {
|
|
t.Fatalf("RebuildRequest: %v", err)
|
|
}
|
|
lease, _ := rebuilder.RebuiltStore().take(draft.RequestRef())
|
|
defer lease.release()
|
|
got, _ := lease.body()
|
|
want := bytes.Replace(body, []byte(`{ "old" : true }`), patch, 1)
|
|
if !bytes.Equal(got, want) {
|
|
t.Fatalf("schema rebuild changed non-target bytes:\n got=%s\nwant=%s", got, want)
|
|
}
|
|
}
|
|
|
|
func TestOpenAIRequestRebuilderPatchPlusOutputPeakOverflow(t *testing.T) {
|
|
body := []byte(`{"model":"m","messages":[]}`)
|
|
patch := json.RawMessage(`[{"role":"user","content":"larger"}]`)
|
|
patchPlan, err := planTopLevelJSONPatches(body, []topLevelJSONPatch{{name: "messages", value: patch}})
|
|
if err != nil {
|
|
t.Fatalf("planTopLevelJSONPatches: %v", err)
|
|
}
|
|
maxBytes := int64(len(body) + len(patch) + patchPlan.outputSize - 1)
|
|
ingress, rebuilder, ref := newOpenAIRebuilderFixture(t, openAIRebuildEndpointChat, body, maxBytes)
|
|
if err := rebuilder.PatchStore().PutSchema("schema.chat", "patch.messages", patch); err != nil {
|
|
t.Fatalf("PutSchema: %v", err)
|
|
}
|
|
directive, _ := streamgate.NewRecoveryDirectiveSchema("schema.chat", "patch.messages")
|
|
plan := mustOpenAIRecoveryPlan(t, "plan.overflow", streamgate.RecoveryStrategySchemaRepair, directive)
|
|
_, err = rebuilder.RebuildRequest(context.Background(), ref, plan)
|
|
if !errors.Is(err, streamgate.ErrIngressSnapshotRebuildOverflow) {
|
|
t.Fatalf("error = %v, want rebuild overflow", err)
|
|
}
|
|
if len(rebuilder.RebuiltStore().leases) != 0 {
|
|
t.Fatal("overflow retained a dispatchable request lease")
|
|
}
|
|
if len(rebuilder.PatchStore().schema) != 0 {
|
|
t.Fatal("overflow retained a one-shot patch")
|
|
}
|
|
accessor, accessErr := ingress.accessor()
|
|
if accessErr != nil {
|
|
t.Fatalf("accessor after pre-allocation overflow: %v", accessErr)
|
|
}
|
|
if got := accessor.ReservedTempBytes(); got != 0 {
|
|
t.Fatalf("reserved bytes after overflow = %d, want 0", got)
|
|
}
|
|
}
|
|
|
|
func TestProviderRequestPatchesPreserveUnknownOrderAndWhitespace(t *testing.T) {
|
|
chat := []byte(`{ "z" : 1, "model" : "alias", "max_completion_tokens" : 7, "unknown" : { "a" : 2 } }`)
|
|
maxTokens := 9
|
|
gotChat, err := rewriteChatCompletionModel(chat, "served", chatCompletionRequest{MaxTokens: &maxTokens})
|
|
if err != nil {
|
|
t.Fatalf("rewriteChatCompletionModel: %v", err)
|
|
}
|
|
if !bytes.Contains(gotChat, []byte(`"z" : 1`)) || !bytes.Contains(gotChat, []byte(`"unknown" : { "a" : 2 }`)) {
|
|
t.Fatalf("chat unknown byte ranges changed: %s", gotChat)
|
|
}
|
|
if bytes.Contains(gotChat, []byte("max_completion_tokens")) || !bytes.Contains(gotChat, []byte(`"max_tokens":9`)) {
|
|
t.Fatalf("chat policy patch mismatch: %s", gotChat)
|
|
}
|
|
|
|
responses := []byte("{\n \"input\" : \"hello\", \"model\" : \"alias\", \"future\" : {\"x\":true}\n}")
|
|
gotResponses, err := rewriteResponsesModel(responses, "served")
|
|
if err != nil {
|
|
t.Fatalf("rewriteResponsesModel: %v", err)
|
|
}
|
|
wantResponses := bytes.Replace(responses, []byte(`"alias"`), []byte(`"served"`), 1)
|
|
if !bytes.Equal(gotResponses, wantResponses) {
|
|
t.Fatalf("responses non-target bytes changed:\n%s", gotResponses)
|
|
}
|
|
}
|
|
|
|
func TestOpenAIRequestRebuilderRejectsReferenceMismatch(t *testing.T) {
|
|
body := []byte(`{"model":"m","messages":[]}`)
|
|
_, rebuilder, ref := newOpenAIRebuilderFixture(t, openAIRebuildEndpointChat, body, 1024)
|
|
directive, _ := streamgate.NewRecoveryDirectiveExact(ref.SnapshotRef())
|
|
plan := mustOpenAIRecoveryPlan(t, "plan.mismatch", streamgate.RecoveryStrategyExactReplay, directive)
|
|
wrong, err := streamgate.NewRecoveryRequestSnapshotRef("openai.ingress.wrong", ref.RetainedBytes(), ref.PeakBytes(), ref.MaxBytes())
|
|
if err != nil {
|
|
t.Fatalf("NewRecoveryRequestSnapshotRef: %v", err)
|
|
}
|
|
if _, err := rebuilder.RebuildRequest(context.Background(), wrong, plan); err == nil {
|
|
t.Fatal("mismatched snapshot reference was accepted")
|
|
}
|
|
}
|
|
|
|
func Example_openAIRequestRebuilder() {
|
|
fmt.Println(openAIRebuildFamily)
|
|
// Output: openai.json
|
|
}
|
|
|
|
type cancelAfterFirstErrContext struct {
|
|
context.Context
|
|
calls int
|
|
}
|
|
|
|
func (c *cancelAfterFirstErrContext) Err() error {
|
|
c.calls++
|
|
if c.calls > 1 {
|
|
return context.Canceled
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func TestOpenAIRequestRebuilderActualOwnedPeakBoundaries(t *testing.T) {
|
|
body := []byte(`{"model":"alias","input":{"old":true},"custom":7}`)
|
|
patch := json.RawMessage(`[{"role":"user","content":"fixed"}]`)
|
|
patchPlan, err := planTopLevelJSONPatches(body, []topLevelJSONPatch{{name: "input", value: patch}})
|
|
if err != nil {
|
|
t.Fatalf("planTopLevelJSONPatches: %v", err)
|
|
}
|
|
maxBytes := int64(len(body) + len(patch) + patchPlan.outputSize)
|
|
ingress, rebuilder, ref := newOpenAIRebuilderFixture(t, openAIRebuildEndpointResponses, body, maxBytes)
|
|
if err := rebuilder.PatchStore().PutSchema("schema.actual", "patch.input", patch); err != nil {
|
|
t.Fatalf("PutSchema: %v", err)
|
|
}
|
|
directive, _ := streamgate.NewRecoveryDirectiveSchema("schema.actual", "patch.input")
|
|
plan := mustOpenAIRecoveryPlan(t, "plan.actual", streamgate.RecoveryStrategySchemaRepair, directive)
|
|
draft, err := rebuilder.RebuildRequest(context.Background(), ref, plan)
|
|
if err != nil {
|
|
t.Fatalf("RebuildRequest at exact owned peak: %v", err)
|
|
}
|
|
wantRetained := uint64(len(body) + patchPlan.outputSize)
|
|
if draft.RetainedBytes() != wantRetained || draft.PeakBytes() != uint64(maxBytes) || draft.MaxBytes() != uint64(maxBytes) {
|
|
t.Fatalf("draft accounting = retained:%d peak:%d max:%d, want %d/%d/%d",
|
|
draft.RetainedBytes(), draft.PeakBytes(), draft.MaxBytes(), wantRetained, maxBytes, maxBytes)
|
|
}
|
|
accessor, err := ingress.accessor()
|
|
if err != nil {
|
|
t.Fatalf("accessor: %v", err)
|
|
}
|
|
if got := accessor.ReservedTempBytes(); got != 0 {
|
|
t.Fatalf("patch reservation after rebuild = %d, want 0", got)
|
|
}
|
|
|
|
lease, err := rebuilder.RebuiltStore().take(draft.RequestRef())
|
|
if err != nil {
|
|
t.Fatalf("take rebuilt lease: %v", err)
|
|
}
|
|
got, err := lease.body()
|
|
if err != nil {
|
|
t.Fatalf("lease body: %v", err)
|
|
}
|
|
typedAlias, err := lease.rebuilt.Accessor().TypedViewAlias(openAIRebuiltBodyViewName)
|
|
if err != nil {
|
|
t.Fatalf("TypedViewAlias: %v", err)
|
|
}
|
|
if len(got) == 0 || &got[0] != &typedAlias[0] {
|
|
t.Fatal("dispatch lease did not retain the committed owned output alias")
|
|
}
|
|
guard := lease.guard
|
|
lease.release()
|
|
if !lease.isReleased() || !guard.IsReleased() {
|
|
t.Fatal("rebuilt output lease did not release snapshot and guard")
|
|
}
|
|
}
|
|
|
|
func TestOpenAIRequestRebuilderPatchStoreBoundedOneShotRelease(t *testing.T) {
|
|
body := []byte(`{"model":"m","messages":[]}`)
|
|
patch := json.RawMessage(`["owned"]`)
|
|
ingress, rebuilder, ref := newOpenAIRebuilderFixture(t, openAIRebuildEndpointChat, body, int64(len(body)+128))
|
|
store := rebuilder.PatchStore()
|
|
if err := store.PutContinuation("snapshot.once", 9, patch); err != nil {
|
|
t.Fatalf("PutContinuation: %v", err)
|
|
}
|
|
accessor, _ := ingress.accessor()
|
|
if got := accessor.ReservedTempBytes(); got != int64(len(patch)) {
|
|
t.Fatalf("reserved patch bytes = %d, want %d", got, len(patch))
|
|
}
|
|
store.mu.Lock()
|
|
stored := store.continuation["snapshot.once"]
|
|
store.mu.Unlock()
|
|
if stored == nil || len(stored.value) == 0 || &stored.value[0] != &patch[0] {
|
|
t.Fatal("json.RawMessage patch was copied instead of ownership-transferred")
|
|
}
|
|
if err := store.PutContinuation("snapshot.once", 9, json.RawMessage(`["duplicate"]`)); !errors.Is(err, errOpenAIRecoveryPatchDuplicate) {
|
|
t.Fatalf("duplicate PutContinuation = %v, want duplicate", err)
|
|
}
|
|
if got := accessor.ReservedTempBytes(); got != int64(len(patch)) {
|
|
t.Fatalf("duplicate changed reservation to %d", got)
|
|
}
|
|
entry, err := store.takeContinuation("snapshot.once", 9)
|
|
if err != nil {
|
|
t.Fatalf("takeContinuation: %v", err)
|
|
}
|
|
if _, err := store.takeContinuation("snapshot.once", 9); err == nil {
|
|
t.Fatal("one-shot continuation patch was available twice")
|
|
}
|
|
entry.release()
|
|
if got := accessor.ReservedTempBytes(); got != 0 {
|
|
t.Fatalf("reservation after one-shot release = %d, want 0", got)
|
|
}
|
|
|
|
cancelPatch := json.RawMessage(`{"cancelled":true}`)
|
|
if err := store.PutSchema("schema.cancel", "patch.cancel", cancelPatch); err != nil {
|
|
t.Fatalf("PutSchema(cancel): %v", err)
|
|
}
|
|
directive, _ := streamgate.NewRecoveryDirectiveSchema("schema.cancel", "patch.cancel")
|
|
plan := mustOpenAIRecoveryPlan(t, "plan.cancel", streamgate.RecoveryStrategySchemaRepair, directive)
|
|
cancelCtx := &cancelAfterFirstErrContext{Context: context.Background()}
|
|
if _, err := rebuilder.RebuildRequest(cancelCtx, ref, plan); !errors.Is(err, context.Canceled) {
|
|
t.Fatalf("cancelled RebuildRequest = %v, want context canceled", err)
|
|
}
|
|
if got := accessor.ReservedTempBytes(); got != 0 {
|
|
t.Fatalf("reservation after cancellation = %d, want 0", got)
|
|
}
|
|
if len(store.schema) != 0 || len(rebuilder.RebuiltStore().leases) != 0 {
|
|
t.Fatal("cancellation retained a patch or dispatch lease")
|
|
}
|
|
|
|
closePatch := json.RawMessage(`{"close":true}`)
|
|
if err := store.PutSchema("schema.close", "patch.close", closePatch); err != nil {
|
|
t.Fatalf("PutSchema(close): %v", err)
|
|
}
|
|
rebuilder.Close()
|
|
if got := accessor.ReservedTempBytes(); got != 0 {
|
|
t.Fatalf("reservation after Close = %d, want 0", got)
|
|
}
|
|
if !store.closed || len(store.schema) != 0 || len(store.continuation) != 0 {
|
|
t.Fatal("Close did not empty and close the patch store")
|
|
}
|
|
|
|
oversizedIngress, oversizedRebuilder, _ := newOpenAIRebuilderFixture(t, openAIRebuildEndpointChat, body, int64(len(body)+2))
|
|
if err := oversizedRebuilder.PatchStore().PutSchema("schema.large", "patch.large", json.RawMessage(`{"large":true}`)); !errors.Is(err, streamgate.ErrIngressSnapshotRebuildOverflow) {
|
|
t.Fatalf("oversized PutSchema = %v, want rebuild overflow", err)
|
|
}
|
|
oversizedAccessor, _ := oversizedIngress.accessor()
|
|
if got := oversizedAccessor.ReservedTempBytes(); got != 0 {
|
|
t.Fatalf("oversized patch left %d reserved bytes", got)
|
|
}
|
|
if len(oversizedRebuilder.PatchStore().schema) != 0 {
|
|
t.Fatal("oversized patch entered the store")
|
|
}
|
|
}
|
|
|
|
func TestOpenAIProviderBodyLeaseRelease(t *testing.T) {
|
|
body := []byte(`{"model":"alias","input":"hello","custom":true}`)
|
|
modelJSON, _ := json.Marshal("served")
|
|
patchPlan, err := planTopLevelJSONPatches(body, []topLevelJSONPatch{{name: "model", value: modelJSON}})
|
|
if err != nil {
|
|
t.Fatalf("planTopLevelJSONPatches: %v", err)
|
|
}
|
|
ingress, err := buildOpenAIIngressSnapshot(int64(len(body)+patchPlan.outputSize), body, json.RawMessage(body))
|
|
if err != nil {
|
|
t.Fatalf("buildOpenAIIngressSnapshot: %v", err)
|
|
}
|
|
defer ingress.Close()
|
|
builder := newOpenAIProviderBodyBuilder(func(target string) (*openAIRebuiltLease, error) {
|
|
return rewriteResponsesModelFromIngress(ingress, target)
|
|
})
|
|
got, err := builder.BuildBody("served")
|
|
if err != nil {
|
|
t.Fatalf("BuildBody: %v", err)
|
|
}
|
|
if !bytes.Contains(got, []byte(`"model":"served"`)) || !bytes.Contains(got, []byte(`"custom":true`)) {
|
|
t.Fatalf("provider body rewrite mismatch: %s", got)
|
|
}
|
|
builder.mu.Lock()
|
|
lease := builder.lease
|
|
builder.mu.Unlock()
|
|
if lease == nil || lease.guard == nil || lease.isReleased() || lease.guard.IsReleased() {
|
|
t.Fatal("provider body lease was not live during synchronous submission")
|
|
}
|
|
guard := lease.guard
|
|
builder.Close()
|
|
if !lease.isReleased() || !guard.IsReleased() {
|
|
t.Fatal("provider body lease was not released after synchronous submission")
|
|
}
|
|
if _, err := builder.BuildBody("served-again"); err == nil {
|
|
t.Fatal("provider body builder allowed a second build")
|
|
}
|
|
}
|