refactor: introduce TokenSource trait and stop copying tokens as String by rawkode · Pull Request #30 · datum-cloud/connect

rawkode · 2026-09-18T22:28:15Z

Summary

DatumCloudClient and ProjectControlPlaneClient were bound to the concrete ExternalTokenSource. Every test that needed a client had to write a shell script to disk, mutate DATUM_CREDENTIALS_HELPER and DATUM_SESSION under a global mutex, and exec the script, even when the helper had nothing to do with what was under test. That harness would have been the foundation for the tunnel-service tests planned next, so it is replaced first.

Changes

New datum_cloud/token_source.rs with the TokenSource trait (token, watch, force_refresh) and StaticTokenSource, an in-memory implementation for embedders that already hold a token and for tests. It records force_refresh calls so a test can assert a refresh was requested.
DatumCloudClient and ProjectControlPlaneClient hold Arc<dyn TokenSource>. with_token_source is the new constructor; with_external_token_source stays as a wrapper so the binary is untouched. DatumCloudClient::token_source() exposes the shared source for sibling clients.
ExternalTokenSource implements the trait. swap_token and start_refresh are unchanged.
The trait returns SecretString and the canonical watch channel carries SecretString.
Tests in datum_cloud/mod.rs, heartbeat.rs and project_control_plane.rs use test_util::static_token_source() and no longer take ENV_LOCK. Only external_token_source.rs keeps the script harness, since the helper exec path is what it tests.

Follow-up fixes (second commit)

Three real issues, found by reading the actual code rather than trusting an unverifiable external report:

Retained clients missed rotations. ProjectControlPlaneClient::new() — the path every production caller uses via DatumCloudClient::project_control_plane_client — never wired up a token watch, so a client that outlives one call site continued using its original token even after the source rotated. new() now always shares datum.token_source().watch(). Covered by a #[tokio::test] that rotates a StaticTokenSource and asserts a retained client picks it up on its own.
watch::Sender::send silently discarded rotations with zero receivers. Confirmed against the vendored tokio source: send checks receiver_count() before writing the value at all, so a rotation that happened while nobody was subscribed was lost, not just unnotified — a later subscriber saw the stale value. StaticTokenSource::set and ExternalTokenSource::swap_token now use send_replace, which writes unconditionally. StaticTokenSource also drops its separate ArcSwap cache in favour of reading directly from the watch sender, so there is exactly one place the current value lives.
Public API compatibility. DatumCloudClient::token, ProjectControlPlaneClient::access_token, and ExternalTokenSource::{token, watch, force_refresh} changed return type or moved from inherent methods to a trait impl (which changes behaviour under plain dot-call syntax without importing the trait). All four keep their old String-returning, no-import-required inherent form, backed by new _secret-suffixed methods for the SecretString form. ProjectControlPlaneClient::new_with_token_source is restored to take a concrete ExternalTokenSource; a new new_with_shared_token_source takes Arc<dyn TokenSource> for the normal factory path and embedders holding only a trait object.

While adding the regression test for (1), the pre-existing #[cfg(feature = "integration-tests")] tests in project_control_plane.rs turned out to have never actually run in an environment without a pre-installed rustls CryptoProvider: kube::Client::try_from panics (not Err) without one, which those tests' own if let Ok(pcp) = pcp { .. } skip idiom can't catch, and separately they were plain #[test] functions calling a path that spawns a tokio task internally, which panics outside a runtime. Both fixed: rustls (matching bin/'s existing version and ring feature) added as a lib dev-dependency and installed once per test binary, and the affected tests converted to #[tokio::test]. All 7 now pass for real instead of silently no-op'ing.

Behaviour

The binary still constructs ExternalTokenSource::from_env, calls start_refresh, and passes it to with_external_token_source — unchanged.

Verification

Check	Result
`cargo fmt --all --check`	clean
`cargo clippy --workspace --all-targets`	exit 0, with and without `--features integration-tests`
`cargo test --workspace`	84 passed, 0 failed (77 default + 7 under `--features integration-tests`)
Repeated `--test-threads=8` runs	no flakiness across 3 runs
`grep -rln ENV_LOCK lib/src`	only `repo.rs`, `env.rs` and `external_token_source.rs`, each of which tests env vars directly

Copilot

Copilot review overview

🟡 Changes recommended

Token propagation and published API compatibility issues remain unresolved.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 6

Open (6)

Retain rotations for late token subscribers · New
Preserve ExternalTokenSource API compatibility · New
Preserve DatumCloudClient token accessor compatibility · New
Subscribe retained clients to shared token rotations · New
Retain token updates when no subscribers exist · New
Preserve constructor compatibility for external token sources · New

What changed in this PR

Introduces a shared token-source abstraction for Datum clients while retaining secrets in SecretString.

Changes:

Adds TokenSource and in-memory StaticTokenSource.
Refactors clients to share token sources and observe rotations.
Simplifies tests by removing unnecessary environment-backed helpers.

