heicode-mananger/heicode/controller/agent_approval_test.go

package controller

import (
	"bytes"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/gin-gonic/gin"
	"github.com/glebarez/sqlite"
	"github.com/heicode/manager/common"
	"github.com/heicode/manager/model"
	"github.com/stretchr/testify/require"
	"gorm.io/gorm"
)

func setupAgentApprovalTestDB(t *testing.T) *gorm.DB {
	t.Helper()
	gin.SetMode(gin.TestMode)
	common.UsingSQLite = true
	common.UsingMySQL = false
	common.UsingPostgreSQL = false
	common.RedisEnabled = false

	dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_"))
	db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{})
	require.NoError(t, err)
	model.DB = db
	model.LOG_DB = db
	require.NoError(t, db.AutoMigrate(
		&model.AgentApprovalRequest{},
		&model.AgentCredentialLease{},
		&model.AgentAuditEvent{},
		&model.AgentDeployment{},
	))
	t.Cleanup(func() {
		sqlDB, err := db.DB()
		if err == nil {
			_ = sqlDB.Close()
		}
	})
	return db
}

func TestAgentApprovalApproveNotifiesRuntimeDecision(t *testing.T) {
	db := setupAgentApprovalTestDB(t)
	var runtimeBody string
	var runtimeAuth string
	var runtimePath string
	runtimeServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		runtimePath = r.URL.Path
		runtimeAuth = r.Header.Get("Authorization")
		body, _ := io.ReadAll(r.Body)
		runtimeBody = string(body)
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"success":true,"data":{"status":"accepted"}}`))
	}))
	defer runtimeServer.Close()
	t.Setenv("AGENT_RUNTIME_ENABLED", "true")
	t.Setenv("AGENT_RUNTIME_BASE_URL", runtimeServer.URL)
	t.Setenv("AGENT_RUNTIME_SERVICE_TOKEN", "runtime-token")
	t.Setenv("AGENT_RUNTIME_APPROVAL_DECISION_PATH", "/api/swarms/{swarm_id}/approvals/{approval_id}")

	require.NoError(t, db.Create(&model.AgentDeployment{
		DeploymentID:        "dep_runtime_approval",
		UserID:              "7",
		BindingScope:        "project-alpha",
		CorrelationID:       "corr-runtime-approval",
		RuntimeDeploymentID: "runtime-dep-1",
		RuntimeSwarmID:      "swarm-runtime-1",
		Status:              "running",
		RuntimeState:        "running",
		PlanJSON:            `{"user_context":{"user_id":"7"},"metadata":{"correlation_id":"corr-runtime-approval"}}`,
	}).Error)

	body := `{
		"deployment_id":"dep_runtime_approval",
		"binding_scope":"project-alpha",
		"operation":"write_repository",
		"resource_id":"repo-main",
		"resource_type":"git",
		"resource_scope":"https://example.invalid/acme/repo#main",
		"target_role":"backend",
		"risk_level":"high",
		"requires_credential":true,
		"secret_ref":"azkv://heicode-kv.vault.azure.net/secrets/repo-main",
		"ttl_seconds":600
	}`

	_, approvalID := createAgentApprovalForTest(t, 7, body)
	w := performAgentApprovalRequest(ApproveAgentApprovalRequest, 7, http.MethodPost, "/approvals/:approval_id/approve", "/approvals/"+approvalID+"/approve", `{"reason":"允许写入"}`)
	require.Equal(t, http.StatusOK, w.Code)
	require.Equal(t, "/api/swarms/swarm-runtime-1/approvals/"+approvalID, runtimePath)
	require.Equal(t, "Bearer runtime-token", runtimeAuth)
	require.Contains(t, runtimeBody, `"decision":"approved"`)
	require.Contains(t, runtimeBody, `"approval_id":"`+approvalID+`"`)
	require.Contains(t, runtimeBody, `"manager_deployment_id":"dep_runtime_approval"`)
	require.Contains(t, runtimeBody, `"runtime_deployment_id":"runtime-dep-1"`)
	require.Contains(t, runtimeBody, `"swarm_id":"swarm-runtime-1"`)
	require.Contains(t, runtimeBody, `"credential_ref":"lease://agent/`)
	require.NotContains(t, runtimeBody, "azkv://")

	var auditRows []model.AgentAuditEvent
	require.NoError(t, db.Where("deployment_id = ?", "dep_runtime_approval").Order("id asc").Find(&auditRows).Error)
	require.NotEmpty(t, auditRows)
	require.Equal(t, "runtime.approval_decision.accepted", auditRows[len(auditRows)-1].Event)
}

