Skip to content

Commit 7473d1f

Browse files
Ambient Code Botclaude
andcommitted
feat(mcp-sidecar): implement RSA-OAEP token exchange for dynamic token refresh
Replace the static AMBIENT_TOKEN injection with the same RSA-OAEP token exchange protocol the runner uses, so the MCP sidecar can fetch and periodically refresh its own API token from the CP token server. - Add tokenexchange package: RSA-OAEP encrypt session ID, fetch token from CP /token endpoint with retry/backoff, background refresh every 5 minutes - Make client.Client thread-safe (sync.RWMutex) with SetToken() for hot-swapping tokens on refresh - Refactor mention.Resolver to use TokenFunc instead of static string so it always reads the current token - Inject SESSION_ID into MCP sidecar env in buildMCPSidecar() - Default MCP_API_SERVER_URL to AMBIENT_API_SERVER_URL so the sidecar inherits the correct API endpoint automatically - Remove hardcoded MCP_API_SERVER_URL from mpp-openshift overlay - Add unit tests for tokenexchange, client, and mention packages 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
1 parent ad938e2 commit 7473d1f

File tree

12 files changed

+990
-36
lines changed

12 files changed

+990
-36
lines changed

components/ambient-cli/cmd/acpctl/session/send.go

Lines changed: 14 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"context"
66
"encoding/json"
77
"fmt"
8+
"io"
89
"os"
910
"os/signal"
1011
"strings"
@@ -56,6 +57,19 @@ func runSend(cmd *cobra.Command, args []string) error {
5657
ctx, cancel := context.WithTimeout(cmd.Context(), cfg.GetRequestTimeout())
5758
defer cancel()
5859

60+
streamCtx, streamCancel := signal.NotifyContext(cmd.Context(), os.Interrupt)
61+
defer streamCancel()
62+
63+
var stream io.ReadCloser
64+
if sendFollow {
65+
s, err := client.Sessions().StreamEvents(streamCtx, sessionID)
66+
if err != nil {
67+
return fmt.Errorf("stream events: %w", err)
68+
}
69+
stream = s
70+
defer stream.Close()
71+
}
72+
5973
msg, err := client.Sessions().PushMessage(ctx, sessionID, payload)
6074
if err != nil {
6175
return fmt.Errorf("send message: %w", err)
@@ -67,15 +81,6 @@ func runSend(cmd *cobra.Command, args []string) error {
6781
return nil
6882
}
6983

70-
streamCtx, streamCancel := signal.NotifyContext(cmd.Context(), os.Interrupt)
71-
defer streamCancel()
72-
73-
stream, err := client.Sessions().StreamEvents(streamCtx, sessionID)
74-
if err != nil {
75-
return fmt.Errorf("stream events: %w", err)
76-
}
77-
defer stream.Close()
78-
7984
out := cmd.OutOrStdout()
8085
scanner := bufio.NewScanner(stream)
8186
var reasoningBuf strings.Builder

components/ambient-control-plane/internal/config/config.go

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,13 +70,17 @@ func Load() (*ControlPlaneConfig, error) {
7070
VertexSecretNamespace: envOrDefault("VERTEX_SECRET_NAMESPACE", "ambient-code"),
7171
RunnerImageNamespace: os.Getenv("RUNNER_IMAGE_NAMESPACE"),
7272
MCPImage: os.Getenv("MCP_IMAGE"),
73-
MCPAPIServerURL: envOrDefault("MCP_API_SERVER_URL", "http://ambient-api-server.ambient-code.svc:8000"),
73+
MCPAPIServerURL: envOrDefault("MCP_API_SERVER_URL", ""),
7474
RunnerLogLevel: envOrDefault("RUNNER_LOG_LEVEL", "info"),
7575
ProjectKubeTokenFile: os.Getenv("PROJECT_KUBE_TOKEN_FILE"),
7676
CPTokenListenAddr: envOrDefault("CP_TOKEN_LISTEN_ADDR", ":8080"),
7777
CPTokenURL: os.Getenv("CP_TOKEN_URL"),
7878
}
7979

80+
if cfg.MCPAPIServerURL == "" {
81+
cfg.MCPAPIServerURL = cfg.APIServerURL
82+
}
83+
8084
if cfg.APIToken == "" && (cfg.OIDCClientID == "" || cfg.OIDCClientSecret == "") {
8185
return nil, fmt.Errorf("either AMBIENT_API_TOKEN or both OIDC_CLIENT_ID and OIDC_CLIENT_SECRET must be set; set AMBIENT_API_TOKEN for k8s SA token auth or OIDC_CLIENT_ID+OIDC_CLIENT_SECRET for OIDC")
8286
}

components/ambient-control-plane/internal/reconciler/kube_reconciler.go

Lines changed: 3 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -441,12 +441,7 @@ func (r *SimpleKubeReconciler) ensurePod(ctx context.Context, namespace string,
441441
}
442442

443443
if useMCPSidecar {
444-
ambientToken, err := r.factory.Token(ctx)
445-
if err != nil {
446-
r.logger.Warn().Err(err).Str("session_id", session.ID).Msg("failed to fetch token for MCP sidecar; sidecar will start without AMBIENT_TOKEN")
447-
ambientToken = ""
448-
}
449-
containers = append(containers, r.buildMCPSidecar(ambientToken))
444+
containers = append(containers, r.buildMCPSidecar(session.ID))
450445
r.logger.Info().Str("session_id", session.ID).Msg("MCP sidecar enabled for session")
451446
}
452447

@@ -816,7 +811,7 @@ func boolToStr(b bool) string {
816811
return "false"
817812
}
818813

819-
func (r *SimpleKubeReconciler) buildMCPSidecar(ambientToken string) interface{} {
814+
func (r *SimpleKubeReconciler) buildMCPSidecar(sessionID string) interface{} {
820815
mcpImage := r.cfg.MCPImage
821816
imagePullPolicy := "Always"
822817
if strings.HasPrefix(mcpImage, "localhost/") {
@@ -828,9 +823,7 @@ func (r *SimpleKubeReconciler) buildMCPSidecar(ambientToken string) interface{}
828823
envVar("AMBIENT_API_URL", r.cfg.MCPAPIServerURL),
829824
envVar("AMBIENT_CP_TOKEN_URL", r.cfg.CPTokenURL),
830825
envVar("AMBIENT_CP_TOKEN_PUBLIC_KEY", r.cfg.CPTokenPublicKey),
831-
}
832-
if ambientToken != "" {
833-
env = append(env, envVar("AMBIENT_TOKEN", ambientToken))
826+
envVar("SESSION_ID", sessionID),
834827
}
835828
return map[string]interface{}{
836829
"name": "ambient-mcp",

components/ambient-mcp/client/client.go

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,12 +9,14 @@ import (
99
"net/http"
1010
"net/url"
1111
"strings"
12+
"sync"
1213
"time"
1314
)
1415

1516
type Client struct {
1617
httpClient *http.Client
1718
baseURL string
19+
mu sync.RWMutex
1820
token string
1921
}
2022

@@ -27,7 +29,18 @@ func New(baseURL, token string) *Client {
2729
}
2830

2931
func (c *Client) BaseURL() string { return c.baseURL }
30-
func (c *Client) Token() string { return c.token }
32+
33+
func (c *Client) Token() string {
34+
c.mu.RLock()
35+
defer c.mu.RUnlock()
36+
return c.token
37+
}
38+
39+
func (c *Client) SetToken(token string) {
40+
c.mu.Lock()
41+
defer c.mu.Unlock()
42+
c.token = token
43+
}
3144

3245
func (c *Client) do(ctx context.Context, method, path string, body []byte, result interface{}, expectedStatuses ...int) error {
3346
reqURL := c.baseURL + "/api/ambient/v1" + path
@@ -42,7 +55,7 @@ func (c *Client) do(ctx context.Context, method, path string, body []byte, resul
4255
if body != nil {
4356
req.Header.Set("Content-Type", "application/json")
4457
}
45-
req.Header.Set("Authorization", "Bearer "+c.token)
58+
req.Header.Set("Authorization", "Bearer "+c.Token())
4659
req.Header.Set("Accept", "application/json")
4760

4861
resp, err := c.httpClient.Do(req)
Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
package client
2+
3+
import (
4+
"context"
5+
"encoding/json"
6+
"net/http"
7+
"net/http/httptest"
8+
"sync"
9+
"testing"
10+
)
11+
12+
func TestNew(t *testing.T) {
13+
c := New("http://localhost:8080/", "my-token")
14+
if c.BaseURL() != "http://localhost:8080" {
15+
t.Errorf("BaseURL() = %q, want trailing slash stripped", c.BaseURL())
16+
}
17+
if c.Token() != "my-token" {
18+
t.Errorf("Token() = %q, want %q", c.Token(), "my-token")
19+
}
20+
}
21+
22+
func TestSetToken(t *testing.T) {
23+
c := New("http://localhost:8080", "initial")
24+
c.SetToken("refreshed")
25+
if c.Token() != "refreshed" {
26+
t.Errorf("Token() after SetToken = %q, want %q", c.Token(), "refreshed")
27+
}
28+
}
29+
30+
func TestSetToken_ConcurrentAccess(t *testing.T) {
31+
c := New("http://localhost:8080", "initial")
32+
var wg sync.WaitGroup
33+
34+
for i := range 100 {
35+
wg.Add(2)
36+
go func(n int) {
37+
defer wg.Done()
38+
c.SetToken("token-" + string(rune('A'+n%26)))
39+
}(i)
40+
go func() {
41+
defer wg.Done()
42+
_ = c.Token()
43+
}()
44+
}
45+
wg.Wait()
46+
47+
got := c.Token()
48+
if got == "" {
49+
t.Error("Token() is empty after concurrent access")
50+
}
51+
}
52+
53+
func TestGet_SendsBearerToken(t *testing.T) {
54+
var receivedAuth string
55+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
56+
receivedAuth = r.Header.Get("Authorization")
57+
w.Header().Set("Content-Type", "application/json")
58+
json.NewEncoder(w).Encode(map[string]string{"status": "ok"})
59+
}))
60+
defer srv.Close()
61+
62+
c := New(srv.URL, "test-bearer")
63+
var result map[string]string
64+
err := c.Get(context.Background(), "/healthz", &result)
65+
if err != nil {
66+
t.Fatalf("Get: %v", err)
67+
}
68+
if receivedAuth != "Bearer test-bearer" {
69+
t.Errorf("Authorization = %q, want %q", receivedAuth, "Bearer test-bearer")
70+
}
71+
}
72+
73+
func TestGet_UsesRefreshedToken(t *testing.T) {
74+
var receivedAuth string
75+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
76+
receivedAuth = r.Header.Get("Authorization")
77+
w.Header().Set("Content-Type", "application/json")
78+
json.NewEncoder(w).Encode(map[string]string{"status": "ok"})
79+
}))
80+
defer srv.Close()
81+
82+
c := New(srv.URL, "old-token")
83+
c.SetToken("new-token")
84+
85+
var result map[string]string
86+
err := c.Get(context.Background(), "/healthz", &result)
87+
if err != nil {
88+
t.Fatalf("Get: %v", err)
89+
}
90+
if receivedAuth != "Bearer new-token" {
91+
t.Errorf("Authorization = %q, want %q", receivedAuth, "Bearer new-token")
92+
}
93+
}
94+
95+
func TestGet_UnexpectedStatus(t *testing.T) {
96+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
97+
http.Error(w, "not found", http.StatusNotFound)
98+
}))
99+
defer srv.Close()
100+
101+
c := New(srv.URL, "token")
102+
err := c.Get(context.Background(), "/missing", nil)
103+
if err == nil {
104+
t.Fatal("expected error for 404 response")
105+
}
106+
}
107+
108+
func TestPost_SendsBody(t *testing.T) {
109+
var receivedBody map[string]string
110+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
111+
json.NewDecoder(r.Body).Decode(&receivedBody)
112+
if r.Header.Get("Content-Type") != "application/json" {
113+
t.Errorf("Content-Type = %q, want application/json", r.Header.Get("Content-Type"))
114+
}
115+
w.WriteHeader(http.StatusCreated)
116+
json.NewEncoder(w).Encode(map[string]string{"id": "123"})
117+
}))
118+
defer srv.Close()
119+
120+
c := New(srv.URL, "token")
121+
body := map[string]string{"name": "test"}
122+
var result map[string]string
123+
err := c.Post(context.Background(), "/items", body, &result, http.StatusCreated)
124+
if err != nil {
125+
t.Fatalf("Post: %v", err)
126+
}
127+
if receivedBody["name"] != "test" {
128+
t.Errorf("received body name = %q, want %q", receivedBody["name"], "test")
129+
}
130+
if result["id"] != "123" {
131+
t.Errorf("result id = %q, want %q", result["id"], "123")
132+
}
133+
}

