diff --git a/cmd/root.go b/cmd/root.go index a6f3a3e..886c73d 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -160,6 +160,11 @@ func init() { } cmd.SetContext(contextWithClient(ctx, c)) + // Automatic sandbox→live self-heal: if we're on a sandbox token and the + // account has gone live, promote it transparently before the command + // runs. Best-effort and throttled; never blocks the command. + maybePromoteSandboxToken(ctx, c, saToken, readOnly) + return nil } } diff --git a/cmd/selfheal.go b/cmd/selfheal.go new file mode 100644 index 0000000..8546158 --- /dev/null +++ b/cmd/selfheal.go @@ -0,0 +1,75 @@ +package cmd + +import ( + "context" + "encoding/json" + "fmt" + "os" + "time" + + "github.com/customerio/cli/internal/client" +) + +// sandboxPromoteThrottle bounds how often the CLI probes the promote endpoint +// while still on a sandbox token. Promotion only succeeds after go-live; until +// then each attempt is a no-op, so we space them out. The post-go-live grace +// window is days, so up to an hour of heal latency is immaterial. +const sandboxPromoteThrottle = time.Hour + +// maybePromoteSandboxToken performs the automatic sandbox→live self-heal. +// +// When the CLI is authenticated with a stored sandbox token and the account has +// gone live, the server's POST /promote_sandbox_token swaps it for a live token +// (inheriting the sandbox token's permissions) and revokes the sandbox token. +// This persists the new token, points the in-memory client at it, and lets the +// user's command proceed on the live token — all transparently. +// +// It is best-effort: any failure is swallowed (only a throttle timestamp is +// recorded) so the user's command always runs. It acts only on a sandbox token +// the CLI itself persisted — never an env/flag token, which is not ours to +// rewrite — and never in read-only mode, where the POST would be blocked. +func maybePromoteSandboxToken(ctx context.Context, c *client.Client, saToken string, readOnly bool) { + if readOnly || !client.IsSandboxServiceAccountToken(saToken) { + return + } + + creds, err := client.ReadCredentials() + if err != nil { + return // token came from env/flag, not our config — leave it alone + } + if creds.ServiceAccountToken != saToken || creds.AccountID == "" { + return + } + if time.Since(creds.SandboxPromoteCheckedAt) < sandboxPromoteThrottle { + return + } + + path := fmt.Sprintf("/v1/accounts/%s/promote_sandbox_token", creds.AccountID) + resp, err := c.Do(ctx, "POST", path, nil, nil) + if err != nil { + // 403 (still in sandbox), 404 (already promoted), or transient — record + // the attempt so we don't probe again until the throttle elapses. + creds.SandboxPromoteCheckedAt = time.Now() + _ = client.WriteCredentials(creds) + return + } + + var promoted struct { + Token string `json:"token"` + } + if err := json.Unmarshal(resp, &promoted); err != nil || !client.IsServiceAccountToken(promoted.Token) { + creds.SandboxPromoteCheckedAt = time.Now() + _ = client.WriteCredentials(creds) + return + } + + creds.ServiceAccountToken = promoted.Token + creds.AccessToken = "" + creds.AccessTokenExpiresAt = time.Time{} + creds.SandboxPromoteCheckedAt = time.Time{} + if err := client.WriteCredentials(creds); err != nil { + return + } + c.SetServiceAccountToken(promoted.Token) + fmt.Fprintln(os.Stderr, "cio: account is live — upgraded the CLI to a live token.") +} diff --git a/cmd/selfheal_test.go b/cmd/selfheal_test.go new file mode 100644 index 0000000..c178bea --- /dev/null +++ b/cmd/selfheal_test.go @@ -0,0 +1,129 @@ +package cmd + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/customerio/cli/internal/client" +) + +// newClientFor builds a client pointed at srv with a pre-set access token so +// Do() skips the OAuth exchange (the bogus JWT gets a far-future expiry). +func newClientFor(url, saToken string) *client.Client { + return client.New(client.Config{ + BaseURL: url, + ServiceAccountToken: saToken, + AccessToken: "test-jwt", + }) +} + +func TestMaybePromoteSandboxToken_PromotesWhenLive(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + + var gotPath, gotMethod string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath, gotMethod = r.URL.Path, r.Method + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"token":"sa_live_promoted999","id":1,"name":"live-bootstrap"}`)) + })) + defer srv.Close() + + saToken := "sa_sandbox_bootstrap123" + if err := client.WriteCredentials(&client.Credentials{ + ServiceAccountToken: saToken, + AccountID: "42", + Region: "us", + }); err != nil { + t.Fatal(err) + } + + c := newClientFor(srv.URL, saToken) + maybePromoteSandboxToken(context.Background(), c, saToken, false) + + if gotMethod != http.MethodPost || gotPath != "/v1/accounts/42/promote_sandbox_token" { + t.Fatalf("unexpected request: %s %s", gotMethod, gotPath) + } + if c.ServiceAccountToken() != "sa_live_promoted999" { + t.Fatalf("in-memory client token not swapped: %q", c.ServiceAccountToken()) + } + creds, err := client.ReadCredentials() + if err != nil { + t.Fatal(err) + } + if creds.ServiceAccountToken != "sa_live_promoted999" { + t.Fatalf("stored token not swapped: %q", creds.ServiceAccountToken) + } + if creds.AccessToken != "" { + t.Fatal("cached access token should be cleared after promotion") + } +} + +func TestMaybePromoteSandboxToken_Throttled(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + + called := false + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { called = true })) + defer srv.Close() + + saToken := "sa_sandbox_x" + if err := client.WriteCredentials(&client.Credentials{ + ServiceAccountToken: saToken, + AccountID: "1", + SandboxPromoteCheckedAt: time.Now(), + }); err != nil { + t.Fatal(err) + } + + maybePromoteSandboxToken(context.Background(), newClientFor(srv.URL, saToken), saToken, false) + if called { + t.Fatal("should not probe promote within the throttle window") + } +} + +func TestMaybePromoteSandboxToken_SkipsLiveAndReadOnly(t *testing.T) { + called := false + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { called = true })) + defer srv.Close() + + // Live token: nothing to promote. + maybePromoteSandboxToken(context.Background(), newClientFor(srv.URL, "sa_live_x"), "sa_live_x", false) + // Sandbox token but read-only: the POST would be blocked, so skip. + maybePromoteSandboxToken(context.Background(), newClientFor(srv.URL, "sa_sandbox_x"), "sa_sandbox_x", true) + if called { + t.Fatal("should not call promote for a live token or in read-only mode") + } +} + +func TestMaybePromoteSandboxToken_SwallowsStillSandbox(t *testing.T) { + t.Setenv("HOME", t.TempDir()) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"errors":[{"detail":"still in sandbox"}]}`)) + })) + defer srv.Close() + + saToken := "sa_sandbox_y" + if err := client.WriteCredentials(&client.Credentials{ + ServiceAccountToken: saToken, + AccountID: "7", + }); err != nil { + t.Fatal(err) + } + + maybePromoteSandboxToken(context.Background(), newClientFor(srv.URL, saToken), saToken, false) + + creds, err := client.ReadCredentials() + if err != nil { + t.Fatal(err) + } + if creds.ServiceAccountToken != saToken { + t.Fatalf("token must be unchanged on 403, got %q", creds.ServiceAccountToken) + } + if creds.SandboxPromoteCheckedAt.IsZero() { + t.Fatal("throttle timestamp should be recorded after a failed attempt") + } +} diff --git a/internal/client/auth.go b/internal/client/auth.go index 06cc17e..24d55e9 100644 --- a/internal/client/auth.go +++ b/internal/client/auth.go @@ -47,6 +47,10 @@ type Credentials struct { ReadOnly bool `json:"read_only,omitempty"` // Scopes holds additional OAuth scope values that were requested. Scopes []string `json:"scopes,omitempty"` + // SandboxPromoteCheckedAt throttles the automatic sandbox→live promotion so + // the CLI does not probe the promote endpoint on every command while the + // account is still in sandbox. + SandboxPromoteCheckedAt time.Time `json:"sandbox_promote_checked_at,omitempty"` } // ScopeReadOnly is the OAuth 2.0 scope value that requests a read-only session. @@ -418,3 +422,10 @@ func IsServiceAccountToken(token string) bool { } return false } + +// IsSandboxServiceAccountToken reports whether the token is a Builder sandbox +// token (sa_sandbox_), which the CLI promotes to a live token once the account +// has gone live. +func IsSandboxServiceAccountToken(token string) bool { + return strings.HasPrefix(token, SandboxServiceAccountTokenPrefix) +} diff --git a/internal/client/auth_test.go b/internal/client/auth_test.go index 528de39..85e90f6 100644 --- a/internal/client/auth_test.go +++ b/internal/client/auth_test.go @@ -285,6 +285,18 @@ func TestIsServiceAccountToken(t *testing.T) { } } +func TestIsSandboxServiceAccountToken(t *testing.T) { + if !IsSandboxServiceAccountToken("sa_sandbox_abc123") { + t.Error("expected true for sa_sandbox_ prefix") + } + if IsSandboxServiceAccountToken("sa_live_abc123") { + t.Error("expected false for sa_live_ prefix") + } + if IsSandboxServiceAccountToken("") { + t.Error("expected false for empty string") + } +} + func TestBaseURLForRegion(t *testing.T) { tests := []struct { region string diff --git a/internal/client/client.go b/internal/client/client.go index 495af7c..8271cbf 100644 --- a/internal/client/client.go +++ b/internal/client/client.go @@ -27,12 +27,18 @@ import ( // - X-CIO-Agent: 1 — set only when the CIO_AGENT env var is "1". The // sandbox that runs the CLI on behalf of an AI agent sets this so // downstream metrics can attribute traffic to the agent. +// - X-CIO-Source — identifies the originating tool for server-side analytics +// (e.g. the "Design Studio: Email Created" source). "AI Assistant" when run +// by the agent sandbox (CIO_AGENT=1), otherwise "CLI". // - X-CIO-Capability-Grant — forwarded from $X_CIO_CAPABILITY_GRANT so the env carries a session-scoped grant without a CLI flag. func setStandardHeaders(req *http.Request) { req.Header.Set("User-Agent", useragent.Get()) req.Header.Set("X-Validate", "strict") if os.Getenv("CIO_AGENT") == "1" { req.Header.Set("X-CIO-Agent", "1") + req.Header.Set("X-CIO-Source", "AI Assistant") + } else { + req.Header.Set("X-CIO-Source", "CLI") } if grant := os.Getenv("X_CIO_CAPABILITY_GRANT"); grant != "" { req.Header.Set("X-CIO-Capability-Grant", grant) @@ -574,6 +580,14 @@ func (c *Client) ServiceAccountToken() string { return c.serviceAccountToken } +// SetServiceAccountToken replaces the in-memory service account credential and +// clears any cached access token, so the next request re-authenticates with the +// new token. Used by the sandbox→live self-heal after promoting the token. +func (c *Client) SetServiceAccountToken(token string) { + c.serviceAccountToken = token + c.clearAccessToken() +} + // AccessToken returns the current JWT access token. func (c *Client) AccessToken() string { return c.accessToken diff --git a/internal/client/client_test.go b/internal/client/client_test.go index 4a54294..c9ef9b6 100644 --- a/internal/client/client_test.go +++ b/internal/client/client_test.go @@ -754,15 +754,16 @@ func TestDoTrack_StandardHeaders(t *testing.T) { func TestClient_Do_AgentHeader(t *testing.T) { cases := []struct { - name string - envValue string - envSet bool - want string // expected X-CIO-Agent header value on the request + name string + envValue string + envSet bool + want string // expected X-CIO-Agent header value on the request + wantSource string // expected X-CIO-Source header value on the request }{ - {"env unset", "", false, ""}, - {"env empty", "", true, ""}, - {"env 1", "1", true, "1"}, - {"env other value", "true", true, ""}, + {"env unset", "", false, "", "CLI"}, + {"env empty", "", true, "", "CLI"}, + {"env 1", "1", true, "1", "AI Assistant"}, + {"env other value", "true", true, "", "CLI"}, } for _, tc := range cases { @@ -775,9 +776,10 @@ func TestClient_Do_AgentHeader(t *testing.T) { _ = os.Unsetenv("CIO_AGENT") } - var got string + var got, gotSource string server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { got = r.Header.Get("X-CIO-Agent") + gotSource = r.Header.Get("X-CIO-Source") _, _ = w.Write([]byte(`{"ok":true}`)) })) defer server.Close() @@ -793,6 +795,9 @@ func TestClient_Do_AgentHeader(t *testing.T) { if got != tc.want { t.Errorf("X-CIO-Agent: got %q, want %q", got, tc.want) } + if gotSource != tc.wantSource { + t.Errorf("X-CIO-Source: got %q, want %q", gotSource, tc.wantSource) + } }) } }