iop/apps/node/internal/adapters/openai_compat/thinking_policy_test.go

package openai_compat

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"go.uber.org/zap"

	"iop/apps/node/internal/runtime"
	"iop/packages/go/config"
)

func TestOpenAICompatExecuteMapsThinkFalseForVLLMChatTemplateKwargs(t *testing.T) {
	var body map[string]any
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
			t.Fatalf("decode: %v", err)
		}
		w.Header().Set("Content-Type", "text/event-stream")
		_, _ = fmt.Fprintf(w, "data: %s\n\n", `{"choices":[{"delta":{"content":"ok"}}]}`)
		_, _ = fmt.Fprintf(w, "data: [DONE]\n\n")
	}))
	defer server.Close()

	adapter := New(config.OpenAICompatConf{
		Provider: "vllm",
		Endpoint: server.URL,
	}, zap.NewNop())

	sink := &fakeSink{}
	err := adapter.Execute(context.Background(), runtime.ExecutionSpec{
		RunID:  "run-vllm",
		Target: "qwen-model",
		Input: map[string]any{
			"prompt": "hello",
			"think":  false,
		},
	}, sink)
	if err != nil {
		t.Fatalf("Execute failed: %v", err)
	}

	ctk, ok := body["chat_template_kwargs"].(map[string]any)
	if !ok {
		t.Fatalf("expected chat_template_kwargs to be a map, got %T", body["chat_template_kwargs"])
	}
	if ctk["enable_thinking"] != false {
		t.Fatalf("expected enable_thinking to be false, got %v", ctk["enable_thinking"])
	}
}

func TestOpenAICompatExecuteVLLMMLXUnsupportedThinkDisable(t *testing.T) {
	// The vLLM-MLX runtime cannot produce a reasoning-free streaming response,
	// so think=false / reasoning_effort=none must fail loudly instead of being
	// silently accepted as a reasoning stream.
	for _, disable := range []map[string]any{
		{"prompt": "hello", "think": false},
		{"prompt": "hello", "reasoning_effort": "none"},
	} {
		adapter := New(config.OpenAICompatConf{
			Provider: "vllm-mlx",
			Endpoint: "http://localhost:8000",
		}, zap.NewNop())

		sink := &fakeSink{}
		err := adapter.Execute(context.Background(), runtime.ExecutionSpec{
			RunID:  "run-vllm-mlx-disable",
			Target: "qwen-model",
			Input:  disable,
		}, sink)
		if err == nil {
			t.Fatalf("expected error for vllm-mlx think disable %v, got nil", disable)
		}
		if !strings.Contains(err.Error(), "unsupported think control") {
			t.Fatalf("expected error message to contain 'unsupported think control' for %v, got: %v", disable, err)
		}
	}
}

func TestOpenAICompatExecuteMapsThinkFalseForLemonadeChatTemplateKwargs(t *testing.T) {
	var body map[string]any
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
			t.Fatalf("decode: %v", err)
		}
		w.Header().Set("Content-Type", "text/event-stream")
		_, _ = fmt.Fprintf(w, "data: %s\n\n", `{"choices":[{"delta":{"content":"ok"}}]}`)
		_, _ = fmt.Fprintf(w, "data: [DONE]\n\n")
	}))
	defer server.Close()

	adapter := New(config.OpenAICompatConf{
		Provider: "lemonade",
		Endpoint: server.URL,
	}, zap.NewNop())

	sink := &fakeSink{}
	err := adapter.Execute(context.Background(), runtime.ExecutionSpec{
		RunID:  "run-lemonade",
		Target: "lemonade-model",
		Input: map[string]any{
			"prompt": "hello",
			"think":  false,
		},
	}, sink)
	if err != nil {
		t.Fatalf("Execute failed: %v", err)
	}

	if _, ok := body["think"]; ok {
		t.Fatalf("expected top-level think to be omitted for lemonade, got %v", body["think"])
	}
	ctk, ok := body["chat_template_kwargs"].(map[string]any)
	if !ok {
		t.Fatalf("expected chat_template_kwargs to be a map, got %T", body["chat_template_kwargs"])
	}
	if ctk["enable_thinking"] != false {
		t.Fatalf("expected enable_thinking to be false, got %v", ctk["enable_thinking"])
	}
}

