From 25379767223c114e024d6c91a2f9311da88a3e72 Mon Sep 17 00:00:00 2001 From: "tao.chen" <93983997+taochen-ct@users.noreply.github.com> Date: Mon, 13 Jul 2026 15:15:21 +0800 Subject: [PATCH] =?UTF-8?q?rm:=20list=5Fapplications=20=E5=8A=A0=20limit/q?= =?UTF-8?q?ueue=20+=20=E6=97=A5=E5=BF=97=20fallback=20=E9=93=BE=E8=A1=A5?= =?UTF-8?q?=20NM=20=E7=9B=B4=E8=BF=9E?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit P0 (Python parity 缺口): - rm.Client.ListApps 新增 queue + limit 查询参数 - list_applications MCP tool schema 暴露 queue (string) 和 limit (number, default 100), handler 透传给 RM client - 限流原因: 生产 RM 上无 limit 会拉回 N MB JSON 撞 MaxResponseBytes - limit <= 0 不发送 query 参数, 跟 RM 默认行为一致; 老的调用方 (TestListApplications_EndToEnd 等) 不传参时行为不变 - 4 个旧 ListApps 测试调用点跟着更新, 加 3 个新测试: TestListApps_WithLimit, TestListApps_WithQueue, TestListApps_LimitZeroNoParam P1.3 (Python fallback 行为补全): - rm.Client.GetLogs 加第 4 步 fallback: 当前 3 步 (amContainerLogs/aggregated-logs/logs) 全失败时, 调 GetApp 解析 app.amContainerLogs 字段, GET 该 URL 直连 NodeManager - 新 source 名 'amContainerLogs-direct' 区分 RM-level endpoint - 提取 doBytesWithHosts(extraHosts...) 支持 per-call host 白名单, NM host 动态加进 allowedHosts, SSRF 保护不破 (URL 是 RM 响应里 回来的, 不是 LLM 任意填的) - 新增 TestGetLogs_FallbackToAMDirect (成功路径) 和 TestGetLogs_FallbackFailsWhenAMFieldMissing (amContainerLogs 字段为空时正常返回 error) P1.4 (零代码改动 + 文档化): - get_application_logs / get_application_status 工具 description 更新, 提示 LLM raw RM JSON 里包含 amContainerLogs 字段, 可在 aggregated-logs 全部失败时自己用 fetch_url 直连 NM - get_application_status 函数体不变 (本就是透传 raw JSON) 验证: - go build / go vet / go test 全部通过 - 5 个新测试全 PASS, 老的 TestListApplications_EndToEnd 仍 PASS - 单 ListApps 调用点 (list_applications.go:81) 编译过 未提交: spark-mcp-linux-amd64 (本地 build 产物, 当前 .gitignore 没拦, 需另行决定) Co-Authored-By: Claude --- internal/mcp/tools/get_application_logs.go | 2 +- internal/mcp/tools/get_application_status.go | 2 +- internal/mcp/tools/list_applications.go | 22 +++- internal/rm/client.go | 43 ++++++- internal/rm/client_test.go | 120 ++++++++++++++++++- 5 files changed, 176 insertions(+), 13 deletions(-) diff --git a/internal/mcp/tools/get_application_logs.go b/internal/mcp/tools/get_application_logs.go index 3e0fcf9..50f3e88 100644 --- a/internal/mcp/tools/get_application_logs.go +++ b/internal/mcp/tools/get_application_logs.go @@ -15,7 +15,7 @@ const defaultTailBytes = 1 << 20 // NewGetApplicationLogsTool returns the schema for the get_application_logs MCP Tool. func NewGetApplicationLogsTool() mcp.Tool { return mcp.NewTool(GetApplicationLogsName, - mcp.WithDescription("Fetch YARN application logs. Tries amContainerLogs first, then aggregated-logs, then the legacy /logs endpoint. Returns {source, text}."), + mcp.WithDescription("Fetch YARN application logs. Tries amContainerLogs (RM endpoint, often 307 to a NodeManager) first, then aggregated-logs, then the legacy /logs endpoint, and finally the AM container's direct NodeManager URL parsed from the app response. Returns {source, text}."), mcp.WithString("cluster_id", mcp.Required(), mcp.Description("ID of the configured cluster (from list_clusters)"), diff --git a/internal/mcp/tools/get_application_status.go b/internal/mcp/tools/get_application_status.go index a881d55..115763c 100644 --- a/internal/mcp/tools/get_application_status.go +++ b/internal/mcp/tools/get_application_status.go @@ -13,7 +13,7 @@ const GetApplicationStatusName = "get_application_status" // NewGetApplicationStatusTool returns the schema for the get_application_status MCP Tool. func NewGetApplicationStatusTool() mcp.Tool { return mcp.NewTool(GetApplicationStatusName, - mcp.WithDescription("Get the detailed status of a single YARN application from the ResourceManager. Returns the raw RM JSON for the app."), + mcp.WithDescription("Get the detailed status of a single YARN application from the ResourceManager. Returns the raw RM JSON for the app, including the amContainerLogs field (a direct NodeManager URL usable with fetch_url as a last-resort log source if all get_application_logs fallback endpoints fail)."), mcp.WithString("cluster_id", mcp.Required(), mcp.Description("ID of the configured cluster (from list_clusters)"), diff --git a/internal/mcp/tools/list_applications.go b/internal/mcp/tools/list_applications.go index 4374a15..df65cf9 100644 --- a/internal/mcp/tools/list_applications.go +++ b/internal/mcp/tools/list_applications.go @@ -24,6 +24,13 @@ func NewListApplicationsTool() mcp.Tool { mcp.WithString("user", mcp.Description("Filter by submitting user."), ), + mcp.WithString("queue", + mcp.Description("Filter by YARN queue name. Matches the YARN RM ?queue= parameter."), + ), + mcp.WithNumber("limit", + mcp.Description("Maximum number of applications to return. Defaults to 100 if not provided; set to 0 to request the RM default (no limit)."), + mcp.DefaultNumber(100), + ), ) } @@ -43,11 +50,24 @@ func (d *Deps) ListApplicationsHandler(ctx context.Context, req mcp.CallToolRequ if v, ok := args["user"].(string); ok { user = v } + queue := "" + if v, ok := args["queue"].(string); ok { + queue = v + } + limit := 0 + if v, ok := args["limit"].(float64); ok { + limit = int(v) + } + if limit < 0 { + limit = 0 + } callLog := startToolCall(ctx, d.Logger, ListApplicationsName, map[string]any{ "cluster_id": clusterID, "state": state, "user": user, + "queue": queue, + "limit": limit, }) defer callLog.End() @@ -58,7 +78,7 @@ func (d *Deps) ListApplicationsHandler(ctx context.Context, req mcp.CallToolRequ } rmc := rm.New(d.HTTPClient, cl) - raw, err := rmc.ListApps(ctx, state, user) + raw, err := rmc.ListApps(ctx, state, user, queue, limit) if err != nil { callLog.WithError(err) return errResult("list_applications: " + err.Error()), nil diff --git a/internal/rm/client.go b/internal/rm/client.go index a75468b..a2071bf 100644 --- a/internal/rm/client.go +++ b/internal/rm/client.go @@ -7,6 +7,7 @@ import ( "io" "net/http" "net/url" + "strconv" "strings" "spark-mcp-go/internal/cluster" @@ -24,8 +25,8 @@ func New(hc *httpclient.Client, c *cluster.Cluster) *Client { return &Client{hc: hc, cluster: c} } -// ListApps returns raw JSON from GET /ws/v1/cluster/apps[?state=&user=]. -func (r *Client) ListApps(ctx context.Context, state, user string) (json.RawMessage, error) { +// ListApps returns raw JSON from GET /ws/v1/cluster/apps[?state=&user=&queue=&limit=]. +func (r *Client) ListApps(ctx context.Context, state, user, queue string, limit int) (json.RawMessage, error) { base := r.cluster.RMURL + "/ws/v1/cluster/apps" u, err := url.Parse(base) if err != nil { @@ -38,6 +39,12 @@ func (r *Client) ListApps(ctx context.Context, state, user string) (json.RawMess if user != "" { q.Set("user", user) } + if queue != "" { + q.Set("queue", queue) + } + if limit > 0 { + q.Set("limit", strconv.Itoa(limit)) + } u.RawQuery = q.Encode() return r.do(ctx, http.MethodGet, u.String(), nil) } @@ -55,9 +62,10 @@ func (r *Client) KillApp(ctx context.Context, appID string) (json.RawMessage, er } // GetLogs walks a fallback chain to retrieve YARN application logs: -// 1. GET /ws/v1/cluster/apps/{id}/amContainerLogs (RM often 307 to a NodeManager) +// 1. GET /ws/v1/cluster/apps/{id}/amContainerLogs (RM endpoint, often 307 to a NodeManager) // 2. GET /ws/v1/cluster/apps/{id}/aggregated-logs?container=... // 3. GET /ws/v1/cluster/apps/{id}/logs +// 4. GET the NodeManager URL from the app's amContainerLogs field. // // The first successful endpoint wins. The returned source names the endpoint // that produced the body. @@ -69,12 +77,31 @@ func (r *Client) GetLogs(ctx context.Context, appID, container string) (body []b } var lastErr error for _, ep := range endpoints { - b, e := r.doBytes(ctx, http.MethodGet, r.cluster.RMURL+ep.path, nil) + b, e := r.doBytesWithHosts(ctx, http.MethodGet, r.cluster.RMURL+ep.path, nil) if e == nil { return b, ep.name, nil } lastErr = e } + // Try fetching the NodeManager's stdout URL directly. + app, err := r.GetApp(ctx, appID) + if err == nil { + var resp struct { + App struct { + AMContainerLogs string `json:"amContainerLogs"` + } `json:"app"` + } + if err := json.Unmarshal(app, &resp); err == nil && resp.App.AMContainerLogs != "" { + nmHost := extractHost(resp.App.AMContainerLogs) + b, e := r.doBytesWithHosts(ctx, http.MethodGet, resp.App.AMContainerLogs, nil, nmHost) + if e == nil { + return b, "amContainerLogs-direct", nil + } + lastErr = e + } + } else { + lastErr = err + } return nil, "", fmt.Errorf("rm: all log endpoints failed: %w", lastErr) } @@ -92,7 +119,7 @@ func (r *Client) do(ctx context.Context, method, rawURL string, body io.Reader) return raw, nil } -func (r *Client) doBytes(ctx context.Context, method, rawURL string, body io.Reader) ([]byte, error) { +func (r *Client) doBytesWithHosts(ctx context.Context, method, rawURL string, body io.Reader, extraHosts ...string) ([]byte, error) { req, err := http.NewRequestWithContext(ctx, method, rawURL, body) if err != nil { return nil, err @@ -100,7 +127,7 @@ func (r *Client) doBytes(ctx context.Context, method, rawURL string, body io.Rea if err := httpclient.ApplyAuth(req, "", r.cluster); err != nil { return nil, err } - resp, err := r.hc.DoWithRedirect(ctx, req, 5, r.allowedHosts()...) + resp, err := r.hc.DoWithRedirect(ctx, req, 5, append(r.allowedHosts(), extraHosts...)...) if err != nil { return nil, err } @@ -111,6 +138,10 @@ func (r *Client) doBytes(ctx context.Context, method, rawURL string, body io.Rea return io.ReadAll(resp.Body) } +func (r *Client) doBytes(ctx context.Context, method, rawURL string, body io.Reader) ([]byte, error) { + return r.doBytesWithHosts(ctx, method, rawURL, body) +} + func (r *Client) allowedHosts() []string { hosts := []string{extractHost(r.cluster.RMURL)} for _, pat := range r.cluster.URLAllowlist { diff --git a/internal/rm/client_test.go b/internal/rm/client_test.go index 2331d29..737d82a 100644 --- a/internal/rm/client_test.go +++ b/internal/rm/client_test.go @@ -3,6 +3,7 @@ package rm import ( "context" "encoding/base64" + "fmt" "io" "net/http" "net/http/httptest" @@ -30,7 +31,7 @@ func TestListApps(t *testing.T) { defer srv.Close() cl := &cluster.Cluster{RMURL: srv.URL} - raw, err := testClient(srv, cl).ListApps(context.Background(), "", "") + raw, err := testClient(srv, cl).ListApps(context.Background(), "", "", "", 0) if err != nil { t.Fatalf("ListApps: %v", err) } @@ -54,7 +55,7 @@ func TestListApps_WithQuery(t *testing.T) { defer srv.Close() cl := &cluster.Cluster{RMURL: srv.URL} - raw, err := testClient(srv, cl).ListApps(context.Background(), "RUNNING", "yarn") + raw, err := testClient(srv, cl).ListApps(context.Background(), "RUNNING", "yarn", "", 0) if err != nil { t.Fatalf("ListApps: %v", err) } @@ -212,7 +213,7 @@ func TestApplyAuth_SimpleUserName(t *testing.T) { defer srv.Close() cl := &cluster.Cluster{RMURL: srv.URL, AuthType: cluster.AuthSimple, AuthUsername: "alice"} - _, err := testClient(srv, cl).ListApps(context.Background(), "", "") + _, err := testClient(srv, cl).ListApps(context.Background(), "", "", "", 0) if err != nil { t.Fatalf("ListApps: %v", err) } @@ -227,7 +228,7 @@ func TestApplyAuth_BasicHeader(t *testing.T) { defer srv.Close() cl := &cluster.Cluster{RMURL: srv.URL, AuthType: cluster.AuthBasic, AuthUsername: "u", AuthPassword: "p"} - _, err := testClient(srv, cl).ListApps(context.Background(), "", "") + _, err := testClient(srv, cl).ListApps(context.Background(), "", "", "", 0) if err != nil { t.Fatalf("ListApps: %v", err) } @@ -236,3 +237,114 @@ func TestApplyAuth_BasicHeader(t *testing.T) { t.Errorf("Authorization=%q, want %q", captured, want) } } + +func TestListApps_WithLimit(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Query().Get("limit") != "50" { + t.Errorf("limit=%q, want 50", r.URL.Query().Get("limit")) + } + w.Write([]byte(`{"apps":{"app":[]}}`)) + })) + defer srv.Close() + cl := &cluster.Cluster{RMURL: srv.URL} + _, err := testClient(srv, cl).ListApps(context.Background(), "", "", "", 50) + if err != nil { + t.Fatalf("ListApps: %v", err) + } +} + +func TestListApps_WithQueue(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Query().Get("queue") != "root.users.alice" { + t.Errorf("queue=%q, want root.users.alice", r.URL.Query().Get("queue")) + } + w.Write([]byte(`{"apps":{"app":[]}}`)) + })) + defer srv.Close() + cl := &cluster.Cluster{RMURL: srv.URL} + _, err := testClient(srv, cl).ListApps(context.Background(), "", "", "root.users.alice", 0) + if err != nil { + t.Fatalf("ListApps: %v", err) + } +} + +func TestListApps_LimitZeroNoParam(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Query().Get("limit") != "" { + t.Errorf("limit should not be sent when 0; got %q", r.URL.Query().Get("limit")) + } + w.Write([]byte(`{"apps":{"app":[]}}`)) + })) + defer srv.Close() + cl := &cluster.Cluster{RMURL: srv.URL} + _, err := testClient(srv, cl).ListApps(context.Background(), "", "", "", 0) + if err != nil { + t.Fatalf("ListApps: %v", err) + } +} + +func TestGetLogs_FallbackToAMDirect(t *testing.T) { + // Mock NM that we want GetLogs to hit after standard endpoints fail. + nmSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte("nm direct logs")) + })) + defer nmSrv.Close() + nmHost := nmSrv.Listener.Addr().String() + + rmSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.URL.Path == "/ws/v1/cluster/apps/app_123/amContainerLogs": + w.WriteHeader(http.StatusInternalServerError) + case r.URL.Path == "/ws/v1/cluster/apps/app_123/aggregated-logs": + w.WriteHeader(http.StatusNotImplemented) + case r.URL.Path == "/ws/v1/cluster/apps/app_123/logs": + w.WriteHeader(http.StatusNotFound) + case r.URL.Path == "/ws/v1/cluster/apps/app_123" && r.Method == http.MethodGet: + // GetApp returns amContainerLogs pointing to the mock NM. + w.Write([]byte(fmt.Sprintf(`{"app":{"id":"app_123","amContainerLogs":"%s/node/containerlogs/container_1/root"}}`, nmSrv.URL))) + default: + t.Errorf("unexpected RM path %q", r.URL.Path) + } + })) + defer rmSrv.Close() + + cl := &cluster.Cluster{ + RMURL: rmSrv.URL, + URLAllowlist: []string{nmHost}, // allow NM host for SSRF + } + body, source, err := testClient(rmSrv, cl).GetLogs(context.Background(), "app_123", "container_1") + if err != nil { + t.Fatalf("GetLogs: %v", err) + } + if source != "amContainerLogs-direct" { + t.Errorf("source=%q, want amContainerLogs-direct", source) + } + if string(body) != "nm direct logs" { + t.Errorf("body=%q", string(body)) + } +} + +func TestGetLogs_FallbackFailsWhenAMFieldMissing(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.URL.Path == "/ws/v1/cluster/apps/app_123/amContainerLogs": + w.WriteHeader(http.StatusInternalServerError) + case r.URL.Path == "/ws/v1/cluster/apps/app_123/aggregated-logs": + w.WriteHeader(http.StatusNotImplemented) + case r.URL.Path == "/ws/v1/cluster/apps/app_123/logs": + w.WriteHeader(http.StatusNotFound) + case r.URL.Path == "/ws/v1/cluster/apps/app_123" && r.Method == http.MethodGet: + // amContainerLogs is empty + w.Write([]byte(`{"app":{"id":"app_123"}}`)) + default: + t.Errorf("unexpected path %q", r.URL.Path) + } + })) + defer srv.Close() + + cl := &cluster.Cluster{RMURL: srv.URL} + _, _, err := testClient(srv, cl).GetLogs(context.Background(), "app_123", "container_1") + if err == nil { + t.Fatal("expected error when all endpoints fail and amContainerLogs is empty") + } +}