package openai

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"
	"time"

	edgeservice "iop/apps/edge/internal/service"
	"iop/packages/go/config"
	"iop/packages/go/streamgate"
	iop "iop/proto/gen/iop"
)

type hotPathSequenceSource struct {
	events []streamgate.NormalizedEvent
	index  int
}

func TestHotPathOuterTurnIntegrationKeepsRemainingStageBudget(t *testing.T) {
	outer := newHotPathOuterTurn("turn-integration")
	outer.recordCollectedStage(normalizedStageOutput{
		ResponseID:  "selector-response",
		OpenAIUsage: &openAIUsage{PromptTokens: 3, CompletionTokens: 4, TotalTokens: 7},
	})
	if got := hotPathRemainingOutputTokens(10, outer); got != 6 {
		t.Fatalf("remaining output tokens = %d, want 6", got)
	}
	limited := newHotPathCallerCappedOuterTurn("turn-limited", 10)
	limited.recordCollectedStage(normalizedStageOutput{
		ResponseID: "limited-stage", OpenAIUsage: &openAIUsage{CompletionTokens: 4},
	})
	if state := limited.outputBudget(); !state.Limited || state.Exhausted || state.Remaining != 6 {
		t.Fatalf("positive budget state = %+v, want limited remaining 6", state)
	}
	unlimited := newHotPathCallerCappedOuterTurn("turn-unlimited", 0).outputBudget()
	if unlimited.Limited || unlimited.Exhausted || unlimited.Remaining != 0 {
		t.Fatalf("unlimited budget state = %+v", unlimited)
	}
	limited.recordCollectedStage(normalizedStageOutput{
		ResponseID: "exhausting-stage", OpenAIUsage: &openAIUsage{CompletionTokens: 6},
	})
	if state := limited.outputBudget(); !state.Limited || !state.Exhausted || state.Remaining != 0 {
		t.Fatalf("exhausted budget state = %+v", state)
	}

	for _, test := range []struct {
		name string
		body func() ([]byte, error)
	}{
		{
			name: "chat",
			body: func() ([]byte, error) {
				return hotPathChatStageBody(hotPathDispatchSnapshot{
					Stage:        config.ExecutionRouteStage{Options: map[string]any{"max_tokens": 999}},
					OutputBudget: hotPathOutputBudget{Limited: true, Remaining: 6},
				}, "continue", "stage-model")
			},
		},
		{
			name: "anthropic",
			body: func() ([]byte, error) {
				return hotPathAnthropicStageBody(hotPathDispatchSnapshot{
					Stage:        config.ExecutionRouteStage{Options: map[string]any{"max_tokens": 999}},
					OutputBudget: hotPathOutputBudget{Limited: true, Remaining: 6},
				}, "continue", "stage-model")
			},
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			body, err := test.body()
			if err != nil {
				t.Fatal(err)
			}
			var decoded map[string]any
			if err := json.Unmarshal(body, &decoded); err != nil {
				t.Fatal(err)
			}
			if got := decoded["max_tokens"]; got != float64(6) {
				t.Fatalf("max_tokens = %#v, want 6", got)
			}
		})
	}
	runInput := hotPathStageRunInput(hotPathDispatchSnapshot{
		Stage:        config.ExecutionRouteStage{Options: map[string]any{"max_tokens": 999}},
		OutputBudget: hotPathOutputBudget{Limited: true, Remaining: 6},
	}, "continue")
	options, ok := runInput["options"].(map[string]any)
	if !ok || options["max_tokens"] != 6 {
		t.Fatalf("normalized options = %#v, want reserved max_tokens 6", runInput["options"])
	}
}

func TestHotPathOuterTurnBudgetProjectionAndPostTerminalStop(t *testing.T) {
	outer := newHotPathCallerCappedOuterTurn("turn-projection", 20)
	stage := normalizedStageOutput{
		ResponseID: "provider-stage", Created: 77, Content: "content", Reasoning: "reason",
		ToolCalls:      []normalizedToolCall{{ID: "provider-tool", ProviderCallID: "provider-tool", Name: "read_file", RawArgs: `{"path":"README.md"}`}},
		TerminalReason: "tool_calls", Usage: json.RawMessage(`{"prompt_tokens":3,"completion_tokens":4,"total_tokens":7,"provider_extra":true}`),
		OpenAIUsage: &openAIUsage{PromptTokens: 3, CompletionTokens: 4, TotalTokens: 7},
	}
	if err := runHotPathCollectedStage(context.Background(), outer, "stage-one", stage); err != nil {
		t.Fatal(err)
	}
	if err := outer.projectToolIdentities([]normalizedToolCall{{ID: "public-tool", ProviderCallID: "provider-tool", Name: "read_file"}}); err != nil {
		t.Fatal(err)
	}
	outer.commitTerminalSuccess(stage.TerminalReason)
	projected := hotPathCompatibilityOutput(outer, stage, "openai")
	if projected.ResponseID != "provider-stage" || projected.Created != 77 || projected.Content != "content" || projected.Reasoning != "reason" ||
		len(projected.ToolCalls) != 1 || projected.ToolCalls[0].ID != "public-tool" || projected.ToolCalls[0].ProviderCallID != "provider-tool" || projected.TerminalReason != "tool_calls" {
		t.Fatalf("compatibility projection = %+v", projected)
	}
	var usage map[string]any
	if err := json.Unmarshal(projected.Usage, &usage); err != nil {
		t.Fatal(err)
	}
	if usage["prompt_tokens"] != float64(3) || usage["completion_tokens"] != float64(4) || usage["provider_extra"] != true {
		t.Fatalf("projected usage = %#v", usage)
	}
	if outer.commitTerminalError("api_error", "late") {
		t.Fatal("post-terminal error won the terminal race")
	}
	late, err := streamgate.NewReleaseTextDeltaEvent(streamGateChannelDefault, "late", time.Now())
	if err != nil {
		t.Fatal(err)
	}
	if err := outer.releaseDelta(2, late); !errors.Is(err, errHotPathTurnTerminal) {
		t.Fatalf("post-terminal release err=%v, want terminal guard", err)
	}
}