File	Description
connect-lib/lib/src/test_util.rs	Adds a reusable static token source fixture.
connect-lib/lib/src/project_control_plane.rs	Adopts shared secret-bearing token sources.
connect-lib/lib/src/lib.rs	Re-exports the new token-source API.
connect-lib/lib/src/heartbeat.rs	Uses the in-memory source in tests.
connect-lib/lib/src/datum_cloud/token_source.rs	Defines the trait and static implementation.
connect-lib/lib/src/datum_cloud/mod.rs	Stores and exposes shared token sources.
connect-lib/lib/src/datum_cloud/external_token_source.rs	Implements the trait for external credentials.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

        let token = self.token_source.token();
-        self.project_control_plane_client_with_token(project_id, &token)
+        self.project_control_plane_client_with_token(project_id, token.expose_secret())


DatumCloudClient and ProjectControlPlaneClient were bound to the concrete ExternalTokenSource, so every test that needed a client had to write a shell script to disk, mutate DATUM_CREDENTIALS_HELPER and DATUM_SESSION under a global mutex, and exec the script, even when the helper was nothing to do with what was being tested. Add a TokenSource trait (token, watch, force_refresh) and hold it as Arc<dyn TokenSource> in both clients. ExternalTokenSource implements it unchanged. StaticTokenSource is an in-memory implementation for embedders that already hold a token and for tests; it records force_refresh calls so callers can assert a refresh was requested. DatumCloudClient::with_external_token_source stays as a wrapper so the binary is untouched. The trait returns SecretString and the watch channel carries SecretString. ExternalTokenSource previously unwrapped the secret to a plain String at its boundary and broadcast it that way, and ProjectControlPlaneClient stored a plain String copy. Both now keep the secrecy wrapper until the kube config needs the bytes. Tests in datum_cloud/mod.rs, heartbeat.rs and project_control_plane.rs switch to static_token_source() and drop ENV_LOCK. Only external_token_source.rs keeps the script harness, since the helper exec path is what it tests. This commit was created with the assistance of a LLM.

…ch updates, restore public API compatibility Three real issues in the TokenSource refactor, found by independently reading the actual code (the patch/commit this was reported against did not exist anywhere in this environment or its git history, so nothing was applied blindly). 1. A ProjectControlPlaneClient built via the normal factory (DatumCloudClient::project_control_plane_client, used by every production call site) was constructed through ProjectControlPlaneClient ::new(), which always set token_rx: None. With no token_rx, the background auth-watch task fell back to watching login_state/ auth_update instead, and auth_update_watch() hands out a receiver whose sender is dropped before it's returned, so that fallback task exits almost immediately. A client that is *called fresh* on every use (which is every current internal call site) is unaffected, since it reads the token source directly at construction time; a *retained* client from an embedder is not. new() now always shares datum.token_source().watch(), which every DatumCloudClient has unconditionally, so a retained client's kube Client rebuilds itself on rotation instead of only ever seeing the token it was built with. Covered by a new #[tokio::test] that rotates a StaticTokenSource after building a client via the normal factory and asserts the retained client's access_token() updates on its own. 2. tokio::sync::watch::Sender::send returns early without writing the value at all when there are zero receivers (confirmed against the vendored tokio 1.51 source: it checks receiver_count() before calling send_replace, not after). StaticTokenSource::set and ExternalTokenSource::swap_token both called send and discarded the Result, so a rotation that happened while nobody was subscribed was silently lost — not just unnotified, but never stored — and a later subscriber would see the stale pre-rotation value. Both now use send_replace, which writes unconditionally. StaticTokenSource also drops its separate ArcSwap<SecretString> cache in favour of reading token() from the watch sender directly, so there is exactly one place the current value lives. Covered by new regression tests in both token_source.rs and external_token_source.rs. 3. The refactor silently changed several public signatures that connect-lib (an embeddable library per its own README) exposed before: DatumCloudClient::token and ProjectControlPlaneClient::access_token returned String and are now SecretString; ExternalTokenSource::token/ watch/force_refresh moved from inherent methods to a trait impl, which changes their behavior under plain dot-call syntax without importing the trait; ProjectControlPlaneClient::new_with_token_source changed from taking a concrete ExternalTokenSource to Arc<dyn TokenSource>. All four now keep their old String-returning / no-import-required inherent form for compatibility, backed by new _secret-suffixed or token_secret() methods for callers that want the SecretString. Every inherent compat method reads the same underlying storage as its trait counterpart (no parallel cache that could drift). watch() needed an actual second channel, since a watch::Receiver<T> is fixed to one T at creation; it is written in the same call as the canonical channel, so it cannot diverge from it. new_with_token_source(ExternalTokenSource) is restored as a thin wrapper around a new new_with_shared_token_source(Arc<dyn TokenSource>), which is what the normal factory path and embedders holding only a trait object now use. While adding a regression test for (1), the existing #[cfg(feature = "integration-tests")] tests in project_control_plane.rs turned out to have never actually run in an environment without a pre-installed rustls CryptoProvider: kube::Client::try_from panics (not Err) without one, which the tests' own `if let Ok(pcp) = pcp { .. }` skip idiom cannot catch, and separately they were plain #[test] functions calling a path that spawns a tokio task internally, which panics outside a runtime. Both are fixed: `rustls` (same version and `ring` feature `bin/` already depends on) is added as a lib dev-dependency and installed once per test binary, and the affected tests are converted to #[tokio::test]. All 7 integration-tests-gated tests now pass for real instead of silently no-op'ing, confirmed with repeated parallel runs. Verification: cargo fmt --all --check clean; cargo clippy --workspace --all-targets exit 0 with and without --features integration-tests; cargo test --workspace 84 passed, 0 failed (77 default + 7 under --features integration-tests), repeated three times under --test-threads=8 with no flakiness. This commit was created with the assistance of a LLM.

