heicode-mananger/heicode/controller/agent_template_handlers.go

package controller

import (
	"strconv"
	"strings"
	"time"

	"github.com/gin-gonic/gin"

	"github.com/heicode/manager/common"
	"github.com/heicode/manager/model"
)

// ─────────────────────────────────────────────────────────────────────────────
// Template-agent client API (new model).
//
// User deploys a template agent from the web console; AM starts it with the
// selected bound resources injected as env and returns a unique subdomain +
// access token. The desktop client reads the agent list here, then connects to
// the subdomain directly over SSE. HM is NOT in the agent conversation path.
// ─────────────────────────────────────────────────────────────────────────────

// sealAgentToken encrypts the AM-issued access token before it is persisted
// (AES-256-GCM via CryptoSecret). Falls back to plaintext only if CryptoSecret
// is unset, so the token is never lost.
func sealAgentToken(token string) string {
	if strings.TrimSpace(token) == "" {
		return ""
	}
	sealed, err := common.SealWithCryptoSecret([]byte(token))
	if err != nil {
		common.SysLog("WARNING: agent access_token stored unencrypted (CryptoSecret unset?): " + err.Error())
		return token
	}
	return sealed
}

// unsealAgentToken reverses sealAgentToken. If the stored value is not a sealed
// blob (legacy plaintext / CryptoSecret unset), it is returned as-is.
func unsealAgentToken(stored string) string {
	if strings.TrimSpace(stored) == "" {
		return ""
	}
	plain, err := common.UnsealWithCryptoSecret(stored)
	if err != nil {
		return stored
	}
	return string(plain)
}

// templateAgentResponse maps a deployed template-agent record to the client view.
func templateAgentResponse(row model.AgentDeployment) gin.H {
	var bindingIDs []int
	if strings.TrimSpace(row.BindingIDsJSON) != "" {
		_ = common.UnmarshalJsonStr(row.BindingIDsJSON, &bindingIDs)
	}
	return gin.H{
		"agent_id":     row.DeploymentID,
		"template_id":  row.TemplateID,
		"subdomain":    row.Subdomain,
		"access_token": unsealAgentToken(row.AccessToken),
		"binding_ids":  bindingIDs,
		"status":       row.Status,
		"runtime_id":   row.RuntimeDeploymentID,
		"created_at":   row.CreatedAtText,
		"updated_at":   row.UpdatedAtText,
	}
}

// HeicodeListAgentTemplates is defined in agent_template_library.go (reads the
// HM-maintained agent_templates table, Chinese display).