func TestHotPathOuterTurnCapTerminalContinuity(t *testing.T) {
	t.Run("reported exhaustion preserves current tool terminal", func(t *testing.T) {
		outer := newHotPathCallerCappedOuterTurn("turn-cap-tool", 4)
		stage := normalizedStageOutput{
			ResponseID: "provider-cap-tool",
			ToolCalls: []normalizedToolCall{{
				ID: "provider-tool", ProviderCallID: "provider-tool", Name: "read",
				RawArgs: `{"p":"x"}`,
			}},
			TerminalReason: "tool_calls",
			OpenAIUsage:    &openAIUsage{CompletionTokens: 4, TotalTokens: 4},
		}
		if err := runHotPathCollectedStage(context.Background(), outer, "stage-tool", stage); err != nil {
			t.Fatal(err)
		}
		if budget := outer.outputBudget(); !budget.Exhausted || budget.Remaining != 0 {
			t.Fatalf("tool-stage budget = %+v, want exhausted", budget)
		}
		if err := outer.projectToolIdentities([]normalizedToolCall{{
			ID: "public-tool", ProviderCallID: "provider-tool", Name: "read",
		}}); err != nil {
			t.Fatal(err)
		}
		if !outer.commitTerminalSuccess(stage.TerminalReason) {
			t.Fatal("tool terminal did not commit")
		}
		visible := hotPathCompatibilityOutput(outer, stage, "openai")
		if visible.TerminalReason != "tool_calls" || len(visible.ToolCalls) != 1 || visible.ToolCalls[0].ID != "public-tool" {
			t.Fatalf("cap-at-tool output = %+v", visible)
		}
	})

	t.Run("content exhaustion remains length terminal", func(t *testing.T) {
		outer := newHotPathCallerCappedOuterTurn("turn-cap-content", 4)
		stage := normalizedStageOutput{
			ResponseID: "provider-cap-content", Content: "done", TerminalReason: "stop",
			OpenAIUsage: &openAIUsage{CompletionTokens: 4, TotalTokens: 4},
		}
		if err := runHotPathCollectedStage(context.Background(), outer, "stage-content", stage); err != nil {
			t.Fatal(err)
		}
		if budget := outer.outputBudget(); !budget.Exhausted {
			t.Fatalf("content-stage budget = %+v, want exhausted", budget)
		}
		outer.commitLengthTerminal()
		if visible := hotPathCompatibilityOutput(outer, stage, "openai"); visible.TerminalReason != "length" || len(visible.ToolCalls) != 0 {
			t.Fatalf("content cap output = %+v", visible)
		}
	})

	t.Run("usage-less unicode is preserved and blocks later provider dispatch", func(t *testing.T) {
		content, reasoning, name, args := "한", "글", "도구", `{"값":"✓"}`
		outer := newHotPathCallerCappedOuterTurn("turn-cap-unicode", 1)
		stage := normalizedStageOutput{
			ResponseID: "provider-cap-unicode", Content: content, Reasoning: reasoning,
			ToolCalls: []normalizedToolCall{{
				ID: "provider-unicode", ProviderCallID: "provider-unicode", Name: name, RawArgs: args,
			}},
			TerminalReason: "tool_calls",
		}
		if err := runHotPathCollectedStage(context.Background(), outer, "stage-unicode", stage); err != nil {
			t.Fatal(err)
		}
		if budget := outer.outputBudget(); budget.Exhausted || budget.Remaining != 1 || !budget.MissingUsage {
			t.Fatalf("usage-less Unicode budget = %+v, want preserved cap with missing-usage gate", budget)
		}
		outer.commitTerminalSuccess("tool_use")
		if visible := hotPathCompatibilityOutput(outer, stage, "anthropic"); visible.Content != content || visible.Reasoning != reasoning ||
			visible.TerminalReason != "tool_use" || len(visible.ToolCalls) != 1 || visible.ToolCalls[0].RawArgs != args {
			t.Fatalf("usage-less Unicode tool terminal = %+v", visible)
		}
	})
}

func (s *hotPathSequenceSource) NextEvent(context.Context) (streamgate.NormalizedEvent, error) {
	if s.index >= len(s.events) {
		return streamgate.NormalizedEvent{}, errors.New("hot path test source exhausted")
	}
	event := s.events[s.index]
	s.index++
	return event, nil
}

type hotPathContextSource struct{}

func (hotPathContextSource) NextEvent(ctx context.Context) (streamgate.NormalizedEvent, error) {
	return streamgate.NormalizedEvent{}, ctx.Err()
}

type hotPathFixedUsage struct{ usage hotPathStageUsage }

func (p hotPathFixedUsage) stageUsage() (hotPathStageUsage, bool) { return p.usage, p.usage.Reported }

type hotPathCountingController struct {
	mu     sync.Mutex
	aborts int
	closes int
}

func (c *hotPathCountingController) AbortAttempt(context.Context) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.aborts++
	return nil
}

func (c *hotPathCountingController) CloseAttempt(context.Context) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.closes++
	return nil
}

func (c *hotPathCountingController) counts() (aborts, closes int) {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.aborts, c.closes
}

func hotPathTestEvent(t *testing.T, build func() (streamgate.NormalizedEvent, error)) streamgate.NormalizedEvent {
	t.Helper()
	event, err := build()
	if err != nil {
		t.Fatalf("build normalized event: %v", err)
	}
	return event
}

func hotPathTestRelease(t *testing.T, build func() (streamgate.ReleaseEvent, error)) streamgate.ReleaseEvent {
	t.Helper()
	event, err := build()
	if err != nil {
		t.Fatalf("build release event: %v", err)
	}
	return event
}

