CLIProxyAPI/internal/runtime/executor/helps/payload_finalizer_test.go

package helps

import (
	"bytes"
	"context"
	"mime/multipart"
	"net/http"
	"testing"

	"github.com/router-for-me/CLIProxyAPI/v8/internal/config"
	cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executor"
	sdktranslator "github.com/router-for-me/CLIProxyAPI/v8/sdk/translator"
	"github.com/tidwall/gjson"
	"github.com/tidwall/sjson"
	"google.golang.org/protobuf/encoding/protowire"
)

func TestPayloadFinalizerDefaultsUseOriginalAndNormalizeBeforeRules(t *testing.T) {
	cfg := &config.Config{Payload: config.PayloadConfig{
		Default:  []config.PayloadRule{{Models: []config.PayloadModelRule{{Name: "alias", Protocol: "openai", FromProtocol: "claude", Headers: map[string]string{"X-Test": "yes"}}}, Params: map[string]any{"missing": "user default", "present": "not applied"}}},
		Override: []config.PayloadRule{{Models: []config.PayloadModelRule{{Name: "alias"}}, Params: map[string]any{"tools.0.function.parameters.properties.count.type": "number"}}},
		Filter:   []config.PayloadFilterRule{{Models: []config.PayloadModelRule{{Name: "alias"}}, Params: []string{"late"}}},
	}}
	req := cliproxyexecutor.Request{Model: "alias"}
	opts := cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatClaude, Headers: http.Header{"X-Test": {"yes"}, "User-Agent": {"codex-cli/0.1"}}}
	finalize := NewPayloadFinalizer(cfg, "openai", "upstream", "openai", "", []byte(`{"present":"caller"}`), req, opts)
	body := []byte(`{"missing":"built-in","present":"caller","late":"injected","tools":[{"type":"function","function":{"name":"f","parameters":{"type":"object","properties":{"count":{"type":"integer"}}}}}]}`)
	out := FinalizePayload(WithPayloadFinalizer(context.Background(), finalize), body)
	if gjson.GetBytes(out, "missing").String() != "user default" || gjson.GetBytes(out, "present").String() != "caller" || gjson.GetBytes(out, "late").Exists() || gjson.GetBytes(out, "tools.0.function.parameters.properties.count.type").String() != "number" {
		t.Fatalf("unexpected final payload: %s", out)
	}
}

func TestPayloadFinalizerDevinBusinessFieldsAndCredentials(t *testing.T) {
	wire := BuildDevinGetChatMessageRequest("secret-token", "device", "model", "system", []DevinPrompt{{MessageID: "id", Source: 1, Content: "hello"}}, []DevinTool{{Name: "f", Parameters: []byte(`{"type":"object"}`)}}, nil, 1024, "session", "cascade", nil)
	unchanged, _, errUnchanged := FinalizeDevinPayload(wire, func(body []byte) []byte { return body })
	if errUnchanged != nil || !bytes.Equal(unchanged, wire) {
		t.Fatalf("no-op changed protobuf: %v", errUnchanged)
	}
	final, view, errFinalize := FinalizeDevinPayload(wire, func(body []byte) []byte {
		body, _ = sjson.DeleteBytes(body, "system_prompt")
		body, _ = sjson.DeleteBytes(body, "tools")
		body, _ = sjson.SetBytes(body, "completion_config.max_tokens", 0)
		body, _ = sjson.SetBytes(body, "prompts.0.content", "configured")
		return body
	})
	if errFinalize != nil {
		t.Fatal(errFinalize)
	}
	if gjson.GetBytes(view, "system_prompt").Exists() || gjson.GetBytes(view, "tools").Exists() {
		t.Fatalf("filter undone: %s", view)
	}
	decoded, errDecode := decodeDevinPayload(final, devinPayloadFields)
	if errDecode != nil {
		t.Fatal(errDecode)
	}
	if decoded["completion_config"].(map[string]any)["max_tokens"] != uint64(0) {
		t.Fatalf("default was re-injected: %v", decoded)
	}
	fieldOne := func(wire []byte) []byte {
		for len(wire) > 0 {
			num, kind, tagLen := protowire.ConsumeTag(wire)
			n := protowire.ConsumeFieldValue(num, kind, wire[tagLen:])
			if num == 1 {
				return wire[:tagLen+n]
			}
			wire = wire[tagLen+n:]
		}
		return nil
	}
	if !bytes.Equal(fieldOne(wire), fieldOne(final)) {
		t.Fatal("payload rules changed credentials")
	}
}

