iop/apps/edge/internal/openai/openai_request_rebuilder_test.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 TestOpenAIRequestRebuilderOverflowCreatesNoDispatchLease(t *testing.T) {
	body := []byte(`{"model":"m","messages":[]}`)
	_, rebuilder, ref := newOpenAIRebuilderFixture(t, openAIRebuildEndpointChat, body, int64(len(body)))
	if err := rebuilder.PatchStore().PutSchema("schema.chat", "patch.messages", json.RawMessage(`[{"role":"user","content":"larger"}]`)); 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")
	}
}

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
}