// HeicodeDeployAgent: POST /api/heicode/agents
// Body: {template_id (= template_key), binding_ids:[...]}.
func HeicodeDeployAgent(c *gin.Context) {
	userID := c.GetInt("id")
	if userID <= 0 {
		agentError(c, "POLICY_REJECTED", "authentication required")
		return
	}
	var req struct {
		TemplateID string `json:"template_id"`
		BindingIDs []int  `json:"binding_ids"`
	}
	if err := common.UnmarshalBodyReusable(c, &req); err != nil {
		agentError(c, "POLICY_REJECTED", "invalid request body")
		return
	}
	req.TemplateID = strings.TrimSpace(req.TemplateID)
	if req.TemplateID == "" {
		agentError(c, "POLICY_REJECTED", "template_id is required")
		return
	}
	if req.BindingIDs == nil {
		req.BindingIDs = []int{}
	}
	if model.DB == nil {
		agentError(c, "DEPLOYMENT_PERSIST_FAILED", "database not initialised")
		return
	}

	// Load the HM-maintained template (its .md definition is sent to AM).
	tpl, ok := loadAgentTemplate(req.TemplateID)
	if !ok {
		agentError(c, "POLICY_REJECTED", "unknown template_id")
		return
	}

	// Resolve selected bindings -> env (non-secret config + KV-resolved secrets).
	// NEVER log env: it can contain plaintext secrets.
	env, err := buildAgentEnvFromBindings(userID, req.BindingIDs)
	if err != nil {
		agentError(c, "RESOURCE_BINDING_INVALID", err.Error())
		return
	}

	deploymentID := "dep_" + common.GetUUID()[:12]
	bindingIDsJSON, _ := common.Marshal(req.BindingIDs)

	// Ask AM to start the agent with the template definition (.md) + env injected.
	result, err := amStartTemplateAgent(c.Request.Context(), amStartArgs{
		ManagerDeploymentID: deploymentID,
		UserID:              strconv.Itoa(userID),
		TemplateKey:         tpl.TemplateKey,
		AgentDefinition:     tpl.Definition,
		Model:               tpl.Model,
		Env:                 env,
		CallbackURL:         agentRuntimeCallbackURL(),
	})
	if err != nil {
		agentError(c, "RUNTIME_UNAVAILABLE", "failed to start agent: "+err.Error())
		return
	}

	now := agentNow()
	nowMs := time.Now().UnixMilli()
	row := model.AgentDeployment{
		DeploymentID:        deploymentID,
		UserID:              strconv.Itoa(userID),
		TemplateID:          req.TemplateID,
		Subdomain:           result.Subdomain,
		AccessToken:         sealAgentToken(result.AccessToken),
		BindingIDsJSON:      string(bindingIDsJSON),
		RuntimeDeploymentID: result.RuntimeID,
		Status:              firstNonEmpty(result.Status, "running"),
		CreatedAtText:       now,
		UpdatedAtText:       now,
		CreatedAtMs:         nowMs,
		UpdatedAtMs:         nowMs,
	}
	if err := model.DB.Create(&row).Error; err != nil {
		common.SysLog("HeicodeDeployAgent persist: " + err.Error())
		// We started an agent in AM but failed to record it: roll back the
		// orphan so it does not leak/keep running with no Manager record.
		if strings.TrimSpace(result.RuntimeID) != "" {
			if delErr := amDeleteTemplateAgent(c.Request.Context(), result.RuntimeID); delErr != nil {
				common.SysLog("HeicodeDeployAgent orphan cleanup failed: " + delErr.Error())
			}
		}
		agentError(c, "DEPLOYMENT_PERSIST_FAILED", "failed to persist agent")
		return
	}
	common.ApiSuccess(c, templateAgentResponse(row))
}

// findUserTemplateAgent loads a deployed template agent owned by the caller.
func findUserTemplateAgent(c *gin.Context) (model.AgentDeployment, bool) {
	var row model.AgentDeployment
	userID := c.GetInt("id")
	deploymentID := strings.TrimSpace(c.Param("deployment_id"))
	if userID <= 0 || deploymentID == "" || model.DB == nil {
		agentError(c, "POLICY_REJECTED", "deployment_id and authentication required")
		return row, false
	}
	// template_id <> '' guards against touching a legacy task deployment of the
	// same user through the new template-agent endpoints.
	if err := model.DB.Where("deployment_id = ? AND user_id = ? AND template_id <> ''", deploymentID, strconv.Itoa(userID)).First(&row).Error; err != nil {
		agentError(c, "DEPLOYMENT_CONFLICT", "agent not found")
		return row, false
	}
	return row, true
}

// HeicodeListAgents: GET /api/heicode/agents
func HeicodeListAgents(c *gin.Context) {
	userID := c.GetInt("id")
	if userID <= 0 || model.DB == nil {
		agentError(c, "POLICY_REJECTED", "authentication required")
		return
	}
	var rows []model.AgentDeployment
	// Only template-agent records (TemplateID set), not legacy task deployments.
	if err := model.DB.Where("user_id = ? AND template_id <> ''", strconv.Itoa(userID)).
		Order("created_at_ms desc").Find(&rows).Error; err != nil {
		agentError(c, "DEPLOYMENT_CONFLICT", "failed to list agents")
		return
	}
	items := make([]gin.H, 0, len(rows))
	for _, row := range rows {
		items = append(items, templateAgentResponse(row))
	}
	common.ApiSuccess(c, gin.H{"items": items, "total": len(items)})
}

