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
	}

	// Enforce the per-user deployed-agent cap (产品文档「个人5/团队8」). The cap
	// source is env HEICODE_MAX_AGENTS_PER_USER (default 5 = personal tier); the
	// team tier (8) is a follow-up that depends on subscription-tier modelling
	// (see issue #7). Only the user's own non-stopped template agents count — a
	// stopped agent consumes no runtime, so it is excluded. LOWER(status) keeps
	// the comparison portable across SQLite/MySQL/PostgreSQL. maxAgents<=0 means
	// "unlimited" (escape hatch). This closes the over-deployment hole where a
	// user could deploy unbounded agents across repeated calls.
	if maxAgents := common.GetEnvOrDefault("HEICODE_MAX_AGENTS_PER_USER", 5); maxAgents > 0 {
		var active int64
		if err := model.DB.Model(&model.AgentDeployment{}).
			Where("user_id = ? AND template_id <> '' AND LOWER(status) <> ?", strconv.Itoa(userID), "stopped").
			Count(&active).Error; err != nil {
			agentError(c, "DEPLOYMENT_PERSIST_FAILED", "failed to check agent quota")
			return
		}
		if active >= int64(maxAgents) {
			agentError(c, "POLICY_REJECTED", "agent deployment limit reached ("+strconv.Itoa(maxAgents)+"); stop or delete an existing agent first")
			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

	// Persist the agent as Pending and return immediately; start it in AM in the
	// background. Starting an agent in AM takes ~30s — longer than the Azure
	// gateway's ~20s timeout — so doing it inside the request returned 504, which
	// cancelled the request context, aborted the AM call and rolled the deploy
	// back. Now the client gets the Pending agent at once and polls
	// GET /api/heicode/agents (which refreshes live status from AM), seeing
	// Pending -> running once AM is up (or -> failed if AM start failed).
	now := agentNow()
	nowMs := time.Now().UnixMilli()
	row := model.AgentDeployment{
		DeploymentID:   deploymentID,
		UserID:         strconv.Itoa(userID),
		TemplateID:     req.TemplateID,
		AccessToken:    sealAgentToken(agentAccessToken),
		BindingIDsJSON: string(bindingIDsJSON),
		Status:         "Pending",
		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)
		agentError(c, "DEPLOYMENT_PERSIST_FAILED", "failed to persist agent")
		return
	}

	startTemplateAgentAsync(deploymentID, modelTokenID, amStartArgs{
		ManagerDeploymentID: deploymentID,
		UserID:              strconv.Itoa(userID),
		TemplateKey:         tpl.TemplateKey,
		AgentDefinition:     tpl.Definition,
		Model:               tpl.Model,
		Env:                 env,
		CallbackURL:         agentRuntimeCallbackURL(),
	})

	common.ApiSuccess(c, templateAgentResponse(row))
}

