heicode-mananger/heicode/controller/agent_swarm_query_test.go

package controller

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/heicode/manager/model"
	"github.com/stretchr/testify/require"
)

// #45 复审 #1:事件 payload 必须脱敏 —— 递归剔除 secret_ref / credentials / 大字段,
// 并对结果再跑 RedactText 兜底,绝不把 azkv:// secret_ref 或 sk-/Bearer 下发给客户端。
func TestSanitizeSwarmPayload_StripsSecrets(t *testing.T) {
	raw := `{
	  "task_id":"t1",
	  "approval":{"secret_ref":"azkv://heicode-kv.vault.azure.net/secrets/git-pat","note":"deploy"},
	  "credentials":{"access_key":"AKIA123","secret_access_key":"xxx"},
	  "headers":{"authorization":"Bearer aZ09tokenVALUE"},
	  "api_key":"sk-abcDEF1234567890",
	  "stack":["plain frame","key=sk-leak0987654321ABCD"],
	  "ok":true
	}`
	out := sanitizeSwarmPayload(raw)
	b, err := json.Marshal(out)
	require.NoError(t, err)
	s := string(b)

	// 敏感键被递归剔除
	require.NotContains(t, s, "secret_ref")
	require.NotContains(t, s, "azkv://")
	require.NotContains(t, s, "git-pat")
	require.NotContains(t, s, "credentials")
	require.NotContains(t, s, "AKIA123")
	require.NotContains(t, s, "api_key")
	// RedactText 兜底:残留在普通字段里的 sk-/Bearer 也被打码
	require.NotContains(t, s, "sk-leak0987654321ABCD")
	require.NotContains(t, s, "aZ09tokenVALUE")
	// 非敏感内容保留
	require.Contains(t, s, "t1")
	require.Contains(t, s, "ok")
}

func TestSanitizeSwarmPayload_EmptyAndPlain(t *testing.T) {
	require.Empty(t, sanitizeSwarmPayload(""))
	out := sanitizeSwarmPayload(`{"route":"chat","n":3}`)
	require.Equal(t, "chat", out["route"])
}

// #45/§4.1: 状态展示映射 —— blocked→degraded,其余直通(不臆造 preparing/verifying)。
func TestSwarmDisplayStatus(t *testing.T) {
	require.Equal(t, "degraded", swarmDisplayStatus("blocked"))
	require.Equal(t, "degraded", swarmDisplayStatus("BLOCKED"))
	require.Equal(t, "running", swarmDisplayStatus("running"))
	require.Equal(t, "completed", swarmDisplayStatus("completed"))
	require.Equal(t, "waiting_approval", swarmDisplayStatus("waiting_approval"))
}

// #15: artifact.created → 扁平视图 {uri,checksum,task_id,created_at}(size 缺省省略);绝不泄 secret_ref。
func TestSwarmArtifactView(t *testing.T) {
	e := model.AgentCallbackEvent{
		EventID: "evt-a1", Sequence: 5, TaskID: "t9", OccurredAt: "2026-06-10T00:00:00Z",
		PayloadJSON: `{"uri":"azblob://bucket/x.zip","checksum":"sha256:abc","task_id":"t9","secret_ref":"azkv://kv/secrets/s"}`,
	}
	v := swarmArtifactView(e)
	require.Equal(t, "azblob://bucket/x.zip", v["uri"])
	require.Equal(t, "sha256:abc", v["checksum"])
	require.Equal(t, "t9", v["task_id"])
	require.Equal(t, "2026-06-10T00:00:00Z", v["created_at"]) // payload 无 created_at → 回退 occurred_at
	_, hasSize := v["size_bytes"]
	require.False(t, hasSize, "size 未知不应出现")
	b, _ := json.Marshal(v)
	require.NotContains(t, string(b), "secret_ref")
	require.NotContains(t, string(b), "azkv://")

	// 带 size_bytes 时保留
	e2 := model.AgentCallbackEvent{EventID: "evt-a2", PayloadJSON: `{"uri":"u","checksum":"c","size_bytes":1234,"created_at":"2026-06-10T01:00:00Z"}`}
	v2 := swarmArtifactView(e2)
	require.EqualValues(t, 1234, v2["size_bytes"])
	require.Equal(t, "2026-06-10T01:00:00Z", v2["created_at"])
}

