heicode-mananger/heicode/controller/agent_template_test.go

package controller

import (
	"context"
	"io"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"

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

func TestBuildAgentEnvFromBindings_BuiltinBlobMetadataOnly(t *testing.T) {
	setupResourceControllerTestDB(t)

	// No secret_ref -> AZURE_BLOB_ACCOUNT_KEY is skipped, not an error.
	b := model.ResourceBinding{
		UserId:       7,
		Name:         "blob1",
		ResourceType: "blob",
		Provider:     "azure",
		Metadata:     `{"account":"myacct","container":"uploads"}`,
	}
	require.NoError(t, model.DB.Create(&b).Error)

	env, err := buildAgentEnvFromBindings(7, []int{b.Id})
	require.NoError(t, err)
	require.Equal(t, "myacct", env["AZURE_BLOB_ACCOUNT_NAME"])
	require.Equal(t, "uploads", env["AZURE_BLOB_CONTAINER"])
	_, hasKey := env["AZURE_BLOB_ACCOUNT_KEY"]
	require.False(t, hasKey) // no secret bound
}

func TestBuildAgentEnvFromBindings_DatabaseProviderPrefix(t *testing.T) {
	setupResourceControllerTestDB(t)
	cases := []struct {
		provider string
		wantHost string // env var that should carry the host (AM names)
	}{
		{"mysql", "MYSQL_HOST"},
		{"postgres", "POSTGRES_HOST"},
		{"postgresql", "POSTGRES_HOST"}, // alias normalises to postgres
		{"pg", "POSTGRES_HOST"},         // alias normalises to postgres
	}
	for _, tc := range cases {
		b := model.ResourceBinding{
			UserId:       7,
			Name:         "db-" + tc.provider,
			ResourceType: "database",
			Provider:     tc.provider,
			// form fields: host/port/db_name/username (+ secret database_password)
			Metadata: `{"host":"db.example","port":"5432","db_name":"app","username":"u"}`,
		}
		require.NoError(t, model.DB.Create(&b).Error)
		env, err := buildAgentEnvFromBindings(7, []int{b.Id})
		require.NoError(t, err, tc.provider)
		require.Equal(t, "db.example", env[tc.wantHost], "provider %s -> %s", tc.provider, tc.wantHost)
	}
}

func TestBuildAgentEnvFromBindings_UnsupportedDBProviderRejected(t *testing.T) {
	setupResourceControllerTestDB(t)
	b := model.ResourceBinding{UserId: 7, Name: "r", ResourceType: "database", Provider: "redis", Status: "active", Metadata: `{"host":"h"}`}
	require.NoError(t, model.DB.Create(&b).Error)
	_, err := buildAgentEnvFromBindings(7, []int{b.Id}) // AM supports only mysql/postgres
	require.Error(t, err)
}

func TestBuildAgentEnvFromBindings_GitNamesProviderAgnostic(t *testing.T) {
	setupResourceControllerTestDB(t)
	for _, provider := range []string{"github", "gitea", "gitlab"} {
		b := model.ResourceBinding{
			UserId:       7,
			Name:         "git-" + provider,
			ResourceType: "git",
			Provider:     provider,
			Metadata:     `{"repo_url":"https://x/owner/repo","default_branch":"main"}`,
		}
		require.NoError(t, model.DB.Create(&b).Error)
		env, err := buildAgentEnvFromBindings(7, []int{b.Id})
		require.NoError(t, err)
		require.Equal(t, "https://x/owner/repo", env["GIT_REPO_URL"])
		require.Equal(t, "main", env["GIT_DEFAULT_BRANCH"])
		require.Equal(t, provider, env["GIT_PROVIDER"]) // sourced from binding.Provider
	}
}

func TestBuildAgentEnvFromBindings_SameTypeCollisionRejected(t *testing.T) {
	setupResourceControllerTestDB(t)
	b1 := model.ResourceBinding{UserId: 7, Name: "g1", ResourceType: "git", Provider: "github", Status: "active", Metadata: `{"repo_url":"https://x/a"}`}
	b2 := model.ResourceBinding{UserId: 7, Name: "g2", ResourceType: "git", Provider: "gitea", Status: "active", Metadata: `{"repo_url":"https://y/b"}`}
	require.NoError(t, model.DB.Create(&b1).Error)
	require.NoError(t, model.DB.Create(&b2).Error)

	_, err := buildAgentEnvFromBindings(7, []int{b1.Id, b2.Id}) // both set GIT_REPO_URL
	require.Error(t, err)
}

func TestBuildAgentEnvFromBindings_InactiveRejected(t *testing.T) {
	setupResourceControllerTestDB(t)
	b := model.ResourceBinding{UserId: 7, Name: "x", ResourceType: "blob", Provider: "blob", Status: "revoked", Metadata: `{"account":"a"}`}
	require.NoError(t, model.DB.Create(&b).Error)

	_, err := buildAgentEnvFromBindings(7, []int{b.Id}) // not active
	require.Error(t, err)
}

func TestBuildAgentEnvFromBindings_OwnershipEnforced(t *testing.T) {
	setupResourceControllerTestDB(t)
	b := model.ResourceBinding{UserId: 7, Name: "x", ResourceType: "blob", Provider: "blob", Metadata: `{}`}
	require.NoError(t, model.DB.Create(&b).Error)

	_, err := buildAgentEnvFromBindings(99, []int{b.Id}) // different user
	require.Error(t, err)
}

func TestBuildAgentEnvFromBindings_UnsupportedType(t *testing.T) {
	setupResourceControllerTestDB(t)
	b := model.ResourceBinding{UserId: 7, Name: "x", ResourceType: "quantum", Provider: "weird", Metadata: `{}`}
	require.NoError(t, model.DB.Create(&b).Error)

	_, err := buildAgentEnvFromBindings(7, []int{b.Id})
	require.Error(t, err)
}

func TestBuildAgentEnvFromBindings_Empty(t *testing.T) {
	setupResourceControllerTestDB(t)
	env, err := buildAgentEnvFromBindings(7, nil)
	require.NoError(t, err)
	require.Empty(t, env)
}

func TestAgentTemplatePathSubstitution(t *testing.T) {
	require.Contains(t, agentTemplateStartPath(), "/agents")
	require.Contains(t, agentTemplateAgentStopPath("rt-9"), "/agents/rt-9/stop")
	require.Contains(t, agentTemplateAgentPath("rt-9"), "/agents/rt-9")
}

func TestAMStartTemplateAgent_RoundTrip(t *testing.T) {
	var gotBody string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.Equal(t, http.MethodPost, r.Method)
		require.Equal(t, "/agents", r.URL.Path) // AM's create endpoint
		b, _ := io.ReadAll(r.Body)
		gotBody = string(b)
		w.Header().Set("Content-Type", "application/json")
		// AM-style response: access_info.domain, namespace; no access_token.
		_, _ = w.Write([]byte(`{"success":true,"data":{"namespace":"ns-1","access_info":{"domain":"https://abc.agents.example"}}}`))
	}))
	defer srv.Close()
	t.Setenv("AGENT_RUNTIME_BASE_URL", srv.URL)

	res, err := amStartTemplateAgent(context.Background(), amStartArgs{
		ManagerDeploymentID: "dep_1",
		UserID:              "22",
		TemplateKey:         "architect",
		AgentDefinition:     "---\nname: architect\n---\nbody",
		Model:               "gpt-5.4",
		Env:                 map[string]string{"GIT_REPO_URL": "https://x/r"},
		CallbackURL:         "https://cb",
	})
	require.NoError(t, err)
	// payload mapped onto AM's contract
	require.Contains(t, gotBody, `"template":"coding_a2a_agent"`)
	require.Contains(t, gotBody, "AGENT_INSTRUCTION_TEXT") // .md folded into env
	require.Contains(t, gotBody, "AGENT_ROLE_NAME")
	require.Contains(t, gotBody, "OPENAI_BASE_URL")
	require.Contains(t, gotBody, "manager_deployment_id")
	// response parsed
	require.Equal(t, "ns-1", res.RuntimeID)
	require.Equal(t, "https://abc.agents.example", res.Subdomain)
	require.Equal(t, "running", res.Status) // defaulted
}