components/ambient-mcp/main.go

Lines changed: 39 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88
"github.com/mark3labs/mcp-go/server"
99

1010
"github.com/ambient-code/platform/components/ambient-mcp/client"
11+
"github.com/ambient-code/platform/components/ambient-mcp/tokenexchange"
1112
)
1213

1314
func main() {
@@ -16,18 +17,50 @@ func main() {
1617
apiURL = "http://localhost:8080"
1718
}
1819

19-
token := os.Getenv("AMBIENT_TOKEN")
20-
if token == "" {
21-
fmt.Fprintln(os.Stderr, "AMBIENT_TOKEN is required")
22-
os.Exit(1)
23-
}
24-
2520
transport := os.Getenv("MCP_TRANSPORT")
2621
if transport == "" {
2722
transport = "stdio"
2823
}
2924

25+
cpTokenURL := os.Getenv("AMBIENT_CP_TOKEN_URL")
26+
cpPublicKey := os.Getenv("AMBIENT_CP_TOKEN_PUBLIC_KEY")
27+
sessionID := os.Getenv("SESSION_ID")
28+
29+
var token string
30+
var exchanger *tokenexchange.Exchanger
31+
32+
if cpTokenURL != "" && cpPublicKey != "" && sessionID != "" {
33+
var err error
34+
exchanger, err = tokenexchange.New(cpTokenURL, cpPublicKey, sessionID)
35+
if err != nil {
36+
fmt.Fprintf(os.Stderr, "token exchange init failed: %v\n", err)
37+
os.Exit(1)
38+
}
39+
token, err = exchanger.FetchToken()
40+
if err != nil {
41+
fmt.Fprintf(os.Stderr, "initial token fetch failed: %v\n", err)
42+
os.Exit(1)
43+
}
44+
fmt.Fprintln(os.Stderr, "bootstrapped token via CP token exchange")
45+
} else {
46+
token = os.Getenv("AMBIENT_TOKEN")
47+
if token == "" {
48+
fmt.Fprintln(os.Stderr, "AMBIENT_TOKEN is required when CP token exchange env vars are not set")
49+
os.Exit(1)
50+
}
51+
fmt.Fprintln(os.Stderr, "using static AMBIENT_TOKEN (no CP token exchange)")
52+
}
53+
3054
c := client.New(apiURL, token)
55+
56+
if exchanger != nil {
57+
exchanger.OnRefresh(func(freshToken string) {
58+
c.SetToken(freshToken)
59+
})
60+
exchanger.StartBackgroundRefresh()
61+
defer exchanger.Stop()
62+
}
63+
3164
s := newServer(c, transport)
3265

3366
switch transport {

components/ambient-mcp/mention/resolve.go

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -18,16 +18,18 @@ type Doer interface {
1818
Do(req *http.Request) (*http.Response, error)
1919
}
2020

21+
type TokenFunc func() string
22+
2123
type Resolver struct {
2224
baseURL string
23-
token string
25+
tokenFn TokenFunc
2426
http *http.Client
2527
}
2628

27-
func NewResolver(baseURL, token string) *Resolver {
29+
func NewResolver(baseURL string, tokenFn TokenFunc) *Resolver {
2830
return &Resolver{
2931
baseURL: strings.TrimSuffix(baseURL, "/"),
30-
token: token,
32+
tokenFn: tokenFn,
3133
http: &http.Client{},
3234
}
3335
}
@@ -43,7 +45,7 @@ func (r *Resolver) Resolve(ctx context.Context, projectID, identifier string) (s
4345
if uuidPattern.MatchString(strings.ToLower(identifier)) {
4446
path := r.baseURL + "/api/ambient/v1/projects/" + url.PathEscape(projectID) + "/agents/" + url.PathEscape(identifier)
4547
req, _ := http.NewRequestWithContext(ctx, http.MethodGet, path, nil)
46-
req.Header.Set("Authorization", "Bearer "+r.token)
48+
req.Header.Set("Authorization", "Bearer "+r.tokenFn())
4749
resp, err := r.http.Do(req)
4850
if err != nil {
4951
return "", fmt.Errorf("lookup agent by ID: %w", err)
@@ -66,7 +68,7 @@ func (r *Resolver) Resolve(ctx context.Context, projectID, identifier string) (s
6668

6769
path := r.baseURL + "/api/ambient/v1/projects/" + url.PathEscape(projectID) + "/agents?search=name='" + url.QueryEscape(identifier) + "'"
6870
req, _ := http.NewRequestWithContext(ctx, http.MethodGet, path, nil)
69-
req.Header.Set("Authorization", "Bearer "+r.token)
71+
req.Header.Set("Authorization", "Bearer "+r.tokenFn())
7072
resp, err := r.http.Do(req)
7173
if err != nil {
7274
return "", fmt.Errorf("search agent by name: %w", err)

0 commit comments

Comments
 (0)