diff --git a/heicode/controller/agent_template_handlers.go b/heicode/controller/agent_template_handlers.go index c4678c87..9a5008e2 100644 --- a/heicode/controller/agent_template_handlers.go +++ b/heicode/controller/agent_template_handlers.go @@ -48,6 +48,42 @@ func unsealAgentToken(stored string) string { 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-") 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 @@ -117,6 +153,15 @@ func HeicodeDeployAgent(c *gin.Context) { 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 + // Ask AM to start the agent with the template definition (.md) + env injected. result, err := amStartTemplateAgent(c.Request.Context(), amStartArgs{ ManagerDeploymentID: deploymentID, @@ -128,6 +173,7 @@ func HeicodeDeployAgent(c *gin.Context) { CallbackURL: agentRuntimeCallbackURL(), }) if err != nil { + revokeAgentModelToken(modelTokenID) // don't leak the minted key agentError(c, "RUNTIME_UNAVAILABLE", "failed to start agent: "+err.Error()) return } @@ -143,6 +189,7 @@ func HeicodeDeployAgent(c *gin.Context) { BindingIDsJSON: string(bindingIDsJSON), RuntimeDeploymentID: result.RuntimeID, Status: firstNonEmpty(result.Status, "running"), + ModelTokenID: modelTokenID, CreatedAtText: now, UpdatedAtText: now, CreatedAtMs: nowMs, @@ -150,6 +197,7 @@ func HeicodeDeployAgent(c *gin.Context) { } 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) != "" { @@ -292,6 +340,7 @@ func HeicodeDeleteAgent(c *gin.Context) { 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", diff --git a/heicode/model/agent_deployment.go b/heicode/model/agent_deployment.go index fd51175f..b4321001 100644 --- a/heicode/model/agent_deployment.go +++ b/heicode/model/agent_deployment.go @@ -35,6 +35,7 @@ type AgentDeployment struct { Subdomain string `gorm:"type:varchar(255)" json:"subdomain"` // AM-assigned public address AccessToken string `gorm:"type:varchar(512)" json:"access_token"` // client presents this to the agent; HM stores a copy BindingIDsJSON string `gorm:"type:text" json:"binding_ids_json"` // JSON array of resource_binding ids attached to this agent + ModelTokenID int `gorm:"index" json:"model_token_id"` // sk- token minted for this agent's model calls (OPENAI_API_KEY); revoked on delete } func (AgentDeployment) TableName() string {