Copilot

Copilot review overview

🟡 Changes recommended

Concurrent token rebuilds can leave the cached token and kube client inconsistent.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 1 · 9

Open (10)

Atomically update client and token during rebuilds · New
Preserve SecretString through the production client factory · New
Synchronize environment access in shared-token tests · New
Require successful shared-token client construction · New
Unwrap construction before asserting project ID · New
Require construction before checking token accessors · New
Require construction before checking server URL · New
Require construction before checking plugin mode · New
Unwrap compatibility constructor result · New
Subscribe retained clients to shared token rotations

Resolved since last review (5)

Preserve constructor compatibility for external token sources
Retain token updates when no subscribers exist
Preserve DatumCloudClient token accessor compatibility
Preserve ExternalTokenSource API compatibility
Retain rotations for late token subscribers

+        // ever seeing the token it was constructed with. Every
+        // `DatumCloudClient` has a `TokenSource` unconditionally now, so
+        // this is always available.
+        let token_rx = Some(datum.token_source().watch());


        let token = self.token_source.token();
-        self.project_control_plane_client_with_token(project_id, &token)
+        self.project_control_plane_client_with_token(project_id, token.expose_secret())


-        let datum = DatumCloudClient::with_external_token_source(
+        let client = Self::build_kube_client(&server_url, initial_token.expose_secret())?;
+        let datum = DatumCloudClient::with_token_source(
            crate::ApiEnv::from_env_with_host_override(),


+        let result = ProjectControlPlaneClient::new_with_shared_token_source(
            "test-project".to_string(),
            "https://api.datum.net/apis/resourcemanager.miloapis.com/v1alpha1/projects/test-project/control-plane".to_string(),
            token_source,
        );
        let _ = result;


@@ -218,27 +272,30 @@ mod tests {
        }


+        let pcp = ProjectControlPlaneClient::new_with_shared_token_source(
            "test-project".to_string(),
            "https://api.datum.net/apis/resourcemanager.miloapis.com/v1alpha1/projects/test-project/control-plane".to_string(),
            token_source,
        );
        if let Ok(pcp) = pcp {
            assert_eq!(pcp.access_token(), expected_token);
+            assert_eq!(pcp.access_token_secret().expose_secret(), expected_token);
        }


@@ -248,11 +305,12 @@ mod tests {
        }


@@ -261,4 +319,58 @@ mod tests {
            assert!(pcp.datum.is_plugin_mode());
        }


+        let pcp = ProjectControlPlaneClient::new_with_token_source(
+            "test-project".to_string(),
+            "https://api.datum.net/apis/resourcemanager.miloapis.com/v1alpha1/projects/test-project/control-plane".to_string(),
+            token_source,
+        );
+        if let Ok(pcp) = pcp {
+            assert_eq!(pcp.access_token(), expected);
+        }


rawkode requested a balanced review from Copilot September 18, 2026 22:40

Copilot started reviewing on behalf of rawkode September 18, 2026 22:40 View session

Copilot AI reviewed Sep 18, 2026

View reviewed changes

rawkode force-pushed the claude/token-source-trait branch from 8dc401c to ac80b39 Compare September 19, 2026 15:58

rawkode marked this pull request as ready for review September 19, 2026 18:04

rawkode requested a balanced review from Copilot September 19, 2026 18:05

Copilot started reviewing on behalf of rawkode September 19, 2026 18:06 View session

Copilot AI reviewed Sep 19, 2026

View reviewed changes

rawkode merged commit b160478 into develop Sep 19, 2026
4 checks passed

Navigation Menu

refactor: introduce TokenSource trait and stop copying tokens as String - #30

refactor: introduce TokenSource trait and stop copying tokens as String#30
rawkode merged 2 commits into
developfrom
claude/token-source-trait

rawkode commented Sep 18, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		@@ -261,4 +319,58 @@ mod tests {
		assert!(pcp.datum.is_plugin_mode());
		}

Conversation

rawkode commented Sep 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

Follow-up fixes (second commit)

Behaviour

Verification

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

rawkode commented Sep 18, 2026 •

edited

Loading