func TestOpenAICompatExecuteOmitsThinkControlWhenRequestOmitsIt(t *testing.T) {
	for _, provider := range []string{"vllm", "vllm-mlx", "lemonade", ""} {
		t.Run(provider, func(t *testing.T) {
			var body map[string]any
			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
					t.Fatalf("decode: %v", err)
				}
				w.Header().Set("Content-Type", "text/event-stream")
				_, _ = fmt.Fprintf(w, "data: %s\n\n", `{"choices":[{"delta":{"content":"ok"}}]}`)
				_, _ = fmt.Fprintf(w, "data: [DONE]\n\n")
			}))
			defer server.Close()

			adapter := New(config.OpenAICompatConf{
				Provider: provider,
				Endpoint: server.URL,
			}, zap.NewNop())

			sink := &fakeSink{}
			err := adapter.Execute(context.Background(), runtime.ExecutionSpec{
				RunID:  "run-omit",
				Target: "qwen-model",
				Input: map[string]any{
					"prompt": "hello",
				},
			}, sink)
			if err != nil {
				t.Fatalf("Execute failed: %v", err)
			}

			for _, key := range []string{"think", "reasoning_effort", "thinking_token_budget", "chat_template_kwargs"} {
				if _, ok := body[key]; ok {
					t.Fatalf("expected key %q to be omitted for provider %q", key, provider)
				}
			}
		})
	}
}

func TestOpenAICompatExecuteVLLMUnsupportedReasoningEffort(t *testing.T) {
	adapter := New(config.OpenAICompatConf{
		Provider: "vllm",
		Endpoint: "http://localhost:8000",
	}, zap.NewNop())

	sink := &fakeSink{}
	err := adapter.Execute(context.Background(), runtime.ExecutionSpec{
		RunID:  "run-unsupported",
		Target: "qwen-model",
		Input: map[string]any{
			"prompt":           "hello",
			"reasoning_effort": "low",
		},
	}, sink)
	if err == nil {
		t.Fatalf("expected error for unsupported reasoning_effort, got nil")
	}
	if !strings.Contains(err.Error(), "unsupported think control") {
		t.Fatalf("expected error message to contain 'unsupported think control', got: %v", err)
	}
}

func TestOpenAICompatExecuteVLLMUnsupportedNegativeBudget(t *testing.T) {
	adapter := New(config.OpenAICompatConf{
		Provider: "vllm",
		Endpoint: "http://localhost:8000",
	}, zap.NewNop())

	sink := &fakeSink{}
	err := adapter.Execute(context.Background(), runtime.ExecutionSpec{
		RunID:  "run-unsupported",
		Target: "qwen-model",
		Input: map[string]any{
			"prompt":                "hello",
			"thinking_token_budget": -10,
		},
	}, sink)
	if err == nil {
		t.Fatalf("expected error for negative budget, got nil")
	}
	if !strings.Contains(err.Error(), "unsupported think control") {
		t.Fatalf("expected error message to contain 'unsupported think control', got: %v", err)
	}
}

func TestOpenAICompatExecuteMapsReasoningEffortNone(t *testing.T) {
	t.Run("vllm", func(t *testing.T) {
		var body map[string]any
		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
				t.Fatalf("decode: %v", err)
			}
			w.Header().Set("Content-Type", "text/event-stream")
			_, _ = fmt.Fprintf(w, "data: %s\n\n", `{"choices":[{"delta":{"content":"ok"}}]}`)
			_, _ = fmt.Fprintf(w, "data: [DONE]\n\n")
		}))
		defer server.Close()

		adapter := New(config.OpenAICompatConf{
			Provider: "vllm",
			Endpoint: server.URL,
		}, zap.NewNop())

		sink := &fakeSink{}
		err := adapter.Execute(context.Background(), runtime.ExecutionSpec{
			RunID:  "run-vllm-none",
			Target: "qwen-model",
			Input: map[string]any{
				"prompt":           "hello",
				"reasoning_effort": "none",
			},
		}, sink)
		if err != nil {
			t.Fatalf("Execute failed: %v", err)
		}

		ctk, ok := body["chat_template_kwargs"].(map[string]any)
		if !ok {
			t.Fatalf("expected chat_template_kwargs to be a map, got %T", body["chat_template_kwargs"])
		}
		if ctk["enable_thinking"] != false {
			t.Fatalf("expected enable_thinking to be false, got %v", ctk["enable_thinking"])
		}
	})

	t.Run("lemonade", func(t *testing.T) {
		var body map[string]any
		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
				t.Fatalf("decode: %v", err)
			}
			w.Header().Set("Content-Type", "text/event-stream")
			_, _ = fmt.Fprintf(w, "data: %s\n\n", `{"choices":[{"delta":{"content":"ok"}}]}`)
			_, _ = fmt.Fprintf(w, "data: [DONE]\n\n")
		}))
		defer server.Close()

		adapter := New(config.OpenAICompatConf{
			Provider: "lemonade",
			Endpoint: server.URL,
		}, zap.NewNop())

		sink := &fakeSink{}
		err := adapter.Execute(context.Background(), runtime.ExecutionSpec{
			RunID:  "run-lemonade-none",
			Target: "lemonade-model",
			Input: map[string]any{
				"prompt":           "hello",
				"reasoning_effort": "none",
			},
		}, sink)
		if err != nil {
			t.Fatalf("Execute failed: %v", err)
		}

		if _, ok := body["reasoning_effort"]; ok {
			t.Fatalf("expected reasoning_effort to be omitted, got %v", body["reasoning_effort"])
		}
		ctk, ok := body["chat_template_kwargs"].(map[string]any)
		if !ok {
			t.Fatalf("expected chat_template_kwargs to be a map, got %T", body["chat_template_kwargs"])
		}
		if ctk["enable_thinking"] != false {
			t.Fatalf("expected enable_thinking to be false, got %v", ctk["enable_thinking"])
		}
	})
}

