iop/apps/edge/internal/openai/stream_gate_ingress_test.go

package openai

import (
	"bytes"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"iop/packages/go/config"
	"iop/packages/go/streamgate"
	iop "iop/proto/gen/iop"
)

func TestReadOpenAIIngressBodyBoundary(t *testing.T) {
	for _, tc := range []struct {
		name    string
		size    int
		wantErr bool
	}{
		{name: "limit minus one", size: 7},
		{name: "exact limit", size: 8},
		{name: "limit plus one", size: 9, wantErr: true},
	} {
		t.Run(tc.name, func(t *testing.T) {
			recorder := httptest.NewRecorder()
			request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(bytes.Repeat([]byte("x"), tc.size)))
			body, err := readOpenAIIngressBody(recorder, request, 8)
			if tc.wantErr {
				if !errors.Is(err, errOpenAIIngressTooLarge) {
					t.Fatalf("error = %v, want ingress-too-large", err)
				}
				return
			}
			if err != nil || len(body) != tc.size {
				t.Fatalf("body len/error = %d/%v, want %d/nil", len(body), err, tc.size)
			}
		})
	}
}

func TestOpenAIIngressSnapshotSharedBackingAndRelease(t *testing.T) {
	body := []byte(`{"model":"m","messages":[]}`)
	ingress, err := buildOpenAIIngressSnapshot(int64(len(body)), body, json.RawMessage(body))
	if err != nil {
		t.Fatalf("buildOpenAIIngressSnapshot: %v", err)
	}
	accessor, err := ingress.accessor()
	if err != nil {
		t.Fatalf("accessor: %v", err)
	}
	if accessor.RetainedCount() != 1 || accessor.RetainedBytes() != int64(len(body)) {
		t.Fatalf("retained count/bytes = %d/%d", accessor.RetainedCount(), accessor.RetainedBytes())
	}
	ingress.Close()
	ingress.Close()
	if !ingress.isClosed() {
		t.Fatal("ingress snapshot was not released")
	}
}

func TestOpenAIIngressTypedViewOverflow(t *testing.T) {
	body := []byte(`{"model":"m"}`)
	_, err := buildOpenAIIngressSnapshot(int64(len(body)), body, struct {
		Model string `json:"model"`
		Extra string `json:"extra"`
	}{Model: "m", Extra: "typed-view-expansion"})
	if !errors.Is(err, streamgate.ErrIngressSnapshotLimitExceeded) {
		t.Fatalf("error = %v, want typed-view overflow", err)
	}
}

func TestOpenAIIngressExcludesAuthorization(t *testing.T) {
	body := []byte(`{"model":"m","messages":[]}`)
	ingress, err := buildOpenAIIngressSnapshot(1024, body, json.RawMessage(body))
	if err != nil {
		t.Fatalf("buildOpenAIIngressSnapshot: %v", err)
	}
	defer ingress.Close()
	request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
	request.Header.Set("X-Provider-Token", "secret-never-snapshot")
	canonical, _ := ingress.canonicalBody()
	semantic, _ := ingress.semanticBody()
	if bytes.Contains(canonical, []byte("secret-never-snapshot")) || bytes.Contains(semantic, []byte("secret-never-snapshot")) {
		t.Fatal("provider auth leaked into ingress snapshot")
	}
}

func TestChatIngressOverflowDoesNotDispatch(t *testing.T) {
	service := &fakeRunService{}
	server := NewServer(config.EdgeOpenAIConf{
		Adapter: "ollama",
		Target:  "model",
		StreamEvidenceGate: config.StreamEvidenceGateConf{
			MaxIngressSnapshotBytes: 32,
		},
	}, service, nil)
	body := `{"model":"model","messages":[{"role":"user","content":"` + strings.Repeat("x", 64) + `"}]}`
	request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(body))
	recorder := httptest.NewRecorder()

	server.handleChatCompletions(recorder, request)

	if recorder.Code != http.StatusRequestEntityTooLarge {
		t.Fatalf("status = %d, want 413; body=%s", recorder.Code, recorder.Body.String())
	}
	if len(service.reqsSnapshot()) != 0 {
		t.Fatal("overflow request reached provider dispatch")
	}
	if strings.Contains(recorder.Body.String(), "32") || strings.Contains(recorder.Body.String(), "streamgate") {
		t.Fatalf("error leaked internal limit details: %s", recorder.Body.String())
	}
}

func TestChatIngressTypedOverflowDoesNotDispatch(t *testing.T) {
	body := `{"model":"model", "messages":[{"role":"user","content":"hello"}]}`
	service := &fakeRunService{}
	server := NewServer(config.EdgeOpenAIConf{
		Adapter: "ollama",
		Target:  "model",
		StreamEvidenceGate: config.StreamEvidenceGateConf{
			// The canonical body fits exactly. Its normalized typed view differs,
			// so the combined ledger must fail before identity/admission.
			MaxIngressSnapshotBytes: len(body),
		},
	}, service, nil)
	request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(body))
	recorder := httptest.NewRecorder()

	server.handleChatCompletions(recorder, request)

	if recorder.Code != http.StatusRequestEntityTooLarge {
		t.Fatalf("status = %d, want 413; body=%s", recorder.Code, recorder.Body.String())
	}
	if len(service.reqsSnapshot()) != 0 {
		t.Fatal("typed-view overflow reached provider dispatch")
	}
}

func TestResponsesIngressBoundaryAndTypedOverflowZeroDispatch(t *testing.T) {
	for _, tc := range []struct {
		name    string
		size    int
		wantErr bool
	}{
		{name: "limit-minus-one", size: 7},
		{name: "limit", size: 8},
		{name: "limit-plus-one", size: 9, wantErr: true},
	} {
		t.Run(tc.name, func(t *testing.T) {
			request := httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(bytes.Repeat([]byte("x"), tc.size)))
			_, err := readOpenAIIngressBody(httptest.NewRecorder(), request, 8)
			if errors.Is(err, errOpenAIIngressTooLarge) != tc.wantErr {
				t.Fatalf("error=%v wantErr=%t", err, tc.wantErr)
			}
		})
	}

	body := `{"model":"m","input":"x"}`
	for _, tc := range []struct {
		name  string
		limit int
	}{
		{name: "raw overflow", limit: len(body) - 1},
		{name: "typed view overflow", limit: len(body)},
	} {
		t.Run(tc.name, func(t *testing.T) {
			service := &fakeRunService{events: bufferedRunEvents(&iop.RunEvent{Type: "complete"})}
			server := NewServer(config.EdgeOpenAIConf{Adapter: "ollama", Target: "served", StreamEvidenceGate: config.StreamEvidenceGateConf{MaxIngressSnapshotBytes: tc.limit}}, service, nil)
			request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(body))
			recorder := httptest.NewRecorder()
			server.handleResponses(recorder, request)
			if recorder.Code != http.StatusRequestEntityTooLarge || len(service.reqsSnapshot()) != 0 {
				t.Fatalf("status=%d dispatches=%d body=%s", recorder.Code, len(service.reqsSnapshot()), recorder.Body.String())
			}
		})
	}
}