func performAgentApprovalRequest(handler gin.HandlerFunc, userID int, method string, routePath string, requestPath string, body string) *httptest.ResponseRecorder {
	r := gin.New()
	r.Handle(method, routePath, func(c *gin.Context) {
		c.Set("id", userID)
		handler(c)
	})
	req := httptest.NewRequest(method, requestPath, bytes.NewBufferString(body))
	req.Header.Set("Content-Type", "application/json")
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)
	return w
}

func decodeAgentApprovalEnvelope(t *testing.T, w *httptest.ResponseRecorder) map[string]any {
	t.Helper()
	var envelope map[string]any
	require.NoError(t, common.Unmarshal(w.Body.Bytes(), &envelope))
	return envelope
}

func createAgentApprovalForTest(t *testing.T, userID int, body string) (map[string]any, string) {
	t.Helper()
	w := performAgentApprovalRequest(CreateAgentApprovalRequest, userID, http.MethodPost, "/approvals", "/approvals", body)
	require.Equal(t, http.StatusOK, w.Code)
	envelope := decodeAgentApprovalEnvelope(t, w)
	require.Equal(t, true, envelope["success"])
	data := envelope["data"].(map[string]any)
	approvalID := data["approval_id"].(string)
	require.NotEmpty(t, approvalID)
	return data, approvalID
}

func TestAgentApprovalApproveCreatesShortLivedLeaseAndAudit(t *testing.T) {
	db := setupAgentApprovalTestDB(t)
	body := `{
		"deployment_id":"dep_approval_1",
		"binding_scope":"project-alpha",
		"operation":"write_repository",
		"resource_id":"repo-main",
		"resource_type":"git",
		"resource_scope":"https://example.invalid/acme/repo#main",
		"target_role":"backend",
		"risk_level":"high",
		"requires_credential":true,
		"secret_ref":"azkv://heicode-kv.vault.azure.net/secrets/repo-main",
		"ttl_seconds":600,
		"reason":"需要写入功能分支"
	}`

	_, approvalID := createAgentApprovalForTest(t, 7, body)
	w := performAgentApprovalRequest(ApproveAgentApprovalRequest, 7, http.MethodPost, "/approvals/:approval_id/approve", "/approvals/"+approvalID+"/approve", `{"reason":"允许本次任务"}`)
	require.Equal(t, http.StatusOK, w.Code)
	envelope := decodeAgentApprovalEnvelope(t, w)
	require.Equal(t, true, envelope["success"])
	data := envelope["data"].(map[string]any)
	require.Equal(t, "approved", data["status"])
	lease := data["credential_lease"].(map[string]any)
	require.Equal(t, "active", lease["status"])
	require.Contains(t, lease["credential_ref"], "lease://agent/")
	require.NotContains(t, w.Body.String(), "do-not-store")
	require.NotContains(t, w.Body.String(), "azkv://", "API responses must not expose internal Key Vault secret_ref")

	var storedLease model.AgentCredentialLease
	require.NoError(t, db.Where("approval_id = ?", approvalID).First(&storedLease).Error)
	require.Equal(t, "azkv://heicode-kv.vault.azure.net/secrets/repo-main", storedLease.SecretRef)
	require.Equal(t, "active", storedLease.Status)
	require.Greater(t, storedLease.ExpiresAt, storedLease.CreatedAt)
	require.NotContains(t, storedLease.CredentialRef, "azkv://")

	var auditRows []model.AgentAuditEvent
	require.NoError(t, db.Order("id asc").Find(&auditRows).Error)
	require.Len(t, auditRows, 3)
	require.Equal(t, "approval.requested", auditRows[0].Event)
	require.Equal(t, "approval.approved", auditRows[1].Event)
	require.Equal(t, "credential_lease.created", auditRows[2].Event)
	require.NotContains(t, auditRows[2].DetailsJSON, "repo-main-secret-value")
}