func TestOpenAICompatExecuteMapsThinkTrueForVLLMChatTemplateKwargs(t *testing.T) {
	var body map[string]any
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
			t.Fatalf("decode: %v", err)
		}
		w.Header().Set("Content-Type", "text/event-stream")
		_, _ = fmt.Fprintf(w, "data: %s\n\n", `{"choices":[{"delta":{"content":"ok"}}]}`)
		_, _ = fmt.Fprintf(w, "data: [DONE]\n\n")
	}))
	defer server.Close()

	adapter := New(config.OpenAICompatConf{
		Provider: "vllm",
		Endpoint: server.URL,
	}, zap.NewNop())

	sink := &fakeSink{}
	err := adapter.Execute(context.Background(), runtime.ExecutionSpec{
		RunID:  "run-vllm",
		Target: "qwen-model",
		Input: map[string]any{
			"prompt": "hello",
			"think":  true,
		},
	}, sink)
	if err != nil {
		t.Fatalf("Execute failed: %v", err)
	}

	ctk, ok := body["chat_template_kwargs"].(map[string]any)
	if !ok {
		t.Fatalf("expected chat_template_kwargs to be a map, got %T", body["chat_template_kwargs"])
	}
	if ctk["enable_thinking"] != true {
		t.Fatalf("expected enable_thinking to be true, got %v", ctk["enable_thinking"])
	}
}

func TestOpenAICompatExecuteMapsThinkTrueForVLLMMLXChatTemplateKwargs(t *testing.T) {
	var body map[string]any
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
			t.Fatalf("decode: %v", err)
		}
		w.Header().Set("Content-Type", "text/event-stream")
		_, _ = fmt.Fprintf(w, "data: %s\n\n", `{"choices":[{"delta":{"content":"ok"}}]}`)
		_, _ = fmt.Fprintf(w, "data: [DONE]\n\n")
	}))
	defer server.Close()

	adapter := New(config.OpenAICompatConf{
		Provider: "vllm-mlx",
		Endpoint: server.URL,
	}, zap.NewNop())

	sink := &fakeSink{}
	err := adapter.Execute(context.Background(), runtime.ExecutionSpec{
		RunID:  "run-vllm-mlx",
		Target: "qwen-model",
		Input: map[string]any{
			"prompt": "hello",
			"think":  true,
		},
	}, sink)
	if err != nil {
		t.Fatalf("Execute failed: %v", err)
	}

	ctk, ok := body["chat_template_kwargs"].(map[string]any)
	if !ok {
		t.Fatalf("expected chat_template_kwargs to be a map, got %T", body["chat_template_kwargs"])
	}
	if ctk["enable_thinking"] != true {
		t.Fatalf("expected enable_thinking to be true, got %v", ctk["enable_thinking"])
	}
}