func TestHotPathStageRuntime(t *testing.T) {
	now := time.Now()
	source := &hotPathSequenceSource{events: []streamgate.NormalizedEvent{
		hotPathTestEvent(t, func() (streamgate.NormalizedEvent, error) {
			return streamgate.NewResponseStartEvent(streamGateChannelDefault, 200, nil, now)
		}),
		hotPathTestEvent(t, func() (streamgate.NormalizedEvent, error) {
			return streamgate.NewTextDeltaEvent(streamGateChannelDefault, "released-before-terminal", now)
		}),
		hotPathTestEvent(t, func() (streamgate.NormalizedEvent, error) {
			return streamgate.NewTerminalEvent(streamGateChannelDefault, now)
		}),
	}}
	outer := newHotPathOuterTurn("turn-stage")
	usage := hotPathFixedUsage{usage: hotPathStageUsage{ResponseID: "provider-response", InputTokens: 3, OutputTokens: 5, Reported: true}}

	controller := &hotPathCountingController{}
	term, err := runHotPathStage(context.Background(), outer, hotPathStageMeta{StageID: "selector", Model: "selector-model", Provider: "provider-a", AttemptID: "attempt-a"}, source, usage, controller)
	if err != nil {
		t.Fatalf("run stage: %v", err)
	}
	if !term.Success || !term.HasUsage {
		t.Fatalf("terminal = %#v, want successful held terminal with usage", term)
	}
	if outer.isTerminalCommitted() {
		t.Fatal("stage terminal committed the public turn terminal")
	}
	released := outer.releasedDeltas()
	if len(released) != 1 || released[0].Text != "released-before-terminal" {
		t.Fatalf("released deltas = %#v, want progressive stage delta", released)
	}
	if usage := outer.turnUsage(); usage.InputTokens != 3 || usage.OutputTokens != 5 || !usage.Reported {
		t.Fatalf("turn usage = %#v", usage)
	}
	if !outer.commitTerminalSuccess("stop") || !outer.isTerminalCommitted() {
		t.Fatal("outer terminal was not independently committed")
	}
	if aborts, closes := controller.counts(); aborts != 0 || closes != 1 {
		t.Fatalf("controller calls = aborts:%d closes:%d, want graceful close once", aborts, closes)
	}
}

func TestHotPathStageTransportOwnership(t *testing.T) {
	now := time.Now()
	tests := []struct {
		name         string
		events       []streamgate.NormalizedEvent
		cancel       bool
		wantRunError bool
		wantSuccess  bool
		wantAborts   int
		wantCloses   int
	}{
		{
			name: "success closes gracefully once",
			events: []streamgate.NormalizedEvent{
				hotPathTestEvent(t, func() (streamgate.NormalizedEvent, error) {
					return streamgate.NewResponseStartEvent(streamGateChannelDefault, 200, nil, now)
				}),
				hotPathTestEvent(t, func() (streamgate.NormalizedEvent, error) {
					return streamgate.NewTerminalEvent(streamGateChannelDefault, now)
				}),
			},
			wantSuccess: true,
			wantCloses:  1,
		},
		{
			name: "provider error aborts once",
			events: []streamgate.NormalizedEvent{
				hotPathTestEvent(t, func() (streamgate.NormalizedEvent, error) {
					return streamgate.NewResponseStartEvent(streamGateChannelDefault, 200, nil, now)
				}),
				hotPathTestEvent(t, func() (streamgate.NormalizedEvent, error) {
					return newOpenAIProviderErrorEvent(streamGateErrorTunnelFailed)
				}),
			},
			wantAborts: 1,
		},
		{
			name:         "cancellation aborts once",
			cancel:       true,
			wantRunError: true,
			wantAborts:   1,
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			controller := &hotPathCountingController{}
			outer := newHotPathOuterTurn("turn-ownership")
			var source streamgate.NormalizedEventSource
			ctx := context.Background()
			if test.cancel {
				cancelCtx, cancel := context.WithCancel(context.Background())
				cancel()
				ctx = cancelCtx
				source = hotPathContextSource{}
			} else {
				source = &hotPathSequenceSource{events: test.events}
			}

			rt, sink, err := newHotPathStageRuntime(outer, hotPathStageMeta{StageID: "ownership", Model: "model", Provider: "provider", AttemptID: test.name}, source, nil, controller)
			if err != nil {
				t.Fatalf("new stage runtime: %v", err)
			}
			runErr := rt.Run(ctx)
			if (runErr != nil) != test.wantRunError {
				t.Fatalf("run error = %v, want error=%t", runErr, test.wantRunError)
			}
			term, committed := sink.stageTerminal()
			graceful := runErr == nil && committed && term.Success
			if err := rt.CloseRequestResources(context.Background(), graceful); err != nil {
				t.Fatalf("close request resources: %v", err)
			}
			if err := rt.CloseRequestResources(context.Background(), graceful); err != nil {
				t.Fatalf("duplicate close request resources: %v", err)
			}
			if committed && term.Success != test.wantSuccess {
				t.Fatalf("terminal = %#v, want success=%t", term, test.wantSuccess)
			}
			if aborts, closes := controller.counts(); aborts != test.wantAborts || closes != test.wantCloses {
				t.Fatalf("controller calls = aborts:%d closes:%d, want aborts:%d closes:%d", aborts, closes, test.wantAborts, test.wantCloses)
			}
		})
	}
}