// #45/#15: SSE 终态判定 —— 仅 swarm.completed/failed/stopped 收流,其余继续。
func TestSwarmEventIsTerminal(t *testing.T) {
	for _, et := range []string{"swarm.completed", "swarm.failed", "swarm.stopped", "SWARM.STOPPED"} {
		require.True(t, swarmEventIsTerminal(et), "%s 应为终态", et)
	}
	for _, et := range []string{"swarm.started", "handoff.created", "approval.approved", "", "running"} {
		require.False(t, swarmEventIsTerminal(et), "%s 不应为终态", et)
	}
}

// #45: stop 真实接入 —— 命中冻结契约路径/鉴权/幂等头,并解析 {success,data} 信封。
func TestCallSwarmRuntimeStop(t *testing.T) {
	var gotPath, gotAuth, gotIdem, gotCorr string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotPath = r.URL.Path
		gotAuth = r.Header.Get("Authorization")
		gotIdem = r.Header.Get("X-Idempotency-Key")
		gotCorr = r.Header.Get("X-Correlation-ID")
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"success":true,"data":{"runtime_status":"stopping"}}`))
	}))
	defer srv.Close()

	cfg := agentRuntimeConfig{
		Enabled:  true,
		BaseURL:  srv.URL,
		Token:    "svc-token",
		StopPath: "/api/agent/swarm/deployments/{deployment_id}/stop",
		Timeout:  5 * time.Second,
	}
	dep := model.AgentDeployment{
		DeploymentID:        "dep-1",
		RuntimeDeploymentID: "rt-9",
		RuntimeSwarmID:      "sw-7",
		CorrelationID:       "cor-3",
	}
	res, err := callSwarmRuntimeStop(context.Background(), cfg, dep, "user requested")
	require.NoError(t, err)
	require.Equal(t, "stopping", res.RuntimeStatus)
	require.Equal(t, "/api/agent/swarm/deployments/rt-9/stop", gotPath) // runtime_deployment_id 优先
	require.Equal(t, "Bearer svc-token", gotAuth)
	require.Equal(t, "manager-stop-dep-1", gotIdem) // 幂等键按 manager deployment_id 派生
	require.Equal(t, "cor-3", gotCorr)
}

// runtime_deployment_id 缺失时回退 runtime_swarm_id(契约三映射);无任何 runtime id 则报错不发请求。
func TestCallSwarmRuntimeStop_RuntimeIDFallbackAndMissing(t *testing.T) {
	var gotPath string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotPath = r.URL.Path
		_, _ = w.Write([]byte(`{"success":true,"data":{}}`))
	}))
	defer srv.Close()
	cfg := agentRuntimeConfig{
		Enabled: true, BaseURL: srv.URL, Token: "t",
		StopPath: "/api/agent/swarm/deployments/{deployment_id}/stop", Timeout: 5 * time.Second,
	}
	_, err := callSwarmRuntimeStop(context.Background(), cfg, model.AgentDeployment{DeploymentID: "d", RuntimeSwarmID: "sw-only"}, "")
	require.NoError(t, err)
	require.Equal(t, "/api/agent/swarm/deployments/sw-only/stop", gotPath)

	_, err = callSwarmRuntimeStop(context.Background(), cfg, model.AgentDeployment{DeploymentID: "d"}, "")
	require.Error(t, err) // 无 runtime id → 不发请求
}

// 运行时非 2xx → 返回错误(handler 据此回 RUNTIME_UNAVAILABLE,不伪造受理)。
func TestCallSwarmRuntimeStop_RuntimeError(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusBadGateway)
		_, _ = w.Write([]byte(`upstream down`))
	}))
	defer srv.Close()
	cfg := agentRuntimeConfig{
		Enabled: true, BaseURL: srv.URL, Token: "t",
		StopPath: "/api/agent/swarm/deployments/{deployment_id}/stop", Timeout: 5 * time.Second,
	}
	_, err := callSwarmRuntimeStop(context.Background(), cfg, model.AgentDeployment{DeploymentID: "d", RuntimeDeploymentID: "rt"}, "")
	require.Error(t, err)
}

// #15: 6 类新事件已注册(类别 + 必填字段两张表)。
func TestSwarmFrozenEventTypesRegistered(t *testing.T) {
	for _, et := range []string{
		"swarm.completed", "swarm.failed", "swarm.stopped",
		"approval.approved", "approval.rejected", "handoff.created",
	} {
		_, inCat := agentCallbackEventCategories[et]
		_, inReq := agentCallbackEventRequiredFields[et]
		require.True(t, inCat, "event_type %s 应在 categories 注册", et)
		require.True(t, inReq, "event_type %s 应在 requiredFields 注册", et)
	}
}

// #28(@Mem0ried 消费需求):goal_summary 从 plan_json 顶层 objective 提取 —— 折叠空白为单行、
// 截断、脱敏;无 plan / 无 objective / 坏 JSON 返回空串(不臆造)。
func TestSwarmGoalSummary(t *testing.T) {
	require.Empty(t, swarmGoalSummary(""))
	require.Empty(t, swarmGoalSummary("not-json"))
	require.Empty(t, swarmGoalSummary(`{"sub_mode":"swarm"}`)) // 无 objective

	// 多行/多空格折叠成单行(JSON 里的 \n 是合法转义,用原始串避免再转义)
	got := swarmGoalSummary(`{"objective":"  迁移   支付服务\n  到 K8s  "}`)
	require.Equal(t, "迁移 支付服务 到 K8s", got)

	// 误入的 sk- 被 RedactText 兜底
	red := swarmGoalSummary(`{"objective":"deploy with key sk-abcdEFGH1234567890XYZ now"}`)
	require.NotContains(t, red, "sk-abcdEFGH1234567890XYZ")

	// 超长截断带省略号
	long := `{"objective":"` + strings.Repeat("x", 400) + `"}`
	out := swarmGoalSummary(long)
	require.LessOrEqual(t, len([]rune(out)), 201) // 200 + …
	require.True(t, strings.HasSuffix(out, "…"))
}

// #66 里程碑07:resource_grants 脱敏视图 —— 绝不下发 secret_ref 值,只标 has_secret。
func TestSwarmResourceGrantsView(t *testing.T) {
	plan := `{"resource_grants":[
	  {"grant_id":"g1","resource_id":"r1","resource_type":"git","binding_scope":"git:repo","target_role":"impl","permission_scope":["read","write"],"status":"active","secret_ref":"azkv://heicode-kv.vault.azure.net/secrets/git-pat"},
	  {"grant_id":"g2","resource_id":"r2","resource_type":"sk","status":"active"}
	]}`
	out := swarmResourceGrantsView(plan)
	require.Len(t, out, 2)
	require.Equal(t, "git", out[0]["resource_type"])
	require.Equal(t, true, out[0]["has_secret"])
	require.Equal(t, false, out[1]["has_secret"])
	b, _ := json.Marshal(out)
	require.NotContains(t, string(b), "secret_ref")
	require.NotContains(t, string(b), "azkv://")
	require.NotContains(t, string(b), "git-pat")
	require.Empty(t, swarmResourceGrantsView(""))
	require.Empty(t, swarmResourceGrantsView("bad-json"))
}

// #66 里程碑08:budget 上限;used_* 本批不返。
func TestSwarmBudgetView(t *testing.T) {
	b := swarmBudgetView(`{"budget":{"max_tokens":100000,"max_cost_usd":5.5,"max_duration_sec":3600}}`)
	require.EqualValues(t, 100000, b["max_tokens"])
	require.EqualValues(t, 5.5, b["max_cost_usd"])
	require.EqualValues(t, 3600, b["max_duration_sec"])
	_, hasUsed := b["used_model_cost"]
	require.False(t, hasUsed, "used_* 本批不应臆造")
	empty := swarmBudgetView("")
	require.EqualValues(t, 0, empty["max_tokens"])
}