func TestPayloadFinalizerMultipartPreservesFiles(t *testing.T) {
	var buffer bytes.Buffer
	writer := multipart.NewWriter(&buffer)
	if errWrite := writer.WriteField("model", "original"); errWrite != nil {
		t.Fatal(errWrite)
	}
	file, errFile := writer.CreateFormFile("image", "input.png")
	if errFile != nil {
		t.Fatal(errFile)
	}
	if _, errWrite := file.Write([]byte{0, 255, 13, 10}); errWrite != nil {
		t.Fatal(errWrite)
	}
	if errClose := writer.Close(); errClose != nil {
		t.Fatal(errClose)
	}
	cfg := &config.Config{Payload: config.PayloadConfig{Override: []config.PayloadRule{{Models: []config.PayloadModelRule{{Name: "original"}}, Params: map[string]any{"model": "configured"}}}}}
	req := cliproxyexecutor.Request{Model: "original", Payload: buffer.Bytes()}
	opts := cliproxyexecutor.Options{Headers: http.Header{"Content-Type": {writer.FormDataContentType()}}}
	body, contentType, errFinalize := ApplyMediaPayloadConfig(cfg, "openai", "original", "openai", req.Payload, writer.FormDataContentType(), req, opts)
	if errFinalize != nil {
		t.Fatal(errFinalize)
	}
	view, _, errView := mediaPayloadJSON(body, contentType)
	if errView != nil {
		t.Fatal(errView)
	}
	if gjson.GetBytes(view, "model").String() != "configured" || gjson.GetBytes(view, "image.data").String() != "AP8NCg==" {
		t.Fatalf("unexpected multipart business payload: %s", view)
	}
}

func TestPayloadRulesMatchFinalBodyAndTrackPaths(t *testing.T) {
	cfg := &config.Config{Payload: config.PayloadConfig{
		Override: []config.PayloadRule{{Models: []config.PayloadModelRule{{Name: "*", Match: []map[string]any{{"max_tokens": 100}}}}, Params: map[string]any{"unexpected": true}},
			{Models: []config.PayloadModelRule{{Name: "*", Match: []map[string]any{{"max_tokens": 300}}}}, Params: map[string]any{"diagnostics.user": true}}},
		Filter: []config.PayloadFilterRule{{Models: []config.PayloadModelRule{{Name: "*", Match: []map[string]any{{"max_tokens": 300}}}}, Params: []string{"context_management", "messages.0"}}},
	}}
	original := []byte(`{"max_tokens":100,"messages":[{},{}]}`)
	body, touched := ApplyPayloadConfigWithTrackedPaths(cfg, "model", "claude", "claude", "", []byte(`{"max_tokens":300,"context_management":{"builtin":true},"messages":[{},{},{}]}`), original, "model", "", nil, "diagnostics", "context_management")
	if gjson.GetBytes(body, "unexpected").Exists() || !gjson.GetBytes(body, "diagnostics.user").Bool() || gjson.GetBytes(body, "context_management").Exists() || gjson.GetBytes(body, "messages.#").Int() != 2 {
		t.Fatalf("rules did not match final body: %s", body)
	}
	if !touched["diagnostics"] || !touched["context_management"] {
		t.Fatalf("final tracking lost: %v", touched)
	}
}