// refreshAgentStatus best-effort pulls the live status from AM and persists it.
// On any AM error it keeps the last-known status (never blocks the read).
func refreshAgentStatus(c *gin.Context, row *model.AgentDeployment) {
	if strings.TrimSpace(row.RuntimeDeploymentID) == "" {
		return
	}
	status, err := amGetAgentStatus(c.Request.Context(), row.RuntimeDeploymentID)
	if err != nil || strings.TrimSpace(status) == "" || status == row.Status {
		return
	}
	now := agentNow()
	nowMs := time.Now().UnixMilli()
	row.Status = status
	row.UpdatedAtText = now
	row.UpdatedAtMs = nowMs
	// Field-level update (not Save of the whole row) so a concurrent stop/delete
	// or other column write isn't clobbered by a stale full-row save.
	if model.DB != nil {
		_ = model.DB.Model(&model.AgentDeployment{}).
			Where("deployment_id = ?", row.DeploymentID).
			Updates(map[string]any{"status": status, "updated_at_text": now, "updated_at_ms": nowMs}).Error
	}
}

// HeicodeGetAgent: GET /api/heicode/agents/:deployment_id
func HeicodeGetAgent(c *gin.Context) {
	row, ok := findUserTemplateAgent(c)
	if !ok {
		return
	}
	refreshAgentStatus(c, &row)
	common.ApiSuccess(c, templateAgentResponse(row))
}

// HeicodeGetAgentStatus: GET /api/heicode/agents/:deployment_id/status
// Pulls the live status from AM (is the agent running / crashed?).
func HeicodeGetAgentStatus(c *gin.Context) {
	row, ok := findUserTemplateAgent(c)
	if !ok {
		return
	}
	refreshAgentStatus(c, &row)
	common.ApiSuccess(c, gin.H{
		"agent_id":   row.DeploymentID,
		"status":     row.Status,
		"updated_at": row.UpdatedAtText,
	})
}

// HeicodeStopAgent: POST /api/heicode/agents/:deployment_id/stop
func HeicodeStopAgent(c *gin.Context) {
	row, ok := findUserTemplateAgent(c)
	if !ok {
		return
	}
	if strings.TrimSpace(row.RuntimeDeploymentID) != "" {
		if err := amStopTemplateAgent(c.Request.Context(), row.RuntimeDeploymentID); err != nil {
			agentError(c, "RUNTIME_UNAVAILABLE", "failed to stop agent: "+err.Error())
			return
		}
	}
	row.Status = "stopped"
	row.UpdatedAtText = agentNow()
	row.UpdatedAtMs = time.Now().UnixMilli()
	if err := model.DB.Save(&row).Error; err != nil {
		agentError(c, "DEPLOYMENT_PERSIST_FAILED", "failed to persist stop")
		return
	}
	common.ApiSuccess(c, templateAgentResponse(row))
}

// HeicodeDeleteAgent: DELETE /api/heicode/agents/:deployment_id
func HeicodeDeleteAgent(c *gin.Context) {
	row, ok := findUserTemplateAgent(c)
	if !ok {
		return
	}
	if strings.TrimSpace(row.RuntimeDeploymentID) != "" {
		if err := amDeleteTemplateAgent(c.Request.Context(), row.RuntimeDeploymentID); err != nil {
			agentError(c, "RUNTIME_UNAVAILABLE", "failed to delete agent: "+err.Error())
			return
		}
	}
	if err := model.DB.Delete(&row).Error; err != nil {
		agentError(c, "DEPLOYMENT_PERSIST_FAILED", "failed to delete agent")
		return
	}
	common.ApiSuccess(c, gin.H{"deployment_id": row.DeploymentID, "status": "deleted"})
}