func TestHotPathStageProtocolFragments(t *testing.T) {
	tests := []struct {
		name      string
		protocol  string
		frames    [][]byte
		wantText  string
		wantTool  string
		wantInput int
	}{
		{
			name:     "openai chat fragments",
			protocol: "openai",
			frames: [][]byte{
				[]byte("data: {\"id\":\"chat-stage\",\"choices\":[{\"delta\":{\"content\":\"hel"),
				[]byte("lo\",\"tool_calls\":[{\"index\":0,\"id\":\"call-a\",\"function\":{\"name\":\"write\",\"arguments\":\"{\\\"x\\\":\"}}]}}]}\n\n"),
				[]byte("data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"1}\"}}]}}],\"usage\":{\"prompt_tokens\":2,\"completion_tokens\":4}}\n\n"),
			},
			wantText:  "hello",
			wantTool:  "{\"x\":1}",
			wantInput: 2,
		},
		{
			name:     "anthropic messages fragments",
			protocol: "anthropic",
			frames: [][]byte{
				[]byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg-stage\",\"usage\":{\"input_tokens\":3}}}\n\n"),
				[]byte("event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"tool-a\",\"name\":\"write\"}}\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"x\\\":\"}}\n\n"),
				[]byte("event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"1}\"}}\n\nevent: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":1,\"content_block\":{\"type\":\"text\",\"text\":\"hello\"}}\n\nevent: message_delta\ndata: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":4}}\n\n"),
			},
			wantText:  "hello",
			wantTool:  "{\"x\":1}",
			wantInput: 3,
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			frames := make(chan *iop.ProviderTunnelFrame, len(test.frames)+2)
			frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START, StatusCode: 200}
			for _, body := range test.frames {
				frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY, Body: body}
			}
			frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END}
			close(frames)

			source := newHotPathTunnelStageSource(edgeservice.ProviderTunnelStream{Frames: frames}, time.Second, newHotPathStageDecoderForProtocol(test.protocol))
			outer := newHotPathOuterTurn("turn-" + test.protocol)
			term, err := runHotPathStage(context.Background(), outer, hotPathStageMeta{StageID: test.protocol, Protocol: test.protocol, Model: "model", Provider: "provider", AttemptID: "attempt"}, source, source, &hotPathCountingController{})
			if err != nil {
				t.Fatalf("run %s stage: %v", test.protocol, err)
			}
			if !term.Success || outer.isTerminalCommitted() {
				t.Fatalf("terminal = %#v, outer committed = %t", term, outer.isTerminalCommitted())
			}
			out := outer.accumulator()
			if out.Content != test.wantText || len(out.ToolCalls) != 1 || out.ToolCalls[0].RawArgs != test.wantTool {
				t.Fatalf("accumulator = %#v", out)
			}
			if out.OpenAIUsage == nil || out.OpenAIUsage.PromptTokens != test.wantInput || out.OpenAIUsage.CompletionTokens != 4 {
				t.Fatalf("usage = %#v", out.OpenAIUsage)
			}
		})
	}
}

func TestHotPathStageTunnelFraming(t *testing.T) {
	tests := []struct {
		name        string
		frames      []*iop.ProviderTunnelFrame
		wantSuccess bool
	}{
		{
			name: "explicit response start body and end succeeds",
			frames: []*iop.ProviderTunnelFrame{
				{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START, StatusCode: 200},
				{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY, Body: []byte("data: {\"id\":\"chatcmpl-framing\",\"choices\":[{\"delta\":{\"content\":\"ok\"}}]}\n\n")},
				{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END},
			},
			wantSuccess: true,
		},
		{
			name:   "body before response start fails closed",
			frames: []*iop.ProviderTunnelFrame{{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY, Body: []byte("data: ignored\n\n")}},
		},
		{
			name:   "end before response start fails closed",
			frames: []*iop.ProviderTunnelFrame{{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END}},
		},
		{
			name:   "channel close before explicit end fails closed",
			frames: []*iop.ProviderTunnelFrame{{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START, StatusCode: 200}},
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			frames := make(chan *iop.ProviderTunnelFrame, len(test.frames))
			for _, frame := range test.frames {
				frames <- frame
			}
			close(frames)
			source := newHotPathTunnelStageSource(edgeservice.ProviderTunnelStream{Frames: frames}, time.Second, newOpenAIChatStageDecoder())
			outer := newHotPathOuterTurn("turn-framing")
			term, err := runHotPathStage(context.Background(), outer, hotPathStageMeta{StageID: "framing", Model: "model", Provider: "provider", AttemptID: test.name}, source, source, &hotPathCountingController{})
			if err != nil {
				t.Fatalf("run stage: %v", err)
			}
			if term.Success != test.wantSuccess {
				t.Fatalf("terminal = %#v, want success=%t", term, test.wantSuccess)
			}
			if outer.isTerminalCommitted() {
				t.Fatal("stage framing committed a public terminal")
			}
		})
	}
}

func TestHotPathOuterTurnOrderingAndAggregation(t *testing.T) {
	now := time.Now()
	outer := newHotPathOuterTurn("turn-order")
	if err := outer.openResponse(streamgate.ResponseStart{}); err != nil {
		t.Fatalf("open response: %v", err)
	}
	first, second := outer.beginStage(), outer.beginStage()
	for _, item := range []struct {
		stage int
		event streamgate.ReleaseEvent
	}{
		{first, hotPathTestRelease(t, func() (streamgate.ReleaseEvent, error) {
			return streamgate.NewReleaseTextDeltaEvent(streamGateChannelDefault, "one", now)
		})},
		{first, hotPathTestRelease(t, func() (streamgate.ReleaseEvent, error) {
			return streamgate.NewReleaseToolCallFragmentEvent(streamGateChannelDefault, "duplicate", "write", "{", now)
		})},
		{second, hotPathTestRelease(t, func() (streamgate.ReleaseEvent, error) {
			return streamgate.NewReleaseReasoningDeltaEvent(streamGateChannelDefault, "think", now)
		})},
		{second, hotPathTestRelease(t, func() (streamgate.ReleaseEvent, error) {
			return streamgate.NewReleaseToolCallFragmentEvent(streamGateChannelDefault, "duplicate", "write", "}", now)
		})},
	} {
		if err := outer.releaseDelta(item.stage, item.event); err != nil {
			t.Fatalf("release delta: %v", err)
		}
	}
	outer.recordStageTerminal(hotPathStageTerminal{HasUsage: true, Usage: hotPathStageUsage{ResponseID: "shared", InputTokens: 2, OutputTokens: 3, Reported: true}})
	outer.recordStageTerminal(hotPathStageTerminal{HasUsage: true, Usage: hotPathStageUsage{ResponseID: "shared", InputTokens: 99, OutputTokens: 99, Reported: true}})
	outer.recordStageTerminal(hotPathStageTerminal{HasUsage: true, Usage: hotPathStageUsage{ResponseID: "other", InputTokens: 5, OutputTokens: 7, Reported: true}})
	if !outer.commitTerminalSuccess("") {
		t.Fatal("initial terminal must win")
	}
	out := outer.accumulator()
	if out.Content != "one" || out.Reasoning != "think" || out.TerminalReason != "tool_calls" {
		t.Fatalf("accumulator = %#v", out)
	}
	if len(out.ToolCalls) != 2 || out.ToolCalls[0].ID == out.ToolCalls[1].ID || out.ToolCalls[0].RawArgs != "{" || out.ToolCalls[1].RawArgs != "}" {
		t.Fatalf("remapped tools = %#v", out.ToolCalls)
	}
	if usage := outer.turnUsage(); usage.InputTokens != 7 || usage.OutputTokens != 10 {
		t.Fatalf("usage = %#v, want deduplicated aggregate", usage)
	}
}

