From 988ea49a033cb7aa839c57b196526e848feb51a3 Mon Sep 17 00:00:00 2001
From: mahongran <623893204@qq.com>
Date: Fri, 3 Jan 2025 14:13:06 +0800
Subject: [PATCH] feat: enhance WorkAccessToken to include AgentID for improved
 token management

- Added AgentID field to WorkAccessToken struct.
- Updated NewWorkAccessToken function to accept AgentID as a parameter.
- Modified access token cache key to incorporate AgentID, ensuring unique cache entries per agent.

This change improves the handling of access tokens in a multi-agent environment.
---
 credential/default_access_token.go | 8 ++++++--
 work/kf/client.go                  | 2 +-
 work/work.go                       | 2 +-
 3 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/credential/default_access_token.go b/credential/default_access_token.go
index ec88a89..9f5ba82 100644
--- a/credential/default_access_token.go
+++ b/credential/default_access_token.go
@@ -189,19 +189,21 @@ func (ak *StableAccessToken) GetAccessTokenDirectly(ctx context.Context, forceRe
 type WorkAccessToken struct {
 	CorpID          string
 	CorpSecret      string
+	AgentID         string
 	cacheKeyPrefix  string
 	cache           cache.Cache
 	accessTokenLock *sync.Mutex
 }
 
 // NewWorkAccessToken new WorkAccessToken
-func NewWorkAccessToken(corpID, corpSecret, cacheKeyPrefix string, cache cache.Cache) AccessTokenContextHandle {
+func NewWorkAccessToken(corpID, corpSecret, agentID, cacheKeyPrefix string, cache cache.Cache) AccessTokenContextHandle {
 	if cache == nil {
 		panic("cache the not exist")
 	}
 	return &WorkAccessToken{
 		CorpID:          corpID,
 		CorpSecret:      corpSecret,
+		AgentID:         agentID,
 		cache:           cache,
 		cacheKeyPrefix:  cacheKeyPrefix,
 		accessTokenLock: new(sync.Mutex),
@@ -218,7 +220,9 @@ func (ak *WorkAccessToken) GetAccessTokenContext(ctx context.Context) (accessTok
 	// 加上lock，是为了防止在并发获取token时，cache刚好失效，导致从微信服务器上获取到不同token
 	ak.accessTokenLock.Lock()
 	defer ak.accessTokenLock.Unlock()
-	accessTokenCacheKey := fmt.Sprintf("%s_access_token_%s", ak.cacheKeyPrefix, ak.CorpID)
+
+	// 修改缓存key，加入agentID
+	accessTokenCacheKey := fmt.Sprintf("%s_access_token_%s_%s", ak.cacheKeyPrefix, ak.CorpID, ak.AgentID)
 	val := ak.cache.Get(accessTokenCacheKey)
 	if val != nil {
 		accessToken = val.(string)
diff --git a/work/kf/client.go b/work/kf/client.go
index d83a5fc..3358198 100644
--- a/work/kf/client.go
+++ b/work/kf/client.go
@@ -24,7 +24,7 @@ func NewClient(cfg *config.Config) (client *Client, err error) {
 	}
 
 	// 初始化 AccessToken Handle
-	defaultAkHandle := credential.NewWorkAccessToken(cfg.CorpID, cfg.CorpSecret, credential.CacheKeyWorkPrefix, cfg.Cache)
+	defaultAkHandle := credential.NewWorkAccessToken(cfg.CorpID, cfg.CorpSecret, cfg.AgentID, credential.CacheKeyWorkPrefix, cfg.Cache)
 	ctx := &context.Context{
 		Config:            cfg,
 		AccessTokenHandle: defaultAkHandle,
diff --git a/work/work.go b/work/work.go
index 24c5773..69a2bae 100644
--- a/work/work.go
+++ b/work/work.go
@@ -24,7 +24,7 @@ type Work struct {
 
 // NewWork init work
 func NewWork(cfg *config.Config) *Work {
-	defaultAkHandle := credential.NewWorkAccessToken(cfg.CorpID, cfg.CorpSecret, credential.CacheKeyWorkPrefix, cfg.Cache)
+	defaultAkHandle := credential.NewWorkAccessToken(cfg.CorpID, cfg.CorpSecret, cfg.AgentID, credential.CacheKeyWorkPrefix, cfg.Cache)
 	ctx := &context.Context{
 		Config:            cfg,
 		AccessTokenHandle: defaultAkHandle,