func TestAgentApprovalRejectDoesNotCreateLease(t *testing.T) {
	db := setupAgentApprovalTestDB(t)
	body := `{
		"deployment_id":"dep_approval_2",
		"binding_scope":"project-alpha",
		"operation":"delete_resource",
		"resource_id":"vm-prod",
		"resource_type":"cloud_resource",
		"resource_scope":"/subscriptions/sub/resourceGroups/rg/providers/Microsoft.Compute/virtualMachines/prod",
		"target_role":"ops",
		"risk_level":"critical",
		"requires_credential":true,
		"secret_ref":"azkv://heicode-kv.vault.azure.net/secrets/azure-sp",
		"ttl_seconds":300
	}`

	_, approvalID := createAgentApprovalForTest(t, 8, body)
	w := performAgentApprovalRequest(RejectAgentApprovalRequest, 8, http.MethodPost, "/approvals/:approval_id/reject", "/approvals/"+approvalID+"/reject", `{"reason":"风险过高"}`)
	require.Equal(t, http.StatusOK, w.Code)
	envelope := decodeAgentApprovalEnvelope(t, w)
	require.Equal(t, true, envelope["success"])
	data := envelope["data"].(map[string]any)
	require.Equal(t, "rejected", data["status"])

	var leases int64
	require.NoError(t, db.Model(&model.AgentCredentialLease{}).Count(&leases).Error)
	require.Equal(t, int64(0), leases)
}

func TestAgentApprovalExpiredRequestCannotBeApproved(t *testing.T) {
	db := setupAgentApprovalTestDB(t)
	body := `{
		"deployment_id":"dep_approval_3",
		"binding_scope":"project-alpha",
		"operation":"write_repository",
		"resource_id":"repo-main",
		"resource_type":"git",
		"resource_scope":"https://example.invalid/acme/repo#main",
		"target_role":"backend",
		"risk_level":"high",
		"requires_credential":true,
		"secret_ref":"azkv://heicode-kv.vault.azure.net/secrets/repo-main",
		"ttl_seconds":1
	}`

	_, approvalID := createAgentApprovalForTest(t, 9, body)
	require.NoError(t, db.Model(&model.AgentApprovalRequest{}).Where("approval_id = ?", approvalID).Updates(map[string]any{
		"expires_at": 1,
	}).Error)

	w := performAgentApprovalRequest(ApproveAgentApprovalRequest, 9, http.MethodPost, "/approvals/:approval_id/approve", "/approvals/"+approvalID+"/approve", `{}`)
	require.Equal(t, http.StatusOK, w.Code)
	envelope := decodeAgentApprovalEnvelope(t, w)
	require.Equal(t, false, envelope["success"])
	require.Contains(t, envelope["message"], "approval request expired")

	var approval model.AgentApprovalRequest
	require.NoError(t, db.Where("approval_id = ?", approvalID).First(&approval).Error)
	require.Equal(t, "expired", approval.Status)
}

func TestAgentCredentialLeaseCanBeRevoked(t *testing.T) {
	db := setupAgentApprovalTestDB(t)
	body := `{
		"deployment_id":"dep_approval_4",
		"binding_scope":"project-alpha",
		"operation":"write_repository",
		"resource_id":"repo-main",
		"resource_type":"git",
		"resource_scope":"https://example.invalid/acme/repo#main",
		"target_role":"backend",
		"risk_level":"high",
		"requires_credential":true,
		"secret_ref":"azkv://heicode-kv.vault.azure.net/secrets/repo-main",
		"ttl_seconds":600
	}`

	_, approvalID := createAgentApprovalForTest(t, 10, body)
	w := performAgentApprovalRequest(ApproveAgentApprovalRequest, 10, http.MethodPost, "/approvals/:approval_id/approve", "/approvals/"+approvalID+"/approve", `{}`)
	require.Equal(t, http.StatusOK, w.Code)
	var storedLease model.AgentCredentialLease
	require.NoError(t, db.Where("approval_id = ?", approvalID).First(&storedLease).Error)

	revokePath := "/credential-leases/" + storedLease.LeaseID + "/revoke"
	w = performAgentApprovalRequest(RevokeAgentCredentialLease, 10, http.MethodPost, "/credential-leases/:lease_id/revoke", revokePath, `{"reason":"任务结束"}`)
	require.Equal(t, http.StatusOK, w.Code)
	envelope := decodeAgentApprovalEnvelope(t, w)
	require.Equal(t, true, envelope["success"])
	data := envelope["data"].(map[string]any)
	require.Equal(t, "revoked", data["status"])

	require.NoError(t, db.Where("lease_id = ?", storedLease.LeaseID).First(&storedLease).Error)
	require.Equal(t, "revoked", storedLease.Status)
	require.Greater(t, storedLease.RevokedAt, int64(0))
}