func TestHotPathOuterTurnOutputCap(t *testing.T) {
	now := time.Now()
	outer := newHotPathCallerCappedOuterTurn("turn-cap", 7)
	stage := outer.beginStage()
	longUnicode := "한글과 UTF-8 payload length are unrelated to provider token usage"
	if err := outer.releaseDelta(stage, hotPathTestRelease(t, func() (streamgate.ReleaseEvent, error) {
		return streamgate.NewReleaseTextDeltaEvent(streamGateChannelDefault, longUnicode, now)
	})); err != nil {
		t.Fatalf("release delta: %v", err)
	}
	outer.recordStageTerminal(hotPathStageTerminal{Success: true, HasUsage: true, Usage: hotPathStageUsage{
		ResponseID: "provider-cap", OutputTokens: 2, Reported: true,
	}})
	if budget := outer.outputBudget(); budget.Exhausted || budget.Remaining != 5 || budget.MissingUsage {
		t.Fatalf("provider-token budget = %+v, want remaining 5", budget)
	}
	if !outer.commitTerminalSuccess("stop") {
		t.Fatal("provider terminal did not commit")
	}
	out := outer.accumulator()
	if out.Content != longUnicode || out.TerminalReason != "stop" {
		t.Fatalf("within-cap result = %#v", out)
	}
	if got := outer.releasedDeltas(); len(got) != 1 || got[0].Text != longUnicode {
		t.Fatalf("released deltas = %#v, want unmodified text", got)
	}
	if err := outer.releaseDelta(stage, hotPathTestRelease(t, func() (streamgate.ReleaseEvent, error) {
		return streamgate.NewReleaseReasoningDeltaEvent(streamGateChannelDefault, "after-terminal", now)
	})); !errors.Is(err, errHotPathTurnTerminal) {
		t.Fatalf("post-terminal release error = %v", err)
	}
}

func TestHotPathOuterTurnTerminalRace(t *testing.T) {
	outer := newHotPathOuterTurn("turn-race")
	const racers = 64
	var wg sync.WaitGroup
	results := make(chan bool, racers)
	for i := 0; i < racers; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			if i%2 == 0 {
				results <- outer.commitTerminalSuccess("stop")
				return
			}
			results <- outer.commitTerminalError("api_error", "race")
		}(i)
	}
	wg.Wait()
	close(results)
	wins := 0
	for won := range results {
		if won {
			wins++
		}
	}
	if wins != 1 || !outer.isTerminalCommitted() {
		t.Fatalf("terminal winners = %d, committed = %t", wins, outer.isTerminalCommitted())
	}
	event := hotPathTestRelease(t, func() (streamgate.ReleaseEvent, error) {
		return streamgate.NewReleaseTextDeltaEvent(streamGateChannelDefault, "late", time.Now())
	})
	if err := outer.releaseDelta(outer.beginStage(), event); !errors.Is(err, errHotPathTurnTerminal) {
		t.Fatalf("post-terminal release error = %v", err)
	}
}

func TestHotPathTerminalDispositionClosedSet(t *testing.T) {
	tests := []struct {
		name string
		act  func(*hotPathOuterTurn)
		want hotPathDispositionKind
	}{
		{name: "success", act: func(outer *hotPathOuterTurn) { outer.commitTerminalSuccess("stop") }, want: hotPathDispositionSuccess},
		{name: "tool turn", act: func(outer *hotPathOuterTurn) { outer.commitTerminalSuccess("tool_calls") }, want: hotPathDispositionToolTurn},
		{name: "length", act: func(outer *hotPathOuterTurn) { outer.commitLengthTerminal() }, want: hotPathDispositionLength},
		{name: "provider error", act: func(outer *hotPathOuterTurn) { outer.commitTerminalError("api_error", "upstream") }, want: hotPathDispositionProviderError},
		{name: "validation error", act: func(outer *hotPathOuterTurn) { outer.commitTerminalError("invalid_request_error", "validation") }, want: hotPathDispositionValidationError},
		{name: "timeout", act: func(outer *hotPathOuterTurn) {
			outer.selectDisposition(hotPathTerminalDisposition{Kind: hotPathDispositionTimeout, Source: "test", Cause: "deadline"})
		}, want: hotPathDispositionTimeout},
		{name: "caller cancel", act: func(outer *hotPathOuterTurn) {
			outer.cancelActiveStage(hotPathDispositionCallerCancel, "test", context.Canceled)
		}, want: hotPathDispositionCallerCancel},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			outer := newHotPathOuterTurn("turn-disposition")
			test.act(outer)
			disposition, ok := outer.terminalDisposition()
			if !ok || !disposition.valid() || disposition.Kind != test.want || disposition.Source == "" {
				t.Fatalf("disposition = %+v, present=%t, want %q", disposition, ok, test.want)
			}
			if outer.selectDisposition(hotPathTerminalDisposition{Kind: hotPathDispositionProviderError, Source: "duplicate"}) {
				t.Fatal("duplicate disposition replaced the winner")
			}
			preserved, _ := outer.terminalDisposition()
			if preserved != disposition {
				t.Fatalf("winner changed: before=%+v after=%+v", disposition, preserved)
			}
		})
	}
}

