Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: enhance WorkAccessToken to include AgentID for improved token m…
Browse files Browse the repository at this point in the history
…anagement

- 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.
mahongran committed Jan 3, 2025

Verified

This commit was signed with the committer’s verified signature.
Byron Sebastian Thiel
1 parent d38e750 commit 988ea49
Showing 3 changed files with 8 additions and 4 deletions.
8 changes: 6 additions & 2 deletions credential/default_access_token.go
Original file line number Diff line number Diff line change
@@ -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)
2 changes: 1 addition & 1 deletion work/kf/client.go
Original file line number Diff line number Diff line change
@@ -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,
2 changes: 1 addition & 1 deletion work/work.go
Original file line number Diff line number Diff line change
@@ -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,

0 comments on commit 988ea49

Please sign in to comment.