Skip to content

Commit b7e9d6b

Browse files
committed
feat: heartbeat manual mode so CLI tunnel does not fan out across projects
HeartbeatAgent::start() auto-enrolls every project the user has access to, which is correct for the UI (it surfaces tunnels across projects) but wrong for the CLI tunnel-listen command, which owns exactly one project. Today's logs showed the fan-out clearly: heartbeat: no connector yet project_id=drewr-y4nd1b heartbeat: no connector yet project_id=drewr3-ceu4gt The CLI silently maintained presence in drewr3-ceu4gt — a project the user never mentioned — for the lifetime of the tunnel. Harmless today, but it makes logs misleading, multiplies API load, and would create real risk if a misconfigured token granted access to a project that shouldn't be touched. Add start_manual() that skips the watcher entirely. The CLI now starts manual mode and explicitly registers its single project. Per-project loops still handle 401s via force-refresh, so transient auth blips are tolerated; the CLI's own login-state watch surfaces permanent logout. Keeps start() unchanged so the UI continues to auto-enroll. The new entry point is documented with a pointer to start_manual for callers in the CLI-style pattern.
1 parent 6264818 commit b7e9d6b

2 files changed

Lines changed: 78 additions & 4 deletions

File tree

cli/src/main.rs

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -648,10 +648,12 @@ async fn main() -> n0_error::Result<()> {
648648
tunnel.id
649649
};
650650

651-
heartbeat.start().await;
652-
if let Some(ctx) = datum.selected_context() {
653-
heartbeat.register_project(ctx.project_id).await;
654-
}
651+
// Manual mode: only heartbeat the project the tunnel
652+
// lives in. Auto-enroll would silently keep presence in
653+
// every other project the user can see, which is both
654+
// wasteful and confusing in logs.
655+
heartbeat.start_manual().await;
656+
heartbeat.register_project(project_id.clone()).await;
655657

656658
service.set_enabled_active(&tunnel_id, true).await?;
657659
println!();

lib/src/heartbeat.rs

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,13 @@ impl HeartbeatAgent {
8686
}
8787
}
8888

89+
/// Start in auto-enroll mode: watch login + projects state and keep
90+
/// heartbeats running for every project the user has access to.
91+
/// Intended for multi-project consumers like the UI.
92+
///
93+
/// For the CLI tunnel use case where there is exactly one project of
94+
/// interest, prefer [`Self::start_manual`] — auto-enroll silently
95+
/// maintains presence in projects the user didn't ask about.
8996
pub async fn start(&self) {
9097
let mut guard = self.inner.login_task.lock().await;
9198
if guard.is_some() {
@@ -134,6 +141,24 @@ impl HeartbeatAgent {
134141
*guard = Some(AbortOnDropHandle::new(task));
135142
}
136143

144+
/// Start in manual mode: do not watch login state and do not auto-enroll
145+
/// projects. Callers are responsible for [`Self::register_project`] /
146+
/// [`Self::deregister_project`] for the projects they want heartbeats
147+
/// for. Per-project loops still handle 401s via their own
148+
/// force-refresh logic, so transient auth blips are tolerated; a
149+
/// permanent logout is surfaced separately by the CLI's own login
150+
/// watcher.
151+
pub async fn start_manual(&self) {
152+
let mut guard = self.inner.login_task.lock().await;
153+
if guard.is_some() {
154+
return;
155+
}
156+
// Park a completed task so future start() / start_manual() calls
157+
// remain no-ops, matching start()'s "single-start" contract.
158+
let task = tokio::spawn(async {});
159+
*guard = Some(AbortOnDropHandle::new(task));
160+
}
161+
137162
pub async fn register_project(&self, project_id: impl Into<String>) {
138163
let project_id = project_id.into();
139164
let mut projects = self.inner.projects.lock().await;
@@ -726,6 +751,53 @@ mod tests {
726751
assert_eq!(count, 0);
727752
}
728753

754+
#[tokio::test]
755+
async fn start_manual_does_not_auto_enroll() {
756+
// Manual mode is the CLI tunnel-listen path: only the project the
757+
// caller explicitly registers should get a heartbeat task. Auto-
758+
// enroll would have probed `orgs_and_projects()` on bootstrap and
759+
// registered every accessible project — we verify it didn't by
760+
// checking the projects map stays empty until we register one.
761+
let repo = crate::Repo::open_or_create(test_repo_path()).await.unwrap();
762+
let datum = crate::datum_cloud::DatumCloudClient::with_repo(
763+
crate::datum_cloud::ApiEnv::Staging,
764+
repo,
765+
)
766+
.await
767+
.unwrap();
768+
let provider = Arc::new(TestProvider {
769+
endpoint_id: "test-endpoint".to_string(),
770+
});
771+
let runner: ProjectRunner = Arc::new(|_project_id, _datum, _provider, cancel| {
772+
tokio::spawn(async move {
773+
cancel.cancelled().await;
774+
})
775+
});
776+
let agent = HeartbeatAgent::new_with_runner(datum, provider, runner);
777+
778+
agent.start_manual().await;
779+
// Give any background bootstrap a chance to run; manual mode
780+
// shouldn't have spawned one, but if it did this would expose it.
781+
tokio::task::yield_now().await;
782+
assert_eq!(
783+
agent.inner.projects.lock().await.len(),
784+
0,
785+
"manual mode must not auto-enroll any project"
786+
);
787+
788+
agent.register_project("explicit-project").await;
789+
assert_eq!(
790+
agent.inner.projects.lock().await.len(),
791+
1,
792+
"register_project still works in manual mode"
793+
);
794+
795+
// start_manual is idempotent (matches start()'s contract): a
796+
// second call is a no-op rather than tearing down and replacing.
797+
agent.start_manual().await;
798+
assert_eq!(agent.inner.projects.lock().await.len(), 1);
799+
}
800+
729801
#[test]
730802
fn renewal_interval_in_range() {
731803
for lease_duration_seconds in [1, 2, 10, 60] {

0 commit comments

Comments
 (0)