Skip to content

Commit c9912fe

Browse files
DumbrisDumbris
authored
Merge pull request #109 from smart-mcp-proxy/feature/update-mcp-go
Upgrade mcp-go to v0.42.0
2 parents 9b0c035 + 03f8120 commit c9912fe

File tree

7 files changed

+45
-26
lines changed

7 files changed

+45
-26
lines changed

‎go.mod‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ require (
88
github.com/blevesearch/bleve/v2 v2.5.2
99
github.com/go-chi/chi/v5 v5.2.3
1010
github.com/inconshreveable/go-update v0.0.0-20160112193335-8152e7eb6ccf
11-
github.com/mark3labs/mcp-go v0.38.0
11+
github.com/mark3labs/mcp-go v0.42.0
1212
github.com/pkoukk/tiktoken-go v0.1.8
1313
github.com/prometheus/client_golang v1.23.2
1414
github.com/spf13/cobra v1.9.1

‎go.sum‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,8 @@ github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0
112112
github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc=
113113
github.com/mark3labs/mcp-go v0.38.0 h1:E5tmJiIXkhwlV0pLAwAT0O5ZjUZSISE/2Jxg+6vpq4I=
114114
github.com/mark3labs/mcp-go v0.38.0/go.mod h1:T7tUa2jO6MavG+3P25Oy/jR7iCeJPHImCZHRymCn39g=
115+
github.com/mark3labs/mcp-go v0.42.0 h1:gk/8nYJh8t3yroCAOBhNbYsM9TCKvkM13I5t5Hfu6Ls=
116+
github.com/mark3labs/mcp-go v0.42.0/go.mod h1:YnJfOL382MIWDx1kMY+2zsRHU/q78dBg9aFb8W6Thdw=
115117
github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
116118
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg=
117119
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=

‎internal/appctx/adapters.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ func (o *OAuthTokenManagerImpl) GetToken(serverName string) (interface{}, error)
5757

5858
// Use persistent token store to get token
5959
tokenStore := oauth.NewPersistentTokenStore(serverName, "", o.storage)
60-
return tokenStore.GetToken()
60+
return tokenStore.GetToken(context.Background())
6161
}
6262

6363
// SaveToken saves a token for the given server

‎internal/oauth/config.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -257,7 +257,7 @@ func CreateOAuthConfig(serverConfig *config.ServerConfig, storage *storage.BoltD
257257
zap.String("storage", "BBolt database"))
258258

