diff --git a/sdk/api/handlers/handlers_context.go b/sdk/api/handlers/handlers_context.go index 7926be0f3..6362218fc 100644 --- a/sdk/api/handlers/handlers_context.go +++ b/sdk/api/handlers/handlers_context.go @@ -54,8 +54,11 @@ func (h *BaseAPIHandler) PrepareStreamModelRoute(ctx context.Context, handlerTyp return ctx, hasOverride } -func preparedModelRouteFromContext(ctx context.Context) (modelRouteDecision, bool) { - if ctx == nil { +func preparedModelRouteFromContext(ctx context.Context, skipRouterPluginID string) (modelRouteDecision, bool) { + // A host.model.execute_stream callback is a nested execution. Its caller is + // excluded from model routing, so an outer prepared route cannot be reused: + // it may point straight back at that caller. + if ctx == nil || strings.TrimSpace(skipRouterPluginID) != "" { return modelRouteDecision{}, false } decision, ok := ctx.Value(preparedModelRouteContextKey{}).(modelRouteDecision) diff --git a/sdk/api/handlers/handlers_model_router_test.go b/sdk/api/handlers/handlers_model_router_test.go index ad012c6e6..76bb4ddd3 100644 --- a/sdk/api/handlers/handlers_model_router_test.go +++ b/sdk/api/handlers/handlers_model_router_test.go @@ -93,6 +93,20 @@ type handlerDirectExecutorRouteHost struct { stream func(context.Context, string, coreexecutor.Request, coreexecutor.Options) (*coreexecutor.StreamResult, error) } +type handlerSkipAwareDirectExecutorRouteHost struct { + handlerDirectExecutorRouteHost + routeSkip string +} + +func (h *handlerSkipAwareDirectExecutorRouteHost) RouteModelExcept(ctx context.Context, req pluginapi.ModelRouteRequest, skipPluginID string) (pluginapi.ModelRouteResponse, bool) { + h.routeSkip = skipPluginID + return pluginapi.ModelRouteResponse{}, false +} + +func (h *handlerSkipAwareDirectExecutorRouteHost) HasModelRoutersExcept(string) bool { + return h != nil && h.hasRouters +} + func (h *handlerDirectExecutorRouteHost) ExecutePluginExecutor(ctx context.Context, pluginID string, req coreexecutor.Request, opts coreexecutor.Options) (coreexecutor.Response, error) { h.lastPluginID = pluginID h.lastRequest = req @@ -496,6 +510,42 @@ func TestPrepareStreamModelRouteReusesDecisionDuringExecution(t *testing.T) { } } +func TestExecuteModelStreamDoesNotReusePreparedRouteWhenRouterPluginSkipped(t *testing.T) { + const originalModel = "prepared-router-model" + const mappedModel = "mapped-upstream-model" + const originPluginID = "origin-plugin" + host := &handlerSkipAwareDirectExecutorRouteHost{} + host.hasRouters = true + host.route = func(context.Context, pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) { + return pluginapi.ModelRouteResponse{Handled: true, TargetKind: pluginapi.ModelRouteTargetExecutor, Target: originPluginID}, true + } + handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, nil) + handler.SetModelRouterHost(host) + body := []byte(`{"model":"prepared-router-model","stream":true}`) + ctx, routedToPlugin := handler.PrepareStreamModelRoute(context.Background(), "openai-response", originalModel, body) + if !routedToPlugin { + t.Fatal("PrepareStreamModelRoute() did not detect plugin executor route") + } + + _, errMsg := handler.ExecuteModelStream(ctx, ModelExecutionRequest{ + EntryProtocol: "openai-response", + ExitProtocol: "openai-response", + Model: mappedModel, + Stream: true, + Body: []byte(`{"model":"mapped-upstream-model","stream":true}`), + SkipRouterPluginID: originPluginID, + }) + if host.routeSkip != originPluginID { + t.Fatalf("router skip id = %q, want %q", host.routeSkip, originPluginID) + } + if host.lastPluginID == originPluginID { + t.Fatalf("plugin executor %q was re-entered despite SkipRouterPluginID", host.lastPluginID) + } + if errMsg == nil { + t.Fatal("ExecuteModelStream() error = nil, want normal provider resolution failure with empty auth manager") + } +} + func TestExecuteModelPropagatesRouterSkipPluginID(t *testing.T) { model := "model-execution-router-skip-model" requestBody := []byte(fmt.Sprintf(`{"model":%q}`, model)) diff --git a/sdk/api/handlers/handlers_stream.go b/sdk/api/handlers/handlers_stream.go index 4daa2e987..6669e0833 100644 --- a/sdk/api/handlers/handlers_stream.go +++ b/sdk/api/handlers/handlers_stream.go @@ -227,7 +227,7 @@ func (h *BaseAPIHandler) executeStreamWithAuthManager(ctx context.Context, handl func (h *BaseAPIHandler) executeStreamWithAuthManagerFormats(ctx context.Context, entryProtocol, exitProtocol, modelName string, rawJSON []byte, alt string, allowImageModel bool, execOptions modelExecutionOptions) (<-chan []byte, http.Header, <-chan *interfaces.ErrorMessage) { originalRequestedModel := modelName - routeDecision, preparedRoute := preparedModelRouteFromContext(ctx) + routeDecision, preparedRoute := preparedModelRouteFromContext(ctx, execOptions.SkipRouterPluginID) if !preparedRoute { routeDecision = h.applyModelRouter(ctx, entryProtocol, modelName, rawJSON, true, execOptions) }