heicode-mananger/heicode/controller/agent_template_handlers.go

package controller

import (
	"context"
	"crypto/subtle"
	"strconv"
	"strings"
	"sync"
	"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)
}

// mintAgentModelToken creates a hidden, unlimited-quota sk- token for the user
// so the started agent can call HM /v1/* (OPENAI_API_KEY) and bill to the user.
// Returns the bearer value ("sk-<key>") and the token id (stored for cleanup).
func mintAgentModelToken(userID int, deploymentID string) (string, int, error) {
	rawKey, err := common.GenerateKey()
	if err != nil {
		return "", 0, err
	}
	now := common.GetTimestamp()
	tok := model.Token{
		UserId:         userID,
		Name:           "agent:" + deploymentID,
		Key:            rawKey,
		Status:         common.TokenStatusEnabled,
		CreatedTime:    now,
		AccessedTime:   now,
		ExpiredTime:    -1,   // never naturally
		UnlimitedQuota: true, // bills straight from User.Quota
		HideFromUserUI: true, // not a user-created key; don't clutter the token list
	}
	if err := tok.Insert(); err != nil {
		return "", 0, err
	}
	return "sk-" + rawKey, tok.Id, nil
}

// revokeAgentModelToken deletes the agent's minted model token (best-effort).
func revokeAgentModelToken(tokenID int) {
	if tokenID <= 0 || model.DB == nil {
		return
	}
	if err := model.DB.Where("id = ?", tokenID).Delete(&model.Token{}).Error; err != nil {
		common.SysLog("revokeAgentModelToken: " + err.Error())
	}
}

// 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)

	// Mint a model key so the agent can call HM /v1/* (OPENAI_API_KEY).
	modelKey, modelTokenID, err := mintAgentModelToken(userID, deploymentID)
	if err != nil {
		common.SysLog("HeicodeDeployAgent mint model token: " + err.Error())
		agentError(c, "DEPLOYMENT_PERSIST_FAILED", "failed to provision model key")
		return
	}
	env["OPENAI_API_KEY"] = modelKey

	// Per-agent client↔agent access token. HM mints a random secret, injects it
	// into the agent env (AGENT_ACCESS_TOKEN) AND returns it to the deploying
	// client (the agent list's access_token). Only the user who deployed this
	// agent receives the token, so only they can drive it: AM can verify the
	// caller's token == its own env AGENT_ACCESS_TOKEN locally (no HM round-trip),
	// or call POST /api/heicode/agent-access/verify to let HM be the authority.
	// AM may opt out of this entirely. HEICODE_AGENT_ID lets the agent name itself
	// when calling the verify endpoint.
	agentAccessToken := common.GetUUID()
	env["AGENT_ACCESS_TOKEN"] = agentAccessToken
	env["HEICODE_AGENT_ID"] = deploymentID

	// 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 {
		revokeAgentModelToken(modelTokenID) // don't leak the minted key
		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(agentAccessToken),
		BindingIDsJSON:      string(bindingIDsJSON),
		RuntimeDeploymentID: result.RuntimeID,
		Status:              firstNonEmpty(result.Status, "running"),
		ModelTokenID:        modelTokenID,
		CreatedAtText:       now,
		UpdatedAtText:       now,
		CreatedAtMs:         nowMs,
		UpdatedAtMs:         nowMs,
	}
	if err := model.DB.Create(&row).Error; err != nil {
		common.SysLog("HeicodeDeployAgent persist: " + err.Error())
		revokeAgentModelToken(modelTokenID)
		// 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
	}
	// Refresh live status from AM for every non-terminal agent BEFORE returning,
	// so the list (and the desktop client that polls it) reflects reality without
	// needing the detail panel to be opened. Previously only the detail/status
	// endpoints refreshed, so a freshly-deployed agent stayed "Pending" in the
	// list forever. Bounded: per-user agent count is capped (≤8) and the refreshes
	// run concurrently under a short deadline, so the list never hangs on a slow AM.
	refreshAgentStatusBatch(c.Request.Context(), rows)

	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)})
}

// refreshAgentStatusBatch concurrently refreshes the live status of every
// non-terminal agent in rows (best-effort) under a single bounded deadline.
// Each row is updated in place + persisted by refreshAgentStatus.
func refreshAgentStatusBatch(ctx context.Context, rows []model.AgentDeployment) {
	ctx, cancel := context.WithTimeout(ctx, 8*time.Second)
	defer cancel()
	var wg sync.WaitGroup
	for i := range rows {
		row := &rows[i]
		if strings.TrimSpace(row.RuntimeDeploymentID) == "" ||
			strings.EqualFold(strings.TrimSpace(row.Status), "stopped") {
			continue
		}
		wg.Add(1)
		go func(r *model.AgentDeployment) {
			defer wg.Done()
			refreshAgentStatus(ctx, r)
		}(row)
	}
	wg.Wait()
}