259259
// Check if token exists in persistent storage
260-
existingToken, err := tokenStore.GetToken()
260+
existingToken, err := tokenStore.GetToken(context.Background())
261261
if err != nil {
262262
logger.Info("🔍 No existing token found in persistent storage",
263263
zap.String("server", serverConfig.Name),

‎internal/oauth/persistent_token_store.go‎

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package oauth
22

33
import (
4+
"context"
45
"crypto/sha256"
56
"encoding/hex"
67
"fmt"
@@ -10,6 +11,7 @@ import (
1011
"mcpproxy-go/internal/storage"
1112

1213
"github.com/mark3labs/mcp-go/client"
14+
transport "github.com/mark3labs/mcp-go/client/transport"
1315
"go.uber.org/zap"
1416
)
1517

@@ -46,7 +48,14 @@ func generateServerKey(serverName, serverURL string) string {
4648
}
4749

4850
// GetToken retrieves the OAuth token from persistent storage
49-
func (p *PersistentTokenStore) GetToken() (*client.Token, error) {
51+
func (p *PersistentTokenStore) GetToken(ctx context.Context) (*client.Token, error) {
52+
if ctx == nil {
53+
ctx = context.Background()
54+
}
55+
if err := ctx.Err(); err != nil {
56+
return nil, err
57+
}
58+
5059
p.logger.Info("🔍 Loading OAuth token from persistent storage",
5160
zap.String("server_key", p.serverKey))
5261

@@ -55,7 +64,7 @@ func (p *PersistentTokenStore) GetToken() (*client.Token, error) {
5564
p.logger.Info("❌ No stored OAuth token found",
5665
zap.String("server_key", p.serverKey),
5766
zap.Error(err))
58-
return nil, err
67+
return nil, transport.ErrNoToken
5968
}
6069

6170
// Check if token is expired
@@ -88,7 +97,14 @@ func (p *PersistentTokenStore) GetToken() (*client.Token, error) {
8897
}
8998

9099
// SaveToken stores the OAuth token to persistent storage
91-
func (p *PersistentTokenStore) SaveToken(token *client.Token) error {
100+
func (p *PersistentTokenStore) SaveToken(ctx context.Context, token *client.Token) error {
101+
if ctx == nil {
102+
ctx = context.Background()
103+
}
104+
if err := ctx.Err(); err != nil {
105+
return err
106+
}
107+
92108
p.logger.Info("💾 Saving OAuth token to persistent storage",
93109
zap.String("token_type", token.TokenType),
94110
zap.Time("expires_at", token.ExpiresAt))

‎internal/oauth/persistent_token_store_test.go‎

Lines changed: 19 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package oauth
22

33
import (
4+
"context"
45
"testing"
56
"time"
67

@@ -28,7 +29,7 @@ func TestPersistentTokenStore(t *testing.T) {
2829
tokenStore := NewPersistentTokenStore("test-server", "https://test.example.com/mcp", db)
2930

3031
// Test case 1: Get non-existent token
31-
_, err = tokenStore.GetToken()
32+
_, err = tokenStore.GetToken(context.Background())
3233
if err == nil {
3334
t.Error("Expected error when getting non-existent token")
3435
}
@@ -42,13 +43,13 @@ func TestPersistentTokenStore(t *testing.T) {
4243
Scope: "mcp.read mcp.write",
4344
}
4445

45-
err = tokenStore.SaveToken(originalToken)
46+
err = tokenStore.SaveToken(context.Background(), originalToken)
4647
if err != nil {
4748
t.Fatalf("Failed to save token: %v", err)
4849
}
4950

5051
// Test case 3: Retrieve the saved token
51-
retrievedToken, err := tokenStore.GetToken()
52+
retrievedToken, err := tokenStore.GetToken(context.Background())
5253
if err != nil {
5354
t.Fatalf("Failed to get token: %v", err)
5455
}
@@ -79,12 +80,12 @@ func TestPersistentTokenStore(t *testing.T) {
7980
Scope: "mcp.read mcp.write admin",
8081
}
8182

82-
err = tokenStore.SaveToken(updatedToken)
83+
err = tokenStore.SaveToken(context.Background(), updatedToken)
8384
if err != nil {
8485
t.Fatalf("Failed to save updated token: %v", err)
8586
}
8687

87-
retrievedUpdatedToken, err := tokenStore.GetToken()
88+
retrievedUpdatedToken, err := tokenStore.GetToken(context.Background())
8889
if err != nil {
8990
t.Fatalf("Failed to get updated token: %v", err)
9091
}
@@ -100,7 +101,7 @@ func TestPersistentTokenStore(t *testing.T) {
100101
t.Fatalf("Failed to clear token: %v", err)
101102
}
102103

103-
_, err = tokenStore.GetToken()
104+
_, err = tokenStore.GetToken(context.Background())
104105
if err == nil {
105106
t.Error("Expected error when getting cleared token")
106107
}
@@ -114,12 +115,12 @@ func TestPersistentTokenStore(t *testing.T) {
114115
Scope: "mcp.read",
115116
}
116117

117-
err = tokenStore.SaveToken(expiredToken)
118+
err = tokenStore.SaveToken(context.Background(), expiredToken)
118119
if err != nil {
119120
t.Fatalf("Failed to save expired token: %v", err)
120121
}
121122

122-
retrievedExpiredToken, err := tokenStore.GetToken()
123+
retrievedExpiredToken, err := tokenStore.GetToken(context.Background())
123124
if err != nil {
124125
t.Fatalf("Failed to get expired token: %v", err)
125126
}
@@ -162,23 +163,23 @@ func TestPersistentTokenStoreMultipleServers(t *testing.T) {
162163
Scope: "mcp.write",
163164
}
164165

165-
err = tokenStore1.SaveToken(token1)
166+
err = tokenStore1.SaveToken(context.Background(), token1)
166167
if err != nil {
167168
t.Fatalf("Failed to save token1: %v", err)
168169
}
169170

170-
err = tokenStore2.SaveToken(token2)
171+
err = tokenStore2.SaveToken(context.Background(), token2)
171172
if err != nil {
172173
t.Fatalf("Failed to save token2: %v", err)
173174
}
174175

175176
// Retrieve tokens and verify they are separate
176-
retrievedToken1, err := tokenStore1.GetToken()
177+
retrievedToken1, err := tokenStore1.GetToken(context.Background())
177178
if err != nil {
178179
t.Fatalf("Failed to get token1: %v", err)
179180
}
180181

181-
retrievedToken2, err := tokenStore2.GetToken()
182+
retrievedToken2, err := tokenStore2.GetToken(context.Background())
182183
if err != nil {
183184
t.Fatalf("Failed to get token2: %v", err)
184185
}
@@ -233,23 +234,23 @@ func TestPersistentTokenStoreSameNameDifferentURL(t *testing.T) {
233234
Scope: "mcp.write",
234235
}
235236

236-
err = tokenStore1.SaveToken(token1)
237+
err = tokenStore1.SaveToken(context.Background(), token1)
237238
if err != nil {
238239
t.Fatalf("Failed to save token1: %v", err)
239240
}
240241

241-
err = tokenStore2.SaveToken(token2)
242+
err = tokenStore2.SaveToken(context.Background(), token2)
242243
if err != nil {
243244
t.Fatalf("Failed to save token2: %v", err)
244245
}
245246

246247
// Retrieve tokens and verify they are separate despite same server name
247-
retrievedToken1, err := tokenStore1.GetToken()
248+
retrievedToken1, err := tokenStore1.GetToken(context.Background())
248249
if err != nil {
249250
t.Fatalf("Failed to get token1: %v", err)
250251
}
251252

252-
retrievedToken2, err := tokenStore2.GetToken()
253+
retrievedToken2, err := tokenStore2.GetToken(context.Background())
253254
if err != nil {
254255
t.Fatalf("Failed to get token2: %v", err)
255256
}
@@ -279,13 +280,13 @@ func TestPersistentTokenStoreSameNameDifferentURL(t *testing.T) {
279280
}
280281

281282
// Token 1 should be gone
282-
_, err = tokenStore1.GetToken()
283+
_, err = tokenStore1.GetToken(context.Background())
283284
if err == nil {
284285
t.Error("Expected error when getting cleared token1")
285286
}
286287

287288
// Token 2 should still exist
288-
retrievedToken2Again, err := tokenStore2.GetToken()
289+
retrievedToken2Again, err := tokenStore2.GetToken(context.Background())
289290
if err != nil {
290291
t.Fatalf("Failed to get token2 after clearing token1: %v", err)
291292
}

‎internal/upstream/manager.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -796,7 +796,7 @@ func (m *Manager) RetryConnection(serverName string) error {
796796
var tokenExpires time.Time
797797
if m.storage != nil {
798798
ts := oauth.NewPersistentTokenStore(client.Config.Name, client.Config.URL, m.storage)
799-
if tok, err := ts.GetToken(); err == nil && tok != nil {
799+
if tok, err := ts.GetToken(context.Background()); err == nil && tok != nil {
800800
hasToken = true
801801
tokenExpires = tok.ExpiresAt
802802
}
@@ -961,7 +961,7 @@ func (m *Manager) scanForNewTokens() {
961961

962962
// Check for a persisted token
963963
ts := oauth.NewPersistentTokenStore(cfg.Name, cfg.URL, m.storage)
964-
tok, err := ts.GetToken()
964+
tok, err := ts.GetToken(context.Background())
965965
if err != nil || tok == nil {
966966
continue
967967
}

0 commit comments

Comments
 (0)