diff --git a/pkg/llm/service/handler.go b/pkg/llm/service/handler.go index ebd08b1a0e..6f921d8ad8 100644 --- a/pkg/llm/service/handler.go +++ b/pkg/llm/service/handler.go @@ -304,9 +304,21 @@ func handleDefaultMcpTools(ctx context.Context, w http.ResponseWriter, r *http.R appsrv.SendJSON(w, result) } +func registerLLMRouterAgentRoute(app *appsrv.Application) { + app.AddHandler2("POST", "/llm_router_agents//route", auth.Authenticate(handleLLMRouterAgentRoute), nil, "llm_router_agent_route", nil) +} + func handleLLMRouterAgentRoute(ctx context.Context, w http.ResponseWriter, r *http.Request) { + userCred := auth.FetchUserCredential(ctx, policy.FilterPolicyCredential) + if userCred == nil { + httperrors.UnauthorizedError(ctx, w, "Unauthorized") + return + } params, _, body := appsrv.FetchEnv(ctx, w, r) - id := params[""] + id := "" + if params != nil { + id = params[""] + } if id == "" { httperrors.MissingParameterError(ctx, w, "id") return @@ -325,12 +337,16 @@ func handleLLMRouterAgentRoute(ctx context.Context, w http.ResponseWriter, r *ht httperrors.InvalidInputError(ctx, w, "invalid input: %v", err) return } - obj, err := models.GetLLMRouterAgentManager().FetchByIdOrName(ctx, nil, id) + obj, err := models.GetLLMRouterAgentManager().FetchByIdOrName(ctx, userCred, id) if err != nil { httperrors.GeneralServerError(ctx, w, err) return } agent := obj.(*models.SLLMRouterAgent) + if err := db.IsObjectRbacAllowed(ctx, agent, userCred, policy.PolicyActionPerform, "route"); err != nil { + httperrors.GeneralServerError(ctx, w, err) + return + } out, err := agent.Route(ctx, input) if err != nil { httperrors.GeneralServerError(ctx, w, err) @@ -378,7 +394,7 @@ func InitHandlers(app *appsrv.Application, isSlave bool) { // 默认 MCP 服务器 tools:仅使用 options.MCPServerURL,不依赖 mcp_agent 条目 app.AddHandler2("GET", "/mcp_agents/default-mcp-tools", auth.Authenticate(handleDefaultMcpTools), nil, "default_mcp_tools", nil) - app.AddHandler2("POST", "/llm_router_agents//route", handleLLMRouterAgentRoute, nil, "llm_router_agent_route", nil) + registerLLMRouterAgentRoute(app) for _, manager := range []db.IModelManager{ taskman.TaskManager, diff --git a/pkg/llm/service/handler_route_test.go b/pkg/llm/service/handler_route_test.go new file mode 100644 index 0000000000..85a89180ce --- /dev/null +++ b/pkg/llm/service/handler_route_test.go @@ -0,0 +1,47 @@ +// Copyright 2019 Yunion +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package service + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "yunion.io/x/onecloud/pkg/appsrv" +) + +func TestLLMRouterAgentRouteRequiresAuth(t *testing.T) { + app := appsrv.NewApplication("test-llm-router-route", 1, 1, false) + registerLLMRouterAgentRoute(app) + + req := httptest.NewRequest("POST", "/llm_router_agents/some-id/route", strings.NewReader(`{"prompt":"hi"}`)) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + app.ServeHTTP(w, req) + if w.Code != http.StatusUnauthorized { + t.Fatalf("POST /llm_router_agents//route: status = %d, want %d", w.Code, http.StatusUnauthorized) + } +} + +func TestHandleLLMRouterAgentRouteUnauthorized(t *testing.T) { + req := httptest.NewRequest("POST", "/llm_router_agents/some-id/route", strings.NewReader(`{"prompt":"hi"}`)) + w := httptest.NewRecorder() + handleLLMRouterAgentRoute(context.Background(), w, req) + if w.Code != http.StatusUnauthorized { + t.Fatalf("no session: status = %d, want %d", w.Code, http.StatusUnauthorized) + } +}