fix(gateway): avoid repeated model probe during running status polling

This commit is contained in:
Alix-007 2026-03-30 12:21:17 +08:00
parent 93f4c4a843
commit 26d7ee2394
2 changed files with 144 additions and 2 deletions

View file

@ -145,6 +145,12 @@ func (h *Handler) TryAutoStartGateway() {
// gatewayStartReady validates whether current config can start the gateway.
func (h *Handler) gatewayStartReady() (bool, string, error) {
return h.gatewayStartReadyWithRuntimeProbe(true)
}
// gatewayStartReadyWithRuntimeProbe validates whether current config can start
// the gateway, with optional local runtime reachability probing.
func (h *Handler) gatewayStartReadyWithRuntimeProbe(allowRuntimeProbe bool) (bool, string, error) {
cfg, err := config.LoadConfig(h.configPath)
if err != nil {
return false, "", fmt.Errorf("failed to load config: %w", err)
@ -163,7 +169,7 @@ func (h *Handler) gatewayStartReady() (bool, string, error) {
if !hasModelConfiguration(modelCfg) {
return false, fmt.Sprintf("default model %q has no credentials configured", modelName), nil
}
if requiresRuntimeProbe(modelCfg) && !probeLocalModelAvailability(modelCfg) {
if allowRuntimeProbe && requiresRuntimeProbe(modelCfg) && !probeLocalModelAvailability(modelCfg) {
return false, fmt.Sprintf("default model %q is not reachable", modelName), nil
}
@ -878,7 +884,8 @@ func (h *Handler) gatewayStatusData() map[string]any {
gatewayStatus,
)
ready, reason, readyErr := h.gatewayStartReady()
// Avoid repeated local model probes in steady-state running status polling.
ready, reason, readyErr := h.gatewayStartReadyWithRuntimeProbe(gatewayStatus != "running")
if readyErr != nil {
data["gateway_start_allowed"] = false
data["gateway_start_reason"] = readyErr.Error()

View file

@ -385,6 +385,141 @@ func TestGatewayStatusIncludesStartConditionWhenNotReady(t *testing.T) {
}
}
func TestGatewayStatusRunningSkipsRuntimeProbeForStartCondition(t *testing.T) {
resetGatewayTestState(t)
resetModelProbeHooks(t)
configPath, cleanup := setupOAuthTestEnv(t)
defer cleanup()
cfg, err := config.LoadConfig(configPath)
if err != nil {
t.Fatalf("LoadConfig() error = %v", err)
}
cfg.ModelList = []*config.ModelConfig{{
ModelName: "local-vllm",
Model: "vllm/custom-model",
APIBase: "http://127.0.0.1:8000/v1",
}}
cfg.Agents.Defaults.ModelName = "local-vllm"
if err := config.SaveConfig(configPath, cfg); err != nil {
t.Fatalf("SaveConfig() error = %v", err)
}
h := NewHandler(configPath)
mux := http.NewServeMux()
h.RegisterRoutes(mux)
process, err := os.FindProcess(os.Getpid())
if err != nil {
t.Fatalf("FindProcess() error = %v", err)
}
gateway.mu.Lock()
gateway.cmd = &exec.Cmd{Process: process}
gateway.bootDefaultModel = "local-vllm"
setGatewayRuntimeStatusLocked("running")
gateway.mu.Unlock()
probeCalls := 0
probeOpenAICompatibleModelFunc = func(apiBase, modelID, apiKey string) bool {
probeCalls++
return false
}
gatewayHealthGet = func(string, time.Duration) (*http.Response, error) {
return mockGatewayHealthResponse(http.StatusOK, os.Getpid()), nil
}
rec := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/api/gateway/status", nil)
mux.ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("status = %d, want %d", rec.Code, http.StatusOK)
}
var body map[string]any
if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil {
t.Fatalf("unmarshal response: %v", err)
}
if got := body["gateway_status"]; got != "running" {
t.Fatalf("gateway_status = %#v, want %q", got, "running")
}
if got := body["gateway_start_allowed"]; got != true {
t.Fatalf("gateway_start_allowed = %#v, want true", got)
}
if probeCalls != 0 {
t.Fatalf("runtime probe calls = %d, want 0 while running status is healthy", probeCalls)
}
}
func TestGatewayStatusStoppedStillUsesRuntimeProbeForStartCondition(t *testing.T) {
resetGatewayTestState(t)
resetModelProbeHooks(t)
configPath, cleanup := setupOAuthTestEnv(t)
defer cleanup()
cfg, err := config.LoadConfig(configPath)
if err != nil {
t.Fatalf("LoadConfig() error = %v", err)
}
cfg.ModelList = []*config.ModelConfig{{
ModelName: "local-vllm",
Model: "vllm/custom-model",
APIBase: "http://127.0.0.1:8000/v1",
}}
cfg.Agents.Defaults.ModelName = "local-vllm"
if err := config.SaveConfig(configPath, cfg); err != nil {
t.Fatalf("SaveConfig() error = %v", err)
}
h := NewHandler(configPath)
mux := http.NewServeMux()
h.RegisterRoutes(mux)
gateway.mu.Lock()
gateway.cmd = nil
gateway.bootDefaultModel = ""
setGatewayRuntimeStatusLocked("stopped")
gateway.mu.Unlock()
probeCalls := 0
probeOpenAICompatibleModelFunc = func(apiBase, modelID, apiKey string) bool {
probeCalls++
return false
}
gatewayHealthGet = func(string, time.Duration) (*http.Response, error) {
return nil, errors.New("no gateway running")
}
rec := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/api/gateway/status", nil)
mux.ServeHTTP(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("status = %d, want %d", rec.Code, http.StatusOK)
}
var body map[string]any
if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil {
t.Fatalf("unmarshal response: %v", err)
}
if got := body["gateway_status"]; got != "stopped" {
t.Fatalf("gateway_status = %#v, want %q", got, "stopped")
}
if got := body["gateway_start_allowed"]; got != false {
t.Fatalf("gateway_start_allowed = %#v, want false", got)
}
if reason, _ := body["gateway_start_reason"].(string); !strings.Contains(reason, "not reachable") {
t.Fatalf("gateway_start_reason = %#v, want contains %q", body["gateway_start_reason"], "not reachable")
}
if probeCalls != 1 {
t.Fatalf("runtime probe calls = %d, want 1 while stopped status checks start-readiness", probeCalls)
}
}
func TestGatewayStatusKeepsRunningWhenHealthProbeFailsAfterRunning(t *testing.T) {
resetGatewayTestState(t)