func TestOpenAICompatExecuteMapsThinkTrueForLemonadeChatTemplateKwargs(t *testing.T) {
	var body map[string]any
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
			t.Fatalf("decode: %v", err)
		}
		w.Header().Set("Content-Type", "text/event-stream")
		_, _ = fmt.Fprintf(w, "data: %s\n\n", `{"choices":[{"delta":{"content":"ok"}}]}`)
		_, _ = fmt.Fprintf(w, "data: [DONE]\n\n")
	}))
	defer server.Close()

	adapter := New(config.OpenAICompatConf{
		Provider: "lemonade",
		Endpoint: server.URL,
	}, zap.NewNop())

	sink := &fakeSink{}
	err := adapter.Execute(context.Background(), runtime.ExecutionSpec{
		RunID:  "run-lemonade",
		Target: "lemonade-model",
		Input: map[string]any{
			"prompt":                "hello",
			"think":                 true,
			"thinking_token_budget": 1024,
		},
	}, sink)
	if err != nil {
		t.Fatalf("Execute failed: %v", err)
	}

	if _, ok := body["think"]; ok {
		t.Fatalf("expected top-level think to be omitted for lemonade, got %v", body["think"])
	}
	if _, ok := body["thinking_token_budget"]; ok {
		t.Fatalf("expected top-level thinking_token_budget to be omitted for lemonade, got %v", body["thinking_token_budget"])
	}
	ctk, ok := body["chat_template_kwargs"].(map[string]any)
	if !ok {
		t.Fatalf("expected chat_template_kwargs to be a map, got %T", body["chat_template_kwargs"])
	}
	if ctk["enable_thinking"] != true {
		t.Fatalf("expected enable_thinking to be true, got %v", ctk["enable_thinking"])
	}
	if ctk["thinking_token_budget"] != float64(1024) {
		t.Fatalf("expected thinking_token_budget to be float64 1024, got %v", ctk["thinking_token_budget"])
	}
}

func TestOpenAICompatExecuteMapsThinkingTokenBudgetForVLLMChatTemplateKwargs(t *testing.T) {
	var body map[string]any
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
			t.Fatalf("decode: %v", err)
		}
		w.Header().Set("Content-Type", "text/event-stream")
		_, _ = fmt.Fprintf(w, "data: %s\n\n", `{"choices":[{"delta":{"content":"ok"}}]}`)
		_, _ = fmt.Fprintf(w, "data: [DONE]\n\n")
	}))
	defer server.Close()

	adapter := New(config.OpenAICompatConf{
		Provider: "vllm",
		Endpoint: server.URL,
	}, zap.NewNop())

	sink := &fakeSink{}
	err := adapter.Execute(context.Background(), runtime.ExecutionSpec{
		RunID:  "run-vllm-budget",
		Target: "qwen-model",
		Input: map[string]any{
			"prompt":                "hello",
			"thinking_token_budget": 2048,
		},
	}, sink)
	if err != nil {
		t.Fatalf("Execute failed: %v", err)
	}

	if _, ok := body["thinking_token_budget"]; ok {
		t.Fatalf("expected top-level thinking_token_budget to be omitted for vllm, got %v", body["thinking_token_budget"])
	}

	ctk, ok := body["chat_template_kwargs"].(map[string]any)
	if !ok {
		t.Fatalf("expected chat_template_kwargs to be a map, got %T", body["chat_template_kwargs"])
	}
	if ctk["enable_thinking"] != true {
		t.Fatalf("expected enable_thinking to be true for budget-only request, got %v", ctk["enable_thinking"])
	}
	if ctk["thinking_token_budget"] != float64(2048) {
		t.Fatalf("expected thinking_token_budget to be 2048, got %v", ctk["thinking_token_budget"])
	}
}

func TestOpenAICompatExecuteVLLMUnsupportedInvalidBudget(t *testing.T) {
	adapter := New(config.OpenAICompatConf{
		Provider: "vllm",
		Endpoint: "http://localhost:8000",
	}, zap.NewNop())

	sink := &fakeSink{}
	invalidBudgets := []any{"invalid", -5, 12.34}
	for _, val := range invalidBudgets {
		err := adapter.Execute(context.Background(), runtime.ExecutionSpec{
			RunID:  "run-vllm-invalid-budget",
			Target: "qwen-model",
			Input: map[string]any{
				"prompt":                "hello",
				"thinking_token_budget": val,
			},
		}, sink)
		if err == nil {
			t.Fatalf("expected error for budget val %v, got nil", val)
		}
		if !strings.Contains(err.Error(), "unsupported think control") {
			t.Fatalf("expected error message to contain 'unsupported think control' for val %v, got: %v", val, err)
		}
	}
}