func TestHotPathActiveStageCancelTargetsCurrentGeneration(t *testing.T) {
	outer := newHotPathOuterTurn("turn-active-stage")
	firstController := &hotPathCountingController{}
	first, err := outer.registerActiveStage("local", firstController)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := outer.registerActiveStage("review", &hotPathCountingController{}); err == nil {
		t.Fatal("active stage was replaced before prior closure")
	}
	if err := first.CloseAttempt(context.Background()); err != nil {
		t.Fatal(err)
	}

	secondController := &hotPathCountingController{}
	second, err := outer.registerActiveStage("review", secondController)
	if err != nil {
		t.Fatal(err)
	}
	staleSink := &hotPathStageReleaseSink{outer: outer, active: first, tools: make(map[string]*hotPathProjectedTool)}
	stale, err := streamgate.NewReleaseTextDeltaEvent(streamGateChannelDefault, "stale", time.Now())
	if err != nil {
		t.Fatal(err)
	}
	if _, err := staleSink.Release(context.Background(), stale); err != nil {
		t.Fatalf("stale callback returned error: %v", err)
	}
	if got := outer.releasedDeltas(); len(got) != 0 {
		t.Fatalf("stale callback released output: %+v", got)
	}

	if !outer.cancelActiveStage(hotPathDispositionTimeout, "stage_timer", errRunTimedOut) {
		t.Fatal("timeout did not win terminal disposition")
	}
	if outer.cancelActiveStage(hotPathDispositionCallerCancel, "duplicate", context.Canceled) {
		t.Fatal("duplicate cancellation replaced timeout")
	}
	if err := second.AbortAttempt(context.Background()); err != nil {
		t.Fatal(err)
	}
	if err := first.AbortAttempt(context.Background()); err != nil {
		t.Fatal(err)
	}
	if aborts, closes := firstController.counts(); aborts != 0 || closes != 1 {
		t.Fatalf("prior stage calls = aborts:%d closes:%d, want close once", aborts, closes)
	}
	if aborts, closes := secondController.counts(); aborts != 1 || closes != 0 {
		t.Fatalf("active stage calls = aborts:%d closes:%d, want exact abort", aborts, closes)
	}
	disposition, ok := outer.terminalDisposition()
	if !ok || disposition.Kind != hotPathDispositionTimeout || disposition.StageID != "review" || disposition.Generation != second.generation {
		t.Fatalf("timeout ownership = %+v, present=%t", disposition, ok)
	}
}

func TestHotPathActiveStageCancelUsesExactCancelRunTarget(t *testing.T) {
	service := &fakeRunService{}
	outer := newHotPathOuterTurn("turn-exact-cancel")
	firstDispatch := edgeservice.RunDispatch{
		RunID: "run-local", NodeID: "node-local", Adapter: "adapter-local", Target: "target-local", SessionID: "session-local",
	}
	first, err := outer.registerActiveStage("local", newHotPathStageTransportController(service, firstDispatch, func() {}))
	if err != nil {
		t.Fatal(err)
	}
	if err := first.CloseAttempt(context.Background()); err != nil {
		t.Fatal(err)
	}

	secondDispatch := edgeservice.RunDispatch{
		RunID: "run-review", NodeID: "node-review", Adapter: "adapter-review", Target: "target-review", SessionID: "session-review",
	}
	if _, err := outer.registerActiveStage("review", newHotPathStageTransportController(service, secondDispatch, func() {})); err != nil {
		t.Fatal(err)
	}
	if !outer.cancelActiveStage(hotPathDispositionTimeout, "stage_timer", errRunTimedOut) {
		t.Fatal("timeout did not cancel the active review run")
	}
	outer.cancelActiveStage(hotPathDispositionTimeout, "duplicate", errRunTimedOut)
	calls := service.cancelCallsSnapshot()
	if len(calls) != 1 || calls[0] != (edgeservice.CancelRunRequest{
		NodeRef: secondDispatch.NodeID, RunID: secondDispatch.RunID,
	}) {
		t.Fatalf("CancelRun calls = %+v, want exact active review target once", calls)
	}
	if wire := edgeservice.BuildCancelRunRequest(calls[0]); wire.GetRunId() != secondDispatch.RunID {
		t.Fatalf("cancel wire = %+v, want run_id %q", wire, secondDispatch.RunID)
	}
}

func TestHotPathCancelCompleteRaceHasOneWinner(t *testing.T) {
	for iteration := 0; iteration < 128; iteration++ {
		outer := newHotPathOuterTurn("turn-cancel-complete")
		controller := &hotPathCountingController{}
		active, err := outer.registerActiveStage("review", controller)
		if err != nil {
			t.Fatal(err)
		}
		start := make(chan struct{})
		var wg sync.WaitGroup
		wg.Add(2)
		go func() {
			defer wg.Done()
			<-start
			_ = active.CloseAttempt(context.Background())
			outer.commitTerminalSuccess("stop")
		}()
		go func() {
			defer wg.Done()
			<-start
			outer.cancelActiveStage(hotPathDispositionCallerCancel, "caller_context", context.Canceled)
		}()
		close(start)
		wg.Wait()

		disposition, ok := outer.terminalDisposition()
		if !ok || (disposition.Kind != hotPathDispositionSuccess && disposition.Kind != hotPathDispositionCallerCancel) {
			t.Fatalf("iteration %d disposition = %+v, present=%t", iteration, disposition, ok)
		}
		aborts, closes := controller.counts()
		if aborts+closes != 1 {
			t.Fatalf("iteration %d transport actions = aborts:%d closes:%d, want exactly one", iteration, aborts, closes)
		}
		if outer.commitTerminalError("api_error", "late") {
			t.Fatalf("iteration %d accepted a second public terminal", iteration)
		}
	}
}

// rejectFixturedRun is a tiny fake run handle that records Close calls.
type rejectFixturedRun struct {
	dispatch edgeservice.RunDispatch
	closeMu  sync.Mutex
	closes   int
}

func (r *rejectFixturedRun) Dispatch() edgeservice.RunDispatch { return r.dispatch }
func (r *rejectFixturedRun) Close() {
	r.closeMu.Lock()
	defer r.closeMu.Unlock()
	r.closes++
}
func (r *rejectFixturedRun) Stream() edgeservice.RunStream { return edgeservice.RunStream{} }
func (r *rejectFixturedRun) WaitTimeout() time.Duration    { return 0 }
func (r *rejectFixturedRun) count() int {
	r.closeMu.Lock()
	defer r.closeMu.Unlock()
	return r.closes
}

// rejectFixturedTunnel is a tiny fake tunnel handle that records Close calls.
type rejectFixturedTunnel struct {
	dispatch edgeservice.RunDispatch
	closeMu  sync.Mutex
	closes   int
}