// startTemplateAgentAsync starts the template agent in AM off the request path
// and records the outcome. The deploy handler returns the Pending agent at once
// so it never blocks ~30s on AM (which exceeded the Azure gateway timeout → 504
// + rollback). Uses context.Background() because the request context is gone
// once the handler returned; the AM HTTP call is still bounded by amTemplateDo's
// own client timeout. The client observes Pending -> running (or -> failed) by
// polling the agent list, which refreshes live status from AM.
func startTemplateAgentAsync(deploymentID string, modelTokenID int, args amStartArgs) {
	go func() {
		result, err := amStartTemplateAgent(context.Background(), args)
		if model.DB == nil {
			return
		}
		now := agentNow()
		nowMs := time.Now().UnixMilli()
		if err != nil {
			common.SysLog("startTemplateAgentAsync: AM start failed for " + deploymentID + ": " + err.Error())
			reason := err.Error()
			if len(reason) > 480 {
				reason = strings.ToValidUTF8(reason[:480], "")
			}
			// Only flip Pending->failed. If the user stopped/deleted/cancelled during
			// the AM start window the status is no longer Pending (or the row is gone),
			// and we must NOT clobber that terminal intent. The status guard makes the
			// transition own-or-nothing; LOWER() because the row is stored as "Pending".
			res := model.DB.Model(&model.AgentDeployment{}).
				Where("deployment_id = ? AND LOWER(status) = ?", deploymentID, "pending").
				Updates(map[string]any{
					"status":          "failed",
					"failure_reason":  reason,
					"updated_at_text": now,
					"updated_at_ms":   nowMs,
				})
			// Revoke the minted key ONLY when we owned the Pending->failed transition
			// (the agent never started). If the user already acted, leave the token to
			// that path: delete revokes it, stop intentionally retains it. Avoids both
			// double-revoke and revoking a key the stop path means to keep.
			if res.Error == nil && res.RowsAffected > 0 {
				revokeAgentModelToken(modelTokenID)
			}
			return
		}
		// Success: flip Pending->running and fill in the runtime. Same status guard —
		// if the user stopped/deleted/cancelled while AM was starting, this updates 0
		// rows and we must clean up the now-orphaned runtime instead of resurrecting it.
		res := model.DB.Model(&model.AgentDeployment{}).
			Where("deployment_id = ? AND LOWER(status) = ?", deploymentID, "pending").
			Updates(map[string]any{
				"subdomain":             result.Subdomain,
				"runtime_deployment_id": result.RuntimeID,
				"status":                firstNonEmpty(result.Status, "running"),
				"updated_at_text":       now,
				"updated_at_ms":         nowMs,
			})
		if res.Error != nil {
			common.SysLog("startTemplateAgentAsync: persist result failed for " + deploymentID + ": " + res.Error.Error())
			return
		}
		if res.RowsAffected == 0 {
			// No longer Pending: the user stopped/deleted/cancelled while AM was still
			// starting. Re-read to log which case, then delete the orphan runtime AM
			// just started — the action is identical either way (clean up, never
			// resurrect), so the record is never flipped back to running. Token handling
			// is left to the user's stop/delete path (delete revoked it; stop keeps it).
			var cur model.AgentDeployment
			if model.DB.Where("deployment_id = ?", deploymentID).First(&cur).Error != nil {
				common.SysLog("startTemplateAgentAsync: " + deploymentID + " deleted during AM start; cleaning orphan runtime")
			} else {
				common.SysLog("startTemplateAgentAsync: " + deploymentID + " no longer Pending (status=" + cur.Status + ") during AM start; cleaning orphan runtime, preserving user state")
			}
			if strings.TrimSpace(result.RuntimeID) != "" {
				if delErr := amDeleteTemplateAgent(context.Background(), result.RuntimeID); delErr != nil {
					common.SysLog("startTemplateAgentAsync orphan cleanup failed for " + deploymentID + ": " + delErr.Error())
				}
			}
		}
	}()
}

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

// HeicodeGetAgentUsage: GET /api/heicode/agents/:deployment_id/usage
// Per-agent model-usage rollup (#9): the agent calls HM /v1/* with its minted
// token named "agent:<deployment_id>", so its consumption is the sum of consume
// logs under that token name. Optional ?start=&end= unix-second window. Returns
// raw quota + quota_per_unit (caller converts, same contract as /api/heicode/self).
func HeicodeGetAgentUsage(c *gin.Context) {
	row, ok := findUserTemplateAgent(c)
	if !ok {
		return
	}
	start, _ := strconv.ParseInt(strings.TrimSpace(c.Query("start")), 10, 64)
	end, _ := strconv.ParseInt(strings.TrimSpace(c.Query("end")), 10, 64)
	usage, err := model.SumAgentUsage(c.GetInt("id"), "agent:"+row.DeploymentID, start, end)
	if err != nil {
		agentError(c, "DEPLOYMENT_CONFLICT", "failed to aggregate agent usage")
		return
	}
	common.ApiSuccess(c, gin.H{
		"agent_id":          row.DeploymentID,
		"quota":             usage.Quota,
		"prompt_tokens":     usage.PromptTokens,
		"completion_tokens": usage.CompletionTokens,
		"call_count":        usage.CallCount,
		"quota_per_unit":    common.QuotaPerUnit,
	})
}

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