heicode-mananger/heicode/controller/swarm_model_key.go

package controller

import (
	"errors"
	"fmt"
	"strings"

	"gorm.io/gorm"

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

// #60 模型 key 注入(方案A) —— 对接 agent_swarm PR#43 定死的参数:
//
//   - A.1 粒度/命名:HM 为每个用户 mint 一把 per-user `sk-`,跨该用户所有 swarm run 复用;
//     KV 密文名 = `swarm-model-key-<user_id>`。
//   - A.2 KV value:JSON `{"openai_api_key":"sk-..."}`,对齐 Swarm
//     `orchestrator/agent_launcher._extract_model_key` 的解析字段。
//   - A.4 OPENAI_API_BASE:Swarm 部署常量(=HM 网关 /v1),HM **不**经 billing_context 下发。
//   - A.5 吊销:事件驱动 —— 用户全部 run 被 stop 时 Swarm 恰好发一次 `swarm.pool_terminated`,
//     HM 据此删 KV 密文 + 删 token(见 agent_callback.go handleSwarmPoolTerminated)。
//
// create 时把 putSecret 返回的 azkv:// secret_ref 放进 billing_context.secret_ref 下发;
// Swarm 侧凭 KV 读权限(A.3,运维授权 pending)解析后注入 agent 环境。
// 明文 `sk-` 仅经 secret_ref 服务端解析,绝不入代码/日志/事件/argv。

func swarmModelKeyTokenName(userID int) string {
	return fmt.Sprintf("swarm:user:%d", userID)
}

func swarmModelKeySecretName(userID int) string {
	return fmt.Sprintf("swarm-model-key-%d", userID)
}

// getOrMintSwarmModelToken 返回该用户长存的 swarm 模型 token 的 `sk-` bearer。
// 已存在(未软删)则复用(A.1:跨 run 复用、稳定 key),否则新建一把隐藏、无限额度、
// 不自然过期的系统托管 token(计费直接走 User.Quota,与 mintAgentModelToken 一致)。
func getOrMintSwarmModelToken(userID int) (string, error) {
	if userID <= 0 || model.DB == nil {
		return "", errors.New("invalid user for swarm model key")
	}
	name := swarmModelKeyTokenName(userID)
	var tok model.Token
	err := model.DB.Where("user_id = ? AND name = ?", userID, name).First(&tok).Error
	if err == nil && strings.TrimSpace(tok.Key) != "" {
		return "sk-" + tok.Key, nil
	}
	if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) {
		return "", err
	}
	rawKey, err := common.GenerateKey()
	if err != nil {
		return "", err
	}
	now := common.GetTimestamp()
	tok = model.Token{
		UserId:         userID,
		Name:           name,
		Key:            rawKey,
		Status:         common.TokenStatusEnabled,
		CreatedTime:    now,
		AccessedTime:   now,
		ExpiredTime:    -1,   // never naturally
		UnlimitedQuota: true, // bills straight from User.Quota
		HideFromUserUI: true, // 系统托管令牌,不在用户令牌列表展示
	}
	if err := tok.Insert(); err != nil {
		return "", err
	}
	return "sk-" + rawKey, nil
}

// provisionSwarmModelKey 确保用户的 swarm 模型 key 已就位(token + KV 密文),返回 azkv:// secret_ref。
// 仅在 Azure Key Vault 已配置时可用;未配置/不可达返回 error,由调用方按需降级
// (联调前 secret_ref 可留空,不阻断 swarm create)。
func provisionSwarmModelKey(userID int) (string, error) {
	store, err := newSecretStoreClientFromEnv()
	if err != nil {
		return "", err
	}
	bearer, err := getOrMintSwarmModelToken(userID)
	if err != nil {
		return "", err
	}
	// A.2:KV value = JSON {"openai_api_key":"sk-..."}。putSecret 幂等(同名新建版本)。
	secretRef, err := store.putSecret(swarmModelKeySecretName(userID), map[string]any{
		"openai_api_key": bearer,
	})
	if err != nil {
		return "", err
	}
	return secretRef, nil
}

// revokeSwarmModelKey 吊销用户的 swarm 模型 key:删 KV 密文 + 软删 token(均尽力而为)。
// 由 swarm.pool_terminated 回调触发(A.5)。token 软删后,下次 swarm create 会重新 mint 一把新 key。
func revokeSwarmModelKey(userID int) {
	if userID <= 0 {
		return
	}
	if store, err := newSecretStoreClientFromEnv(); err == nil {
		if err := store.deleteSecret(store.secretRef(swarmModelKeySecretName(userID))); err != nil {
			common.SysLog("revokeSwarmModelKey delete KV secret: " + err.Error())
		}
	}
	if model.DB != nil {
		if err := model.DB.Where("user_id = ? AND name = ?", userID, swarmModelKeyTokenName(userID)).Delete(&model.Token{}).Error; err != nil {
			common.SysLog("revokeSwarmModelKey delete token: " + err.Error())
		}
	}
}