func (t *rejectFixturedTunnel) Dispatch() edgeservice.RunDispatch { return t.dispatch }
func (t *rejectFixturedTunnel) Close() {
	t.closeMu.Lock()
	defer t.closeMu.Unlock()
	t.closes++
}
func (t *rejectFixturedTunnel) Stream() edgeservice.ProviderTunnelStream {
	return edgeservice.ProviderTunnelStream{}
}
func (t *rejectFixturedTunnel) WaitTimeout() time.Duration   { return 0 }
func (t *rejectFixturedTunnel) SetHeaders(map[string]string) {}
func (t *rejectFixturedTunnel) count() int {
	t.closeMu.Lock()
	defer t.closeMu.Unlock()
	return t.closes
}

func assertExactRejectedDispatch(t *testing.T, calls []edgeservice.CancelRunRequest, dispatch edgeservice.RunDispatch) {
	t.Helper()
	if len(calls) != 1 {
		t.Fatalf("cancel calls=%d, want 1", len(calls))
	}
	want := edgeservice.CancelRunRequest{
		NodeRef: dispatch.NodeID, RunID: dispatch.RunID,
	}
	if calls[0] != want {
		t.Fatalf("cancel=%+v, want %+v", calls[0], want)
	}
	if wire := edgeservice.BuildCancelRunRequest(calls[0]); wire.GetRunId() != dispatch.RunID {
		t.Fatalf("cancel wire=%+v, want run_id %q", wire, dispatch.RunID)
	}
}

func assertRejectedHandleCloseCounts(t *testing.T, result *edgeservice.ProviderPoolDispatchResult) {
	t.Helper()
	if handle, ok := result.Run.(*rejectFixturedRun); ok && handle.count() != 1 {
		t.Fatalf("run close count=%d, want 1", handle.count())
	}
	if handle, ok := result.Tunnel.(*rejectFixturedTunnel); ok && handle.count() != 1 {
		t.Fatalf("tunnel close count=%d, want 1", handle.count())
	}
}

func TestHotPathRejectedDispatchExactOnceMatrix(t *testing.T) {
	for _, tc := range []struct {
		name    string
		path    string
		withRun bool
		withTun bool
	}{
		{name: "normalized", path: "normalized", withRun: true},
		{name: "tunnel", path: "provider_tunnel", withTun: true},
		{name: "malformed_both_handles", path: "normalized", withRun: true, withTun: true},
	} {
		t.Run(tc.name, func(t *testing.T) {
			dispatch := edgeservice.RunDispatch{RunID: "run-" + tc.name, NodeID: "node-" + tc.name, Adapter: "adapter", Target: "target", SessionID: "session"}
			result := &edgeservice.ProviderPoolDispatchResult{DispatchInfo: dispatch}
			if tc.path == "normalized" {
				result.Path = edgeservice.ProviderPoolPathNormalized
			} else {
				result.Path = edgeservice.ProviderPoolPathTunnel
			}
			if tc.withRun {
				result.Run = &rejectFixturedRun{dispatch: dispatch}
			}
			if tc.withTun {
				result.Tunnel = &rejectFixturedTunnel{dispatch: dispatch}
			}
			svc := &rejectPoolService{}
			srv := NewServer(config.EdgeOpenAIConf{Adapter: "test", Target: "t", TimeoutSec: 5}, svc, nil)
			owner := srv.newHotPathRejectedDispatchOwner(result)
			srv.abortHotPathRejectedDispatch(owner)
			srv.abortHotPathRejectedDispatch(owner)
			assertExactRejectedDispatch(t, svc.cancelSnapshot(), dispatch)
			assertRejectedHandleCloseCounts(t, result)
		})
	}
}

// rejectPoolService returns a scripted dispatch result whose validation fails
// because the RunID is empty.
type rejectPoolService struct {
	result      *edgeservice.ProviderPoolDispatchResult
	cancelCalls []edgeservice.CancelRunRequest
	closeMu     sync.Mutex
}

func (s *rejectPoolService) SubmitProviderPool(context.Context, edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
	return s.result, nil
}
func (s *rejectPoolService) SubmitRun(context.Context, edgeservice.SubmitRunRequest) (edgeservice.RunResult, error) {
	return nil, errors.New("not expected")
}
func (s *rejectPoolService) SubmitProviderTunnel(context.Context, edgeservice.SubmitProviderTunnelRequest) (edgeservice.ProviderTunnelResult, error) {
	return nil, errors.New("not expected")
}
func (s *rejectPoolService) OllamaAPI(context.Context, edgeservice.OllamaAPIRequest) (edgeservice.OllamaAPIView, error) {
	return edgeservice.OllamaAPIView{StatusCode: http.StatusOK}, nil
}
func (s *rejectPoolService) CancelRun(_ context.Context, req edgeservice.CancelRunRequest) (edgeservice.CommandResult, error) {
	s.closeMu.Lock()
	defer s.closeMu.Unlock()
	s.cancelCalls = append(s.cancelCalls, req)
	return edgeservice.CommandResult{NodeID: req.NodeRef}, nil
}
func (s *rejectPoolService) cancelSnapshot() []edgeservice.CancelRunRequest {
	s.closeMu.Lock()
	defer s.closeMu.Unlock()
	out := append([]edgeservice.CancelRunRequest(nil), s.cancelCalls...)
	return out
}