func TestAMGetAgentStatus_RoundTrip(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.Equal(t, "/agents/rt-1", r.URL.Path)
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"success":true,"data":{"status":"running"}}`))
	}))
	defer srv.Close()
	t.Setenv("AGENT_RUNTIME_BASE_URL", srv.URL)

	status, err := amGetAgentStatus(context.Background(), "rt-1")
	require.NoError(t, err)
	require.Equal(t, "running", status)
}

func TestParseAgentFrontmatter(t *testing.T) {
	md := "---\nname: architect\ndescription: Strategic Architecture Advisor\nmodel: opus\ndisallowedTools: Write, Edit\n---\n\n<Agent_Prompt>body</Agent_Prompt>"
	fm := parseAgentFrontmatter(md)
	require.Equal(t, "architect", fm["name"])
	require.Equal(t, "opus", fm["model"])
	require.Equal(t, "Strategic Architecture Advisor", fm["description"])
}

func TestSeedAndLoadAgentTemplate(t *testing.T) {
	db := setupResourceControllerTestDB(t)
	require.NoError(t, db.AutoMigrate(&model.AgentTemplate{}))
	// reset the once guard so seeding runs against this fresh test DB
	seedAgentTemplatesOnce = sync.Once{}

	tpl, ok := loadAgentTemplate("architect")
	require.True(t, ok)
	require.Equal(t, "架构顾问", tpl.NameZh)
	require.NotEmpty(t, tpl.Definition)
	require.Contains(t, tpl.Definition, "name: architect")

	var count int64
	db.Model(&model.AgentTemplate{}).Count(&count)
	require.Equal(t, int64(19), count) // all presets seeded
}

func TestAgentTokenSealRoundTrip(t *testing.T) {
	old := common.CryptoSecret
	common.CryptoSecret = "unit-test-crypto-secret"
	defer func() { common.CryptoSecret = old }()

	sealed := sealAgentToken("super-secret-agent-token")
	require.NotEqual(t, "super-secret-agent-token", sealed) // encrypted at rest
	require.NotEmpty(t, sealed)
	require.Equal(t, "super-secret-agent-token", unsealAgentToken(sealed)) // recovered

	require.Equal(t, "", sealAgentToken(""))
	require.Equal(t, "plain-legacy", unsealAgentToken("plain-legacy")) // non-sealed passes through
}

func TestTemplateAgentResponse(t *testing.T) {
	row := model.AgentDeployment{
		DeploymentID:   "dep_abc",
		TemplateID:     "tpl1",
		Subdomain:      "https://abc.agents.example",
		AccessToken:    "tok",
		BindingIDsJSON: "[1,2]",
		Status:         "running",
	}
	resp := templateAgentResponse(row)
	require.Equal(t, "dep_abc", resp["agent_id"])
	require.Equal(t, "tpl1", resp["template_id"])
	require.Equal(t, []int{1, 2}, resp["binding_ids"])
	require.Equal(t, "running", resp["status"])
}