1078 lines
43 KiB
Go
1078 lines
43 KiB
Go
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
|
|
}
|