func TestHotPathRejectedDispatchSelectorMatrix(t *testing.T) {
	for _, tc := range []struct {
		name       string
		live       bool
		path       string
		withRun    bool
		withTunnel bool
	}{
		{name: "buffered_normalized_validation", path: "normalized", withRun: true},
		{name: "buffered_tunnel_validation", path: "provider_tunnel", withTunnel: true},
		{name: "live_normalized_validation", live: true, path: "normalized", withRun: true},
		{name: "live_tunnel_validation", live: true, path: "provider_tunnel", withTunnel: true},
		{name: "buffered_unsupported", path: "unknown", withRun: true},
		{name: "live_unsupported", live: true, path: "unknown", withTunnel: true},
		{name: "buffered_malformed_both_handles", path: "normalized", withRun: true, withTunnel: true},
		{name: "live_malformed_both_handles", live: true, path: "provider_tunnel", withRun: true, withTunnel: true},
	} {
		t.Run(tc.name, func(t *testing.T) {
			dispatch := edgeservice.RunDispatch{RunID: "run-" + tc.name, NodeID: "node-" + tc.name, Adapter: "adapter", Target: "target", SessionID: "session", ModelGroupKey: "group", ProviderID: "provider", ExecutionPath: string(tc.path)}
			result := &edgeservice.ProviderPoolDispatchResult{DispatchInfo: dispatch}
			switch tc.path {
			case "normalized":
				result.Path = edgeservice.ProviderPoolPathNormalized
			case "provider_tunnel":
				result.Path = edgeservice.ProviderPoolPathTunnel
			default:
				result.Path = "unknown"
			}
			mismatch := dispatch
			mismatch.ProviderID = "other-provider"
			if tc.withRun {
				result.Run = &rejectFixturedRun{dispatch: mismatch}
			}
			if tc.withTunnel {
				result.Tunnel = &rejectFixturedTunnel{dispatch: mismatch}
			}
			svc := &rejectPoolService{}
			srv := NewServer(config.EdgeOpenAIConf{Adapter: "test", Target: "t", TimeoutSec: 5}, svc, nil)
			var err error
			if tc.live {
				_, _, err = srv.runLivePresetSelectorResult(context.Background(), routeDispatch{}, "openai", "selector", result, newHotPathOuterTurn("selector"))
			} else {
				_, _, err = srv.collectPresetSelectorResult(context.Background(), routeDispatch{}, "openai", result)
			}
			if err == nil {
				t.Fatal("expected selector rejection")
			}
			assertExactRejectedDispatch(t, svc.cancelSnapshot(), dispatch)
			assertRejectedHandleCloseCounts(t, result)
		})
	}
}

func rejectedStageSnapshot(stream bool) hotPathDispatchSnapshot {
	paths := newReservedPaths("req-stage-reject")
	selector := hotPathStageCorrelation{StageID: "stg-s", ResponseID: "r:s/1", RunID: "run-s", ProviderID: "p", Terminal: "t"}
	return hotPathDispatchSnapshot{
		Protocol: "openai", Stream: stream, StageID: "stage-r", Stage: config.ExecutionRouteStage{Model: "m"},
		Input: buildLocalStageInput("immutable user task", paths, selector),
		Route: routeDispatch{NodeRef: "node-stage", ProviderID: "p", Adapter: "a-stage", Target: "t-stage", SessionID: "s-stage", TimeoutSec: 5, ProviderPool: true},
	}
}

func rejectedStageRequest() *http.Request {
	reqBody, _ := json.Marshal(map[string]any{"model": "m", "messages": []map[string]any{{"role": "user", "content": "hi"}}, "stream": false})
	return httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(reqBody))
}

func TestHotPathRejectedDispatchStageMatrix(t *testing.T) {
	for _, tc := range []struct {
		name       string
		stream     bool
		path       string
		withRun    bool
		withTunnel bool
		invalid    bool
	}{
		{name: "buffered_normalized_validation", path: "normalized", withRun: true, invalid: true},
		{name: "progressive_tunnel_validation", stream: true, path: "provider_tunnel", withTunnel: true, invalid: true},
		{name: "buffered_normalized_no_handle", path: "normalized"},
		{name: "progressive_normalized_no_handle", stream: true, path: "normalized"},
		{name: "buffered_normalized_opposite_handle", path: "normalized", withTunnel: true},
		{name: "progressive_normalized_opposite_handle", stream: true, path: "normalized", withTunnel: true},
		{name: "buffered_tunnel_no_handle", path: "provider_tunnel"},
		{name: "progressive_tunnel_no_handle", stream: true, path: "provider_tunnel"},
		{name: "buffered_tunnel_opposite_handle", path: "provider_tunnel", withRun: true},
		{name: "progressive_tunnel_opposite_handle", stream: true, path: "provider_tunnel", withRun: true},
		{name: "buffered_unsupported", path: "unknown", withRun: true},
		{name: "progressive_unsupported", stream: true, path: "unknown", withTunnel: true},
		{name: "buffered_malformed_both_handles", path: "normalized", withRun: true, withTunnel: true},
	} {
		t.Run(tc.name, func(t *testing.T) {
			dispatch := edgeservice.RunDispatch{RunID: "run-" + tc.name, NodeID: "node-stage", Adapter: "a-stage", Target: "t-stage", SessionID: "s-stage", ModelGroupKey: "m", ProviderID: "p", ExecutionPath: string(tc.path)}
			if tc.invalid {
				dispatch.ModelGroupKey = "wrong-model"
			}
			result := &edgeservice.ProviderPoolDispatchResult{DispatchInfo: dispatch}
			switch tc.path {
			case "normalized":
				result.Path = edgeservice.ProviderPoolPathNormalized
			case "provider_tunnel":
				result.Path = edgeservice.ProviderPoolPathTunnel
			default:
				result.Path = "unknown"
			}
			if tc.withRun {
				result.Run = &rejectFixturedRun{dispatch: dispatch}
			}
			if tc.withTunnel {
				result.Tunnel = &rejectFixturedTunnel{dispatch: dispatch}
			}
			svc := &rejectPoolService{result: result}
			srv := NewServer(config.EdgeOpenAIConf{Adapter: "test", Target: "t", TimeoutSec: 5}, svc, nil)
			_, _, err := srv.submitHotPathStage(context.Background(), rejectedStageRequest(), rejectedStageSnapshot(tc.stream), newHotPathOuterTurn("stage"))
			if err == nil {
				t.Fatal("expected stage rejection")
			}
			disposition, ok := hotPathDispositionFromError(err)
			if !ok || disposition.Kind != hotPathDispositionValidationError {
				t.Fatalf("disposition=%+v, typed=%t, want validation_error", disposition, ok)
			}
			assertExactRejectedDispatch(t, svc.cancelSnapshot(), dispatch)
			assertRejectedHandleCloseCounts(t, result)
		})
	}
}

// svcCancelSnapshot extracts cancel calls from a server's service when the
// service implements the cancel-snapshot accessor.
func svcCancelSnapshot(t *testing.T, srv *Server) []edgeservice.CancelRunRequest {
	t.Helper()
	if s, ok := srv.service.(*rejectPoolService); ok {
		return s.cancelSnapshot()
	}
	if s, ok := srv.service.(*fakeRunService); ok {
		return s.cancelCallsSnapshot()
	}
	t.Fatalf("unexpected service type %T", srv.service)
	return nil
}
