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 1/2] 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, From dd5f6801612802d0a50c067bfea463c4cd6d7cb7 Mon Sep 17 00:00:00 2001 From: mahongran <623893204@qq.com> Date: Fri, 3 Jan 2025 18:22:44 +0800 Subject: [PATCH 2/2] refactor: enhance WorkAccessToken to improve cache key handling - Updated the AgentID field in WorkAccessToken struct to clarify its optional nature for distinguishing applications. - Modified the access token cache key construction to support both new and legacy formats based on the presence of AgentID. - Added comments for better understanding of the cache key logic and its compatibility with historical versions. This change improves the flexibility and clarity of access token management in multi-agent scenarios. --- credential/default_access_token.go | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/credential/default_access_token.go b/credential/default_access_token.go index 9f5ba82..39e055b 100644 --- a/credential/default_access_token.go +++ b/credential/default_access_token.go @@ -189,7 +189,7 @@ func (ak *StableAccessToken) GetAccessTokenDirectly(ctx context.Context, forceRe type WorkAccessToken struct { CorpID string CorpSecret string - AgentID string + AgentID string // 可选，用于区分不同应用 cacheKeyPrefix string cache cache.Cache accessTokenLock *sync.Mutex @@ -203,7 +203,7 @@ func NewWorkAccessToken(corpID, corpSecret, agentID, cacheKeyPrefix string, cach return &WorkAccessToken{ CorpID: corpID, CorpSecret: corpSecret, - AgentID: agentID, + AgentID: agentID, // agentID可以为空，兼容历史版本 cache: cache, cacheKeyPrefix: cacheKeyPrefix, accessTokenLock: new(sync.Mutex), @@ -221,8 +221,18 @@ func (ak *WorkAccessToken) GetAccessTokenContext(ctx context.Context) (accessTok ak.accessTokenLock.Lock() defer ak.accessTokenLock.Unlock() - // 修改缓存key，加入agentID - accessTokenCacheKey := fmt.Sprintf("%s_access_token_%s_%s", ak.cacheKeyPrefix, ak.CorpID, ak.AgentID) + // 构建缓存key + var accessTokenCacheKey string + // 每个应用有独立的secret，获取到的access_token只能本应用使用，所以每个应用的access_token应该分开来获取 + // API文档:https://developer.work.weixin.qq.com/document/path/91039 + if ak.AgentID != "" { + // 如果设置了AgentID，使用新的key格式 + accessTokenCacheKey = fmt.Sprintf("%s_access_token_%s_%s", ak.cacheKeyPrefix, ak.CorpID, ak.AgentID) + } else { + // 兼容历史版本的key格式 + accessTokenCacheKey = fmt.Sprintf("%s_access_token_%s", ak.cacheKeyPrefix, ak.CorpID) + } + val := ak.cache.Get(accessTokenCacheKey) if val != nil { accessToken = val.(string)