From e217698d1035cb435defc48292fde2cab1c8b511 Mon Sep 17 00:00:00 2001
From: mahongran <623893204@qq.com>
Date: Tue, 7 Jan 2025 17:53:46 +0800
Subject: [PATCH] feat: enhance WorkAccessToken with new constructor for
 AgentID support

- Introduced NewWorkAccessTokenWithAgentID function to maintain backward compatibility while allowing for AgentID usage.
- Updated NewWorkAccessToken to call the new constructor, ensuring seamless integration.
- Improved error handling in GetAccessTokenContext by checking for cache availability and handling potential errors during cache operations.

This change enhances the flexibility of access token management, particularly in multi-agent scenarios, while ensuring compatibility with existing implementations.
---
 credential/default_access_token.go | 17 ++++++++++++-----
 work/config/config.go              | 11 +++++------
 2 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/credential/default_access_token.go b/credential/default_access_token.go
index 39e055b..70e2971 100644
--- a/credential/default_access_token.go
+++ b/credential/default_access_token.go
@@ -195,15 +195,21 @@ type WorkAccessToken struct {
 	accessTokenLock *sync.Mutex
 }
 
-// NewWorkAccessToken new WorkAccessToken
+// NewWorkAccessToken new WorkAccessToken (保持向后兼容)
 func NewWorkAccessToken(corpID, corpSecret, agentID, cacheKeyPrefix string, cache cache.Cache) AccessTokenContextHandle {
+	// 调用新方法，保持兼容性
+	return NewWorkAccessTokenWithAgentID(corpID, corpSecret, agentID, cacheKeyPrefix, cache)
+}
+
+// NewWorkAccessTokenWithAgentID new WorkAccessToken with agentID
+func NewWorkAccessTokenWithAgentID(corpID, corpSecret, agentID, cacheKeyPrefix string, cache cache.Cache) AccessTokenContextHandle {
 	if cache == nil {
-		panic("cache the not exist")
+		panic("cache is needed")
 	}
 	return &WorkAccessToken{
 		CorpID:          corpID,
 		CorpSecret:      corpSecret,
-		AgentID:         agentID, // agentID可以为空，兼容历史版本
+		AgentID:         agentID,
 		cache:           cache,
 		cacheKeyPrefix:  cacheKeyPrefix,
 		accessTokenLock: new(sync.Mutex),
@@ -223,8 +229,6 @@ func (ak *WorkAccessToken) GetAccessTokenContext(ctx context.Context) (accessTok
 
 	// 构建缓存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)
@@ -248,6 +252,9 @@ func (ak *WorkAccessToken) GetAccessTokenContext(ctx context.Context) (accessTok
 
 	expires := resAccessToken.ExpiresIn - 1500
 	err = ak.cache.Set(accessTokenCacheKey, resAccessToken.AccessToken, time.Duration(expires)*time.Second)
+	if err != nil {
+		return
+	}
 
 	accessToken = resAccessToken.AccessToken
 	return
diff --git a/work/config/config.go b/work/config/config.go
index 84aef7c..ac15e0a 100644
--- a/work/config/config.go
+++ b/work/config/config.go
@@ -7,12 +7,11 @@ import (
 
 // Config for 企业微信
 type Config struct {
-	CorpID        string `json:"corp_id"`     // corp_id
-	CorpSecret    string `json:"corp_secret"` // corp_secret,如果需要获取会话存档实例，当前参数请填写聊天内容存档的Secret，可以在企业微信管理端--管理工具--聊天内容存档查看
-	AgentID       string `json:"agent_id"`    // agent_id
-	Cache         cache.Cache
-	RasPrivateKey string // 消息加密私钥，可以在企业微信管理端--管理工具--消息加密公钥查看对用公钥，私钥一般由自己保存
-
+	CorpID         string `json:"corp_id"`     // corp_id
+	CorpSecret     string `json:"corp_secret"` // corp_secret,如果需要获取会话存档实例，当前参数请填写聊天内容存档的Secret，可以在企业微信管理端--管理工具--聊天内容存档查看
+	AgentID        string `json:"agent_id"`    // agent_id
+	Cache          cache.Cache
+	RasPrivateKey  string // 消息加密私钥，可以在企业微信管理端--管理工具--消息加密公钥查看对用公钥，私钥一般由自己保存
 	Token          string `json:"token"`            // 微信客服回调配置，用于生成签名校验回调请求的合法性
 	EncodingAESKey string `json:"encoding_aes_key"` // 微信客服回调p配置，用于解密回调消息内容对应的密文
 }