// 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(ctx context.Context, row *model.AgentDeployment) {
	if strings.TrimSpace(row.RuntimeDeploymentID) == "" {
		return
	}
	// A user-initiated "stopped" is a terminal intent: don't let AM's eventually-
	// consistent live status (which may still report running/pending while the pod
	// drains) resurrect it. Only delete (which removes the row) leaves "stopped".
	if strings.EqualFold(strings.TrimSpace(row.Status), "stopped") {
		return
	}
	status, err := amGetAgentStatus(ctx, 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.Request.Context(), &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.Request.Context(), &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
		}
	}
	now := agentNow()
	nowMs := time.Now().UnixMilli()
	row.Status = "stopped"
	row.UpdatedAtText = now
	row.UpdatedAtMs = nowMs
	// Field-level update (not a full-row Save of a possibly-stale snapshot) — same
	// discipline as refreshAgentStatus, so concurrent column writes aren't clobbered.
	if err := model.DB.Model(&model.AgentDeployment{}).
		Where("deployment_id = ?", row.DeploymentID).
		Updates(map[string]any{"status": "stopped", "updated_at_text": now, "updated_at_ms": nowMs}).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
	}
	runtimeCleanup := "ok"
	if strings.TrimSpace(row.RuntimeDeploymentID) != "" {
		if err := amDeleteTemplateAgent(c.Request.Context(), row.RuntimeDeploymentID); err != nil {
			// Best-effort: still remove the local record so the user is never
			// stuck with an undeletable row when AM's delete is unavailable/buggy.
			runtimeCleanup = "failed"
			common.SysLog("HeicodeDeleteAgent AM delete failed (removing local record anyway): " + err.Error())
		}
	}
	if err := model.DB.Delete(&row).Error; err != nil {
		agentError(c, "DEPLOYMENT_PERSIST_FAILED", "failed to delete agent")
		return
	}
	revokeAgentModelToken(row.ModelTokenID) // invalidate the agent's model key
	common.ApiSuccess(c, gin.H{
		"deployment_id":   row.DeploymentID,
		"status":          "deleted",
		"runtime_cleanup": runtimeCleanup, // "failed" => AM may still hold an orphan
	})
}

// HeicodeVerifyAgentAccess: POST /api/heicode/agent-access/verify (public).
//
// The client↔agent access-control primitive HM provides for AM (optional). When
// the desktop client connects to an agent's subdomain it presents the per-agent
// access_token it got from the agent list. AM's agent can authorize the caller
// in one of two ways:
//
//	① Local (no HM call): compare the caller's token to its own env
//	   AGENT_ACCESS_TOKEN — they are the same secret HM injected at deploy.
//	② Authoritative: POST here with {agent_id, access_token}. HM constant-time
//	   compares against the stored token and returns {valid, user_id} so AM also
//	   learns which HM user owns the agent.
//
// Only the user who deployed the agent ever receives its token, so a valid match
// means "this user is allowed to call this agent". No user identity is leaked on
// a miss. AM may ignore this entirely and roll its own scheme.
func HeicodeVerifyAgentAccess(c *gin.Context) {
	var req struct {
		AgentID     string `json:"agent_id"`
		AccessToken string `json:"access_token"`
	}
	if err := common.UnmarshalBodyReusable(c, &req); err != nil {
		agentError(c, "INVALID_REQUEST", "invalid request body")
		return
	}
	req.AgentID = strings.TrimSpace(req.AgentID)
	req.AccessToken = strings.TrimSpace(req.AccessToken)
	if req.AgentID == "" || req.AccessToken == "" || model.DB == nil {
		common.ApiSuccess(c, gin.H{"valid": false})
		return
	}
	var row model.AgentDeployment
	if err := model.DB.Where("deployment_id = ? AND template_id <> ''", req.AgentID).First(&row).Error; err != nil {
		common.ApiSuccess(c, gin.H{"valid": false})
		return
	}
	stored := unsealAgentToken(row.AccessToken)
	if stored == "" || subtle.ConstantTimeCompare([]byte(stored), []byte(req.AccessToken)) != 1 {
		common.ApiSuccess(c, gin.H{"valid": false})
		return
	}
	common.ApiSuccess(c, gin.H{
		"valid":       true,
		"agent_id":    row.DeploymentID,
		"user_id":     row.UserID,
		"template_id": row.TemplateID,
	})
}