diff --git a/.github/workflows/component-fixtures.yml b/.github/workflows/component-fixtures.yml index 7c6fe21639108b..a168e69233791c 100644 --- a/.github/workflows/component-fixtures.yml +++ b/.github/workflows/component-fixtures.yml @@ -30,8 +30,8 @@ jobs: with: # Need enough history for the merge-base lookup below to succeed even # when the target branch has advanced since the PR was opened. Full - # clone would be wasteful for this large repo, so cap at 50. - fetch-depth: 50 + # clone would be wasteful for this large repo, so cap at 150. + fetch-depth: 150 - name: Setup Node.js uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7.0.0 @@ -184,8 +184,11 @@ jobs: if [ "${{ github.event_name }}" = "pull_request" ]; then # For PRs, diff against the merge-base with the target branch. TARGET_REF="origin/$BASE_REF" - git fetch --no-tags --depth=50 origin "$BASE_REF" - BASE_SHA=$(git merge-base "$EVENT_SHA" "$TARGET_REF") + git fetch --no-tags --depth=150 origin "$BASE_REF" + if ! BASE_SHA=$(git merge-base "$EVENT_SHA" "$TARGET_REF"); then + echo "::warning::Unable to find a merge base between $EVENT_SHA and $TARGET_REF. The depth-150 shallow history may not contain their common ancestor; skipping screenshot comparison." + exit 0 + fi else # For push events, diff against the parent commit. BASE_SHA=$(git rev-parse "$EVENT_SHA^") @@ -229,6 +232,7 @@ jobs: - name: Fetch base commit manifest id: base_manifest + if: steps.base.outputs.base_sha != '' env: BASE_SHA: ${{ steps.base.outputs.base_sha }} run: | @@ -250,7 +254,7 @@ jobs: - name: Diff screenshots id: diff - if: always() + if: always() && steps.base.outputs.base_sha != '' run: | node build/lib/screenshotDiffReport.ts \ https://hediet-screenshots.azurewebsites.net \ diff --git a/.github/workflows/pr-darwin-test.yml b/.github/workflows/pr-darwin-test.yml index cd0712c52d7b64..4f41f59687bfb3 100644 --- a/.github/workflows/pr-darwin-test.yml +++ b/.github/workflows/pr-darwin-test.yml @@ -13,7 +13,10 @@ on: remote_tests: type: boolean default: false - unit_and_integration_tests: + unit_tests: + type: boolean + default: true + integration_tests: type: boolean default: true smoke_tests: @@ -25,7 +28,7 @@ jobs: name: ${{ inputs.job_name }} runs-on: macos-26-xlarge env: - ARTIFACT_NAME: ${{ (inputs.electron_tests && 'electron') || (inputs.browser_tests && 'browser') || (inputs.remote_tests && 'remote') || 'unknown' }}${{ (!inputs.unit_and_integration_tests && inputs.smoke_tests) && '-smoke' || '' }} + ARTIFACT_NAME: ${{ (inputs.electron_tests && 'electron') || (inputs.browser_tests && 'browser') || (inputs.remote_tests && 'remote') || 'unknown' }}${{ (inputs.unit_tests && !inputs.integration_tests && '-unit') || (!inputs.unit_tests && !inputs.integration_tests && inputs.smoke_tests && '-smoke') || '' }} NPM_ARCH: arm64 VSCODE_ARCH: arm64 steps: @@ -35,7 +38,7 @@ jobs: lfs: true - name: Detect Agent Host E2E changes - if: ${{ inputs.electron_tests && inputs.unit_and_integration_tests }} + if: ${{ inputs.electron_tests && inputs.integration_tests }} id: agent-host-e2e-changes uses: ./.github/actions/detect-agent-host-e2e-changes with: @@ -104,12 +107,13 @@ jobs: - name: Transpile client and extensions run: npm run gulp transpile-client-esbuild transpile-extensions - - name: Download Electron and Playwright + - name: Download Electron + if: ${{ inputs.electron_tests || inputs.remote_tests }} run: | set -e for i in {1..3}; do # try 3 times (matching retryCountOnTaskFailure: 3) - if npm exec -- npm-run-all2 -lp "electron ${{ env.VSCODE_ARCH }}" "playwright-install"; then + if npm run electron -- ${{ env.VSCODE_ARCH }}; then echo "Download successful on attempt $i" break fi @@ -125,26 +129,52 @@ jobs: env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + - name: Compile browser integration test runner + if: ${{ inputs.browser_tests && inputs.integration_tests }} + working-directory: test/integration/browser + run: npm run compile + + - name: Install Playwright Chromium and WebKit + if: ${{ inputs.browser_tests }} + run: | + set -e + + for i in {1..3}; do + if npm exec -- playwright install --only-shell chromium webkit; then + echo "Install successful on attempt $i" + break + fi + + if [ $i -eq 3 ]; then + echo "Install failed after 3 attempts" >&2 + exit 1 + fi + + echo "Install failed on attempt $i, retrying..." + sleep 5 + done + - name: 🧪 Run unit tests (Electron) - if: ${{ inputs.electron_tests && inputs.unit_and_integration_tests }} + if: ${{ inputs.electron_tests && inputs.unit_tests }} timeout-minutes: 15 run: ./scripts/test.sh --tfs "Unit Tests" env: VSCODE_SKIP_PRELAUNCH: '1' - name: 🧪 Run unit tests (node.js) - if: ${{ inputs.electron_tests && inputs.unit_and_integration_tests }} + if: ${{ inputs.electron_tests && inputs.unit_tests }} timeout-minutes: 15 run: npm run test-node - name: 🧪 Run unit tests (Browser, Webkit) - if: ${{ inputs.browser_tests && inputs.unit_and_integration_tests }} + if: ${{ inputs.browser_tests && inputs.unit_tests }} timeout-minutes: 30 run: npm run test-browser-no-install -- --browser webkit --tfs "Browser Unit Tests" env: DEBUG: "*browser*" - name: Compile extensions for integration tests & smoke tests + if: ${{ inputs.integration_tests || inputs.smoke_tests }} run: | set -e npm run gulp \ @@ -167,7 +197,7 @@ jobs: compile-extension:vscode-test-resolver - name: 🧪 Run integration tests (Electron) - if: ${{ inputs.electron_tests && inputs.unit_and_integration_tests }} + if: ${{ inputs.electron_tests && inputs.integration_tests }} timeout-minutes: 20 run: ./scripts/test-integration.sh --tfs "Integration Tests" env: @@ -175,12 +205,12 @@ jobs: VSCODE_SKIP_PRELAUNCH: '1' - name: 🧪 Run integration tests (Browser, Webkit) - if: ${{ inputs.browser_tests && inputs.unit_and_integration_tests }} + if: ${{ inputs.browser_tests && inputs.integration_tests }} timeout-minutes: 20 run: ./scripts/test-web-integration.sh --browser webkit - name: 🧪 Run integration tests (Remote) - if: ${{ inputs.remote_tests && inputs.unit_and_integration_tests }} + if: ${{ inputs.remote_tests && inputs.integration_tests }} timeout-minutes: 20 run: ./scripts/test-remote-integration.sh env: diff --git a/.github/workflows/pr-linux-test.yml b/.github/workflows/pr-linux-test.yml index 7e2cf2b1c3b00e..1eb53e5de9f00f 100644 --- a/.github/workflows/pr-linux-test.yml +++ b/.github/workflows/pr-linux-test.yml @@ -13,7 +13,10 @@ on: remote_tests: type: boolean default: false - unit_and_integration_tests: + unit_tests: + type: boolean + default: true + integration_tests: type: boolean default: true smoke_tests: @@ -25,7 +28,7 @@ jobs: name: ${{ inputs.job_name }} runs-on: ubuntu-24.04 env: - ARTIFACT_NAME: ${{ (inputs.electron_tests && 'electron') || (inputs.browser_tests && 'browser') || (inputs.remote_tests && 'remote') || 'unknown' }}${{ (!inputs.unit_and_integration_tests && inputs.smoke_tests) && '-smoke' || '' }} + ARTIFACT_NAME: ${{ (inputs.electron_tests && 'electron') || (inputs.browser_tests && 'browser') || (inputs.remote_tests && 'remote') || 'unknown' }}${{ (inputs.unit_tests && !inputs.integration_tests && '-unit') || (!inputs.unit_tests && !inputs.integration_tests && inputs.smoke_tests && '-smoke') || '' }} NPM_ARCH: x64 VSCODE_ARCH: x64 steps: @@ -35,7 +38,7 @@ jobs: lfs: true - name: Detect Agent Host E2E changes - if: ${{ inputs.electron_tests && inputs.unit_and_integration_tests }} + if: ${{ inputs.electron_tests && inputs.integration_tests }} id: agent-host-e2e-changes uses: ./.github/actions/detect-agent-host-e2e-changes with: @@ -304,12 +307,13 @@ jobs: - name: Transpile client and extensions run: npm run gulp transpile-client-esbuild transpile-extensions - - name: Download Electron and Playwright + - name: Download Electron + if: ${{ inputs.electron_tests || inputs.remote_tests }} run: | set -e for i in {1..3}; do # try 3 times (matching retryCountOnTaskFailure: 3) - if npm exec -- npm-run-all2 -lp "electron ${{ env.VSCODE_ARCH }}" "playwright-install"; then + if npm run electron -- ${{ env.VSCODE_ARCH }}; then echo "Download successful on attempt $i" break fi @@ -325,8 +329,33 @@ jobs: env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + - name: Compile browser integration test runner + if: ${{ inputs.browser_tests && inputs.integration_tests }} + working-directory: test/integration/browser + run: npm run compile + + - name: Install Playwright Chromium + if: ${{ inputs.browser_tests }} + run: | + set -e + + for i in {1..3}; do + if npm exec -- playwright install --only-shell chromium; then + echo "Install successful on attempt $i" + break + fi + + if [ $i -eq 3 ]; then + echo "Install failed after 3 attempts" >&2 + exit 1 + fi + + echo "Install failed on attempt $i, retrying..." + sleep 5 + done + - name: 🧪 Run unit tests (Electron) - if: ${{ inputs.electron_tests && inputs.unit_and_integration_tests }} + if: ${{ inputs.electron_tests && inputs.unit_tests }} timeout-minutes: 15 run: ./scripts/test.sh --tfs "Unit Tests" env: @@ -334,18 +363,19 @@ jobs: VSCODE_SKIP_PRELAUNCH: '1' - name: 🧪 Run unit tests (node.js) - if: ${{ inputs.electron_tests && inputs.unit_and_integration_tests }} + if: ${{ inputs.electron_tests && inputs.unit_tests }} timeout-minutes: 15 run: npm run test-node - name: 🧪 Run unit tests (Browser, Chromium) - if: ${{ inputs.browser_tests && inputs.unit_and_integration_tests }} + if: ${{ inputs.browser_tests && inputs.unit_tests }} timeout-minutes: 30 run: npm run test-browser-no-install -- --browser chromium --tfs "Browser Unit Tests" env: DEBUG: "*browser*" - name: Compile extensions for integration tests & smoke tests + if: ${{ inputs.integration_tests || inputs.smoke_tests }} run: | set -e npm run gulp \ @@ -368,7 +398,7 @@ jobs: compile-extension:vscode-test-resolver - name: 🧪 Run integration tests (Electron) - if: ${{ inputs.electron_tests && inputs.unit_and_integration_tests }} + if: ${{ inputs.electron_tests && inputs.integration_tests }} timeout-minutes: 20 run: ./scripts/test-integration.sh --tfs "Integration Tests" env: @@ -377,12 +407,12 @@ jobs: VSCODE_SKIP_PRELAUNCH: '1' - name: 🧪 Run integration tests (Browser, Chromium) - if: ${{ inputs.browser_tests && inputs.unit_and_integration_tests }} + if: ${{ inputs.browser_tests && inputs.integration_tests }} timeout-minutes: 20 run: ./scripts/test-web-integration.sh --browser chromium - name: 🧪 Run integration tests (Remote) - if: ${{ inputs.remote_tests && inputs.unit_and_integration_tests }} + if: ${{ inputs.remote_tests && inputs.integration_tests }} timeout-minutes: 20 run: ./scripts/test-remote-integration.sh env: diff --git a/.github/workflows/pr-win32-test.yml b/.github/workflows/pr-win32-test.yml index 981bf097e7d679..9edb6eb1343f32 100644 --- a/.github/workflows/pr-win32-test.yml +++ b/.github/workflows/pr-win32-test.yml @@ -13,7 +13,10 @@ on: remote_tests: type: boolean default: false - unit_and_integration_tests: + unit_tests: + type: boolean + default: true + integration_tests: type: boolean default: true smoke_tests: @@ -25,7 +28,7 @@ jobs: name: ${{ inputs.job_name }} runs-on: [ self-hosted, 1ES.Pool=1es-vscode-oss-windows-2022-x64, "JobId=windows-test-${{ inputs.job_name }}-${{ github.run_id }}-${{ github.run_number }}-${{ github.run_attempt }}" ] env: - ARTIFACT_NAME: ${{ (inputs.electron_tests && 'electron') || (inputs.browser_tests && 'browser') || (inputs.remote_tests && 'remote') || 'unknown' }}${{ (!inputs.unit_and_integration_tests && inputs.smoke_tests) && '-smoke' || '' }} + ARTIFACT_NAME: ${{ (inputs.electron_tests && 'electron') || (inputs.browser_tests && 'browser') || (inputs.remote_tests && 'remote') || 'unknown' }}${{ (inputs.unit_tests && !inputs.integration_tests && '-unit') || (!inputs.unit_tests && !inputs.integration_tests && inputs.smoke_tests && '-smoke') || '' }} NPM_ARCH: x64 VSCODE_ARCH: x64 steps: @@ -35,7 +38,7 @@ jobs: lfs: true - name: Detect Agent Host E2E changes - if: ${{ inputs.electron_tests && inputs.unit_and_integration_tests }} + if: ${{ inputs.electron_tests && inputs.integration_tests }} id: agent-host-e2e-changes uses: ./.github/actions/detect-agent-host-e2e-changes with: @@ -112,7 +115,8 @@ jobs: shell: pwsh run: npm run gulp "transpile-client-esbuild" "transpile-extensions" - - name: Download Electron and Playwright + - name: Download Electron + if: ${{ inputs.electron_tests || inputs.remote_tests }} shell: pwsh run: | . build/azure-pipelines/win32/exec.ps1 @@ -120,7 +124,7 @@ jobs: for ($i = 1; $i -le 3; $i++) { try { - exec { npm exec -- npm-run-all2 -lp "electron ${{ env.VSCODE_ARCH }}" "playwright-install" } + exec { npm run electron -- ${{ env.VSCODE_ARCH }} } break } catch { @@ -135,8 +139,35 @@ jobs: env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + - name: Compile browser integration test runner + if: ${{ inputs.browser_tests && inputs.integration_tests }} + working-directory: test/integration/browser + run: npm run compile + + - name: Install Playwright Chromium + if: ${{ inputs.browser_tests }} + shell: pwsh + run: | + . build/azure-pipelines/win32/exec.ps1 + $ErrorActionPreference = "Stop" + + for ($i = 1; $i -le 3; $i++) { + try { + exec { npm exec -- playwright install --only-shell chromium } + break + } + catch { + if ($i -eq 3) { + Write-Error "Install failed after 3 attempts" + throw + } + Write-Host "Install failed attempt $i, retrying..." + Start-Sleep -Seconds 2 + } + } + - name: 🧪 Run unit tests (Electron) - if: ${{ inputs.electron_tests && inputs.unit_and_integration_tests }} + if: ${{ inputs.electron_tests && inputs.unit_tests }} timeout-minutes: 15 shell: pwsh run: .\scripts\test.bat --tfs "Unit Tests" @@ -144,13 +175,13 @@ jobs: VSCODE_SKIP_PRELAUNCH: '1' - name: 🧪 Run unit tests (node.js) - if: ${{ inputs.electron_tests && inputs.unit_and_integration_tests }} + if: ${{ inputs.electron_tests && inputs.unit_tests }} timeout-minutes: 15 shell: pwsh run: npm run test-node - name: 🧪 Run unit tests (Browser, Chromium) - if: ${{ inputs.browser_tests && inputs.unit_and_integration_tests }} + if: ${{ inputs.browser_tests && inputs.unit_tests }} timeout-minutes: 20 shell: pwsh run: node test/unit/browser/index.js --browser chromium --tfs "Browser Unit Tests" @@ -158,6 +189,7 @@ jobs: DEBUG: "*browser*" - name: Compile extensions for integration tests & smoke tests + if: ${{ inputs.integration_tests || inputs.smoke_tests }} shell: pwsh run: | . build/azure-pipelines/win32/exec.ps1 @@ -183,13 +215,13 @@ jobs: } - name: Diagnostics before integration test runs - if: ${{ inputs.unit_and_integration_tests && always() }} + if: ${{ inputs.integration_tests && always() }} shell: pwsh run: .\build\azure-pipelines\win32\listprocesses.bat continue-on-error: true - name: 🧪 Run integration tests (Electron) - if: ${{ inputs.electron_tests && inputs.unit_and_integration_tests }} + if: ${{ inputs.electron_tests && inputs.integration_tests }} timeout-minutes: 20 shell: pwsh run: .\scripts\test-integration.bat --tfs "Integration Tests" @@ -198,13 +230,13 @@ jobs: VSCODE_SKIP_PRELAUNCH: '1' - name: 🧪 Run integration tests (Browser, Chromium) - if: ${{ inputs.browser_tests && inputs.unit_and_integration_tests }} + if: ${{ inputs.browser_tests && inputs.integration_tests }} timeout-minutes: 20 shell: pwsh run: .\scripts\test-web-integration.bat --browser chromium - name: 🧪 Run integration tests (Remote) - if: ${{ inputs.remote_tests && inputs.unit_and_integration_tests }} + if: ${{ inputs.remote_tests && inputs.integration_tests }} timeout-minutes: 20 shell: pwsh run: .\scripts\test-remote-integration.bat @@ -212,7 +244,7 @@ jobs: VSCODE_SKIP_PRELAUNCH: '1' - name: Diagnostics after integration test runs - if: ${{ inputs.unit_and_integration_tests && always() }} + if: ${{ inputs.integration_tests && always() }} shell: pwsh run: .\build\azure-pipelines\win32\listprocesses.bat continue-on-error: true diff --git a/.github/workflows/pr.yml b/.github/workflows/pr.yml index b04ecfe533cbf4..f196afce554fb7 100644 --- a/.github/workflows/pr.yml +++ b/.github/workflows/pr.yml @@ -113,12 +113,22 @@ jobs: job_name: CLI rustup_toolchain: 1.88 + linux-electron-unit-tests: + name: Linux + uses: ./.github/workflows/pr-linux-test.yml + with: + job_name: Electron-Unit + electron_tests: true + integration_tests: false + smoke_tests: false + linux-electron-tests: name: Linux uses: ./.github/workflows/pr-linux-test.yml with: job_name: Electron electron_tests: true + unit_tests: false smoke_tests: false linux-electron-smoke-tests: @@ -127,7 +137,8 @@ jobs: with: job_name: Electron-Smoke electron_tests: true - unit_and_integration_tests: false + unit_tests: false + integration_tests: false linux-browser-tests: name: Linux @@ -143,12 +154,22 @@ jobs: job_name: Remote remote_tests: true + macos-electron-unit-tests: + name: macOS + uses: ./.github/workflows/pr-darwin-test.yml + with: + job_name: Electron-Unit + electron_tests: true + integration_tests: false + smoke_tests: false + macos-electron-tests: name: macOS uses: ./.github/workflows/pr-darwin-test.yml with: job_name: Electron electron_tests: true + unit_tests: false smoke_tests: false macos-electron-smoke-tests: @@ -157,7 +178,8 @@ jobs: with: job_name: Electron-Smoke electron_tests: true - unit_and_integration_tests: false + unit_tests: false + integration_tests: false macos-browser-tests: name: macOS @@ -173,12 +195,22 @@ jobs: job_name: Remote remote_tests: true + windows-electron-unit-tests: + name: Windows + uses: ./.github/workflows/pr-win32-test.yml + with: + job_name: Electron-Unit + electron_tests: true + integration_tests: false + smoke_tests: false + windows-electron-tests: name: Windows uses: ./.github/workflows/pr-win32-test.yml with: job_name: Electron electron_tests: true + unit_tests: false smoke_tests: false windows-electron-smoke-tests: @@ -187,7 +219,8 @@ jobs: with: job_name: Electron-Smoke electron_tests: true - unit_and_integration_tests: false + unit_tests: false + integration_tests: false windows-browser-tests: name: Windows @@ -407,14 +440,6 @@ jobs: working-directory: extensions/copilot run: npm ci - - name: TypeScript type checking - working-directory: extensions/copilot - run: npm run typecheck - - - name: Lint - working-directory: extensions/copilot - run: npm run lint - - name: Compile working-directory: extensions/copilot run: npm run compile diff --git a/cli/src/download_cache.rs b/cli/src/download_cache.rs index 87ca1924a798cc..a88d7789ababb5 100644 --- a/cli/src/download_cache.rs +++ b/cli/src/download_cache.rs @@ -4,20 +4,27 @@ *--------------------------------------------------------------------------------------------*/ use std::{ - fs::create_dir_all, + fs::{create_dir, create_dir_all, OpenOptions}, path::{Path, PathBuf}, }; use futures::Future; -use tokio::fs::remove_dir_all; +use uuid::Uuid; use crate::{ state::PersistedState, - util::errors::{wrap, AnyError, WrappedError}, + util::{ + errors::{wrap, AnyError, WrappedError}, + file_lock::{FileLock, Lock}, + }, }; const KEEP_LRU: usize = 5; const STAGING_SUFFIX: &str = ".staging"; +const LOCKS_DIRECTORY: &str = ".locks"; +const LOCK_WAIT_INITIAL_DELAY: std::time::Duration = std::time::Duration::from_millis(200); +const LOCK_WAIT_MAX_DELAY: std::time::Duration = std::time::Duration::from_secs(2); +const LOCK_WAIT_HEARTBEAT_INTERVAL: std::time::Duration = std::time::Duration::from_secs(5); const RENAME_ATTEMPTS: u32 = 20; const RENAME_DELAY: std::time::Duration = std::time::Duration::from_millis(200); const PERSISTED_STATE_FILE_NAME: &str = "lru.json"; @@ -28,6 +35,16 @@ pub struct DownloadCache { state: PersistedState>, } +struct StagingDirectory(PathBuf); + +impl Drop for StagingDirectory { + fn drop(&mut self) { + // Drop cannot await, so use blocking cleanup to also remove staging directories + // when the creating future is cancelled. + let _ = std::fs::remove_dir_all(&self.0); + } +} + impl DownloadCache { pub fn new(path: PathBuf) -> DownloadCache { DownloadCache { @@ -90,20 +107,92 @@ impl DownloadCache { return Ok(target_dir); } - let temp_dir = self.path.join(format!("{name}{STAGING_SUFFIX}")); - let _ = remove_dir_all(&temp_dir).await; // cleanup any existing + create_dir_all(&self.path).map_err(|e| wrap(e, "error creating server directory"))?; + + let lock_path = self.path.join(LOCKS_DIRECTORY).join(name); + if let Some(lock_parent) = lock_path.parent() { + create_dir_all(lock_parent) + .map_err(|e| wrap(e, "error creating server download lock"))?; + } + + let mut lock_wait_started = None; + let mut lock_wait_delay = LOCK_WAIT_INITIAL_DELAY; + let mut next_lock_wait_heartbeat = LOCK_WAIT_HEARTBEAT_INTERVAL; + let _lock = loop { + let lock_file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + // The file is a lock holder, not a data file: its contents are + // never read or written, so it must not be truncated out from + // under another process already holding the lock. + .truncate(false) + .open(&lock_path) + .map_err(|e| wrap(e, "error creating server download lock"))?; + + match FileLock::acquire(lock_file) + .map_err(|e| wrap(e, "error acquiring server download lock"))? + { + Lock::Acquired(lock) => break lock, + Lock::AlreadyLocked(_) if target_dir.exists() => { + let _ = self.touch(name.to_string()); + return Ok(target_dir); + } + Lock::AlreadyLocked(_) => { + let first_wait = lock_wait_started.is_none(); + let wait_started = + lock_wait_started.get_or_insert_with(std::time::Instant::now); + let elapsed = wait_started.elapsed(); + if first_wait { + log::info!( + "Another instance is already downloading the server; waiting for it to finish" + ); + } else if elapsed >= next_lock_wait_heartbeat { + log::info!( + "Another instance is still downloading the server; waited {} seconds", + elapsed.as_secs() + ); + next_lock_wait_heartbeat = elapsed + LOCK_WAIT_HEARTBEAT_INTERVAL; + } - create_dir_all(&temp_dir).map_err(|e| wrap(e, "error creating server directory"))?; - do_create(temp_dir.clone()).await?; + tokio::time::sleep(lock_wait_delay).await; + lock_wait_delay = + std::cmp::min(lock_wait_delay.saturating_mul(2), LOCK_WAIT_MAX_DELAY); + } + } + }; + + if target_dir.exists() { + let _ = self.touch(name.to_string()); + return Ok(target_dir); + } + + // Holding the lock for `name` means no other process is staging this + // entry, so any `{name}.staging-*` left behind belongs to an attempt + // that died before its cleanup guard ran. Nothing else reaps these, and + // each one is a partial server download, so drop them here rather than + // leaking disk on every crash. The `{name}{STAGING_SUFFIX}-` prefix + // cannot match another entry's staging directory. + self.remove_orphaned_staging_directories(name); + + let temp_dir = self + .path + .join(format!("{name}{STAGING_SUFFIX}-{}", Uuid::new_v4())); + create_dir(&temp_dir).map_err(|e| wrap(e, "error creating server directory"))?; + let temp_dir = StagingDirectory(temp_dir); + do_create(temp_dir.0.clone()).await?; let _ = self.touch(name.to_string()); // retry the rename, it seems on WoA sometimes it takes a second for the // directory to be 'unlocked' after doing file/process operations in it. for attempt_no in 0..=RENAME_ATTEMPTS { - match std::fs::rename(&temp_dir, &target_dir) { + match std::fs::rename(&temp_dir.0, &target_dir) { Ok(_) => { break; } + Err(_) if target_dir.exists() => { + return Ok(target_dir); + } Err(e) if attempt_no == RENAME_ATTEMPTS => { return Err(wrap(e, "error renaming downloaded server").into()) } @@ -116,6 +205,20 @@ impl DownloadCache { Ok(target_dir) } + /// Removes staging directories left by earlier attempts at `name`. Only safe + /// while holding that entry's download lock — see the call site. + fn remove_orphaned_staging_directories(&self, name: &str) { + let prefix = format!("{name}{STAGING_SUFFIX}-"); + let Ok(entries) = std::fs::read_dir(&self.path) else { + return; + }; + for entry in entries.flatten() { + if entry.file_name().to_string_lossy().starts_with(&prefix) { + let _ = std::fs::remove_dir_all(entry.path()); + } + } + } + fn touch(&self, name: String) -> Result<(), AnyError> { self.state.update(|l| { if let Some(index) = l.iter().position(|s| s == &name) { @@ -138,3 +241,122 @@ impl DownloadCache { Ok(()) } } + +#[cfg(test)] +mod tests { + use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }; + + use super::*; + + fn staging_directories(cache: &DownloadCache, name: &str) -> Vec { + std::fs::read_dir(cache.path()) + .unwrap() + .filter_map(Result::ok) + .map(|entry| entry.path()) + .filter(|path| { + path.file_name() + .unwrap() + .to_string_lossy() + .starts_with(&format!("{name}{STAGING_SUFFIX}")) + }) + .collect() + } + + #[tokio::test] + async fn test_concurrent_create_runs_creator_once() { + let dir = tempfile::tempdir().unwrap(); + let cache = DownloadCache::new(dir.path().join("cache")); + let create_count = Arc::new(AtomicUsize::new(0)); + + let first_count = create_count.clone(); + let first = cache.create("server", move |path| { + first_count.fetch_add(1, Ordering::SeqCst); + async move { + std::fs::write(path.join("created"), "").unwrap(); + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + Ok(()) + } + }); + let second_count = create_count.clone(); + let second = cache.create("server", move |_| { + second_count.fetch_add(1, Ordering::SeqCst); + async { Ok(()) } + }); + + let (first, second) = tokio::join!(first, second); + assert_eq!(first.unwrap(), second.unwrap()); + assert_eq!(create_count.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn test_reaps_staging_directories_left_by_a_dead_attempt() { + let dir = tempfile::tempdir().unwrap(); + let cache = DownloadCache::new(dir.path().join("cache")); + std::fs::create_dir_all(cache.path()).unwrap(); + // A staging directory whose process died before its cleanup guard ran, + // plus a sibling entry's staging directory that must survive. + let orphan = cache.path().join(format!("server{STAGING_SUFFIX}-dead")); + let other = cache.path().join(format!("server-2{STAGING_SUFFIX}-live")); + std::fs::create_dir_all(&orphan).unwrap(); + std::fs::create_dir_all(&other).unwrap(); + + cache + .create("server", |path| async move { + std::fs::write(path.join("created"), "").unwrap(); + Ok(()) + }) + .await + .unwrap(); + + assert_eq!( + ( + orphan.exists(), + other.exists(), + staging_directories(&cache, "server").len() + ), + (false, true, 0) + ); + } + + #[tokio::test] + async fn test_failed_create_removes_staging_directory() { + let dir = tempfile::tempdir().unwrap(); + let cache = DownloadCache::new(dir.path().join("cache")); + + let result = cache + .create("server", |_| async { + Err::<(), AnyError>( + wrap(std::io::Error::other("expected failure"), "test failure").into(), + ) + }) + .await; + + assert!(result.is_err()); + assert!(staging_directories(&cache, "server").is_empty()); + } + + #[tokio::test] + async fn test_lost_rename_race_returns_existing_target() { + let dir = tempfile::tempdir().unwrap(); + let cache = DownloadCache::new(dir.path().join("cache")); + let target_dir = cache.path().join("server"); + + let result = cache + .create("server", move |path| { + let target_dir = target_dir.clone(); + async move { + std::fs::write(path.join("created"), "").unwrap(); + std::fs::create_dir(&target_dir).unwrap(); + std::fs::write(target_dir.join("winner"), "").unwrap(); + Ok(()) + } + }) + .await; + + assert_eq!(result.unwrap(), cache.path().join("server")); + assert!(staging_directories(&cache, "server").is_empty()); + } +} diff --git a/extensions/copilot/package.json b/extensions/copilot/package.json index eab066dfb408d9..5eba41442252af 100644 --- a/extensions/copilot/package.json +++ b/extensions/copilot/package.json @@ -2707,6 +2707,12 @@ "icon": "$(inspect)", "category": "Developer" }, + { + "command": "github.copilot.debug.logTypeScriptContainers", + "title": "%github.copilot.command.logTypeScriptContainers%", + "enablement": "editorLangId == typescript || editorLangId == javascript", + "category": "Developer" + }, { "command": "github.copilot.debug.validateNesRename", "title": "%github.copilot.command.validateNesRename%", @@ -4951,16 +4957,6 @@ "onExp" ] }, - "github.copilot.chat.inlineEdits.nextCursorPrediction.displayLine": { - "type": "boolean", - "default": true, - "markdownDescription": "%github.copilot.config.inlineEdits.nextCursorPrediction.displayLine%", - "tags": [ - "advanced", - "experimental", - "onExp" - ] - }, "github.copilot.chat.inlineEdits.nextCursorPrediction.currentFileMaxTokens": { "type": "number", "default": 3000, diff --git a/extensions/copilot/package.nls.json b/extensions/copilot/package.nls.json index a6d7eac6a176d4..0839456250b8a3 100644 --- a/extensions/copilot/package.nls.json +++ b/extensions/copilot/package.nls.json @@ -103,6 +103,7 @@ "github.copilot.command.showChatLogView": "Show Chat Debug View", "github.copilot.command.showOutputChannel": "Show Output Channel", "github.copilot.command.showContextInspectorView": "Inspect Language Context", + "github.copilot.command.logTypeScriptContainers": "Log TypeScript Containers", "github.copilot.command.validateNesRename": "Validate NES Rename", "github.copilot.command.resetVirtualToolGroups": "Reset Virtual Tool Groups", "github.copilot.command.extensionState": "Log Extension State", @@ -440,7 +441,6 @@ "github.copilot.config.cloudAgent.enabled": "Enable the Cloud Agent. When disabled, the Cloud Agent will not be available in 'Continue In' context menus.", "github.copilot.config.gpt5AlternativePatch": "Enable GPT-5 alternative patch format.", "github.copilot.config.inlineEdits.triggerOnEditorChangeAfterSeconds": "Trigger inline edits after editor has been idle for this many seconds.", - "github.copilot.config.inlineEdits.nextCursorPrediction.displayLine": "Display predicted cursor line for next edit suggestions.", "github.copilot.config.inlineEdits.nextCursorPrediction.currentFileMaxTokens": "Maximum tokens for current file in next cursor prediction.", "github.copilot.config.inlineEdits.renameSymbolSuggestions": "Enable rename symbol suggestions in inline edits.", "github.copilot.config.nextEditSuggestions.preferredModel": "Preferred model for next edit suggestions.", diff --git a/extensions/copilot/src/extension/extension/vscode-node/services.ts b/extensions/copilot/src/extension/extension/vscode-node/services.ts index 15ca799eec7c00..9271573284f753 100644 --- a/extensions/copilot/src/extension/extension/vscode-node/services.ts +++ b/extensions/copilot/src/extension/extension/vscode-node/services.ts @@ -150,6 +150,10 @@ import { registerServices as registerCommonServices } from '../vscode/services'; import { PromptsServiceImpl } from '../../../platform/promptFiles/vscode-node/promptsServiceImpl'; import { IPromptsService } from '../../../platform/promptFiles/common/promptsService'; import { AutomaticInstructionsCollector, IAutomaticInstructionsCollector } from '../../../platform/promptFiles/node/automaticInstructionsCollector'; +import { GrepResultService, IGrepResultService } from '../../tools/node/grepResultService'; +import { IRegionContextProviderService } from '../../../platform/languageContextProvider/common/regionContextProvider'; +import { ContainerContextProviderService } from '../../typescriptContext/vscode-node/regionContextProvider'; + // ########################################################################################### // ### ### @@ -167,6 +171,8 @@ export function registerServices(builder: IInstantiationServiceBuilder, extensio builder.define(IAutomodeService, new SyncDescriptor(AutomodeService)); builder.define(IConversationStore, new SyncDescriptor(ConversationStore)); builder.define(IDiffService, new DiffServiceImpl()); + builder.define(IGrepResultService, new SyncDescriptor(GrepResultService)); + builder.define(IRegionContextProviderService, new SyncDescriptor(ContainerContextProviderService)); builder.define(ITokenizerProvider, new SyncDescriptor(TokenizerProvider, [true])); builder.define(IToolsService, new SyncDescriptor(ToolsService)); builder.define(IToolDeferralService, new ToolDeferralService()); diff --git a/extensions/copilot/src/extension/tools/node/findTextInFilesTool.tsx b/extensions/copilot/src/extension/tools/node/findTextInFilesTool.tsx index d4f1c27fee3128..f74dfcd84fed75 100644 --- a/extensions/copilot/src/extension/tools/node/findTextInFilesTool.tsx +++ b/extensions/copilot/src/extension/tools/node/findTextInFilesTool.tsx @@ -31,6 +31,7 @@ import { ToolName } from '../common/toolNames'; import { CopilotToolMode, ICopilotTool, ToolRegistry } from '../common/toolsRegistry'; import { checkCancellation, InputGlobResult, inputGlobToPattern, patternContainsWorkspaceFolderPath } from './toolUtils'; import { IExperimentationService } from '../../../lib/node/chatLibMain'; +import { IGrepResultService } from './grepResultService'; interface IFindTextInFilesToolParams { query: string; @@ -44,6 +45,7 @@ interface IFindTextInFilesToolParams { interface FileMatch { path: string; + uri: vscode.Uri; matches: vscode.TextSearchMatch2[]; elidedMatches?: number; } @@ -70,6 +72,7 @@ export class FindTextInFilesTool implements ICopilotTool, token: CancellationToken) { @@ -186,6 +189,9 @@ Then if you want to include those files you can call the tool again by setting " if (!groupedMatches) { return this.errorResult(noMatchInstructions ? `No matches found. ${noMatchInstructions}` : 'No matches found.'); } + if (options.chatRequestId !== undefined) { + this.grepResultService.addGrepResult(options.chatRequestId, groupedMatches); + } const prompt = await renderPromptElementJSON(this.instantiationService, FindTextInFilesGrepResult, { grouped: groupedMatches, query: options.input.query }, @@ -207,7 +213,7 @@ Then if you want to include those files you can call the tool again by setting " const path = this.promptPathRepresentationService.getFilePath(textMatch.uri, true); let fileMatch = groupedByFile.get(path); if (fileMatch === undefined) { - fileMatch = { path, matches: [] }; + fileMatch = { path, uri: textMatch.uri, matches: [] }; groupedByFile.set(path, fileMatch); } fileMatch.matches.push(textMatch); diff --git a/extensions/copilot/src/extension/tools/node/grepResultService.ts b/extensions/copilot/src/extension/tools/node/grepResultService.ts new file mode 100644 index 00000000000000..ef0bfd0db7235e --- /dev/null +++ b/extensions/copilot/src/extension/tools/node/grepResultService.ts @@ -0,0 +1,115 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import type * as vscode from 'vscode'; + +import { createServiceIdentifier } from '../../../util/common/services'; +import { LRUCache } from '../../../util/vs/base/common/map'; + +export const IGrepResultService = createServiceIdentifier('IGrepResultService'); + +interface FileMatch { + uri: vscode.Uri; + matches: vscode.TextSearchMatch2[]; +} + +interface MatchResult { + files: FileMatch[]; +} + +export interface IGrepResultService { + readonly _serviceBrand: undefined; + + addGrepResult(requestId: string, result: MatchResult): void; + getGrepResult(requestId: string, uri: vscode.Uri, startLine: number, endLine: number): vscode.Range[] | undefined; +} + +export class NullGrepResultService implements IGrepResultService { + declare readonly _serviceBrand: undefined; + + addGrepResult(requestId: string, result: MatchResult): void { + // No-op + } + + getGrepResult(requestId: string, uri: vscode.Uri, startLine: number, endLine: number): vscode.Range[] | undefined { + return undefined; + } +} + +interface Matches { + files: Map; +} + +export class GrepResultService implements IGrepResultService { + readonly _serviceBrand: undefined; + + private readonly cache: LRUCache; + + constructor() { + this.cache = new LRUCache(10); + } + + addGrepResult(requestId: string, result: MatchResult): void { + let matches: Matches | undefined = this.cache.get(requestId); + if (matches === undefined) { + matches = { files: new Map() }; + for (const file of result.files) { + matches.files.set(file.uri.toString(), file.matches.map(m => m.ranges[0].sourceRange)); + } + this.cache.set(requestId, matches); + } else { + for (const file of result.files) { + const existingRanges = matches.files.get(file.uri.toString()); + if (existingRanges === undefined) { + matches.files.set(file.uri.toString(), file.matches.map(m => m.ranges[0].sourceRange)); + } else { + const existingRangesSet = new Set(existingRanges.map(r => r.start.line)); + for (const match of file.matches) { + const line = match.ranges[0].sourceRange.start.line; + if (!existingRangesSet.has(line)) { + existingRanges.push(match.ranges[0].sourceRange); + existingRangesSet.add(line); + } + } + existingRanges.sort((a, b) => a.start.line - b.start.line); + matches.files.set(file.uri.toString(), existingRanges); + } + } + } + } + + getGrepResult(requestId: string, uri: vscode.Uri, startLine: number, endLine: number): vscode.Range[] | undefined { + const matches = this.cache.get(requestId); + if (!matches) { + return undefined; + } + const fileMatches = matches.files.get(uri.toString()); + if (!fileMatches) { + return undefined; + } + + let low = 0; + let high = fileMatches.length; + while (low < high) { + const mid = low + Math.floor((high - low) / 2); + if (fileMatches[mid].start.line < startLine) { + low = mid + 1; + } else { + high = mid; + } + } + + const result: vscode.Range[] = []; + for (let i = low; i < fileMatches.length; i++) { + const match = fileMatches[i]; + if (match.start.line > endLine) { + break; + } + result.push(match); + } + + return result; + } +} diff --git a/extensions/copilot/src/extension/tools/node/readFileTool.tsx b/extensions/copilot/src/extension/tools/node/readFileTool.tsx index a234134dcf1dc7..363800cf662222 100644 --- a/extensions/copilot/src/extension/tools/node/readFileTool.tsx +++ b/extensions/copilot/src/extension/tools/node/readFileTool.tsx @@ -36,6 +36,8 @@ import { ICopilotTool, ToolRegistry } from '../common/toolsRegistry'; import { formatUriForFileWidget } from '../common/toolUtils'; import { getImageMimeType } from './imageToolUtils'; import { assertFileNotContentExcluded, isFileExternalAndNeedsConfirmation, resolveToolInputPath } from './toolUtils'; +import { IGrepResultService } from './grepResultService'; +import { IRegionContextProviderService } from '../../../platform/languageContextProvider/common/regionContextProvider'; export const getReadFileV2Description = (orig: vscode.LanguageModelToolInformation): vscode.LanguageModelToolInformation => ({ name: ToolName.ReadFile, @@ -133,6 +135,8 @@ export class ReadFileTool implements ICopilotTool { @ICustomInstructionsService private readonly customInstructionsService: ICustomInstructionsService, @IFileSystemService private readonly fileSystemService: IFileSystemService, @IExtensionsService private readonly extensionsService: IExtensionsService, + @IGrepResultService private readonly grepResultService: IGrepResultService, + @IRegionContextProviderService private readonly regionContextProvider: IRegionContextProviderService ) { } async invoke(options: vscode.LanguageModelToolInvocationOptions, token: vscode.CancellationToken) { @@ -180,6 +184,41 @@ export class ReadFileTool implements ICopilotTool { const documentSnapshot = await this.getSnapshot(uri); ranges = getParamRanges(options.input, documentSnapshot); + const languageId = documentSnapshot.languageId; + if (options.chatRequestId !== undefined && uri.scheme === 'file' && (languageId === 'typescript' || languageId === 'javascript')) { + const startLine = ranges.start - 1; + const endLine = ranges.end - 1; + try { + const grepResultMatches = this.grepResultService.getGrepResult(options.chatRequestId, uri, startLine, endLine); + if (grepResultMatches !== undefined && grepResultMatches.length > 0 && documentSnapshot.version === documentSnapshot.document.version) { + const regions = await this.regionContextProvider.getRegions(documentSnapshot.uri, documentSnapshot.languageId, grepResultMatches, { start: startLine, end: endLine}); + if (regions !== undefined && regions.length > 0 && documentSnapshot.version === documentSnapshot.document.version) { + this.sendAdjustedRegionTelemetry(options, startLine, endLine, regions[0].range.start, regions[0].range.end); + // const saving = (ranges.end - ranges.start) - (regions[0].range.end - regions[0].range.start); + // this.logService.info(`Saving ${saving} lines reading ${documentSnapshot.uri.fsPath}. Requests [${ranges.start}-${ranges.end}], Grep matches: [${grepResultMatches.map(m => m.start.line + 1).join(',')}], region [${regions[0].range.start + 1}-${regions[0].range.end + 1}]`); + } else { + if (documentSnapshot.version === documentSnapshot.document.version) { + this.sendAdjustingFailedTelemetry(options, startLine, endLine, 'noGrepRegions'); + // this.logService.info(`No regions found for grep result match in file ${documentSnapshot.uri.fsPath} at lines [${grepResultMatches.map(m => m.start.line + 1).join(',')}]`); + } else { + this.sendAdjustingFailedTelemetry(options, startLine, endLine, 'documentVersionChanged'); + // this.logService.info(`Document version changed for requestId ${options.chatRequestId}`); + } + } + } else { + if (documentSnapshot.version === documentSnapshot.document.version) { + this.sendAdjustingFailedTelemetry(options, startLine, endLine, 'noGrep'); + // this.logService.info(`No grep result match found for requestId ${options.chatRequestId}`); + } else { + this.sendAdjustingFailedTelemetry(options, startLine, endLine, 'documentVersionChanged'); + // this.logService.info(`Document version changed for requestId ${options.chatRequestId}`); + } + } + } catch (err) { + this.sendAdjustingFailedTelemetry(options, startLine, endLine, 'exception'); + // this.logService.error(`Error processing grep result for requestId ${options.chatRequestId}: ${err}`); + } + } void this.sendReadFileTelemetry('success', options, ranges, uri, documentSnapshot); const useCodeFences = this.configurationService.getExperimentBasedConfig(ConfigKey.TeamInternal.ReadFileCodeFences, this.experimentationService); @@ -391,6 +430,51 @@ export class ReadFileTool implements ICopilotTool { } } + private async sendAdjustedRegionTelemetry(options: Pick, 'model' | 'chatRequestId' | 'input'>, originalStart: number, originalEnd: number, adjustedStart: number, adjustedEnd: number) { + /* __GDPR__ + "readFileRegionAdjusted" : { + "owner": "dbaeumer", + "comment": "Information about the clipping of the requested region to read", + "requestId": { "classification": "SystemMetaData", "purpose": "FeatureInsight", "comment": "The id of the current request turn." }, + "originalLines": { "classification": "SystemMetaData", "purpose": "FeatureInsight", "comment": "The number of original lines of the requested region", "isMeasurement": true }, + "adjustedLines": { "classification": "SystemMetaData", "purpose": "FeatureInsight", "comment": "The number of lines after the requested region has been adjusted", "isMeasurement": true }, + "deltaStart": { "classification": "SystemMetaData", "purpose": "FeatureInsight", "comment": "The difference between the original start line and the adjusted start line", "isMeasurement": true }, + "deltaEnd": { "classification": "SystemMetaData", "purpose": "FeatureInsight", "comment": "The difference between the original end line and the adjusted end line", "isMeasurement": true } + } + */ + this.telemetryService.sendMSFTTelemetryEvent('readFileRegionAdjusted', + { + requestId: options.chatRequestId, + }, + { + originalLines: originalEnd - originalStart + 1, + adjustedLines: adjustedEnd - adjustedStart + 1, + deltaStart: adjustedStart - originalStart, + deltaEnd: originalEnd - adjustedEnd, + } + ); + } + + private async sendAdjustingFailedTelemetry(options: Pick, 'model' | 'chatRequestId' | 'input'>, startLine: number, endLine: number, reason: 'noGrep' | 'noGrepRegions' | 'documentVersionChanged' | 'exception') { + /* __GDPR__ + "readFileRegionAdjustingFailed" : { + "owner": "dbaeumer", + "comment": "Information about the failure to adjust the requested region to read", + "requestId": { "classification": "SystemMetaData", "purpose": "FeatureInsight", "comment": "The id of the current request turn." }, + "lines": { "classification": "SystemMetaData", "purpose": "FeatureInsight", "comment": "The number of line to read", "isMeasurement": true }, + "reason": { "classification": "SystemMetaData", "purpose": "FeatureInsight", "comment": "The reason why adjusting the requested region failed" } + } + */ + this.telemetryService.sendMSFTTelemetryEvent('readFileRegionAdjustingFailed', + { + requestId: options.chatRequestId, + reason, + }, { + lines: endLine - startLine + 1 + } + ); + } + async resolveInput(input: IReadFileParamsV1, promptContext: IBuildPromptContext): Promise { this._promptContext = promptContext; return input; diff --git a/extensions/copilot/src/extension/tools/node/test/findTextInFilesResult.spec.tsx b/extensions/copilot/src/extension/tools/node/test/findTextInFilesResult.spec.tsx index 0e8255bd90e226..428b1948fc17eb 100644 --- a/extensions/copilot/src/extension/tools/node/test/findTextInFilesResult.spec.tsx +++ b/extensions/copilot/src/extension/tools/node/test/findTextInFilesResult.spec.tsx @@ -207,6 +207,7 @@ suite('FindTextInFilesGrepResult', () => { files: [ { path: '/src/a.ts', + uri: URI.file('/src/a.ts'), matches: [lineMatch(URI.file('/src/a.ts'), 5, 'const a = 1;'), lineMatch(URI.file('/src/a.ts'), 9, 'const b = 2;')], }, ], @@ -225,10 +226,12 @@ suite('FindTextInFilesGrepResult', () => { files: [ { path: '/src/a.ts', + uri: URI.file('/src/a.ts'), matches: [lineMatch(URI.file('/src/a.ts'), 5, 'alpha')], }, { path: '/src/b.ts', + uri: URI.file('/src/b.ts'), matches: [lineMatch(URI.file('/src/b.ts'), 1, 'beta'), lineMatch(URI.file('/src/b.ts'), 3, 'gamma')], elidedMatches: 1, }, @@ -256,6 +259,7 @@ suite('FindTextInFilesGrepResult', () => { files: [ { path: '/src/big.ts', + uri: URI.file('/src/big.ts'), matches: [{ uri, previewText, @@ -284,6 +288,7 @@ suite('FindTextInFilesGrepResult', () => { files: [ { path: '/src/big.ts', + uri: URI.file('/src/big.ts'), matches: [{ uri, previewText, diff --git a/extensions/copilot/src/extension/tools/node/test/findTextInFilesTool.spec.tsx b/extensions/copilot/src/extension/tools/node/test/findTextInFilesTool.spec.tsx index 062d53f50d3c92..d7933eadb1d811 100644 --- a/extensions/copilot/src/extension/tools/node/test/findTextInFilesTool.spec.tsx +++ b/extensions/copilot/src/extension/tools/node/test/findTextInFilesTool.spec.tsx @@ -309,4 +309,4 @@ class RecordingSearchService extends AbstractSearchService { override async findFiles(filePattern: vscode.GlobPattern, options?: vscode.FindFiles2Options | undefined, token?: vscode.CancellationToken | undefined): Promise { throw new Error('Method not implemented.'); } -} \ No newline at end of file +} diff --git a/extensions/copilot/src/extension/tools/node/test/grepResultService.spec.ts b/extensions/copilot/src/extension/tools/node/test/grepResultService.spec.ts new file mode 100644 index 00000000000000..7786a1b017c60a --- /dev/null +++ b/extensions/copilot/src/extension/tools/node/test/grepResultService.spec.ts @@ -0,0 +1,45 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import type * as vscode from 'vscode'; +import { expect, suite, test } from 'vitest'; +import { URI } from '../../../../util/vs/base/common/uri'; +import { Range } from '../../../../vscodeTypes'; +import { GrepResultService, NullGrepResultService } from '../grepResultService'; + +suite('GrepResultService', () => { + const uri = URI.file('/file.ts'); + + function createMatch(range: vscode.Range): vscode.TextSearchMatch2 { + return { + uri, + previewText: '', + ranges: [{ + previewRange: range, + sourceRange: range, + }] + }; + } + + test('returns all ranges within the inclusive line bounds', () => { + const before = new Range(3, 0, 3, 1); + const first = new Range(4, 2, 4, 5); + const second = new Range(8, 1, 8, 7); + const after = new Range(9, 0, 9, 1); + const service = new GrepResultService(); + service.addGrepResult('request', { + files: [{ uri, matches: [before, first, second, after].map(createMatch) }] + }); + + expect(service.getGrepResult('request', uri, 4, 8)).toEqual([first, second]); + }); + + test('returns undefined when no results are available', () => { + const service = new GrepResultService(); + + expect(service.getGrepResult('unknown', uri, 0, 10)).toBeUndefined(); + expect(new NullGrepResultService().getGrepResult('request', uri, 0, 10)).toBeUndefined(); + }); +}); diff --git a/extensions/copilot/src/extension/typescriptContext/common/serverProtocol.ts b/extensions/copilot/src/extension/typescriptContext/common/serverProtocol.ts index 86ba81a2439fd0..462cf6684a83da 100644 --- a/extensions/copilot/src/extension/typescriptContext/common/serverProtocol.ts +++ b/extensions/copilot/src/extension/typescriptContext/common/serverProtocol.ts @@ -45,6 +45,17 @@ export type Range = { end: Position; }; +export type LineRange = { + start: number; + end: number; +}; + +export type Region = { + kind: string; + name?: string; + range: LineRange; +}; + export type WithinRangeCacheScope = { kind: CacheScopeKind.WithinRange; range: Range; @@ -439,6 +450,35 @@ export namespace CustomResponse { } } +export interface RegionContextRequestArgs extends tt.server.protocol.FileLocationRequestArgs { + ranges: readonly Range[]; + requested?: LineRange; +} + +export interface RegionContextRequest extends tt.server.protocol.Request { + arguments?: RegionContextRequestArgs; +} + +export namespace RegionContextResponse { + export type OK = { + regions: Region[]; + }; + + export type Failed = CustomResponse.Failed; + + export function isOk(response: RegionContextResponse | undefined): response is Omit & { body: OK } { + return response?.type === 'response' && Array.isArray((response.body as OK | undefined)?.regions); + } + + export function isError(response: RegionContextResponse | undefined): response is Omit & { body: Failed } { + return response?.type === 'response' && CustomResponse.isError(response); + } +} + +export type RegionContextResponse = (tt.server.protocol.Response & { + body: RegionContextResponse.OK | RegionContextResponse.Failed; +}) | { type: 'cancelled' }; + export namespace ComputeContextResponse { export type OK = ContextRequestResult; diff --git a/extensions/copilot/src/extension/typescriptContext/serverPlugin/src/common/protocol.ts b/extensions/copilot/src/extension/typescriptContext/serverPlugin/src/common/protocol.ts index df8032f437f79f..5240908cc9b42a 100644 --- a/extensions/copilot/src/extension/typescriptContext/serverPlugin/src/common/protocol.ts +++ b/extensions/copilot/src/extension/typescriptContext/serverPlugin/src/common/protocol.ts @@ -45,6 +45,17 @@ export type Range = { end: Position; }; +export type LineRange = { + start: number; + end: number; +}; + +export type Region = { + kind: string; + name?: string; + range: LineRange; +}; + export type WithinRangeCacheScope = { kind: CacheScopeKind.WithinRange; range: Range; @@ -439,6 +450,35 @@ export namespace CustomResponse { } } +export interface RegionContextRequestArgs extends tt.server.protocol.FileLocationRequestArgs { + ranges: readonly Range[]; + requested?: LineRange; +} + +export interface RegionContextRequest extends tt.server.protocol.Request { + arguments?: RegionContextRequestArgs; +} + +export namespace RegionContextResponse { + export type OK = { + regions: Region[]; + }; + + export type Failed = CustomResponse.Failed; + + export function isOk(response: RegionContextResponse | undefined): response is Omit & { body: OK } { + return response?.type === 'response' && Array.isArray((response.body as OK | undefined)?.regions); + } + + export function isError(response: RegionContextResponse | undefined): response is Omit & { body: Failed } { + return response?.type === 'response' && CustomResponse.isError(response); + } +} + +export type RegionContextResponse = (tt.server.protocol.Response & { + body: RegionContextResponse.OK | RegionContextResponse.Failed; +}) | { type: 'cancelled' }; + export namespace ComputeContextResponse { export type OK = ContextRequestResult; diff --git a/extensions/copilot/src/extension/typescriptContext/serverPlugin/src/common/regionContextProvider.ts b/extensions/copilot/src/extension/typescriptContext/serverPlugin/src/common/regionContextProvider.ts new file mode 100644 index 00000000000000..f225ce4dea312e --- /dev/null +++ b/extensions/copilot/src/extension/typescriptContext/serverPlugin/src/common/regionContextProvider.ts @@ -0,0 +1,283 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ +import type tt from 'typescript/lib/tsserverlibrary'; +import TS from './typescript'; +const ts = TS(); + +import type { LineRange, Range, Region } from './protocol'; +import tss from './typescripts'; + +type StructuralEntity = { kind: string; name?: string; rangeNode: tt.Node | [tt.Node, tt.Node]; includeJsDoc?: boolean; continueWith?: tt.Node }; + +export class RegionContextProvider { + + public getRegions(sourceFile: tt.SourceFile, ranges: readonly Range[], requested?: LineRange | undefined): Region[] | undefined { + if (ranges.length === 0) { + return undefined; + } + + if (ranges.length === 1) { + return this.findEnclosingScopes(sourceFile, ranges[0].start.line, ranges[0].start.character, requested); + } + + const containersList: Region[][] = []; + for (const range of ranges) { + const containers = this.findEnclosingScopes(sourceFile, range.start.line, range.start.character, requested); + if (containers !== undefined && containers.length > 0) { + containersList.push(containers.reverse()); + } + } + if (containersList.length === 0) { + return undefined; + } + + const longestContainers = containersList.reduce((longest, containers) => containers.length > longest.length ? containers : longest); + const commonContainers = longestContainers.slice(); + for (const containers of containersList) { + if (containers === longestContainers) { + continue; + } + let commonLength = 0; + while (commonLength < commonContainers.length && commonLength < containers.length) { + const commonContainer = commonContainers[commonLength]; + const container = containers[commonLength]; + if (commonContainer.kind !== container.kind + || commonContainer.name !== container.name + || commonContainer.range.start !== container.range.start + || commonContainer.range.end !== container.range.end) { + break; + } + commonLength++; + } + commonContainers.length = commonLength; + } + + const tailContainers = containersList.map(containers => containers[containers.length - 1]); + if (tailContainers.length > 0) { + const container: Region = { + kind: 'merged', + range: { + start: Math.min(...tailContainers.map(container => container.range.start)), + end: Math.max(...tailContainers.map(container => container.range.end)) + } + }; + const lastContainer = commonContainers[commonContainers.length - 1]; + if (lastContainer !== undefined && container.range.end - container.range.start < lastContainer.range.end - lastContainer.range.start) { + commonContainers.push(container); + } + } + + return commonContainers.reverse(); + } + + private findEnclosingScopes(sourceFile: tt.SourceFile, line: number, column: number, requested?: LineRange | undefined): Region[] | undefined { + const position = sourceFile.getPositionOfLineAndCharacter(line, column); + const tokenInfo = tss.getRelevantTokens(sourceFile, position); + const node = tokenInfo.touching ?? tokenInfo.token; + if (node === undefined) { + return undefined; + } + + const result: Region[] = []; + for (let current: tt.Node | undefined = node; current; current = current.parent) { + if (ts.isSourceFile(current)) { + const endLine = sourceFile.getLineAndCharacterOfPosition(sourceFile.getEnd()).line; + result.push({ + kind: 'sourceFile', + name: this.getBaseFileName(sourceFile.fileName), + range: { start: 0, end: endLine } + }); + break; + } + + const structuralEntity = this.getStructuralEntity(sourceFile, current, requested); + if (structuralEntity !== undefined) { + const { kind, name, rangeNode, includeJsDoc, continueWith } = structuralEntity; + const rangeStartNode = Array.isArray(rangeNode) ? rangeNode[0] : rangeNode; + const rangeEndNode = Array.isArray(rangeNode) ? rangeNode[1] : rangeNode; + result.push({ + kind, + name, + range: { + start: sourceFile.getLineAndCharacterOfPosition(rangeStartNode.getStart(sourceFile, includeJsDoc)).line, + end: sourceFile.getLineAndCharacterOfPosition(rangeEndNode.getEnd()).line + } + }); + current = continueWith ?? current; + } + } + return result.length > 0 ? result : undefined; + } + + private getStructuralEntity(sourceFile: tt.SourceFile, node: tt.Node, requested?: LineRange | undefined): StructuralEntity | undefined { + const parent = node.parent; + let name: string | undefined; + switch (node.kind) { + case ts.SyntaxKind.JSDoc: { + const parentEntity = this.getStructuralEntity(sourceFile, parent, requested); + if (parentEntity !== undefined) { + parentEntity.includeJsDoc = true; + parentEntity.continueWith ??= parent; + } + return parentEntity; + } + case ts.SyntaxKind.ImportDeclaration: + name = (node as tt.ImportDeclaration).moduleSpecifier.getText(); + return { kind: 'import', name, rangeNode: node }; + case ts.SyntaxKind.ExportDeclaration: + name = (node as tt.ExportDeclaration).moduleSpecifier?.getText(); + return { kind: 'export', name, rangeNode: node }; + case ts.SyntaxKind.FunctionDeclaration: + name = (node as tt.FunctionDeclaration).name?.text; + if (name === undefined) { + if (ts.isPropertyAssignment(parent) && ts.isIdentifier(parent.name)) { + name = parent.name.text; + } else if (ts.isVariableDeclaration(parent) && ts.isIdentifier(parent.name)) { + name = parent.name.text; + } + } + return { kind: 'function', name, rangeNode: node }; + case ts.SyntaxKind.Constructor: + return { kind: 'constructor', name: 'constructor', rangeNode: node }; + case ts.SyntaxKind.MethodDeclaration: + name = (node as tt.MethodDeclaration).name.getText(); + return { kind: 'method', name, rangeNode: node }; + case ts.SyntaxKind.MethodSignature: + name = (node as tt.MethodSignature).name.getText(); + return { kind: 'method', name, rangeNode: node }; + case ts.SyntaxKind.ArrowFunction: + if (ts.isPropertyAssignment(parent) && ts.isIdentifier(parent.name)) { + name = parent.name.text; + return { kind: 'function', name, rangeNode: parent }; + } else if (ts.isVariableDeclaration(parent) && ts.isIdentifier(parent.name)) { + name = parent.name.text; + return { kind: 'arrow-function', name, rangeNode: parent }; + } else if (ts.isCallExpression(parent)) { + return { kind: 'arrow-function', rangeNode: parent }; + } + return { kind: 'arrow-function', rangeNode: node }; + case ts.SyntaxKind.PropertyDeclaration: + return this.handleProperty(sourceFile, node as tt.PropertyDeclaration, requested); + case ts.SyntaxKind.PropertyAssignment: + return this.handleProperty(sourceFile, node as tt.PropertyAssignment, requested); + case ts.SyntaxKind.PropertySignature: + return this.handleProperty(sourceFile, node as tt.PropertySignature, requested); + case ts.SyntaxKind.GetAccessor: + name = (node as tt.GetAccessorDeclaration).name.getText(); + return { kind: 'getter', name, rangeNode: node }; + case ts.SyntaxKind.SetAccessor: + name = (node as tt.SetAccessorDeclaration).name.getText(); + return { kind: 'setter', name, rangeNode: node }; + case ts.SyntaxKind.ClassDeclaration: + name = (node as tt.ClassDeclaration).name?.text; + return { kind: 'class', name, rangeNode: node }; + case ts.SyntaxKind.InterfaceDeclaration: + name = (node as tt.InterfaceDeclaration).name.text; + return { kind: 'interface', name, rangeNode: node }; + case ts.SyntaxKind.ModuleDeclaration: + name = (node as tt.ModuleDeclaration).name.text; + return { kind: 'module', name, rangeNode: node }; + case ts.SyntaxKind.TypeAliasDeclaration: + name = (node as tt.TypeAliasDeclaration).name.text; + return { kind: 'type-alias', name, rangeNode: node }; + default: + return undefined; + } + } + + private handleProperty(sourceFile: tt.SourceFile, node: tt.PropertyDeclaration | tt.PropertyAssignment | tt.PropertySignature, requested?: LineRange | undefined): StructuralEntity | undefined { + const name = node.name.getText(); + if (ts.isPropertyDeclaration(node) || ts.isPropertyAssignment(node)) { + const initializeKind = node.initializer?.kind; + if (initializeKind === ts.SyntaxKind.FunctionType || initializeKind === ts.SyntaxKind.FunctionDeclaration || initializeKind === ts.SyntaxKind.FunctionExpression || initializeKind === ts.SyntaxKind.ArrowFunction) { + return { kind: 'function', name, rangeNode: node }; + } + } + const parent = node.parent; + if (requested !== undefined) { + const info = this.getMemberInfo(parent); + if (info === undefined) { + return undefined; + } + const { items, kind, memberKind, name } = info; + const range = this.calculateRange(sourceFile, parent, node, items, requested); + if (range === undefined) { + return undefined; + } + if (Array.isArray(range)) { + const [startIndex, endIndex] = range; + return { + kind: memberKind, + name, + rangeNode: [items[startIndex], items[endIndex]], + continueWith: parent + }; + } else { + return { + kind, + name, + rangeNode: parent, + continueWith: parent + }; + } + } + return undefined; + } + + private getMemberInfo(parent: tt.ClassLikeDeclaration | tt.ObjectLiteralExpression| tt.InterfaceDeclaration | tt.TypeLiteralNode): { items: tt.NodeArray; kind: string; memberKind: string; name?: string | undefined } | undefined { + if (ts.isClassDeclaration(parent)) { + return { items: parent.members, kind: 'class', memberKind: 'class-members', name: parent.name?.text }; + } else if (ts.isInterfaceDeclaration(parent)) { + return { items: parent.members, kind: 'interface', memberKind: 'interface-members', name: parent.name?.text }; + } else if (ts.isObjectLiteralExpression(parent)) { + return { items: parent.properties, kind: 'object-literal', memberKind: 'object-literal-members' }; + } else if (ts.isTypeLiteralNode(parent)) { + return { items: parent.members, kind: 'type-literal', memberKind: 'type-literal-members' }; + } + return undefined; + } + + private calculateRange(sourceFile: tt.SourceFile, parent: tt.Node, node: tt.Node, items: tt.NodeArray, requested: LineRange): [number, number] | tt.Node | undefined { + const startLine = sourceFile.getLineAndCharacterOfPosition(parent.getStart(sourceFile)).line; + const endLine = sourceFile.getLineAndCharacterOfPosition(parent.getEnd()).line; + if (requested.start <= startLine && requested.end >= endLine) { + return parent; + } + + const index = items.indexOf(node); + if (index === -1) { + return undefined; + } + + let startIndex = Math.max(0, index - 1); + while (index - startIndex < 3 && startIndex > 0) { + const member = items[startIndex - 1]; + if (!this.isInsideRequestedRange(sourceFile, member, requested)) { + break; + } + startIndex--; + } + + let endIndex = Math.min(items.length - 1, index + 1); + while (endIndex - index < 3 && endIndex < items.length - 1) { + const member = items[endIndex + 1]; + if (!this.isInsideRequestedRange(sourceFile, member, requested)) { + break; + } + endIndex++; + } + return [startIndex, endIndex]; + } + + private isInsideRequestedRange(sourceFile: tt.SourceFile, member: tt.Node, requested: LineRange): boolean { + const memberStartLine = sourceFile.getLineAndCharacterOfPosition(member.getStart(sourceFile)).line; + const memberEndLine = sourceFile.getLineAndCharacterOfPosition(member.getEnd()).line; + return requested.start <= memberStartLine && requested.end >= memberEndLine; + } + + private getBaseFileName(fileName: string): string { + return fileName.substring(Math.max(fileName.lastIndexOf('/'), fileName.lastIndexOf('\\')) + 1); + } +} diff --git a/extensions/copilot/src/extension/typescriptContext/serverPlugin/src/node/create.ts b/extensions/copilot/src/extension/typescriptContext/serverPlugin/src/node/create.ts index 44799003358f6f..8872e1fc05f3c6 100644 --- a/extensions/copilot/src/extension/typescriptContext/serverPlugin/src/node/create.ts +++ b/extensions/copilot/src/extension/typescriptContext/serverPlugin/src/node/create.ts @@ -5,7 +5,8 @@ import type tt from 'typescript/lib/tsserverlibrary'; import { computeContext, nesRename, prepareNesRename } from '../common/api'; import { CharacterBudget, ComputeContextSession, ContextResult, NullLogger, RequestContext, TokenBudgetExhaustedError, type Logger } from '../common/contextProvider'; -import { ErrorCode, RenameKind, type CachedContextRunnableResult, type ComputeContextRequest, type ComputeContextResponse, type ContextRunnableResultId, type CustomResponse, type NesRenameRequest, type NesRenameResponse, type PingResponse, type PrepareNesRenameRequest, type PrepareNesRenameResponse, type Range, type RenameGroup } from '../common/protocol'; +import { ErrorCode, RenameKind, type CachedContextRunnableResult, type ComputeContextRequest, type ComputeContextResponse, type ContextRunnableResultId, type CustomResponse, type NesRenameRequest, type NesRenameResponse, type PingResponse, type PrepareNesRenameRequest, type PrepareNesRenameResponse, type Range, type RegionContextRequest, type RegionContextResponse, type RenameGroup } from '../common/protocol'; +import { RegionContextProvider } from '../common/regionContextProvider'; import { CancellationTokenWithTimer, Sessions } from '../common/typescripts'; const ts = TS(); @@ -101,6 +102,10 @@ interface NesRenameHandlerResponse extends tt.server.HandlerResponse { response: NesRenameResponse.OK | NesRenameResponse.Failed; } +interface RegionContextHandlerResponse extends tt.server.HandlerResponse { + response: RegionContextResponse.OK | RegionContextResponse.Failed; +} + let installAttempted: boolean = false; let languageServerSession: LanguageServerSession | undefined = undefined; let languageServiceHost: tt.LanguageServiceHost | undefined = undefined; @@ -202,6 +207,24 @@ const computeContextHandler = (request: ComputeContextRequest): ComputeContextHa return { response: result.toJson(), responseRequired: true }; }; +const regionContextHandler = (request: RegionContextRequest): RegionContextHandlerResponse => { + const input = resolveInput(request.arguments, 0); + if (FailedHandlerResponse.is(input)) { + return input; + } + + try { + const sourceFile = input.program.getSourceFile(input.file); + const regions = sourceFile === undefined ? [] : new RegionContextProvider().getRegions(sourceFile, request.arguments!.ranges, request.arguments!.requested) ?? []; + return { response: { regions }, responseRequired: true }; + } catch (error) { + if (error instanceof Error) { + return { response: { error: ErrorCode.exception, message: error.message, stack: error.stack }, responseRequired: true }; + } + return { response: { error: ErrorCode.exception, message: 'Unknown error' }, responseRequired: true }; + } +}; + const prepareNesRenameHandler = (request: PrepareNesRenameRequest): PrepareNesRenameHandlerResponse => { const input = resolveInput(request.arguments, 50); if (FailedHandlerResponse.is(input)) { @@ -271,6 +294,7 @@ export function create(info: tt.server.PluginCreateInfo): tt.LanguageService { languageServerSession = new LanguageServerSession(info.session, info.languageServiceHost, new NodeHost()); languageServiceHost = info.languageServiceHost; info.session.addProtocolHandler('_.copilot.context', computeContextHandler); + info.session.addProtocolHandler('_.copilot.regionContext', regionContextHandler); info.session.addProtocolHandler('_.copilot.prepareNesRename', prepareNesRenameHandler); info.session.addProtocolHandler('_.copilot.postNesRename', nesRenameHandler); } @@ -310,4 +334,4 @@ function isSupportedVersion(): boolean { } catch (e) { return false; } -} \ No newline at end of file +} diff --git a/extensions/copilot/src/extension/typescriptContext/serverPlugin/src/node/test/regionContext.spec.ts b/extensions/copilot/src/extension/typescriptContext/serverPlugin/src/node/test/regionContext.spec.ts new file mode 100644 index 00000000000000..9ebda67e380363 --- /dev/null +++ b/extensions/copilot/src/extension/typescriptContext/serverPlugin/src/node/test/regionContext.spec.ts @@ -0,0 +1,84 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ +import assert from 'assert'; +import { beforeAll, suite, test } from 'vitest'; + +import ts from 'typescript'; + +import type { LineRange, Range, Region } from '../../common/protocol'; +import type * as regionContextProvider from '../../common/regionContextProvider'; + +let RegionContextProvider: typeof regionContextProvider.RegionContextProvider; + +beforeAll(async () => { + const TS = await import('../../common/typescript'); + TS.default.install(ts); + RegionContextProvider = (await import('../../common/regionContextProvider')).RegionContextProvider; +}); + +function getRegionContext(sourceFile: ts.SourceFile, ranges: readonly Range[], requested?: LineRange): Region[] | undefined { + return new RegionContextProvider().getRegions(sourceFile, ranges, requested); +} + +function range(line: number, character: number = 0): Range { + return { + start: { line, character }, + end: { line, character } + }; +} + +suite('Region context', () => { + test('returns enclosing structural regions', () => { + const sourceFile = ts.createSourceFile('C:\\workspace\\regions.ts', [ + 'class Container {', + '\tmethod(): void {', + '\t\tconst callback = () => {', + '\t\t\treturn;', + '\t\t};', + '\t}', + '}', + ].join('\n'), ts.ScriptTarget.Latest, true); + + assert.deepStrictEqual(getRegionContext(sourceFile, [range(3)]), [ + { kind: 'arrow-function', name: 'callback', range: { start: 2, end: 4 } }, + { kind: 'method', name: 'method', range: { start: 1, end: 5 } }, + { kind: 'class', name: 'Container', range: { start: 0, end: 6 } }, + { kind: 'sourceFile', name: 'regions.ts', range: { start: 0, end: 6 } }, + ] satisfies Region[]); + }); + + test('merges distinct innermost regions', () => { + const sourceFile = ts.createSourceFile('regions.ts', [ + 'class Container {', + '\tfirst(): void {', + '\t\treturn;', + '\t}', + '\tsecond(): void {', + '\t\treturn;', + '\t}', + '}', + ].join('\n'), ts.ScriptTarget.Latest, true); + + assert.deepStrictEqual(getRegionContext(sourceFile, [range(2), range(5)]), [ + { kind: 'merged', range: { start: 1, end: 6 } }, + { kind: 'class', name: 'Container', range: { start: 0, end: 7 } }, + { kind: 'sourceFile', name: 'regions.ts', range: { start: 0, end: 7 } }, + ] satisfies Region[]); + }); + + test('groups property signatures within the requested range', () => { + const sourceFile = ts.createSourceFile('regions.ts', [ + 'interface Result {', + '\tvalue: number;', + '\tmessage: string;', + '}', + ].join('\n'), ts.ScriptTarget.Latest, true); + + assert.deepStrictEqual(getRegionContext(sourceFile, [range(1, 1), range(2, 1)], { start: 1, end: 2 }), [ + { kind: 'interface-members', name: 'Result', range: { start: 1, end: 2 } }, + { kind: 'sourceFile', name: 'regions.ts', range: { start: 0, end: 3 } }, + ] satisfies Region[]); + }); +}); diff --git a/extensions/copilot/src/extension/typescriptContext/vscode-node/languageContextService.ts b/extensions/copilot/src/extension/typescriptContext/vscode-node/languageContextService.ts index 30d94eff4d5034..df95c45bf4ee9c 100644 --- a/extensions/copilot/src/extension/typescriptContext/vscode-node/languageContextService.ts +++ b/extensions/copilot/src/extension/typescriptContext/vscode-node/languageContextService.ts @@ -7,6 +7,7 @@ import * as vscode from 'vscode'; import { ConfigKey, IConfigurationService } from '../../../platform/configuration/common/configurationService'; import { Copilot } from '../../../platform/inlineCompletions/common/api'; +import { IRegionContextProviderService } from '../../../platform/languageContextProvider/common/regionContextProvider'; import { ILanguageContextProviderService, ProviderTarget } from '../../../platform/languageContextProvider/common/languageContextProviderService'; import { ContextKind, ILanguageContextService, KnownSources, TriggerKind, type ContextItem, type RequestContext } from '../../../platform/languageServer/common/languageContextService'; import { ILogService } from '../../../platform/log/common/logService'; @@ -18,7 +19,7 @@ import { generateUuid } from '../../../util/vs/base/common/uuid'; import { InspectorDataProvider } from './inspector'; import { ThrottledDebouncer } from './throttledDebounce'; import { ContextItemSummary, ErrorLocation, ErrorPart, type OnCachePopulatedEvent, type OnContextComputedEvent, type OnContextComputedOnTimeoutEvent } from './types'; -import { TS6LanguageContextService } from './tsc6/tsContextService'; +import { TS6LanguageContextService } from './ts6/tsContextService'; import { TS7LanguageContextService } from './ts7/tsContextService'; import { currentTokenBudget, NullTSLanguageContextService, type TSLanguageContextService } from './tsContextService'; import { TypeScript } from './tsService'; @@ -333,6 +334,7 @@ export class InlineCompletionContribution implements vscode.Disposable, TokenBud @ILogService private readonly logService: ILogService, @ILanguageContextService private readonly languageContextService: ILanguageContextService, @ILanguageContextProviderService private readonly languageContextProviderService: ILanguageContextProviderService, + @IRegionContextProviderService private readonly containerContextProviderService: IRegionContextProviderService, ) { this.registrations = undefined; this.telemetrySender = new TelemetrySender(telemetryService, logService); @@ -346,6 +348,22 @@ export class InlineCompletionContribution implements vscode.Disposable, TokenBud })); this.disposables.add(vscode.window.registerTreeDataProvider('context-inspector', new InspectorDataProvider(languageContextService))); } + this.disposables.add(vscode.commands.registerCommand('github.copilot.debug.logTypeScriptContainers', async () => { + const editor = vscode.window.activeTextEditor; + const languageId = editor?.document.languageId; + if (!editor || (languageId !== 'typescript' && languageId !== 'typescriptreact' && languageId !== 'javascript' && languageId !== 'javascriptreact')) { + return; + } + + const positions = editor.selections.map(selection => selection.active); + const containers = await this.containerContextProviderService.getRegions( + editor.document.uri, + editor.document.languageId, + positions.map(position => new vscode.Range(position, position)) + ); + const locations = positions.map(position => `${editor.document.uri.toString()}:${position.line + 1}:${position.character + 1}`).join(', '); + this.logService.info(`[ContainerContextProvider] Containers at ${locations}: ${JSON.stringify(containers, undefined, 2)}`); + })); // Check if there are any TypeScript files open in the workspace. const open = vscode.workspace.textDocuments.some((document) => document.languageId === 'typescript' || document.languageId === 'typescriptreact'); diff --git a/extensions/copilot/src/extension/typescriptContext/vscode-node/nesRenameService.ts b/extensions/copilot/src/extension/typescriptContext/vscode-node/nesRenameService.ts index 95d6a71ea163da..954fbc4c51848a 100644 --- a/extensions/copilot/src/extension/typescriptContext/vscode-node/nesRenameService.ts +++ b/extensions/copilot/src/extension/typescriptContext/vscode-node/nesRenameService.ts @@ -9,7 +9,7 @@ import { ITelemetryService } from '../../../platform/telemetry/common/telemetry' import { DisposableStore } from '../../../util/vs/base/common/lifecycle'; import * as protocol from '../common/serverProtocol'; import { TS7NesRenameService } from './ts7/nesRenameService'; -import { TS6NesRenameService } from './tsc6/nesRenameService'; +import { TS6NesRenameService } from './ts6/nesRenameService'; import { TypeScript } from './tsService'; type TextChange = { diff --git a/extensions/copilot/src/extension/typescriptContext/vscode-node/regionContextProvider.ts b/extensions/copilot/src/extension/typescriptContext/vscode-node/regionContextProvider.ts new file mode 100644 index 00000000000000..9ba4f7d05814ed --- /dev/null +++ b/extensions/copilot/src/extension/typescriptContext/vscode-node/regionContextProvider.ts @@ -0,0 +1,74 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ +import type * as vscode from 'vscode'; + +import { type IRegionContextProviderService, type Region, type LineRange, NullRegionContextProviderService } from '../../../platform/languageContextProvider/common/regionContextProvider'; +import { ConfigKey, IConfigurationService } from '../../../platform/configuration/common/configurationService'; +import { ILogService } from '../../../platform/log/common/logService'; +import { TypeScript } from './tsService'; +import { TS7RegionContextProvider } from './ts7/regionContextProvider'; +import { TS6RegionContextProvider } from './ts6/regionContextProvider'; +import { DisposableStore } from '../../../util/vs/base/common/lifecycle'; + +export class ContainerContextProviderService implements IRegionContextProviderService { + + readonly _serviceBrand: undefined; + + private readonly disposables: DisposableStore; + private provider: Omit; + + constructor( + @ILogService private readonly logService: ILogService, + @IConfigurationService private readonly configurationService: IConfigurationService + ) { + this.disposables = new DisposableStore(); + this.disposables.add(this.configurationService.onDidChangeConfiguration(e => { + if (TypeScript.affectsVersion(e) || e.affectsConfiguration(ConfigKey.TypeScript7LanguageContext.fullyQualifiedId)) { + this.updateProvider(); + } + })); + this.provider = this.createProvider(); + } + + dispose(): void { + this.provider.dispose(); + this.disposables.dispose(); + } + + getRegions(document: vscode.Uri, languageId: string, ranges: vscode.Range[], requested?: LineRange): Promise { + return this.provider.getRegions(document, languageId, ranges, requested); + } + + private createProvider(): Omit { + if (!TypeScript.runsVersion7()) { + return new TS6RegionContextProvider(); + } + return TypeScript.isVersion7SupportEnabled(this.configurationService) + ? new TS7RegionContextProvider(this.logService) + : new NullRegionContextProviderService(); + } + + private updateProvider(): void { + const runsTS7 = TypeScript.runsVersion7(); + const enableTS7 = TypeScript.isVersion7SupportEnabled(this.configurationService); + const oldProvider = this.provider; + if (runsTS7) { + if (oldProvider instanceof TS6RegionContextProvider) { + this.provider = enableTS7 + ? new TS7RegionContextProvider(this.logService) + : new NullRegionContextProviderService(); + } else if (oldProvider instanceof TS7RegionContextProvider && !enableTS7) { + this.provider = new NullRegionContextProviderService(); + } else if (oldProvider instanceof NullRegionContextProviderService && enableTS7) { + this.provider = new TS7RegionContextProvider(this.logService); + } + } else if (!(oldProvider instanceof TS6RegionContextProvider)) { + this.provider = new TS6RegionContextProvider(); + } + if (oldProvider !== this.provider) { + oldProvider.dispose(); + } + } +} diff --git a/extensions/copilot/src/extension/typescriptContext/vscode-node/tsc6/nesRenameService.ts b/extensions/copilot/src/extension/typescriptContext/vscode-node/ts6/nesRenameService.ts similarity index 100% rename from extensions/copilot/src/extension/typescriptContext/vscode-node/tsc6/nesRenameService.ts rename to extensions/copilot/src/extension/typescriptContext/vscode-node/ts6/nesRenameService.ts diff --git a/extensions/copilot/src/extension/typescriptContext/vscode-node/ts6/regionContextProvider.ts b/extensions/copilot/src/extension/typescriptContext/vscode-node/ts6/regionContextProvider.ts new file mode 100644 index 00000000000000..3d8daa0a59f5fd --- /dev/null +++ b/extensions/copilot/src/extension/typescriptContext/vscode-node/ts6/regionContextProvider.ts @@ -0,0 +1,59 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ +import * as vscode from 'vscode'; + +import type { IRegionContextProviderService, Region, LineRange } from '../../../../platform/languageContextProvider/common/regionContextProvider'; +import * as protocol from '../../common/serverProtocol'; + +enum ExecutionTarget { + Semantic, + Syntax +} + +type ExecConfig = { + readonly executionTarget?: ExecutionTarget; +}; + +type RegionContextRequestArgs = Omit & { + file: vscode.Uri; + line: number; + offset: number; +}; + +export class TS6RegionContextProvider implements Omit, vscode.Disposable { + private static readonly ExecConfig: ExecConfig = { executionTarget: ExecutionTarget.Semantic }; + + async getRegions(document: vscode.Uri, languageId: string, ranges: vscode.Range[], requested?: LineRange): Promise { + if (document.scheme !== 'file' || (languageId !== 'typescript' && languageId !== 'javascript')) { + return undefined; + } + if (ranges.length === 0) { + return undefined; + } + + const firstPosition = ranges[0].start; + const args: RegionContextRequestArgs = { + file: document, + line: firstPosition.line + 1, + offset: firstPosition.character + 1, + ranges: ranges.map(range => ({ + start: { line: range.start.line, character: range.start.character }, + end: { line: range.end.line, character: range.end.character } + })), + requested + }; + const response = await vscode.commands.executeCommand( + 'typescript.tsserverRequest', + '_.copilot.regionContext', + args, + TS6RegionContextProvider.ExecConfig + ); + return protocol.RegionContextResponse.isOk(response) && response.body.regions.length > 0 ? response.body.regions : undefined; + } + + dispose(): void { + // No resources to dispose for the TS6 implementation + } +} diff --git a/extensions/copilot/src/extension/typescriptContext/vscode-node/tsc6/tsContextService.ts b/extensions/copilot/src/extension/typescriptContext/vscode-node/ts6/tsContextService.ts similarity index 100% rename from extensions/copilot/src/extension/typescriptContext/vscode-node/tsc6/tsContextService.ts rename to extensions/copilot/src/extension/typescriptContext/vscode-node/ts6/tsContextService.ts diff --git a/extensions/copilot/src/extension/typescriptContext/vscode-node/ts7/regionContextProvider.ts b/extensions/copilot/src/extension/typescriptContext/vscode-node/ts7/regionContextProvider.ts new file mode 100644 index 00000000000000..8c46ca8572b94b --- /dev/null +++ b/extensions/copilot/src/extension/typescriptContext/vscode-node/ts7/regionContextProvider.ts @@ -0,0 +1,331 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ +import * as vscode from 'vscode'; + +import type { Snapshot } from '@typescript/native/unstable/async'; +import * as ts from '@typescript/native/unstable/ast'; + +import type { ILogService } from '../../../../platform/log/common/logService'; +import { type IRegionContextProviderService, type Region, type LineRange } from '../../../../platform/languageContextProvider/common/regionContextProvider'; +import { TypeScript7Api } from './ts7Api'; +import { DisposableStore } from '../../../../util/vs/base/common/lifecycle'; +import tss from './typescripts'; + +type StructuralEntity = { kind: string; name?: string; rangeNode: ts.Node | [ts.Node, ts.Node]; includeJsDoc?: boolean; continueWith?: ts.Node }; + +interface RegionContextApi { + clearSourceFileCache(): void; + updateSnapshot(): Promise; +} + +interface RegionContextApiProvider extends vscode.Disposable { + getApi(): Promise; +} + +export class TS7RegionContextProvider implements Omit, vscode.Disposable { + + private readonly disposables: DisposableStore; + private readonly nativeApi: RegionContextApiProvider; + + constructor(readonly logService: ILogService, nativeApi: RegionContextApiProvider = new TypeScript7Api(logService)) { + this.disposables = new DisposableStore(); + this.nativeApi = this.disposables.add(nativeApi); + } + + async getRegions(document: vscode.Uri, languageId: string, ranges: vscode.Range[], requested?: LineRange): Promise { + if (document.scheme !== 'file' || (languageId !== 'typescript' && languageId !== 'javascript')) { + return undefined; + } + if (ranges.length === 0) { + return undefined; + } + + const api = await this.nativeApi.getApi(); + if (api === undefined) { + return undefined; + } + api.clearSourceFileCache(); + const snapshot = await api.updateSnapshot(); + try { + + const project = await snapshot.getDefaultProjectForFile(document.fsPath); + if (project === undefined) { + return undefined; + } + const sourceFile = await project.program.getSourceFile(document.fsPath); + if (sourceFile === undefined) { + return undefined; + } + + if (ranges.length === 1) { + return this.findEnclosingScopes(sourceFile, ranges[0].start.line, ranges[0].start.character, requested); + } else { + const containersList: Region[][] = []; + for (const range of ranges) { + const containers = await this.findEnclosingScopes(sourceFile, range.start.line, range.start.character, requested); + if (containers !== undefined && containers.length > 0) { + containersList.push(containers.reverse()); + } + } + if (containersList.length === 0) { + return undefined; + } + + const longestContainers = containersList.reduce((longest, containers) => containers.length > longest.length ? containers : longest); + const commonContainers = longestContainers.slice(); + for (const containers of containersList) { + if (containers === longestContainers) { + continue; + } + let commonLength = 0; + while (commonLength < commonContainers.length && commonLength < containers.length) { + const commonContainer = commonContainers[commonLength]; + const container = containers[commonLength]; + if (commonContainer.kind !== container.kind + || commonContainer.name !== container.name + || commonContainer.range.start !== container.range.start + || commonContainer.range.end !== container.range.end) { + break; + } + commonLength++; + } + commonContainers.length = commonLength; + } + + const tailContainers = containersList.map(containers => containers[containers.length - 1]); + if (tailContainers.length > 0) { + const container: Region = { + kind: 'merged', + range: { + start: Math.min(...tailContainers.map(container => container.range.start)), + end: Math.max(...tailContainers.map(container => container.range.end)) + } + }; + const lastContainer = commonContainers[commonContainers.length - 1]; + if (lastContainer !== undefined && container.range.end - container.range.start < lastContainer.range.end - lastContainer.range.start) { + commonContainers.push(container); + } + } + + return commonContainers.reverse(); + } + } finally { + await snapshot.dispose(); + } + } + + private async findEnclosingScopes(sourceFile: ts.SourceFile, line: number, column: number, requested?: LineRange | undefined): Promise { + const position = sourceFile.getPositionOfLineAndCharacter(line, column); + const tokenInfo = tss.getRelevantTokens(sourceFile, position); + const node = tokenInfo.touching ?? tokenInfo.token; + if (node === undefined) { + return undefined; + } + + const result: Region[] = []; + for (let current: ts.Node | undefined = node; current; current = current.parent) { + if (ts.isSourceFile(current)) { + const endLine = sourceFile.getLineAndCharacterOfPosition(sourceFile.getEnd()).line; + result.push({ + kind: 'sourceFile', + name: this.getBaseFileName(sourceFile.fileName), + range: { start: 0, end: endLine } + }); + break; + } + + const structuralEntity = this.getStructuralEntity(sourceFile, current, requested); + if (structuralEntity !== undefined) { + const { kind, name, rangeNode, includeJsDoc, continueWith } = structuralEntity; + const rangeStartNode = Array.isArray(rangeNode) ? rangeNode[0] : rangeNode; + const rangeEndNode = Array.isArray(rangeNode) ? rangeNode[1] : rangeNode; + result.push({ + kind, + name, + range: { + start: sourceFile.getLineAndCharacterOfPosition(rangeStartNode.getStart(sourceFile, includeJsDoc)).line, + end: sourceFile.getLineAndCharacterOfPosition(rangeEndNode.getEnd()).line + } + }); + current = continueWith ?? current; + } + } + return result.length > 0 ? result : undefined; + } + + private getStructuralEntity(sourceFile: ts.SourceFile, node: ts.Node, requested?: LineRange | undefined): StructuralEntity | undefined { + const parent = node.parent; + let name: string | undefined; + switch (node.kind) { + case ts.SyntaxKind.JSDoc: { + const parentEntity = this.getStructuralEntity(sourceFile, parent, requested); + if (parentEntity !== undefined) { + parentEntity.includeJsDoc = true; + parentEntity.continueWith ??= parent; + } + return parentEntity; + } + case ts.SyntaxKind.ImportDeclaration: + name = (node as ts.ImportDeclaration).moduleSpecifier.getText(); + return { kind: 'import', name, rangeNode: node }; + case ts.SyntaxKind.ExportDeclaration: + name = (node as ts.ExportDeclaration).moduleSpecifier?.getText(); + return { kind: 'export', name, rangeNode: node }; + case ts.SyntaxKind.FunctionDeclaration: + name = (node as ts.FunctionDeclaration).name?.text; + if (name === undefined) { + if (ts.isPropertyAssignment(parent) && ts.isIdentifier(parent.name)) { + name = parent.name.text; + } else if (ts.isVariableDeclaration(parent) && ts.isIdentifier(parent.name)) { + name = parent.name.text; + } + } + return { kind: 'function', name, rangeNode: node }; + case ts.SyntaxKind.Constructor: + return { kind: 'constructor', name: 'constructor', rangeNode: node }; + case ts.SyntaxKind.MethodDeclaration: + name = (node as ts.MethodDeclaration).name.getText(); + return { kind: 'method', name, rangeNode: node }; + case ts.SyntaxKind.MethodSignature: + name = (node as ts.MethodSignatureDeclaration).name.getText(); + return { kind: 'method', name, rangeNode: node }; + case ts.SyntaxKind.ArrowFunction: + if (ts.isPropertyAssignment(parent) && ts.isIdentifier(parent.name)) { + name = parent.name.text; + return { kind: 'function', name, rangeNode: parent, continueWith: parent }; + } else if (ts.isVariableDeclaration(parent) && ts.isIdentifier(parent.name)) { + name = parent.name.text; + return { kind: 'arrow-function', name, rangeNode: parent, continueWith: parent }; + } else if (ts.isCallExpression(parent)) { + return { kind: 'arrow-function', rangeNode: parent, continueWith: parent }; + } + return { kind: 'arrow-function', rangeNode: node }; + case ts.SyntaxKind.PropertyDeclaration: + return this.handleProperty(sourceFile, node as ts.PropertyDeclaration, requested); + case ts.SyntaxKind.PropertyAssignment: + return this.handleProperty(sourceFile, node as ts.PropertyAssignment, requested); + case ts.SyntaxKind.PropertySignature: + return this.handleProperty(sourceFile, node as ts.PropertySignatureDeclaration, requested); + case ts.SyntaxKind.GetAccessor: + name = (node as ts.GetAccessorDeclaration).name.getText(); + return { kind: 'getter', name, rangeNode: node }; + case ts.SyntaxKind.SetAccessor: + name = (node as ts.SetAccessorDeclaration).name.getText(); + return { kind: 'setter', name, rangeNode: node }; + case ts.SyntaxKind.ClassDeclaration: + name = (node as ts.ClassDeclaration).name?.text; + return { kind: 'class', name, rangeNode: node }; + case ts.SyntaxKind.InterfaceDeclaration: + name = (node as ts.InterfaceDeclaration).name.text; + return { kind: 'interface', name, rangeNode: node }; + case ts.SyntaxKind.ModuleDeclaration: + name = (node as ts.ModuleDeclaration).name.text; + return { kind: 'module', name, rangeNode: node }; + case ts.SyntaxKind.TypeAliasDeclaration: + name = (node as ts.TypeAliasDeclaration).name.text; + return { kind: 'type-alias', name, rangeNode: node }; + default: + return undefined; + } + } + + private handleProperty(sourceFile: ts.SourceFile, node: ts.PropertyDeclaration | ts.PropertyAssignment | ts.PropertySignatureDeclaration, requested?: LineRange | undefined): StructuralEntity | undefined { + const name = node.name.getText(); + if (ts.isPropertyDeclaration(node) || ts.isPropertyAssignment(node)) { + const initializeKind = node.initializer?.kind; + if (initializeKind === ts.SyntaxKind.FunctionType || initializeKind === ts.SyntaxKind.FunctionDeclaration || initializeKind === ts.SyntaxKind.FunctionExpression || initializeKind === ts.SyntaxKind.ArrowFunction) { + return { kind: 'function', name, rangeNode: node }; + } + } + const parent = node.parent; + if (requested !== undefined) { + const info = this.getMemberInfo(parent); + if (info === undefined) { + return undefined; + } + const { items, kind, memberKind, name } = info; + const range = this.calculateRange(sourceFile, parent, node, items, requested); + if (range === undefined) { + return undefined; + } + if (Array.isArray(range)) { + const [startIndex, endIndex] = range; + return { + kind: memberKind, + name, + rangeNode: [items[startIndex], items[endIndex]], + continueWith: parent + }; + } else { + return { + kind, + name, + rangeNode: parent, + continueWith: parent + }; + } + } + return undefined; + } + + private getMemberInfo(parent: ts.Node): { items: ts.NodeArray; kind: string; memberKind: string; name?: string | undefined } | undefined { + if (ts.isClassDeclaration(parent)) { + return { items: parent.members, kind: 'class', memberKind: 'class-members', name: parent.name?.text }; + } else if (ts.isInterfaceDeclaration(parent)) { + return { items: parent.members, kind: 'interface', memberKind: 'interface-members', name: parent.name?.text }; + } else if (ts.isObjectLiteralExpression(parent)) { + return { items: parent.properties, kind: 'object-literal', memberKind: 'object-literal-members' }; + } else if (ts.isTypeLiteralNode(parent)) { + return { items: parent.members, kind: 'type-literal', memberKind: 'type-literal-members' }; + } + return undefined; + } + + private calculateRange(sourceFile: ts.SourceFile, parent: ts.Node, node: ts.Node, items: ts.NodeArray, requested: LineRange): [number, number] | ts.Node | undefined { + const startLine = sourceFile.getLineAndCharacterOfPosition(parent.getStart(sourceFile)).line; + const endLine = sourceFile.getLineAndCharacterOfPosition(parent.getEnd()).line; + if (requested.start <= startLine && requested.end >= endLine) { + return parent; + } + + const index = items.indexOf(node); + if (index === -1) { + return undefined; + } + + let startIndex = Math.max(0, index - 1); + while (index - startIndex < 3 && startIndex > 0) { + const member = items[startIndex - 1]; + if (!this.isInsideRequestedRange(sourceFile, member, requested)) { + break; + } + startIndex--; + } + + let endIndex = Math.min(items.length - 1, index + 1); + while (endIndex - index < 3 && endIndex < items.length - 1) { + const member = items[endIndex + 1]; + if (!this.isInsideRequestedRange(sourceFile, member, requested)) { + break; + } + endIndex++; + } + return [startIndex, endIndex]; + } + + private isInsideRequestedRange(sourceFile: ts.SourceFile, member: ts.Node, requested: LineRange): boolean { + const memberStartLine = sourceFile.getLineAndCharacterOfPosition(member.getStart(sourceFile)).line; + const memberEndLine = sourceFile.getLineAndCharacterOfPosition(member.getEnd()).line; + return requested.start <= memberStartLine && requested.end >= memberEndLine; + } + + private getBaseFileName(fileName: string): string { + return fileName.substring(Math.max(fileName.lastIndexOf('/'), fileName.lastIndexOf('\\')) + 1); + } + + dispose(): void { + this.disposables.dispose(); + } +} diff --git a/extensions/copilot/src/extension/typescriptContext/vscode-node/ts7/test/regionContext.spec.ts b/extensions/copilot/src/extension/typescriptContext/vscode-node/ts7/test/regionContext.spec.ts new file mode 100644 index 00000000000000..59bd35008caa0c --- /dev/null +++ b/extensions/copilot/src/extension/typescriptContext/vscode-node/ts7/test/regionContext.spec.ts @@ -0,0 +1,84 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import assert from 'node:assert'; +import path from 'node:path'; + +import { API } from '@typescript/native/unstable/async'; +import * as vscode from 'vscode'; +import { afterAll, beforeAll, suite, test } from 'vitest'; + +import type { LineRange, Region } from '../../../../../platform/languageContextProvider/common/regionContextProvider'; +import { TestLogService } from '../../../../../platform/testing/common/testLogService'; +import { TS7RegionContextProvider } from '../regionContextProvider'; + +const fixtures = path.join(__dirname, '../../../serverPlugin/fixtures/context'); +const projectDirectory = path.join(fixtures, 'p14'); +const configFile = path.join(projectDirectory, 'tsconfig.json'); +const fileName = path.join(projectDirectory, 'source/f1.ts'); + +suite('TypeScript 7 region context', () => { + let api: API; + + beforeAll(() => { + api = new API({ cwd: process.cwd() }); + }); + + afterAll(async () => { + await api.close(); + }); + + async function getRegions(ranges: vscode.Range[], requested?: LineRange): Promise { + const provider = new TS7RegionContextProvider(new TestLogService(), new TestTypeScript7Api(api, configFile)); + try { + return await provider.getRegions(vscode.Uri.file(fileName), 'typescript', ranges, requested); + } finally { + provider.dispose(); + } + } + + test('returns enclosing structural regions', async () => { + assert.deepStrictEqual(await getRegions([range(9, 2)]), [ + { kind: 'constructor', name: 'constructor', range: { start: 8, end: 10 } }, + { kind: 'class', name: 'Calculator', range: { start: 5, end: 23 } }, + { kind: 'sourceFile', name: 'f1.ts', range: { start: 0, end: 32 } }, + ] satisfies Region[]); + }); + + test('merges distinct innermost regions', async () => { + assert.deepStrictEqual(await getRegions([range(13, 2), range(18, 2)]), [ + { kind: 'merged', range: { start: 12, end: 22 } }, + { kind: 'class', name: 'Calculator', range: { start: 5, end: 23 } }, + { kind: 'sourceFile', name: 'f1.ts', range: { start: 0, end: 32 } }, + ] satisfies Region[]); + }); + + test('groups property signatures within the requested range', async () => { + assert.deepStrictEqual(await getRegions([range(1, 1), range(2, 1)], { start: 1, end: 2 }), [ + { kind: 'interface-members', name: 'Result', range: { start: 1, end: 2 } }, + { kind: 'sourceFile', name: 'f1.ts', range: { start: 0, end: 32 } }, + ] satisfies Region[]); + }); +}); + +class TestTypeScript7Api { + constructor( + private readonly api: API, + private readonly configFile: string, + ) { } + + async getApi() { + return { + clearSourceFileCache: () => this.api.clearSourceFileCache(), + updateSnapshot: () => this.api.updateSnapshot({ openProjects: [this.configFile] }), + }; + } + + dispose(): void { } +} + +function range(line: number, character: number = 0): vscode.Range { + return new vscode.Range(line, character, line, character); +} diff --git a/extensions/copilot/src/platform/configuration/common/configurationService.ts b/extensions/copilot/src/platform/configuration/common/configurationService.ts index b3c94e4f9dc8ea..9cd52e957eca50 100644 --- a/extensions/copilot/src/platform/configuration/common/configurationService.ts +++ b/extensions/copilot/src/platform/configuration/common/configurationService.ts @@ -817,7 +817,6 @@ export namespace ConfigKey { export const BackgroundTodoAgentEnabled = defineSetting('chat.agent.backgroundTodoAgent.enabled', ConfigType.ExperimentBased, false); export const InlineEditsTriggerOnEditorChangeAfterSeconds = defineAndMigrateExpSetting('chat.advanced.inlineEdits.triggerOnEditorChangeAfterSeconds', 'chat.inlineEdits.triggerOnEditorChangeAfterSeconds', 10); - export const InlineEditsNextCursorPredictionDisplayLine = defineAndMigrateExpSetting('chat.advanced.inlineEdits.nextCursorPrediction.displayLine', 'chat.inlineEdits.nextCursorPrediction.displayLine', true); export const InlineEditsNextCursorPredictionCurrentFileMaxTokens = defineAndMigrateExpSetting('chat.advanced.inlineEdits.nextCursorPrediction.currentFileMaxTokens', 'chat.inlineEdits.nextCursorPrediction.currentFileMaxTokens', 3000); export const InlineEditsRenameSymbolSuggestions = defineSetting('chat.inlineEdits.renameSymbolSuggestions', ConfigType.ExperimentBased, true); export const InlineEditsPreferredModel = defineSetting('nextEditSuggestions.preferredModel', ConfigType.ExperimentBased, 'none'); diff --git a/extensions/copilot/src/platform/languageContextProvider/common/regionContextProvider.ts b/extensions/copilot/src/platform/languageContextProvider/common/regionContextProvider.ts new file mode 100644 index 00000000000000..2c7454f6d9c995 --- /dev/null +++ b/extensions/copilot/src/platform/languageContextProvider/common/regionContextProvider.ts @@ -0,0 +1,39 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ +import type * as vscode from 'vscode'; + +import { createServiceIdentifier } from '../../../util/common/services'; + +export interface LineRange { + start: number; + end: number; +} + +export interface Region { + kind: string; + name?: string; + range: LineRange; +} + +export const IRegionContextProviderService = createServiceIdentifier('IRegionContextProviderService'); + +export interface IRegionContextProviderService extends vscode.Disposable { + readonly _serviceBrand: undefined; + + getRegions(document: vscode.Uri, languageId: string, ranges: vscode.Range[]): Promise; + getRegions(document: vscode.Uri, languageId: string, ranges: vscode.Range[], requested?: LineRange): Promise; +} + +export class NullRegionContextProviderService implements IRegionContextProviderService { + readonly _serviceBrand: undefined; + + async getRegions(): Promise { + return undefined; + } + + dispose(): void { + // No resources to dispose for the Null implementation + } +} diff --git a/extensions/copilot/src/platform/test/node/services.ts b/extensions/copilot/src/platform/test/node/services.ts index 178f0ae9ccc429..c20a90fbb05ef4 100644 --- a/extensions/copilot/src/platform/test/node/services.ts +++ b/extensions/copilot/src/platform/test/node/services.ts @@ -100,6 +100,8 @@ import { SnapshotSearchService, TestingTabsAndEditorsService } from './simulatio import { TestChatAgentService } from './testChatAgentService'; import { TestWorkbenchService } from './testWorkbenchService'; import { TestWorkspaceService } from './testWorkspaceService'; +import { IGrepResultService, NullGrepResultService } from '../../../extension/tools/node/grepResultService'; +import { IRegionContextProviderService, NullRegionContextProviderService } from '../../languageContextProvider/common/regionContextProvider'; /** * Collects descriptors for services to use in testing. @@ -267,6 +269,8 @@ export function createPlatformServices(disposables: Pick testingServiceCollection.define(IImageService, nullImageService); testingServiceCollection.define(ILanguageContextService, NullLanguageContextService); testingServiceCollection.define(ILanguageContextProviderService, new SyncDescriptor(NullLanguageContextProviderService)); + testingServiceCollection.define(IGrepResultService, new SyncDescriptor(NullGrepResultService)); + testingServiceCollection.define(IRegionContextProviderService, new SyncDescriptor(NullRegionContextProviderService)); testingServiceCollection.define(ILanguageDiagnosticsService, new SyncDescriptor(TestLanguageDiagnosticsService)); testingServiceCollection.define(IPromptPathRepresentationService, new SyncDescriptor(TestPromptPathRepresentationService)); testingServiceCollection.define(IRequestLogger, new SyncDescriptor(NullRequestLogger)); diff --git a/src/vs/platform/agentHost/AGENTS.md b/src/vs/platform/agentHost/AGENTS.md index ce5eadb7faf9cf..0fd4811a053b43 100644 --- a/src/vs/platform/agentHost/AGENTS.md +++ b/src/vs/platform/agentHost/AGENTS.md @@ -220,7 +220,7 @@ a chat URI. New provider code must consume the seams. `register` takes the resolved provenance and whether to check tombstones. Explicit `AgentService.createSession` calls skip the tombstone check and clear any tombstone for that session URI; restore and discovery calls atomically decline to register if the session is or concurrently becomes tombstoned. An explicit row is never rewritten by catalog discovery. A migration-time host-owned marker can correct a previously discovered row back to internal provenance. -Providers own discovery lifecycle and push unknown chats with provider-classified provenance through `onDidDiscoverChats`. Claude, Codex, and Copilot classify their unknown native chats as external, except that Copilot keeps an unknown *legacy extension-host* chat internal because it is adoptable in place rather than someone else's session. Agent Service preserves that classification when it additively registers the event payload. Every provider starts one memoized initial attempt when the first discovery-event listener is attached; that attempt retries internally, but once it settles it is not re-armed by SDK readiness, so the only later trigger is an explicit one (for Copilot, the migrate-legacy toggle). Ordinary list refreshes never enumerate provider catalogs. External discovery has no migration marker or Copilot migrate-legacy gate; only the adoptable legacy extension-host half of Copilot's payload is withheld while migrate-legacy is off. Discovery never prunes a registry row when a provider later omits it and filters subagents and marked internal chat backings. +Providers own discovery lifecycle and push unknown chats with provider-classified provenance through `onDidDiscoverChats`. Claude, Codex, and Copilot classify their unknown native chats as external, except that Copilot keeps an unknown *legacy extension-host* chat internal because it is adoptable in place rather than someone else's session. Agent Service preserves that classification when it additively registers the event payload. Agent Service always attaches the event listener and queues each provider's external-session discovery through `_runWhenStartupSettled`, so the request waits for both Agent Host startup and the first successful session listing. Providers registered after that barrier opens run their queued work immediately, and a later transition from `none` starts discovery directly. Adopt-in-place legacy migration remains an independent provider-initialization trigger immediately after the discovery listener is attached, and another catalog consumer may also trigger discovery after it enumerates the provider catalog. This keeps `showExternalSessions: none` from initiating native discovery while allowing independently triggered discovery to populate the hidden registry normally. Ordinary list refreshes never enumerate provider catalogs. External discovery has no migration marker or Copilot migrate-legacy gate; only the adoptable legacy extension-host half of Copilot's payload is withheld while migrate-legacy is off. Discovery never prunes a registry row when a provider later omits it and filters subagents and marked internal chat backings. Discovery is registry-first: Agent Service hands each provider an optional `setKnownSessionsFilter` seam that answers, for a whole candidate set in one registry query, which sessions the host already owns. A provider drops those candidates before any per-session database open, and Copilot additionally skips adoptable legacy classification work (project/Git resolution) while migrate-legacy is off, since those candidates would not be emitted. Agent Service in turn rejects an already-registered candidate before `_isChatBacking()` or any other per-session I/O; provenance of a registered row stays owned by the explicit create/restore paths. Tombstoned sessions are absent from the registry and therefore never reported as known, so an explicitly deleted session still reaches `register`, whose atomic tombstone check declines it. @@ -228,7 +228,7 @@ Claude and Codex each use one memoized initial path: resolve/download the SDK, e If a provider cannot enumerate yet, its initial discovery attempt emits nothing; once ready, it emits the resulting chats through `onDidDiscoverChats`. Registry provenance is projected into `IAgentSessionMetadata._meta` with `readSessionExternal` / `withSessionExternal`, and the normal AHP listSessions round trip carries it to the Sessions provider. There is no external-specific UI behavior. -`listSessions()` coalesces concurrent computations per external-sessions mode, so the burst of calls a multi-window restore produces shares one registry traversal instead of one per window. The shared entry records the registry epoch it started at and is invalidated by every registry mutation, so a caller arriving after a mutation starts a fresh pass rather than joining a possibly pre-mutation one; each caller receives its own array. +`listSessions()` coalesces concurrent computations per external-sessions mode, so the burst of calls a multi-window restore produces shares one registry traversal instead of one per window. Registry mutations advance an epoch without removing an active computation. A caller arriving after a mutation shares one trailing computation that starts after the active one settles, preventing expensive provider, database, and Git work from overlapping for the same mode. Further invalidations before that trailing computation starts are absorbed by it; invalidations during it can schedule at most one subsequent computation for later callers. Each caller receives its own array, and no caller recursively follows more than the computation it joined. Legacy registry migration uses the `listChatsToMigrate()` contract. An array is authoritative even when empty. `undefined` means the catalog is unavailable and must not advance migration markers; Agent Service retries an unavailable registration-time catalog once before listing, and persistent unavailability rejects aggregate `listSessions()` with a typed provider-catalog error so clients preserve their last successful snapshots. `AgentChatMigrationDeferred` means the catalog cannot be enumerated until an external readiness action, such as downloading an optional SDK: it does not advance the provider marker and does not block healthy providers' aggregate listing. A provider's later discovery signal force-retries its migration partition before additively registering the signal's unknown/external entries, and a subsequent list refresh can retry a still-deferred provider. `BaseAgentHostSessionsProvider` retries failures with exponential backoff; `AgentHostSessionListStore` leaves its cache invalid and retries on the next controller, lifecycle, or workspace refresh trigger. Replacement retry ownership is compare-and-swap single-flight: overlapping list computations that observed the same failed attempt await the first caller's installed retry rather than queueing another provider enumeration. Successful providers retain their completed migration state when a sibling provider is unavailable. diff --git a/src/vs/platform/agentHost/browser/agentHostProtocolClient.ts b/src/vs/platform/agentHost/browser/agentHostProtocolClient.ts index 0d1c45a11c7ef2..398c837bc265f3 100644 --- a/src/vs/platform/agentHost/browser/agentHostProtocolClient.ts +++ b/src/vs/platform/agentHost/browser/agentHostProtocolClient.ts @@ -97,6 +97,15 @@ function transportLostError(address: string): ProtocolError { return new ProtocolError(AHP_CLIENT_CONNECTION_CLOSED, `Transport lost (reconnecting): ${address}`); } +/** + * Whether an error means the transport went away rather than the request + * being rejected on its merits. Such a failure is transient and must stay + * recoverable, so it is never reclassified as a terminal condition. + */ +function isConnectionClosedError(error: unknown): boolean { + return error instanceof ProtocolError && error.code === AHP_CLIENT_CONNECTION_CLOSED; +} + interface IRemoteAgentHostExtensionNotificationMap { 'setClientManagedSettingsPermissions': { params: { permissions: IAgentHostManagedSettingsPermissions } }; } @@ -172,6 +181,16 @@ export interface IAgentHostProtocolClientOptions { readonly clientInfo?: Implementation; /** How a dropped transport is restored. Defaults to {@link DEFAULT_RECONNECT_POLICY}. */ readonly reconnectPolicy?: IRemoteAgentHostReconnectPolicy; + /** Resolves authentication to restore immediately after every fresh initialize. */ + readonly resolveInitialAuthentication?: () => Promise; +} + +/** An initial authentication resolver failed after a successful initialize. */ +export class InitialAuthenticationError extends Error { + constructor(error: unknown) { + super(`Initial authentication failed: ${error instanceof Error ? error.message : String(error)}`); + this.name = 'InitialAuthenticationError'; + } } /** @@ -285,6 +304,7 @@ export class AgentHostProtocolClient extends Disposable implements IAgentConnect private readonly _loadEstimator: ILoadEstimator; private readonly _clientInfo: Implementation | undefined; private readonly _reconnectPolicy: IRemoteAgentHostReconnectPolicy; + private readonly _resolveInitialAuthentication: (() => Promise) | undefined; /** * URIs we have already granted implicit read access for on this connection. @@ -339,6 +359,7 @@ export class AgentHostProtocolClient extends Disposable implements IAgentConnect this._loadEstimator = options?.loadEstimator ?? LoadEstimator.getInstance(); this._clientInfo = options?.clientInfo; this._reconnectPolicy = options?.reconnectPolicy ?? DEFAULT_RECONNECT_POLICY; + this._resolveInitialAuthentication = options?.resolveInitialAuthentication; if (typeof transportOrFactory === 'function') { this._transportFactory = transportOrFactory; @@ -481,6 +502,12 @@ export class AgentHostProtocolClient extends Disposable implements IAgentConnect initialSubscriptions: [ROOT_STATE_URI], }, { bypassInitializeQueue: true }); this._applyInitializeResult(result); + if (this._resolveInitialAuthentication || this._authentication.size > 0) { + await this._restoreAuthenticationAfterFreshInitialize(AgentHostClientState.Connecting); + if (this._state.kind !== AgentHostClientState.Connecting) { + throw transportLostError(this._address); + } + } // Hydrate root state from the initial snapshot for (const snapshot of result.snapshots ?? []) { @@ -501,7 +528,7 @@ export class AgentHostProtocolClient extends Disposable implements IAgentConnect const protocolError = error instanceof ProtocolError ? error : new ProtocolError(AHP_CLIENT_CONNECTION_CLOSED, error instanceof Error ? error.message : String(error)); - if (protocolError.code === AhpErrorCodes.UnsupportedProtocolVersion) { + if (protocolError.code === AhpErrorCodes.UnsupportedProtocolVersion || error instanceof InitialAuthenticationError) { this._cancelLivenessTimers(); if (this._state.kind === AgentHostClientState.Connecting) { this._state.outbox.length = 0; @@ -699,7 +726,10 @@ export class AgentHostProtocolClient extends Disposable implements IAgentConnect this._applyReconnectResult(result, freshInitialize); this._updateManagedSettingsPermissions(true); if (freshInitialize && result.type === ReconnectResultType.Snapshot) { - await this._restoreAuthenticationAfterFreshInitialize(); + await this._restoreAuthenticationAfterFreshInitialize(AgentHostClientState.Reconnecting); + if (this._state.kind !== AgentHostClientState.Reconnecting) { + return; + } await this._restoreSubscriptionsAfterFreshInitialize(result.snapshots); } if (this._state.kind !== AgentHostClientState.Reconnecting) { @@ -735,6 +765,14 @@ export class AgentHostProtocolClient extends Disposable implements IAgentConnect this._handleFatalClose(protocolError); return; } + if (err instanceof InitialAuthenticationError) { + const protocolError = new ProtocolError(AHP_CLIENT_CONNECTION_CLOSED, err.message); + this._cancelLivenessTimers(); + this._rejectPendingRequests(protocolError); + reconnect.gate.error(err); + this._transitionTo({ kind: AgentHostClientState.Incompatible, error: protocolError }); + return; + } // Replace the gate so awaiting callers see the failure but new // callers gate on the next attempt instead of slipping through onto // the dead transport. Outbox carries forward to the next attempt. @@ -812,12 +850,40 @@ export class AgentHostProtocolClient extends Disposable implements IAgentConnect ]); } - private async _restoreAuthenticationAfterFreshInitialize(): Promise { - await Promise.all([...this._authentication.values()].map(params => this._dispatchRequest('authenticate', { - channel: ROOT_STATE_URI, - ...params, - scopes: params.scopes ? [...params.scopes] : undefined, - }, { bypassReconnectGate: true }))); + private async _restoreAuthenticationAfterFreshInitialize(expectedState: AgentHostClientState.Connecting | AgentHostClientState.Reconnecting): Promise { + let resolvedInitialAuthentication = false; + if (this._resolveInitialAuthentication) { + try { + const initialAuthentication = await this._resolveInitialAuthentication(); + if (initialAuthentication) { + const normalizedParams = this._normalizeAuthenticationParams(initialAuthentication); + this._authentication.set(this._authenticationKey(normalizedParams), normalizedParams); + resolvedInitialAuthentication = true; + } + } catch (error) { + throw new InitialAuthenticationError(error); + } + if (this._state.kind !== expectedState) { + return; + } + } + try { + await Promise.all([...this._authentication.values()].map(params => this._dispatchRequest('authenticate', { + channel: ROOT_STATE_URI, + ...params, + scopes: params.scopes ? [...params.scopes] : undefined, + }, this._state.kind === AgentHostClientState.Connecting + ? { bypassInitializeQueue: true, bypassReconnectGate: true } + : { bypassReconnectGate: true }))); + } catch (error) { + // A dropped transport is not an authentication failure. Wrapping it + // would classify a momentary blip as terminally incompatible and + // permanently stop recovery, so let it stay a reconnectable error. + if (resolvedInitialAuthentication && !isConnectionClosedError(error)) { + throw new InitialAuthenticationError(error); + } + throw error; + } } private _clientMeta(): Record { @@ -1143,16 +1209,13 @@ export class AgentHostProtocolClient extends Disposable implements IAgentConnect * Authenticate with the remote agent host using a specific scheme. */ async authenticate(params: AuthenticateParams): Promise { - const normalizedParams: AuthenticateParams = { - ...params, - scopes: params.scopes ? [...new Set(params.scopes)].sort() : undefined, - }; + const normalizedParams = this._normalizeAuthenticationParams(params); await this._sendRequest('authenticate', { channel: ROOT_STATE_URI, ...normalizedParams, scopes: normalizedParams.scopes ? [...normalizedParams.scopes] : undefined, }); - const key = `${normalizedParams.resource}\0${JSON.stringify(normalizedParams.scopes ?? [])}`; + const key = this._authenticationKey(normalizedParams); if (params.token) { this._authentication.set(key, normalizedParams); } else { @@ -1161,6 +1224,17 @@ export class AgentHostProtocolClient extends Disposable implements IAgentConnect return { authenticated: true }; } + private _normalizeAuthenticationParams(params: AuthenticateParams): AuthenticateParams { + return { + ...params, + scopes: params.scopes ? [...new Set(params.scopes)].sort() : undefined, + }; + } + + private _authenticationKey(params: AuthenticateParams): string { + return `${params.resource}\0${JSON.stringify(params.scopes ?? [])}`; + } + /** * Gracefully shut down all sessions on the remote host. */ diff --git a/src/vs/platform/agentHost/browser/remoteAgentHostServiceImpl.ts b/src/vs/platform/agentHost/browser/remoteAgentHostServiceImpl.ts index a2f73cb658f4f2..f2a76d4dd681b5 100644 --- a/src/vs/platform/agentHost/browser/remoteAgentHostServiceImpl.ts +++ b/src/vs/platform/agentHost/browser/remoteAgentHostServiceImpl.ts @@ -6,35 +6,42 @@ // Service implementation that manages remote agent host connections from // entries supplied by registered connection factories. -import { Emitter, Event } from '../../../base/common/event.js'; +import { Emitter } from '../../../base/common/event.js'; +import { isCancellationError } from '../../../base/common/errors.js'; import { Disposable, DisposableStore, IDisposable, toDisposable } from '../../../base/common/lifecycle.js'; import { DeferredPromise, raceTimeout } from '../../../base/common/async.js'; -import { autorun, derived, IObservable, observableFromEvent, observableValue } from '../../../base/common/observable.js'; +import { autorun, derived, IObservable, observableValue } from '../../../base/common/observable.js'; import { IConfigurationService } from '../../configuration/common/configuration.js'; import { IEnvironmentService } from '../../environment/common/environment.js'; import { IInstantiationService } from '../../instantiation/common/instantiation.js'; import { ILabelService } from '../../label/common/label.js'; import { ILogService } from '../../log/common/log.js'; +import { observableConfigValue } from '../../observable/common/platformObservableUtils.js'; import { hasKey } from '../../../base/common/types.js'; import { AgentHostAhpJsonlLoggingSettingId, type IAgentConnection } from '../common/agentService.js'; import { IRemoteAgentHostService, RemoteAgentHostConnectionStatus, + RemoteAgentHostAutoConnectSettingId, RemoteAgentHostsEnabledSettingId, RemoteAgentHostsSettingId, + WEBSOCKET_ENTRY_TYPE_CONFIG, getEntryTypeConfig, - readWebSocketRemoteAgentHostEntries, + isLegacySshRawEntry, + isRawRemoteAgentHostEntry, type IRemoteAgentHostConnectionFactory, type IRemoteAgentHostConnectionInfo, type IRemoteAgentHostConnectOptions, type IRemoteAgentHostCreatedConnection, type IRemoteAgentHostEntry, + type IRawRemoteAgentHostEntry, type IRemoteAgentHostProtocolClient, RemoteAgentHostEntryType, } from '../common/remoteAgentHostService.js'; import { computeReconnectDelay, hasExhaustedReconnectAttempts } from '../common/reconnectPolicy.js'; -import { AgentHostProtocolClient, AgentHostClientState } from './agentHostProtocolClient.js'; +import { NonReconnectableTransportError } from '../common/state/sessionTransport.js'; +import { AgentHostProtocolClient, InitialAuthenticationError } from './agentHostProtocolClient.js'; import { WebSocketClientTransport } from './webSocketClientTransport.js'; import { AGENT_HOST_LABEL_FORMATTER, AGENT_HOST_SCHEME, agentHostAuthority, normalizeRemoteAgentHostAddress } from '../common/agentHostUri.js'; import { PROTOCOL_VERSION } from '../common/state/protocol/version/registry.js'; @@ -53,6 +60,8 @@ interface IConnectionEntry { * disconnect the freshly-established tunnel as a side effect. */ readonly transportDisposable?: IDisposable; + /** Whether a replacement connection assumes transport teardown ownership. */ + readonly reconnectTransfersTransportOwnership: boolean; connected: boolean; /** Current connection status for UI display. */ status: RemoteAgentHostConnectionStatus; @@ -63,10 +72,23 @@ function disposeEntry(entry: IConnectionEntry): void { entry.transportDisposable?.dispose(); } +/** + * Whether a failed connection attempt must not be retried automatically. + * + * A transport that declared itself non-reconnectable, and anything the user + * cancelled or refused, are decisions rather than transient faults. Retrying + * them re-prompts the person who just declined — the SSH host-key flow surfaces + * both, as a cancellation and as a denial converted by its factory. + */ +function isTerminalConnectError(err: unknown): boolean { + return err instanceof NonReconnectableTransportError || isCancellationError(err); +} + /** Builds WebSocket-backed protocol clients without performing their handshake. */ class WebSocketConnectionFactory extends Disposable implements IRemoteAgentHostConnectionFactory { readonly kind = RemoteAgentHostEntryType.WebSocket; readonly entries: IObservable; + private readonly _rawEntries: IObservable; constructor( private readonly _instantiationService: IInstantiationService, @@ -75,13 +97,15 @@ class WebSocketConnectionFactory extends Disposable implements IRemoteAgentHostC private readonly _clientInfo: () => typeof editorWindowAgentHostClientInfo, ) { super(); - this.entries = observableFromEvent( - this, - Event.filter( - this._configurationService.onDidChangeConfiguration, - event => event.affectsConfiguration(RemoteAgentHostsSettingId), - ), - () => this._getEntries(), + this._rawEntries = observableConfigValue( + RemoteAgentHostsSettingId, + [], + this._configurationService, + ); + this.entries = derived(this, reader => this._rawEntries.read(reader) + .filter(isRawRemoteAgentHostEntry) + .filter(entry => !isLegacySshRawEntry(entry)) + .map(entry => WEBSOCKET_ENTRY_TYPE_CONFIG.fromRaw(entry)) ); } @@ -105,9 +129,6 @@ class WebSocketConnectionFactory extends Disposable implements IRemoteAgentHostC return Promise.resolve({ connection }); } - private _getEntries(): IRemoteAgentHostEntry[] { - return readWebSocketRemoteAgentHostEntries(this._configurationService); - } } export class RemoteAgentHostService extends Disposable implements IRemoteAgentHostService { @@ -127,6 +148,8 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo private readonly _entries = new Map(); private readonly _connectionFactories = new Map(); private readonly _connectionFactoriesObservable = observableValue(this, [] as readonly IRemoteAgentHostConnectionFactory[]); + private readonly _remoteAgentHostsEnabled: IObservable; + private readonly _remoteAgentHostsAutoConnect: IObservable; private readonly _configuredEntries = derived(this, reader => { let entries: IRemoteAgentHostEntry[] = []; for (const factory of this._connectionFactoriesObservable.read(reader)) { @@ -140,14 +163,9 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo private readonly _pendingConnects = new Map>(); private readonly _names = new Map(); private readonly _tokens = new Map(); - /** - * Stores the original {@link IRemoteAgentHostEntry} for connections - * registered via {@link addManagedConnection}. This is needed because - * tunnel entries are not persisted to settings and therefore don't - * appear in {@link configuredEntries}. - */ - private readonly _registeredEntries = new Map(); private readonly _pendingConnectionWaits = new Map>(); + /** Errors from reconnects that could not start a dial. */ + private readonly _failedReconnects = new Map(); /** Pending reconnect timeouts, keyed by normalized address. */ private readonly _reconnectTimeouts = new Map>(); /** Current reconnect attempt count per address for exponential backoff. */ @@ -173,6 +191,9 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo ) { super(); + this._remoteAgentHostsEnabled = observableConfigValue(RemoteAgentHostsEnabledSettingId, true, this._configurationService); + this._remoteAgentHostsAutoConnect = observableConfigValue(RemoteAgentHostAutoConnectSettingId, true, this._configurationService); + // The service creates these built-in factories, so it owns their // lifetime too; `registerConnectionFactory` only manages registry // membership so externally-supplied factories stay owned by their producer. @@ -184,13 +205,10 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo )))); this._register(autorun(reader => { this._configuredEntries.read(reader); + this._remoteAgentHostsEnabled.read(reader); + this._remoteAgentHostsAutoConnect.read(reader); this._reconcileConnections(); })); - this._register(this._configurationService.onDidChangeConfiguration(e => { - if (e.affectsConfiguration(RemoteAgentHostsEnabledSettingId)) { - this._reconcileConnections(); - } - })); } @@ -260,13 +278,6 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo getEntryByAddress(address: string): IRemoteAgentHostEntry | undefined { const normalized = normalizeRemoteAgentHostAddress(address); - // Check dynamically registered entries first (e.g. tunnel connections - // that are not persisted to settings). - const registered = this._registeredEntries.get(normalized); - if (registered) { - return registered; - } - // Fall back to configured entries from settings. return this.configuredEntries.find( entry => this._entryAddress(entry) === normalized ); @@ -295,12 +306,35 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo } reconnect(address: string, userInitiated = true): void { + if (this._store.isDisposed) { + return; + } const normalized = normalizeRemoteAgentHostAddress(address); + // A dial already in flight is itself a fresh attempt, so neither a + // retry nor a user request gains anything by tearing it down and + // starting a second one — that is what produced concurrent remote + // bootstraps. Join it instead. A user-initiated request still restores + // the retry budget, so pressing reconnect while a slow bootstrap runs + // is not silently useless if that bootstrap ultimately fails. + if (this._pendingConnects.has(normalized)) { + if (userInitiated) { + this._failedReconnects.delete(normalized); + this._cancelReconnect(normalized); + this._reconnectAttempts.delete(normalized); + } + return; + } + this._failedReconnects.delete(normalized); const configuredEntry = this._configuredEntries.get().find( entry => this._entryAddress(entry) === normalized ); - if (!configuredEntry || !getEntryTypeConfig(configuredEntry.connection.type).dialableByService) { + if (!configuredEntry) { + this._failedReconnects.set(normalized, new Error(`No remote agent host entry is staged for ${normalized}.`)); + return; + } + if (!this._connectionFactories.has(configuredEntry.connection.type)) { + this._failedReconnects.set(normalized, new Error(`No connection factory is registered for ${configuredEntry.connection.type}.`)); return; } @@ -311,20 +345,20 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo // Cancel any pending reconnect this._cancelReconnect(normalized); - this._reconnectAttempts.delete(normalized); + if (userInitiated) { + // An automatic retry must not resurrect its own exhausted attempt budget. + this._reconnectAttempts.delete(normalized); + } // Tear down existing connection if present const entry = this._entries.get(normalized); if (entry) { this._entries.delete(normalized); - // SSH reconnects replace the relay in the shared process using the - // same connection id. Disposing its previous transport here would - // race that replacement and disconnect the fresh relay. The SSH - // factory transfers teardown ownership to the new entry. entry.store.dispose(); - if (configuredEntry.connection.type !== RemoteAgentHostEntryType.SSH) { + if (!entry.reconnectTransfersTransportOwnership) { entry.transportDisposable?.dispose(); } + this._onDidChangeConnections.fire(); } // Start fresh connection attempt @@ -332,7 +366,10 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo } async waitForConnection(address: string): Promise { - if (!this._configurationService.getValue(RemoteAgentHostsEnabledSettingId)) { + if (this._store.isDisposed) { + throw new Error('Remote agent host service is disposed.'); + } + if (!this._remoteAgentHostsEnabled.get()) { throw new Error('Remote agent host connections are not enabled.'); } @@ -341,6 +378,10 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo if (existingConnection) { return existingConnection; } + const reconnectFailure = this._failedReconnects.get(normalizedAddress); + if (reconnectFailure) { + throw reconnectFailure; + } const wait = this._getOrCreateConnectionWait(normalizedAddress); @@ -371,90 +412,13 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo return connection; } - async addManagedConnection(entry: IRemoteAgentHostEntry, connection: IAgentConnection, transportDisposable?: IDisposable, status = RemoteAgentHostConnectionStatus.connected): Promise { - if (!this._configurationService.getValue(RemoteAgentHostsEnabledSettingId)) { - throw new Error('Remote agent host connections are not enabled.'); - } - - const address = this._entryAddress(entry); - - // Dispose any existing entry for this address to avoid leaking - // old protocol clients and relay transports on reconnect. - // - // CRITICAL: we deliberately do NOT run the existing entry's - // transportDisposable. On a reconnect to the same address, the - // shared-process tunnel keyed by connectionId is already owned by - // the new connection we just established. Running the old teardown - // would call _mainService.disconnect(connectionId) and immediately - // kill the brand-new tunnel. - const existingEntry = this._entries.get(address); - if (existingEntry) { - this._entries.delete(address); - existingEntry.store.dispose(); - } - - const store = new DisposableStore(); - - // Create a connection entry wrapping the pre-connected client - const protocolClient = connection as AgentHostProtocolClient; - store.add(protocolClient); - const connEntry: IConnectionEntry = { store, client: protocolClient, transportDisposable, connected: RemoteAgentHostConnectionStatus.isConnected(status), status }; - this._entries.set(address, connEntry); - this._names.set(address, entry.name); - this._registeredEntries.set(address, entry); - this._updateHostLabelFormatter(address, entry.name); - if (entry.connectionToken) { - this._tokens.set(address, entry.connectionToken); - } - - store.add(protocolClient.onDidClose(() => { - if (this._entries.get(address) === connEntry) { - connEntry.connected = false; - connEntry.status = RemoteAgentHostConnectionStatus.disconnected; - this._onDidChangeConnections.fire(); - } - })); - - store.add(protocolClient.onDidChangeConnectionState(state => { - if (this._entries.get(address) !== connEntry) { - return; - } - switch (state) { - case AgentHostClientState.Reconnecting: - connEntry.connected = false; - connEntry.status = RemoteAgentHostConnectionStatus.reconnecting; - this._onDidChangeConnections.fire(); - break; - case AgentHostClientState.Connected: - connEntry.connected = true; - connEntry.status = RemoteAgentHostConnectionStatus.connected; - this._onDidChangeConnections.fire(); - break; - case AgentHostClientState.Connecting: - case AgentHostClientState.Incompatible: - case AgentHostClientState.Closed: - break; - } - })); - - this._onDidChangeConnections.fire(); - - return { - address, - name: entry.name, - clientId: protocolClient.clientId, - defaultDirectory: protocolClient.defaultDirectory, - status, - }; - } - async removeRemoteAgentHost(address: string): Promise { const normalized = normalizeRemoteAgentHostAddress(address); // Eagerly clear in-memory state so the UI updates immediately // (the config change listener will reconcile, but this is instant). this._names.delete(normalized); this._tokens.delete(normalized); - this._registeredEntries.delete(normalized); + this._failedReconnects.delete(normalized); this._clearHostLabelFormatter(normalized); this._cancelReconnect(normalized); this._reconnectAttempts.delete(normalized); @@ -465,7 +429,6 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo const entry = this._entries.get(address); if (entry) { this._entries.delete(address); - this._registeredEntries.delete(address); disposeEntry(entry); this._rejectPendingConnectionWait(address, new Error(`Connection closed: ${address}`)); this._onDidChangeConnections.fire(); @@ -491,7 +454,7 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo return; } - if (!this._configurationService.getValue(RemoteAgentHostsEnabledSettingId)) { + if (!this._remoteAgentHostsEnabled.get()) { // Disconnect all when disabled for (const address of [...this._entries.keys()]) { this._cancelReconnect(address); @@ -500,14 +463,9 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo this._names.clear(); this._tokens.clear(); this._reconnectAttempts.clear(); - // Drop label formatters for entries no longer represented by an - // active connection or a dynamically registered entry. Connections - // added via {@link addManagedConnection} (e.g. tunnels) live outside - // the configured-entries set and must keep their formatter. + // Drop label formatters for entries no longer represented by an active connection. for (const address of [...this._labelFormatters.keys()]) { - if (!this._registeredEntries.has(address)) { - this._clearHostLabelFormatter(address); - } + this._clearHostLabelFormatter(address); } return; } @@ -523,14 +481,6 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo const oldNames = new Map(this._names); this._names.clear(); this._tokens.clear(); - // Runtime-registered connections are not part of the persisted set, so - // seed their metadata first; without this a live tunnel/WSL/cloud - // connection survives reconcile but reports its address as its name, - // which downstream provider reconciliation treats as a rename. - for (const [address, entry] of this._registeredEntries) { - this._names.set(address, entry.name); - this._tokens.set(address, entry.connectionToken); - } for (const { entry, address } of entriesWithAddress) { this._names.set(address, entry.name); this._tokens.set(address, entry.connectionToken); @@ -540,17 +490,16 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo } } - // Drop formatters for addresses that are no longer configured and - // not dynamically registered. + // Drop formatters for addresses that are no longer configured. for (const address of [...this._labelFormatters.keys()]) { - if (!desired.has(address) && !this._registeredEntries.has(address)) { + if (!desired.has(address)) { this._clearHostLabelFormatter(address); } } - // Remove connections no longer in the setting + // Remove connections no longer exposed by a factory. for (const address of [...this._entries.keys()]) { - if (!desired.has(address) && !this._registeredEntries.has(address)) { + if (!desired.has(address)) { this._logService.info(`[RemoteAgentHost] Disconnecting from ${address}`); this._cancelReconnect(address); this._reconnectAttempts.delete(address); @@ -558,10 +507,10 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo } } - // Add entries that this service owns. + // Add entry-driven connection kinds. for (const { entry, address } of entriesWithAddress) { // This gate becomes redundant once every entry type has a registered factory. - if (!this._entries.has(address) && !this._pendingConnects.has(address) && getEntryTypeConfig(entry.connection.type).dialableByService) { + if (!this._entries.has(address) && !this._pendingConnects.has(address) && this._shouldAutoConnect(entry)) { void this._connectTo(entry, { userInitiated: false }); } } @@ -573,6 +522,9 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo } private _connectTo(entryToConnect: IRemoteAgentHostEntry, options: IRemoteAgentHostConnectOptions): Promise { + if (this._store.isDisposed) { + return Promise.resolve(); + } const entryToCreate = this._normalizeEntry(entryToConnect); const address = this._entryAddress(entryToCreate); const existingPendingConnect = this._pendingConnects.get(address); @@ -598,13 +550,16 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo } private async _createAndConnect(entryToCreate: IRemoteAgentHostEntry, address: string, options: IRemoteAgentHostConnectOptions): Promise { - if (!this._configurationService.getValue(RemoteAgentHostsEnabledSettingId)) { + if (this._store.isDisposed || !this._remoteAgentHostsEnabled.get()) { return; } const factory = this._connectionFactories.get(entryToCreate.connection.type); if (!factory) { - this._logService.error(`[RemoteAgentHost] No connection factory registered for ${entryToCreate.connection.type} at ${address}`); + const error = new Error(`No connection factory is registered for ${entryToCreate.connection.type}.`); + this._logService.error(`[RemoteAgentHost] ${error.message} at ${address}`); + this._failedReconnects.set(address, error); + this._rejectPendingConnectionWait(address, error); return; } @@ -622,7 +577,7 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo } catch (err) { this._logService.error(`[RemoteAgentHost] Failed to create a connection to ${address}. Verify address and connectionToken`, err); this._rejectPendingConnectionWait(address, err); - if (!this._store.isDisposed && this._configurationService.getValue(RemoteAgentHostsEnabledSettingId)) { + if (!isTerminalConnectError(err) && !this._store.isDisposed && this._remoteAgentHostsEnabled.get()) { this._scheduleReconnect(address, entryToCreate.connectionToken); } return; @@ -630,12 +585,13 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo if ( this._store.isDisposed - || !this._configurationService.getValue(RemoteAgentHostsEnabledSettingId) + || !this._remoteAgentHostsEnabled.get() || !this._configuredEntries.get().some(entry => this._entryAddress(entry) === address) || this._entries.has(address) ) { createdConnection.connection.dispose(); createdConnection.transportDisposable?.dispose(); + this._rejectPendingConnectionWait(address, new Error(`Connection attempt for ${address} was discarded because it is no longer active.`)); return; } @@ -645,6 +601,7 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo store, client, transportDisposable: createdConnection.transportDisposable, + reconnectTransfersTransportOwnership: createdConnection.reconnectTransfersTransportOwnership ?? false, connected: false, status: RemoteAgentHostConnectionStatus.connecting, }; @@ -692,9 +649,14 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo this._onDidChangeConnections.fire(); break; case 'connecting': - case 'incompatible': case 'closed': break; + case 'incompatible': + entry.connected = false; + entry.status = RemoteAgentHostConnectionStatus.incompatible('Authentication failed during connection initialization.', [PROTOCOL_VERSION]); + this._reconnectAttempts.delete(address); + this._onDidChangeConnections.fire(); + break; } })); @@ -729,7 +691,9 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo // reconnect attempts would just spin until the user upgrades // either side, so leave recovery to the manual `Reconnect` // action in the picker. - const incompatible = RemoteAgentHostConnectionStatus.fromConnectError(err, [PROTOCOL_VERSION]); + const incompatible = err instanceof InitialAuthenticationError + ? RemoteAgentHostConnectionStatus.incompatible(err.message, [PROTOCOL_VERSION]) + : RemoteAgentHostConnectionStatus.fromConnectError(err, [PROTOCOL_VERSION]); if (incompatible) { this._logService.warn(`[RemoteAgentHost] Incompatible with ${address}: ${incompatible.kind === 'incompatible' ? incompatible.message : ''}`); entry.status = incompatible; @@ -764,10 +728,10 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo /** * Schedule a reconnect attempt with exponential backoff. - * Only reconnects if the address is still in the configured entries. + * Only reconnects if the address remains exposed by a configuration or factory. */ private _scheduleReconnect(address: string, connectionToken?: string): void { - // Don't reconnect if the address was removed from settings. + // Don't reconnect if the address is no longer exposed. const configuredEntry = this._configuredEntries.get().find(entry => this._entryAddress(entry) === address); if (!configuredEntry) { this._logService.info(`[RemoteAgentHost] Not reconnecting to ${address}: no longer configured`); @@ -779,6 +743,10 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo this._logService.info(`[RemoteAgentHost] Not reconnecting to ${address}: automatic restore is disabled`); return; } + if (getEntryTypeConfig(configuredEntry.connection.type).dialedFromEntries && !this._shouldAutoConnect(configuredEntry)) { + this._logService.info(`[RemoteAgentHost] Not reconnecting to ${address}: automatic connection is disabled`); + return; + } // Check the recorded count before adding this attempt, so a policy of // `maxAttempts: n` actually performs n attempts rather than n - 1. @@ -799,7 +767,7 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo const timeout = setTimeout(() => { this._reconnectTimeouts.delete(address); const currentEntry = this._configuredEntries.get().find(entry => this._entryAddress(entry) === address); - if (currentEntry) { + if (currentEntry && (!getEntryTypeConfig(currentEntry.connection.type).dialedFromEntries || this._shouldAutoConnect(currentEntry))) { void this._connectTo({ ...currentEntry, connectionToken: connectionToken ?? this._tokens.get(address) ?? currentEntry.connectionToken, @@ -809,6 +777,12 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo this._reconnectTimeouts.set(address, timeout); } + private _shouldAutoConnect(entry: IRemoteAgentHostEntry): boolean { + const config = getEntryTypeConfig(entry.connection.type); + return config.dialedFromEntries + && (!config.autoConnectGated || this._remoteAgentHostsAutoConnect.get()); + } + /** Cancel a pending reconnect timeout for the given address. */ private _cancelReconnect(address: string): void { const timeout = this._reconnectTimeouts.get(address); diff --git a/src/vs/platform/agentHost/common/agent.ts b/src/vs/platform/agentHost/common/agent.ts index f1d9ee618eed1b..7a800b1387c707 100644 --- a/src/vs/platform/agentHost/common/agent.ts +++ b/src/vs/platform/agentHost/common/agent.ts @@ -1206,6 +1206,9 @@ export interface IAgent { /** Provides chats that are ready to be registered as Agent Host sessions. */ readonly onDidDiscoverChats: Event; + /** Starts the provider's memoized native chat discovery pass. */ + startChatDiscovery?(): Promise; + /** Lets discovery drop registered candidates before per-session I/O. */ setKnownSessionsFilter?(filter: IAgentKnownSessionsFilter): void; diff --git a/src/vs/platform/agentHost/common/agentHostSchema.ts b/src/vs/platform/agentHost/common/agentHostSchema.ts index 8cc3f0de13c064..bd4e2dd1b72a3b 100644 --- a/src/vs/platform/agentHost/common/agentHostSchema.ts +++ b/src/vs/platform/agentHost/common/agentHostSchema.ts @@ -824,7 +824,7 @@ export const platformRootSchema = createSchema({ enum: [ChatExternalSessionsMode.None, ChatExternalSessionsMode.Recent, ChatExternalSessionsMode.Last24Hours, ChatExternalSessionsMode.Last7Days, ChatExternalSessionsMode.Last30Days], enumDescriptions: [ localize('agentHost.config.showExternalSessions.none', "Do not show external sessions."), - localize('agentHost.config.showExternalSessions.recent', "Show up to the 2 most recent external sessions updated in the last 7 days. Once at least 2 local sessions exist, external sessions older than the second-newest local session are hidden."), + localize('agentHost.config.showExternalSessions.recent', "Show up to the 2 most recent external sessions updated in the last 7 days. At startup, external sessions older than the second-most-recently updated local session are hidden."), localize('agentHost.config.showExternalSessions.last24Hours', "Show external sessions updated in the last 24 hours."), localize('agentHost.config.showExternalSessions.last7Days', "Show external sessions updated in the last 7 days."), localize('agentHost.config.showExternalSessions.last30Days', "Show external sessions updated in the last 30 days."), diff --git a/src/vs/platform/agentHost/common/remoteAgentHostService.ts b/src/vs/platform/agentHost/common/remoteAgentHostService.ts index 595b29772b628d..5a03828ea833b6 100644 --- a/src/vs/platform/agentHost/common/remoteAgentHostService.ts +++ b/src/vs/platform/agentHost/common/remoteAgentHostService.ts @@ -114,10 +114,7 @@ export const RemoteAgentHostsSettingId = 'chat.remoteAgentHosts'; /** Configuration key to enable remote agent host connections. */ export const RemoteAgentHostsEnabledSettingId = 'chat.remoteAgentHostsEnabled'; -/** - * Configuration key that controls whether online dev tunnels and - * WSL remote agent hosts are auto-connected at startup. - */ +/** Configuration key that controls whether online dev tunnels, configured SSH remote agent hosts, and WSL remote agent hosts are auto-connected at startup. */ export const RemoteAgentHostAutoConnectSettingId = 'chat.remoteAgentHostsAutoConnect'; export const enum RemoteAgentHostEntryType { @@ -213,8 +210,7 @@ export interface IRemoteAgentHostCloudSandboxConnection { /** * A runtime-only connection to an agent host running inside a Dev Container. - * The owning Dev Container integration establishes the transport and registers - * the connected client through {@link IRemoteAgentHostService.addManagedConnection}. + * The Dev Container integration stages its transport for its connection factory. */ export interface IRemoteAgentHostDevContainerConnection { readonly type: RemoteAgentHostEntryType.DevContainer; @@ -272,6 +268,11 @@ export interface IRemoteAgentHostCreatedConnection { * (e.g. a shared-process relay channel). Disposed with the connection entry. */ readonly transportDisposable?: IDisposable; + /** + * Whether a redial transfers transport teardown ownership to the new connection. + * Defaults to `false`. + */ + readonly reconnectTransfersTransportOwnership?: boolean; } /** Builds agent host connections of one {@link RemoteAgentHostEntryType}. */ @@ -347,10 +348,15 @@ export type RemoteAgentHostEntryStore = 'settings' | 'storage' | 'runtime'; interface IRemoteAgentHostEntryTypeConfigBase { readonly type: TConnection['type']; /** - * Whether RemoteAgentHostService can dial this entry from its address alone. - * When `false`, an owning transport service registers the connection. + * Whether this entry-driven kind is dialed during reconciliation from the factory's entries. + * On-demand kinds set this to `false`, but an explicit {@link IRemoteAgentHostService.reconnect} still dials their staged entries. */ - readonly dialableByService: boolean; + readonly dialedFromEntries: boolean; + /** + * Whether background dialing is controlled by {@link RemoteAgentHostAutoConnectSettingId}. + * Defaults to `false`. + */ + readonly autoConnectGated?: boolean; /** Whether the address is subject to `normalizeRemoteAgentHostAddress`. */ readonly normalizedAddress: boolean; /** Policy for restoring a dropped transport. */ @@ -384,7 +390,8 @@ export type IRemoteAgentHostEntryTypeConfig = { type: RemoteAgentHostEntryType.WebSocket, store: 'settings', - dialableByService: true, + dialedFromEntries: true, + autoConnectGated: false, normalizedAddress: true, reconnect: DEFAULT_RECONNECT_POLICY, address: connection => connection.address, @@ -397,7 +404,8 @@ export const WEBSOCKET_ENTRY_TYPE_CONFIG: IPersistedEntryTypeConfig = { type: RemoteAgentHostEntryType.SSH, store: 'storage', - dialableByService: true, + dialedFromEntries: true, + autoConnectGated: true, normalizedAddress: true, reconnect: DEFAULT_RECONNECT_POLICY, address: connection => connection.address, @@ -572,11 +580,19 @@ export function removeSSHRemoteAgentHostEntry(storageService: IStorageService, a } function runtimeEntryTypeConfig(type: TConnection['type'], normalizedAddress: boolean, address: (connection: TConnection) => string, reconnect: IRemoteAgentHostReconnectPolicy = DEFAULT_RECONNECT_POLICY): IRuntimeEntryTypeConfig { - return { type, store: 'runtime', dialableByService: false, normalizedAddress, reconnect, address }; + return { type, store: 'runtime', dialedFromEntries: false, normalizedAddress, reconnect, address }; } -const WSL_ENTRY_TYPE_CONFIG = runtimeEntryTypeConfig(RemoteAgentHostEntryType.WSL, true, connection => connection.address); -const TUNNEL_ENTRY_TYPE_CONFIG = runtimeEntryTypeConfig(RemoteAgentHostEntryType.Tunnel, false, connection => `${TUNNEL_ADDRESS_PREFIX}${connection.tunnelId}`); +const WSL_ENTRY_TYPE_CONFIG: IRemoteAgentHostEntryTypeConfig = { + ...runtimeEntryTypeConfig(RemoteAgentHostEntryType.WSL, true, connection => connection.address), + dialedFromEntries: true, + autoConnectGated: true, +}; +const TUNNEL_ENTRY_TYPE_CONFIG: IRemoteAgentHostEntryTypeConfig = { + ...runtimeEntryTypeConfig(RemoteAgentHostEntryType.Tunnel, false, connection => `${TUNNEL_ADDRESS_PREFIX}${connection.tunnelId}`), + dialedFromEntries: true, + autoConnectGated: true, +}; const CLOUD_SANDBOX_ENTRY_TYPE_CONFIG = runtimeEntryTypeConfig(RemoteAgentHostEntryType.CloudSandbox, true, connection => connection.address); // Relay failures are cheap, but a cold container can make `devcontainer up` rebuild Docker for minutes; retry slower and favor explicit recovery. const DEV_CONTAINER_RECONNECT_POLICY: IRemoteAgentHostReconnectPolicy = { @@ -634,10 +650,8 @@ export type RemoteAgentHostInputParseResult = export const IRemoteAgentHostService = createDecorator('remoteAgentHostService'); /** - * Manages connections to one or more remote agent host processes. Each - * connection is identified by its address string and - * exposed as an {@link IAgentConnection}, the same interface used for - * the local agent host. + * Owns factory-built remote agent host connections, including handshake, status, + * retry, and disposal. Each connection is identified by address and exposed as an {@link IAgentConnection}. */ export interface IRemoteAgentHostService { readonly _serviceBrand: undefined; @@ -648,7 +662,7 @@ export interface IRemoteAgentHostService { /** Currently connected remote addresses with metadata. */ readonly connections: readonly IRemoteAgentHostConnectionInfo[]; - /** All configured remote agent host entries, regardless of connection status. */ + /** All remote agent host entries exposed by registered factories, regardless of connection status. */ readonly configuredEntries: readonly IRemoteAgentHostEntry[]; /** Registers a factory for one connection kind. Throws if that kind already has one. */ @@ -687,28 +701,6 @@ export interface IRemoteAgentHostService { */ reconnect(address: string, userInitiated?: boolean): void; - /** - * Register a pre-connected agent connection. - * Used by transport services that do not yet provide a connection factory - * to inject relay-backed connections. - * - * The optional `transportDisposable` represents the underlying transport - * (e.g. an SSH tunnel relay or tunnel-relay session) and is owned by this - * service for the lifetime of the entry. It will be disposed when: - * - the entry is removed via {@link removeRemoteAgentHost} - * - the entry is reconciled away (config-driven removal) - * - this service itself is disposed - * Callers should put any teardown that needs to happen on entry removal - * (e.g. closing the shared-process tunnel, dropping renderer-side handles) - * into this disposable, so a single removal path tears down the whole stack. - * - * `status` defaults to `connected`. Pass `incompatible` when the managed - * transport is alive but the protocol handshake rejected the client version; - * this keeps recovery actions (such as server upgrade) addressable without - * exposing the connection as ready for session traffic. - */ - addManagedConnection(entry: IRemoteAgentHostEntry, connection: IAgentConnection, transportDisposable?: IDisposable, status?: RemoteAgentHostConnectionStatus): Promise; - /** * Force the protocol client at `address` (if any) to treat its * transport as closed. Used by services that learn about a @@ -726,8 +718,7 @@ export interface IRemoteAgentHostService { /** * Look up the {@link IRemoteAgentHostEntry} for a given address. - * Checks both configured entries from settings and dynamically - * registered entries (e.g. tunnel connections). + * Entries are supplied by registered connection factories. */ getEntryByAddress(address: string): IRemoteAgentHostEntry | undefined; @@ -773,9 +764,6 @@ export class NullRemoteAgentHostService implements IRemoteAgentHostService { async removeRemoteAgentHost(_address: string): Promise { } reconnect(_address: string, _userInitiated?: boolean): void { } notifyConnectionClosed(_address: string): void { } - async addManagedConnection(): Promise { - throw new Error('Remote agent host connections are not supported in this environment.'); - } getEntryByAddress(): IRemoteAgentHostEntry | undefined { return undefined; } async triggerServerUpgrade(): Promise { throw new Error('Remote agent host connections are not supported in this environment.'); diff --git a/src/vs/platform/agentHost/common/state/sessionState.ts b/src/vs/platform/agentHost/common/state/sessionState.ts index 31dcd0d3863773..2740e9ce216742 100644 --- a/src/vs/platform/agentHost/common/state/sessionState.ts +++ b/src/vs/platform/agentHost/common/state/sessionState.ts @@ -2047,6 +2047,33 @@ export function withSessionEhcliAdopted(meta: SessionSummaryMeta | undefined, ad return Object.keys(next).length > 0 ? next : undefined; } +/** + * Session-DB key recording the id of the final turn that existed when a legacy + * Copilot CLI session was adopted. It marks the boundary between the migrated + * (checkpoint-less) history and any turns added after adoption, so a consumer + * that substitutes the session-wide changeset for a migrated turn's absent + * per-turn changeset (see the chat editor fallback) can target exactly that + * turn and never a post-adoption one. + */ +export const AH_META_EHCLI_LAST_TURN_DB_KEY = 'agentHost.ehcliLastMigratedTurn'; + +/** `_meta` key mirroring {@link AH_META_EHCLI_LAST_TURN_DB_KEY} on a summary. */ +export const SESSION_META_EHCLI_LAST_TURN_KEY = 'ehcliLastMigratedTurn'; + +/** The id of the last turn migrated when the legacy Copilot CLI session was adopted, if recorded. */ +export function readSessionEhcliLastMigratedTurn(meta: SessionSummaryMeta | undefined): string | undefined { + const value = meta?.[SESSION_META_EHCLI_LAST_TURN_KEY]; + return typeof value === 'string' && value ? value : undefined; +} + +/** Returns a copy of `meta` with the last-migrated-turn marker set, or unchanged when `turnId` is empty. */ +export function withSessionEhcliLastMigratedTurn(meta: SessionSummaryMeta | undefined, turnId: string | undefined): SessionSummaryMeta | undefined { + if (!turnId) { + return meta; + } + return { ...meta, [SESSION_META_EHCLI_LAST_TURN_KEY]: turnId }; +} + /** * Whether a session should be matched against a workspace folder by its project * (repository) root in addition to its working directories. True only for diff --git a/src/vs/platform/agentHost/common/tunnelAgentHost.ts b/src/vs/platform/agentHost/common/tunnelAgentHost.ts index 5baccfbf750963..a17b0f2ed53d52 100644 --- a/src/vs/platform/agentHost/common/tunnelAgentHost.ts +++ b/src/vs/platform/agentHost/common/tunnelAgentHost.ts @@ -81,6 +81,8 @@ export interface ICachedTunnel { readonly tunnelId: string; readonly clusterId: string; readonly name: string; + /** Protocol version at cache time. Optional because entries from older builds do not contain it. */ + readonly protocolVersion?: number; readonly authProvider?: 'github' | 'microsoft'; } @@ -506,13 +508,22 @@ export interface ITunnelAgentHostService { /** Remove a tunnel from the cache. */ removeCachedTunnel(tunnelId: string): void; - /** Whether startup/background auto-connect should skip this tunnel because the user disconnected it. */ + /** Whether the user dismissed this tunnel from the remote-host picker. */ + isTunnelDismissed(tunnelId: string): boolean; + + /** Persist that the user dismissed this tunnel from the remote-host picker. */ + dismissTunnel(tunnelId: string): void; + + /** Clear a previous picker-dismissal after the user explicitly reconnects this tunnel. */ + clearTunnelDismissal(tunnelId: string): void; + + /** Whether startup/background auto-connect should skip this tunnel, because this machine hosts it. */ isAutoConnectSuppressed(tunnelId: string): boolean; - /** Remember that the user explicitly disconnected this tunnel, so startup/background auto-connect skips it. */ + /** Remember that startup/background auto-connect must skip this tunnel, because this machine hosts it. */ suppressAutoConnect(tunnelId: string): void; - /** Clear a previous user-disconnect marker after the user explicitly reconnects this tunnel. */ + /** Clear a previous auto-connect suppression once this machine no longer hosts the tunnel. */ clearAutoConnectSuppression(tunnelId: string): void; /** diff --git a/src/vs/platform/agentHost/common/tunnelGatewaySelection.ts b/src/vs/platform/agentHost/common/tunnelGatewaySelection.ts index ae52add775513d..6970ab57212ce3 100644 --- a/src/vs/platform/agentHost/common/tunnelGatewaySelection.ts +++ b/src/vs/platform/agentHost/common/tunnelGatewaySelection.ts @@ -8,7 +8,6 @@ import { type IDialogService } from '../../dialogs/common/dialogs.js'; import { type IProductService } from '../../product/common/productService.js'; import { type IRemoteAgentHostLocationPreferenceService } from './remoteAgentHostLocationPreference.js'; import { promptRemoteAgentHostLocationPreference } from './remoteAgentHostLocationPreferenceDialog.js'; -import { type IRemoteAgentHostService } from './remoteAgentHostService.js'; import { type ITunnelGatewayEndpoint, type ITunnelGatewayInventory, type ITunnelGatewaySelection, type TunnelGatewayServerType } from './tunnelAgentHost.js'; /** Endpoints of `type`, sorted deterministically by `instanceId`. */ @@ -134,9 +133,7 @@ export async function resolveGatewaySelection( } /** - * Decide whether a tunnel-failover notification should be shown after a - * connection attempt's {@link IRemoteAgentHostService.addManagedConnection} - * has already succeeded. Fires in two cases, both of which mean the editor + * Decide whether a tunnel-failover notification should be shown after a successful factory-built connection. Fires in two cases, both of which mean the editor * process that used to host the connection is gone and a dedicated agent * host silently took its place: * @@ -173,9 +170,7 @@ export function shouldNotifyTunnelFailover( * Retains the last successfully registered endpoint's server type per * stable tunnel address (`tunnel:`) so a later automatic * reconnect for the same tunnel can detect a silent editor → standalone - * failover via {@link shouldNotifyTunnelFailover}. Entries are only ever - * written after a successful {@link IRemoteAgentHostService.addManagedConnection} - * registration and are deliberately never cleared on relay closure, so the + * failover via {@link shouldNotifyTunnelFailover}. Server types are recorded only after a successful factory-built connection and are deliberately never cleared on relay closure, so the * comparison survives disconnect/reconnect cycles for the tunnel's * lifetime. Exported (and kept free of any IPC/protocol dependencies) so * the retention + decision behavior can be unit tested in isolation. diff --git a/src/vs/platform/agentHost/common/wslRemoteAgentHost.ts b/src/vs/platform/agentHost/common/wslRemoteAgentHost.ts index cb7e47eaf312da..27a057bf483f57 100644 --- a/src/vs/platform/agentHost/common/wslRemoteAgentHost.ts +++ b/src/vs/platform/agentHost/common/wslRemoteAgentHost.ts @@ -39,6 +39,8 @@ export interface IWSLAgentHostConfig { readonly name: string; /** Dev override: custom command to start the remote agent host. See SSH equivalent. */ readonly remoteAgentHostCommand?: string; + /** Whether an explicit user action initiated the connection. */ + readonly userInitiated?: boolean; } export interface IWSLConnectProgress { @@ -64,11 +66,10 @@ export interface IWSLAgentHostConnection extends IDisposable { /** * A WSL distro the user has connected to during this or a previous window. - * Persisted by {@link IWSLRemoteAgentHostService} so the startup - * auto-reconnect loop knows which running distros to re-attach to. This is - * the WSL analogue of the tunnel service's cached-tunnels list — WSL - * connections are managed in-memory and are never written to the remote - * agent hosts setting. + * Persisted by {@link IWSLRemoteAgentHostService} so its connection factory + * can supply startup entries. This is the WSL analogue of the tunnel + * service's cached-tunnels list — WSL connections are managed in-memory and + * are never written to the remote agent hosts setting. */ export interface IWSLCachedDistro { readonly distro: string; @@ -96,12 +97,12 @@ export interface IWSLRemoteAgentHostService { listRunningDistros(): Promise; connect(config: IWSLAgentHostConfig): Promise; disconnect(distro: string): Promise; - /** Used by the contribution's auto-reconnect loop on startup. */ - reconnect(distro: string, name: string): Promise; + /** Reconnect a cached distro, optionally as an automatic recovery attempt. */ + reconnect(distro: string, name: string, userInitiated?: boolean): Promise; /** * Distros the user has connected to, persisted across windows. Drives the - * startup auto-reconnect loop. WSL connections themselves live in-memory, - * mirroring how tunnels are handled. + * remote agent host service's startup auto-connect. WSL connections + * themselves live in-memory, mirroring how tunnels are handled. */ getCachedDistros(): readonly IWSLCachedDistro[]; } @@ -110,8 +111,8 @@ export const IWSLRemoteAgentHostMainService = createDecorator; connect(config: IWSLAgentHostConfig): Promise; disconnect(distro: string): Promise; - reconnect(distro: string, name: string, remoteAgentHostCommand?: string): Promise; + reconnect(distro: string, name: string, remoteAgentHostCommand?: string, userInitiated?: boolean): Promise; } diff --git a/src/vs/platform/agentHost/electron-browser/sshRemoteAgentHostServiceImpl.ts b/src/vs/platform/agentHost/electron-browser/sshRemoteAgentHostServiceImpl.ts index ea6b97acd96533..d54829f1076774 100644 --- a/src/vs/platform/agentHost/electron-browser/sshRemoteAgentHostServiceImpl.ts +++ b/src/vs/platform/agentHost/electron-browser/sshRemoteAgentHostServiceImpl.ts @@ -38,6 +38,7 @@ import { SSH_REMOTE_AGENT_HOST_CHANNEL, computeSSHConnectionKey, isSSHHostKeyDeniedError, + SSH_HOST_KEY_DENIED_ERROR_NAME, SSHAuthMethod, type ISSHAgentHostConfig, type ISSHAgentHostConnection, @@ -193,25 +194,40 @@ class SSHConnectionFactory extends Disposable implements IRemoteAgentHostConnect const stagedConfig = this._stagedConfigurations.get(entry.connection.address); this._stagedConfigurations.delete(entry.connection.address); - const result = stagedConfig - ? await this._mainService.connect(this._augmentConfig({ ...stagedConfig, userInitiated: stagedConfig.userInitiated ?? options.userInitiated })) - : entry.connection.sshConfigHost - ? await this._mainService.reconnect( - entry.connection.sshConfigHost, - entry.name, - this._getRemoteAgentHostCommand(), - this._isSSHAgentForwardingEnabled(), - options.userInitiated, - this._locationPreferenceService.getPreference(computeSSHConnectionKey({ sshConfigHost: entry.connection.sshConfigHost })), - ) - : await this._mainService.connect(this._augmentConfig({ - host: entry.connection.hostName, - port: entry.connection.port, - username: entry.connection.user ?? entry.connection.hostName, - authMethod: SSHAuthMethod.Agent, - name: entry.name, - userInitiated: options.userInitiated, - })); + let result; + try { + result = stagedConfig + ? await this._mainService.connect(this._augmentConfig({ ...stagedConfig, userInitiated: stagedConfig.userInitiated ?? options.userInitiated })) + : entry.connection.sshConfigHost + ? await this._mainService.reconnect( + entry.connection.sshConfigHost, + entry.name, + this._getRemoteAgentHostCommand(), + this._isSSHAgentForwardingEnabled(), + options.userInitiated, + this._locationPreferenceService.getPreference(computeSSHConnectionKey({ sshConfigHost: entry.connection.sshConfigHost })), + ) + : await this._mainService.connect(this._augmentConfig({ + host: entry.connection.hostName, + port: entry.connection.port, + username: entry.connection.user ?? entry.connection.hostName, + authMethod: SSHAuthMethod.Agent, + name: entry.name, + userInitiated: options.userInitiated, + })); + } catch (error) { + // A refused host key is the user's decision, not a transient fault. + // Report it in the shared vocabulary for "do not retry" while keeping + // the host-key-denial name, which `isSSHHostKeyDeniedError` matches + // across IPC — telemetry and the contribution's pause policy both + // depend on that identity surviving. + if (isSSHHostKeyDeniedError(error)) { + const denied = new NonReconnectableTransportError(error instanceof Error ? error.message : String(error)); + denied.name = SSH_HOST_KEY_DENIED_ERROR_NAME; + throw denied; + } + throw error; + } this._logService.trace(`[SSHRemoteAgentHost] SSH tunnel established, connectionId=${result.connectionId}`); const existing = this._connections.get(result.connectionId); @@ -238,6 +254,7 @@ class SSHConnectionFactory extends Disposable implements IRemoteAgentHostConnect return { connection: this._createRelayClient(result), transportDisposable: this._createTransportDisposable(result.connectionId, existing, this._observeSuccessfulConnection(result, options.userInitiated)), + reconnectTransfersTransportOwnership: true, }; } this._logService.info(`[SSHRemoteAgentHost] Replacing stale connection handle for ${result.address}`); @@ -267,6 +284,7 @@ class SSHConnectionFactory extends Disposable implements IRemoteAgentHostConnect return { connection: this._createRelayClient(result), transportDisposable: this._createTransportDisposable(result.connectionId, handle, endpointSelectionObserver), + reconnectTransfersTransportOwnership: true, }; } catch (err) { this._logService.error('[SSHRemoteAgentHost] Connection setup failed', err); diff --git a/src/vs/platform/agentHost/electron-browser/wslRemoteAgentHostServiceImpl.ts b/src/vs/platform/agentHost/electron-browser/wslRemoteAgentHostServiceImpl.ts index b22bd7c540b081..45a48caf665f31 100644 --- a/src/vs/platform/agentHost/electron-browser/wslRemoteAgentHostServiceImpl.ts +++ b/src/vs/platform/agentHost/electron-browser/wslRemoteAgentHostServiceImpl.ts @@ -6,13 +6,14 @@ import { Emitter, Event } from '../../../base/common/event.js'; import { localize } from '../../../nls.js'; import { Disposable, IDisposable, toDisposable } from '../../../base/common/lifecycle.js'; +import { IObservable, observableFromEvent } from '../../../base/common/observable.js'; import { ILogService } from '../../log/common/log.js'; import { IConfigurationService } from '../../configuration/common/configuration.js'; import { IEnvironmentService } from '../../environment/common/environment.js'; import { ISharedProcessService } from '../../ipc/electron-browser/services.js'; import { IStorageService, StorageScope, StorageTarget } from '../../storage/common/storage.js'; import { ProxyChannel } from '../../../base/parts/ipc/common/ipc.js'; -import { IRemoteAgentHostService, RemoteAgentHostEntryType, RemoteAgentHostsEnabledSettingId, type IRemoteAgentHostEntry } from '../common/remoteAgentHostService.js'; +import { IRemoteAgentHostService, RemoteAgentHostConnectionStatus, RemoteAgentHostEntryType, RemoteAgentHostsEnabledSettingId, getEntryAddress, type IRemoteAgentHostConnectOptions, type IRemoteAgentHostConnectionFactory, type IRemoteAgentHostCreatedConnection, type IRemoteAgentHostEntry } from '../common/remoteAgentHostService.js'; import { createDecorator, IInstantiationService } from '../../instantiation/common/instantiation.js'; import { AhpJsonlLogger } from '../common/ahpJsonlLogger.js'; import { AgentHostAhpJsonlLoggingSettingId } from '../common/agentService.js'; @@ -31,6 +32,7 @@ import { type IWSLConnectResult, type IWSLDistro, type IWSLRemoteAgentHostMainService, + WSL_ADDRESS_PREFIX, } from '../common/wslRemoteAgentHost.js'; export const IWSLRelayClientFactory = createDecorator('wslRelayClientFactory'); @@ -66,7 +68,11 @@ export class WSLRelayClientFactory implements IWSLRelayClientFactory { } try { - const result = await mainService.reconnect(config.distro, config.name, config.remoteAgentHostCommand); + const runningDistros = await mainService.listRunningDistros().catch((): string[] => []); + if (!runningDistros.includes(config.distro)) { + throw new NonReconnectableTransportError(`WSL distro '${config.distro}' is not running.`); + } + const result = await mainService.reconnect(config.distro, config.name, config.remoteAgentHostCommand, false); return { connectionId: result.connectionId, }; @@ -108,15 +114,244 @@ export class WSLRelayClientFactory implements IWSLRelayClientFactory { */ const CACHED_WSL_DISTROS_KEY = 'agentHost.wsl.cachedDistros'; +function readCachedWSLDistros(storageService: IStorageService): readonly IWSLCachedDistro[] { + const raw = storageService.get(CACHED_WSL_DISTROS_KEY, StorageScope.APPLICATION); + if (!raw) { + return []; + } + try { + const parsed: unknown = JSON.parse(raw); + if (!Array.isArray(parsed)) { + return []; + } + return parsed.filter((item): item is IWSLCachedDistro => + !!item && typeof item.distro === 'string' && typeof item.name === 'string'); + } catch { + return []; + } +} + +function storeCachedWSLDistros(storageService: IStorageService, distros: readonly IWSLCachedDistro[]): void { + if (distros.length === 0) { + storageService.remove(CACHED_WSL_DISTROS_KEY, StorageScope.APPLICATION); + } else { + storageService.store(CACHED_WSL_DISTROS_KEY, JSON.stringify(distros), StorageScope.APPLICATION, StorageTarget.USER); + } +} + +/** Creates WSL relay clients for {@link WSLRemoteAgentHostService}. */ +class WSLConnectionFactory extends Disposable implements IRemoteAgentHostConnectionFactory { + readonly kind = RemoteAgentHostEntryType.WSL; + readonly entries: IObservable; + + private readonly _stagedConfigurations = new Map(); + + constructor( + private readonly _storageService: IStorageService, + private readonly _mainService: IWSLRemoteAgentHostMainService, + private readonly _remoteAgentHostService: IRemoteAgentHostService, + private readonly _relayClientFactory: IWSLRelayClientFactory, + private readonly _connections: Map, + private readonly _onDidChangeConnections: () => void, + private readonly _onDidReportConnectProgress: (progress: IWSLConnectProgress) => void, + private readonly _getRemoteAgentHostCommand: () => string | undefined, + private readonly _createTransportDisposable: (connectionId: string, distro: string, handle: WSLAgentHostConnectionHandle) => IDisposable, + private readonly _logService: ILogService, + ) { + super(); + this.entries = observableFromEvent( + this, + this._storageService.onDidChangeValue(StorageScope.APPLICATION, CACHED_WSL_DISTROS_KEY, this._store), + () => this._getEntries(), + ); + } + + stageConfiguration(config: IWSLAgentHostConfig): IRemoteAgentHostEntry { + const entry = this._createEntry(config.distro, config.name); + this._stagedConfigurations.set(getEntryAddress(entry), { config, isInitialConnection: true }); + this._storeEntry(entry); + return entry; + } + + stageEntry(distro: string, name: string, userInitiated = true): IRemoteAgentHostEntry { + const entry = this._createEntry(distro, name); + this._stagedConfigurations.set(getEntryAddress(entry), { + config: { distro, name, remoteAgentHostCommand: this._getRemoteAgentHostCommand(), userInitiated }, + isInitialConnection: false, + }); + this._storeEntry(entry); + return entry; + } + + async createConnection(entry: IRemoteAgentHostEntry, options: IRemoteAgentHostConnectOptions): Promise { + if (entry.connection.type !== RemoteAgentHostEntryType.WSL) { + throw new Error(`WSL factory cannot create a ${entry.connection.type} connection.`); + } + + const address = getEntryAddress(entry); + let stagedConnection = this._stagedConfigurations.get(address); + this._stagedConfigurations.delete(address); + let config = stagedConnection?.config ?? { + distro: entry.connection.distro, + name: entry.name, + remoteAgentHostCommand: this._getRemoteAgentHostCommand(), + }; + let userInitiated = config.userInitiated ?? options.userInitiated; + if (!userInitiated) { + try { + await this._ensureDistroIsRunning(config.distro); + } catch (err) { + const userStagedConnection = this._stagedConfigurations.get(address); + if (!userStagedConnection) { + throw err; + } + this._stagedConfigurations.delete(address); + stagedConnection = userStagedConnection; + config = stagedConnection.config; + userInitiated = config.userInitiated ?? options.userInitiated; + } + // A user action may have arrived while the background precondition ran. + const userStagedConnection = this._stagedConfigurations.get(address); + if (userStagedConnection) { + this._stagedConfigurations.delete(address); + stagedConnection = userStagedConnection; + config = stagedConnection.config; + userInitiated = config.userInitiated ?? options.userInitiated; + } + } + + const result = stagedConnection?.isInitialConnection + ? await this._mainService.connect({ ...config, userInitiated }) + : await this._mainService.reconnect(config.distro, config.name, config.remoteAgentHostCommand, userInitiated); + this._logService.trace(`[WSLRemoteAgentHost] WSL relay established, connectionId=${result.connectionId}`); + return this._setupConnection(result, config.remoteAgentHostCommand); + } + + private _createEntry(distro: string, name: string): IRemoteAgentHostEntry { + return { + name, + connection: { + type: RemoteAgentHostEntryType.WSL, + address: `${WSL_ADDRESS_PREFIX}${distro}`, + distro, + }, + }; + } + + private _storeEntry(entry: IRemoteAgentHostEntry): void { + if (entry.connection.type !== RemoteAgentHostEntryType.WSL) { + return; + } + // Bind the narrowed connection before the closure: TypeScript does not + // carry the discriminant narrowing into the filter callback below. + const connection = entry.connection; + const cached = readCachedWSLDistros(this._storageService).filter(distro => distro.distro !== connection.distro); + storeCachedWSLDistros(this._storageService, [{ distro: connection.distro, name: entry.name }, ...cached]); + } + + private _getEntries(): readonly IRemoteAgentHostEntry[] { + return readCachedWSLDistros(this._storageService).map(({ distro, name }) => ({ + name, + connection: { + type: RemoteAgentHostEntryType.WSL, + address: `${WSL_ADDRESS_PREFIX}${distro}`, + distro, + }, + })); + } + + private async _ensureDistroIsRunning(distro: string): Promise { + const runningDistros = await this._mainService.listRunningDistros(); + if (!runningDistros.includes(distro)) { + throw new NonReconnectableTransportError(`WSL distro '${distro}' is not running.`); + } + } + + private _setupConnection(result: IWSLConnectResult, remoteAgentHostCommand: string | undefined): IRemoteAgentHostCreatedConnection { + const existing = this._connections.get(result.connectionId); + if (existing) { + if (this._remoteAgentHostService.getConnection(result.address)) { + this._logService.trace(`[WSLRemoteAgentHost] Returning existing connection handle for ${result.address}, connectionId=${result.connectionId}`); + return this._createConnection(result, remoteAgentHostCommand, existing); + } + this._logService.info(`[WSLRemoteAgentHost] Replacing stale connection handle for ${result.address}, connectionId=${result.connectionId}`); + this._connections.delete(result.connectionId); + existing.fireClose(); + existing.dispose(); + this._onDidChangeConnections(); + } + + const handle = new WSLAgentHostConnectionHandle( + result.distro, + result.address, + result.name, + () => this._mainService.disconnect(result.distro), + ); + try { + this._connections.set(result.connectionId, handle); + this._onDidChangeConnections(); + return this._createConnection(result, remoteAgentHostCommand, handle); + } catch (err) { + if (this._connections.get(result.connectionId) === handle) { + this._connections.delete(result.connectionId); + this._onDidChangeConnections(); + } + handle.dispose(); + this._mainService.disconnect(result.distro).catch(() => { /* best effort */ }); + throw err; + } + } + + private _createConnection(result: IWSLConnectResult, remoteAgentHostCommand: string | undefined, handle: WSLAgentHostConnectionHandle): IRemoteAgentHostCreatedConnection { + this._onDidReportConnectProgress({ + connectionKey: result.address, + message: localize('wslProgressHandshake', "Establishing connection to {0}...", result.name), + }); + const completionObserver = this._observeSuccessfulConnection(result); + const transportDisposable = this._createTransportDisposable(result.connectionId, result.distro, handle); + try { + return { + connection: this._relayClientFactory.createClient(this._mainService, result.connectionId, result.address, result, remoteAgentHostCommand), + transportDisposable: toDisposable(() => { + completionObserver.dispose(); + transportDisposable.dispose(); + }), + reconnectTransfersTransportOwnership: true, + }; + } catch (err) { + completionObserver.dispose(); + transportDisposable.dispose(); + throw err; + } + } + + private _observeSuccessfulConnection(result: IWSLConnectResult): IDisposable { + const listener = this._remoteAgentHostService.onDidChangeConnections(() => { + const status = this._remoteAgentHostService.connections.find(connection => connection.address === result.address)?.status; + if (RemoteAgentHostConnectionStatus.isConnected(status)) { + listener?.dispose(); + this._onDidReportConnectProgress({ + connectionKey: result.address, + message: localize('wslProgressFinalizing', "Provisioning agent host in {0}...", result.name), + }); + } else if (!status || RemoteAgentHostConnectionStatus.isIncompatible(status)) { + listener?.dispose(); + } + }); + return listener; + } +} + /** * Renderer-side implementation of {@link IWSLRemoteAgentHostService} that * delegates the actual WSL work to the main process via IPC, then registers - * the resulting connection with the renderer-local {@link IRemoteAgentHostService}. + * a WSL connection factory with the renderer-local {@link IRemoteAgentHostService}. */ export class WSLRemoteAgentHostService extends Disposable implements IWSLRemoteAgentHostService { declare readonly _serviceBrand: undefined; private readonly _mainService: IWSLRemoteAgentHostMainService; + private readonly _connectionFactory: WSLConnectionFactory; private readonly _onDidChangeConnections = this._register(new Emitter()); readonly onDidChangeConnections: Event = this._onDidChangeConnections.event; @@ -141,6 +376,19 @@ export class WSLRemoteAgentHostService extends Disposable implements IWSLRemoteA ); this.onDidReportConnectProgress = Event.any(this._mainService.onDidReportConnectProgress, this._onDidReportLocalConnectProgress.event); + this._connectionFactory = this._register(new WSLConnectionFactory( + this._storageService, + this._mainService, + this._remoteAgentHostService, + this._relayClientFactory, + this._connections, + () => this._onDidChangeConnections.fire(), + progress => this._onDidReportLocalConnectProgress.fire(progress), + () => this._getRemoteAgentHostCommand(), + (connectionId, distro, handle) => this._createTransportDisposable(connectionId, distro, handle), + this._logService, + )); + this._register(this._remoteAgentHostService.registerConnectionFactory(this._connectionFactory)); this._register(this._mainService.onDidCloseConnection(connectionId => { this._logService.info(`[WSLRemoteAgentHost] onDidCloseConnection: connectionId=${connectionId}`); @@ -187,11 +435,12 @@ export class WSLRemoteAgentHostService extends Disposable implements IWSLRemoteA throw new Error('Remote agent host connections are not enabled.'); } - const augmentedConfig = this._augmentConfig(config); + const entry = this._connectionFactory.stageConfiguration(this._augmentConfig({ ...config, userInitiated: config.userInitiated ?? true })); + const address = getEntryAddress(entry); this._logService.info(`[WSLRemoteAgentHost] Connecting to distro ${config.distro}`); - const result = await this._mainService.connect(augmentedConfig); - this._logService.trace(`[WSLRemoteAgentHost] WSL relay established, connectionId=${result.connectionId}`); - return this._setupConnection(result, augmentedConfig.remoteAgentHostCommand); + this._remoteAgentHostService.reconnect(address, true); + await this._remoteAgentHostService.waitForConnection(address); + return this._getConnectionHandle(address); } async disconnect(distro: string): Promise { @@ -199,119 +448,28 @@ export class WSLRemoteAgentHostService extends Disposable implements IWSLRemoteA await this._mainService.disconnect(distro); } - async reconnect(distro: string, name: string): Promise { + async reconnect(distro: string, name: string, userInitiated = true): Promise { if (!this._configurationService.getValue(RemoteAgentHostsEnabledSettingId)) { throw new Error('Remote agent host connections are not enabled.'); } - const commandOverride = this._getRemoteAgentHostCommand(); + const entry = this._connectionFactory.stageEntry(distro, name, userInitiated); + const address = getEntryAddress(entry); this._logService.info(`[WSLRemoteAgentHost] Reconnecting to distro ${distro}`); - const result = await this._mainService.reconnect(distro, name, commandOverride); - return this._setupConnection(result, commandOverride); - } - - /** - * Build the renderer-side handle, do the protocol handshake, and register - * with IRemoteAgentHostService. Any failure after the shared-process tunnel - * was established tears it back down so we don't leak it. - */ - private async _setupConnection(result: IWSLConnectResult, remoteAgentHostCommand: string | undefined): Promise { - const existing = this._connections.get(result.connectionId); - if (existing) { - if (this._remoteAgentHostService.getConnection(result.address)) { - this._logService.trace(`[WSLRemoteAgentHost] Returning existing connection handle for ${result.address}, connectionId=${result.connectionId}`); - return existing; - } - this._logService.info(`[WSLRemoteAgentHost] Replacing stale connection handle for ${result.address}, connectionId=${result.connectionId}`); - this._connections.delete(result.connectionId); - existing.fireClose(); - existing.dispose(); - this._onDidChangeConnections.fire(); - } - - let protocolClient: AgentHostProtocolClient | undefined; - let handle: WSLAgentHostConnectionHandle | undefined; - let registeredHandle = false; - try { - this._onDidReportLocalConnectProgress.fire({ - connectionKey: result.address, - message: localize('wslProgressHandshake', "Establishing connection to {0}...", result.name), - }); - protocolClient = this._relayClientFactory.createClient(this._mainService, result.connectionId, result.address, result, remoteAgentHostCommand); - await protocolClient.connect(); - this._logService.trace('[WSLRemoteAgentHost] Protocol handshake completed'); - - this._onDidReportLocalConnectProgress.fire({ - connectionKey: result.address, - message: localize('wslProgressFinalizing', "Provisioning agent host in {0}...", result.name), - }); - - handle = new WSLAgentHostConnectionHandle( - result.distro, - result.address, - result.name, - () => this._mainService.disconnect(result.distro), - ); - - this._connections.set(result.connectionId, handle); - registeredHandle = true; - this._onDidChangeConnections.fire(); - - const entry: IRemoteAgentHostEntry = { - name: result.name, - connectionToken: result.connectionToken, - connection: { - type: RemoteAgentHostEntryType.WSL, - address: result.address, - distro: result.distro, - }, - }; - - this._cacheDistro(result.distro, result.name); - - await this._remoteAgentHostService.addManagedConnection(entry, protocolClient, this._createTransportDisposable(result.connectionId, result.distro, handle)); - - return handle; - } catch (err) { - this._logService.error('[WSLRemoteAgentHost] Connection setup failed', err); - if (registeredHandle && this._connections.get(result.connectionId) === handle) { - this._connections.delete(result.connectionId); - this._onDidChangeConnections.fire(); - } - handle?.dispose(); - protocolClient?.dispose(); - this._mainService.disconnect(result.distro).catch(() => { /* best effort */ }); - throw err; - } + this._remoteAgentHostService.reconnect(address, userInitiated); + await this._remoteAgentHostService.waitForConnection(address); + return this._getConnectionHandle(address); } getCachedDistros(): readonly IWSLCachedDistro[] { - const raw = this._storageService.get(CACHED_WSL_DISTROS_KEY, StorageScope.APPLICATION); - if (!raw) { - return []; - } - try { - const parsed: unknown = JSON.parse(raw); - if (!Array.isArray(parsed)) { - return []; - } - return parsed.filter((item): item is IWSLCachedDistro => - !!item && typeof item.distro === 'string' && typeof item.name === 'string'); - } catch { - return []; - } - } - - private _cacheDistro(distro: string, name: string): void { - const cached = this.getCachedDistros().filter(d => d.distro !== distro); - this._storeCachedDistros([{ distro, name }, ...cached]); + return readCachedWSLDistros(this._storageService); } private _removeCachedDistro(distro: string): void { const cached = this.getCachedDistros(); const filtered = cached.filter(d => d.distro !== distro); if (filtered.length !== cached.length) { - this._storeCachedDistros(filtered); + storeCachedWSLDistros(this._storageService, filtered); } } @@ -328,16 +486,16 @@ export class WSLRemoteAgentHostService extends Disposable implements IWSLRemoteA const cached = this.getCachedDistros(); const filtered = cached.filter(d => existing.has(d.distro)); if (filtered.length !== cached.length) { - this._storeCachedDistros(filtered); + storeCachedWSLDistros(this._storageService, filtered); } } - private _storeCachedDistros(distros: readonly IWSLCachedDistro[]): void { - if (distros.length === 0) { - this._storageService.remove(CACHED_WSL_DISTROS_KEY, StorageScope.APPLICATION); - } else { - this._storageService.store(CACHED_WSL_DISTROS_KEY, JSON.stringify(distros), StorageScope.APPLICATION, StorageTarget.USER); + private _getConnectionHandle(address: string): WSLAgentHostConnectionHandle { + const handle = [...this._connections.values()].find(candidate => candidate.localAddress === address); + if (!handle) { + throw new Error(`WSL connection handle not found for ${address}.`); } + return handle; } /** diff --git a/src/vs/platform/agentHost/node/agentService.ts b/src/vs/platform/agentHost/node/agentService.ts index 0ffe5c601a06c1..99b387f4b5fdc6 100644 --- a/src/vs/platform/agentHost/node/agentService.ts +++ b/src/vs/platform/agentHost/node/agentService.ts @@ -37,7 +37,7 @@ import type { InvokeChangesetOperationParams, InvokeChangesetOperationResult } f import { AhpErrorCodes, AHP_SESSION_NOT_FOUND, ContentEncoding, JSON_RPC_INTERNAL_ERROR, ProtocolError, ResourceChangeType, ResourceType, ResourceWriteMode, type CreateResourceWatchParams, type CreateResourceWatchResult, type DirectoryEntry, type ResourceCopyParams, type ResourceCopyResult, type ResourceDeleteParams, type ResourceDeleteResult, type ResourceListResult, type ResourceMkdirParams, type ResourceMkdirResult, type ResourceMoveParams, type ResourceMoveResult, type ResourceReadResult, type ResourceResolveParams, type ResourceResolveResult, type ResourceWatchState, type ResourceWriteParams, type ResourceWriteResult, type IStateSnapshot } from '../common/state/sessionProtocol.js'; import { ChangesSummary, ChatInteractivity, ChatOriginKind, MessageAttachmentKind, type Annotation, type AnnotationEntry, type AnnotationOrigin, type AnnotationsState, type ChatOrigin, type Customization, type Message, type MessageAttachment, type MessageResourceAttachment, type TextRange } from '../common/state/protocol/state.js'; import type { ChatPendingMessageSetAction, ChatTurnStartedAction, SessionConfigChangedAction } from '../common/state/protocol/actions.js'; -import { isAhpAutomationCatalogChannel, isAhpAutomationRunChannel, ISessionGitHubState, ISessionGitState, MessageKind, ResponsePartKind, SESSION_META_GITHUB_KEY, SESSION_META_GIT_KEY, SESSION_META_MULTI_ROOT_KEY, SESSION_META_SOURCE_CONTROL_KEY, AH_META_CREATED_BY_SESSION_DB_KEY, readSessionCreationReference, readSessionSpawnDepth, withSessionSpawnDepth, withSessionCreationReference, parseSessionCreationReference, SessionLifecycle, SessionStatus, ToolCallStatus, ToolResultContentType, TurnState, AH_META_WORKSPACELESS_DB_KEY, AH_META_EHCLI_ADOPTED_DB_KEY, AH_META_IS_ARCHIVED_DB_KEY, AH_META_IS_DONE_DB_KEY, AH_META_IS_READ_DB_KEY, buildChatUri, buildDefaultChatUri, buildResourceWatchChannelUri, buildSubagentChatUri, buildSubagentSessionUriPrefix, getErrorResponsePart, isAhpChatChannel, isChatReadOnly, isDefaultChatUri, isSubagentChatUri, isSubagentSession, needsSessionGitStateRefresh, parseChatUri, parseDefaultChatUri, parseRequiredSessionUriFromChatUri, parseResourceWatchChannelUri, parseSessionMultiRootMetadata, parseSubagentSessionUri, readSessionExternal, readSessionGitHubState, readSessionGitState, readSessionMultiRootMetadata, readSessionSourceControlState, readSessionWorkspaceless, withMessageHiddenFromTranscript, withSessionExternal, withSessionGitHubState, withSessionGitState, withSessionMultiRootMetadata, withSessionSourceControlState, withSessionStatusFlag, withSessionWorkspaceless, withSessionEhcliAdopted, withSessionFolderPickerDecision, readSessionFolderPickerDecision, parseSessionFolderPickerDecision, SESSION_META_FOLDER_PICKER_KEY, readSessionEhcliAdoptable, type ISessionSourceControlState, type SessionConfigState, type SessionSummary, type ToolResultSubagentContent, type Turn } from '../common/state/sessionState.js'; +import { isAhpAutomationCatalogChannel, isAhpAutomationRunChannel, ISessionGitHubState, ISessionGitState, MessageKind, ResponsePartKind, SESSION_META_GITHUB_KEY, SESSION_META_GIT_KEY, SESSION_META_MULTI_ROOT_KEY, SESSION_META_SOURCE_CONTROL_KEY, AH_META_CREATED_BY_SESSION_DB_KEY, readSessionCreationReference, readSessionSpawnDepth, withSessionSpawnDepth, withSessionCreationReference, parseSessionCreationReference, SessionLifecycle, SessionStatus, ToolCallStatus, ToolResultContentType, TurnState, AH_META_WORKSPACELESS_DB_KEY, AH_META_EHCLI_ADOPTED_DB_KEY, AH_META_IS_ARCHIVED_DB_KEY, AH_META_IS_DONE_DB_KEY, AH_META_IS_READ_DB_KEY, buildChatUri, buildDefaultChatUri, buildResourceWatchChannelUri, buildSubagentChatUri, buildSubagentSessionUriPrefix, getErrorResponsePart, isAhpChatChannel, isChatReadOnly, isDefaultChatUri, isSubagentChatUri, isSubagentSession, needsSessionGitStateRefresh, parseChatUri, parseDefaultChatUri, parseRequiredSessionUriFromChatUri, parseResourceWatchChannelUri, parseSessionMultiRootMetadata, parseSubagentSessionUri, readSessionExternal, readSessionGitHubState, readSessionGitState, readSessionMultiRootMetadata, readSessionSourceControlState, readSessionWorkspaceless, withMessageHiddenFromTranscript, withSessionExternal, withSessionGitHubState, withSessionGitState, withSessionMultiRootMetadata, withSessionSourceControlState, withSessionStatusFlag, withSessionWorkspaceless, withSessionEhcliAdopted, withSessionEhcliLastMigratedTurn, AH_META_EHCLI_LAST_TURN_DB_KEY, withSessionFolderPickerDecision, readSessionFolderPickerDecision, parseSessionFolderPickerDecision, SESSION_META_FOLDER_PICKER_KEY, readSessionEhcliAdoptable, type ISessionSourceControlState, type SessionConfigState, type SessionSummary, type ToolResultSubagentContent, type Turn } from '../common/state/sessionState.js'; import { readToolCallMeta } from '../common/meta/agentToolCallMeta.js'; import { isHostSnapshotAttachment, toHostSnapshotAttachmentMeta } from '../common/meta/agentSnapshotAttachmentMeta.js'; import { readEphemeralSessionMeta, withEphemeralSessionMeta } from '../common/meta/agentEphemeralSessionMeta.js'; @@ -94,6 +94,7 @@ import { IAgentHostChangesetService, CHANGESET_DB_METADATA_KEYS, META_CHANGES_SU import { GIT_DB_METADATA_KEYS, IAgentHostGitStateService, META_GIT_STATE, META_GITHUB_STATE, META_SOURCE_CONTROL_STATE } from '../common/agentHostGitStateService.js'; import { IAgentHostChangesetOperationService } from '../common/agentHostChangesetOperationService.js'; import { IAgentHostChatContributions } from '../common/agentHostChatContributionsService.js'; +import { IAgentHostStorageService } from './agentHostStorageService.js'; /** * Grace period before an empty, unsubscribed session is garbage-collected @@ -105,14 +106,23 @@ const SESSION_GC_GRACE_MS = 30_000; const DAY_MS = 24 * 60 * 60 * 1000; const EXTERNAL_SESSION_MAX_AGE_MS = 30 * DAY_MS; const RECENT_EXTERNAL_SESSION_LIMIT = 2; -/** - * How many locally created sessions must postdate an external session's last - * update before {@link AgentHostExternalSessionsMode.Recent} stops surfacing it. - */ -const RECENT_EXTERNAL_SUPERSEDING_LOCAL_LIMIT = 2; +const RECENT_LOCAL_SESSION_UPDATE_LIMIT = 2; +const RECENT_LOCAL_SESSION_UPDATES_STORAGE_KEY = 'recentLocalSessionUpdates'; /** A catalog pass slower than this is logged at info, since it delays every session-list refresh. */ const SLOW_LIST_SESSIONS_THRESHOLD_MS = 1_000; +/** A recent update to one local Agent Host session. */ +interface IRecentLocalSessionUpdate { + readonly session: string; + readonly modifiedTime: number; +} + +interface ISessionListComputation { + readonly epoch: number; + readonly promise: Promise; + trailing?: Promise; +} + type AgentHostLegacyMigrationEvent = { provider: string; outcome: 'migrated' | 'skipped' | 'failed'; @@ -204,6 +214,12 @@ function isRecord(value: unknown): value is Record { return typeof value === 'object' && value !== null; } +function isRecentLocalSessionUpdate(value: unknown): value is IRecentLocalSessionUpdate { + return isRecord(value) + && typeof value.session === 'string' + && Number.isFinite(value.modifiedTime); +} + function isPersistedAnnotationEntry(value: unknown): value is AnnotationEntry { if (!isRecord(value) || typeof value.id !== 'string') { return false; @@ -444,6 +460,8 @@ export class AgentService extends Disposable implements IAgentService { private readonly _orchestratorDatabase: IAgentHostDatabase; /** Serializes durable last-modified advances emitted by live session state. */ private _sessionModifiedTimeWrites: Promise = Promise.resolve(); + private readonly _recentLocalSessionUpdateSnapshot: readonly IRecentLocalSessionUpdate[]; + private _recentLocalSessionUpdates: readonly IRecentLocalSessionUpdate[]; private readonly _providerMigrations = new Map(); private readonly _initialProviderMigrations = new Map>(); @@ -593,6 +611,7 @@ export class AgentService extends Disposable implements IAgentService { @IInstantiationService instantiationService: IInstantiationService, @IAgentHostWorktreeIsolation private readonly _worktree: IAgentHostWorktreeIsolation, @IAgentHostProviderService private readonly _providerService: IAgentHostProviderService, + @IAgentHostStorageService private readonly _storageService: IAgentHostStorageService, ) { super(); this._authService = core.authenticationService; @@ -601,6 +620,8 @@ export class AgentService extends Disposable implements IAgentService { this._sessionRegistry = core.sessionRegistry; this._stateManager = core.stateManager; this._configurationService = core.configurationService; + this._recentLocalSessionUpdateSnapshot = this._readRecentLocalSessionUpdates(); + this._recentLocalSessionUpdates = this._recentLocalSessionUpdateSnapshot; this.onMcpNotification = this._providerService.onMcpNotification; this._gitHubEndpointService = collaborators.gitHubEndpointService; this._gitStateService = collaborators.gitStateService; @@ -694,7 +715,14 @@ export class AgentService extends Disposable implements IAgentService { this._register(this._stateManager.onDidChangeSessionSummary(({ session, changes }) => { const meta = this._stateManager.getSessionSummary(session)?._meta; if (changes.modifiedAt !== undefined) { - this._writeSessionModifiedTime(URI.parse(session), Date.parse(changes.modifiedAt)); + const modifiedTime = Date.parse(changes.modifiedAt); + if (!readSessionExternal(meta) + && !isSubagentSession(session) + && !this._stateManager.isEphemeralSession(session) + && !this._stateManager.isIdleProvisionalSession(session)) { + this._recordRecentLocalSessionUpdate(URI.parse(session), modifiedTime); + } + this._writeSessionModifiedTime(URI.parse(session), modifiedTime); } if (changes.modifiedAt !== undefined && this._getExternalSessionsMode() === AgentHostExternalSessionsMode.Recent @@ -712,10 +740,12 @@ export class AgentService extends Disposable implements IAgentService { if (nextMode !== externalSessionsMode) { const previousMode = externalSessionsMode; externalSessionsMode = nextMode; - // The only point past startup where `Recent` re-measures the - // superseding local sessions. - this._invalidateRecentSupersedingCutoff(); this._logService.info(`[AgentService] ${AgentHostShowExternalSessionsConfigKey} changed '${previousMode}' -> '${nextMode}'; queueing session list reconciliation`); + if (this._startupSettled.isOpen() && this._hidesAllExternalSessions(previousMode) && !this._hidesAllExternalSessions(nextMode)) { + for (const provider of this._providerService.getProviders()) { + this._startChatDiscovery(provider, 'external sessions were enabled'); + } + } this._queueSessionListReconciliation(previousMode); } const nextAgentMergeEnabled = this._isAgentMergeEnabled(); @@ -758,6 +788,9 @@ export class AgentService extends Disposable implements IAgentService { * ambient timer of its own. */ markStartupComplete(): void { + if (this._hostStartupComplete) { + return; + } this._hostStartupComplete = true; this._openStartupSettled(); } @@ -774,7 +807,7 @@ export class AgentService extends Disposable implements IAgentService { * compete with startup — pruning stale external sessions, titling external * sessions a provider surfaced without a title, and similar. */ - private _runWhenStartupSettled(name: string, work: () => Promise): void { + private _runWhenStartupSettled(name: string, work: () => void | Promise): void { this._deferredWork = this._deferredWork .then(() => this._startupSettled.wait()) .then(() => this._store.isDisposed ? undefined : work()) @@ -1040,6 +1073,7 @@ export class AgentService extends Disposable implements IAgentService { void this._migrateAndRegisterDiscoveredChats(provider, chats).catch(err => this._logService.warn(`[AgentService] registering discovered chats for provider ${provider.id} failed`, err)); })); + this._setupChatDiscoveryForProvider(provider); subscriptions.add(provider.onDidChangeChatData(e => this._onChatDataChanged(e))); subscriptions.add(provider.onDidSpawnChat(e => this._onChatSpawned(e))); this._providerSubscriptions.set(provider.id, subscriptions); @@ -1054,6 +1088,18 @@ export class AgentService extends Disposable implements IAgentService { } } + private _setupChatDiscoveryForProvider(provider: IAgent): void { + if (this._migrateLegacyEnabledSnapshot === true && provider.ensureChatAdopted) { + this._startChatDiscovery(provider, 'legacy chat migration is enabled'); + } else { + this._runWhenStartupSettled(`external session discovery for ${provider.id}`, () => { + if (!this._hidesAllExternalSessions(this._getExternalSessionsMode())) { + this._startChatDiscovery(provider, 'Agent Host startup settled with external sessions enabled'); + } + }); + } + } + private _onDidRegisterProvider(provider: IAgent): void { this._registerSkillCompletionProvider(); const initialMigration = this._ensureLegacyChatsMigrated(provider); @@ -1869,6 +1915,7 @@ export class AgentService extends Disposable implements IAgentService { await this._sessionRegistry.markProviderBackfilled(provider.id); this._deferredProviderMigrations.delete(provider.id); this._readableProviderCatalogs.add(provider.id); + this._startChatDiscovery(provider, 'legacy migration enumerated the provider catalog'); if (registeredExternal) { this._queueSessionListReconciliation(); } @@ -2007,32 +2054,43 @@ export class AgentService extends Disposable implements IAgentService { } } - /** In-flight list computations, shared per mode until they settle or the registry changes. */ - private readonly _inFlightListSessions = new Map }>(); + /** Active list computations and their optional trailing refresh, shared per mode. */ + private readonly _inFlightListSessions = new Map(); private _registryEpoch = 0; private _invalidateSessionList(): void { this._registryEpoch++; - this._inFlightListSessions.clear(); } async listSessions(mode = this._getExternalSessionsMode()): Promise { const epoch = this._registryEpoch; const inFlight = this._inFlightListSessions.get(mode); - if (inFlight && inFlight.epoch === epoch) { - // Callers own their array; the shared result must not be mutable by one of them. + if (!inFlight) { + return [...await this._startSessionListComputation(mode).promise]; + } + if (inFlight.epoch === epoch) { return [...await inFlight.promise]; } - const promise = this._computeSessions(mode, epoch); - const entry = { epoch, promise }; + if (!inFlight.trailing) { + const startTrailing = () => this._startSessionListComputation(mode).promise; + inFlight.trailing = inFlight.promise.then(startTrailing, startTrailing); + } + return [...await inFlight.trailing]; + } + + private _startSessionListComputation(mode: AgentHostExternalSessionsMode): ISessionListComputation { + const entry: ISessionListComputation = { + epoch: this._registryEpoch, + promise: this._computeSessions(mode), + }; this._inFlightListSessions.set(mode, entry); const clear = () => { - if (this._inFlightListSessions.get(mode) === entry) { + if (!entry.trailing && this._inFlightListSessions.get(mode) === entry) { this._inFlightListSessions.delete(mode); } }; - void promise.then( + void entry.promise.then( () => { clear(); // Only a served listing ends startup: a failed one is retried, and @@ -2042,10 +2100,10 @@ export class AgentService extends Disposable implements IAgentService { }, clear, ); - return [...await promise]; + return entry; } - private async _computeSessions(mode: AgentHostExternalSessionsMode, epoch = this._registryEpoch): Promise { + private async _computeSessions(mode: AgentHostExternalSessionsMode): Promise { this._logService.trace('[AgentService] listSessions computation started'); const startedAt = Date.now(); // The first list waits for registration-time legacy migration if it is still in flight. @@ -2115,8 +2173,8 @@ export class AgentService extends Disposable implements IAgentService { const sessionStr = s.session.toString(); const changesetKeys = this._changesetCoordinator.getListMetadataKeys(sessionStr); const metadataKeys: Record = changesetKeys - ? { customTitle: true, [AH_META_IS_READ_DB_KEY]: true, [AH_META_IS_ARCHIVED_DB_KEY]: true, [AH_META_IS_DONE_DB_KEY]: true, [AH_META_CREATED_BY_SESSION_DB_KEY]: true, [AH_META_WORKSPACELESS_DB_KEY]: true, [AH_META_EHCLI_ADOPTED_DB_KEY]: true, [SESSION_META_MULTI_ROOT_KEY]: true, [SESSION_META_FOLDER_PICKER_KEY]: true, [SESSION_ARTIFACTS_KEY]: true, [CHAT_BACKING_METADATA_KEY]: true, [WORKTREE_META_REPOSITORY_ROOT]: true, ...GIT_DB_METADATA_KEYS, ...changesetKeys } - : { customTitle: true, [AH_META_IS_READ_DB_KEY]: true, [AH_META_IS_ARCHIVED_DB_KEY]: true, [AH_META_IS_DONE_DB_KEY]: true, [AH_META_CREATED_BY_SESSION_DB_KEY]: true, [AH_META_WORKSPACELESS_DB_KEY]: true, [AH_META_EHCLI_ADOPTED_DB_KEY]: true, [SESSION_META_MULTI_ROOT_KEY]: true, [SESSION_META_FOLDER_PICKER_KEY]: true, [SESSION_ARTIFACTS_KEY]: true, [CHAT_BACKING_METADATA_KEY]: true, [WORKTREE_META_REPOSITORY_ROOT]: true, ...GIT_DB_METADATA_KEYS }; + ? { customTitle: true, [AH_META_IS_READ_DB_KEY]: true, [AH_META_IS_ARCHIVED_DB_KEY]: true, [AH_META_IS_DONE_DB_KEY]: true, [AH_META_CREATED_BY_SESSION_DB_KEY]: true, [AH_META_WORKSPACELESS_DB_KEY]: true, [AH_META_EHCLI_ADOPTED_DB_KEY]: true, [AH_META_EHCLI_LAST_TURN_DB_KEY]: true, [SESSION_META_MULTI_ROOT_KEY]: true, [SESSION_META_FOLDER_PICKER_KEY]: true, [SESSION_ARTIFACTS_KEY]: true, [CHAT_BACKING_METADATA_KEY]: true, [WORKTREE_META_REPOSITORY_ROOT]: true, ...GIT_DB_METADATA_KEYS, ...changesetKeys } + : { customTitle: true, [AH_META_IS_READ_DB_KEY]: true, [AH_META_IS_ARCHIVED_DB_KEY]: true, [AH_META_IS_DONE_DB_KEY]: true, [AH_META_CREATED_BY_SESSION_DB_KEY]: true, [AH_META_WORKSPACELESS_DB_KEY]: true, [AH_META_EHCLI_ADOPTED_DB_KEY]: true, [AH_META_EHCLI_LAST_TURN_DB_KEY]: true, [SESSION_META_MULTI_ROOT_KEY]: true, [SESSION_META_FOLDER_PICKER_KEY]: true, [SESSION_ARTIFACTS_KEY]: true, [CHAT_BACKING_METADATA_KEY]: true, [WORKTREE_META_REPOSITORY_ROOT]: true, ...GIT_DB_METADATA_KEYS }; const m = await ref.object.getMetadataObject(metadataKeys); // This session is an internal peer-chat backing (e.g. a // Claude peer chat's SDK session, enumerated by the agent's @@ -2173,6 +2231,9 @@ export class AgentService extends Disposable implements IAgentService { if (m[AH_META_EHCLI_ADOPTED_DB_KEY] !== undefined) { updated = { ...updated, _meta: withSessionEhcliAdopted(updated._meta, m[AH_META_EHCLI_ADOPTED_DB_KEY] === 'true') }; } + if (m[AH_META_EHCLI_LAST_TURN_DB_KEY] !== undefined) { + updated = { ...updated, _meta: withSessionEhcliLastMigratedTurn(updated._meta, m[AH_META_EHCLI_LAST_TURN_DB_KEY]) }; + } const multiRoot = parseSessionMultiRootMetadata(m[SESSION_META_MULTI_ROOT_KEY]); if (multiRoot) { updated = { ...updated, _meta: withSessionMultiRootMetadata(updated._meta, multiRoot) }; @@ -2270,7 +2331,7 @@ export class AgentService extends Disposable implements IAgentService { const combined = additions.length > 0 ? [...withStatus, ...additions] : withStatus; const now = Date.now(); const recentSessionKeys = mode === AgentHostExternalSessionsMode.Recent - ? this._getRecentSessionKeys(combined, now, this._resolveRecentSupersedingCutoff(allRegistered, epoch)) + ? this._getRecentSessionKeys(combined, now) : undefined; const visible: IAgentSessionMetadata[] = []; // Adoptable-legacy rows are withheld by migrate-legacy, not by the external mode. @@ -2326,16 +2387,22 @@ export class AgentService extends Disposable implements IAgentService { return this._configurationService.getRootValue(platformRootSchema, AgentHostShowExternalSessionsConfigKey) ?? AgentHostExternalSessionsMode.None; } + private _startChatDiscovery(provider: IAgent, reason: string): void { + void provider.startChatDiscovery?.().catch(error => + this._logService.warn(`[AgentService] Chat discovery for provider ${provider.id} failed after ${reason}`, error)); + } + private _isExternalSessionOlderThanMaxAge(modifiedTime: number, now: number): boolean { return modifiedTime < now - EXTERNAL_SESSION_MAX_AGE_MS; } - private _getRecentSessionKeys(sessions: readonly IAgentSessionMetadata[], now: number, supersededBefore: number | undefined): ReadonlySet { + private _getRecentSessionKeys(sessions: readonly IAgentSessionMetadata[], now: number): ReadonlySet { + const supersededBefore = this._getRecentLocalSessionUpdateCutoff(now); const recentExternalSessions = sessions .filter(session => readSessionExternal(session._meta) && !readSessionEhcliAdoptable(session._meta) && session.modifiedTime >= now - 7 * DAY_MS - && (supersededBefore === undefined || session.modifiedTime >= supersededBefore)) + && session.modifiedTime >= supersededBefore) .sort((a, b) => { const timeDifference = b.modifiedTime - a.modifiedTime; if (timeDifference !== 0) { @@ -2349,47 +2416,60 @@ export class AgentService extends Disposable implements IAgentService { return new Set(recentExternalSessions.map(session => session.session.toString())); } - /** - * Start time of the {@link RECENT_EXTERNAL_SUPERSEDING_LOCAL_LIMIT}-th most - * recently created local session, or `undefined` while fewer exist. `Recent` - * drops external sessions last updated before it. - */ - private _recentSupersedingCutoff: number | undefined; - private _hasRecentSupersedingCutoff = false; + private _getRecentLocalSessionUpdateCutoff(now: number): number { + return this._recentLocalSessionUpdateSnapshot[RECENT_LOCAL_SESSION_UPDATE_LIMIT - 1]?.modifiedTime ?? now - 7 * DAY_MS; + } - /** - * Snapshots the cutoff from the registry, which — unlike the hydrated - * metadata — never drops a local session because its provider is - * unavailable or its metadata read failed. Sending a first message - * materializes a local session, so a per-listing cutoff would rotate an - * external row out of the list mid-use. Committed only while `epoch` still - * holds, so a discarded pass cannot freeze an undercounted value. - */ - private _resolveRecentSupersedingCutoff(registered: readonly IRegisteredSession[], epoch: number): number | undefined { - if (this._hasRecentSupersedingCutoff) { - return this._recentSupersedingCutoff; - } - // Idle provisional sessions are the composer's eagerly-created - // placeholder, not sessions the user started. - const localStartTimes = registered - .filter(entry => !entry.external - && Number.isFinite(entry.startTime) - && !this._stateManager.isIdleProvisionalSession(entry.session.toString())) - .map(entry => entry.startTime) - .sort((a, b) => b - a); - const cutoff = localStartTimes.length >= RECENT_EXTERNAL_SUPERSEDING_LOCAL_LIMIT - ? localStartTimes[RECENT_EXTERNAL_SUPERSEDING_LOCAL_LIMIT - 1] - : undefined; - if (epoch === this._registryEpoch) { - this._recentSupersedingCutoff = cutoff; - this._hasRecentSupersedingCutoff = true; + private _recordRecentLocalSessionUpdate(session: URI, modifiedTime: number): void { + if (!Number.isFinite(modifiedTime)) { + return; + } + + const sessionKey = session.toString(); + const existing = this._recentLocalSessionUpdates.find(entry => entry.session === sessionKey); + if (existing && existing.modifiedTime >= modifiedTime) { + return; + } + + const next = [ + ...this._recentLocalSessionUpdates.filter(entry => entry.session !== sessionKey), + { session: sessionKey, modifiedTime }, + ] + .sort((a, b) => b.modifiedTime - a.modifiedTime || a.session.localeCompare(b.session)) + .slice(0, RECENT_LOCAL_SESSION_UPDATE_LIMIT); + if (next.length === this._recentLocalSessionUpdates.length + && next.every((entry, index) => entry.session === this._recentLocalSessionUpdates[index].session + && entry.modifiedTime === this._recentLocalSessionUpdates[index].modifiedTime)) { + return; + } + + this._recentLocalSessionUpdates = next; + if (!this._storageService.loadError) { + this._storageService.set(RECENT_LOCAL_SESSION_UPDATES_STORAGE_KEY, next); } - return cutoff; } - private _invalidateRecentSupersedingCutoff(): void { - this._hasRecentSupersedingCutoff = false; - this._recentSupersedingCutoff = undefined; + private _readRecentLocalSessionUpdates(): readonly IRecentLocalSessionUpdate[] { + if (this._storageService.loadError) { + this._logService.warn('[AgentService] Recent local session updates could not be restored because Agent Host storage failed to load.'); + return []; + } + const stored = this._storageService.get(RECENT_LOCAL_SESSION_UPDATES_STORAGE_KEY); + if (stored === undefined) { + return []; + } + if (!Array.isArray(stored) + || stored.length > RECENT_LOCAL_SESSION_UPDATE_LIMIT + || !stored.every(isRecentLocalSessionUpdate)) { + this._logService.warn('[AgentService] Ignoring invalid persisted recent local session updates.'); + return []; + } + const updates: readonly IRecentLocalSessionUpdate[] = stored; + if (new Set(updates.map(entry => entry.session)).size !== updates.length) { + this._logService.warn('[AgentService] Ignoring persisted recent local session updates with duplicate sessions.'); + return []; + } + return updates.toSorted((a, b) => b.modifiedTime - a.modifiedTime || a.session.localeCompare(b.session)); } private _shouldIncludeSession( @@ -2530,7 +2610,7 @@ export class AgentService extends Disposable implements IAgentService { previouslyExposed.add(session); } const listed = previousMode !== undefined - ? await this._resolveModeChangeVisibility(await this.listSessions(AgentHostExternalSessionsMode.Last30Days), previousMode, previouslyExposed) + ? this._resolveModeChangeVisibility(await this.listSessions(AgentHostExternalSessionsMode.Last30Days), previousMode, previouslyExposed) : await this.listSessions(); const visible = new Set(); let published = 0; @@ -2584,20 +2664,15 @@ export class AgentService extends Disposable implements IAgentService { * mode and the mode is just a parameter to {@link _shouldIncludeSession}. * Adds what `previousMode` had exposed into `previouslyExposed`. */ - private async _resolveModeChangeVisibility( + private _resolveModeChangeVisibility( superset: readonly IAgentSessionMetadata[], previousMode: AgentHostExternalSessionsMode, previouslyExposed: Set, - ): Promise { + ): IAgentSessionMetadata[] { const now = Date.now(); const mode = this._getExternalSessionsMode(); - // The pass above ran as `Last30Days`, so it never snapshotted the cutoff. - const epoch = this._registryEpoch; - const supersededBefore = previousMode === AgentHostExternalSessionsMode.Recent || mode === AgentHostExternalSessionsMode.Recent - ? this._resolveRecentSupersedingCutoff(await this._listRegisteredSessions(), epoch) - : undefined; const recentKeysFor = (candidate: AgentHostExternalSessionsMode) => candidate === AgentHostExternalSessionsMode.Recent - ? this._getRecentSessionKeys(superset, now, supersededBefore) + ? this._getRecentSessionKeys(superset, now) : undefined; const previousRecentKeys = recentKeysFor(previousMode); @@ -2710,6 +2785,7 @@ export class AgentService extends Disposable implements IAgentService { this._createProviderSession(provider, config, deferWorktreeCreation), ]); const session = created.session; + const isIdleProvisional = created.provisional === true && !config?.importConversation; this._logService.trace(`[AgentService] createSession: initialization complete`); const creationReference = readSessionCreationReference(config?._meta); if (creationReference && !isEphemeral) { @@ -2728,7 +2804,9 @@ export class AgentService extends Disposable implements IAgentService { () => this._sessionRegistry.tombstone(session), `tombstoning ephemeral session ${session.toString()}`, ); - this._invalidateSessionList(); + if (!isIdleProvisional) { + this._invalidateSessionList(); + } } catch (err) { await this._rollbackProviderSession(provider, session); throw err; @@ -2740,7 +2818,9 @@ export class AgentService extends Disposable implements IAgentService { () => this._sessionRegistry.register(session, { provider: provider.id, startTime: registeredAt, modifiedTime: registeredAt, source: 'explicit' }, { checkTombstone: false }), `registration for ${session.toString()}`, ); - this._invalidateSessionList(); + if (!isIdleProvisional) { + this._invalidateSessionList(); + } } catch (err) { await this._rollbackProviderSession(provider, session); throw err; @@ -2780,7 +2860,7 @@ export class AgentService extends Disposable implements IAgentService { // updates while resolving that snapshot; without a state entry those // actions are rejected as targeting an unknown session and custom agents // can disappear from the picker permanently. - const provisionalState = created.provisional && !config?.importConversation + const provisionalState = isIdleProvisional ? (() => { const summary = this._buildInitialSummary(provider, session, config, created, ''); const state = this._stateManager.createSession(summary, { emitNotification: false }); @@ -3876,6 +3956,7 @@ export class AgentService extends Disposable implements IAgentService { const sessionKey = session.toString(); this._cancelPendingSessionGc(session); const isEphemeral = this._stateManager.isEphemeralSession(sessionKey); + const isIdleProvisional = this._stateManager.isIdleProvisionalSession(sessionKey); this._stateManager.invalidateSessionChatResolutions(session.toString()); const sessionChats = this._stateManager.getSessionState(session.toString())?.chats ?? []; for (const chat of sessionChats) { @@ -3900,7 +3981,9 @@ export class AgentService extends Disposable implements IAgentService { `unregistration for ${session.toString()}`, ); } - this._invalidateSessionList(); + if (!isIdleProvisional) { + this._invalidateSessionList(); + } if (provider) { this._providerService.releaseSession(session.toString()); this._clearDownloadProgressInterest(session.toString()); @@ -5217,9 +5300,11 @@ export class AgentService extends Disposable implements IAgentService { // worktree-isolated sessions. No-op for folder / primary-checkout cwds. let adoptedWorktree = false; if (adopted && this._worktree.supported) { - // The predecessor recorded this worktree but its checkout is gone, so it - // cannot be probed; seed the same metadata a native session persists at - // creation and let resume recreate it. + // The predecessor recorded this worktree; seed the same metadata a native + // session persists at creation. When its checkout is gone this is the only + // way to recover it (resume recreates it from the branch); when the checkout + // still exists this carries the authoritatively recorded base branch, which + // the probe-based bridge below could not recover without a remote (#333642). if (adoptionWorktree) { try { await this._worktree.recordAdoptedWorktreeMetadata(session, adoptionWorktree); @@ -5319,6 +5404,7 @@ export class AgentService extends Disposable implements IAgentService { configValues: true, [AH_META_WORKSPACELESS_DB_KEY]: true, [AH_META_EHCLI_ADOPTED_DB_KEY]: true, + [AH_META_EHCLI_LAST_TURN_DB_KEY]: true, [AH_META_CREATED_BY_SESSION_DB_KEY]: true, [SESSION_META_MULTI_ROOT_KEY]: true, [SESSION_ARTIFACTS_KEY]: true, @@ -5383,6 +5469,9 @@ export class AgentService extends Disposable implements IAgentService { if (m[AH_META_EHCLI_ADOPTED_DB_KEY] !== undefined) { sessionMetadata = withSessionEhcliAdopted(sessionMetadata, m[AH_META_EHCLI_ADOPTED_DB_KEY] === 'true'); } + if (m[AH_META_EHCLI_LAST_TURN_DB_KEY] !== undefined) { + sessionMetadata = withSessionEhcliLastMigratedTurn(sessionMetadata, m[AH_META_EHCLI_LAST_TURN_DB_KEY]); + } const creationReference = parseSessionCreationReference(m[AH_META_CREATED_BY_SESSION_DB_KEY]); if (creationReference) { sessionMetadata = withSessionCreationReference(sessionMetadata, creationReference); diff --git a/src/vs/platform/agentHost/node/claude/claudeAgent.ts b/src/vs/platform/agentHost/node/claude/claudeAgent.ts index 67111ec4d61964..5d05bf9af346b0 100644 --- a/src/vs/platform/agentHost/node/claude/claudeAgent.ts +++ b/src/vs/platform/agentHost/node/claude/claudeAgent.ts @@ -429,12 +429,7 @@ export class ClaudeAgent extends Disposable implements IAgent { private readonly _onDidSpawnChat = this._register(new Emitter()); readonly onDidSpawnChat: Event = this._onDidSpawnChat.event; - private readonly _onDidDiscoverChats = this._register(new Emitter({ - // Discovery is provider-owned and only has observable value once the host - // subscribes. Registered chats remain independently available through - // listChatsToMigrate(). - onDidAddFirstListener: () => { void this._startClaudeCodeChatDiscovery(); }, - })); + private readonly _onDidDiscoverChats = this._register(new Emitter()); readonly onDidDiscoverChats = this._onDidDiscoverChats.event; private _claudeCodeChatDiscovery: Promise | undefined; @@ -2054,6 +2049,10 @@ export class ClaudeAgent extends Disposable implements IAgent { })); } + startChatDiscovery(): Promise { + return this._startClaudeCodeChatDiscovery(); + } + async listChatsToMigrate(): Promise { if (!(await this._sdkService.canLoadWithoutDownload())) { this._logService.info('[Claude] SDK not downloaded yet; deferring the migratable chat list'); diff --git a/src/vs/platform/agentHost/node/codex/codexAgent.ts b/src/vs/platform/agentHost/node/codex/codexAgent.ts index 8c22ffc22bda5d..b2b7c61c777fea 100644 --- a/src/vs/platform/agentHost/node/codex/codexAgent.ts +++ b/src/vs/platform/agentHost/node/codex/codexAgent.ts @@ -1118,10 +1118,9 @@ export class CodexAgent extends Disposable implements IAgent { private _transientAccountConnection: IConnectionReady | undefined; /** Owns a one-off connection even while its initialize handshake is pending. */ private _transientConnectionCancellation: CancellationTokenSource | undefined; - private readonly _onDidDiscoverChats = this._register(new Emitter({ - onDidAddFirstListener: () => { void this._startCodexChatDiscovery(); }, - })); + private readonly _onDidDiscoverChats = this._register(new Emitter()); readonly onDidDiscoverChats = this._onDidDiscoverChats.event; + private _chatDiscoveryRequested = false; private _codexChatDiscovery: Promise | undefined; private _modelsRefreshPromise: Promise | undefined; private readonly _modelRefreshSequencer = new Sequencer(); @@ -2113,7 +2112,7 @@ export class CodexAgent extends Disposable implements IAgent { // flight may have observed the inactive state and skipped Codex models. void this._queueModelRefresh(); void this._refreshProviderConfiguration(); - if (this._onDidDiscoverChats.hasListeners()) { + if (this._chatDiscoveryRequested) { void this._startCodexChatDiscovery(); } } @@ -6591,11 +6590,10 @@ export class CodexAgent extends Disposable implements IAgent { } async listChatsToMigrate(): Promise { - // Registration-time migration is ambient. Report an empty initial catalog - // so provider registration can finish without starting Codex; activated - // discovery later emits both known (internal) and unknown (external) chats. + // Registration-time migration is ambient. Defer until explicit Codex use + // rather than claiming an authoritative empty catalog without enumerating. if (!this._activated) { - return []; + return AgentChatMigrationDeferred; } if (!(await this._isSdkResolvableWithoutDownload())) { this._logService.info('[Codex] SDK not downloaded yet; deferring the migratable chat list'); @@ -6612,6 +6610,11 @@ export class CodexAgent extends Disposable implements IAgent { return known.filter((chat): chat is IAgentChatMetadata => chat !== undefined); } + startChatDiscovery(): Promise { + this._chatDiscoveryRequested = true; + return this._startCodexChatDiscovery(); + } + private _startCodexChatDiscovery(): Promise { if (this._isShuttingDown || this._store.isDisposed || !this._activated) { return Promise.resolve(); diff --git a/src/vs/platform/agentHost/node/copilot/copilotAgent.ts b/src/vs/platform/agentHost/node/copilot/copilotAgent.ts index a2e689442aac85..357da2d57ed8b9 100644 --- a/src/vs/platform/agentHost/node/copilot/copilotAgent.ts +++ b/src/vs/platform/agentHost/node/copilot/copilotAgent.ts @@ -59,7 +59,7 @@ import type { ErrorInfo } from '../../common/state/protocol/common/state.js'; import { ProtectedResourceMetadata, type AgentSelection, type ChildCustomizationType, type ConfigPropertySchema, type ConfigSchema, type CustomizationEnablement, type ModelSelection, type ToolDefinition } from '../../common/state/protocol/state.js'; import { ActionType, AuthRequiredReason, type AuthRequiredParams, type SessionAction } from '../../common/state/sessionActions.js'; import { areAdditionalWorkingDirectoriesEqual } from '../../common/state/sessionWorkingDirectories.js'; -import { AgentCustomization, CustomizationLoadStatus, CustomizationType, RuleCustomization, ChatInputResponseKind, SkillCustomization, customizationId, buildChatUri, buildDefaultChatUri, AH_META_WORKSPACELESS_DB_KEY, AH_META_IS_ARCHIVED_DB_KEY, AH_META_EHCLI_ADOPTED_DB_KEY, AH_META_IS_READ_DB_KEY, isDefaultChatUri, withSessionEhcliAdoptable, type ChildCustomization, type ClientPluginCustomization, type Customization, type DirectoryCustomization, type HookCustomization, type ISessionFolderPickerDecision, type MessageAttachment, type PendingMessage, type PluginCustomization, type PolicyState, type ChatInputAnswer, type ToolCallResult, type Turn, type UsageInfo } from '../../common/state/sessionState.js'; +import { AgentCustomization, CustomizationLoadStatus, CustomizationType, RuleCustomization, ChatInputResponseKind, SkillCustomization, customizationId, buildChatUri, buildDefaultChatUri, AH_META_WORKSPACELESS_DB_KEY, AH_META_IS_ARCHIVED_DB_KEY, AH_META_EHCLI_ADOPTED_DB_KEY, AH_META_EHCLI_LAST_TURN_DB_KEY, AH_META_IS_READ_DB_KEY, isDefaultChatUri, withSessionEhcliAdoptable, type ChildCustomization, type ClientPluginCustomization, type Customization, type DirectoryCustomization, type HookCustomization, type ISessionFolderPickerDecision, type MessageAttachment, type PendingMessage, type PluginCustomization, type PolicyState, type ChatInputAnswer, type ToolCallResult, type Turn, type UsageInfo } from '../../common/state/sessionState.js'; import { getByokLmAgentModelId, resolveByokLmEnablement } from '../../common/agentHostByokLm.js'; import { isCustomizationEnabled } from '../../common/customizationEnablement.js'; import { ActiveClientToolSet, structuralToolsEqual } from '../activeClientState.js'; @@ -67,7 +67,7 @@ import { IAgentConfigurationService } from '../agentConfigurationService.js'; import { IAgentHostManagedSettingsService } from '../agentHostManagedSettingsService.js'; import { IAgentHostGitHubEndpointService } from '../agentHostGitHubEndpointService.js'; import { IAgentHostCompletions } from '../agentHostCompletions.js'; -import { IAgentHostGitService } from '../../common/agentHostGitService.js'; +import { IAgentHostGitService, META_DIFF_BASE_BRANCH } from '../../common/agentHostGitService.js'; import { applyMcpServerEnablement, buildMcpTopLevelCustomizationId, type IMcpServerRuntimeState } from '../shared/mcpCustomizationController.js'; import { IAgentHostCustomizationEnablementService } from '../agentHostCustomizationEnablementService.js'; import { getSdkMcpServerEnablement, isCustomizationSdkEligible, resolveCustomizationEnablement } from '../shared/customizationEnablementGate.js'; @@ -720,9 +720,7 @@ export class CopilotAgent extends Disposable implements IAgent { * Fires when the native chat catalog may have changed. The {@link AgentService} * responds with an additive discovery pass. */ - private readonly _onDidDiscoverChats = this._register(new Emitter({ - onDidAddFirstListener: () => { void this._startCopilotChatDiscovery(); }, - })); + private readonly _onDidDiscoverChats = this._register(new Emitter()); readonly onDidDiscoverChats = this._onDidDiscoverChats.event; /** * Per-session MCP notifications, fanned in from every active @@ -2459,6 +2457,10 @@ export class CopilotAgent extends Disposable implements IAgent { this._knownSessionsFilter = filter; } + startChatDiscovery(): Promise { + return this._startCopilotChatDiscovery(); + } + /** * One memoized initial discovery attempt, mirroring Claude and Codex. The * CLI client may still be starting when the first discovery listener @@ -3372,16 +3374,34 @@ export class CopilotAgent extends Disposable implements IAgent { } /** - * Worktree identity the extension host recorded, when its checkout is gone but - * the repository remains. Resume recreates the worktree from this, matching how - * a natively worktree-isolated session recovers. + * Worktree identity the extension host recorded, so the migrated session diffs + * against the same base branch the worktree was branched from. Returned when the + * repository still exists, covering two cases: + * + * - The checkout is gone: resume recreates the worktree from this, matching how + * a natively worktree-isolated session recovers. + * - The checkout still exists but the marker carries a base branch: the recorded + * base is authoritative and independent of `refs/remotes/origin/HEAD`, which is + * the only source the probe-based bridge ({@link IAgentHostWorktreeIsolation.adoptExistingWorktreeMetadata}) + * has. Without this, a worktree session in a repository with no remote (or no + * `origin/HEAD`) persists no base branch, so its Branch Changes diff falls back + * to `HEAD` and every committed-on-branch change is invisible (#333642). + * + * A still-existing checkout whose marker has no base branch is left to the + * probe-based bridge so the pre-existing `origin/HEAD` fallback is preserved. */ private async _extensionHostCliAdoptedWorktree(sessionId: string): Promise { const worktree = (await this._readExtensionHostCliMarker(sessionId))?.worktreeProperties; if (!worktree?.worktreePath || !worktree.repositoryPath || !worktree.branchName) { return undefined; } - if (await this._isExistingDirectory(worktree.worktreePath) || !(await this._isExistingDirectory(worktree.repositoryPath))) { + if (!(await this._isExistingDirectory(worktree.repositoryPath))) { + return undefined; + } + // The checkout still exists: only take over from the probe-based bridge when + // the marker gives us an authoritative base branch to persist; otherwise let + // the probe resolve it (e.g. from `origin/HEAD`) exactly as before. + if (await this._isExistingDirectory(worktree.worktreePath) && !worktree.baseBranchName) { return undefined; } return { @@ -3393,10 +3413,16 @@ export class CopilotAgent extends Disposable implements IAgent { } /** - * Records the durable adopted-legacy marker on a session adopted by a build - * that predates it. Without this those sessions keep the extension-host marker - * but no provenance, so a worktree one stays filtered out of the window opened - * on its repository. Keyed off the marker, so it never claims a native session. + * Repairs durable metadata on a legacy Copilot CLI session adopted by a build + * that predates it. Keyed off the extension-host marker, so it never claims a + * native session. Backfills, when missing: + * - the adopted-legacy provenance marker (without it a worktree session stays + * filtered out of the window opened on its repository); + * - the Branch Changes base branch from the marker's recorded worktree base, so + * a session migrated before this was persisted (e.g. a no-remote repo whose + * `origin/HEAD` could not answer) stops anchoring its diff to `HEAD` (#333642); + * - the last-migrated-turn boundary, so the chat editor can surface the + * session-wide changes on the migrated turn. */ private async _backfillAdoptedLegacyMarker(session: URI, sessionId: string): Promise { const ref = await this._sessionDataService.tryOpenDatabase(session); @@ -3404,16 +3430,40 @@ export class CopilotAgent extends Disposable implements IAgent { return; } try { - if (await ref.object.getMetadata(AH_META_EHCLI_ADOPTED_DB_KEY) !== undefined) { + const [existingMarker, existingBaseBranch, existingLastTurn] = await Promise.all([ + ref.object.getMetadata(AH_META_EHCLI_ADOPTED_DB_KEY), + ref.object.getMetadata(META_DIFF_BASE_BRANCH), + ref.object.getMetadata(AH_META_EHCLI_LAST_TURN_DB_KEY), + ]); + if (existingMarker !== undefined && existingBaseBranch !== undefined && existingLastTurn !== undefined) { return; } if (!(await this._isExtensionHostCliSession(sessionId))) { return; } - await ref.object.setMetadata(AH_META_EHCLI_ADOPTED_DB_KEY, 'true'); - this._logService.info(`[Copilot] Backfilled the adopted-legacy marker for ${sessionId}, migrated before it was recorded`); + const work: Promise[] = []; + if (existingMarker === undefined) { + work.push(ref.object.setMetadata(AH_META_EHCLI_ADOPTED_DB_KEY, 'true')); + } + if (existingBaseBranch === undefined) { + const recordedBase = (await this._readExtensionHostCliMarker(sessionId))?.worktreeProperties?.baseBranchName; + if (recordedBase) { + work.push(ref.object.setMetadata(META_DIFF_BASE_BRANCH, recordedBase)); + } + } + if (existingLastTurn === undefined) { + const lastMigratedTurnId = await this._readExtensionHostCliLastTurnId(sessionId); + if (lastMigratedTurnId) { + work.push(ref.object.setMetadata(AH_META_EHCLI_LAST_TURN_DB_KEY, lastMigratedTurnId)); + } + } + if (work.length === 0) { + return; + } + await Promise.all(work); + this._logService.info(`[Copilot] Backfilled durable metadata for ${sessionId} (marker=${existingMarker === undefined} baseBranch=${existingBaseBranch === undefined} lastTurn=${existingLastTurn === undefined}), migrated before it was recorded`); } catch (err) { - this._logService.warn(`[Copilot] Failed to backfill the adopted-legacy marker for ${sessionId}`, err); + this._logService.warn(`[Copilot] Failed to backfill durable metadata for ${sessionId}`, err); } finally { ref.dispose(); } @@ -3449,7 +3499,9 @@ export class CopilotAgent extends Disposable implements IAgent { const sdkWorkingDirectory = typeof sdkMetadata?.context?.workingDirectory === 'string' ? sdkMetadata.context.workingDirectory : undefined; // A deleted worktree is recoverable the same way a native session recovers // one: keep it as the working directory and let resume recreate it from the - // recorded branch. + // recorded branch. A worktree whose checkout still exists is also bridged + // (when the marker records its base branch) so the recorded base is + // persisted for the Branch Changes diff even without a remote (#333642). const adoptedWorktree = await this._extensionHostCliAdoptedWorktree(sessionId); const workingDirectory = adoptedWorktree?.worktreePath ?? (sdkWorkingDirectory && await this._isExistingDirectory(sdkWorkingDirectory) ? URI.file(sdkWorkingDirectory) : undefined) @@ -3460,7 +3512,8 @@ export class CopilotAgent extends Disposable implements IAgent { this._logService.warn(`[Copilot] Adoption skipped for ${sessionId}: no usable working directory (sdk='${sdkWorkingDirectory ?? '(none)'}' exists=${sdkWorkingDirectory ? await this._isExistingDirectory(sdkWorkingDirectory) : false}, no recorded worktree, no marker fallback). The session stays on the legacy provider.`); return { adopted: false, eligible: true, reason: 'workingDirectoryMissing' }; } - this._logService.info(`[Copilot] Adopting legacy session ${sessionId} in place (reusing on-disk events.jsonl): cwd=${workingDirectory.fsPath}${adoptedWorktree ? ` worktree=${adoptedWorktree.worktreePath.fsPath} branch=${adoptedWorktree.branchName} base=${adoptedWorktree.baseBranch ?? '(none)'} repo=${adoptedWorktree.repositoryRoot.fsPath} (checkout missing, will be recreated on resume)` : ''}`); + const worktreeCheckoutMissing = adoptedWorktree ? !(await this._isExistingDirectory(adoptedWorktree.worktreePath.fsPath)) : false; + this._logService.info(`[Copilot] Adopting legacy session ${sessionId} in place (reusing on-disk events.jsonl): cwd=${workingDirectory.fsPath}${adoptedWorktree ? ` worktree=${adoptedWorktree.worktreePath.fsPath} branch=${adoptedWorktree.branchName} base=${adoptedWorktree.baseBranch ?? '(none)'} repo=${adoptedWorktree.repositoryRoot.fsPath}${worktreeCheckoutMissing ? ' (checkout missing, will be recreated on resume)' : ' (checkout present)'}` : ''}`); // Resolve the project from the SDK-derived cwd (authoritative) — the // caller may not have supplied a working directory (e.g. the chat // editor), so we cannot trust a hint. @@ -3490,7 +3543,11 @@ export class CopilotAgent extends Disposable implements IAgent { // `isolation: 'folder'` keeps the session in place in the reused cwd — // a git repo would otherwise default to worktree and show a spurious // "Creating worktree…". - await this._storeSessionMetadata(session, undefined, workingDirectory, [workingDirectory], workingDirectory, project, project !== undefined, { [SessionConfigKey.Isolation]: 'folder' }, adoptedTitle, /* markRead */ true, archived, /* ehcliAdopted */ true); + // The migration boundary: the id of the last turn recorded on disk, so the + // chat editor can substitute the session-wide changeset for that migrated + // (checkpoint-less) turn without misattributing it to a later, post-adoption turn. + const lastMigratedTurnId = await this._readExtensionHostCliLastTurnId(sessionId); + await this._storeSessionMetadata(session, undefined, workingDirectory, [workingDirectory], workingDirectory, project, project !== undefined, { [SessionConfigKey.Isolation]: 'folder' }, adoptedTitle, /* markRead */ true, archived, /* ehcliAdopted */ true, lastMigratedTurnId); await this._adoptLegacyTurnUsage(session, sessionId); this._logService.info(`[Copilot] Adopted legacy session ${sessionId}: project=${project ? project.uri.fsPath : '(unresolved)'} archived=${archived} title=${adoptedTitle !== undefined ? (cliName ? 'name' : customTitle ? 'custom' : 'summary') : 'none'} worktreeBridged=${!!adoptedWorktree}`); return { adopted: true, eligible: true, reason: 'adopted', ...(adoptedWorktree ? { worktree: adoptedWorktree } : {}) }; @@ -3553,6 +3610,35 @@ export class CopilotAgent extends Disposable implements IAgent { } } + /** + * The id of the final turn recorded in the extension host's request sidecar — + * the migration boundary. Best-effort: absent for sessions predating credit + * tracking, in which case the chat editor simply keeps no migrated-turn fallback. + */ + private async _readExtensionHostCliLastTurnId(sessionId: string): Promise { + const raw = await fs.readFile(this._extensionHostCliSidecarPath(sessionId, 'vscode.requests.metadata.json'), 'utf8').catch(() => undefined); + if (raw === undefined) { + return undefined; + } + try { + const parsed: unknown = JSON.parse(raw); + if (!Array.isArray(parsed)) { + return undefined; + } + // Entries are in turn order; the last valid `copilotRequestId` is the id + // `mapSessionEvents` restores the final turn under. + for (let i = parsed.length - 1; i >= 0; i--) { + const turnId = (parsed[i] as IExtensionHostCliRequestDetails | undefined)?.copilotRequestId; + if (typeof turnId === 'string' && turnId) { + return turnId; + } + } + } catch { + // Malformed sidecar: treat as no recorded boundary. + } + return undefined; + } + /** Materializes a provisional chat into a real SDK session immediately before first send. */ private async _materializeProvisional(sessionId: string, resolvedWorkingDirectories?: readonly URI[]): Promise { const provisional = this._provisionalSessions.get(sessionId); @@ -5118,7 +5204,7 @@ export class CopilotAgent extends Disposable implements IAgent { } - private async _storeSessionMetadata(session: URI, model: ModelSelection | undefined, workingDirectory: URI | undefined, workingDirectories: readonly URI[] | undefined, customizationDirectory: URI | undefined, project: IAgentSessionProjectInfo | undefined, projectResolved = project !== undefined, configValues?: Record, customTitle?: string, markRead?: boolean, archived?: boolean, ehcliAdopted?: boolean): Promise { + private async _storeSessionMetadata(session: URI, model: ModelSelection | undefined, workingDirectory: URI | undefined, workingDirectories: readonly URI[] | undefined, customizationDirectory: URI | undefined, project: IAgentSessionProjectInfo | undefined, projectResolved = project !== undefined, configValues?: Record, customTitle?: string, markRead?: boolean, archived?: boolean, ehcliAdopted?: boolean, lastMigratedTurnId?: string): Promise { const dbRef = this._sessionDataService.openDatabase(session); const db = dbRef.object; try { @@ -5140,6 +5226,12 @@ export class CopilotAgent extends Disposable implements IAgent { if (ehcliAdopted) { work.push(db.setMetadata(AH_META_EHCLI_ADOPTED_DB_KEY, 'true')); } + // The migration boundary: the last turn that existed at adoption. Lets the + // chat editor substitute the session-wide changeset only for that turn (a + // migrated turn has no per-turn checkpoint) and never a post-adoption one. + if (lastMigratedTurnId) { + work.push(db.setMetadata(AH_META_EHCLI_LAST_TURN_DB_KEY, lastMigratedTurnId)); + } if (workingDirectory) { work.push(db.setMetadata(CopilotAgent._META_CWD, workingDirectory.toString())); } diff --git a/src/vs/platform/agentHost/node/devContainerAgentHostService.ts b/src/vs/platform/agentHost/node/devContainerAgentHostService.ts index 2d1775583627ad..6a0bd0d8b17057 100644 --- a/src/vs/platform/agentHost/node/devContainerAgentHostService.ts +++ b/src/vs/platform/agentHost/node/devContainerAgentHostService.ts @@ -30,6 +30,7 @@ import { buildAgentHostSpawnCommand, buildAgentRelayCommand, filterLiveAgentHostEndpoints, + getNewAgentHostRegistrationTimeoutMs, getRemoteCLIDataDir, ISshExec, resolveRemotePlatform, @@ -144,7 +145,7 @@ export class DevContainerAgentHostMainService extends Disposable implements IDev const serverDataFolderName = this._productService.serverDataFolderName ?? '.vscode-server-oss'; const quality = this._productService.quality || 'insider'; - const cliBin = await ensureRemoteAgentHostCliInstalled(exec, platform, { + const cliInstallation = await ensureRemoteAgentHostCliInstalled(exec, platform, { serverDataFolderName, quality, commit: this._productService.commit, @@ -152,6 +153,7 @@ export class DevContainerAgentHostMainService extends Disposable implements IDev logService: this._logService, logPrefix: LOG_PREFIX, }); + const { cliBin } = cliInstallation; const cliDataDir = getRemoteCLIDataDir(serverDataFolderName); const initial = await runAgentEndpoints(exec, cliBin, cliDataDir); const live = await filterLiveAgentHostEndpoints(exec, initial.endpoints); @@ -168,13 +170,18 @@ export class DevContainerAgentHostMainService extends Disposable implements IDev void exec(spawnCommand, { ignoreExitCode: true }).catch(error => { this._logService.warn(`${LOG_PREFIX} Agent Host spawn command failed`, error); }); + this._logService.info(`${LOG_PREFIX} Waiting for the new agent host to register...`); endpoint = await waitForNewStandaloneEndpoint( exec, cliBin, cliDataDir, initial.userDataPath, live, - { token: tokenSource.token }, + { + timeoutMs: getNewAgentHostRegistrationTimeoutMs(cliInstallation.installed), + token: tokenSource.token, + progress: elapsedMs => this._logService.info(`${LOG_PREFIX} Waiting for the new agent host to register... (${Math.floor(elapsedMs / 1000)} seconds elapsed)`), + }, ); } diff --git a/src/vs/platform/agentHost/node/remoteAgentHostCliInstaller.ts b/src/vs/platform/agentHost/node/remoteAgentHostCliInstaller.ts index e0fa1413b22558..ccf678096ece6c 100644 --- a/src/vs/platform/agentHost/node/remoteAgentHostCliInstaller.ts +++ b/src/vs/platform/agentHost/node/remoteAgentHostCliInstaller.ts @@ -24,6 +24,12 @@ export interface IRemoteAgentHostCliInstallOptions { readonly logPrefix?: string; } +/** The resolved CLI path and whether this invocation installed it. */ +export interface IRemoteAgentHostCliInstallResult { + readonly cliBin: string; + readonly installed: boolean; +} + /** * Ensure that a VS Code CLI suitable for launching an Agent Host is installed * on a remote execution target. @@ -32,7 +38,7 @@ export async function ensureRemoteAgentHostCliInstalled( exec: ISshExec, platform: { readonly os: string; readonly arch: string }, options: IRemoteAgentHostCliInstallOptions, -): Promise { +): Promise { return options.commit ? ensurePinnedCliInstalled(exec, platform, options, options.commit) : ensureLooseCliInstalled(exec, platform, options); @@ -43,7 +49,7 @@ async function ensurePinnedCliInstalled( platform: { readonly os: string; readonly arch: string }, options: IRemoteAgentHostCliInstallOptions, commit: string, -): Promise { +): Promise { const cliBin = getRemoteCLIBin(options.serverDataFolderName, options.quality, commit); const installRoot = getRemoteCLIInstallRoot(options.serverDataFolderName); const logPrefix = options.logPrefix ?? '[RemoteAgentHostCliInstaller]'; @@ -56,7 +62,7 @@ async function ensurePinnedCliInstalled( } else { options.logService.warn(`${logPrefix} Skipping CLI retention cleanup: touch exited ${touchCode}`); } - return cliBin; + return { cliBin, installed: false }; } options.reportInstalling(); @@ -78,14 +84,14 @@ async function ensurePinnedCliInstalled( } options.logService.info(`${logPrefix} Installed remote CLI at ${cliBin}`); await exec(buildCleanupOldCLIsCommand(options.serverDataFolderName, options.quality), { ignoreExitCode: true }); - return cliBin; + return { cliBin, installed: true }; } catch (error) { const message = error instanceof Error ? error.message : String(error); options.logService.warn(`${logPrefix} Could not install matching CLI for commit ${commit}: ${message}. Looking for a fallback CLI...`); const fallback = await findFallbackCli(exec, options); if (fallback) { options.logService.warn(`${logPrefix} Using fallback CLI at ${fallback} (does not match desktop commit ${commit}).`); - return fallback; + return { cliBin: fallback, installed: false }; } throw error; } @@ -95,7 +101,7 @@ async function ensureLooseCliInstalled( exec: ISshExec, platform: { readonly os: string; readonly arch: string }, options: IRemoteAgentHostCliInstallOptions, -): Promise { +): Promise { const cliBin = getRemoteCLIBin(options.serverDataFolderName, options.quality); const installRoot = getRemoteCLIInstallRoot(options.serverDataFolderName); const logPrefix = options.logPrefix ?? '[RemoteAgentHostCliInstaller]'; @@ -110,7 +116,7 @@ async function ensureLooseCliInstalled( options.logService.warn(`${logPrefix} Could not refresh the dev-build remote CLI at ${cliBin}; reusing the existing executable: update exited ${updateExitCode}`); } options.logService.info(`${logPrefix} Reusing remote CLI at ${cliBin} (dev build, latest-version refresh attempted)`); - return cliBin; + return { cliBin, installed: false }; } options.reportInstalling(); @@ -121,7 +127,7 @@ async function ensureLooseCliInstalled( `chmod +x ${cliBin}`, ].join(' && ')); options.logService.info(`${logPrefix} Installed remote CLI at ${cliBin}`); - return cliBin; + return { cliBin, installed: true }; } async function findFallbackCli(exec: ISshExec, options: IRemoteAgentHostCliInstallOptions): Promise { diff --git a/src/vs/platform/agentHost/node/shared/worktreeIsolation.ts b/src/vs/platform/agentHost/node/shared/worktreeIsolation.ts index 98a8d3b6f72ad5..dd8606fde0d439 100644 --- a/src/vs/platform/agentHost/node/shared/worktreeIsolation.ts +++ b/src/vs/platform/agentHost/node/shared/worktreeIsolation.ts @@ -986,10 +986,12 @@ export class WorktreeIsolation extends Disposable implements IAgentHostWorktreeI } /** - * Records worktree identity supplied by a predecessor for an adopted session whose - * checkout is gone, so resume recreates it exactly like a native worktree session. - * Values come from the predecessor's own record rather than probing the (missing) - * directory, which is what {@link adoptExistingWorktreeMetadata} requires. + * Records worktree identity supplied by a predecessor for an adopted session, so + * resume treats it exactly like a native worktree session. Values come from the + * predecessor's own record rather than probing the directory, which is what + * {@link adoptExistingWorktreeMetadata} requires. Used both when the checkout is + * gone (resume recreates it) and when it still exists but the predecessor recorded + * a base branch that could not otherwise be recovered without a remote (#333642). */ async recordAdoptedWorktreeMetadata(sessionUri: URI, metadata: { readonly branchName: string; readonly baseBranch: string | undefined; readonly worktreePath: URI; readonly repositoryRoot: URI }): Promise { this._logService.info(`[${this._logLabel}:${AgentSession.id(sessionUri)}] Recorded adopted worktree metadata: worktree='${metadata.worktreePath.fsPath}' branch='${metadata.branchName}' base='${metadata.baseBranch ?? '(none)'}' repo='${metadata.repositoryRoot.fsPath}'`); diff --git a/src/vs/platform/agentHost/node/sshRemoteAgentHostHelpers.ts b/src/vs/platform/agentHost/node/sshRemoteAgentHostHelpers.ts index 9adcaf51b2d7ff..25a5e7d5334c28 100644 --- a/src/vs/platform/agentHost/node/sshRemoteAgentHostHelpers.ts +++ b/src/vs/platform/agentHost/node/sshRemoteAgentHostHelpers.ts @@ -5,6 +5,7 @@ import { timeout } from '../../../base/common/async.js'; import { CancellationToken } from '../../../base/common/cancellation.js'; +import { CancellationError } from '../../../base/common/errors.js'; import { vArray, vObj, vString, vUnknown } from '../../../base/common/validation.js'; import { TelemetryConfiguration } from '../../telemetry/common/telemetry.js'; import { getAgentHostEndpointIdentityKey, IAgentHostEndpointMetadata, parseAgentHostEndpointRegistry } from '../common/agentHostEndpointRegistry.js'; @@ -487,17 +488,32 @@ export function findNewAgentHostEndpoint(before: readonly IAgentHostEndpointMeta } export interface IWaitForNewEndpointOptions { - /** Maximum number of `agent endpoints` polls before giving up. Defaults to 20. */ - readonly attempts?: number; - /** Delay between polls, in milliseconds. Defaults to 500. */ + /** + * Overall deadline for endpoint registration in milliseconds. When omitted, + * the deadline is twenty initial polling intervals (10 seconds by default). + */ + readonly timeoutMs?: number; + /** Initial delay between polls in milliseconds. Defaults to 500. */ readonly intervalMs?: number; readonly token?: CancellationToken; + /** Called periodically while endpoint registration is still pending. */ + readonly progress?: (elapsedMs: number) => void; +} + +const DEFAULT_ENDPOINT_REGISTRATION_POLL_COUNT = 20; +const MAX_ENDPOINT_REGISTRATION_POLL_INTERVAL_MS = 5_000; +const ENDPOINT_REGISTRATION_PROGRESS_INTERVAL_MS = 10_000; +const COLD_AGENT_HOST_REGISTRATION_TIMEOUT_MS = 300_000; + +/** Gets the endpoint-registration deadline for a newly installed CLI. */ +export function getNewAgentHostRegistrationTimeoutMs(installedCLI: boolean): number | undefined { + return installedCLI ? COLD_AGENT_HOST_REGISTRATION_TIMEOUT_MS : undefined; } /** * Poll `code agent endpoints` until a newly spawned standalone entry shows - * up (see {@link findNewAgentHostEndpoint}), or throw once the attempt - * budget is exhausted. The spawn command itself is fire-and-forget (its + * up (see {@link findNewAgentHostEndpoint}), or throw once the deadline + * expires. The spawn command itself is fire-and-forget (its * process is not tied to the SSH exec channel that launched it — see * {@link buildAgentHostSpawnCommand}), so this is the only way to learn * the freshly assigned TCP address/token/instanceId. @@ -510,21 +526,42 @@ export async function waitForNewStandaloneEndpoint( before: readonly IAgentHostEndpointMetadata[], options?: IWaitForNewEndpointOptions, ): Promise { - const attempts = options?.attempts ?? 20; - const intervalMs = options?.intervalMs ?? 500; - for (let attempt = 0; attempt < attempts; attempt++) { + const initialIntervalMs = options?.intervalMs ?? 500; + const timeoutMs = options?.timeoutMs ?? DEFAULT_ENDPOINT_REGISTRATION_POLL_COUNT * initialIntervalMs; + const startTime = Date.now(); + const deadline = startTime + timeoutMs; + let polls = 0; + let nextProgressReport = ENDPOINT_REGISTRATION_PROGRESS_INTERVAL_MS; + + while (true) { + if (options?.token?.isCancellationRequested) { + throw new CancellationError(); + } const { endpoints } = await runAgentEndpoints(exec, cliBin, cliDataDir, userDataPath); const found = findNewAgentHostEndpoint(before, endpoints); if (found) { return found; } - if (attempt < attempts - 1) { - if (options?.token) { - await timeout(intervalMs, options.token); - } else { - await timeout(intervalMs); - } + + polls++; + const elapsedMs = Date.now() - startTime; + if (elapsedMs >= nextProgressReport) { + options?.progress?.(elapsedMs); + nextProgressReport = (Math.floor(elapsedMs / ENDPOINT_REGISTRATION_PROGRESS_INTERVAL_MS) + 1) * ENDPOINT_REGISTRATION_PROGRESS_INTERVAL_MS; + } + if (Date.now() >= deadline) { + throw new Error(`Timed out waiting for the newly spawned agent host to register itself after ${Date.now() - startTime}ms (deadline ${timeoutMs}ms)`); + } + + const intervalMs = Math.min( + initialIntervalMs * 2 ** Math.floor((polls - 1) / 10), + MAX_ENDPOINT_REGISTRATION_POLL_INTERVAL_MS, + deadline - Date.now(), + ); + if (options?.token) { + await timeout(intervalMs, options.token); + } else { + await timeout(intervalMs); } } - throw new Error(`Timed out waiting for the newly spawned agent host to register itself (checked ${attempts} times, ~${Math.round(attempts * intervalMs / 1000)}s)`); } diff --git a/src/vs/platform/agentHost/node/sshRemoteAgentHostService.ts b/src/vs/platform/agentHost/node/sshRemoteAgentHostService.ts index e27480563d87d0..6aaa253288ad30 100644 --- a/src/vs/platform/agentHost/node/sshRemoteAgentHostService.ts +++ b/src/vs/platform/agentHost/node/sshRemoteAgentHostService.ts @@ -61,6 +61,7 @@ import { buildAgentRelayCommand, extractAgentHostWebSocketURL, filterLiveAgentHostEndpoints, + getNewAgentHostRegistrationTimeoutMs, getRemoteCLIDataDir, redactToken, resolveRemotePlatform, @@ -69,7 +70,7 @@ import { validateAgentHostTelemetryLevel, waitForNewStandaloneEndpoint, } from './sshRemoteAgentHostHelpers.js'; -import { ensureRemoteAgentHostCliInstalled } from './remoteAgentHostCliInstaller.js'; +import { ensureRemoteAgentHostCliInstalled, type IRemoteAgentHostCliInstallResult } from './remoteAgentHostCliInstaller.js'; import { parseSSHConfigHostEntries, parseSSHGOutput, stripSSHComment } from '../common/sshConfigParsing.js'; import { removeAnsiEscapeCodes } from '../../../base/common/strings.js'; @@ -939,7 +940,8 @@ export class SSHRemoteAgentHostMainService extends Disposable implements ISSHRem } this._logService.info(`${LOG_PREFIX} Remote platform: ${platform.os}-${platform.arch}`); reportProgress(localize('sshProgressInstallingCLI', "Checking remote CLI installation...")); - cliBin = await this._ensureCLIInstalled(sshClient, platform, reportProgress); + const cliInstallation = await this._ensureCLIInstalled(sshClient, platform, reportProgress); + cliBin = cliInstallation.cliBin; cliDataDir = getRemoteCLIDataDir(this._serverDataFolderName); // 3. Discover every live endpoint on the remote via the shared registry. @@ -963,7 +965,10 @@ export class SSHRemoteAgentHostMainService extends Disposable implements ISSHRem this._logService.warn(`${LOG_PREFIX} Spawn command for dedicated agent host reported an error: ${err instanceof Error ? err.message : String(err)}`); }); reportProgress(localize('sshProgressAwaitingAgent', "Waiting for the new agent host to register...")); - return waitForNewStandaloneEndpoint(exec, cliBin, cliDataDir, userDataPath, live); + return waitForNewStandaloneEndpoint(exec, cliBin, cliDataDir, userDataPath, live, { + timeoutMs: getNewAgentHostRegistrationTimeoutMs(cliInstallation.installed), + progress: elapsedMs => reportProgress(localize('sshProgressStillAwaitingAgent', "Waiting for the new agent host to register... ({0} seconds elapsed)", Math.floor(elapsedMs / 1000))), + }); }; // Deterministic dedicated (standalone) selection: reuse a live @@ -2079,9 +2084,9 @@ export class SSHRemoteAgentHostMainService extends Disposable implements ISSHRem * at `~//`. Existing CLIs self-update * against the latest release before reuse. * - * Returns the resolved CLI binary path to run. + * Returns the resolved CLI binary path and its install outcome. */ - private async _ensureCLIInstalled(client: SSHClient, platform: { os: string; arch: string }, reportProgress: (message: string) => void): Promise { + private async _ensureCLIInstalled(client: SSHClient, platform: { os: string; arch: string }, reportProgress: (message: string) => void): Promise { return ensureRemoteAgentHostCliInstalled(bindSshExec(client), platform, { serverDataFolderName: this._serverDataFolderName, quality: this._quality, diff --git a/src/vs/platform/agentHost/node/wslRemoteAgentHostHelpers.ts b/src/vs/platform/agentHost/node/wslRemoteAgentHostHelpers.ts index 575c63377cc4d8..e862b86244d9a7 100644 --- a/src/vs/platform/agentHost/node/wslRemoteAgentHostHelpers.ts +++ b/src/vs/platform/agentHost/node/wslRemoteAgentHostHelpers.ts @@ -283,7 +283,8 @@ export function composeAgentHostBootstrapScript(args: IComposeAgentHostBootstrap const cliBin = getRemoteCLIBin(args.serverDataFolderName, args.quality, args.commit); const cliDataDir = getRemoteCLIDataDir(args.serverDataFolderName); const url = buildCLIDownloadUrl(args.os, args.arch, args.quality, args.commit); - const launch = `exec ${buildAgentHostBaseCommand(cliBin, cliDataDir, telemetryLevel)}`; + const agentHostCommand = buildAgentHostBaseCommand(cliBin, cliDataDir, telemetryLevel); + const launch = buildWslAgentHostLaunch(agentHostCommand); if (args.commit) { // Pinned-install path. Mirrors SSH's _ensureCLIInstalledPinned: stage @@ -318,6 +319,16 @@ export function composeAgentHostBootstrapScript(args: IComposeAgentHostBootstrap ].join(' && '); } +/** + * Build the WSL launch command with the CLI's disconnected-host reaper. + */ +function buildWslAgentHostLaunch(command: string, idleTimeoutSec = 300): string { + if (!Number.isSafeInteger(idleTimeoutSec) || idleTimeoutSec <= 0) { + throw new Error(`Unsafe idle timeout value for shell interpolation: ${JSON.stringify(idleTimeoutSec)}`); + } + return `exec ${command} --idle-timeout ${idleTimeoutSec}`; +} + /** * Validate that a string is safe to interpolate as a `wsl.exe -d ` * argument. WSL distro names are user-creatable so they could in principle diff --git a/src/vs/platform/agentHost/node/wslRemoteAgentHostService.ts b/src/vs/platform/agentHost/node/wslRemoteAgentHostService.ts index 0f6d052bd2d3b5..fd43925e4cf137 100644 --- a/src/vs/platform/agentHost/node/wslRemoteAgentHostService.ts +++ b/src/vs/platform/agentHost/node/wslRemoteAgentHostService.ts @@ -37,8 +37,11 @@ import { const LOG_PREFIX = '[WSLRemoteAgentHost]'; -/** Max time to wait for `code agent host` inside the distro to print its `ws://` URL. */ -const AGENT_HOST_READY_TIMEOUT_MS = 60_000; +/** Max time `code agent host` may be silent before printing its `ws://` URL. */ +const AGENT_HOST_OUTPUT_IDLE_TIMEOUT_MS = 60_000; + +/** Absolute upper bound for bootstrap, including CLI and server downloads. */ +const AGENT_HOST_READY_OVERALL_TIMEOUT_MS = 10 * 60_000; /** Max time to wait for the host-side WebSocket to complete its handshake. */ const WEBSOCKET_OPEN_TIMEOUT_MS = 30_000; @@ -76,6 +79,7 @@ export class WSLRemoteAgentHostMainService extends Disposable implements IWSLRem private readonly _connections = new Map(); private readonly _distroToConnectionId = new Map(); + private readonly _pendingConnects = new Map>(); private _nativeRequire: NodeJS.Require | undefined; @@ -160,7 +164,7 @@ export class WSLRemoteAgentHostMainService extends Disposable implements IWSLRem } } - async connect(config: IWSLAgentHostConfig): Promise { + connect(config: IWSLAgentHostConfig): Promise { const distro = validateDistroName(config.distro); // Idempotent: a second `connect` for an already-live distro returns @@ -171,16 +175,34 @@ export class WSLRemoteAgentHostMainService extends Disposable implements IWSLRem if (existingId) { const existing = this._connections.get(existingId); if (existing) { - return { + return Promise.resolve({ connectionId: existing.connectionId, address: existing.address, distro: existing.distro, name: existing.name, connectionToken: existing.connectionToken, - }; + }); } } + const existingPendingConnect = this._pendingConnects.get(distro); + if (existingPendingConnect) { + return existingPendingConnect; + } + + // Reserve synchronously, before _connectUnguarded reaches its first + // await, so simultaneous callers cannot start concurrent downloads. + const pendingConnect = this._connectUnguarded(config, distro); + this._pendingConnects.set(distro, pendingConnect); + void pendingConnect.finally(() => { + if (this._pendingConnects.get(distro) === pendingConnect) { + this._pendingConnects.delete(distro); + } + }).catch(() => { /* The caller observes the original rejection. */ }); + return pendingConnect; + } + + private async _connectUnguarded(config: IWSLAgentHostConfig, distro: string): Promise { const connectionKey = `wsl:${distro}`; const reportProgress = (message: string) => { this._onDidReportConnectProgress.fire({ connectionKey, message }); @@ -209,10 +231,7 @@ export class WSLRemoteAgentHostMainService extends Disposable implements IWSLRem // agent host's stdout/stderr, which is already valid UTF-8 from a // Linux process. Keeping the bytes untouched also avoids surprising // the URL/PID regex. - const child = cp.spawn(getWslExePath(), ['-d', distro, '-e', 'bash', '-lc', script], { - windowsHide: true, - stdio: ['ignore', 'pipe', 'pipe'], - }); + const child = this._spawnAgentHost(distro, script); let url: string | undefined; let urlResolve: ((value: { url: string; token: string | undefined }) => void) | undefined; @@ -232,6 +251,34 @@ export class WSLRemoteAgentHostMainService extends Disposable implements IWSLRem } }; + let outputIdleTimeoutHandle: ReturnType | undefined; + let overallTimeoutHandle: ReturnType | undefined; + const clearReadyTimeouts = () => { + if (outputIdleTimeoutHandle !== undefined) { + clearTimeout(outputIdleTimeoutHandle); + outputIdleTimeoutHandle = undefined; + } + if (overallTimeoutHandle !== undefined) { + clearTimeout(overallTimeoutHandle); + overallTimeoutHandle = undefined; + } + }; + const rejectForTimeout = (message: string) => { + clearReadyTimeouts(); + urlReject?.(new Error(`${LOG_PREFIX} ${message}\nOutput: ${outputLines.join('\n')}`)); + }; + const armOutputIdleTimeout = () => { + if (url) { + return; + } + if (outputIdleTimeoutHandle !== undefined) { + clearTimeout(outputIdleTimeoutHandle); + } + outputIdleTimeoutHandle = setTimeout(() => { + rejectForTimeout(`Timed out waiting for agent host in '${distro}' to print its WebSocket URL: no output for ${AGENT_HOST_OUTPUT_IDLE_TIMEOUT_MS}ms.`); + }, AGENT_HOST_OUTPUT_IDLE_TIMEOUT_MS); + }; + const onStreamData = (data: Buffer) => { // `decodeWslOutput` handles both UTF-8 (the agent host's own // stdout when running with `WSL_UTF8` unset, which is what we @@ -244,6 +291,7 @@ export class WSLRemoteAgentHostMainService extends Disposable implements IWSLRem if (!line) { continue; } + armOutputIdleTimeout(); appendLine(line); this._logService.trace(`${LOG_PREFIX} [${distro}] ${redactToken(line)}`); if (!url) { @@ -259,32 +307,37 @@ export class WSLRemoteAgentHostMainService extends Disposable implements IWSLRem child.stdout?.on('data', onStreamData); child.stderr?.on('data', onStreamData); - const childExited = new Promise<{ code: number | null; signal: NodeJS.Signals | null }>((res) => { - child.once('exit', (code, signal) => res({ code, signal })); - }); - - // Race the URL parse against the child dying and the global timeout. + // Race the URL parse against the child dying, output going idle, and + // an overall ceiling. Bootstrap downloads regularly report progress, + // so only a period of silence indicates that it has become stuck. // `outputLines` is already redacted in `appendLine` — no extra wrap needed. - const readyTimeoutHandle = setTimeout(() => { - urlReject?.(new Error(`${LOG_PREFIX} Timed out waiting for agent host in '${distro}' to print its WebSocket URL after ${AGENT_HOST_READY_TIMEOUT_MS}ms.\nOutput: ${outputLines.join('\n')}`)); - }, AGENT_HOST_READY_TIMEOUT_MS); + armOutputIdleTimeout(); + overallTimeoutHandle = setTimeout(() => { + rejectForTimeout(`Timed out waiting for agent host in '${distro}' to print its WebSocket URL: exceeded the overall ${AGENT_HOST_READY_OVERALL_TIMEOUT_MS}ms bootstrap ceiling.`); + }, AGENT_HOST_READY_OVERALL_TIMEOUT_MS); - const earlyExitGuard = childExited.then(({ code, signal }) => { + child.once('exit', (code, signal) => { if (!url) { + clearReadyTimeouts(); urlReject?.(new Error(`${LOG_PREFIX} Agent host in '${distro}' exited (code=${code}, signal=${signal}) before printing its WebSocket URL.\nOutput: ${outputLines.join('\n')}`)); } }); + child.once('error', err => { + if (!url) { + clearReadyTimeouts(); + urlReject?.(new Error(`${LOG_PREFIX} Failed to start agent host in '${distro}': ${err.message}\nOutput: ${outputLines.join('\n')}`)); + } + }); let resolvedUrl: { url: string; token: string | undefined }; try { resolvedUrl = await urlPromise; } catch (err) { - clearTimeout(readyTimeoutHandle); + clearReadyTimeouts(); this._killChild(child); - await earlyExitGuard.catch(() => { /* already surfaced */ }); throw err; } - clearTimeout(readyTimeoutHandle); + clearReadyTimeouts(); reportProgress(localize('wslProgressConnecting', "Connecting to agent host in {0}...", distro)); let ws: WebSocket; @@ -349,12 +402,15 @@ export class WSLRemoteAgentHostMainService extends Disposable implements IWSLRem } } - async reconnect(distro: string, name: string, remoteAgentHostCommand?: string): Promise { + async reconnect(distro: string, name: string, remoteAgentHostCommand?: string, userInitiated?: boolean): Promise { const existingId = this._distroToConnectionId.get(distro); if (existingId) { this._closeConnection(existingId); } - return this.connect({ distro, name, remoteAgentHostCommand }); + // A pending connection is already a fresh bootstrap. Joining it avoids + // starting a competing downloader; callers that reconnect after it + // fails receive that failure and a subsequent reconnect starts anew. + return this.connect({ distro, name, remoteAgentHostCommand, userInitiated }); } async relaySend(connectionId: string, message: string): Promise { @@ -392,12 +448,15 @@ export class WSLRemoteAgentHostMainService extends Disposable implements IWSLRem if (child.exitCode !== null || child.signalCode !== null) { return; } + // A detached distro-side host relies on the bootstrap's --idle-timeout to exit. try { child.kill(); } catch { /* ignore */ } // Escalate to SIGKILL if the process is still alive after 2s. The // `unref` cast avoids the dom/node `setTimeout` typing collision in - // strict mode — we only care that escalation never blocks process exit. + // strict mode — we only care that escalation never blocks process exit, + // so it is optional: outside Node (the unit-test renderer) there is no + // `unref` and keeping the timer referenced is harmless. const escalate = setTimeout(() => { if (child.exitCode === null && child.signalCode === null) { try { @@ -405,11 +464,18 @@ export class WSLRemoteAgentHostMainService extends Disposable implements IWSLRem } catch { /* ignore */ } } }, 2_000) as unknown as NodeJS.Timeout; - escalate.unref(); + escalate.unref?.(); child.once('exit', () => clearTimeout(escalate)); } - private async _resolvePlatform(distro: string): Promise<{ os: string; arch: string }> { + protected _spawnAgentHost(distro: string, script: string): cp.ChildProcess { + return cp.spawn(getWslExePath(), ['-d', distro, '-e', 'bash', '-lc', script], { + windowsHide: true, + stdio: ['ignore', 'pipe', 'pipe'], + }); + } + + protected async _resolvePlatform(distro: string): Promise<{ os: string; arch: string }> { const result = await runWslCommand(['-e', 'uname', '-s', '-m'], { distro, timeout: 10_000 }); if (result.exitCode !== 0) { throw new Error(`${LOG_PREFIX} Failed to detect platform in '${distro}' (exit ${result.exitCode}): ${result.stderr.trim() || result.stdout.trim()}`); @@ -425,7 +491,7 @@ export class WSLRemoteAgentHostMainService extends Disposable implements IWSLRem return resolved; } - private async _openWebSocket(url: string): Promise { + protected async _openWebSocket(url: string): Promise { const nativeRequire = await this._getNativeRequire(); const WS = nativeRequire('ws') as typeof WebSocket; const deadline = Date.now() + WEBSOCKET_OPEN_TIMEOUT_MS; diff --git a/src/vs/platform/agentHost/test/electron-browser/remoteAgentHostService.test.ts b/src/vs/platform/agentHost/test/electron-browser/remoteAgentHostService.test.ts index 03c54135b4cd59..5d87f7057654e2 100644 --- a/src/vs/platform/agentHost/test/electron-browser/remoteAgentHostService.test.ts +++ b/src/vs/platform/agentHost/test/electron-browser/remoteAgentHostService.test.ts @@ -5,7 +5,8 @@ import assert from 'assert'; import { Emitter, Event } from '../../../../base/common/event.js'; -import { Disposable, DisposableStore, toDisposable } from '../../../../base/common/lifecycle.js'; +import { Disposable, DisposableStore, IDisposable, toDisposable } from '../../../../base/common/lifecycle.js'; +import { IObservable, observableValue } from '../../../../base/common/observable.js'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js'; import { ILogService, NullLogService } from '../../../log/common/log.js'; import { IEnvironmentService } from '../../../environment/common/environment.js'; @@ -15,14 +16,23 @@ import { IConfigurationService, type IConfigurationChangeEvent } from '../../../ import { IInstantiationService } from '../../../instantiation/common/instantiation.js'; import { ILabelService, type ResourceLabelFormatter } from '../../../label/common/label.js'; import { AgentsWindowRemoteAgentHostService, RemoteAgentHostService } from '../../browser/remoteAgentHostServiceImpl.js'; -import type { IAgentHostProtocolClientOptions } from '../../browser/agentHostProtocolClient.js'; -import { addSSHRemoteAgentHostEntry, addWebSocketRemoteAgentHostEntry, getEntryTypeConfig, parseRemoteAgentHostInput, removeWebSocketRemoteAgentHostEntry, RemoteAgentHostConnectionStatus, RemoteAgentHostEntryType, RemoteAgentHostsEnabledSettingId, RemoteAgentHostsSettingId, type IRawRemoteAgentHostEntry, type IRemoteAgentHostEntry } from '../../common/remoteAgentHostService.js'; +import { InitialAuthenticationError, type IAgentHostProtocolClientOptions } from '../../browser/agentHostProtocolClient.js'; +import { addSSHRemoteAgentHostEntry, addWebSocketRemoteAgentHostEntry, getEntryAddress, getEntryTypeConfig, parseRemoteAgentHostInput, removeWebSocketRemoteAgentHostEntry, RemoteAgentHostAutoConnectSettingId, RemoteAgentHostConnectionStatus, RemoteAgentHostEntryType, RemoteAgentHostsEnabledSettingId, RemoteAgentHostsSettingId, type IRawRemoteAgentHostEntry, type IRemoteAgentHostConnectionFactory, type IRemoteAgentHostCreatedConnection, type IRemoteAgentHostEntry, type IRemoteAgentHostProtocolClient } from '../../common/remoteAgentHostService.js'; import { AGENT_HOST_SCHEME, agentHostAuthority } from '../../common/agentHostUri.js'; import { DeferredPromise } from '../../../../base/common/async.js'; import { InMemoryStorageService, IStorageService, StorageScope, StorageTarget } from '../../../storage/common/storage.js'; import type { StorageValue } from '../../../../base/parts/storage/common/storage.js'; import type { Implementation } from '../../common/state/protocol/common/commands.js'; import { agentsWindowAgentHostClientInfo, editorWindowAgentHostClientInfo } from '../../common/agentHostClientInfo.js'; +import { PROTOCOL_VERSION } from '../../common/state/protocol/version/registry.js'; +import { computeReconnectDelay } from '../../common/reconnectPolicy.js'; + +interface IRemoteAgentHostServiceTestAccess { + readonly _reconnectAttempts: Map; + readonly _reconnectTimeouts: ReadonlyMap>; + _scheduleReconnect(address: string, connectionToken?: string): void; + _cancelReconnect(address: string): void; +} // ---- Mock transport --------------------------------------------------------- @@ -77,6 +87,47 @@ class MockProtocolClient extends Disposable { } } +class TestConnectionFactory extends Disposable implements IRemoteAgentHostConnectionFactory { + readonly entries: IObservable; + + private readonly _entries = observableValue(this, []); + private readonly _createdConnections = new Map(); + private readonly _onDidCreateConnection = this._register(new Emitter()); + readonly onDidCreateConnection = this._onDidCreateConnection.event; + createdConnectionCount = 0; + + constructor(readonly kind: RemoteAgentHostEntryType) { + super(); + this.entries = this._entries; + } + + stage(entry: IRemoteAgentHostEntry, connection: MockProtocolClient, transportDisposable?: IDisposable, reconnectTransfersTransportOwnership = false): void { + const address = getEntryAddress(entry); + const createdConnections = this._createdConnections.get(address) ?? []; + createdConnections.push({ + connection: connection as unknown as IRemoteAgentHostProtocolClient, + transportDisposable, + reconnectTransfersTransportOwnership, + }); + this._createdConnections.set(address, createdConnections); + this._entries.set([...this._entries.get(), entry], undefined); + } + + createConnection(entry: IRemoteAgentHostEntry): Promise { + if (entry.connection.type !== this.kind) { + return Promise.reject(new Error(`Test factory cannot create a ${entry.connection.type} connection.`)); + } + const address = getEntryAddress(entry); + const connection = this._createdConnections.get(address)?.shift(); + if (!connection) { + return Promise.reject(new Error(`No test connection staged for ${address}.`)); + } + this.createdConnectionCount++; + this._onDidCreateConnection.fire(); + return Promise.resolve(connection); + } +} + // ---- Test configuration service --------------------------------------------- class TestConfigurationService { @@ -85,12 +136,16 @@ class TestConfigurationService { private _entries: IRawRemoteAgentHostEntry[] = []; private _enabled = true; + private _autoConnect = true; updateValueCalls = 0; getValue(key?: string): unknown { if (key === RemoteAgentHostsEnabledSettingId) { return this._enabled; } + if (key === RemoteAgentHostAutoConnectSettingId) { + return this._autoConnect; + } return this._entries; } @@ -635,9 +690,8 @@ suite('RemoteAgentHostService', () => { assert.strictEqual(service.connections.length, 0); }); - suite('addManagedConnection', () => { + suite('factory connections', () => { - // Build a transport disposable that records when it ran. function makeTransportDisposable(): { disposable: { dispose(): void }; disposed: () => boolean } { let disposed = false; return { @@ -646,42 +700,129 @@ suite('RemoteAgentHostService', () => { }; } - // Inject a managed connection (mimicking the SSH/tunnel renderer flow). - async function addManaged(name: string, address: string, transport?: { dispose(): void }) { - const mockClient = disposables.add(new MockProtocolClient(`ws://${address}`)); - return service.addManagedConnection( - { name, connection: { type: RemoteAgentHostEntryType.WebSocket, address } }, - mockClient as unknown as Parameters[1], - transport, - ); + function createFactory(kind = RemoteAgentHostEntryType.CloudSandbox): TestConnectionFactory { + const factory = disposables.add(new TestConnectionFactory(kind)); + disposables.add(service.registerConnectionFactory(factory)); + return factory; } - test('keeps incompatible managed connection addressable for server upgrade', async () => { - const mockClient = disposables.add(new MockProtocolClient('ssh:remote.example')); - await service.addManagedConnection( - { - name: 'SSH Host', - connection: { - type: RemoteAgentHostEntryType.SSH, - address: 'ssh:remote.example', - sshConfigHost: 'remote', - hostName: 'remote.example', - }, - }, - mockClient as unknown as Parameters[1], - undefined, - RemoteAgentHostConnectionStatus.incompatible('Unsupported protocol version', ['0.3.0'], ['^0.2.0'], '_vscodeUpgrade'), - ); + function cloudSandboxEntry(name: string, address: string): IRemoteAgentHostEntry { + return { + name, + connection: { type: RemoteAgentHostEntryType.CloudSandbox, address, environmentId: 'env_test' }, + }; + } + + async function waitForFactoryConnection(factory: TestConnectionFactory, count: number): Promise { + while (factory.createdConnectionCount < count) { + await Event.toPromise(factory.onDidCreateConnection); + } + } - const upgradeResult = await service.triggerServerUpgrade('ssh:remote.example', '_vscodeUpgrade'); + async function reconnectStagedConnection(factory: TestConnectionFactory, entry: IRemoteAgentHostEntry, client: MockProtocolClient, transportDisposable?: IDisposable, reconnectTransfersTransportOwnership = false): Promise { + // Capture the target before staging: `reconnect` dials asynchronously and + // may already have created the connection by the time we start waiting. + const expectedConnectionCount = factory.createdConnectionCount + 1; + factory.stage(entry, client, transportDisposable, reconnectTransfersTransportOwnership); + service.reconnect(getEntryAddress(entry)); + const wait = service.waitForConnection(getEntryAddress(entry)); + await waitForFactoryConnection(factory, expectedConnectionCount); + client.connectDeferred.complete(); + await wait; + } + + test('preserves automatic reconnect attempts while resetting them for a user reconnect', async () => { + const factory = createFactory(); + const entry = cloudSandboxEntry('Cloud Sandbox', 'cloud:reconnect-budget'); + const automaticClient = new MockProtocolClient('cloud:reconnect-budget'); + const address = getEntryAddress(entry); + const internals = service as unknown as IRemoteAgentHostServiceTestAccess; + const reconnectPolicy = getEntryTypeConfig(RemoteAgentHostEntryType.CloudSandbox).reconnect; + internals._reconnectAttempts.set(address, 3); + + factory.stage(entry, automaticClient); + service.reconnect(address, false); + // An automatic retry never spends the budget it depends on, whether + // it starts the dial or joins one already in flight. + service.reconnect(address, false); + assert.deepStrictEqual({ + automaticAttempts: internals._reconnectAttempts.get(address), + automaticCreates: factory.createdConnectionCount, + }, { + automaticAttempts: 3, + automaticCreates: 1, + }); + + service.reconnect(address, true); + + // The user request joins the in-flight dial rather than starting a + // second one, but still restores the budget so a later failure is + // retried instead of being reported as exhausted. + assert.deepStrictEqual({ + automaticAttempts: internals._reconnectAttempts.get(address), + pendingReconnectCreates: factory.createdConnectionCount, + }, { + automaticAttempts: undefined, + pendingReconnectCreates: 1, + }); + + const automaticWait = service.waitForConnection(address); + await waitForFactoryConnection(factory, 1); + automaticClient.connectDeferred.complete(); + await automaticWait; + + const automaticDelays: number[] = []; + for (let attempt = 1; attempt <= reconnectPolicy.maxAttempts; attempt++) { + internals._scheduleReconnect(address); + automaticDelays.push(computeReconnectDelay(reconnectPolicy, attempt)); + internals._cancelReconnect(address); + } + internals._scheduleReconnect(address); + assert.deepStrictEqual({ + delaysForSuccessiveAutomaticFailures: automaticDelays, + attemptsAtLimit: internals._reconnectAttempts.get(address), + hasRetryAtLimit: internals._reconnectTimeouts.has(address), + }, { + delaysForSuccessiveAutomaticFailures: [1000, 2000, 4000, 8000, 16000, 30000, 30000, 30000, 30000, 30000], + attemptsAtLimit: reconnectPolicy.maxAttempts, + hasRetryAtLimit: false, + }); + + const userClient = new MockProtocolClient('cloud:reconnect-budget'); + internals._reconnectAttempts.set(address, 3); + factory.stage(entry, userClient); + service.reconnect(address, true); + + assert.strictEqual(internals._reconnectAttempts.get(address), undefined); + + const userWait = service.waitForConnection(address); + await waitForFactoryConnection(factory, 2); + userClient.connectDeferred.complete(); + await userWait; + }); + + test('keeps an incompatible factory connection addressable for server upgrade', async () => { + const factory = createFactory(); + const entry = cloudSandboxEntry('Cloud Sandbox', 'cloud:incompatible'); + const client = new MockProtocolClient('cloud:incompatible'); + factory.stage(entry, client); + service.reconnect(getEntryAddress(entry)); + const wait = service.waitForConnection(getEntryAddress(entry)); + await waitForFactoryConnection(factory, 1); + const changed = Event.toPromise(service.onDidChangeConnections); + client.connectDeferred.error(new InitialAuthenticationError(new Error('Unsupported protocol version'))); + await changed; + await assert.rejects(() => wait, /Initial authentication failed/); + + const upgradeResult = await service.triggerServerUpgrade('cloud:incompatible', '_vscodeUpgrade'); assert.deepStrictEqual({ status: service.connections[0].status, - connectedConnection: service.getConnection('ssh:remote.example'), - upgradeCalls: mockClient.triggerVscodeUpgradeCalls, + connectedConnection: service.getConnection('cloud:incompatible'), + upgradeCalls: client.triggerVscodeUpgradeCalls, upgradeResult, }, { - status: RemoteAgentHostConnectionStatus.incompatible('Unsupported protocol version', ['0.3.0'], ['^0.2.0'], '_vscodeUpgrade'), + status: RemoteAgentHostConnectionStatus.incompatible('Initial authentication failed: Unsupported protocol version', [PROTOCOL_VERSION]), connectedConnection: undefined, upgradeCalls: ['_vscodeUpgrade'], upgradeResult: { ok: true, upgradeStarted: true }, @@ -689,70 +830,59 @@ suite('RemoteAgentHostService', () => { }); test('disposes transportDisposable when entry is removed via removeRemoteAgentHost', async () => { + const factory = createFactory(); const t = makeTransportDisposable(); - await addManaged('Managed', 'managed:1234', t.disposable); + await reconnectStagedConnection(factory, cloudSandboxEntry('Cloud Sandbox', 'cloud:remove'), new MockProtocolClient('cloud:remove'), t.disposable); assert.strictEqual(t.disposed(), false); - await service.removeRemoteAgentHost('ws://managed:1234'); + await service.removeRemoteAgentHost('cloud:remove'); assert.strictEqual(t.disposed(), true, 'transport disposable runs when entry is removed'); - assert.strictEqual(service.getConnection('ws://managed:1234'), undefined); - }); - - test('throws when disabled', async () => { - configService.setEnabled(false); - - await assert.rejects( - () => addManaged('Managed', 'managed:1234'), - /not enabled/, - ); + assert.strictEqual(service.getConnection('cloud:remove'), undefined); }); - test('does NOT dispose previous transportDisposable when entry is replaced', async () => { - // When the entry is replaced (e.g. on reconnect to the same address), - // the new entry takes ownership of the same underlying connectionId. - // Running the old transportDisposable would call disconnect() on the - // shared-process tunnel keyed by that connectionId and immediately - // tear down the brand-new connection. The new transportDisposable - // inherits responsibility for the underlying tunnel. + test('does not dispose a previous transport when a replacement takes ownership', async () => { + const factory = createFactory(); + const entry = cloudSandboxEntry('Cloud Sandbox', 'cloud:replacement'); const t1 = makeTransportDisposable(); - await addManaged('Managed', 'managed:1234', t1.disposable); + await reconnectStagedConnection(factory, entry, new MockProtocolClient('cloud:replacement'), t1.disposable, true); const t2 = makeTransportDisposable(); - await addManaged('Managed', 'managed:1234', t2.disposable); + await reconnectStagedConnection(factory, entry, new MockProtocolClient('cloud:replacement'), t2.disposable, true); assert.strictEqual(t1.disposed(), false, 'previous transport disposable is not run on replacement'); assert.strictEqual(t2.disposed(), false, 'new transport disposable is still alive'); - await service.removeRemoteAgentHost('ws://managed:1234'); + await service.removeRemoteAgentHost('cloud:replacement'); assert.strictEqual(t2.disposed(), true, 'new transport disposable runs on full removal'); }); test('disposes transportDisposable when service itself is disposed', async () => { + const factory = createFactory(); const t = makeTransportDisposable(); - await addManaged('Managed', 'managed:1234', t.disposable); + await reconnectStagedConnection(factory, cloudSandboxEntry('Cloud Sandbox', 'cloud:dispose'), new MockProtocolClient('cloud:dispose'), t.disposable); service.dispose(); assert.strictEqual(t.disposed(), true, 'transport disposable runs when service is disposed'); }); - test('does not persist runtime managed connections or their removal', async () => { + test('does not persist runtime factory connections or their removal', async () => { + const cloudSandboxFactory = createFactory(RemoteAgentHostEntryType.CloudSandbox); + const devContainerFactory = createFactory(RemoteAgentHostEntryType.DevContainer); const entries: IRemoteAgentHostEntry[] = [ - { name: 'Tunnel', connection: { type: RemoteAgentHostEntryType.Tunnel, tunnelId: 'runtime-tunnel', clusterId: 'cluster' } }, - { name: 'WSL', connection: { type: RemoteAgentHostEntryType.WSL, address: 'wsl:runtime', distro: 'runtime' } }, { name: 'Cloud Sandbox', connection: { type: RemoteAgentHostEntryType.CloudSandbox, address: 'cloud:runtime', environmentId: 'env_runtime' } }, { name: 'Dev Container', connection: { type: RemoteAgentHostEntryType.DevContainer, address: 'devcontainer:runtime', hostPath: '/workspace' } }, ]; - const addresses = ['tunnel:runtime-tunnel', 'wsl:runtime', 'cloud:runtime', 'devcontainer:runtime']; + const factories = [cloudSandboxFactory, devContainerFactory]; for (let index = 0; index < entries.length; index++) { - const client = disposables.add(new MockProtocolClient(addresses[index])); - await service.addManagedConnection(entries[index], client as unknown as Parameters[1]); + const address = getEntryAddress(entries[index]); + await reconnectStagedConnection(factories[index], entries[index], new MockProtocolClient(address)); } - for (const address of addresses) { - await service.removeRemoteAgentHost(address); + for (const entry of entries) { + await service.removeRemoteAgentHost(getEntryAddress(entry)); } assert.deepStrictEqual({ @@ -766,16 +896,15 @@ suite('RemoteAgentHostService', () => { }); }); - test('keeps a registered tunnel connected when WebSocket settings change', async () => { - const tunnel = disposables.add(new MockProtocolClient('tunnel:live')); - await service.addManagedConnection( - { name: 'Tunnel', connection: { type: RemoteAgentHostEntryType.Tunnel, tunnelId: 'live', clusterId: 'cluster' } }, - tunnel as unknown as Parameters[1], - ); + test('keeps a staged on-demand connection connected when WebSocket settings change', async () => { + const factory = createFactory(); + const entry = cloudSandboxEntry('Cloud Sandbox', 'cloud:live'); + const client = new MockProtocolClient('cloud:live'); + await reconnectStagedConnection(factory, entry, client); configService.setEntries([{ name: 'WebSocket', connection: { type: RemoteAgentHostEntryType.WebSocket, address: 'ws://host:8080' } }]); - assert.strictEqual(service.getConnection('tunnel:live'), tunnel); + assert.strictEqual(service.getConnection('cloud:live'), client); }); test('does not surface storage-only SSH entries without an SSH factory', async () => { @@ -805,15 +934,15 @@ suite('RemoteAgentHostService', () => { }); test('keeps runtime connection names across reconciliation', async () => { - const tunnel: IRemoteAgentHostEntry = { name: 'My Tunnel', connection: { type: RemoteAgentHostEntryType.Tunnel, tunnelId: 'tunnel', clusterId: 'cluster' } }; - const client = disposables.add(new MockProtocolClient('tunnel:tunnel')); - await service.addManagedConnection(tunnel, client as unknown as Parameters[1]); + const factory = createFactory(); + const cloudSandbox = cloudSandboxEntry('My Cloud Sandbox', 'cloud:name'); + await reconnectStagedConnection(factory, cloudSandbox, new MockProtocolClient('cloud:name')); configService.setEntries([{ name: 'WebSocket', connection: { type: RemoteAgentHostEntryType.WebSocket, address: 'host1:8080' } }]); assert.deepStrictEqual( - service.connections.find(connection => connection.address === 'tunnel:tunnel')?.name, - 'My Tunnel'); + service.connections.find(connection => connection.address === 'cloud:name')?.name, + 'My Cloud Sandbox'); }); }); diff --git a/src/vs/platform/agentHost/test/electron-browser/wslRemoteAgentHostService.test.ts b/src/vs/platform/agentHost/test/electron-browser/wslRemoteAgentHostService.test.ts new file mode 100644 index 00000000000000..9c9fcbf67f0511 --- /dev/null +++ b/src/vs/platform/agentHost/test/electron-browser/wslRemoteAgentHostService.test.ts @@ -0,0 +1,95 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import assert from 'assert'; +import { Event } from '../../../../base/common/event.js'; +import { DisposableStore, toDisposable } from '../../../../base/common/lifecycle.js'; +import type { IChannel } from '../../../../base/parts/ipc/common/ipc.js'; +import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js'; +import { IConfigurationService } from '../../../configuration/common/configuration.js'; +import { TestInstantiationService } from '../../../instantiation/test/common/instantiationServiceMock.js'; +import { ISharedProcessService } from '../../../ipc/electron-browser/services.js'; +import { ILogService, NullLogService } from '../../../log/common/log.js'; +import { InMemoryStorageService, IStorageService } from '../../../storage/common/storage.js'; +import { IRemoteAgentHostService, type IRemoteAgentHostConnectionFactory } from '../../common/remoteAgentHostService.js'; +import { IWSLRelayClientFactory, WSLRemoteAgentHostService } from '../../electron-browser/wslRemoteAgentHostServiceImpl.js'; + +class MockWSLMainService { + readonly onDidCloseConnection = Event.None; + readonly onDidReportConnectProgress = Event.None; +} + +class MockRemoteAgentHostService { + readonly reconnectCalls: Array<{ readonly address: string; readonly userInitiated: boolean }> = []; + + registerConnectionFactory(_factory: IRemoteAgentHostConnectionFactory) { + return toDisposable(() => undefined); + } + + reconnect(address: string, userInitiated = true): void { + this.reconnectCalls.push({ address, userInitiated }); + } + + async waitForConnection(_address: string): Promise { + throw new Error('Connection was not established in this forwarding test.'); + } +} + +function asChannel(target: object): IChannel { + return { + call: async (method: string, args?: unknown): Promise => { + const fn = (target as Record)[method]; + if (typeof fn !== 'function') { + throw new Error(`MockChannel: no method ${method}`); + } + return (fn as (...a: unknown[]) => Promise).apply(target, (args as unknown[]) ?? []); + }, + listen: (event: string): Event => { + const value = (target as Record)[event]; + if (typeof value !== 'function') { + throw new Error(`MockChannel: no event ${event}`); + } + return value as Event; + }, + }; +} + +suite('WSLRemoteAgentHostService (renderer)', () => { + const disposables = new DisposableStore(); + let remoteAgentHostService: MockRemoteAgentHostService; + let service: WSLRemoteAgentHostService; + + setup(() => { + const mainService = new MockWSLMainService(); + remoteAgentHostService = new MockRemoteAgentHostService(); + const instantiationService = disposables.add(new TestInstantiationService()); + instantiationService.stub(ILogService, new NullLogService()); + instantiationService.stub(IConfigurationService, { + getValue: () => true, + } as Partial); + instantiationService.stub(ISharedProcessService, { + getChannel: () => asChannel(mainService), + } as Partial); + instantiationService.stub(IStorageService, disposables.add(new InMemoryStorageService())); + instantiationService.stub(IRemoteAgentHostService, remoteAgentHostService as Partial); + instantiationService.stub(IWSLRelayClientFactory, { + createClient: () => { throw new Error('Unexpected relay client creation.'); }, + } as Partial); + service = disposables.add(instantiationService.createInstance(WSLRemoteAgentHostService)); + }); + + teardown(() => disposables.clear()); + ensureNoDisposablesAreLeakedInTestSuite(); + + test('forwards whether reconnect was user-initiated', async () => { + await assert.rejects(() => service.reconnect('Ubuntu', 'Ubuntu'), /not established/); + await assert.rejects(() => service.reconnect('Ubuntu', 'Ubuntu', false), /not established/); + + assert.deepStrictEqual(remoteAgentHostService.reconnectCalls, [ + { address: 'wsl:Ubuntu', userInitiated: true }, + { address: 'wsl:Ubuntu', userInitiated: false }, + ]); + }); +}); diff --git a/src/vs/platform/agentHost/test/node/agentService.test.ts b/src/vs/platform/agentHost/test/node/agentService.test.ts index cbad6f85ec65f7..6e5c73db309cf0 100644 --- a/src/vs/platform/agentHost/test/node/agentService.test.ts +++ b/src/vs/platform/agentHost/test/node/agentService.test.ts @@ -48,7 +48,7 @@ import { readAgentMessageDelegationMeta } from '../../common/meta/agentMessageDe import { IProductService } from '../../../product/common/productService.js'; import { AgentService } from '../../node/agentService.js'; import { AgentHostDatabase, IAgentHostDatabase, IAgentHostDatabaseRegisterOptions, IAgentHostDatabaseSession, IAgentHostDatabaseSessionOptions } from '../../node/agentHostDatabase.js'; -import { AgentSessionRegistry, type IRegisteredSession } from '../../node/agentSessionRegistry.js'; +import { AgentSessionRegistry } from '../../node/agentSessionRegistry.js'; import { AgentHostManagementService } from '../../node/agentHostManagementService.js'; import { AGENT_HOST_TITLE_SOURCE_AUTO, SESSION_CUSTOM_TITLE_SOURCE_KEY } from '../../node/shared/persistSessionMetadata.js'; import { MockAgent, ScriptedMockAgent } from './mockAgent.js'; @@ -3186,7 +3186,24 @@ suite('AgentService (node dispatcher)', () => { } } - function createExternalSessionService(sessionDataService = createSessionDataService(), orchestratorDatabase?: IAgentHostDatabase, copilotApiService?: ICopilotApiService): AgentService { + class ControlledDiscoveryAgent extends TimedExternalAgent { + discoveryStarts = 0; + + override async listExternalChats(): Promise { + return []; + } + + async startChatDiscovery(): Promise { + this.discoveryStarts++; + this.fireDiscoveredChats([...this.catalog.values()].map(entry => discoveredChat(entry.session, true, entry.modifiedTime))); + } + + async ensureChatAdopted(): Promise { + return { adopted: false, eligible: false }; + } + } + + function createExternalSessionService(sessionDataService = createSessionDataService(), orchestratorDatabase?: IAgentHostDatabase, copilotApiService?: ICopilotApiService, storageResource?: URI): AgentService { return disposables.add(createTestAgentService( new NullLogService(), fileService, @@ -3200,11 +3217,176 @@ suite('AgentService (node dispatcher)', () => { undefined, [], undefined, - undefined, + storageResource, orchestratorDatabase, )); } + testWithExternalSessionClock('external discovery waits for startup settlement after the setting enables it', async () => { + const database = new TransientRegistryWriteDatabase(); + await database.markProviderBackfilled('copilot'); + const svc = createExternalSessionService(createSessionDataService(), database); + const agent = disposables.add(new ControlledDiscoveryAgent('copilot')); + const external = agent.addSession('setting-enabled-discovery', Date.now()); + registerTestAgentProvider(svc, agent); + + const startsWhileDisabled = agent.discoveryStarts; + setExternalSessionsMode(svc, AgentHostExternalSessionsMode.Last30Days, 1); + const startsBeforeStartupComplete = agent.discoveryStarts; + svc.markStartupComplete(); + const startsAfterStartupComplete = agent.discoveryStarts; + const initiallyVisible = await svc.listSessions(); + await svc.whenDeferredWorkSettled(); + for (let attempt = 0; attempt < 50 && (await svc.getRegisteredSessions()).length === 0; attempt++) { + await timeout(0); + } + await waitForSessionListReconciliation(svc); + + assert.deepStrictEqual({ + initiallyVisible, + startsWhileDisabled, + startsBeforeStartupComplete, + startsAfterStartupComplete, + startsAfterStartupSettled: agent.discoveryStarts, + visibleAfterStartupSettled: (await svc.listSessions()).map(session => session.session.toString()), + }, { + initiallyVisible: [], + startsWhileDisabled: 0, + startsBeforeStartupComplete: 0, + startsAfterStartupComplete: 0, + startsAfterStartupSettled: 1, + visibleAfterStartupSettled: [external.toString()], + }); + }); + + testWithExternalSessionClock('enabling external sessions after startup settlement starts discovery', async () => { + const database = new TransientRegistryWriteDatabase(); + await database.markProviderBackfilled('copilot'); + const svc = createExternalSessionService(createSessionDataService(), database); + const agent = disposables.add(new ControlledDiscoveryAgent('copilot')); + const external = agent.addSession('post-startup-enablement', Date.now()); + registerTestAgentProvider(svc, agent); + svc.markStartupComplete(); + await svc.listSessions(); + await svc.whenDeferredWorkSettled(); + const startsBeforeEnablement = agent.discoveryStarts; + + setExternalSessionsMode(svc, AgentHostExternalSessionsMode.Last30Days, 1); + for (let attempt = 0; attempt < 50 && (await svc.getRegisteredSessions()).length === 0; attempt++) { + await timeout(0); + } + await waitForSessionListReconciliation(svc); + + assert.deepStrictEqual({ + startsBeforeEnablement, + startsAfterEnablement: agent.discoveryStarts, + visible: (await svc.listSessions()).map(session => session.session.toString()), + }, { + startsBeforeEnablement: 0, + startsAfterEnablement: 1, + visible: [external.toString()], + }); + }); + + testWithExternalSessionClock('a provider registered after startup starts external discovery immediately', async () => { + const database = new TransientRegistryWriteDatabase(); + await database.markProviderBackfilled('copilot'); + const svc = createExternalSessionService(createSessionDataService(), database); + setExternalSessionsMode(svc, AgentHostExternalSessionsMode.Last30Days, 1); + svc.markStartupComplete(); + await svc.listSessions(); + await svc.whenDeferredWorkSettled(); + const agent = disposables.add(new ControlledDiscoveryAgent('copilot')); + const external = agent.addSession('late-provider-discovery', Date.now()); + + registerTestAgentProvider(svc, agent); + await svc.whenDeferredWorkSettled(); + for (let attempt = 0; attempt < 50 && (await svc.getRegisteredSessions()).length === 0; attempt++) { + await timeout(0); + } + await waitForSessionListReconciliation(svc); + + assert.deepStrictEqual({ + discoveryStarts: agent.discoveryStarts, + visible: (await svc.listSessions()).map(session => session.session.toString()), + }, { + discoveryStarts: 1, + visible: [external.toString()], + }); + }); + + testWithExternalSessionClock('legacy migration can start discovery while external sessions are hidden', async () => { + const svc = createExternalSessionService(); + const agent = disposables.add(new ControlledDiscoveryAgent('copilot')); + const external = agent.addSession('migration-triggered-discovery', Date.now()); + registerTestAgentProvider(svc, agent); + + await svc.listSessions(); + for (let attempt = 0; attempt < 50 && (await svc.getRegisteredSessions()).length === 0; attempt++) { + await timeout(0); + } + + assert.deepStrictEqual({ + discoveryStarts: agent.discoveryStarts, + registered: (await svc.getRegisteredSessions()).map(session => session.toString()), + visible: await svc.listSessions(), + }, { + discoveryStarts: 1, + registered: [external.toString()], + visible: [], + }); + }); + + testWithExternalSessionClock('enabled legacy migration starts discovery when the provider registry is already backfilled', async () => { + const database = new TransientRegistryWriteDatabase(); + await database.markProviderBackfilled('copilot'); + const svc = createExternalSessionService(createSessionDataService(), database); + getConfigurationService(svc).updateRootConfig({ [AgentHostMigrateLegacyCopilotCliEnabledConfigKey]: true }); + svc.primeMigrateLegacyGate(); + const agent = disposables.add(new ControlledDiscoveryAgent('copilot')); + const external = agent.addSession('backfilled-legacy-discovery', Date.now()); + + registerTestAgentProvider(svc, agent); + for (let attempt = 0; attempt < 50 && (await svc.getRegisteredSessions()).length === 0; attempt++) { + await timeout(0); + } + + assert.deepStrictEqual({ + discoveryStarts: agent.discoveryStarts, + providerBackfilled: await svc.isProviderRegistryBackfilled('copilot'), + registered: (await svc.getRegisteredSessions()).map(session => session.toString()), + visible: await svc.listSessions(), + }, { + discoveryStarts: 1, + providerBackfilled: true, + registered: [external.toString()], + visible: [], + }); + }); + + testWithExternalSessionClock('a deferred migration does not request discovery while external sessions are hidden', async () => { + class DeferredDiscoveryAgent extends ControlledDiscoveryAgent { + override async listChatsToMigrate(): Promise { + return AgentChatMigrationDeferred; + } + } + + const svc = createExternalSessionService(); + const agent = disposables.add(new DeferredDiscoveryAgent('codex')); + registerTestAgentProvider(svc, agent); + svc.markStartupComplete(); + await svc.listSessions(); + await svc.whenDeferredWorkSettled(); + + assert.deepStrictEqual({ + discoveryStarts: agent.discoveryStarts, + providerBackfilled: await svc.isProviderRegistryBackfilled('codex'), + }, { + discoveryStarts: 0, + providerBackfilled: false, + }); + }); + function testWithExternalSessionClock(name: string, fn: () => Promise): void { test(name, () => runWithFakedTimers({ useFakeTimers: true, @@ -3430,8 +3612,7 @@ suite('AgentService (node dispatcher)', () => { }); }); - /** An external session two newer local sessions postdate is no longer recent. */ - test('recent drops external sessions that two newer local sessions superseded', () => { + test('recent keeps its startup snapshot while recording local session updates for the next restart', () => { const hour = 60 * 60 * 1000; const at = (hourOfDay: number) => Date.UTC(2026, 0, 1) + hourOfDay * hour; const now = at(18); @@ -3441,142 +3622,113 @@ suite('AgentService (node dispatcher)', () => { modifiedTime, _meta: withSessionExternal(undefined, true), }); - const local = (id: string, startTime: number): IRegisteredSession => ({ - session: AgentSession.uri('copilot', id), - provider: 'copilot', - startTime, - modifiedTime: startTime, - external: false, - source: 'restore', - }); const catalog = [external('external-morning', at(10)), external('external-afternoon', at(16))]; - // The cutoff is snapshotted per service, so each case needs its own. - const recentIds = (...locals: IRegisteredSession[]) => { - const svc = createExternalSessionService() as unknown as { - _resolveRecentSupersedingCutoff(registered: readonly IRegisteredSession[], epoch: number): number | undefined; - _getRecentSessionKeys(sessions: readonly IAgentSessionMetadata[], now: number, supersededBefore: number | undefined): ReadonlySet; - _registryEpoch: number; - }; - const cutoff = svc._resolveRecentSupersedingCutoff(locals, svc._registryEpoch); - return [...svc._getRecentSessionKeys(catalog, now, cutoff)].map(key => AgentSession.id(URI.parse(key))).sort(); + const svc = createExternalSessionService() as unknown as { + _recordRecentLocalSessionUpdate(session: URI, modifiedTime: number): void; + _getRecentSessionKeys(sessions: readonly IAgentSessionMetadata[], now: number): ReadonlySet; + _recentLocalSessionUpdates: readonly { session: string; modifiedTime: number }[]; }; + const recentIds = () => [...svc._getRecentSessionKeys(catalog, now)].map(key => AgentSession.id(URI.parse(key))).sort(); - assert.deepStrictEqual({ - noLocalSessionsAfter: recentIds(local('local-8am', at(8)), local('local-9am', at(9))), - oneLocalSessionAfter: recentIds(local('local-11am', at(11))), - twoLocalSessionsAfterTheMorningOne: recentIds(local('local-11am', at(11)), local('local-5pm', at(17))), - twoLocalSessionsAfterBoth: recentIds(local('local-5pm', at(17)), local('local-5pm-2', at(17))), + const initial = recentIds(); + svc._recordRecentLocalSessionUpdate(AgentSession.uri('copilot', 'local-first'), at(11)); + const afterOneLocalSession = recentIds(); + svc._recordRecentLocalSessionUpdate(AgentSession.uri('copilot', 'local-first'), at(17)); + const afterSameSessionUpdatesAgain = recentIds(); + svc._recordRecentLocalSessionUpdate(AgentSession.uri('copilot', 'local-second'), at(12)); + const afterTwoDifferentSessions = recentIds(); + svc._recordRecentLocalSessionUpdate(AgentSession.uri('copilot', 'local-third'), at(17)); + const afterThreeDifferentSessions = recentIds(); + + assert.deepStrictEqual({ + initial, + afterOneLocalSession, + afterSameSessionUpdatesAgain, + afterTwoDifferentSessions, + afterThreeDifferentSessions, + recordedSessions: svc._recentLocalSessionUpdates.map(entry => AgentSession.id(URI.parse(entry.session))), }, { - noLocalSessionsAfter: ['external-afternoon', 'external-morning'], - oneLocalSessionAfter: ['external-afternoon', 'external-morning'], - twoLocalSessionsAfterTheMorningOne: ['external-afternoon'], - twoLocalSessionsAfterBoth: [], + initial: ['external-afternoon', 'external-morning'], + afterOneLocalSession: ['external-afternoon', 'external-morning'], + afterSameSessionUpdatesAgain: ['external-afternoon', 'external-morning'], + afterTwoDifferentSessions: ['external-afternoon', 'external-morning'], + afterThreeDifferentSessions: ['external-afternoon', 'external-morning'], + recordedSessions: ['local-first', 'local-third'], }); }); - /** - * The cutoff reads the registry, not the hydrated listing: a local session - * whose provider is unavailable is dropped from the latter, which would - * undercount and leave a superseded external row visible. - */ - testWithExternalSessionClock('recent counts local sessions the provider cannot hydrate', async () => { - const hour = 60 * 60 * 1000; - const now = Date.now(); - const at = (hourOfDay: number) => now - (18 - hourOfDay) * hour; - const database = new TransientRegistryWriteDatabase(); - for (const [id, startTime] of [['external-morning', at(10)], ['external-afternoon', at(16)]] as const) { - await database.registerSession(AgentSession.uri('copilot', id).toString(), { provider: 'copilot', startTime, source: 'discovery' }, { checkTombstone: true }); - } - // Registered under a provider that is never registered with the service. - for (const [id, startTime] of [['local-11am', at(11)], ['local-5pm', at(17)]] as const) { - await database.registerSession(AgentSession.uri('claude', id).toString(), { provider: 'claude', startTime, source: 'restore' }, { checkTombstone: true }); - } - await database.markProviderBackfilled('copilot'); - - const svc = createExternalSessionService(createSessionDataService(), database); - setExternalSessionsMode(svc, AgentHostExternalSessionsMode.Recent, 1); - await waitForSessionListReconciliation(svc); - const agent = disposables.add(new TimedExternalAgent('copilot')); - agent.addSession('external-morning', at(10)); - agent.addSession('external-afternoon', at(16)); + test('recent local session activity follows summary updates outside Recent mode', async () => { + const svc = createExternalSessionService(); + const agent = disposables.add(new MockAgent('copilot')); registerTestAgentProvider(svc, agent); - - const listed = await svc.listSessions(); - - assert.deepStrictEqual({ - visible: listed.map(session => AgentSession.id(session.session)).sort(), - cutoffCountedUnhydratedLocals: (svc as unknown as { _recentSupersedingCutoff: number | undefined })._recentSupersedingCutoff === at(11), - }, { - visible: ['external-afternoon'], - cutoffCountedUnhydratedLocals: true, - }); - }); - - /** A stale pass must not freeze its cutoff: the registry changed under it. */ - test('recent does not commit a superseding cutoff computed for a stale registry epoch', () => { - const at = (hourOfDay: number) => Date.UTC(2026, 0, 1) + hourOfDay * 60 * 60 * 1000; - const svc = createExternalSessionService() as unknown as { - _resolveRecentSupersedingCutoff(registered: readonly IRegisteredSession[], epoch: number): number | undefined; - _hasRecentSupersedingCutoff: boolean; - _registryEpoch: number; + const first = await svc.createSession({ provider: 'copilot' }); + const second = await svc.createSession({ provider: 'copilot' }); + const now = Date.now(); + const updateSession = async (session: URI, modifiedTime: number, turnId: string) => { + const modifiedAt = new Date(modifiedTime).toISOString(); + const changed = Event.toPromise(Event.filter( + getStateManager(svc).onDidChangeSessionSummary, + event => event.session === session.toString() && event.changes.modifiedAt === modifiedAt, + )); + getStateManager(svc).dispatchServerAction(buildDefaultChatUri(session), { + type: ActionType.ChatTurnStarted, + turnId, + startedAt: modifiedAt, + message: { text: 'hello', origin: { kind: MessageKind.User } }, + }); + await changed; }; - const locals: IRegisteredSession[] = [at(11), at(17)].map((startTime, index) => ({ - session: AgentSession.uri('copilot', `local-${index}`), - provider: 'copilot', - startTime, - modifiedTime: startTime, - external: false, - source: 'restore', - })); - const staleCutoff = svc._resolveRecentSupersedingCutoff(locals, svc._registryEpoch - 1); - const committedAfterStalePass = svc._hasRecentSupersedingCutoff; - const currentCutoff = svc._resolveRecentSupersedingCutoff(locals, svc._registryEpoch); + await updateSession(first, now + 60_000, 'turn-first'); + await updateSession(second, now + 120_000, 'turn-second'); - assert.deepStrictEqual({ staleCutoff, committedAfterStalePass, currentCutoff, committedAfterCurrentPass: svc._hasRecentSupersedingCutoff }, { - staleCutoff: at(11), - committedAfterStalePass: false, - currentCutoff: at(11), - committedAfterCurrentPass: true, - }); + const updates = (svc as unknown as { + _recentLocalSessionUpdates: readonly { session: string; modifiedTime: number }[]; + })._recentLocalSessionUpdates; + assert.deepStrictEqual(updates.map(entry => ({ + session: AgentSession.id(URI.parse(entry.session)), + modifiedTime: entry.modifiedTime, + })), [ + { session: AgentSession.id(second), modifiedTime: now + 120_000 }, + { session: AgentSession.id(first), modifiedTime: now + 60_000 }, + ]); }); - /** A first message creates a local session, so the cutoff must not re-measure per listing. */ - testWithExternalSessionClock('recent snapshots the superseding local sessions until the external mode changes', async () => { + test('recent restores local session updates after restart without listing local sessions', async () => { const hour = 60 * 60 * 1000; - const at = (hourOfDay: number) => Date.now() + hourOfDay * hour - 18 * hour; - const now = at(18); - const svc = createExternalSessionService(); - const internals = svc as unknown as { - _resolveRecentSupersedingCutoff(registered: readonly IRegisteredSession[], epoch: number): number | undefined; - _getRecentSessionKeys(sessions: readonly IAgentSessionMetadata[], now: number, supersededBefore: number | undefined): ReadonlySet; - _registryEpoch: number; - }; - const catalog: IAgentSessionMetadata[] = [ - { session: AgentSession.uri('copilot', 'external-morning'), startTime: at(10), modifiedTime: at(10), _meta: withSessionExternal(undefined, true) }, - { session: AgentSession.uri('copilot', 'external-afternoon'), startTime: at(16), modifiedTime: at(16), _meta: withSessionExternal(undefined, true) }, - ]; - const locals: IRegisteredSession[] = []; - const recentIds = () => { - const cutoff = internals._resolveRecentSupersedingCutoff(locals, internals._registryEpoch); - return [...internals._getRecentSessionKeys(catalog, now, cutoff)].map(key => AgentSession.id(URI.parse(key))).sort(); - }; - - const initial = recentIds(); - for (const id of ['local-first', 'local-second']) { - locals.push({ session: AgentSession.uri('copilot', id), provider: 'copilot', startTime: at(17), modifiedTime: at(17), external: false, source: 'restore' }); + const now = Date.now(); + const at = (hourOfDay: number) => now + (hourOfDay - 18) * hour; + const directory = mkdtempSync(join(tmpdir(), 'agent-host-recent-sessions-')); + const storageResource = URI.file(join(directory, 'storage.json')); + try { + const first = createExternalSessionService(createSessionDataService(), undefined, undefined, storageResource) as unknown as { + _recordRecentLocalSessionUpdate(session: URI, modifiedTime: number): void; + _storageService: { whenIdle(): Promise }; + dispose(): void; + }; + first._recordRecentLocalSessionUpdate(AgentSession.uri('copilot', 'local-first'), at(11)); + first._recordRecentLocalSessionUpdate(AgentSession.uri('copilot', 'local-second'), at(17)); + await first._storageService.whenIdle(); + first.dispose(); + + const restored = createExternalSessionService(createSessionDataService(), undefined, undefined, storageResource); + const agent = disposables.add(new TimedExternalAgent('copilot')); + const morning = agent.addSession('external-morning', at(10)); + const afternoon = agent.addSession('external-afternoon', at(16)); + registerTestAgentProvider(restored, agent); + await (restored as unknown as { + _registerDiscoveredChats(provider: IAgent, chats: readonly IAgentDiscoveredChat[]): Promise; + })._registerDiscoveredChats(agent, [ + discoveredChat(morning, true, at(10)), + discoveredChat(afternoon, true, at(16)), + ]); + + const listed = await restored.listSessions(AgentHostExternalSessionsMode.Recent); + + assert.deepStrictEqual(listed.map(session => AgentSession.id(session.session)), ['external-afternoon']); + } finally { + await rm(directory, { recursive: true, force: true }); } - const afterLocalSessionsCreated = recentIds(); - // Invalidation is synchronous; read before the queued reconciliation re-snapshots. - setExternalSessionsMode(svc, AgentHostExternalSessionsMode.Recent, 1); - const afterModeChange = recentIds(); - await waitForSessionListReconciliation(svc); - - assert.deepStrictEqual({ initial, afterLocalSessionsCreated, afterModeChange }, { - initial: ['external-afternoon', 'external-morning'], - afterLocalSessionsCreated: ['external-afternoon', 'external-morning'], - afterModeChange: [], - }); }); testWithExternalSessionClock('filters external sessions in every mode', async () => { @@ -4122,38 +4274,45 @@ suite('AgentService (node dispatcher)', () => { }); }); - test('a caller after a registry mutation does not join an in-flight computation', async () => { + test('callers after registry mutations share one trailing computation', async () => { const svc = disposables.add(createTestAgentService(new NullLogService(), fileService, createSessionDataService(), { _serviceBrand: undefined } as IProductService, createNoopGitService())); const agent = disposables.add(new MockAgent('copilot')); registerTestAgentProvider(svc, agent); const gate = new DeferredPromise(); - const inner = svc as unknown as { _computeSessions(mode: AgentHostExternalSessionsMode, epoch?: number): Promise }; + const inner = svc as unknown as { _computeSessions(mode: AgentHostExternalSessionsMode): Promise }; const original = inner._computeSessions; let computations = 0; - inner._computeSessions = async (mode, epoch) => { + inner._computeSessions = async mode => { computations++; await gate.p; - return original.call(svc, mode, epoch); + return original.call(svc, mode); }; const preInvalidation = svc.listSessions(); await svc.createSession({ provider: 'copilot' }); + await svc.createSession({ provider: 'copilot' }); + const postInvalidation = svc.listSessions(); + const secondPostInvalidation = svc.listSessions(); + const computationsBeforeRelease = computations; gate.complete(); - const preInvalidationCount = (await preInvalidation).length; - const postInvalidationCount = (await svc.listSessions()).length; + const [preInvalidationResult, postInvalidationResult, secondPostInvalidationResult] = await Promise.all([preInvalidation, postInvalidation, secondPostInvalidation]); assert.deepStrictEqual({ + computationsBeforeRelease, computations, - preInvalidation: preInvalidationCount, - postInvalidation: postInvalidationCount, + preInvalidation: preInvalidationResult.length, + postInvalidation: postInvalidationResult.length, + secondPostInvalidation: secondPostInvalidationResult.length, }, { + computationsBeforeRelease: 1, computations: 2, - preInvalidation: 1, - postInvalidation: 1, + preInvalidation: 2, + postInvalidation: 2, + secondPostInvalidation: 2, }); }); - test('provider registration invalidates an in-flight list computation', async () => { + test('provider registration queues a trailing list computation without overlap', async () => { const svc = disposables.add(createTestAgentService(new NullLogService(), fileService, createSessionDataService(), { _serviceBrand: undefined } as IProductService, createNoopGitService())); const gate = new DeferredPromise(); const inner = svc as unknown as { _computeSessions(mode: AgentHostExternalSessionsMode): Promise }; @@ -4169,10 +4328,11 @@ suite('AgentService (node dispatcher)', () => { const agent = disposables.add(new MockAgent('copilot')); registerTestAgentProvider(svc, agent); const afterRegistration = svc.listSessions(); + const computationsBeforeRelease = computations; gate.complete(); await Promise.all([beforeRegistration, afterRegistration]); - assert.strictEqual(computations, 2); + assert.deepStrictEqual({ computationsBeforeRelease, computations }, { computationsBeforeRelease: 1, computations: 2 }); }); test('explicitly created sessions are registered as non-external', async () => { @@ -5944,6 +6104,48 @@ suite('AgentService (node dispatcher)', () => { ); }); + test('idle provisional create and dispose do not invalidate the session list', async () => { + class ConfigurableProvisionalAgent extends MockAgent { + provisional = true; + override readonly chats: IAgentChats = withChatOverrides(getChatSurface(this), base => ({ + createChat: async (chat, context, options) => { + const created = await base.createChat(chat, context, options); + return created && this.provisional ? { ...created, provisional: true } : created; + }, + })); + } + + const localService = disposables.add(createTestAgentService(new NullLogService(), fileService, createSessionDataService(), { _serviceBrand: undefined } as IProductService, createNoopGitService())); + const agent = disposables.add(new ConfigurableProvisionalAgent('copilot')); + registerTestAgentProvider(localService, agent); + const registryEpoch = () => (localService as unknown as { _registryEpoch: number })._registryEpoch; + const initialEpoch = registryEpoch(); + + const provisional = await localService.createSession({ provider: agent.id }); + const afterProvisionalCreate = registryEpoch(); + await localService.disposeSession(provisional); + const afterProvisionalDispose = registryEpoch(); + + agent.provisional = false; + const materialized = await localService.createSession({ provider: agent.id }); + const afterMaterializedCreate = registryEpoch(); + await localService.disposeSession(materialized); + + assert.deepStrictEqual({ + initialEpoch, + afterProvisionalCreate, + afterProvisionalDispose, + afterMaterializedCreate, + afterMaterializedDispose: registryEpoch(), + }, { + initialEpoch, + afterProvisionalCreate: initialEpoch, + afterProvisionalDispose: initialEpoch, + afterMaterializedCreate: initialEpoch + 1, + afterMaterializedDispose: initialEpoch + 2, + }); + }); + test('listSessions overlays live workspace metadata over a stale provider snapshot', async () => { class DelayedListAgent extends MockAgent { readonly listStarted = new DeferredPromise(); @@ -6640,6 +6842,7 @@ suite('AgentService (node dispatcher)', () => { { git: gitState }, ); }); + }); test('subscribe to a registered session changeset URI returns a changeset snapshot', async () => { diff --git a/src/vs/platform/agentHost/test/node/claudeAgent.test.ts b/src/vs/platform/agentHost/test/node/claudeAgent.test.ts index 5a97f830943854..0e5276011547b3 100644 --- a/src/vs/platform/agentHost/test/node/claudeAgent.test.ts +++ b/src/vs/platform/agentHost/test/node/claudeAgent.test.ts @@ -5250,6 +5250,7 @@ suite('ClaudeAgent', () => { const agent = disposables.add(instantiationService.createInstance(ClaudeAgent)); const discoveredChats: number[] = []; disposables.add(agent.onDidDiscoverChats(chats => discoveredChats.push(chats.length))); + void agent.startChatDiscovery(); const sessionUri = AgentSession.uri('claude', 'materialized'); const chat = defaultChatUri(sessionUri); @@ -6228,9 +6229,9 @@ suite('ClaudeAgent — agent SDK setup channel', () => { const ctx = createTestContext(disposables); ctx.sdk.canLoadWithoutDownloadResult = false; ctx.sdk.sessionList = [{ sessionId: 'from-claude-code', summary: 'An existing chat', lastModified: 1000, createdAt: 900 }]; - // Subscribing is what starts discovery. const discovered: number[] = []; disposables.add(ctx.agent.onDidDiscoverChats(chats => discovered.push(chats.length))); + void ctx.agent.startChatDiscovery(); await settle(); const cold = { discovered: [...discovered], diff --git a/src/vs/platform/agentHost/test/node/codex/codexAgent.test.ts b/src/vs/platform/agentHost/test/node/codex/codexAgent.test.ts index 8680b2903e37b7..be0f0153c0b4f1 100644 --- a/src/vs/platform/agentHost/test/node/codex/codexAgent.test.ts +++ b/src/vs/platform/agentHost/test/node/codex/codexAgent.test.ts @@ -471,7 +471,7 @@ suite('CodexAgent', () => { const unavailable = await listChatsToMigrate.call({ ...harness, _listCodexChats: async () => undefined }); assert.deepStrictEqual({ inactive, cold, result, empty, unavailable }, { - inactive: [], + inactive: AgentChatMigrationDeferred, cold: AgentChatMigrationDeferred, result: chats.slice(0, 2), empty: [], diff --git a/src/vs/platform/agentHost/test/node/codex/codexModelRefresh.test.ts b/src/vs/platform/agentHost/test/node/codex/codexModelRefresh.test.ts index 9782a0cc414c07..fd6aa282c0897e 100644 --- a/src/vs/platform/agentHost/test/node/codex/codexModelRefresh.test.ts +++ b/src/vs/platform/agentHost/test/node/codex/codexModelRefresh.test.ts @@ -27,7 +27,7 @@ import { IAgentSdkDownloader } from '../../../node/agentSdkDownloader.js'; import { RecordingAgentSdkDownloader } from '../testAgentSdkDownloader.js'; import { IAgentHostCheckpointService, NULL_CHECKPOINT_SERVICE } from '../../../common/agentHostCheckpointService.js'; import { AGENT_SDK_SETUP_DOWNLOAD_REQUEST_KEY, AGENT_SDK_SETUP_RELOAD_REQUEST_KEY, readAgentSdkSetupInfos } from '../../../common/agentSdkSetup.js'; -import { AgentSession } from '../../../common/agent.js'; +import { AgentChatMigrationDeferred, AgentSession } from '../../../common/agent.js'; import { buildDefaultChatUri } from '../../../common/state/sessionState.js'; import { CodexAgent, toCodexModelSelectionId } from '../../../node/codex/codexAgent.js'; import { ICodexProxyService } from '../../../node/codex/codexProxyService.js'; @@ -137,6 +137,9 @@ function createChatGPTConnection(account: unknown = { type: 'chatgpt', email: 'p if (method === 'model/list') { return modelListResponse; } + if (method === 'thread/list') { + return { data: [], nextCursor: null }; + } throw new Error(`Unexpected request: ${method}`); }, }, @@ -172,7 +175,7 @@ suite('CodexAgent model refresh', () => { assert.deepStrictEqual({ connectionRequested, metadata, migrated, models: agent.models.get() }, { connectionRequested: false, metadata: undefined, - migrated: [], + migrated: AgentChatMigrationDeferred, models: [], }); @@ -189,10 +192,12 @@ suite('CodexAgent model refresh', () => { connectionRequested, // One enumeration, not one per caller that happened to want the connection. enumerations: requests.filter(method => method === 'model/list').length, + discoveries: requests.filter(method => method === 'thread/list').length, models: agent.models.get().map(model => ({ provider: model.provider, id: model.id, name: model.name, meta: model._meta })), }, { connectionRequested: true, enumerations: 1, + discoveries: 0, models: [{ provider: 'codex', id: toCodexModelSelectionId('openai', 'gpt-5.6-sol'), @@ -270,6 +275,29 @@ suite('CodexAgent model refresh', () => { }); }); + test('starts host-requested chat discovery when Codex activates', async () => { + const agent = createAgent(disposables, async () => [], { [AgentHostConfigKey.AllowSignedOutWhenUsable]: true }); + const requests: string[] = []; + const connection = createChatGPTConnection(undefined, requests); + agent['_ensureConnection'] = async () => { + agent['_connection'] = connection as never; + return connection as never; + }; + + await agent.startChatDiscovery(); + const discoveriesBeforeActivation = requests.filter(method => method === 'thread/list').length; + agent['_activate'](); + await agent['_codexChatDiscovery']; + + assert.deepStrictEqual({ + discoveriesBeforeActivation, + discoveriesAfterActivation: requests.filter(method => method === 'thread/list').length, + }, { + discoveriesBeforeActivation: 0, + discoveriesAfterActivation: 1, + }); + }); + test('queues a fresh model refresh when Codex activates during an ambient refresh', async () => { const copilotModels = [{ id: 'copilot-model', name: 'Copilot Model', supported_endpoints: ['/responses'] }] as CCAModel[]; const ambientRefreshStarted = new DeferredPromise(); diff --git a/src/vs/platform/agentHost/test/node/copilotAgent.test.ts b/src/vs/platform/agentHost/test/node/copilotAgent.test.ts index d4c6c296a3df7f..d59ded3036bc46 100644 --- a/src/vs/platform/agentHost/test/node/copilotAgent.test.ts +++ b/src/vs/platform/agentHost/test/node/copilotAgent.test.ts @@ -1034,7 +1034,7 @@ async function collectDiscoveredChats(agent: CopilotAgent): Promise discovered.push(...chats)); try { - await (agent as unknown as { _startCopilotChatDiscovery(): Promise })._startCopilotChatDiscovery(); + await agent.startChatDiscovery(); return discovered.map(chat => ({ id: sessionIdOfChat(chat.chat), external: chat.external, @@ -5644,6 +5644,7 @@ suite('CopilotAgent', () => { const discoveredChats: Array = []; const listener = agent.onDidDiscoverChats(chats => discoveredChats.push(chats)); try { + void agent.startChatDiscovery(); for (let i = 0; i < 10; i++) { await timeout(0); } @@ -5678,6 +5679,7 @@ suite('CopilotAgent', () => { const discoveredChats: Array = []; const listener = agent.onDidDiscoverChats(chats => discoveredChats.push(chats)); try { + void agent.startChatDiscovery(); for (let i = 0; i < 50 && discoveredChats.length === 0; i++) { await timeout(0); } @@ -5708,6 +5710,7 @@ suite('CopilotAgent', () => { const discoveredChats: Array = []; const listener = agent.onDidDiscoverChats(chats => discoveredChats.push(chats)); try { + void agent.startChatDiscovery(); await listStarted.p; // The gate was snapshotted as enabled at startup, so disabling it mid // discovery is ignored: the adoptable chat still surfaces. @@ -5742,6 +5745,7 @@ suite('CopilotAgent', () => { const discoveredChats: Array = []; const listener = agent.onDidDiscoverChats(chats => discoveredChats.push(chats)); try { + void agent.startChatDiscovery(); for (let i = 0; i < 50 && discoveredChats.length === 0; i++) { await timeout(0); } @@ -6075,7 +6079,7 @@ suite('CopilotAgent', () => { const discovered: IAgentDiscoveredChat[] = []; const listener = agent.onDidDiscoverChats(chats => discovered.push(...chats)); try { - await (agent as unknown as { _startCopilotChatDiscovery(): Promise })._startCopilotChatDiscovery(); + await agent.startChatDiscovery(); return discovered.map(chat => ({ id: sessionIdOfChat(chat.chat), workingDirectory: chat.workingDirectories?.[0]?.fsPath, @@ -11788,6 +11792,87 @@ suite('CopilotAgent', () => { } }); + test('bridges an existing worktree checkout so the recorded base branch survives without a remote', async () => { + // #333642: the CLI committed the session's work onto the worktree branch. + // The checkout still exists, so the old bridge skipped it and — with no + // remote to resolve a default branch — persisted no base branch, hiding + // every committed-on-branch change. The marker's recorded base must flow + // through so Branch Changes diffs against the merge-base. + const userHome = URI.file(await fs.mkdtemp(`${os.tmpdir()}/adopt-home-`)); + const repositoryRoot = await fs.mkdtemp(`${os.tmpdir()}/adopt-repo-`); + const worktreePath = join(repositoryRoot, '..', `present.worktrees-${Date.now()}`, 'feature-z'); + await fs.mkdir(worktreePath, { recursive: true }); + const sessionId = 'legacy-worktree-present'; + const session = AgentSession.uri('copilotcli', sessionId); + const sessionDataService = disposables.add(new TestSessionDataService()); + const client = new TestCopilotClient([sdkSession(sessionId, worktreePath)]); + const agent = createTestAgent(disposables, { sessionDataService, copilotClient: client, userHome }); + try { + await agent.authenticate('https://api.github.com', 'token'); + await writeExtensionHostMarker(userHome, sessionId, { + origin: 'vscode', + worktreeProperties: { worktreePath, repositoryPath: repositoryRoot, branchName: 'feature/z', baseBranchName: 'main' }, + }); + + const adopted = await ensureDefaultChatAdopted(agent, session); + + assert.deepStrictEqual( + { + adopted: adopted.adopted, + worktree: adopted.worktree && { + branchName: adopted.worktree.branchName, + baseBranch: adopted.worktree.baseBranch, + worktreePath: adopted.worktree.worktreePath.fsPath, + repositoryRoot: adopted.worktree.repositoryRoot.fsPath, + }, + }, + { + adopted: true, + worktree: { branchName: 'feature/z', baseBranch: 'main', worktreePath: URI.file(worktreePath).fsPath, repositoryRoot: URI.file(repositoryRoot).fsPath }, + }, + ); + } finally { + await fs.rm(userHome.fsPath, { recursive: true, force: true }); + await fs.rm(repositoryRoot, { recursive: true, force: true }); + await fs.rm(join(worktreePath, '..'), { recursive: true, force: true }); + await disposeAgent(agent); + } + }); + + test('leaves an existing worktree checkout without a recorded base branch to the probe-based bridge', async () => { + // An older marker carries no base branch. Taking over here would drop the + // probe's `origin/HEAD` fallback, so the checkout-exists case must defer to + // it (adoption still succeeds, just with no worktree in the result). + const userHome = URI.file(await fs.mkdtemp(`${os.tmpdir()}/adopt-home-`)); + const repositoryRoot = await fs.mkdtemp(`${os.tmpdir()}/adopt-repo-`); + const worktreePath = join(repositoryRoot, '..', `present.worktrees-${Date.now()}-nb`, 'feature-w'); + await fs.mkdir(worktreePath, { recursive: true }); + const sessionId = 'legacy-worktree-present-no-base'; + const session = AgentSession.uri('copilotcli', sessionId); + const sessionDataService = disposables.add(new TestSessionDataService()); + const client = new TestCopilotClient([sdkSession(sessionId, worktreePath)]); + const agent = createTestAgent(disposables, { sessionDataService, copilotClient: client, userHome }); + try { + await agent.authenticate('https://api.github.com', 'token'); + await writeExtensionHostMarker(userHome, sessionId, { + origin: 'vscode', + worktreeProperties: { worktreePath, repositoryPath: repositoryRoot, branchName: 'feature/w' }, + }); + + const adopted = await ensureDefaultChatAdopted(agent, session); + + assert.deepStrictEqual( + { adopted: adopted.adopted, worktree: adopted.worktree }, + { adopted: true, worktree: undefined }, + ); + } finally { + await fs.rm(userHome.fsPath, { recursive: true, force: true }); + await fs.rm(repositoryRoot, { recursive: true, force: true }); + await fs.rm(join(worktreePath, '..'), { recursive: true, force: true }); + await disposeAgent(agent); + } + }); + test('adopts a deleted worktree with the local repository as its project, not the remote', async () => { // Git resolution runs in the (missing) checkout and falls back to the // remote, whose URI is not a path — the session could then never be @@ -11859,6 +11944,82 @@ suite('CopilotAgent', () => { } }); + test('persists the last migrated turn id from the request sidecar on adoption', async () => { + // The chat editor uses this migration boundary to attribute the session's + // committed changes to the final migrated turn and no post-adoption turn. + const userHome = URI.file(await fs.mkdtemp(`${os.tmpdir()}/adopt-home-`)); + const workingDirectory = await fs.mkdtemp(`${os.tmpdir()}/adopt-lastturn-`); + const sessionId = 'legacy-lastturn'; + const session = AgentSession.uri('copilotcli', sessionId); + const sessionDataService = disposables.add(new TestSessionDataService()); + const client = new TestCopilotClient([sdkSession(sessionId, workingDirectory)]); + const agent = createTestAgent(disposables, { sessionDataService, copilotClient: client, userHome }); + try { + await agent.authenticate('https://api.github.com', 'token'); + await writeExtensionHostMarker(userHome, sessionId); + await writeExtensionHostRequestDetails(userHome, sessionId, [ + { copilotRequestId: 'turn-1', creditsUsed: 1 }, + { copilotRequestId: 'turn-2', creditsUsed: 2 }, + ]); + + await ensureDefaultChatAdopted(agent, session); + + const db = await sessionDataService.tryOpenDatabase(session); + const lastTurn = await db?.object.getMetadata('agentHost.ehcliLastMigratedTurn'); + db?.dispose(); + + assert.strictEqual(lastTurn, 'turn-2'); + } finally { + await fs.rm(userHome.fsPath, { recursive: true, force: true }); + await fs.rm(workingDirectory, { recursive: true, force: true }); + await disposeAgent(agent); + } + }); + + test('backfills the base branch and last migrated turn for a session migrated by an older build', async () => { + // A no-remote worktree session migrated by the previous code kept a working + // directory (so adoption short-circuits as `alreadyNative`) but no base + // branch, leaving its diff anchored to HEAD. Repair it in place (#333642). + const userHome = URI.file(await fs.mkdtemp(`${os.tmpdir()}/adopt-home-`)); + const repositoryRoot = await fs.mkdtemp(`${os.tmpdir()}/adopt-old-repo-`); + const workingDirectory = await fs.mkdtemp(`${os.tmpdir()}/adopt-old-wt-`); + const sessionId = 'legacy-old-no-base'; + const session = AgentSession.uri('copilotcli', sessionId); + const sessionDataService = disposables.add(new TestSessionDataService()); + const client = new TestCopilotClient([sdkSession(sessionId, workingDirectory)]); + const agent = createTestAgent(disposables, { sessionDataService, copilotClient: client, userHome }); + try { + await agent.authenticate('https://api.github.com', 'token'); + await writeExtensionHostMarker(userHome, sessionId, { + origin: 'vscode', + worktreeProperties: { worktreePath: workingDirectory, repositoryPath: repositoryRoot, branchName: 'feature/x', baseBranchName: 'main' }, + }); + await writeExtensionHostRequestDetails(userHome, sessionId, [{ copilotRequestId: 'turn-9', creditsUsed: 1 }]); + // Metadata the older build wrote: adopted with a working directory, but no base branch or boundary. + const seed = sessionDataService.openDatabase(session); + await seed.object.setMetadata('copilot.workingDirectory', URI.file(workingDirectory).toString()); + await seed.object.setMetadata('agentHost.ehcliAdopted', 'true'); + seed.dispose(); + + const adopted = await ensureDefaultChatAdopted(agent, session); + + const db = await sessionDataService.tryOpenDatabase(session); + const baseBranch = await db?.object.getMetadata('agentHost.diffBaseBranch'); + const lastTurn = await db?.object.getMetadata('agentHost.ehcliLastMigratedTurn'); + db?.dispose(); + + assert.deepStrictEqual( + { reason: adopted.reason, baseBranch, lastTurn }, + { reason: 'alreadyNative', baseBranch: 'main', lastTurn: 'turn-9' }, + ); + } finally { + await fs.rm(userHome.fsPath, { recursive: true, force: true }); + await fs.rm(repositoryRoot, { recursive: true, force: true }); + await fs.rm(workingDirectory, { recursive: true, force: true }); + await disposeAgent(agent); + } + }); + test('does not backfill the adopted-legacy marker onto a native session', async () => { // No extension-host marker means the session was never a legacy chat. const userHome = URI.file(await fs.mkdtemp(`${os.tmpdir()}/adopt-home-`)); diff --git a/src/vs/platform/agentHost/test/node/sshRemoteAgentHostHelpers.test.ts b/src/vs/platform/agentHost/test/node/sshRemoteAgentHostHelpers.test.ts index 1a0151f771625d..bbc9e340b4e8f2 100644 --- a/src/vs/platform/agentHost/test/node/sshRemoteAgentHostHelpers.test.ts +++ b/src/vs/platform/agentHost/test/node/sshRemoteAgentHostHelpers.test.ts @@ -4,7 +4,9 @@ *--------------------------------------------------------------------------------------------*/ import assert from 'assert'; +import { CancellationTokenSource } from '../../../../base/common/cancellation.js'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js'; +import { NullLogService } from '../../../log/common/log.js'; import { TelemetryConfiguration } from '../../../telemetry/common/telemetry.js'; import { AGENT_HOST_ENDPOINT_REGISTRY_SCHEMA_VERSION, type IAgentHostEndpointMetadata } from '../../common/agentHostEndpointRegistry.js'; import { @@ -17,6 +19,7 @@ import { buildFindFallbackCLICommand, filterLiveAgentHostEndpoints, findNewAgentHostEndpoint, + getNewAgentHostRegistrationTimeoutMs, getRemoteCLIArchiveName, getRemoteCLIBin, getRemoteCLIDataDir, @@ -33,6 +36,7 @@ import { waitForNewStandaloneEndpoint, type ISshExec, } from '../../node/sshRemoteAgentHostHelpers.js'; +import { ensureRemoteAgentHostCliInstalled } from '../../node/remoteAgentHostCliInstaller.js'; suite('SSH Remote Agent Host Helpers', () => { @@ -668,6 +672,72 @@ suite('SSH Remote Agent Host Helpers', () => { }); }); + suite('ensureRemoteAgentHostCliInstalled', () => { + test('reports whether a CLI was reused or installed', async () => { + const cliBin = getRemoteCLIBin('.vscode-server', 'insider'); + const options = { + serverDataFolderName: '.vscode-server', + quality: 'insider', + commit: undefined, + reportInstalling: () => { }, + logService: new NullLogService(), + }; + const commit = '1234567890abcdef1234567890abcdef12345678'; + const pinnedOptions = { ...options, commit }; + const pinnedCliBin = getRemoteCLIBin('.vscode-server', 'insider', commit); + const reused = await ensureRemoteAgentHostCliInstalled( + async () => ({ stdout: '1.0.0\n__vscode_cli_update_exit_code__:0\n', stderr: '', code: 0 }), + { os: 'linux', arch: 'x64' }, + options, + ); + let calls = 0; + const installed = await ensureRemoteAgentHostCliInstalled( + async () => { + calls++; + return { stdout: '', stderr: '', code: calls === 1 ? 1 : 0 }; + }, + { os: 'linux', arch: 'x64' }, + options, + ); + const reusedPinned = await ensureRemoteAgentHostCliInstalled( + async () => ({ stdout: '', stderr: '', code: 0 }), + { os: 'linux', arch: 'x64' }, + pinnedOptions, + ); + calls = 0; + const installedPinned = await ensureRemoteAgentHostCliInstalled( + async () => { + calls++; + return { stdout: '', stderr: '', code: calls === 1 ? 1 : 0 }; + }, + { os: 'linux', arch: 'x64' }, + pinnedOptions, + ); + + assert.deepStrictEqual( + { + reused, + installed, + reusedPinned, + installedPinned, + registrationTimeouts: { + reused: getNewAgentHostRegistrationTimeoutMs(reused.installed), + installed: getNewAgentHostRegistrationTimeoutMs(installed.installed), + reusedPinned: getNewAgentHostRegistrationTimeoutMs(reusedPinned.installed), + installedPinned: getNewAgentHostRegistrationTimeoutMs(installedPinned.installed), + }, + }, + { + reused: { cliBin, installed: false }, + installed: { cliBin, installed: true }, + reusedPinned: { cliBin: pinnedCliBin, installed: false }, + installedPinned: { cliBin: pinnedCliBin, installed: true }, + registrationTimeouts: { reused: undefined, installed: 300_000, reusedPinned: undefined, installedPinned: 300_000 }, + }, + ); + }); + }); + suite('waitForNewStandaloneEndpoint', () => { test('resolves as soon as the new endpoint appears', async () => { const before = [makeEndpoint({ type: 'standalone', pid: 1, instanceId: 'old' })]; @@ -683,13 +753,48 @@ suite('SSH Remote Agent Host Helpers', () => { assert.ok(poll >= 2); }); - test('throws once the attempt budget is exhausted', async () => { + test('uses the default short deadline when no timeout is supplied', async () => { const before = [makeEndpoint({ type: 'standalone', pid: 1, instanceId: 'old' })]; const exec: ISshExec = async () => ({ stdout: JSON.stringify({ userDataPath: '/x', endpoints: before }), stderr: '', code: 0 }); await assert.rejects( - () => waitForNewStandaloneEndpoint(exec, '~/.vscode-server/code', '~/.vscode-server/cli', '/x', before, { attempts: 2, intervalMs: 1 }), - /Timed out waiting/, + () => waitForNewStandaloneEndpoint(exec, '~/.vscode-server/code', '~/.vscode-server/cli', '/x', before, { intervalMs: 1 }), + /deadline 20ms/, ); }); + + test('keeps polling past the default deadline when given a longer deadline', async () => { + const before = [makeEndpoint({ type: 'standalone', pid: 1, instanceId: 'old' })]; + const spawned = makeEndpoint({ type: 'standalone', pid: 2, instanceId: 'new' }); + let polls = 0; + const exec: ISshExec = async () => { + polls++; + const endpoints = polls <= 20 ? before : [...before, spawned]; + return { stdout: JSON.stringify({ userDataPath: '/x', endpoints }), stderr: '', code: 0 }; + }; + + const result = await waitForNewStandaloneEndpoint(exec, '~/.vscode-server/code', '~/.vscode-server/cli', '/x', before, { intervalMs: 1, timeoutMs: getNewAgentHostRegistrationTimeoutMs(true) }); + assert.deepStrictEqual({ result, polls }, { result: spawned, polls: 21 }); + }); + + test('cancels promptly while waiting for registration', async () => { + const before = [makeEndpoint({ type: 'standalone', pid: 1, instanceId: 'old' })]; + const cancellationSource = new CancellationTokenSource(); + let polls = 0; + const exec: ISshExec = async () => { + polls++; + cancellationSource.cancel(); + return { stdout: JSON.stringify({ userDataPath: '/x', endpoints: before }), stderr: '', code: 0 }; + }; + + try { + await assert.rejects( + () => waitForNewStandaloneEndpoint(exec, '~/.vscode-server/code', '~/.vscode-server/cli', '/x', before, { timeoutMs: 60_000, token: cancellationSource.token }), + /Canceled/, + ); + assert.deepStrictEqual(polls, 1); + } finally { + cancellationSource.dispose(); + } + }); }); }); diff --git a/src/vs/platform/agentHost/test/node/wslRemoteAgentHostHelpers.test.ts b/src/vs/platform/agentHost/test/node/wslRemoteAgentHostHelpers.test.ts index 600eee006a8a47..d05ec0926d996c 100644 --- a/src/vs/platform/agentHost/test/node/wslRemoteAgentHostHelpers.test.ts +++ b/src/vs/platform/agentHost/test/node/wslRemoteAgentHostHelpers.test.ts @@ -114,7 +114,7 @@ suite('WSL Remote Agent Host Helpers', () => { telemetryLevel: TelemetryConfiguration.OFF, }); - assert.ok(script.endsWith(`exec ~/.vscode-server/code-${commit} --cli-data-dir ~/.vscode-server/cli --telemetry-level off agent host --port 0`)); + assert.ok(script.endsWith(`exec ~/.vscode-server/code-${commit} --cli-data-dir ~/.vscode-server/cli --telemetry-level off agent host --port 0 --idle-timeout 300`)); }); test('exports telemetry disablement for a custom command', () => { diff --git a/src/vs/platform/agentHost/test/node/wslRemoteAgentHostService.test.ts b/src/vs/platform/agentHost/test/node/wslRemoteAgentHostService.test.ts new file mode 100644 index 00000000000000..95ea58912ec342 --- /dev/null +++ b/src/vs/platform/agentHost/test/node/wslRemoteAgentHostService.test.ts @@ -0,0 +1,170 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import assert from 'assert'; +import * as cp from 'child_process'; +import { EventEmitter } from 'events'; +import { DeferredPromise, timeout } from '../../../../base/common/async.js'; +import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js'; +import { runWithFakedTimers } from '../../../../base/test/common/timeTravelScheduler.js'; +import { NullLogService } from '../../../log/common/log.js'; +import type { IProductService } from '../../../product/common/productService.js'; +import { NullTelemetryService } from '../../../telemetry/common/telemetryUtils.js'; +import type { IWSLConnectResult } from '../../common/wslRemoteAgentHost.js'; +import { WSLRemoteAgentHostMainService } from '../../node/wslRemoteAgentHostService.js'; +import type WebSocket from 'ws'; + +class MockWSLChild extends EventEmitter { + readonly stdout = new EventEmitter(); + readonly stderr = new EventEmitter(); + + exitCode: number | null = null; + signalCode: NodeJS.Signals | null = null; + killCalls = 0; + + kill(_signal?: NodeJS.Signals): boolean { + this.killCalls++; + if (this.exitCode === null && this.signalCode === null) { + this.signalCode = 'SIGTERM'; + queueMicrotask(() => this.emit('exit', null, 'SIGTERM')); + } + return true; + } + + emitStdout(text: string): void { + this.stdout.emit('data', Buffer.from(text)); + } +} + +class MockWebSocket { + on(_event: string, _listener: (...args: never[]) => void): this { + return this; + } + + close(): void { + } +} + +/** + * In-process WSL service double that controls platform detection, process + * output, and WebSocket creation without spawning WSL or loading `ws`. + */ +class TestableWSLRemoteAgentHostMainService extends WSLRemoteAgentHostMainService { + readonly children: MockWSLChild[] = []; + + private readonly _platform = new DeferredPromise<{ os: string; arch: string }>(); + + resolvePlatform(): void { + this._platform.complete({ os: 'linux', arch: 'x64' }); + } + + protected override _spawnAgentHost(_distro: string, _script: string): cp.ChildProcess { + const child = new MockWSLChild(); + this.children.push(child); + return child as unknown as cp.ChildProcess; + } + + protected override _resolvePlatform(_distro: string): Promise<{ os: string; arch: string }> { + return this._platform.p; + } + + protected override async _openWebSocket(_url: string): Promise { + return new MockWebSocket() as never; + } +} + +function createService(): TestableWSLRemoteAgentHostMainService { + const productService: Pick = { + _serviceBrand: undefined, + quality: 'insider', + serverDataFolderName: '.vscode-server', + commit: 'a'.repeat(40), + }; + return new TestableWSLRemoteAgentHostMainService( + new NullLogService(), + productService as IProductService, + NullTelemetryService, + ); +} + +suite('WSL Remote Agent Host Service', () => { + const disposables = ensureNoDisposablesAreLeakedInTestSuite(); + + test('deduplicates simultaneous connects to one distro', async () => { + const service = disposables.add(createService()); + const first = service.connect({ distro: 'Ubuntu', name: 'Ubuntu' }); + const second = service.connect({ distro: 'Ubuntu', name: 'Ubuntu' }); + + assert.strictEqual(first, second); + + service.resolvePlatform(); + await Promise.resolve(); + service.children[0].emitStdout('ws://127.0.0.1:3000?tkn=token\n'); + const [firstResult, secondResult] = await Promise.all([first, second]); + + assert.deepStrictEqual( + { spawnCount: service.children.length, sameResult: firstResult === secondResult, results: [firstResult, secondResult] }, + { + spawnCount: 1, + sameResult: true, + results: [ + { + connectionId: firstResult.connectionId, + address: 'wsl:Ubuntu', + distro: 'Ubuntu', + name: 'Ubuntu', + connectionToken: 'token', + }, + { + connectionId: firstResult.connectionId, + address: 'wsl:Ubuntu', + distro: 'Ubuntu', + name: 'Ubuntu', + connectionToken: 'token', + }, + ], + }, + ); + }); + + test('keeps a chatty bootstrap alive past the output-idle timeout', async () => { + return runWithFakedTimers({ useFakeTimers: true, maxTaskCount: 10_000 }, async () => { + const service = disposables.add(createService()); + const connect = service.connect({ distro: 'Ubuntu', name: 'Ubuntu' }); + service.resolvePlatform(); + await Promise.resolve(); + + const child = service.children[0]; + await timeout(59_000); + child.emitStdout('Downloading server 50%\n'); + await timeout(59_000); + child.emitStdout('ws://127.0.0.1:3000?tkn=token\n'); + + const result = await connect; + assert.deepStrictEqual( + { distro: result.distro, address: result.address, connectionToken: result.connectionToken }, + { distro: 'Ubuntu', address: 'wsl:Ubuntu', connectionToken: 'token' }, + ); + }); + }); + + test('fails a silent bootstrap after the output-idle timeout', async () => { + return runWithFakedTimers({ useFakeTimers: true, maxTaskCount: 10_000 }, async () => { + const service = disposables.add(createService()); + const rejected = service.connect({ distro: 'Ubuntu', name: 'Ubuntu' }).then( + result => result, + error => error instanceof Error ? error : new Error(String(error)), + ); + service.resolvePlatform(); + await Promise.resolve(); + + await timeout(60_001); + const result = await rejected; + + assert.ok(result instanceof Error); + assert.match(result.message, /no output for 60000ms/); + }); + }); +}); diff --git a/src/vs/sessions/common/devContainerAgentHostService.ts b/src/vs/sessions/common/devContainerAgentHostService.ts index a22a74daa66d6a..c10098465ac62f 100644 --- a/src/vs/sessions/common/devContainerAgentHostService.ts +++ b/src/vs/sessions/common/devContainerAgentHostService.ts @@ -6,13 +6,13 @@ import { CancellationToken } from '../../base/common/cancellation.js'; import { IDisposable } from '../../base/common/lifecycle.js'; import { URI } from '../../base/common/uri.js'; -import { IAgentConnection } from '../../platform/agentHost/common/agentService.js'; +import { IProtocolTransport } from '../../platform/agentHost/common/state/sessionTransport.js'; import { createDecorator } from '../../platform/instantiation/common/instantiation.js'; /** Hidden setting that enables Dev Container Agent Host sessions. */ export const DevContainerAgentHostEnabledSettingId = 'chat.agentHost.devContainer.enabled'; -/** Connected Agent Host and workspace mapping produced by a Dev Container connector. */ +/** Agent Host transport and workspace mapping produced by a Dev Container connector. */ export interface IDevContainerAgentHostConnection { /** * Stable address that uniquely identifies this source workspace's running @@ -20,7 +20,7 @@ export interface IDevContainerAgentHostConnection { */ readonly address: string; readonly name: string; - readonly connection: IAgentConnection & IDisposable; + readonly transportFactory: () => IProtocolTransport; readonly transportDisposable?: IDisposable; readonly workspaceUri: URI; readonly defaultDirectory?: string; @@ -30,7 +30,7 @@ export interface IDevContainerAgentHostConnection { export interface IDevContainerAgentHostConnector { /** Whether the workspace has a supported configuration and Docker is available. */ isAvailable(workspaceUri: URI): Promise; - connect(workspaceUri: URI, token: CancellationToken): Promise; + createConnection(workspaceUri: URI, address: string, token: CancellationToken): Promise; } /** Sessions provider and workspace selected after connecting a Dev Container. */ diff --git a/src/vs/sessions/contrib/changes/browser/sessionChangesEditor.ts b/src/vs/sessions/contrib/changes/browser/sessionChangesEditor.ts index fb59f5af11da67..214f63182c7b6c 100644 --- a/src/vs/sessions/contrib/changes/browser/sessionChangesEditor.ts +++ b/src/vs/sessions/contrib/changes/browser/sessionChangesEditor.ts @@ -301,6 +301,9 @@ export class SessionChangesEditor extends AbstractEditorWithViewState { await super.setInput(input, options, context, token); + if (token.isCancellationRequested) { + return; + } const sessionResource = this.sessionChangesService.getSessionResource(input.multiDiffSource); this._inputSessionResource.set(sessionResource, undefined); const viewModel = await input.getViewModel(); diff --git a/src/vs/sessions/contrib/changes/test/browser/sessionChangesEditorInput.test.ts b/src/vs/sessions/contrib/changes/test/browser/sessionChangesEditorInput.test.ts index ff6a18689a7d42..6fff73fd0706bc 100644 --- a/src/vs/sessions/contrib/changes/test/browser/sessionChangesEditorInput.test.ts +++ b/src/vs/sessions/contrib/changes/test/browser/sessionChangesEditorInput.test.ts @@ -4,6 +4,7 @@ *--------------------------------------------------------------------------------------------*/ import assert from 'assert'; +import { CancellationTokenSource } from '../../../../../base/common/cancellation.js'; import { Emitter, Event, ValueWithChangeEvent } from '../../../../../base/common/event.js'; import { URI } from '../../../../../base/common/uri.js'; import { mock } from '../../../../../base/test/common/mock.js'; @@ -112,6 +113,38 @@ suite('SessionChangesEditorInput', () => { }); }); + test('does not resolve a canceled editor input', async () => { + class TestSessionChangesEditorInput extends SessionChangesEditorInput { + viewModelRequested = false; + + override async getViewModel(): Promise { + this.viewModelRequested = true; + throw new Error('Canceled input must not be resolved'); + } + } + + const instantiationService = workbenchInstantiationService(undefined, disposables); + instantiationService.stub(IChangesViewService, {}); + instantiationService.stub(IAgentWorkbenchLayoutService, {}); + instantiationService.stub(ISessionChangesService, {}); + instantiationService.stub(IWorkbenchLayoutService, { + onDidChangePartVisibility: Event.None, + isVisible: () => true, + }); + + const editor = disposables.add(instantiationService.createInstance(SessionChangesEditor, new TestEditorGroupView(1))); + const input = disposables.add(instantiationService.createInstance( + TestSessionChangesEditorInput, + URI.parse('changes-multi-diff-source:?{"sessionResource":"agent-host-copilotcli:/session"}'), + )); + const operation = disposables.add(new CancellationTokenSource()); + operation.cancel(); + + await editor.setInput(input, undefined, {}, operation.token); + + assert.deepStrictEqual(input.viewModelRequested, false); + }); + test('updates managed Changes editor capabilities with editor area visibility', () => { const instantiationService = disposables.add(new TestInstantiationService()); let editorVisible = false; diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/REMOTE_AGENT_HOST_SESSIONS_PROVIDER.md b/src/vs/sessions/contrib/providers/remoteAgentHost/REMOTE_AGENT_HOST_SESSIONS_PROVIDER.md index 88dcf56fcb3765..e79b7d08d49170 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/REMOTE_AGENT_HOST_SESSIONS_PROVIDER.md +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/REMOTE_AGENT_HOST_SESSIONS_PROVIDER.md @@ -10,7 +10,7 @@ Shared Agent Host adaptation is specified in [AGENT_HOST_SESSIONS_PROVIDER.md](. ## Registration -`RemoteAgentHostContribution` observes `IRemoteAgentHostService` connections. It creates and registers one provider per connection and disposes the provider when that connection is removed. +Kind-specific contributions create and register one provider for each remote host they own, disposing it when that host is removed. `RemoteAgentHostContribution` observes connections for shared filesystem, agent-discovery, model, terminal, and authentication wiring. Agent discovery is dynamic. Changes to a host's advertised agents update the provider's session types without recreating the provider. @@ -46,13 +46,12 @@ Grouping changes these behaviors: ## Connection ownership -The remote contribution owns: +The remote Agent Host service owns protocol connection construction, handshake classification, status, retry, and disposal. -- connect, disconnect, and reconnect policy; -- authentication and interactive connection prompts; -- remote filesystem browsing; -- transport diagnostics and connection status; -- connection-scoped listener disposal. +`RemoteAgentHostContribution` owns the workbench integration for a live connection: remote filesystem browsing, agent and model discovery, terminals, authentication, and connection-scoped listener disposal. + +Transport-specific callers own discovery, on-demand staging, credentials, and connection leases. They stage +their context by address, request an explicit reconnect, and wait for the service to report the connection. The provider exposes connection state through `IAgentHostSessionsProvider` and delegates protocol operations to the live connection. Disconnecting clears live state without manufacturing successful operation results. @@ -73,6 +72,7 @@ Concurrent prompts use the shared setup operation where credentials are shared. The remote Agent Host services may remember a user's preferred run location. The owning location-preference service defines its persistence key and selection policy. Providers consume the resolved location; they do not duplicate preference state in session metadata. Transport-specific fallback and retry algorithms belong in the owning SSH, tunnel, or remote-host service and its tests. +Tunnel discovery persists picker dismissals independently from auto-connect suppression; only an explicit user connection clears a dismissal. ## Testing @@ -82,7 +82,7 @@ Focused tests live beside the remote provider and remote-host services. Tests ow `DevContainerAgentHostService` provides the desktop-only connection boundary for an Agent Host running inside a Dev Container. VS Code bundles `@devcontainers/cli` and runs that pinned version through its Electron-as-Node runtime; Docker and related tools are still resolved from the user's shell environment. The desktop connector runs `devcontainer up` for the selected local workspace, installs the matching VS Code remote CLI inside the container, and reuses or launches a dedicated standalone Agent Host. A shared-process relay carries the Agent Host WebSocket protocol over `devcontainer exec` standard input/output. -The service registers the connected client as a runtime-only `DevContainer` managed remote connection and creates a `RemoteAgentHostSessionsProvider` around it. The shared remote Agent Host contribution observes the managed connection and supplies connection-level filesystem, model, terminal, and log integration. Dev Container CLI output is streamed into one stable `Dev Container ()` Output channel per source workspace, which is reused across connection attempts. +The service stages a runtime-only `DevContainer` entry and asks the remote Agent Host service to connect its factory-built client, then creates a `RemoteAgentHostSessionsProvider` around it. The shared remote Agent Host contribution observes the connection and supplies connection-level filesystem, model, terminal, and log integration. Dev Container CLI output is streamed into one stable `Dev Container ()` Output channel per source workspace, which is reused across connection attempts. ## Change policy diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/browserTunnelAgentHostService.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/browserTunnelAgentHostService.ts index 13d6932de17a7c..c1280af507cbc8 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/browserTunnelAgentHostService.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/browserTunnelAgentHostService.ts @@ -5,13 +5,13 @@ import { Emitter, Event } from '../../../../../base/common/event.js'; import { Disposable } from '../../../../../base/common/lifecycle.js'; +import { derived, IObservable, observableSignalFromEvent } from '../../../../../base/common/observable.js'; import { AgentHostProtocolClient } from '../../../../../platform/agentHost/browser/agentHostProtocolClient.js'; import { agentsWindowAgentHostClientInfo } from '../../../../../platform/agentHost/common/agentHostClientInfo.js'; import { AgentHostClientConnectionKind } from '../../../../../platform/agentHost/common/agentHostTelemetry.js'; import { IRemoteAgentHostLocationPreferenceService } from '../../../../../platform/agentHost/common/remoteAgentHostLocationPreference.js'; -import { IRemoteAgentHostService, RemoteAgentHostConnectionStatus, RemoteAgentHostEntryType, RemoteAgentHostsEnabledSettingId } from '../../../../../platform/agentHost/common/remoteAgentHostService.js'; +import { IRemoteAgentHostService, RemoteAgentHostEntryType, RemoteAgentHostsEnabledSettingId, getEntryAddress, type IRemoteAgentHostConnectOptions, type IRemoteAgentHostConnectionFactory, type IRemoteAgentHostCreatedConnection, type IRemoteAgentHostEntry } from '../../../../../platform/agentHost/common/remoteAgentHostService.js'; import { ReconnectingTransport, type IEstablishedTransport } from '../../../../../platform/agentHost/common/reconnectingTransport.js'; -import { PROTOCOL_VERSION } from '../../../../../platform/agentHost/common/state/protocol/version/registry.js'; import type { AhpServerNotification, JsonRpcResponse, ProtocolMessage } from '../../../../../platform/agentHost/common/state/sessionProtocol.js'; import { NonReconnectableTransportError, type IProtocolTransport } from '../../../../../platform/agentHost/common/state/sessionTransport.js'; import { @@ -53,6 +53,82 @@ import { MALFORMED_FRAMES_FORCE_CLOSE_THRESHOLD, MALFORMED_FRAMES_LOG_CAP } from const LOG_PREFIX = '[BrowserTunnelAgentHost]'; +class BrowserTunnelConnectionFactory extends Disposable implements IRemoteAgentHostConnectionFactory { + readonly kind = RemoteAgentHostEntryType.Tunnel; + readonly entries: IObservable; + + private readonly _onDidStageTunnel = this._register(new Emitter()); + private readonly _stagedAuthProviders = new Map(); + /** + * Initiation mode for a staged tunnel, consumed by the first + * {@link createConnection} for that address. Staging publishes the entry + * synchronously, so the service's reconciliation can begin dialing before + * the caller's explicit `reconnect` runs — and that dial would otherwise be + * treated as background, suppressing interactive auth and gateway + * selection for the user's own first connect. + */ + private readonly _stagedUserInitiated = new Map(); + private readonly _onDidStageTunnelSignal = observableSignalFromEvent(this, this._onDidStageTunnel.event); + + constructor( + private readonly _storage: TunnelAgentHostStorage, + private readonly _createConnection: (entry: IRemoteAgentHostEntry, authProvider: 'github' | 'microsoft' | undefined, options: IRemoteAgentHostConnectOptions) => Promise, + ) { + super(); + this.entries = derived(this, reader => { + this._onDidStageTunnelSignal.read(reader); + const autoConnectSuppressedTunnels = this._storage.autoConnectSuppressedTunnels.read(reader); + return this._storage.cachedTunnels.read(reader) + .filter(tunnel => !autoConnectSuppressedTunnels.includes(tunnel.tunnelId)) + .map(tunnel => this._entryForTunnel(tunnel, tunnel.authProvider)); + }); + } + + stageTunnel(tunnel: ITunnelInfo, authProvider?: 'github' | 'microsoft', userInitiated = true): IRemoteAgentHostEntry { + const address = `${TUNNEL_ADDRESS_PREFIX}${tunnel.tunnelId}`; + this._stagedAuthProviders.set(address, authProvider); + this._stagedUserInitiated.set(address, userInitiated); + this._storage.cacheTunnel({ tunnelId: tunnel.tunnelId, clusterId: tunnel.clusterId, name: tunnel.name, protocolVersion: tunnel.protocolVersion, authProvider }); + this._onDidStageTunnel.fire(); + return this._entryForTunnel(tunnel, authProvider); + } + + unstageTunnel(address: string): void { + this._stagedUserInitiated.delete(address); + if (this._stagedAuthProviders.delete(address)) { + this._onDidStageTunnel.fire(); + } + } + + createConnection(entry: IRemoteAgentHostEntry, options: IRemoteAgentHostConnectOptions): Promise { + if (entry.connection.type !== RemoteAgentHostEntryType.Tunnel) { + throw new Error(`Tunnel factory cannot create a ${entry.connection.type} connection.`); + } + const address = getEntryAddress(entry); + const stagedUserInitiated = this._stagedUserInitiated.get(address); + // Consume it: only the connect this staging was for is user-initiated, + // and a later automatic reconnect must not prompt. + this._stagedUserInitiated.delete(address); + const connectOptions = stagedUserInitiated === undefined + ? options + : { ...options, userInitiated: stagedUserInitiated }; + return this._createConnection(entry, this._stagedAuthProviders.has(address) ? this._stagedAuthProviders.get(address) : entry.connection.authProvider, connectOptions); + } + + private _entryForTunnel(tunnel: Pick, authProvider?: 'github' | 'microsoft'): IRemoteAgentHostEntry { + return { + name: tunnel.name, + connection: { + type: RemoteAgentHostEntryType.Tunnel, + tunnelId: tunnel.tunnelId, + clusterId: tunnel.clusterId, + label: tunnel.name, + authProvider, + }, + }; + } +} + /** Creates relay clients directly from the lazily-loaded Dev Tunnels browser SDK. */ export class BrowserTunnelRelayClientFactory implements ITunnelRelayClientFactory { constructor( @@ -142,6 +218,7 @@ export class BrowserTunnelAgentHostService extends Disposable implements ITunnel declare readonly _serviceBrand: undefined; private readonly _storage: TunnelAgentHostStorage; + private readonly _connectionFactory: BrowserTunnelConnectionFactory; readonly onDidChangeTunnels: Event; private readonly _connector: ITunnelAgentHostConnector; @@ -164,6 +241,11 @@ export class BrowserTunnelAgentHostService extends Disposable implements ITunnel super(); this._storage = this._register(new TunnelAgentHostStorage(this._storageService)); this.onDidChangeTunnels = this._storage.onDidChangeTunnels; + this._connectionFactory = this._register(new BrowserTunnelConnectionFactory( + this._storage, + (entry, authProvider, connectOptions) => this._createConnection(entry, authProvider, connectOptions), + )); + this._register(this._remoteAgentHostService.registerConnectionFactory(this._connectionFactory)); const load = options.loadDevTunnelsWeb ?? loadDevTunnelsWeb; this._loadDevTunnelsWeb = load; this._connector = options.connector ?? this._register(new TunnelAgentHostConnector( @@ -213,25 +295,57 @@ export class BrowserTunnelAgentHostService extends Disposable implements ITunnel throw new Error('Remote agent host connections are not enabled.'); } + const entry = this._connectionFactory.stageTunnel(tunnel, authProvider, options?.userInitiated ?? true); + const address = getEntryAddress(entry); + this._remoteAgentHostService.reconnect(address, options?.userInitiated ?? true); + await this._remoteAgentHostService.waitForConnection(address); + } + + private async _createConnection(entry: IRemoteAgentHostEntry, authProvider: 'github' | 'microsoft' | undefined, options: IRemoteAgentHostConnectOptions): Promise { + if (entry.connection.type !== RemoteAgentHostEntryType.Tunnel) { + throw new Error(`Tunnel factory cannot create a ${entry.connection.type} connection.`); + } + // Bind the narrowed connection before the closure: TypeScript does not + // carry the discriminant narrowing into the `find` callback below. + const connection = entry.connection; + const cachedTunnel = this._storage.getCachedTunnels().find(cached => cached.tunnelId === connection.tunnelId); + const tunnel: ITunnelInfo = { + tunnelId: connection.tunnelId, + clusterId: connection.clusterId, + name: connection.label ?? entry.name, + tags: [], + // Legacy cache fallback, not a real capability claim. + protocolVersion: cachedTunnel?.protocolVersion ?? TUNNEL_MIN_PROTOCOL_VERSION, + hostConnectionCount: 0, + }; const auth = authProvider - ? await this._getTokenForProvider(authProvider, false) - : await this._getToken(false); + ? await this._getTokenForProvider(authProvider, !options.userInitiated) + : await this._getToken(!options.userInitiated); if (!auth) { - throw new Error('No authentication available'); + throw new NonReconnectableTransportError('No cached authentication available to connect the tunnel.'); } - const result = await connectThroughTunnelGateway( - this._connector, - this._resolveGatewaySelection, - this._locationPreferenceService, - this._dialogService, - this._productService.nameShort, - auth, - tunnel, - options?.userInitiated ?? true, - ); - if (!result) { - return; + let result: ITunnelConnectResult; + try { + const connected = await connectThroughTunnelGateway( + this._connector, + this._resolveGatewaySelection, + this._locationPreferenceService, + this._dialogService, + this._productService.nameShort, + auth, + tunnel, + options.userInitiated, + ); + if (!connected) { + throw new NonReconnectableTransportError('Tunnel agent host selection requires user interaction.'); + } + result = connected; + } catch (error) { + if (isTunnelNotFoundError(error)) { + throw new NonReconnectableTransportError(error.message); + } + throw error; } let useSeedConnection = true; const establish = async (): Promise => { @@ -282,47 +396,11 @@ export class BrowserTunnelAgentHostService extends Disposable implements ITunnel LOG_PREFIX, AgentHostClientConnectionKind.DevTunnel, ); - const protocolClient = this._instantiationService.createInstance( - AgentHostProtocolClient, result.address, transportFactory, { clientInfo: agentsWindowAgentHostClientInfo }, - ); - - let status: RemoteAgentHostConnectionStatus = RemoteAgentHostConnectionStatus.connected; - let connectError: unknown; - try { - await protocolClient.connect(); - this._logService.info(`${LOG_PREFIX} Protocol handshake completed with ${result.address}`); - } catch (error) { - const incompatible = RemoteAgentHostConnectionStatus.fromConnectError(error, [PROTOCOL_VERSION]); - if (!RemoteAgentHostConnectionStatus.isIncompatible(incompatible)) { - protocolClient.dispose(); - throw error; - } - status = incompatible; - connectError = error; - this._logService.warn(`${LOG_PREFIX} Incompatible with ${result.address}: ${incompatible.message}`); - } - - this.cacheTunnel(tunnel, auth.provider); - try { - await this._remoteAgentHostService.addManagedConnection({ - name: result.name, - connectionToken: result.connectionToken, - connection: { - type: RemoteAgentHostEntryType.Tunnel, - tunnelId: tunnel.tunnelId, - clusterId: tunnel.clusterId, - label: tunnel.name, - authProvider: auth.provider, - }, - }, protocolClient, undefined, status); - } catch (error) { - protocolClient.dispose(); - throw error; - } - - if (connectError) { - throw connectError; - } + return { + connection: this._instantiationService.createInstance( + AgentHostProtocolClient, result.address, transportFactory, { clientInfo: agentsWindowAgentHostClientInfo }, + ), + }; } readonly canDeleteTunnels = true; @@ -338,8 +416,8 @@ export class BrowserTunnelAgentHostService extends Disposable implements ITunnel } async disconnect(address: string): Promise { + this._connectionFactory.unstageTunnel(address); await this._remoteAgentHostService.removeRemoteAgentHost(address); - this._storage.notifyTunnelsChanged(); } async getAuthProvider(options?: { silent?: boolean }): Promise<'github' | 'microsoft' | undefined> { @@ -355,14 +433,28 @@ export class BrowserTunnelAgentHostService extends Disposable implements ITunnel tunnelId: tunnel.tunnelId, clusterId: tunnel.clusterId, name: tunnel.name, + protocolVersion: tunnel.protocolVersion, authProvider, }); } removeCachedTunnel(tunnelId: string): void { + this._connectionFactory.unstageTunnel(`${TUNNEL_ADDRESS_PREFIX}${tunnelId}`); this._storage.removeCachedTunnel(tunnelId); } + isTunnelDismissed(tunnelId: string): boolean { + return this._storage.isTunnelDismissed(tunnelId); + } + + dismissTunnel(tunnelId: string): void { + this._storage.dismissTunnel(tunnelId); + } + + clearTunnelDismissal(tunnelId: string): void { + this._storage.clearTunnelDismissal(tunnelId); + } + isAutoConnectSuppressed(tunnelId: string): boolean { return this._storage.isAutoConnectSuppressed(tunnelId); } diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxAgentHostContribution.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxAgentHostContribution.ts index 65c0157be050c2..d5917f152ef94c 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxAgentHostContribution.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxAgentHostContribution.ts @@ -329,8 +329,8 @@ export class CloudSandboxAgentHostContribution extends Disposable implements IWo /** * Remove the connection (and its credential refresher) for an environment while keeping the * provider and its cached sessions visible in a disconnected state. Disposing the protocol - * client stops the soft-reconnect loop; the {@link CloudSandboxAgentHostService} prunes the - * refresher via `onDidChangeConnections`. + * client stops its soft-reconnect loop and disposes the credential refresher owned by its + * connection factory. */ private async _disconnectEnvironment(address: string): Promise { try { @@ -650,9 +650,6 @@ export class CloudSandboxAgentHostContribution extends Disposable implements IWo void this._disconnectEnvironment(address); throw new CancellationError(); } - // `onDidChangeConnections` fires from addManagedConnection and wires the - // provider; call _wireConnections directly too in case it already fired. - this._wireConnections(); return result; } finally { this._pendingConnects.delete(address); diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxAgentHostService.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxAgentHostService.ts index 6832d81a9cc839..f065cb1fdf0a0f 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxAgentHostService.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxAgentHostService.ts @@ -5,8 +5,9 @@ import { CancellationError, isCancellationError } from '../../../../../base/common/errors.js'; import { CancellationToken } from '../../../../../base/common/cancellation.js'; -import { Disposable, DisposableMap, DisposableStore } from '../../../../../base/common/lifecycle.js'; -import { timeout } from '../../../../../base/common/async.js'; +import { Disposable, DisposableStore, MutableDisposable } from '../../../../../base/common/lifecycle.js'; +import { IObservable, observableValue } from '../../../../../base/common/observable.js'; +import { raceCancellationError, timeout } from '../../../../../base/common/async.js'; import { IProtocolTransport } from '../../../../../platform/agentHost/common/state/sessionTransport.js'; import { AgentHostProtocolClient } from '../../../../../platform/agentHost/browser/agentHostProtocolClient.js'; import { editorWindowAgentHostClientInfo } from '../../../../../platform/agentHost/common/agentHostClientInfo.js'; @@ -24,8 +25,7 @@ import { type CloudSandboxConnectResult, type ICloudSandboxClientToken, } from '../../../../../platform/agentHost/common/cloudSandboxAgentHost.js'; -import { IRemoteAgentHostService, RemoteAgentHostConnectionStatus, RemoteAgentHostEntryType, RemoteAgentHostsEnabledSettingId } from '../../../../../platform/agentHost/common/remoteAgentHostService.js'; -import { PROTOCOL_VERSION } from '../../../../../platform/agentHost/common/state/protocol/version/registry.js'; +import { getEntryAddress, IRemoteAgentHostService, RemoteAgentHostEntryType, RemoteAgentHostsEnabledSettingId, type IRemoteAgentHostConnectOptions, type IRemoteAgentHostConnectionFactory, type IRemoteAgentHostCreatedConnection, type IRemoteAgentHostEntry } from '../../../../../platform/agentHost/common/remoteAgentHostService.js'; import { IConfigurationService } from '../../../../../platform/configuration/common/configuration.js'; import { IEnvironmentService } from '../../../../../platform/environment/common/environment.js'; import { IInstantiationService } from '../../../../../platform/instantiation/common/instantiation.js'; @@ -46,23 +46,142 @@ export const MAX_SEALED_TOKEN_RETRIES = 12; /** Delay between `/connect` re-mints while waiting for complete credentials. */ const SEALED_TOKEN_RETRY_DELAY_MS = 5_000; -/** - * Renderer-side coordinator for Copilot cloud sandbox connections. - * - * Mirrors {@link WebTunnelAgentHostService}: establishes a connection - * out-of-band (mint creds → open a {@link WebPubSubRelayTransport} → drive the - * AHP handshake) and hands the pre-connected {@link AgentHostProtocolClient} - * to {@link IRemoteAgentHostService.addManagedConnection}, so the existing - * remote-agent-host contribution surfaces it as a native, interactive session. - */ +interface IStagedCloudSandboxConnection { + readonly entry: IRemoteAgentHostEntry; + readonly options: ICloudSandboxConnectOptions; + readonly creds: ICloudSandboxCreds; + readonly clientId: string; +} + +/** Builds cloud sandbox protocol clients from credentials staged by the caller. */ +class CloudSandboxConnectionFactory extends Disposable implements IRemoteAgentHostConnectionFactory { + readonly kind = RemoteAgentHostEntryType.CloudSandbox; + readonly entries: IObservable; + + private readonly _stagedConnections = new Map(); + private readonly _entries = observableValue(this, []); + + constructor( + private readonly _instantiationService: IInstantiationService, + private readonly _configurationService: IConfigurationService, + private readonly _environmentService: IEnvironmentService, + ) { + super(); + this.entries = this._entries; + // Staging is cleared only by an explicit `unstageConfiguration`, never by + // observing the connection disappear. The service withdraws an entry + // before arming a retry, so treating that as removal would delete the + // staged credentials the retry needs and leave `_scheduleReconnect` with + // nothing configured — silently turning every scheduled retry into one + // single attempt. `_establish` already unstages on the paths that really + // are terminal. + } + + stageConfiguration(options: ICloudSandboxConnectOptions, clientToken: ICloudSandboxClientToken): IRemoteAgentHostEntry { + const address = cloudSandboxAddress(options.environmentId); + const entry: IRemoteAgentHostEntry = { + name: options.name, + connection: { + type: RemoteAgentHostEntryType.CloudSandbox, + address, + environmentId: options.environmentId, + sessionId: options.sessionId, + }, + }; + this._stagedConnections.set(address, { + entry, + options, + creds: { token: clientToken }, + clientId: clientToken.client_id, + }); + this._updateEntries(); + return entry; + } + + unstageConfiguration(address: string): void { + this._stagedConnections.delete(address); + this._updateEntries(); + } + + getSealedGitHubToken(environmentId: string): string | undefined { + return this._stagedConnections.get(cloudSandboxAddress(environmentId))?.creds.token.encrypted_github_token; + } + + async createConnection(entry: IRemoteAgentHostEntry, _options: IRemoteAgentHostConnectOptions): Promise { + if (entry.connection.type !== RemoteAgentHostEntryType.CloudSandbox) { + throw new Error(`Cloud sandbox factory cannot create a ${entry.connection.type} connection.`); + } + const address = getEntryAddress(entry); + const staged = this._stagedConnections.get(address); + if (!staged) { + throw new Error(`No cloud sandbox connection is staged for ${address}.`); + } + + const ahpLoggingEnabled = !!this._configurationService.getValue(AgentHostAhpJsonlLoggingSettingId); + const transportFactory = (): IProtocolTransport => new WebPubSubRelayTransport({ + url: buildWpsUrl(staged.creds.token), + toHostGroup: staged.creds.token.groups.to_host, + joinGroups: [staged.creds.token.groups.broadcast, staged.creds.token.groups.to_client], + groupValidation: { expected: { cid: staged.creds.token.client_id } }, + ahpLogger: ahpLoggingEnabled + ? this._instantiationService.createInstance(AhpJsonlLogger, { + logsHome: this._environmentService.logsHome, + connectionId: staged.clientId, + transport: 'webpubsub', + }) + : undefined, + }); + const client = this._instantiationService.createInstance( + AgentHostProtocolClient, + address, + transportFactory, + { + clientId: staged.clientId, + clientInfo: editorWindowAgentHostClientInfo, + resolveInitialAuthentication: () => this._resolveInitialAuthentication(address), + }, + ); + const store = new DisposableStore(); + const refresher = store.add(new MutableDisposable()); + store.add(client.onDidChangeConnectionState(state => { + if (state === 'connected' && !refresher.value) { + refresher.value = this._instantiationService.createInstance( + CloudSandboxCredentialRefresher, + address, + { environmentId: staged.options.environmentId, sessionId: staged.options.sessionId }, + staged.clientId, + staged.creds, + ); + } + })); + return { connection: client, transportDisposable: store }; + } + + private async _resolveInitialAuthentication(address: string): Promise<{ readonly resource: string; readonly token: string } | undefined> { + // Throw rather than returning `undefined`: an unusable token must fail + // the connection, not produce one that reports connected and then fails + // every authenticated request. The protocol client classifies this as an + // initial-authentication failure and surfaces it as incompatible. + const sealedToken = this._stagedConnections.get(address)?.creds.token.encrypted_github_token; + if (!sealedToken) { + throw new Error(`Mission Control returned no sealed token for ${address}; the session cannot make authenticated requests.`); + } + if (!isCloudSandboxSealedToken(sealedToken)) { + throw new Error(`Refusing to forward a non-sealed token to ${address}; Mission Control did not return a copilot-sealed envelope.`); + } + return { resource: GITHUB_COPILOT_PROTECTED_RESOURCE.resource, token: sealedToken }; + } + + private _updateEntries(): void { + this._entries.set([...this._stagedConnections.values()].map(connection => connection.entry), undefined); + } +} + +/** Renderer-side coordinator for Copilot cloud sandbox connections. */ export class CloudSandboxAgentHostService extends Disposable implements ICloudSandboxAgentHostService { declare readonly _serviceBrand: undefined; - /** Credential-refresh scheduler per connection address, disposed when the connection is gone. */ - private readonly _managed = this._register(new DisposableMap()); - - /** Current Web PubSub credentials per connection address, including the sealed GitHub token. */ - private readonly _creds = new Map(); + private readonly _connectionFactory: CloudSandboxConnectionFactory; /** Overridable so tests can exercise the re-mint loop without waiting on real delays. */ protected readonly sealedTokenRetryDelayMs: number = SEALED_TOKEN_RETRY_DELAY_MS; @@ -76,19 +195,16 @@ export class CloudSandboxAgentHostService extends Disposable implements ICloudSa @ILogService private readonly _logService: ILogService, ) { super(); - // Stop refreshing credentials once a connection is gone. - this._register(this._remoteAgentHostService.onDidChangeConnections(() => { - for (const address of [...this._managed.keys()]) { - if (!this._remoteAgentHostService.connections.some(c => c.address === address)) { - this._managed.deleteAndDispose(address); - this._creds.delete(address); - } - } - })); + this._connectionFactory = this._register(new CloudSandboxConnectionFactory( + this._instantiationService, + this._configurationService, + this._environmentService, + )); + this._register(this._remoteAgentHostService.registerConnectionFactory(this._connectionFactory)); } getSealedGitHubToken(environmentId: string): string | undefined { - return this._creds.get(cloudSandboxAddress(environmentId))?.token.encrypted_github_token; + return this._connectionFactory.getSealedGitHubToken(environmentId); } async connect(options: ICloudSandboxConnectOptions, token: CancellationToken): Promise { @@ -122,107 +238,29 @@ export class CloudSandboxAgentHostService extends Disposable implements ICloudSa } /** - * Open the relay with an already-minted token, drive the AHP handshake, and register the - * connection. + * Stage already-minted credentials and wait for the remote connection service to handshake it. */ protected async _establish(options: ICloudSandboxConnectOptions, address: string, clientToken: ICloudSandboxClientToken, token: CancellationToken): Promise { - // Mutable holder read by the transport factory: the protocol client re-invokes the factory to - // soft-reconnect, picking up whatever credentials the refresh scheduler last wrote. - const creds: ICloudSandboxCreds = { token: clientToken }; - // Three per-client relay lanes: publish to `to_host`; receive replies on `to_client` and - // unsolicited session state on `broadcast`. `groupValidation` drops inbound frames whose - // group name doesn't carry our own client id. - // Each soft reconnect gets a transport-owned logger keyed by connection id. - const ahpLoggingEnabled = !!this._configurationService.getValue(AgentHostAhpJsonlLoggingSettingId); - const transportFactory = (): IProtocolTransport => new WebPubSubRelayTransport({ - url: buildWpsUrl(creds.token), - toHostGroup: creds.token.groups.to_host, - joinGroups: [creds.token.groups.broadcast, creds.token.groups.to_client], - groupValidation: { expected: { cid: creds.token.client_id } }, - ahpLogger: ahpLoggingEnabled - ? this._instantiationService.createInstance(AhpJsonlLogger, { - logsHome: this._environmentService.logsHome, - connectionId: clientToken.client_id, - transport: 'webpubsub', - }) - : undefined, - }); - - // Mission Control mints the client id and binds the relay lane to it, so the AHP identity - // must match or the host rejects requests on that lane. - const protocolClient = this._instantiationService.createInstance( - AgentHostProtocolClient, address, transportFactory, { clientId: clientToken.client_id, clientInfo: editorWindowAgentHostClientInfo }, - ); - - let status: RemoteAgentHostConnectionStatus = RemoteAgentHostConnectionStatus.connected; - let connectError: unknown; + if (token.isCancellationRequested) { + throw new CancellationError(); + } + this._connectionFactory.stageConfiguration(options, clientToken); try { - await protocolClient.connect(); - this._logService.info(`${LOG_PREFIX} Protocol handshake completed with ${address}`); - } catch (err) { - const incompatible = RemoteAgentHostConnectionStatus.fromConnectError(err, [PROTOCOL_VERSION]); - if (!RemoteAgentHostConnectionStatus.isIncompatible(incompatible)) { - protocolClient.dispose(); - throw err; + if (token.isCancellationRequested) { + throw new CancellationError(); } - this._logService.warn(`${LOG_PREFIX} Incompatible with ${address}: ${incompatible.message}`); - status = incompatible; - connectError = err; - } - - // Push the sealed GitHub token so the host can call api.github.com on the agent's behalf. - // Only a `copilot-sealed.v1.` envelope is forwarded; a plaintext bearer is refused. - if (!connectError && clientToken.encrypted_github_token) { - if (!isCloudSandboxSealedToken(clientToken.encrypted_github_token)) { - this._logService.error(`${LOG_PREFIX} Refusing to forward a non-sealed token to ${address}; Mission Control did not return a copilot-sealed envelope.`); - } else { - try { - await protocolClient.authenticate({ - resource: GITHUB_COPILOT_PROTECTED_RESOURCE.resource, - token: clientToken.encrypted_github_token, - }); - } catch (err) { - this._logService.warn(`${LOG_PREFIX} Sealed-token authenticate failed for ${address}`, err); - } + this._remoteAgentHostService.reconnect(address, true); + await raceCancellationError(this._remoteAgentHostService.waitForConnection(address), token); + this._logService.info(`${LOG_PREFIX} Protocol handshake completed with ${address}`); + return address; + } catch (error) { + const connectionStillRegistered = this._remoteAgentHostService.connections.some(connection => connection.address === address); + if (token.isCancellationRequested || !connectionStillRegistered) { + this._connectionFactory.unstageConfiguration(address); + await this._remoteAgentHostService.removeRemoteAgentHost(address); } - } else if (!connectError) { - // Without an envelope every later request answers `-32007 AuthRequired`. - this._logService.error(`${LOG_PREFIX} Mission Control returned no sealed token for ${address}; this session will not be able to make authenticated requests.`); - } - - try { - await this._remoteAgentHostService.addManagedConnection({ - name: options.name, - connection: { - type: RemoteAgentHostEntryType.CloudSandbox, - address, - environmentId: options.environmentId, - sessionId: options.sessionId, - }, - }, protocolClient, undefined, status); - } catch (err) { - protocolClient.dispose(); - this._logService.error(`${LOG_PREFIX} addManagedConnection failed`, err); - throw err; - } - - // Keep credentials fresh for the life of the connection so reconnects have a valid token. - const store = new DisposableStore(); - store.add(this._instantiationService.createInstance( - CloudSandboxCredentialRefresher, - address, - { environmentId: options.environmentId, sessionId: options.sessionId }, - clientToken.client_id, - creds, - )); - this._managed.set(address, store); - // Expose the sealed GitHub token so the AHP `authenticate` pass can present it to the host. - this._creds.set(address, creds); - - if (connectError) { - throw connectError; + throw error; } - return address; } /** Mint client creds, retrying (bounded) while the environment is waking. */ diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/devContainerAgentHostService.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/devContainerAgentHostService.ts index 46186e44f24a82..13008a55db48fe 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/devContainerAgentHostService.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/devContainerAgentHostService.ts @@ -8,14 +8,18 @@ import { CancellationError } from '../../../../../base/common/errors.js'; import { raceCancellationError, raceTimeout } from '../../../../../base/common/async.js'; import { Event } from '../../../../../base/common/event.js'; import { getComparisonKey } from '../../../../../base/common/resources.js'; +import { StringSHA1 } from '../../../../../base/common/hash.js'; import { Disposable, DisposableMap, DisposableStore, IDisposable, toDisposable } from '../../../../../base/common/lifecycle.js'; +import { IObservable, observableValue } from '../../../../../base/common/observable.js'; import { URI } from '../../../../../base/common/uri.js'; import { localize } from '../../../../../nls.js'; import { AGENT_HOST_SCHEME, agentHostAuthority } from '../../../../../platform/agentHost/common/agentHostUri.js'; -import { IRemoteAgentHostEntry, IRemoteAgentHostService, RemoteAgentHostConnectionStatus, RemoteAgentHostEntryType } from '../../../../../platform/agentHost/common/remoteAgentHostService.js'; +import { agentsWindowAgentHostClientInfo } from '../../../../../platform/agentHost/common/agentHostClientInfo.js'; +import { AgentHostProtocolClient } from '../../../../../platform/agentHost/browser/agentHostProtocolClient.js'; +import { getEntryAddress, getEntryTypeConfig, IRemoteAgentHostEntry, IRemoteAgentHostService, RemoteAgentHostConnectionStatus, RemoteAgentHostEntryType, type IRemoteAgentHostConnectOptions, type IRemoteAgentHostConnectionFactory, type IRemoteAgentHostCreatedConnection } from '../../../../../platform/agentHost/common/remoteAgentHostService.js'; import { IInstantiationService } from '../../../../../platform/instantiation/common/instantiation.js'; import { InstantiationType, registerSingleton } from '../../../../../platform/instantiation/common/extensions.js'; -import { IDevContainerAgentHostConnector, IDevContainerAgentHostService, IDevContainerAgentHostTarget } from '../../../../common/devContainerAgentHostService.js'; +import { IDevContainerAgentHostConnection, IDevContainerAgentHostConnector, IDevContainerAgentHostService, IDevContainerAgentHostTarget } from '../../../../common/devContainerAgentHostService.js'; import { ISessionsProvidersService } from '../../../../services/sessions/browser/sessionsProvidersService.js'; import { RemoteAgentHostSessionsProvider } from './remoteAgentHostSessionsProvider.js'; @@ -31,6 +35,106 @@ interface IPendingDevContainerAgentHost { readonly tokenSource: CancellationTokenSource; } +interface IStagedDevContainerConnection { + readonly entry: IRemoteAgentHostEntry; + readonly connector: IDevContainerAgentHostConnector; + readonly workspaceUri: URI; + initialConnection: IDevContainerAgentHostConnection | undefined; +} + +function devContainerAddress(workspaceUri: URI): string { + const sha = new StringSHA1(); + sha.update(getComparisonKey(workspaceUri)); + return `devcontainer:${sha.digest()}`; +} + +/** Builds Dev Container protocol clients from a staged workspace transport. */ +class DevContainerConnectionFactory extends Disposable implements IRemoteAgentHostConnectionFactory { + readonly kind = RemoteAgentHostEntryType.DevContainer; + readonly entries: IObservable; + + private readonly _stagedConnections = new Map(); + private readonly _entries = observableValue(this, []); + + constructor( + private readonly _instantiationService: IInstantiationService, + ) { + super(); + this.entries = this._entries; + // Staging is cleared only by an explicit `unstageConnection`, never by + // observing the connection disappear. The service withdraws an entry + // before arming a retry, so treating that as removal would delete the + // staged connector the retry needs and leave `_scheduleReconnect` with + // nothing configured — silently turning every scheduled retry into one + // single attempt. + } + + stageConnection(connector: IDevContainerAgentHostConnector, workspaceUri: URI, connection: IDevContainerAgentHostConnection): IRemoteAgentHostEntry { + const entry: IRemoteAgentHostEntry = { + name: connection.name, + connection: { + type: RemoteAgentHostEntryType.DevContainer, + address: connection.address, + hostPath: workspaceUri.fsPath, + }, + }; + this._stagedConnections.set(connection.address, { entry, connector, workspaceUri, initialConnection: connection }); + this._updateEntries(); + return entry; + } + + unstageConnection(address: string): void { + const staged = this._stagedConnections.get(address); + this._stagedConnections.delete(address); + staged?.initialConnection?.transportDisposable?.dispose(); + this._updateEntries(); + } + + async createConnection(entry: IRemoteAgentHostEntry, _options: IRemoteAgentHostConnectOptions): Promise { + if (entry.connection.type !== RemoteAgentHostEntryType.DevContainer) { + throw new Error(`Dev Container factory cannot create a ${entry.connection.type} connection.`); + } + const staged = this._stagedConnections.get(entry.connection.address); + if (!staged) { + throw new Error(`No Dev Container connection is staged for ${entry.connection.address}.`); + } + + const connection = staged.initialConnection ?? await staged.connector.createConnection( + staged.workspaceUri, + entry.connection.address, + CancellationToken.None, + ); + try { + const authority = agentHostAuthority(entry.connection.address); + if (connection.workspaceUri.scheme !== AGENT_HOST_SCHEME || connection.workspaceUri.authority !== authority) { + throw new Error(localize('devContainerAgentHost.invalidWorkspaceUri', "Dev Container workspace URI must use the '{0}' scheme and '{1}' authority.", AGENT_HOST_SCHEME, authority)); + } + + const client = this._instantiationService.createInstance( + AgentHostProtocolClient, + entry.connection.address, + connection.transportFactory, + { clientInfo: agentsWindowAgentHostClientInfo, reconnectPolicy: getEntryTypeConfig(RemoteAgentHostEntryType.DevContainer).reconnect }, + ); + staged.initialConnection = undefined; + return { + connection: client, + transportDisposable: connection.transportDisposable, + }; + } catch (error) { + if (staged.initialConnection === connection) { + staged.initialConnection = undefined; + } + connection.transportDisposable?.dispose(); + throw error; + } + } + + private _updateEntries(): void { + this._entries.set([...this._stagedConnections.values()].map(connection => connection.entry), undefined); + } +} + /** Registers Dev Container Agent Hosts as dynamic remote Sessions providers. */ export class DevContainerAgentHostService extends Disposable implements IDevContainerAgentHostService { declare readonly _serviceBrand: undefined; @@ -38,6 +142,7 @@ export class DevContainerAgentHostService extends Disposable implements IDevCont private readonly _providerStores = this._register(new DisposableMap()); private readonly _activeConnections = new Map(); private readonly _pendingConnections = new Map(); + private readonly _connectionFactory: DevContainerConnectionFactory; private _connector: IDevContainerAgentHostConnector | undefined; constructor( @@ -46,6 +151,8 @@ export class DevContainerAgentHostService extends Disposable implements IDevCont @ISessionsProvidersService private readonly _sessionsProvidersService: ISessionsProvidersService, ) { super(); + this._connectionFactory = this._register(new DevContainerConnectionFactory(this._instantiationService)); + this._register(this._remoteAgentHostService.registerConnectionFactory(this._connectionFactory)); this._register(this._remoteAgentHostService.onDidChangeConnections(() => this._reconcileConnections())); } @@ -114,22 +221,14 @@ export class DevContainerAgentHostService extends Disposable implements IDevCont if (token.isCancellationRequested) { throw new CancellationError(); } - const connected = await connector.connect(workspaceUri, token); + const connected = await connector.createConnection(workspaceUri, devContainerAddress(workspaceUri), token); if (token.isCancellationRequested) { connected.transportDisposable?.dispose(); - connected.connection.dispose(); throw new CancellationError(); } - const authority = agentHostAuthority(connected.address); - if (connected.workspaceUri.scheme !== AGENT_HOST_SCHEME || connected.workspaceUri.authority !== authority) { - connected.transportDisposable?.dispose(); - connected.connection.dispose(); - throw new Error(localize('devContainerAgentHost.invalidWorkspaceUri', "Dev Container workspace URI must use the '{0}' scheme and '{1}' authority.", AGENT_HOST_SCHEME, authority)); - } - const providerStore = new DisposableStore(); - let connectionOwnedByRemoteService = false; + let stagedAddress: string | undefined; try { const provider = providerStore.add(this._createProvider({ address: connected.address, @@ -138,33 +237,38 @@ export class DevContainerAgentHostService extends Disposable implements IDevCont })); providerStore.add(this._sessionsProvidersService.registerProvider(provider)); - const entry: IRemoteAgentHostEntry = { - name: connected.name, - connection: { - type: RemoteAgentHostEntryType.DevContainer, - address: connected.address, - hostPath: workspaceUri.fsPath, - }, - }; - const connectionInfo = await this._remoteAgentHostService.addManagedConnection(entry, connected.connection, connected.transportDisposable); - connectionOwnedByRemoteService = true; - provider.setConnection(connected.connection, connected.defaultDirectory ?? connectionInfo.defaultDirectory); + const entry = this._connectionFactory.stageConnection(connector, workspaceUri, connected); + const address = getEntryAddress(entry); + stagedAddress = address; + if (token.isCancellationRequested) { + throw new CancellationError(); + } + this._remoteAgentHostService.reconnect(address, true); + const connectionInfo = await raceCancellationError(this._remoteAgentHostService.waitForConnection(address), token); + const connection = this._remoteAgentHostService.getConnection(connectionInfo.address); + if (!connection) { + throw new Error(localize('devContainerAgentHost.connectionUnavailable', "Dev Container Agent Host connection was not available after connecting.")); + } + provider.setConnection(connection, connected.defaultDirectory ?? connectionInfo.defaultDirectory); provider.setConnectionStatus(connectionInfo.status); await this._waitForSessionTypes(provider, token); const target = { providerId: provider.id, workspaceUri: connected.workspaceUri }; - const active = { address: connectionInfo.address, provider, target, references: 0 }; + const active = { address, provider, target, references: 0 }; providerStore.add(toDisposable(() => this._activeConnections.delete(key))); this._providerStores.set(key, providerStore); this._activeConnections.set(key, active); return active; } catch (error) { providerStore.dispose(); - if (connectionOwnedByRemoteService) { - await this._remoteAgentHostService.removeRemoteAgentHost(connected.address); + if (stagedAddress !== undefined) { + const connectionStillRegistered = this._remoteAgentHostService.connections.some(connection => connection.address === stagedAddress); + if (token.isCancellationRequested || !connectionStillRegistered) { + this._connectionFactory.unstageConnection(stagedAddress); + await this._remoteAgentHostService.removeRemoteAgentHost(stagedAddress); + } } else { connected.transportDisposable?.dispose(); - connected.connection.dispose(); } throw error; } diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/entryDrivenProviderContribution.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/entryDrivenProviderContribution.ts new file mode 100644 index 00000000000000..46493a09cbb975 --- /dev/null +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/entryDrivenProviderContribution.ts @@ -0,0 +1,154 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { type Event } from '../../../../../base/common/event.js'; +import { Disposable, DisposableMap, DisposableStore, toDisposable } from '../../../../../base/common/lifecycle.js'; +import { type IRemoteAgentHostEntry, IRemoteAgentHostService, RemoteAgentHostConnectionStatus, RemoteAgentHostEntryType, RemoteAgentHostsEnabledSettingId, getEntryAddress } from '../../../../../platform/agentHost/common/remoteAgentHostService.js'; +import { type IConfigurationService } from '../../../../../platform/configuration/common/configuration.js'; +import { type IInstantiationService } from '../../../../../platform/instantiation/common/instantiation.js'; +import { type INotificationService } from '../../../../../platform/notification/common/notification.js'; +import { type IAgentHostConnectProgress } from '../../../../common/agentHostSessionsProvider.js'; +import { type ISessionsProvidersService } from '../../../../services/sessions/browser/sessionsProvidersService.js'; +import { RemoteAgentHostSessionsProvider } from './remoteAgentHostSessionsProvider.js'; +import { watchForIncompatibleNotifications } from './remoteHostOptions.js'; + +/** Options supplied by a remote-host kind when creating its sessions provider. */ +export interface IEntryDrivenProviderOptions { + readonly connectOnDemand?: () => Promise; + readonly disconnectOnDemand?: () => Promise; + readonly onDidReportConnectProgress?: Event; + readonly initialStatus?: RemoteAgentHostConnectionStatus; + readonly preferenceKey?: string; +} + +/** + * Shared provider ownership for remote-host kinds whose providers correspond + * directly to remote-host entries. Subclasses own their entry discovery and + * on-demand connection behavior; this class owns only provider lifecycle. + */ +export abstract class EntryDrivenProviderContribution extends Disposable { + + protected readonly _providerStores = this._register(new DisposableMap()); + protected readonly _providerInstances = new Map(); + private readonly _wiredAddresses = new Set(); + + constructor( + protected readonly _remoteAgentHostService: IRemoteAgentHostService, + protected readonly _configurationService: IConfigurationService, + protected readonly _instantiationService: IInstantiationService, + protected readonly _sessionsProvidersService: ISessionsProvidersService, + protected readonly _notificationService: INotificationService, + ) { + super(); + } + + protected get _enabled(): boolean { + return this._configurationService.getValue(RemoteAgentHostsEnabledSettingId); + } + + /** The entry kind this contribution owns when discovering configured entries. */ + protected abstract readonly _entryType: RemoteAgentHostEntryType; + + /** Supplies all entries owned by this contribution. */ + protected _getProviderEntries(): readonly IRemoteAgentHostEntry[] { + if (!this._enabled) { + return []; + } + return this._remoteAgentHostService.configuredEntries.filter(entry => entry.connection.type === this._entryType); + } + + /** Supplies kind-specific on-demand behavior for an entry's provider. */ + protected abstract _getProviderOptions(entry: IRemoteAgentHostEntry): IEntryDrivenProviderOptions; + + /** + * Whether a vanished connection should clear the provider's active + * connection. Defaults to false to preserve existing WSL behavior. + */ + protected get _clearConnectionOnRemoval(): boolean { + return false; + } + + protected _reconcile(): void { + this._reconcileProviders(); + this._wireConnections(); + this._updateConnectionStatuses(); + } + + protected _reconcileProviders(): void { + const entries = this._getProviderEntries(); + const desiredAddresses = new Set(entries.map(entry => getEntryAddress(entry))); + + for (const [address] of this._providerStores) { + if (!desiredAddresses.has(address)) { + this._providerStores.deleteAndDispose(address); + } + } + + for (const entry of entries) { + const address = getEntryAddress(entry); + const existing = this._providerInstances.get(address); + if (existing && existing.label !== (entry.name || address)) { + this._providerStores.deleteAndDispose(address); + } + if (!this._providerStores.has(address)) { + this._createProvider(address, entry.name, this._getProviderOptions(entry)); + } + } + } + + protected _createProvider(address: string, name: string, options: IEntryDrivenProviderOptions): RemoteAgentHostSessionsProvider { + const store = new DisposableStore(); + const provider = this._instantiationService.createInstance( + RemoteAgentHostSessionsProvider, { + address, + name, + connectOnDemand: options.connectOnDemand, + disconnectOnDemand: options.disconnectOnDemand, + onDidReportConnectProgress: options.onDidReportConnectProgress, + preferenceKey: options.preferenceKey, + }); + if (options.initialStatus !== undefined) { + provider.setConnectionStatus(options.initialStatus); + } + store.add(provider); + store.add(this._sessionsProvidersService.registerProvider(provider)); + store.add(watchForIncompatibleNotifications(provider, this._instantiationService, this._notificationService)); + this._providerInstances.set(address, provider); + store.add(toDisposable(() => { + this._providerInstances.delete(address); + this._wiredAddresses.delete(address); + })); + this._providerStores.set(address, store); + return provider; + } + + private _wireConnections(): void { + for (const [address, provider] of this._providerInstances) { + const connectionInfo = this._remoteAgentHostService.connections.find(connection => connection.address === address); + if (connectionInfo && RemoteAgentHostConnectionStatus.isConnected(connectionInfo.status)) { + const connection = this._remoteAgentHostService.getConnection(address); + if (connection) { + provider.setConnection(connection, connectionInfo.defaultDirectory); + if (this._clearConnectionOnRemoval) { + this._wiredAddresses.add(address); + } + } + } else if (this._clearConnectionOnRemoval && !connectionInfo && this._wiredAddresses.delete(address)) { + provider.clearConnection(); + } + } + } + + private _updateConnectionStatuses(): void { + for (const [address, provider] of this._providerInstances) { + const connectionInfo = this._remoteAgentHostService.connections.find(connection => connection.address === address); + if (connectionInfo) { + provider.setConnectionStatus(connectionInfo.status); + } else if (!RemoteAgentHostConnectionStatus.isIncompatible(provider.connectionStatus.get())) { + provider.setConnectionStatus(RemoteAgentHostConnectionStatus.disconnected); + } + } + } +} diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/managedReconnectAgentHostContribution.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/managedReconnectAgentHostContribution.ts index c8db25ed0eb011..e1b98c1873ca07 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/managedReconnectAgentHostContribution.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/managedReconnectAgentHostContribution.ts @@ -4,19 +4,16 @@ *--------------------------------------------------------------------------------------------*/ import { disposableTimeout } from '../../../../../base/common/async.js'; -import { Event } from '../../../../../base/common/event.js'; -import { Disposable, DisposableMap, DisposableStore, MutableDisposable, toDisposable } from '../../../../../base/common/lifecycle.js'; -import { IRemoteAgentHostService, RemoteAgentHostConnectionStatus, RemoteAgentHostsEnabledSettingId } from '../../../../../platform/agentHost/common/remoteAgentHostService.js'; +import { Disposable, DisposableMap, MutableDisposable } from '../../../../../base/common/lifecycle.js'; +import { type IRemoteAgentHostService, RemoteAgentHostConnectionStatus } from '../../../../../platform/agentHost/common/remoteAgentHostService.js'; import { hasExhaustedReconnectAttempts, type IRemoteAgentHostReconnectPolicy } from '../../../../../platform/agentHost/common/reconnectPolicy.js'; import { PROTOCOL_VERSION } from '../../../../../platform/agentHost/common/state/protocol/version/registry.js'; -import { IConfigurationService } from '../../../../../platform/configuration/common/configuration.js'; -import { IInstantiationService } from '../../../../../platform/instantiation/common/instantiation.js'; -import { ILogService } from '../../../../../platform/log/common/log.js'; -import { INotificationService } from '../../../../../platform/notification/common/notification.js'; -import { IAgentHostConnectProgress } from '../../../../common/agentHostSessionsProvider.js'; -import { ISessionsProvidersService } from '../../../../services/sessions/browser/sessionsProvidersService.js'; -import { RemoteAgentHostSessionsProvider } from './remoteAgentHostSessionsProvider.js'; -import { watchForIncompatibleNotifications } from './remoteHostOptions.js'; +import { type IConfigurationService } from '../../../../../platform/configuration/common/configuration.js'; +import { type IInstantiationService } from '../../../../../platform/instantiation/common/instantiation.js'; +import { type ILogService } from '../../../../../platform/log/common/log.js'; +import { type INotificationService } from '../../../../../platform/notification/common/notification.js'; +import { type ISessionsProvidersService } from '../../../../services/sessions/browser/sessionsProvidersService.js'; +import { EntryDrivenProviderContribution } from './entryDrivenProviderContribution.js'; /** * Per-host auto-reconnect state for a managed (in-renderer relay) remote @@ -33,6 +30,8 @@ export class ManagedReconnectState extends Disposable { paused = false; /** Wall-clock timestamp when {@link paused} was last set to true. */ pausedAt = 0; + /** Whether automatic triggers must not resume this state. */ + requiresUserInitiatedResume = false; get hasPendingTimer(): boolean { return !!this._timer.value; @@ -56,6 +55,15 @@ export class ManagedReconnectState extends Disposable { this.attempts = 0; this.paused = false; this._timer.clear(); + this.requiresUserInitiatedResume = false; + } + + resumeAutomatically(): boolean { + if (!this.paused || this.requiresUserInitiatedResume) { + return false; + } + this.resetForResume(); + return true; } } @@ -73,6 +81,10 @@ export interface IManagedReconnectAttemptOptions { readonly reconnectPolicy: IRemoteAgentHostReconnectPolicy; /** Whether the given error should pause (rather than retry) auto-reconnect. */ readonly shouldPause: (err: unknown) => boolean; + /** Whether the pause must be resumed by an explicit user action. */ + readonly requiresUserInitiatedResume?: (err: unknown) => boolean; + /** Describes why a reconnect was paused for logging. */ + readonly getPauseReason?: (err: unknown) => string; /** * Optional pre-flight gate. Return `{ skip: true }` to bail WITHOUT * incrementing the attempt counter (so a long-unavailable host can't burn @@ -81,21 +93,17 @@ export interface IManagedReconnectAttemptOptions { readonly preCheck?: (userInitiated: boolean) => Promise<{ readonly skip: boolean; readonly reason?: string } | undefined>; /** Perform the actual (re)connect. */ readonly doConnect: () => Promise; - /** Schedule the next retry after a non-terminal failure. */ - readonly schedule: (state: ManagedReconnectState) => void; + /** Schedule the next retry after a non-terminal failure. Omit for on-demand-only reconnects. */ + readonly schedule?: (state: ManagedReconnectState) => void; } /** * Shared base for contributions that own in-renderer relay remote agent hosts - * (WSL, and conceptually SSH/tunnels). Encapsulates the sessions-provider + * (WSL and SSH). Encapsulates the sessions-provider * registry and the managed auto-reconnect state machine so concrete * contributions only implement their type-specific discovery/connect logic. */ -export abstract class ManagedReconnectAgentHostContribution extends Disposable { - - /** Per-address sessions provider stores. */ - protected readonly _providerStores = this._register(new DisposableMap()); - protected readonly _providerInstances = new Map(); +export abstract class ManagedReconnectAgentHostContribution extends EntryDrivenProviderContribution { /** Per-key auto-reconnect state (timer + attempts + paused). */ protected readonly _reconnectStates = this._register(new DisposableMap()); @@ -108,47 +116,14 @@ export abstract class ManagedReconnectAgentHostContribution extends Disposable { protected readonly _pendingReconnects = new Map>(); constructor( - protected readonly _remoteAgentHostService: IRemoteAgentHostService, - protected readonly _configurationService: IConfigurationService, + remoteAgentHostService: IRemoteAgentHostService, + configurationService: IConfigurationService, protected readonly _logService: ILogService, - protected readonly _instantiationService: IInstantiationService, - protected readonly _sessionsProvidersService: ISessionsProvidersService, - protected readonly _notificationService: INotificationService, + instantiationService: IInstantiationService, + sessionsProvidersService: ISessionsProvidersService, + notificationService: INotificationService, ) { - super(); - } - - protected get _enabled(): boolean { - return this._configurationService.getValue(RemoteAgentHostsEnabledSettingId); - } - - // -- Provider registry -- - - protected _createProvider(address: string, name: string, options: { - readonly connectOnDemand?: () => Promise; - readonly disconnectOnDemand?: () => Promise; - readonly onDidReportConnectProgress?: Event; - readonly initialStatus?: RemoteAgentHostConnectionStatus; - }): RemoteAgentHostSessionsProvider { - const store = new DisposableStore(); - const provider = this._instantiationService.createInstance( - RemoteAgentHostSessionsProvider, { - address, - name, - connectOnDemand: options.connectOnDemand, - disconnectOnDemand: options.disconnectOnDemand, - onDidReportConnectProgress: options.onDidReportConnectProgress, - }); - if (options.initialStatus !== undefined) { - provider.setConnectionStatus(options.initialStatus); - } - store.add(provider); - store.add(this._sessionsProvidersService.registerProvider(provider)); - store.add(watchForIncompatibleNotifications(provider, this._instantiationService, this._notificationService)); - this._providerInstances.set(address, provider); - store.add(toDisposable(() => this._providerInstances.delete(address))); - this._providerStores.set(address, store); - return provider; + super(remoteAgentHostService, configurationService, instantiationService, sessionsProvidersService, notificationService); } // -- Managed auto-reconnect -- @@ -170,8 +145,7 @@ export abstract class ManagedReconnectAgentHostContribution extends Disposable { protected _resumeReconnects(logKind: string): number { let resumed = 0; for (const [, state] of this._reconnectStates) { - if (state.paused) { - state.resetForResume(); + if (state.resumeAutomatically()) { resumed++; } } @@ -229,11 +203,12 @@ export abstract class ManagedReconnectAgentHostContribution extends Disposable { provider?.setConnectionStatus(RemoteAgentHostConnectionStatus.disconnected); } if (opts.shouldPause(err)) { - this._logService.info(`[RemoteAgentHost] Pausing ${opts.kind} auto-reconnect for ${opts.key} after user cancellation`); + this._logService.info(`[RemoteAgentHost] Pausing ${opts.kind} auto-reconnect for ${opts.key} after ${opts.getPauseReason?.(err) ?? 'user cancellation'}`); provider?.unpublishCachedSessions(); const liveState = this._getOrCreateReconnectState(opts.key); liveState.paused = true; liveState.pausedAt = Date.now(); + liveState.requiresUserInitiatedResume = opts.requiresUserInitiatedResume?.(err) ?? false; return; } this._logService.error(`[RemoteAgentHost] ${opts.kind} reconnect failed for ${opts.key}`, err); @@ -265,7 +240,7 @@ export abstract class ManagedReconnectAgentHostContribution extends Disposable { if (opts.userInitiated) { return; } - opts.schedule(liveState); + opts.schedule?.(liveState); } })(); this._pendingReconnects.set(opts.key, runPromise); diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/remoteAgentHost.contribution.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/remoteAgentHost.contribution.ts index 0d5110d609ba0c..c975ffebbcea5a 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/remoteAgentHost.contribution.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/remoteAgentHost.contribution.ts @@ -4,30 +4,24 @@ *--------------------------------------------------------------------------------------------*/ import { Event } from '../../../../../base/common/event.js'; -import { Disposable, DisposableMap, DisposableStore, MutableDisposable, toDisposable } from '../../../../../base/common/lifecycle.js'; -import { disposableTimeout, IntervalTimer } from '../../../../../base/common/async.js'; -import { isCancellationError } from '../../../../../base/common/errors.js'; -import { StopWatch } from '../../../../../base/common/stopwatch.js'; +import { Disposable, DisposableMap, DisposableStore, toDisposable } from '../../../../../base/common/lifecycle.js'; import { URI } from '../../../../../base/common/uri.js'; import * as nls from '../../../../../nls.js'; import { agentHostAuthority } from '../../../../../platform/agentHost/common/agentHostUri.js'; import { AgentHostProtocolClient } from '../../../../../platform/agentHost/browser/agentHostProtocolClient.js'; import { type AgentProvider, type AuthenticateParams, type AuthenticateResult } from '../../../../../platform/agentHost/common/agent.js'; import { type IAgentConnection } from '../../../../../platform/agentHost/common/agentService.js'; -import { IRemoteAgentHostConnectionInfo, IRemoteAgentHostEntry, IRemoteAgentHostService, type IRemoteAgentHostSSHConnection, RemoteAgentHostAutoConnectSettingId, RemoteAgentHostConnectionStatus, RemoteAgentHostEntryType, RemoteAgentHostsEnabledSettingId, RemoteAgentHostsSettingId, getEntryAddress } from '../../../../../platform/agentHost/common/remoteAgentHostService.js'; +import { IRemoteAgentHostConnectionInfo, IRemoteAgentHostService, RemoteAgentHostAutoConnectSettingId, RemoteAgentHostConnectionStatus, RemoteAgentHostsEnabledSettingId, RemoteAgentHostsSettingId, getEntryAddress } from '../../../../../platform/agentHost/common/remoteAgentHostService.js'; import { TunnelAgentHostsSettingId } from '../../../../../platform/agentHost/common/tunnelAgentHost.js'; import { CloudSandboxEnabledSettingId } from '../../../../../platform/agentHost/common/cloudSandboxAgentHost.js'; -import { PROTOCOL_VERSION } from '../../../../../platform/agentHost/common/state/protocol/version/registry.js'; import { AgentHostLocalFilePermissionsSettingId } from '../../../../../platform/agentHost/common/agentHostResourceService.js'; import { type ProtectedResourceMetadata } from '../../../../../platform/agentHost/common/state/protocol/state.js'; import { type AgentInfo, type RootState } from '../../../../../platform/agentHost/common/state/sessionState.js'; import { NotificationType, type INotification } from '../../../../../platform/agentHost/common/state/sessionActions.js'; -import { IConfigurationService } from '../../../../../platform/configuration/common/configuration.js'; import { ConfigurationScope, Extensions as ConfigurationExtensions, IConfigurationRegistry } from '../../../../../platform/configuration/common/configurationRegistry.js'; import { IDefaultAccountService } from '../../../../../platform/defaultAccount/common/defaultAccount.js'; import { IInstantiationService, ServicesAccessor } from '../../../../../platform/instantiation/common/instantiation.js'; import { ILogService } from '../../../../../platform/log/common/log.js'; -import { INotificationService } from '../../../../../platform/notification/common/notification.js'; import { Registry } from '../../../../../platform/registry/common/platform.js'; import { IWorkbenchContribution, registerWorkbenchContribution2, WorkbenchPhase } from '../../../../../workbench/common/contributions.js'; import { registerAction2 } from '../../../../../platform/actions/common/actions.js'; @@ -49,11 +43,9 @@ import { RemoteAgentHostLogForwarder } from './remoteAgentHostLogForwarder.js'; import { RemoteAgentHostSessionsProvider } from './remoteAgentHostSessionsProvider.js'; import { IRemoteAgentHostConnectionCustomizationService, RemoteAgentHostConnectionCustomizationService } from './remoteAgentHostConnectionCustomization.js'; import { InstantiationType, registerSingleton } from '../../../../../platform/instantiation/common/extensions.js'; -import { watchForIncompatibleNotifications } from './remoteHostOptions.js'; -import { computeSSHConnectionKey, isSSHHostKeyDeniedError, ISSHRemoteAgentHostService, SSHAuthMethod } from '../../../../../platform/agentHost/common/sshRemoteAgentHost.js'; import { IAgentHostTerminalService } from '../../../../../workbench/contrib/terminal/browser/agentHostTerminalService.js'; import { ITelemetryService } from '../../../../../platform/telemetry/common/telemetry.js'; -import { categorizeSSHConnectError, logSSHConnectAttempt, logTerminalRecovery } from '../../../../common/sessionsTelemetry.js'; +import { logTerminalRecovery } from '../../../../common/sessionsTelemetry.js'; Registry.as(ChatSessionsExtensions.AsyncActivation).register({ matchSessionType: sessionType => isRemoteAgentHostSessionType(sessionType), @@ -115,102 +107,6 @@ function getAddressForSessionType(sessionType: string, remoteAgentHostService: I return authority ? authorities.get(authority) : undefined; } -/** - * How often the periodic provider reconciliation backstop runs. - */ -const SSH_RECONNECT_PERIODIC_INTERVAL_MS = 60_000; // 1 minute - -/** - * Per-host SSH reconnect state used to preserve user-requested pause and - * resume behavior. Owned by {@link RemoteAgentHostContribution._sshReconnectStates}. - */ -export class SSHReconnectState extends Disposable { - private readonly _timer = this._register(new MutableDisposable()); - - /** Consecutive failed reconnect attempts. */ - attempts = 0; - /** True after a reconnect was paused until something resumes it. */ - paused = false; - /** Wall-clock timestamp when {@link paused} was last set to true. */ - pausedAt = 0; - /** Whether only an explicit user reconnect should resume this state. */ - requiresUserInitiatedResume = false; - - get hasPendingTimer(): boolean { - return !!this._timer.value; - } - - scheduleRetry(delayMs: number, handler: () => void): void { - this._timer.value = disposableTimeout(() => { - // Drop the disposable now that the timer has fired so - // `hasPendingTimer` reflects reality even if `handler` returns - // early without scheduling a follow-up attempt. - this._timer.value = undefined; - handler(); - }, delayMs); - } - - cancelTimer(): void { - this._timer.clear(); - } - - resetForResume(): void { - this.attempts = 0; - this.paused = false; - this._timer.clear(); - this.requiresUserInitiatedResume = false; - } - - resumeAutomatically(): boolean { - if (!this.paused || this.requiresUserInitiatedResume) { - return false; - } - this.resetForResume(); - return true; - } -} - -export function shouldPauseSSHReconnectAfterFailure(err: unknown): boolean { - return isCancellationError(err) || isSSHHostKeyDeniedError(err); -} - -/** - * Connection key passed to {@link ISSHRemoteAgentHostService.disconnect} for - * an SSH-backed remote agent host entry. Mirrors the key the SSH service - * itself constructs when it stores the connection. - */ -export function sshConnectionKey(connection: IRemoteAgentHostSSHConnection): string { - return connection.sshConfigHost - ? `ssh:${connection.sshConfigHost}` - : `${connection.user ?? connection.hostName}@${connection.hostName}:${connection.port ?? 22}`; -} - -/** - * Sequence the steps to disconnect an SSH-backed remote agent host entry - * triggered by the user (e.g. clicking X in the workspace picker). - * - * Order matters: `removeRemoteAgentHost` MUST run before the SSH tunnel - * teardown. `sshService.disconnect()` fires `onDidCloseConnection` - * synchronously, which the renderer translates into `onDidChangeConnections` - * and the contribution's reconciliation. If the entry is still in configured - * storage at that point, it can be surfaced again before teardown completes. - * - * `removeRemoteAgentHost` itself runs the entry's transport disposable - * (which calls `_mainService.disconnect(connectionId)`), so the underlying - * SSH tunnel is already closed when this returns. The explicit - * `sshService.disconnect(connectionKey)` is belt-and-suspenders to clear - * the connection by its connection key as well, matching the prior - * teardown behavior. - */ -export async function disconnectSSHEntry( - connection: IRemoteAgentHostSSHConnection, - remoteAgentHostService: Pick, - sshService: Pick, -): Promise { - await remoteAgentHostService.removeRemoteAgentHost(connection.address); - await sshService.disconnect(sshConnectionKey(connection)); -} - /** Per-connection state bundle, disposed when a connection is removed. */ class ConnectionState extends Disposable { readonly store = this._register(new DisposableStore()); @@ -244,19 +140,6 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc /** Per-connection state: client state + per-agent registrations. */ private readonly _connections = this._register(new DisposableMap()); - /** Per-address sessions provider, registered for all configured entries. */ - private readonly _providerStores = this._register(new DisposableMap()); - private readonly _providerInstances = new Map(); - /** - * In-flight reconnect attempts keyed by host id (`sshConfigHost` for SSH, - * `distro` for WSL). Stores the {@link _attemptManagedReconnect} promise - * so concurrent user requests join the existing attempt rather than racing it. - */ - private readonly _pendingSSHReconnects = new Map>(); - - /** Per-host SSH reconnect state (timer + attempts + paused). */ - private readonly _sshReconnectStates = this._register(new DisposableMap()); - constructor( @IRemoteAgentHostService private readonly _remoteAgentHostService: IRemoteAgentHostService, @IChatSessionsService private readonly _chatSessionsService: IChatSessionsService, @@ -265,11 +148,8 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc @IInstantiationService private readonly _instantiationService: IInstantiationService, @IAuthenticationService private readonly _authenticationService: IAuthenticationService, @IDefaultAccountService private readonly _defaultAccountService: IDefaultAccountService, - @INotificationService private readonly _notificationService: INotificationService, @ISessionsProvidersService private readonly _sessionsProvidersService: ISessionsProvidersService, - @IConfigurationService private readonly _configurationService: IConfigurationService, @IAgentHostFileSystemService private readonly _agentHostFileSystemService: IAgentHostFileSystemService, - @ISSHRemoteAgentHostService private readonly _sshService: ISSHRemoteAgentHostService, @ICustomizationHarnessService private readonly _customizationHarnessService: ICustomizationHarnessService, @IAgentHostTerminalService private readonly _agentHostTerminalService: IAgentHostTerminalService, @ITelemetryService private readonly _telemetryService: ITelemetryService, @@ -278,342 +158,15 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc ) { super(); - // Reconcile providers when configured entries change - this._register(this._configurationService.onDidChangeConfiguration(e => { - if (e.affectsConfiguration(RemoteAgentHostsSettingId) || e.affectsConfiguration(RemoteAgentHostsEnabledSettingId)) { - // User changed config — reset any paused on-demand state. - this._resumeSSHReconnects(); - this._reconcile(); - } - })); - - // Reconcile when connections change (added/removed/reconnected) - this._register(this._remoteAgentHostService.onDidChangeConnections(() => { - // New/removed connection gives paused on-demand state a fresh start. - this._resumeSSHReconnects(); - this._reconcile(); - })); - - // Cancel any pending SSH reconnect timers on dispose. - // (Handled automatically by the DisposableMap above; nothing extra needed here.) - - // Push auth token whenever the default account or sessions change + this._register(this._remoteAgentHostService.onDidChangeConnections(() => this._reconcile())); this._register(this._defaultAccountService.onDidChangeDefaultAccount(() => this._authenticateAllConnections())); this._register(this._authenticationService.onDidChangeSessions(() => this._authenticateAllConnections())); - // Initial setup for configured entries and connected remotes this._reconcile(); - - // Periodic backstop: reconcile provider state even if the event-driven - // chain breaks after a sleep/wake cycle. - this._register(new IntervalTimer()).cancelAndSet( - () => { - this._logService.trace('[RemoteAgentHost] Periodic reconcile (backstop)'); - this._reconcile(); - }, - SSH_RECONNECT_PERIODIC_INTERVAL_MS, - ); } private _reconcile(): void { - this._reconcileProviders(); this._reconcileConnections(); - - // Ensure every live connection is wired to its provider. This covers - // the case where a provider was recreated (e.g. name change) while a - // connection for that address already existed. - for (const [address, connState] of this._connections) { - const connectionInfo = this._remoteAgentHostService.connections.find(c => c.address === address); - const provider = this._providerInstances.get(address); - if (provider) { - provider.setConnection(connState.connection, connectionInfo?.defaultDirectory); - } - } - - // Update connection status on all providers (including those - // that are reconnecting and don't have an active connection). - for (const [address, provider] of this._providerInstances) { - const connectionInfo = this._remoteAgentHostService.connections.find(c => c.address === address); - if (connectionInfo) { - // Service has an entry for this address — its status is - // authoritative (including the `incompatible` set by the - // WebSocket connect failure path, and the `connecting` or - // `reconnecting` status of a fresh reconnect attempt). - provider.setConnectionStatus(connectionInfo.status); - } else if (!RemoteAgentHostConnectionStatus.isIncompatible(provider.connectionStatus.get())) { - // No service entry. Preserve incompatible state set by - // the SSH reconnect catch (where the failure happens - // before the service ever sees an entry); otherwise fall - // back to disconnected. - provider.setConnectionStatus(RemoteAgentHostConnectionStatus.disconnected); - } - } - } - - private _reconcileProviders(): void { - const enabled = this._configurationService.getValue(RemoteAgentHostsEnabledSettingId); - const entries = enabled ? this._remoteAgentHostService.configuredEntries : []; - const desiredAddresses = new Set(entries.map(e => getEntryAddress(e))); - - // Remove providers no longer configured - for (const [address] of this._providerStores) { - if (!desiredAddresses.has(address)) { - this._providerStores.deleteAndDispose(address); - } - } - - // Add or recreate providers for configured entries - for (const entry of entries) { - const address = getEntryAddress(entry); - const existing = this._providerInstances.get(address); - if (existing && existing.label !== (entry.name || address)) { - // Name changed — recreate since ISessionsProvider.label is readonly - this._providerStores.deleteAndDispose(address); - } - if (!this._providerStores.has(address)) { - this._createProvider(entry); - } - } - } - - private _createProvider(entry: IRemoteAgentHostEntry): void { - const address = getEntryAddress(entry); - const sshConnection = entry.connection.type === RemoteAgentHostEntryType.SSH ? entry.connection : undefined; - let connectOnDemand: (() => Promise) | undefined; - let disconnectOnDemand: (() => Promise) | undefined; - let preferenceKey: string | undefined; - if (sshConnection) { - connectOnDemand = () => this._connectSSHOnDemand(sshConnection, entry.name, address); - disconnectOnDemand = () => this._disconnectSSHOnDemand(sshConnection); - // The stable key SSHRemoteAgentHostService reads its preference - // by (see computeSSHConnectionKey's docs) - NOT the live - // forwarded `address` above, which changes per-connection. - preferenceKey = computeSSHConnectionKey({ - sshConfigHost: sshConnection.sshConfigHost, - username: sshConnection.user, - host: sshConnection.hostName, - port: sshConnection.port, - }); - } - const store = new DisposableStore(); - const provider = this._instantiationService.createInstance( - RemoteAgentHostSessionsProvider, { address, name: entry.name, connectOnDemand, disconnectOnDemand, preferenceKey }); - store.add(provider); - store.add(this._sessionsProvidersService.registerProvider(provider)); - store.add(watchForIncompatibleNotifications(provider, this._instantiationService, this._notificationService)); - this._providerInstances.set(address, provider); - store.add(toDisposable(() => this._providerInstances.delete(address))); - this._providerStores.set(address, store); - } - - private async _connectSSHOnDemand(connection: IRemoteAgentHostSSHConnection, name: string, address: string): Promise { - const sshConfigHost = connection.sshConfigHost; - if (!sshConfigHost) { - const stopwatch = StopWatch.create(false); - try { - await this._sshService.connect({ - host: connection.hostName, - port: connection.port, - username: connection.user ?? connection.hostName, - authMethod: SSHAuthMethod.Agent, - name, - userInitiated: true, - }); - logSSHConnectAttempt(this._telemetryService, { - operation: 'connect', - userInitiated: true, - attempt: 1, - durationMs: stopwatch.elapsed(), - success: true, - willRetry: false, - }); - } catch (err) { - logSSHConnectAttempt(this._telemetryService, { - operation: 'connect', - userInitiated: true, - attempt: 1, - durationMs: stopwatch.elapsed(), - success: false, - willRetry: false, - errorCategory: categorizeSSHConnectError(err), - }); - throw err; - } - return; - } - if (this._pendingSSHReconnects.has(sshConfigHost)) { - await this._pendingSSHReconnects.get(sshConfigHost)!.catch(() => undefined); - return; - } - this._sshReconnectStates.get(sshConfigHost)?.resetForResume(); - await this._attemptSSHReconnect(sshConfigHost, name, address, { userInitiated: true }); - } - - private async _disconnectSSHOnDemand(connection: IRemoteAgentHostSSHConnection): Promise { - if (connection.sshConfigHost) { - this._sshReconnectStates.deleteAndDispose(connection.sshConfigHost); - } - await disconnectSSHEntry(connection, this._remoteAgentHostService, this._sshService); - } - - private async _attemptSSHReconnect(sshConfigHost: string, name: string, address: string, options: { userInitiated?: boolean } = {}): Promise { - await this._attemptManagedReconnect({ - kind: 'SSH', - key: sshConfigHost, - address, - userInitiated: !!options.userInitiated, - shouldPause: shouldPauseSSHReconnectAfterFailure, - pending: this._pendingSSHReconnects, - states: this._sshReconnectStates, - getOrCreateState: key => this._getOrCreateSSHReconnectState(key), - doConnect: async () => { - this._remoteAgentHostService.reconnect(address, !!options.userInitiated); - await this._remoteAgentHostService.waitForConnection(address); - }, - }); - } - - private _getOrCreateSSHReconnectState(sshConfigHost: string): SSHReconnectState { - let state = this._sshReconnectStates.get(sshConfigHost); - if (!state) { - state = new SSHReconnectState(); - this._sshReconnectStates.set(sshConfigHost, state); - } - return state; - } - - /** - * Reset paused SSH reconnect state after a fresh external trigger. - */ - private _resumeSSHReconnects(): void { - let resumed = 0; - for (const [, state] of this._sshReconnectStates) { - if (state.resumeAutomatically()) { - resumed++; - } - } - if (resumed > 0) { - this._logService.info(`[RemoteAgentHost] Reset SSH reconnect state for ${resumed} paused host(s)`); - } - } - - /** - * Shared retry-loop body for SSH managed-reconnect entries. - * - * Handles `connecting`/`reconnecting`/`disconnected`/`incompatible` provider status, - * cached-session unpublishing on failure, pause-on-cancel, and - * pause-after-max-attempts. An optional pre-check can bail out without - * incrementing the attempt counter (returns `{ skip: true }`). - */ - private async _attemptManagedReconnect(opts: { - readonly kind: 'SSH'; - readonly key: string; - readonly address: string; - readonly userInitiated: boolean; - readonly shouldPause: (err: unknown) => boolean; - readonly pending: Map>; - readonly states: DisposableMap; - readonly getOrCreateState: (key: string) => SSHReconnectState; - readonly preCheck?: (userInitiated: boolean) => Promise<{ readonly skip: boolean; readonly reason?: string } | undefined>; - readonly doConnect: () => Promise; - }): Promise { - // Wrap the body so we can store our own promise in `opts.pending` for - // concurrent on-demand callers to join. The inner IIFE keeps the - // existing control flow intact; only the bookkeeping moves out. - const runPromise = (async () => { - const live = this._remoteAgentHostService.connections.find(connection => connection.address === opts.address); - if (!opts.userInitiated && RemoteAgentHostConnectionStatus.isConnecting(live?.status)) { - return; - } - if (!opts.userInitiated && RemoteAgentHostConnectionStatus.isReconnecting(live?.status)) { - // The protocol client is preserving its state while it reconnects; don't replace it. - this._sshReconnectStates.get(opts.key)?.cancelTimer(); - return; - } - const state = opts.getOrCreateState(opts.key); - const attempt = state.attempts; - const provider = this._providerInstances.get(opts.address); - const stopwatch = StopWatch.create(false); - if (opts.userInitiated) { - provider?.setConnectionStatus(RemoteAgentHostConnectionStatus.connecting); - } - this._logService.info(`[RemoteAgentHost] Re-establishing ${opts.kind} connection for ${opts.key} (attempt ${attempt + 1})`); - try { - if (opts.preCheck) { - const result = await opts.preCheck(opts.userInitiated); - if (result?.skip) { - if (result.reason) { - this._logService.info(`[RemoteAgentHost] ${opts.kind} reconnect for ${opts.key}: ${result.reason}; skipping`); - } - return; - } - } - await opts.doConnect(); - logSSHConnectAttempt(this._telemetryService, { - operation: 'reconnect', - userInitiated: opts.userInitiated, - attempt: attempt + 1, - durationMs: stopwatch.elapsed(), - success: true, - willRetry: false, - }); - opts.states.deleteAndDispose(opts.key); - this._logService.info(`[RemoteAgentHost] ${opts.kind} connection re-established for ${opts.key}`); - } catch (err) { - const enabled = this._configurationService.getValue(RemoteAgentHostsEnabledSettingId); - const pause = opts.shouldPause(err); - const incompatible = RemoteAgentHostConnectionStatus.fromConnectError(err, [PROTOCOL_VERSION]); - logSSHConnectAttempt(this._telemetryService, { - operation: 'reconnect', - userInitiated: opts.userInitiated, - attempt: attempt + 1, - durationMs: stopwatch.elapsed(), - success: false, - willRetry: false, - errorCategory: categorizeSSHConnectError(err), - }); - if (!enabled) { - opts.states.deleteAndDispose(opts.key); - return; - } - if (opts.userInitiated) { - provider?.setConnectionStatus(RemoteAgentHostConnectionStatus.disconnected); - } - if (pause) { - const requiresUserInitiatedResume = isSSHHostKeyDeniedError(err); - this._logService.info(`[RemoteAgentHost] Pausing ${opts.kind} reconnect for ${opts.key} after ${requiresUserInitiatedResume ? 'host key denial' : 'user cancellation'}`); - provider?.unpublishCachedSessions(); - const liveState = opts.getOrCreateState(opts.key); - liveState.paused = true; - liveState.pausedAt = Date.now(); - liveState.requiresUserInitiatedResume = requiresUserInitiatedResume; - return; - } - this._logService.error(`[RemoteAgentHost] ${opts.kind} reconnect failed for ${opts.key}`, err); - // Surface protocol-version mismatches on the provider so the - // workspace picker can show the host's message and the user - // can read it. Other errors stay as the existing disconnected - // state. - if (incompatible) { - provider?.setConnectionStatus(incompatible); - // Don't keep retrying on incompatible — user needs to - // upgrade/downgrade. Drop retry state instead of pausing. - opts.states.deleteAndDispose(opts.key); - return; - } - // Host is unreachable — unpublish any cached sessions we - // were showing so the UI doesn't list stale entries for a - // host we cannot currently reach. - provider?.unpublishCachedSessions(); - return; - } - })(); - opts.pending.set(opts.key, runPromise); - try { - await runPromise; - } finally { - opts.pending.delete(opts.key); - } } private _reconcileConnections(): void { @@ -629,12 +182,10 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc for (const [address] of this._connections) { if (!allAddresses.has(address)) { this._logService.info(`[RemoteAgentHost] Removing contribution for ${address}`); - this._providerInstances.get(address)?.clearConnection(); this._connections.deleteAndDispose(address); } else if (!connectedAddresses.has(address)) { // Connection exists but is not connected (reconnecting or disconnected). - // Keep the contribution state but don't clear the provider — - // the session cache is preserved during reconnect. + // Keep the contribution state while the connection restores. } } @@ -722,11 +273,6 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc this._handleRootStateChange(address, connection, initialRootState); } - // Wire connection to existing sessions provider - const provider = this._providerInstances.get(address); - if (provider) { - provider.setConnection(connection, connectionInfo.defaultDirectory); - } } private _handleRootStateChange(address: string, connection: IAgentConnection, rootState: RootState): void { @@ -1005,7 +551,7 @@ Registry.as(ConfigurationExtensions.Configuration).regis }, [RemoteAgentHostAutoConnectSettingId]: { type: 'boolean', - description: nls.localize('chat.remoteAgentHosts.autoConnect', "Automatically connect to online dev tunnel and WSL remote agent hosts on startup. When disabled, cached sessions are still shown but connections are established only on demand."), + description: nls.localize('chat.remoteAgentHosts.autoConnect', "Automatically connect to online dev tunnel, SSH, and WSL remote agent hosts on startup. When disabled, cached sessions are still shown but connections are established only on demand."), default: true, scope: ConfigurationScope.APPLICATION, tags: ['experimental', 'advanced'], diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/remoteAgentHostActions.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/remoteAgentHostActions.ts index c1ae189ab4728e..4de49bb7ea9aff 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/remoteAgentHostActions.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/remoteAgentHostActions.ts @@ -613,8 +613,7 @@ async function promptForRemoteFolder( const sessionsService = accessor.get(ISessionsService); const sessionsPartService = accessor.get(ISessionsPartService); - // The provider is created synchronously during addManagedConnection's - // onDidChangeConnections event, so it should exist by now. + // The factory-backed entry fires onDidChangeConnections before its handshake completes, so the provider should exist by now. const provider = sessionsProvidersService.getProviders().find((p): p is IAgentHostSessionsProvider => isAgentHostProvider(p) && p.remoteAddress === connection.localAddress); if (!provider) { return; @@ -1028,7 +1027,8 @@ async function promptToConnectViaTunnel( try { // `connect` caches the tunnel internally before wiring the live // connection — no separate `cacheTunnel` call needed here. - await tunnelService.connect(picked.tunnel, authProvider); + tunnelService.clearTunnelDismissal(picked.tunnel.tunnelId); + await tunnelService.connect(picked.tunnel, authProvider, { userInitiated: true }); handle.close(); } catch (err) { handle.close(); diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/sshAgentHost.contribution.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/sshAgentHost.contribution.ts new file mode 100644 index 00000000000000..9169efa9edab13 --- /dev/null +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/sshAgentHost.contribution.ts @@ -0,0 +1,231 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { IntervalTimer } from '../../../../../base/common/async.js'; +import { isCancellationError } from '../../../../../base/common/errors.js'; +import { StopWatch } from '../../../../../base/common/stopwatch.js'; +import { type IRemoteAgentHostEntry, IRemoteAgentHostService, type IRemoteAgentHostSSHConnection, getEntryAddress, getEntryTypeConfig, RemoteAgentHostEntryType, RemoteAgentHostsEnabledSettingId, RemoteAgentHostsSettingId } from '../../../../../platform/agentHost/common/remoteAgentHostService.js'; +import { computeReconnectDelay } from '../../../../../platform/agentHost/common/reconnectPolicy.js'; +import { computeSSHConnectionKey, isSSHHostKeyDeniedError, ISSHRemoteAgentHostService, SSHAuthMethod } from '../../../../../platform/agentHost/common/sshRemoteAgentHost.js'; +import { IConfigurationService } from '../../../../../platform/configuration/common/configuration.js'; +import { IInstantiationService } from '../../../../../platform/instantiation/common/instantiation.js'; +import { ILogService } from '../../../../../platform/log/common/log.js'; +import { INotificationService } from '../../../../../platform/notification/common/notification.js'; +import { ITelemetryService } from '../../../../../platform/telemetry/common/telemetry.js'; +import { IWorkbenchContribution, registerWorkbenchContribution2, WorkbenchPhase } from '../../../../../workbench/common/contributions.js'; +import { categorizeSSHConnectError, logSSHConnectAttempt } from '../../../../common/sessionsTelemetry.js'; +import { ISessionsProvidersService } from '../../../../services/sessions/browser/sessionsProvidersService.js'; +import { ManagedReconnectAgentHostContribution } from './managedReconnectAgentHostContribution.js'; + +const SSH_RECONNECT_PERIODIC_INTERVAL_MS = 60_000; + +/** Returns whether an SSH reconnect failure requires pausing retries. */ +export function shouldPauseSSHReconnectAfterFailure(err: unknown): boolean { + return isCancellationError(err) || isSSHHostKeyDeniedError(err); +} + +/** Returns the SSH service's stable key for a configured connection. */ +export function sshConnectionKey(connection: IRemoteAgentHostSSHConnection): string { + return connection.sshConfigHost + ? `ssh:${connection.sshConfigHost}` + : `${connection.user ?? connection.hostName}@${connection.hostName}:${connection.port ?? 22}`; +} + +/** + * Disconnect an SSH-backed remote agent host at the user's request. + * + * Order matters. `sshService.disconnect` is what drops the persisted SSH + * entry, and that entry is what makes the address "desired" during + * reconciliation. Tearing the connection down first fires + * `onDidChangeConnections` while the entry is still stored, so reconciliation + * sees a desired-but-disconnected host and immediately re-dials it — the host + * reappears moments after the user removed it. Dropping the entry first makes + * the address undesired, so the teardown's own reconcile is a no-op. + */ +export async function disconnectSSHEntry( + connection: IRemoteAgentHostSSHConnection, + remoteAgentHostService: Pick, + sshService: Pick, +): Promise { + await sshService.disconnect(sshConnectionKey(connection)); + await remoteAgentHostService.removeRemoteAgentHost(connection.address); +} + +export class SSHAgentHostContribution extends ManagedReconnectAgentHostContribution implements IWorkbenchContribution { + + static readonly ID = 'sessions.contrib.sshAgentHostContribution'; + + protected readonly _entryType = RemoteAgentHostEntryType.SSH; + + protected override get _clearConnectionOnRemoval(): boolean { + return true; + } + + constructor( + @IRemoteAgentHostService remoteAgentHostService: IRemoteAgentHostService, + @ISSHRemoteAgentHostService private readonly _sshService: ISSHRemoteAgentHostService, + @IConfigurationService configurationService: IConfigurationService, + @ILogService logService: ILogService, + @IInstantiationService instantiationService: IInstantiationService, + @ISessionsProvidersService sessionsProvidersService: ISessionsProvidersService, + @INotificationService notificationService: INotificationService, + @ITelemetryService private readonly _telemetryService: ITelemetryService, + ) { + super(remoteAgentHostService, configurationService, logService, instantiationService, sessionsProvidersService, notificationService); + + this._register(this._remoteAgentHostService.onDidChangeConnections(() => { + this._resumeSSHReconnects(); + this._reconcile(); + })); + + this._register(this._configurationService.onDidChangeConfiguration(e => { + if (e.affectsConfiguration(RemoteAgentHostsSettingId) || e.affectsConfiguration(RemoteAgentHostsEnabledSettingId)) { + this._resumeSSHReconnects(); + this._reconcile(); + } + })); + + this._register(new IntervalTimer()).cancelAndSet(() => { + this._resumeSSHReconnects(); + this._reconcile(); + }, SSH_RECONNECT_PERIODIC_INTERVAL_MS); + + this._reconcile(); + } + + protected override _getProviderOptions(entry: IRemoteAgentHostEntry) { + if (entry.connection.type !== RemoteAgentHostEntryType.SSH) { + return {}; + } + const connection = entry.connection; + const address = getEntryAddress(entry); + return { + connectOnDemand: () => this._connectSSHOnDemand(connection, entry.name, address), + disconnectOnDemand: () => this._disconnectSSHOnDemand(connection), + preferenceKey: computeSSHConnectionKey({ + sshConfigHost: connection.sshConfigHost, + username: connection.user, + host: connection.hostName, + port: connection.port, + }), + }; + } + + private async _connectSSHOnDemand(connection: IRemoteAgentHostSSHConnection, name: string, address: string): Promise { + const sshConfigHost = connection.sshConfigHost; + if (!sshConfigHost) { + const stopwatch = StopWatch.create(false); + try { + await this._sshService.connect({ + host: connection.hostName, + port: connection.port, + username: connection.user ?? connection.hostName, + authMethod: SSHAuthMethod.Agent, + name, + userInitiated: true, + }); + logSSHConnectAttempt(this._telemetryService, { + operation: 'connect', + userInitiated: true, + attempt: 1, + durationMs: stopwatch.elapsed(), + success: true, + willRetry: false, + }); + } catch (err) { + logSSHConnectAttempt(this._telemetryService, { + operation: 'connect', + userInitiated: true, + attempt: 1, + durationMs: stopwatch.elapsed(), + success: false, + willRetry: false, + errorCategory: categorizeSSHConnectError(err), + }); + throw err; + } + return; + } + const pending = this._pendingReconnects.get(sshConfigHost); + if (pending) { + await pending.catch(() => undefined); + return; + } + this._reconnectStates.get(sshConfigHost)?.resetForResume(); + await this._attemptSSHReconnect(sshConfigHost, name, address, true); + } + + private async _disconnectSSHOnDemand(connection: IRemoteAgentHostSSHConnection): Promise { + if (connection.sshConfigHost) { + this._reconnectStates.deleteAndDispose(connection.sshConfigHost); + } + await disconnectSSHEntry(connection, this._remoteAgentHostService, this._sshService); + } + + private async _attemptSSHReconnect(sshConfigHost: string, name: string, address: string, userInitiated: boolean): Promise { + const reconnectPolicy = getEntryTypeConfig(RemoteAgentHostEntryType.SSH).reconnect; + const attempt = (this._reconnectStates.get(sshConfigHost)?.attempts ?? 0) + 1; + const stopwatch = StopWatch.create(false); + await this._attemptManagedReconnect({ + kind: 'SSH', + key: sshConfigHost, + address, + userInitiated, + reconnectPolicy, + shouldPause: shouldPauseSSHReconnectAfterFailure, + requiresUserInitiatedResume: isSSHHostKeyDeniedError, + getPauseReason: err => isSSHHostKeyDeniedError(err) ? 'host key denial' : 'user cancellation', + doConnect: async () => { + try { + this._remoteAgentHostService.reconnect(address, userInitiated); + await this._remoteAgentHostService.waitForConnection(address); + logSSHConnectAttempt(this._telemetryService, { + operation: 'reconnect', + userInitiated, + attempt, + durationMs: stopwatch.elapsed(), + success: true, + willRetry: false, + }); + } catch (err) { + logSSHConnectAttempt(this._telemetryService, { + operation: 'reconnect', + userInitiated, + attempt, + durationMs: stopwatch.elapsed(), + success: false, + willRetry: false, + errorCategory: categorizeSSHConnectError(err), + }); + throw err; + } + }, + schedule: state => { + state.scheduleRetry(computeReconnectDelay(reconnectPolicy, state.attempts), () => { + void this._attemptSSHReconnect(sshConfigHost, name, address, false); + }); + }, + }); + } + + private _resumeSSHReconnects(): void { + let resumed = 0; + for (const entry of this._getProviderEntries()) { + if (entry.connection.type !== RemoteAgentHostEntryType.SSH || !entry.connection.sshConfigHost) { + continue; + } + const state = this._reconnectStates.get(entry.connection.sshConfigHost); + if (state?.resumeAutomatically()) { + resumed++; + void this._attemptSSHReconnect(entry.connection.sshConfigHost, entry.name, getEntryAddress(entry), false); + } + } + if (resumed > 0) { + this._logService.info(`[RemoteAgentHost] Resuming SSH auto-reconnect for ${resumed} paused host(s)`); + } + } +} + +registerWorkbenchContribution2(SSHAgentHostContribution.ID, SSHAgentHostContribution, WorkbenchPhase.AfterRestored); diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/tunnelAgentHost.contribution.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/tunnelAgentHost.contribution.ts index 24978011ed89d1..51cfad7f102e71 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/tunnelAgentHost.contribution.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/tunnelAgentHost.contribution.ts @@ -4,13 +4,9 @@ *--------------------------------------------------------------------------------------------*/ import { Disposable, DisposableMap, DisposableStore, toDisposable } from '../../../../../base/common/lifecycle.js'; -import { isWeb } from '../../../../../base/common/platform.js'; -import { mainWindow } from '../../../../../base/browser/window.js'; import * as nls from '../../../../../nls.js'; -import { IRemoteAgentHostService, RemoteAgentHostAutoConnectSettingId, RemoteAgentHostConnectionStatus, RemoteAgentHostEntryType, RemoteAgentHostsEnabledSettingId, getEntryTypeConfig } from '../../../../../platform/agentHost/common/remoteAgentHostService.js'; -import { computeReconnectDelay, hasExhaustedReconnectAttempts } from '../../../../../platform/agentHost/common/reconnectPolicy.js'; -import { isTunnelHosted, ITunnelAgentHostService, TUNNEL_ADDRESS_PREFIX, type ITunnelInfo } from '../../../../../platform/agentHost/common/tunnelAgentHost.js'; -import { PROTOCOL_VERSION } from '../../../../../platform/agentHost/common/state/protocol/version/registry.js'; +import { IRemoteAgentHostService, RemoteAgentHostAutoConnectSettingId, RemoteAgentHostConnectionStatus, RemoteAgentHostsEnabledSettingId } from '../../../../../platform/agentHost/common/remoteAgentHostService.js'; +import { isTunnelHosted, ITunnelAgentHostService, TUNNEL_ADDRESS_PREFIX, TUNNEL_MIN_PROTOCOL_VERSION, type ITunnelInfo } from '../../../../../platform/agentHost/common/tunnelAgentHost.js'; import { IConfigurationService } from '../../../../../platform/configuration/common/configuration.js'; import { IInstantiationService } from '../../../../../platform/instantiation/common/instantiation.js'; import { ILogService } from '../../../../../platform/log/common/log.js'; @@ -20,7 +16,7 @@ import { IWorkbenchContribution, registerWorkbenchContribution2, WorkbenchPhase import { ITunnelHostService } from '../../../../../workbench/contrib/chat/common/tunnelHost.js'; import { AuthenticationSessionsChangeEvent, IAuthenticationService } from '../../../../../workbench/services/authentication/common/authentication.js'; import { IHostService } from '../../../../../workbench/services/host/browser/host.js'; -import { logTunnelConnectAttempt, logTunnelConnectResolved, logTunnelDiscoveryResult, TunnelConnectErrorCategory, TunnelConnectFailureReason, TunnelDiscoveryTrigger } from '../../../../common/sessionsTelemetry.js'; +import { logTunnelConnectAttempt, logTunnelConnectResolved, logTunnelDiscoveryResult, TunnelDiscoveryTrigger } from '../../../../common/sessionsTelemetry.js'; import { ISessionsProvidersService } from '../../../../services/sessions/browser/sessionsProvidersService.js'; import { IAgentHostFilterService } from '../../../../services/agentHostFilter/common/agentHostFilter.js'; import { RemoteAgentHostSessionsProvider } from './remoteAgentHostSessionsProvider.js'; @@ -29,11 +25,6 @@ import { watchForIncompatibleNotifications } from './remoteHostOptions.js'; /** Minimum interval between silent status checks (5 minutes). */ const STATUS_CHECK_INTERVAL = 5 * 60 * 1000; -/** Minimum gap between event-triggered reconnect resumes. */ -const RESUME_RATE_LIMIT_MS = 10_000; - -type TunnelReconnectTrigger = 'wake' | 'focus' | 'sessionAdded'; - export class TunnelAgentHostContribution extends Disposable implements IWorkbenchContribution { static readonly ID = 'sessions.contrib.tunnelAgentHostContribution'; @@ -42,6 +33,7 @@ export class TunnelAgentHostContribution extends Disposable implements IWorkbenc private readonly _providerInstances = new Map(); private readonly _pendingConnects = new Map>(); private _lastStatusCheck = 0; + private readonly _hostedTunnelSuppressions = new Set(); /** * `false` until the first {@link _silentStatusCheck} resolves. Until then * we keep newly-created providers in the `Connecting` state so the picker @@ -49,30 +41,7 @@ export class TunnelAgentHostContribution extends Disposable implements IWorkbenc */ private _initialStatusChecked = false; - /** Previous connection status per address — used to detect Connected→Disconnected transitions. */ - private readonly _previousStatuses = new Map(); - /** Pending auto-reconnect timer per address. */ - private readonly _reconnectTimeouts = new Map>(); - /** Consecutive failed auto-reconnect attempts per address. */ - private readonly _reconnectAttempts = new Map(); - /** Why auto-reconnect is paused for each address. */ - private readonly _reconnectPauseReasons = new Map(); - /** - * Addresses whose provider currently holds a live connection. Tracked - * separately from {@link _previousStatuses} so a drop is still detected when - * the connection passes through an intermediate `connecting` state on its - * way down. - */ private readonly _wiredAddresses = new Set(); - /** Timestamp of the last focus/wake-triggered resume, to rate-limit rapid tab toggles. */ - private _lastResumeAt = 0; - - /** - * Per-address connect sessions for telemetry. A session starts at the - * first attempt of a connect cycle (initial or reconnect) and ends on - * terminal resolution (connected, host-offline, max-attempts). - */ - private readonly _connectSessions = new Map(); constructor( @ITunnelAgentHostService private readonly _tunnelService: ITunnelAgentHostService, @@ -90,6 +59,7 @@ export class TunnelAgentHostContribution extends Disposable implements IWorkbenc ) { super(); + this._syncHostedTunnelSuppression(); // Create providers for cached tunnels this._reconcileProviders(); @@ -100,27 +70,24 @@ export class TunnelAgentHostContribution extends Disposable implements IWorkbenc // Update connection statuses when connections change this._register(this._remoteAgentHostService.onDidChangeConnections(() => { - this._handleConnectionChanges(); this._updateConnectionStatuses(); this._wireConnections(); })); // Reconcile providers when the tunnel cache changes this._register(this._tunnelService.onDidChangeTunnels(() => { + this._syncHostedTunnelSuppression(); this._reconcileProviders(); - // Stop any reconnect loops for tunnels that no longer exist - this._pruneReconnectState(); })); this._register(this._tunnelHostService.onDidChangeStatus(() => { - this._resetHostedTunnelReconnectState(); - this._silentStatusCheck(); + this._syncHostedTunnelSuppression(); + void this._silentStatusCheck(); })); this._register(this._configurationService.onDidChangeConfiguration(e => { if (e.affectsConfiguration(RemoteAgentHostsEnabledSettingId)) { this._reconcileProviders(); - this._pruneReconnectState(); } })); @@ -136,23 +103,9 @@ export class TunnelAgentHostContribution extends Disposable implements IWorkbenc this._register(this._hostService.onDidChangeFocus(focused => { if (focused) { - this._resumeReconnects('focus'); - } - })); - - // `online` is a browser-only network signal; focus above covers desktop. - if (isWeb) { - const onWake = () => this._resumeReconnects('wake'); - mainWindow.addEventListener('online', onWake); - this._register(toDisposable(() => mainWindow.removeEventListener('online', onWake))); - } - - // Cancel any pending reconnect timers on disposal. - this._register(toDisposable(() => { - for (const timer of this._reconnectTimeouts.values()) { - clearTimeout(timer); + void this._silentStatusCheck(); + this._requestServiceReconnects(); } - this._reconnectTimeouts.clear(); })); // Silently check status of cached tunnels on startup. Routed @@ -198,21 +151,33 @@ export class TunnelAgentHostContribution extends Disposable implements IWorkbenc } private _getProviderTunnels() { - return this._tunnelService.getCachedTunnels().filter(tunnel => !this._tunnelService.isAutoConnectSuppressed(tunnel.tunnelId)); + return this._tunnelService.getCachedTunnels().filter(tunnel => !this._tunnelService.isTunnelDismissed(tunnel.tunnelId)); } private _isHostedTunnel(tunnel: Pick): boolean { return isTunnelHosted(this._tunnelHostService.sharingInfo, tunnel); } - private _resetHostedTunnelReconnectState(): void { + private _syncHostedTunnelSuppression(): void { + const hostedTunnelIds = new Set(); for (const tunnel of this._tunnelService.getCachedTunnels()) { - if (this._isHostedTunnel(tunnel)) { - const address = `${TUNNEL_ADDRESS_PREFIX}${tunnel.tunnelId}`; - this._resetReconnectState(address); - if (this._remoteAgentHostService.connections.some(connection => connection.address === address && RemoteAgentHostConnectionStatus.isConnected(connection.status))) { - this._tunnelService.disconnect(address).catch(() => { /* best effort */ }); - } + if (!this._isHostedTunnel(tunnel)) { + continue; + } + hostedTunnelIds.add(tunnel.tunnelId); + if (!this._tunnelService.isAutoConnectSuppressed(tunnel.tunnelId)) { + this._hostedTunnelSuppressions.add(tunnel.tunnelId); + this._tunnelService.suppressAutoConnect(tunnel.tunnelId); + } + const address = `${TUNNEL_ADDRESS_PREFIX}${tunnel.tunnelId}`; + if (this._remoteAgentHostService.connections.some(connection => connection.address === address && RemoteAgentHostConnectionStatus.isConnected(connection.status))) { + void this._tunnelService.disconnect(address); + } + } + for (const tunnelId of this._hostedTunnelSuppressions) { + if (!hostedTunnelIds.has(tunnelId)) { + this._hostedTunnelSuppressions.delete(tunnelId); + this._tunnelService.clearAutoConnectSuppression(tunnelId); } } } @@ -220,8 +185,8 @@ export class TunnelAgentHostContribution extends Disposable implements IWorkbenc private _createProvider(address: string, name: string): void { const store = new DisposableStore(); const provider = this._instantiateProvider(address, name); - // Surface as "Connecting" until the first silent status check or an - // auto-connect attempt determines the real state; otherwise the picker + // Surface as "Connecting" until the first silent status check determines + // the real state; otherwise the picker // flashes "Offline" for every cached tunnel on startup. provider.setConnectionStatus(RemoteAgentHostConnectionStatus.connecting); store.add(provider); @@ -259,10 +224,7 @@ export class TunnelAgentHostContribution extends Disposable implements IWorkbenc provider.setConnectionStatus(connectionInfo.status); continue; } - // Preserve incompatible state set by `_connectTunnel`'s catch - // (where the failure happens before the service ever has an - // entry) until the user retries — otherwise the `finally` - // block would immediately overwrite it back to `disconnected`. + // The service retains incompatible connections for upgrade support. if (RemoteAgentHostConnectionStatus.isIncompatible(provider.connectionStatus.get())) { continue; } @@ -314,47 +276,14 @@ export class TunnelAgentHostContribution extends Disposable implements IWorkbenc } const tunnelId = address.slice(TUNNEL_ADDRESS_PREFIX.length); - const cached = this._tunnelService.getCachedTunnels().find(t => t.tunnelId === tunnelId); - if (!cached) { - return Promise.resolve(); - } - if (this._isHostedTunnel(cached)) { - this._resetReconnectState(address); - return Promise.resolve(); - } - if (!options.userInitiated && this._tunnelService.isAutoConnectSuppressed(tunnelId)) { - this._logService.info(`[TunnelAgentHost] Skipping background connect for user-disconnected tunnel ${address}`); - return Promise.resolve(); - } - const live = this._remoteAgentHostService.connections.find(connection => connection.address === address); - if (!options.userInitiated && RemoteAgentHostConnectionStatus.isConnecting(live?.status)) { - return Promise.resolve(); - } - if (!options.userInitiated && RemoteAgentHostConnectionStatus.isReconnecting(live?.status)) { - return Promise.resolve(); - } if (options.userInitiated) { - this._tunnelService.clearAutoConnectSuppression(tunnelId); - // Clear any sticky `incompatible` state so this attempt can - // transition through `connecting` and report a fresh result. - const provider = this._providerInstances.get(address); - if (provider && RemoteAgentHostConnectionStatus.isIncompatible(provider.connectionStatus.get())) { - provider.setConnectionStatus(RemoteAgentHostConnectionStatus.connecting); - } + this._tunnelService.clearTunnelDismissal(tunnelId); } - - // A new attempt is starting — cancel any scheduled reconnect timer; - // success/failure of this attempt will drive the next decision. - this._cancelReconnect(address); - - const { attemptNumber, attemptStart, session, isReconnect } = this._beginConnectAttempt(address); - + const cached = this._tunnelService.getCachedTunnels().find(t => t.tunnelId === tunnelId); + const attemptStart = Date.now(); const promise = (async () => { - // Show a progress notification after a short delay so quick - // connects don't flash a notification. Only show for user-initiated - // connects; background auto-connects and reconnects stay silent. let handle: { close(): void } | undefined; - const timer = options.userInitiated ? setTimeout(() => { + const timer = options.userInitiated && cached ? setTimeout(() => { handle = this._notificationService.notify({ severity: Severity.Info, message: nls.localize('tunnelConnecting', "Connecting to tunnel '{0}'...", cached.name), @@ -362,69 +291,26 @@ export class TunnelAgentHostContribution extends Disposable implements IWorkbenc }); }, 1000) : undefined; - this._updateConnectionStatuses(); try { + if (!cached || this._isHostedTunnel(cached)) { + return; + } const tunnelInfo: ITunnelInfo = { tunnelId: cached.tunnelId, clusterId: cached.clusterId, name: cached.name, tags: [], - protocolVersion: 5, + // Legacy cache fallback, not a real capability claim. + protocolVersion: cached.protocolVersion ?? TUNNEL_MIN_PROTOCOL_VERSION, hostConnectionCount: 0, }; await this._tunnelService.connect(tunnelInfo, cached.authProvider, { userInitiated: options.userInitiated }); - if (this._isHostedTunnel(cached)) { - await this._tunnelService.disconnect(address); - this._resetReconnectState(address); - return; - } - // Re-check after the await: the user may have disconnected this - // tunnel while this background connect was already in flight. - if (!options.userInitiated && this._tunnelService.isAutoConnectSuppressed(cached.tunnelId)) { - this._logService.info(`[TunnelAgentHost] Disconnecting background connection for user-disconnected tunnel ${address}`); - await this._tunnelService.disconnect(address); - this._connectSessions.delete(address); - return; - } - this._finishConnectAttempt(address, { success: true, attemptNumber, attemptStart, session, isReconnect }); + logTunnelConnectAttempt(this._telemetryService, { isReconnect: false, attempt: 1, durationMs: Date.now() - attemptStart, success: true }); + logTunnelConnectResolved(this._telemetryService, { isReconnect: false, totalAttempts: 1, totalDurationMs: Date.now() - attemptStart, success: true }); } catch (err) { - this._logService.warn(`[TunnelAgentHost] Connect to ${cached.name} failed:`, err); - const errorCategory = this._categorizeError(err); - this._finishConnectAttempt(address, { success: false, attemptNumber, attemptStart, session, isReconnect, error: err }); - // Clear the pending-connect entry BEFORE deciding what to do - // next; otherwise `_scheduleReconnect`'s in-flight guard - // (`_pendingConnects.has(address)`) would silently bail and - // we'd never re-arm the timer, leaving the tunnel stuck. - this._pendingConnects.delete(address); - - // Protocol version mismatch is a deterministic failure that - // cannot be fixed by retrying. Surface it on the provider so - // the workspace picker can show the host's message, and stop - // scheduling reconnects until the user manually retries via - // the picker's Manage menu. - const incompatible = RemoteAgentHostConnectionStatus.fromConnectError(err, [PROTOCOL_VERSION]); - if (incompatible) { - this._providerInstances.get(address)?.setConnectionStatus(incompatible); - this._resetReconnectState(address); - throw err; - } - - // Auth failures are not worth retrying — a fresh token must - // be acquired by the user or by a session-change event. Pause - // immediately and let `_handleSessionsChange` resume us when - // a new session appears. - if (errorCategory === 'authExpired' || errorCategory === 'auth') { - this._pauseReconnect(address, errorCategory); - throw err; - } - - const hostOnline = await this._probeHostOnline(cached.tunnelId); - if (hostOnline === false) { - this._pauseReconnect(address, 'hostOffline'); - } else { - this._logService.info(`[TunnelAgentHost] Scheduling reconnect for ${address}`); - this._scheduleReconnect(address); - } + this._logService.warn(`[TunnelAgentHost] Connect to ${cached?.name ?? address} failed:`, err); + logTunnelConnectAttempt(this._telemetryService, { isReconnect: false, attempt: 1, durationMs: Date.now() - attemptStart, success: false, errorCategory: 'other' }); + logTunnelConnectResolved(this._telemetryService, { isReconnect: false, totalAttempts: 1, totalDurationMs: Date.now() - attemptStart, success: false }); throw err; } finally { if (timer !== undefined) { @@ -436,401 +322,50 @@ export class TunnelAgentHostContribution extends Disposable implements IWorkbenc } })(); - // Swallow the promise rejection here so unhandled rejection noise - // doesn't bubble up for the background reconnect path; callers that - // await `_connectTunnel` directly will still see it via their own `await`. - promise.catch(() => { /* handled via _scheduleReconnect */ }); - this._pendingConnects.set(address, promise); return promise; } /** - * Tear down the active tunnel relay for {@link address} and cancel any - * pending auto-reconnect. The cached tunnel entry is kept so the user - * can re-connect later; only the live WebSocket is closed. + * Dismiss a tunnel from the remote-host picker and tear down its active relay. */ private async _disconnectTunnel(address: string): Promise { - this._cancelReconnect(address); - this._resetReconnectState(address); - this._tunnelService.suppressAutoConnect(address.slice(TUNNEL_ADDRESS_PREFIX.length)); - // Mark as explicitly disconnected so `_handleConnectionChanges` does - // not treat the impending Connected→(removed) transition as a - // reconnect-worthy drop. - this._previousStatuses.delete(address); + const tunnelId = address.slice(TUNNEL_ADDRESS_PREFIX.length); + this._tunnelService.dismissTunnel(tunnelId); + this._tunnelService.removeCachedTunnel(tunnelId); await this._tunnelService.disconnect(address); } - /** - * Detect tunnel connections that transitioned from Connected to - * Disconnected and schedule an auto-reconnect. - * - * Important: we only trigger on a Connected → Disconnected transition - * where the connection entry is still present. If the entry has been - * removed from the service (e.g. the user clicked "Remove Remote"), - * we do NOT schedule a reconnect — that would override their intent. - */ - private _handleConnectionChanges(): void { - if (!this._configurationService.getValue(RemoteAgentHostsEnabledSettingId)) { - return; - } - - const cachedAddresses = new Set(this._getProviderTunnels().map(t => `${TUNNEL_ADDRESS_PREFIX}${t.tunnelId}`)); - const currentStatuses = new Map(); - for (const conn of this._remoteAgentHostService.connections) { - currentStatuses.set(conn.address, conn.status); - } - - for (const address of cachedAddresses) { - const previous = this._previousStatuses.get(address); - const current = currentStatuses.get(address); - - // Only schedule a reconnect on an explicit Connected→Disconnected - // transition. If the address is absent from the connection list, - // the user (or another code path) removed it — honour that. - const wasConnected = RemoteAgentHostConnectionStatus.isConnected(previous); - const isExplicitlyDisconnected = RemoteAgentHostConnectionStatus.isDisconnected(current); - - if (wasConnected && isExplicitlyDisconnected && !this._pendingConnects.has(address)) { - this._logService.info(`[TunnelAgentHost] Connection lost for ${address}, scheduling reconnect`); - if (!this._connectSessions.has(address)) { - this._connectSessions.set(address, { startedAt: Date.now(), attempts: 0, isReconnect: true }); - } - this._scheduleReconnect(address, /*immediate*/ true); - } - - // Only track previous status while the entry is present so a - // future re-registration starts from a clean slate. If the - // entry disappeared (e.g. user-initiated removal), also cancel - // any already-scheduled reconnect and clear its backoff state - // so the removal is honoured even if a timer was already armed. - if (current !== undefined) { - this._previousStatuses.set(address, current); - } else { - this._previousStatuses.delete(address); - this._resetReconnectState(address); - } - } - - // Drop previous-status entries for addresses no longer cached. - for (const address of [...this._previousStatuses.keys()]) { - if (!cachedAddresses.has(address)) { - this._previousStatuses.delete(address); - } - } - } - - private _scheduleReconnect(address: string, immediate = false): void { - // Respect enablement and tunnel-still-cached. - if (!this._configurationService.getValue(RemoteAgentHostsEnabledSettingId)) { - return; - } - const tunnelId = address.slice(TUNNEL_ADDRESS_PREFIX.length); - const cached = this._tunnelService.getCachedTunnels().find(t => t.tunnelId === tunnelId); - if (!cached) { - return; - } - if (this._isHostedTunnel(cached)) { - this._resetReconnectState(address); - return; - } - - // Already connected or a connect is in flight — nothing to do. - if (this._pendingConnects.has(address)) { - return; - } - const live = this._remoteAgentHostService.connections.find(c => c.address === address); - if (live && RemoteAgentHostConnectionStatus.isConnected(live.status)) { - this._clearReconnectBackoff(address); - return; - } - if (live && RemoteAgentHostConnectionStatus.isConnecting(live.status)) { + private _requestServiceReconnects(): void { + if (!this._configurationService.getValue(RemoteAgentHostAutoConnectSettingId)) { return; } - if (live && RemoteAgentHostConnectionStatus.isReconnecting(live.status)) { - // The protocol client is preserving its state while it reconnects; don't replace it. - return; - } - - // Cancel any existing timer — we're rescheduling. - this._cancelReconnect(address); - - const attempt = this._reconnectAttempts.get(address) ?? 0; - - const reconnectPolicy = getEntryTypeConfig(RemoteAgentHostEntryType.Tunnel).reconnect; - if (hasExhaustedReconnectAttempts(reconnectPolicy, attempt)) { - this._pauseReconnect(address, 'maxAttemptsReached'); - return; - } - - const delay = immediate - ? 0 - : computeReconnectDelay(reconnectPolicy, attempt + 1); - - this._logService.info( - `[TunnelAgentHost] Scheduling reconnect for ${address} in ${delay}ms (attempt ${attempt + 1}/${reconnectPolicy.maxAttempts})` - ); - - const timer = setTimeout(() => { - this._reconnectTimeouts.delete(address); - - // A manual (or other) connect may have started or completed while - // we were waiting. Re-check before counting this as a new attempt, - // otherwise `_connectTunnel` would just return the in-flight promise - // and we'd inflate the backoff counter without really trying again. - if (this._pendingConnects.has(address)) { - return; - } - const live = this._remoteAgentHostService.connections.find(c => c.address === address); - if (live && RemoteAgentHostConnectionStatus.isConnected(live.status)) { - this._clearReconnectBackoff(address); - return; - } - if (live && RemoteAgentHostConnectionStatus.isConnecting(live.status)) { - return; - } - if (live && RemoteAgentHostConnectionStatus.isReconnecting(live.status)) { - // The protocol client is preserving its state while it reconnects; don't replace it. - return; - } - - this._reconnectAttempts.set(address, attempt + 1); - this._connectTunnel(address, { userInitiated: false }).catch(() => { /* _connectTunnel already re-schedules on failure */ }); - }, delay); - this._reconnectTimeouts.set(address, timer); - } - - /** - * Best-effort probe of whether the host backing `tunnelId` is online - * (has any host connections). Returns `undefined` if we couldn't - * determine — caller should treat as "retry normally" in that case. - */ - private async _probeHostOnline(tunnelId: string): Promise { - try { - const tunnels = await this._tunnelService.listTunnels({ silent: true }); - if (!tunnels) { - return undefined; + for (const tunnel of this._tunnelService.getCachedTunnels()) { + if (this._isHostedTunnel(tunnel) || this._tunnelService.isAutoConnectSuppressed(tunnel.tunnelId)) { + continue; } - const info = tunnels.find(t => t.tunnelId === tunnelId); - if (!info) { - return false; + const address = `${TUNNEL_ADDRESS_PREFIX}${tunnel.tunnelId}`; + const status = this._remoteAgentHostService.connections.find(connection => connection.address === address)?.status; + if (RemoteAgentHostConnectionStatus.isConnected(status) + || RemoteAgentHostConnectionStatus.isConnecting(status) + || RemoteAgentHostConnectionStatus.isReconnecting(status) + || RemoteAgentHostConnectionStatus.isIncompatible(status)) { + continue; } - return info.hostConnectionCount > 0; - } catch { - return undefined; - } - } - - private _cancelReconnect(address: string): void { - const timer = this._reconnectTimeouts.get(address); - if (timer !== undefined) { - clearTimeout(timer); - this._reconnectTimeouts.delete(address); + this._remoteAgentHostService.reconnect(address, false); } } - /** Clear retry-backoff and pause state for an address. */ - private _clearReconnectBackoff(address: string): void { - this._reconnectAttempts.delete(address); - this._reconnectPauseReasons.delete(address); - } - - /** Drop all reconnect + telemetry state for an address (e.g. on removal). */ - private _resetReconnectState(address: string): void { - this._cancelReconnect(address); - this._clearReconnectBackoff(address); - this._connectSessions.delete(address); - } - - /** - * React to auth session add/remove. Additions re-run discovery (a fresh - * token may unblock a previously auth-paused tunnel). Removals drop any - * tunnel state that depended on that provider — otherwise we'd sit on a - * stale auth pause forever, or hammer a provider whose session is gone. - */ private _handleSessionsChange(e: { providerId: string; label: string; event: AuthenticationSessionsChangeEvent }): void { - const added = (e.event.added?.length ?? 0) > 0; - const removed = (e.event.removed?.length ?? 0) > 0; - - if (removed) { - const cached = this._tunnelService.getCachedTunnels(); - for (const tunnel of cached) { - if (tunnel.authProvider !== e.providerId) { - continue; + if ((e.event.removed?.length ?? 0) > 0) { + for (const tunnel of this._tunnelService.getCachedTunnels()) { + if (tunnel.authProvider === e.providerId) { + void this._tunnelService.disconnect(`${TUNNEL_ADDRESS_PREFIX}${tunnel.tunnelId}`); } - const address = `${TUNNEL_ADDRESS_PREFIX}${tunnel.tunnelId}`; - this._logService.info( - `[TunnelAgentHost] Auth session removed for ${e.providerId}; tearing down ${address}.` - ); - this._resetReconnectState(address); - // Best-effort disconnect — the transport may already be dead. - this._tunnelService.disconnect(address).catch(() => { /* ignore */ }); } } - - if (added) { - this._logService.info(`[TunnelAgentHost] ${e.providerId} session added; resuming reconnects and rediscovering.`); - this._resumeReconnects('sessionAdded'); - this._silentStatusCheck('sessionChange'); - } - } - - /** - * Stop auto-reconnecting for an address until a recovery signal resumes us. - */ - private _pauseReconnect(address: string, reason: TunnelConnectFailureReason): void { - this._cancelReconnect(address); - this._reconnectAttempts.delete(address); - this._reconnectPauseReasons.set(address, reason); - const resumeCondition = reason === 'hostOffline' - ? 'a status check that confirms the host is online' - : reason === 'auth' || reason === 'authExpired' - ? 'an authentication session change' - : `${isWeb ? 'network-online or ' : ''}window focus`; - this._logService.info( - `[TunnelAgentHost] Pausing auto-reconnect for ${address} (${reason}); ` + - `will resume on ${resumeCondition}.` - ); - const session = this._connectSessions.get(address); - if (session) { - logTunnelConnectResolved(this._telemetryService, { - isReconnect: session.isReconnect, - totalAttempts: session.attempts, - totalDurationMs: Date.now() - session.startedAt, - success: false, - failureReason: reason, - }); - this._connectSessions.delete(address); - } - } - - /** - * Begin (or continue) a connect telemetry session for `address` and - * return the bookkeeping needed to later finish the attempt. A session - * already exists if `_handleConnectionChanges` marked this as a - * reconnect cycle; otherwise this starts a fresh initial-connect session. - */ - private _beginConnectAttempt(address: string): { session: { startedAt: number; attempts: number; isReconnect: boolean }; attemptNumber: number; attemptStart: number; isReconnect: boolean } { - let session = this._connectSessions.get(address); - if (!session) { - session = { startedAt: Date.now(), attempts: 0, isReconnect: false }; - this._connectSessions.set(address, session); - } - session.attempts++; - return { session, attemptNumber: session.attempts, attemptStart: Date.now(), isReconnect: session.isReconnect }; - } - - /** - * Finalize the telemetry for a single connect attempt. On success, also - * clears backoff state and closes the session; on failure, only the - * per-attempt event is emitted (the caller decides whether to retry). - */ - private _finishConnectAttempt(address: string, args: { - success: boolean; - attemptNumber: number; - attemptStart: number; - session: { startedAt: number; attempts: number; isReconnect: boolean }; - isReconnect: boolean; - error?: unknown; - }): void { - const { success, attemptNumber, attemptStart, session, isReconnect, error } = args; - const durationMs = Date.now() - attemptStart; - if (success) { - this._clearReconnectBackoff(address); - logTunnelConnectAttempt(this._telemetryService, { isReconnect, attempt: attemptNumber, durationMs, success: true }); - logTunnelConnectResolved(this._telemetryService, { isReconnect, totalAttempts: attemptNumber, totalDurationMs: Date.now() - session.startedAt, success: true }); - this._connectSessions.delete(address); - } else { - logTunnelConnectAttempt(this._telemetryService, { isReconnect, attempt: attemptNumber, durationMs, success: false, errorCategory: this._categorizeError(error) }); - } - } - - private _categorizeError(err: unknown): TunnelConnectErrorCategory { - const message = err instanceof Error ? err.message : String(err); - // Expired / invalid credential — callers short-circuit this category - // to avoid burning retry budget on a token the user has to refresh. - if (/\b(401|403)\b|token.*expired|expired.*token|invalid[_ -]?grant/i.test(message)) { - return 'authExpired'; - } - // Match authentication-specific language but NOT "connection token" - // or other protocol uses of the word "token". - if (/authenticat|unauthoriz|auth.*(fail|error|invalid)/i.test(message)) { - return 'auth'; - } - if (/WebSocket relay connection failed|failed to connect to relay/i.test(message)) { - return 'relayConnectionFailed'; - } - if (/network|fetch|offline|ECONN|ENOTFOUND|ETIMEDOUT/i.test(message)) { - return 'network'; - } - return 'other'; - } - - /** - * Resume paused reconnects that the given recovery signal can resolve. - * - * Rate-limited: at most one resume per RESUME_RATE_LIMIT_MS so that - * rapid focus/network events cannot start unbounded retry bursts. - */ - private _resumeReconnects(trigger: TunnelReconnectTrigger): void { - if (!this._configurationService.getValue(RemoteAgentHostsEnabledSettingId)) { - return; - } - - const resumableAddresses: string[] = []; - for (const tunnel of this._getProviderTunnels()) { - const address = `${TUNNEL_ADDRESS_PREFIX}${tunnel.tunnelId}`; - const reason = this._reconnectPauseReasons.get(address); - if (!reason || !this._canResumeReconnect(reason, trigger) || this._pendingConnects.has(address)) { - continue; - } - const live = this._remoteAgentHostService.connections.find(connection => connection.address === address); - if (live && RemoteAgentHostConnectionStatus.isReconnecting(live.status)) { - // The protocol client is preserving its state while it reconnects; don't replace it. - continue; - } - if (!live || !RemoteAgentHostConnectionStatus.isConnected(live.status)) { - resumableAddresses.push(address); - } - } - if (resumableAddresses.length === 0) { - return; - } - - if (trigger !== 'sessionAdded') { - const now = Date.now(); - if (now - this._lastResumeAt < RESUME_RATE_LIMIT_MS) { - return; - } - this._lastResumeAt = now; - } - - for (const address of resumableAddresses) { - this._logService.info(`[TunnelAgentHost] Resuming reconnect for ${address} (trigger: ${trigger})`); - this._clearReconnectBackoff(address); - this._scheduleReconnect(address, /*immediate*/ true); - } - } - - private _canResumeReconnect(reason: TunnelConnectFailureReason, trigger: TunnelReconnectTrigger): boolean { - return trigger === 'sessionAdded' - ? reason === 'auth' || reason === 'authExpired' - : reason === 'maxAttemptsReached'; - } - - /** Drop reconnect state for addresses whose tunnel is no longer cached. */ - private _pruneReconnectState(): void { - const cachedAddresses = new Set(this._getProviderTunnels().map(t => `${TUNNEL_ADDRESS_PREFIX}${t.tunnelId}`)); - const tracked = new Set([ - ...this._reconnectTimeouts.keys(), - ...this._reconnectAttempts.keys(), - ...this._reconnectPauseReasons.keys(), - ...this._connectSessions.keys(), - ]); - for (const address of tracked) { - if (!cachedAddresses.has(address)) { - this._resetReconnectState(address); - } + if ((e.event.added?.length ?? 0) > 0) { + void this._silentStatusCheck('sessionChange'); } } @@ -896,16 +431,13 @@ export class TunnelAgentHostContribution extends Disposable implements IWorkbenc // can match these tunnels for teardown on session removal. const cachedIds = new Set(cached.map(t => t.tunnelId)); for (const tunnel of onlineTunnels) { - if (!cachedIds.has(tunnel.tunnelId)) { + if (!cachedIds.has(tunnel.tunnelId) && !this._tunnelService.isTunnelDismissed(tunnel.tunnelId)) { this._tunnelService.cacheTunnel(tunnel, 'github'); } } - // Update online/offline status based on hostConnectionCount. - // For tunnels, Connected means "host is online" (clickable to connect), - // Disconnected means "host is offline". Actual relay connection - // establishment happens when the user clicks the tunnel (or via - // auto-connect below when enabled). + // Update online/offline status based on hostConnectionCount for + // tunnels that do not currently have a service-owned connection. const onlineTunnelMap = new Map(onlineTunnels.map(t => [t.tunnelId, t])); for (const [address, provider] of this._providerInstances) { // Skip tunnels that already have an active relay connection @@ -922,13 +454,6 @@ export class TunnelAgentHostContribution extends Disposable implements IWorkbenc if (info && info.hostConnectionCount > 0) { provider.setConnectionStatus(RemoteAgentHostConnectionStatus.connected); - if (this._reconnectPauseReasons.get(address) === 'hostOffline') { - this._logService.info( - `[TunnelAgentHost] Confirmed host online for paused ${address}; auto-resuming reconnect.` - ); - this._clearReconnectBackoff(address); - this._scheduleReconnect(address, /*immediate*/ true); - } } else { provider.setConnectionStatus(RemoteAgentHostConnectionStatus.disconnected); // Host is not online — drop any cached sessions we were @@ -937,37 +462,6 @@ export class TunnelAgentHostContribution extends Disposable implements IWorkbenc } } - // Auto-connect online tunnels that aren't connected yet when the - // user has opted into auto-connect (default on). This mirrors the - // web embedder behaviour where no workspace picker is available - // to trigger manual connection. - const autoConnect = this._configurationService.getValue(RemoteAgentHostAutoConnectSettingId); - if (autoConnect) { - for (const tunnel of onlineTunnels) { - if (tunnel.hostConnectionCount > 0 && !this._isHostedTunnel(tunnel)) { - const address = `${TUNNEL_ADDRESS_PREFIX}${tunnel.tunnelId}`; - if (this._tunnelService.isAutoConnectSuppressed(tunnel.tunnelId)) { - continue; - } - if (this._reconnectPauseReasons.has(address)) { - continue; - } - // A reconnecting protocol client is already restoring this relay. - const alreadyConnected = this._remoteAgentHostService.connections.some( - c => c.address === address && (RemoteAgentHostConnectionStatus.isConnected(c.status) || RemoteAgentHostConnectionStatus.isReconnecting(c.status)) - ); - if (!alreadyConnected) { - const mode = this._tunnelService.getAutoConnectMode(tunnel); - if (mode === 'prompt') { - this._logService.info(`[TunnelAgentHost] Prompting for the initial agent host location for ${address}`); - this._connectTunnel(address, { userInitiated: true }); - } else { - this._connectTunnel(address, { userInitiated: false }); - } - } - } - } - } } this._initialStatusChecked = true; diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/tunnelAgentHostStorage.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/tunnelAgentHostStorage.ts index 5542ee920e1fa9..de7124299d7248 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/tunnelAgentHostStorage.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/tunnelAgentHostStorage.ts @@ -5,103 +5,128 @@ import { Emitter, Event } from '../../../../../base/common/event.js'; import { Disposable } from '../../../../../base/common/lifecycle.js'; +import { autorun, IObservable } from '../../../../../base/common/observable.js'; import { type ICachedTunnel } from '../../../../../platform/agentHost/common/tunnelAgentHost.js'; +import { observableMemento, ObservableMemento } from '../../../../../platform/observable/common/observableMemento.js'; import { IStorageService, StorageScope, StorageTarget } from '../../../../../platform/storage/common/storage.js'; const CACHED_TUNNELS_KEY = 'tunnelAgentHost.recentTunnels'; +const DISMISSED_TUNNELS_KEY = 'tunnelAgentHost.dismissedTunnels'; const AUTO_CONNECT_SUPPRESSED_TUNNELS_KEY = 'tunnelAgentHost.autoConnectSuppressedTunnels'; -/** Persists the tunnel cache and explicit auto-connect suppressions shared by browser tunnel services. */ +const cachedTunnelMemento = observableMemento({ + defaultValue: [], + key: CACHED_TUNNELS_KEY, + toStorage: tunnels => JSON.stringify(tunnels), + fromStorage: value => JSON.parse(value) as readonly ICachedTunnel[], +}); + +const autoConnectSuppressedTunnelMemento = observableMemento({ + defaultValue: [], + key: AUTO_CONNECT_SUPPRESSED_TUNNELS_KEY, + toStorage: tunnelIds => JSON.stringify(tunnelIds), + fromStorage: value => { + const parsed: unknown = JSON.parse(value); + return Array.isArray(parsed) ? parsed.filter((item): item is string => typeof item === 'string') : []; + }, +}); + +const dismissedTunnelMemento = observableMemento({ + defaultValue: [], + key: DISMISSED_TUNNELS_KEY, + toStorage: tunnelIds => JSON.stringify(tunnelIds), + fromStorage: value => { + const parsed: unknown = JSON.parse(value); + return Array.isArray(parsed) ? parsed.filter((item): item is string => typeof item === 'string') : []; + }, +}); + +/** Persists the tunnel cache, picker dismissals, and auto-connect suppressions shared by browser tunnel services. */ export class TunnelAgentHostStorage extends Disposable { private readonly _onDidChangeTunnels = this._register(new Emitter()); readonly onDidChangeTunnels: Event = this._onDidChangeTunnels.event; + private readonly _cachedTunnels: ObservableMemento; + private readonly _dismissedTunnels: ObservableMemento; + private readonly _autoConnectSuppressedTunnels: ObservableMemento; + + /** Cached tunnels, persisted across windows. */ + readonly cachedTunnels: IObservable; + /** Tunnel IDs explicitly dismissed from the remote-host picker. */ + readonly dismissedTunnels: IObservable; + /** Tunnel IDs whose automatic reconnect is suppressed. */ + readonly autoConnectSuppressedTunnels: IObservable; + constructor( - @IStorageService private readonly _storageService: IStorageService, + @IStorageService storageService: IStorageService, ) { super(); + this._cachedTunnels = this._register(cachedTunnelMemento(StorageScope.APPLICATION, StorageTarget.USER, storageService)); + this._dismissedTunnels = this._register(dismissedTunnelMemento(StorageScope.APPLICATION, StorageTarget.USER, storageService)); + this._autoConnectSuppressedTunnels = this._register(autoConnectSuppressedTunnelMemento(StorageScope.APPLICATION, StorageTarget.USER, storageService)); + this.cachedTunnels = this._cachedTunnels; + this.dismissedTunnels = this._dismissedTunnels; + this.autoConnectSuppressedTunnels = this._autoConnectSuppressedTunnels; + this._register(autorun(reader => { + this.cachedTunnels.read(reader); + this.dismissedTunnels.read(reader); + this.autoConnectSuppressedTunnels.read(reader); + this._onDidChangeTunnels.fire(); + })); } getCachedTunnels(): ICachedTunnel[] { - const raw = this._storageService.get(CACHED_TUNNELS_KEY, StorageScope.APPLICATION); - if (!raw) { - return []; - } - try { - return JSON.parse(raw); - } catch { - return []; - } + return [...this._cachedTunnels.get()]; } cacheTunnel(tunnel: ICachedTunnel): void { - const cached = this.getCachedTunnels(); - const filtered = cached.filter(candidate => candidate.tunnelId !== tunnel.tunnelId); - filtered.unshift(tunnel); + const cached = this._cachedTunnels.get(); this.clearAutoConnectSuppression(tunnel.tunnelId); - this._storeCachedTunnels(filtered); - this._onDidChangeTunnels.fire(); + this._cachedTunnels.set([tunnel, ...cached.filter(candidate => candidate.tunnelId !== tunnel.tunnelId)], undefined); } removeCachedTunnel(tunnelId: string): void { - const cached = this.getCachedTunnels(); - this._storeCachedTunnels(cached.filter(tunnel => tunnel.tunnelId !== tunnelId)); + this._cachedTunnels.set(this._cachedTunnels.get().filter(tunnel => tunnel.tunnelId !== tunnelId), undefined); this.clearAutoConnectSuppression(tunnelId); - this._onDidChangeTunnels.fire(); } - isAutoConnectSuppressed(tunnelId: string): boolean { - return this._getAutoConnectSuppressedTunnels().has(tunnelId); + isTunnelDismissed(tunnelId: string): boolean { + return this._dismissedTunnels.get().includes(tunnelId); } - suppressAutoConnect(tunnelId: string): void { - const suppressed = this._getAutoConnectSuppressedTunnels(); - suppressed.add(tunnelId); - this._storeAutoConnectSuppressedTunnels(suppressed); + dismissTunnel(tunnelId: string): void { + const dismissed = this._dismissedTunnels.get(); + this._dismissedTunnels.set( + dismissed.includes(tunnelId) ? [...dismissed] : [...dismissed, tunnelId], + undefined, + ); } - clearAutoConnectSuppression(tunnelId: string): void { - const suppressed = this._getAutoConnectSuppressedTunnels(); - if (!suppressed.delete(tunnelId)) { + clearTunnelDismissal(tunnelId: string): void { + const dismissed = this._dismissedTunnels.get(); + if (!dismissed.includes(tunnelId)) { return; } - this._storeAutoConnectSuppressedTunnels(suppressed); + this._dismissedTunnels.set(dismissed.filter(id => id !== tunnelId), undefined); } - /** Notifies consumers that a tunnel connection changed without changing its cache entry. */ - notifyTunnelsChanged(): void { - this._onDidChangeTunnels.fire(); - } - - private _storeCachedTunnels(tunnels: ICachedTunnel[]): void { - if (tunnels.length === 0) { - this._storageService.remove(CACHED_TUNNELS_KEY, StorageScope.APPLICATION); - } else { - this._storageService.store(CACHED_TUNNELS_KEY, JSON.stringify(tunnels), StorageScope.APPLICATION, StorageTarget.USER); - } + isAutoConnectSuppressed(tunnelId: string): boolean { + return this._autoConnectSuppressedTunnels.get().includes(tunnelId); } - private _getAutoConnectSuppressedTunnels(): Set { - const raw = this._storageService.get(AUTO_CONNECT_SUPPRESSED_TUNNELS_KEY, StorageScope.APPLICATION); - if (!raw) { - return new Set(); - } - try { - const parsed: unknown = JSON.parse(raw); - if (!Array.isArray(parsed)) { - return new Set(); - } - return new Set(parsed.filter(item => typeof item === 'string')); - } catch { - return new Set(); - } + suppressAutoConnect(tunnelId: string): void { + const suppressed = this._autoConnectSuppressedTunnels.get(); + this._autoConnectSuppressedTunnels.set( + suppressed.includes(tunnelId) ? [...suppressed] : [...suppressed, tunnelId], + undefined, + ); } - private _storeAutoConnectSuppressedTunnels(tunnelIds: Set): void { - if (tunnelIds.size === 0) { - this._storageService.remove(AUTO_CONNECT_SUPPRESSED_TUNNELS_KEY, StorageScope.APPLICATION); - } else { - this._storageService.store(AUTO_CONNECT_SUPPRESSED_TUNNELS_KEY, JSON.stringify([...tunnelIds]), StorageScope.APPLICATION, StorageTarget.USER); + clearAutoConnectSuppression(tunnelId: string): void { + const suppressed = this._autoConnectSuppressedTunnels.get(); + if (!suppressed.includes(tunnelId)) { + return; } + this._autoConnectSuppressedTunnels.set(suppressed.filter(id => id !== tunnelId), undefined); } } diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/webSocketAgentHost.contribution.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/webSocketAgentHost.contribution.ts new file mode 100644 index 00000000000000..fe1f5ced78077e --- /dev/null +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/webSocketAgentHost.contribution.ts @@ -0,0 +1,48 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { type IRemoteAgentHostEntry, IRemoteAgentHostService, RemoteAgentHostEntryType, RemoteAgentHostsEnabledSettingId, RemoteAgentHostsSettingId } from '../../../../../platform/agentHost/common/remoteAgentHostService.js'; +import { IConfigurationService } from '../../../../../platform/configuration/common/configuration.js'; +import { IInstantiationService } from '../../../../../platform/instantiation/common/instantiation.js'; +import { INotificationService } from '../../../../../platform/notification/common/notification.js'; +import { IWorkbenchContribution, registerWorkbenchContribution2, WorkbenchPhase } from '../../../../../workbench/common/contributions.js'; +import { ISessionsProvidersService } from '../../../../services/sessions/browser/sessionsProvidersService.js'; +import { EntryDrivenProviderContribution } from './entryDrivenProviderContribution.js'; + +export class WebSocketAgentHostContribution extends EntryDrivenProviderContribution implements IWorkbenchContribution { + + static readonly ID = 'sessions.contrib.webSocketAgentHostContribution'; + + protected readonly _entryType = RemoteAgentHostEntryType.WebSocket; + + protected override get _clearConnectionOnRemoval(): boolean { + return true; + } + + constructor( + @IRemoteAgentHostService remoteAgentHostService: IRemoteAgentHostService, + @IConfigurationService configurationService: IConfigurationService, + @IInstantiationService instantiationService: IInstantiationService, + @ISessionsProvidersService sessionsProvidersService: ISessionsProvidersService, + @INotificationService notificationService: INotificationService, + ) { + super(remoteAgentHostService, configurationService, instantiationService, sessionsProvidersService, notificationService); + + this._register(this._remoteAgentHostService.onDidChangeConnections(() => this._reconcile())); + this._register(this._configurationService.onDidChangeConfiguration(e => { + if (e.affectsConfiguration(RemoteAgentHostsSettingId) || e.affectsConfiguration(RemoteAgentHostsEnabledSettingId)) { + this._reconcile(); + } + })); + + this._reconcile(); + } + + protected override _getProviderOptions(_entry: IRemoteAgentHostEntry) { + return {}; + } +} + +registerWorkbenchContribution2(WebSocketAgentHostContribution.ID, WebSocketAgentHostContribution, WorkbenchPhase.AfterRestored); diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/webTunnelAgentHostService.contribution.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/webTunnelAgentHostService.contribution.ts index e1fdbc042280af..6db0669ae06d49 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/webTunnelAgentHostService.contribution.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/webTunnelAgentHostService.contribution.ts @@ -102,6 +102,18 @@ class BrowserTunnelAgentHostServiceSelector extends Disposable implements ITunne this._delegate.removeCachedTunnel(tunnelId); } + isTunnelDismissed(tunnelId: string): boolean { + return this._delegate.isTunnelDismissed(tunnelId); + } + + dismissTunnel(tunnelId: string): void { + this._delegate.dismissTunnel(tunnelId); + } + + clearTunnelDismissal(tunnelId: string): void { + this._delegate.clearTunnelDismissal(tunnelId); + } + isAutoConnectSuppressed(tunnelId: string): boolean { return this._delegate.isAutoConnectSuppressed(tunnelId); } diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/webTunnelAgentHostService.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/webTunnelAgentHostService.ts index a53db54d3f795a..43f00e5cfdb140 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/webTunnelAgentHostService.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/webTunnelAgentHostService.ts @@ -5,14 +5,13 @@ import { Emitter, Event } from '../../../../../base/common/event.js'; import { Disposable } from '../../../../../base/common/lifecycle.js'; +import { derived, IObservable, observableSignalFromEvent } from '../../../../../base/common/observable.js'; import { AgentHostProtocolClient } from '../../../../../platform/agentHost/browser/agentHostProtocolClient.js'; import { agentsWindowAgentHostClientInfo } from '../../../../../platform/agentHost/common/agentHostClientInfo.js'; import { AgentHostClientConnectionKind } from '../../../../../platform/agentHost/common/agentHostTelemetry.js'; import { ReconnectingTransport, type IEstablishedTransport } from '../../../../../platform/agentHost/common/reconnectingTransport.js'; import { NonReconnectableTransportError, type IProtocolTransport } from '../../../../../platform/agentHost/common/state/sessionTransport.js'; -import { deriveConnectionToken } from '../../../../../platform/agentHost/common/tunnelAgentHostConnector.js'; -import { RemoteAgentHostEntryType, IRemoteAgentHostService, RemoteAgentHostConnectionStatus, RemoteAgentHostsEnabledSettingId } from '../../../../../platform/agentHost/common/remoteAgentHostService.js'; -import { PROTOCOL_VERSION } from '../../../../../platform/agentHost/common/state/protocol/version/registry.js'; +import { RemoteAgentHostEntryType, IRemoteAgentHostService, RemoteAgentHostsEnabledSettingId, getEntryAddress, type IRemoteAgentHostConnectOptions, type IRemoteAgentHostConnectionFactory, type IRemoteAgentHostCreatedConnection, type IRemoteAgentHostEntry } from '../../../../../platform/agentHost/common/remoteAgentHostService.js'; import type { ProtocolMessage, AhpServerNotification, JsonRpcResponse } from '../../../../../platform/agentHost/common/state/sessionProtocol.js'; import { MALFORMED_FRAMES_FORCE_CLOSE_THRESHOLD, MALFORMED_FRAMES_LOG_CAP } from '../../../../../platform/agentHost/common/transportConstants.js'; import { @@ -36,6 +35,83 @@ import { TunnelAgentHostStorage } from './tunnelAgentHostStorage.js'; const LOG_PREFIX = '[WebTunnelAgentHost]'; +class WebTunnelConnectionFactory extends Disposable implements IRemoteAgentHostConnectionFactory { + readonly kind = RemoteAgentHostEntryType.Tunnel; + readonly entries: IObservable; + + private readonly _onDidStageTunnel = this._register(new Emitter()); + private readonly _stagedAuthProviders = new Map(); + /** + * Initiation mode for a staged tunnel, consumed by the first + * {@link createConnection} for that address. Staging publishes the entry + * synchronously, so the service's reconciliation can begin dialing before + * the caller's explicit `reconnect` runs, and that dial would otherwise be + * reported as background. The embedder's discovery provider owns + * interaction today, so this only keeps the three tunnel factories + * behaving identically. + */ + private readonly _stagedUserInitiated = new Map(); + private readonly _onDidStageTunnelSignal = observableSignalFromEvent(this, this._onDidStageTunnel.event); + + constructor( + private readonly _storage: TunnelAgentHostStorage, + private readonly _createConnection: (entry: IRemoteAgentHostEntry, options: IRemoteAgentHostConnectOptions) => Promise, + ) { + super(); + this.entries = derived(this, reader => { + this._onDidStageTunnelSignal.read(reader); + const autoConnectSuppressedTunnels = this._storage.autoConnectSuppressedTunnels.read(reader); + return this._storage.cachedTunnels.read(reader) + .filter(tunnel => !autoConnectSuppressedTunnels.includes(tunnel.tunnelId)) + .map(tunnel => this._entryForTunnel(tunnel, tunnel.authProvider)); + }); + } + + stageTunnel(tunnel: ITunnelInfo, authProvider?: 'github' | 'microsoft', userInitiated = true): IRemoteAgentHostEntry { + const address = `${TUNNEL_ADDRESS_PREFIX}${tunnel.tunnelId}`; + this._stagedAuthProviders.set(address, authProvider); + this._stagedUserInitiated.set(address, userInitiated); + this._storage.cacheTunnel({ tunnelId: tunnel.tunnelId, clusterId: tunnel.clusterId, name: tunnel.name, protocolVersion: tunnel.protocolVersion, authProvider }); + this._onDidStageTunnel.fire(); + return this._entryForTunnel(tunnel, authProvider); + } + + unstageTunnel(address: string): void { + this._stagedUserInitiated.delete(address); + if (this._stagedAuthProviders.delete(address)) { + this._onDidStageTunnel.fire(); + } + } + + createConnection(entry: IRemoteAgentHostEntry, options: IRemoteAgentHostConnectOptions): Promise { + if (entry.connection.type !== RemoteAgentHostEntryType.Tunnel) { + throw new Error(`Tunnel factory cannot create a ${entry.connection.type} connection.`); + } + const address = getEntryAddress(entry); + const stagedUserInitiated = this._stagedUserInitiated.get(address); + // Consume it: only the connect this staging was for is user-initiated, + // and a later automatic reconnect must not prompt. + this._stagedUserInitiated.delete(address); + const connectOptions = stagedUserInitiated === undefined + ? options + : { ...options, userInitiated: stagedUserInitiated }; + return this._createConnection(entry, connectOptions); + } + + private _entryForTunnel(tunnel: Pick, authProvider?: 'github' | 'microsoft'): IRemoteAgentHostEntry { + return { + name: tunnel.name, + connection: { + type: RemoteAgentHostEntryType.Tunnel, + tunnelId: tunnel.tunnelId, + clusterId: tunnel.clusterId, + label: tunnel.name, + authProvider, + }, + }; + } +} + /** * Web (browser) implementation of {@link ITunnelAgentHostService}. * @@ -52,6 +128,7 @@ export class WebTunnelAgentHostService extends Disposable implements ITunnelAgen declare readonly _serviceBrand: undefined; private readonly _storage: TunnelAgentHostStorage; + private readonly _connectionFactory: WebTunnelConnectionFactory; readonly onDidChangeTunnels: Event; private readonly _discoveryProvider: ITunnelDiscoveryProvider | undefined; @@ -69,6 +146,11 @@ export class WebTunnelAgentHostService extends Disposable implements ITunnelAgen this._storage = this._register(new TunnelAgentHostStorage(this._storageService)); this.onDidChangeTunnels = this._storage.onDidChangeTunnels; this._discoveryProvider = environmentService.options?.tunnelDiscoveryProvider; + this._connectionFactory = this._register(new WebTunnelConnectionFactory( + this._storage, + (entry, options) => this._createConnection(entry, options), + )); + this._register(this._remoteAgentHostService.registerConnectionFactory(this._connectionFactory)); if (!this._discoveryProvider) { this._logService.debug(`${LOG_PREFIX} No tunnelDiscoveryProvider — tunnel discovery disabled`); } @@ -142,39 +224,53 @@ export class WebTunnelAgentHostService extends Disposable implements ITunnelAgen // Connection (via embedder) - async connect(tunnel: ITunnelInfo, authProvider?: 'github' | 'microsoft'): Promise { - if (!this._discoveryProvider) { - throw new Error('No tunnelDiscoveryProvider available'); - } + async connect(tunnel: ITunnelInfo, authProvider?: 'github' | 'microsoft', options?: { readonly userInitiated?: boolean }): Promise { if (!this._configurationService.getValue(RemoteAgentHostsEnabledSettingId)) { throw new Error('Remote agent host connections are not enabled.'); } - const { tunnelId, clusterId } = tunnel; - this._logService.info(`${LOG_PREFIX} Connecting to tunnel '${tunnel.name}' (${tunnelId})`); + const entry = this._connectionFactory.stageTunnel(tunnel, authProvider, options?.userInitiated ?? true); + const address = getEntryAddress(entry); + this._remoteAgentHostService.reconnect(address, options?.userInitiated ?? true); + await this._remoteAgentHostService.waitForConnection(address); + } - // The embedder handles the full connection including auth - const connection = await this._discoveryProvider.connect(tunnelId, clusterId); + private async _createConnection(entry: IRemoteAgentHostEntry, _options: IRemoteAgentHostConnectOptions): Promise { + if (entry.connection.type !== RemoteAgentHostEntryType.Tunnel) { + throw new Error(`Tunnel factory cannot create a ${entry.connection.type} connection.`); + } + const discoveryProvider = this._discoveryProvider; + if (!discoveryProvider) { + throw new NonReconnectableTransportError('No tunnel discovery provider is available to connect.'); + } - // Derive connection token from tunnel ID (same convention as CLI and desktop) - const connectionToken = await deriveConnectionToken(tunnelId); + const { tunnelId, clusterId } = entry.connection; + const address = getEntryAddress(entry); + this._logService.info(`${LOG_PREFIX} Connecting to tunnel '${entry.name}' (${tunnelId})`); + let connection: ITunnelConnection; + try { + connection = await discoveryProvider.connect(tunnelId, clusterId); + } catch (error) { + if (isTunnelNotFoundError(error)) { + throw new NonReconnectableTransportError(error.message); + } + throw error; + } - const address = `${TUNNEL_ADDRESS_PREFIX}${tunnelId}`; let useSeedConnection = true; const establish = async (): Promise => { if (useSeedConnection) { useSeedConnection = false; - // The initial connection is already owned by the transport established for this managed connection. return { transport: new TunnelConnectionTransport(connection, this._logService) }; } - const discoveryProvider = this._discoveryProvider; - if (!discoveryProvider) { + const reconnectProvider = this._discoveryProvider; + if (!reconnectProvider) { throw new NonReconnectableTransportError('No tunnel discovery provider is available to reconnect.'); } try { - const reconnected = await discoveryProvider.connect(tunnelId, clusterId); + const reconnected = await reconnectProvider.connect(tunnelId, clusterId); try { return { transport: new TunnelConnectionTransport(reconnected, this._logService), @@ -197,58 +293,11 @@ export class WebTunnelAgentHostService extends Disposable implements ITunnelAgen LOG_PREFIX, AgentHostClientConnectionKind.DevTunnel, ); - const protocolClient = this._instantiationService.createInstance( - AgentHostProtocolClient, address, transportFactory, { clientInfo: agentsWindowAgentHostClientInfo }, - ); - - // Keep an incompatible handshake from tearing down the relay: the - // protocol client must remain registered with IRemoteAgentHostService - // so `triggerServerUpgrade` can locate it and send `_vscodeUpgrade` - // over the still-open transport. - let status: RemoteAgentHostConnectionStatus = RemoteAgentHostConnectionStatus.connected; - let connectError: unknown; - try { - await protocolClient.connect(); - this._logService.info(`${LOG_PREFIX} Protocol handshake completed with ${address}`); - } catch (err) { - const incompatible = RemoteAgentHostConnectionStatus.fromConnectError(err, [PROTOCOL_VERSION]); - if (!RemoteAgentHostConnectionStatus.isIncompatible(incompatible)) { - protocolClient.dispose(); - this._logService.error(`${LOG_PREFIX} Connection setup failed`, err); - throw err; - } - this._logService.warn(`${LOG_PREFIX} Incompatible with ${address}: ${incompatible.message}`); - status = incompatible; - connectError = err; - } - - // Cache before announcing the live connection so the contribution's - // `onDidChangeTunnels` handler has created the provider by the time - // `onDidChangeConnections` fires from `addManagedConnection` and - // wires the connection. Also fires `onDidChangeTunnels`. - this.cacheTunnel(tunnel, authProvider); - - try { - await this._remoteAgentHostService.addManagedConnection({ - name: tunnel.name, - connectionToken, - connection: { - type: RemoteAgentHostEntryType.Tunnel, - tunnelId, - clusterId, - label: tunnel.name, - authProvider, - }, - }, protocolClient, undefined, status); - } catch (err) { - protocolClient.dispose(); - this._logService.error(`${LOG_PREFIX} addManagedConnection failed`, err); - throw err; - } - - if (connectError) { - throw connectError; - } + return { + connection: this._instantiationService.createInstance( + AgentHostProtocolClient, address, transportFactory, { clientInfo: agentsWindowAgentHostClientInfo }, + ), + }; } get canDeleteTunnels(): boolean { @@ -266,8 +315,8 @@ export class WebTunnelAgentHostService extends Disposable implements ITunnelAgen } async disconnect(address: string): Promise { + this._connectionFactory.unstageTunnel(address); await this._remoteAgentHostService.removeRemoteAgentHost(address); - this._storage.notifyTunnelsChanged(); } // Auth @@ -293,14 +342,28 @@ export class WebTunnelAgentHostService extends Disposable implements ITunnelAgen tunnelId: tunnel.tunnelId, clusterId: tunnel.clusterId, name: tunnel.name, + protocolVersion: tunnel.protocolVersion, authProvider, }); } removeCachedTunnel(tunnelId: string): void { + this._connectionFactory.unstageTunnel(`${TUNNEL_ADDRESS_PREFIX}${tunnelId}`); this._storage.removeCachedTunnel(tunnelId); } + isTunnelDismissed(tunnelId: string): boolean { + return this._storage.isTunnelDismissed(tunnelId); + } + + dismissTunnel(tunnelId: string): void { + this._storage.dismissTunnel(tunnelId); + } + + clearTunnelDismissal(tunnelId: string): void { + this._storage.clearTunnelDismissal(tunnelId); + } + isAutoConnectSuppressed(tunnelId: string): boolean { return this._storage.isAutoConnectSuppressed(tunnelId); } diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/wslAgentHost.contribution.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/wslAgentHost.contribution.ts index 730fd5ade527cc..e161c7685a3656 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/wslAgentHost.contribution.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/wslAgentHost.contribution.ts @@ -5,9 +5,7 @@ import { IntervalTimer } from '../../../../../base/common/async.js'; import { isCancellationError } from '../../../../../base/common/errors.js'; -import { isWindows } from '../../../../../base/common/platform.js'; -import { IRemoteAgentHostService, RemoteAgentHostAutoConnectSettingId, RemoteAgentHostConnectionStatus, RemoteAgentHostEntryType, RemoteAgentHostsEnabledSettingId, getEntryTypeConfig } from '../../../../../platform/agentHost/common/remoteAgentHostService.js'; -import { computeReconnectDelay } from '../../../../../platform/agentHost/common/reconnectPolicy.js'; +import { type IRemoteAgentHostEntry, IRemoteAgentHostService, RemoteAgentHostConnectionStatus, RemoteAgentHostEntryType, RemoteAgentHostsEnabledSettingId, getEntryAddress, getEntryTypeConfig } from '../../../../../platform/agentHost/common/remoteAgentHostService.js'; import { IWSLRemoteAgentHostService, WSL_ADDRESS_PREFIX } from '../../../../../platform/agentHost/common/wslRemoteAgentHost.js'; import { IConfigurationService } from '../../../../../platform/configuration/common/configuration.js'; import { IInstantiationService } from '../../../../../platform/instantiation/common/instantiation.js'; @@ -15,38 +13,31 @@ import { ILogService } from '../../../../../platform/log/common/log.js'; import { INotificationService } from '../../../../../platform/notification/common/notification.js'; import { IWorkbenchContribution, registerWorkbenchContribution2, WorkbenchPhase } from '../../../../../workbench/common/contributions.js'; import { ISessionsProvidersService } from '../../../../services/sessions/browser/sessionsProvidersService.js'; -import { ManagedReconnectAgentHostContribution, ManagedReconnectState } from './managedReconnectAgentHostContribution.js'; - -/** After this much wall-clock time, a paused auto-reconnect is auto-resumed. */ -const WSL_RECONNECT_PAUSE_AUTO_RESUME_MS = 5 * 60 * 1000; -/** - * Background poll for `wsl --list --running` so a user-initiated WSL boot can - * be detected and a cached distro reconnected without waiting for an unrelated - * event. - */ -const WSL_RUNNING_POLL_MS = 5 * 60 * 1000; +import { ManagedReconnectAgentHostContribution } from './managedReconnectAgentHostContribution.js'; export function shouldPauseWSLReconnectAfterFailure(err: unknown): boolean { return isCancellationError(err); } /** - * Manages sessions providers and auto-reconnect for WSL-backed remote agent - * hosts. Mirrors {@link TunnelAgentHostContribution}: providers are sourced - * from the WSL service's in-memory cache ({@link IWSLRemoteAgentHostService.getCachedDistros}) - * rather than from persisted settings, and live connections are wired back to - * their providers as connection events arrive. + * How often to look for cached distros that have started since the last check. * - * The per-connection agent registration (chat sessions, language models) is - * handled by {@link RemoteAgentHostContribution} reacting to - * `onDidChangeConnections` — exactly as it does for tunnels. + * A stopped distro fails to connect terminally, so no retry stays armed for it. + * WSL raises no event when a distro boots, and the user may well start one + * outside VS Code, so this poll is the only way a cached host recovers without + * a manual action or a reload. + */ +const WSL_RUNNING_POLL_MS = 5 * 60 * 1000; + +/** + * Manages session providers for WSL-backed remote agent hosts. The remote + * agent host service owns automatic dialing and retry of cached distros. */ export class WSLAgentHostContribution extends ManagedReconnectAgentHostContribution implements IWorkbenchContribution { static readonly ID = 'sessions.contrib.wslAgentHostContribution'; - /** Distros that were running at the last poll; used to detect newly-running distros. */ - private _lastKnownRunningDistros = new Set(); + protected readonly _entryType = RemoteAgentHostEntryType.WSL; constructor( @IRemoteAgentHostService remoteAgentHostService: IRemoteAgentHostService, @@ -59,268 +50,76 @@ export class WSLAgentHostContribution extends ManagedReconnectAgentHostContribut ) { super(remoteAgentHostService, configurationService, logService, instantiationService, sessionsProvidersService, notificationService); - // Reconcile providers when connections change (added/removed/reconnected). this._register(this._remoteAgentHostService.onDidChangeConnections(() => { - // New/removed connection — paused auto-reconnect may have been - // caused by a transient outage that's now resolved. this._resumeReconnects('WSL'); this._reconcile(); })); - // Reconcile when enablement / auto-connect config changes. this._register(this._configurationService.onDidChangeConfiguration(e => { - if (e.affectsConfiguration(RemoteAgentHostsEnabledSettingId) || e.affectsConfiguration(RemoteAgentHostAutoConnectSettingId)) { + if (e.affectsConfiguration(RemoteAgentHostsEnabledSettingId)) { this._resumeReconnects('WSL'); this._reconcile(); } })); - // Initial setup for cached distros and connected remotes. - this._reconcile(); - - // Periodic backstop: catches user-initiated WSL boots even when no - // other event fires. Cheap (`wsl --list --running --quiet`) so the - // 5-minute cadence has no measurable cost. - this._register(new IntervalTimer()).cancelAndSet( - () => void this._reconnectWSLEntriesIfRunning(), - WSL_RUNNING_POLL_MS, - ); - } - - private _reconcile(): void { - this._reconcileProviders(); - this._wireConnections(); - this._updateConnectionStatuses(); - void this._reconnectWSLEntriesIfRunning(); - } - - // -- Provider management -- - - private _reconcileProviders(): void { - const entries = this._enabled ? this._getCachedWSLEntries() : []; - const desiredAddresses = new Set(entries.map(e => e.address)); - - // Remove providers whose distro is no longer cached. - for (const [address] of this._providerStores) { - if (!desiredAddresses.has(address)) { - this._providerStores.deleteAndDispose(address); - } - } - - // Add or recreate providers for cached distros. - for (const entry of entries) { - const existing = this._providerInstances.get(entry.address); - if (existing && existing.label !== (entry.name || entry.address)) { - // Name changed — recreate since ISessionsProvider.label is readonly. - this._providerStores.deleteAndDispose(entry.address); - } - if (!this._providerStores.has(entry.address)) { - this._createProvider(entry.address, entry.name, { - // WSL: an explicit user click should boot a stopped distro - // (`wsl.exe -d ` boots it). The "never auto-boot" - // rule only applies to the periodic auto-reconnect path. - connectOnDemand: () => this._connectWSLOnDemand(entry.distro, entry.name, entry.address), - disconnectOnDemand: () => this._disconnectWSLOnDemand(entry.distro, entry.address), - onDidReportConnectProgress: this._wslService.onDidReportConnectProgress, - }); - } - } - } - - /** Wire live connections to their providers so session operations work. */ - private _wireConnections(): void { - for (const [address, provider] of this._providerInstances) { - const connectionInfo = this._remoteAgentHostService.connections.find( - c => c.address === address && RemoteAgentHostConnectionStatus.isConnected(c.status) - ); - if (connectionInfo) { - const connection = this._remoteAgentHostService.getConnection(address); - if (connection) { - provider.setConnection(connection, connectionInfo.defaultDirectory); - } - } - } - } - - private _updateConnectionStatuses(): void { - for (const [address, provider] of this._providerInstances) { - const connectionInfo = this._remoteAgentHostService.connections.find(c => c.address === address); - if (connectionInfo) { - // Service has an entry for this address — its status is - // authoritative (including `incompatible` from the WebSocket - // connect failure path and `connecting` or `reconnecting`). - provider.setConnectionStatus(connectionInfo.status); - } else if (this._pendingReconnects.has(this._distroForAddress(address))) { - provider.setConnectionStatus(RemoteAgentHostConnectionStatus.connecting); - } else if (!RemoteAgentHostConnectionStatus.isIncompatible(provider.connectionStatus.get())) { - // No service entry. Preserve incompatible state set by the - // reconnect catch; otherwise fall back to disconnected. - provider.setConnectionStatus(RemoteAgentHostConnectionStatus.disconnected); - } - } - } + this._register(new IntervalTimer()).cancelAndSet(() => this._reconnectNewlyRunningDistros(), WSL_RUNNING_POLL_MS); - private _distroForAddress(address: string): string { - return address.startsWith(WSL_ADDRESS_PREFIX) ? address.slice(WSL_ADDRESS_PREFIX.length) : address; - } - - private _getCachedWSLEntries(): readonly { distro: string; name: string; address: string }[] { - return this._wslService.getCachedDistros().map(({ distro, name }) => ({ - distro, - name, - address: `${WSL_ADDRESS_PREFIX}${distro}`, - })); + this._reconcile(); } - // -- Auto-reconnect -- - /** - * Re-establish WSL connections for cached distros that are already - * running. Never auto-boots a distro; only acts on user-initiated boots - * observed via {@link IWSLRemoteAgentHostService.listRunningDistros}. + * Ask the service to redial cached distros that are running but not + * connected. Discovery is this contribution's job; the dial itself stays + * with the service, which owns every connection's lifecycle. */ - private async _reconnectWSLEntriesIfRunning(): Promise { - if (!isWindows) { - return; - } + private async _reconnectNewlyRunningDistros(): Promise { if (!this._enabled) { - this._reconnectStates.clearAndDisposeAll(); return; } - - const running = new Set(await this._wslService.listRunningDistros().catch(() => [])); - const newlyRunning: string[] = []; - for (const distro of running) { - if (!this._lastKnownRunningDistros.has(distro)) { - newlyRunning.push(distro); - } - } - this._lastKnownRunningDistros = running; - if (newlyRunning.length > 0) { - this._logService.info(`[WSLAgentHost] Newly running WSL distro(s): ${newlyRunning.join(', ')}`); + const entries = this._getProviderEntries(); + if (entries.length === 0) { + return; } - - const autoConnect = this._configurationService.getValue(RemoteAgentHostAutoConnectSettingId); - const entries = this._getCachedWSLEntries(); - const stillCached = new Set(); + const running = new Set(await this._wslService.listRunningDistros().catch(() => [])); for (const entry of entries) { - stillCached.add(entry.distro); - if (!running.has(entry.distro)) { - continue; - } - const connection = this._remoteAgentHostService.connections.find(c => c.address === entry.address); - if (connection && RemoteAgentHostConnectionStatus.isConnected(connection.status)) { - this._reconnectStates.deleteAndDispose(entry.distro); - continue; - } - if (connection && RemoteAgentHostConnectionStatus.isConnecting(connection.status)) { - continue; - } - if (connection && RemoteAgentHostConnectionStatus.isReconnecting(connection.status)) { - // The protocol client is preserving its state while it reconnects; don't replace it. - this._reconnectStates.get(entry.distro)?.cancelTimer(); - continue; - } - if (this._pendingReconnects.has(entry.distro)) { - this._logService.trace(`[WSLAgentHost] WSL reconnect for ${entry.distro}: reconnect already in progress, skipping`); + if (entry.connection.type !== RemoteAgentHostEntryType.WSL || !running.has(entry.connection.distro)) { continue; } - const state = this._reconnectStates.get(entry.distro); - if (state?.hasPendingTimer) { - this._logService.trace(`[WSLAgentHost] WSL reconnect for ${entry.distro}: retry timer already scheduled, skipping`); + const address = getEntryAddress(entry); + if (this._remoteAgentHostService.connections.some(connection => connection.address === address)) { continue; } - if (state?.paused) { - const pausedMs = Date.now() - state.pausedAt; - if (pausedMs < WSL_RECONNECT_PAUSE_AUTO_RESUME_MS) { - this._logService.trace(`[WSLAgentHost] WSL reconnect for ${entry.distro}: paused (${Math.round(pausedMs / 1000)}s ago), skipping`); - continue; - } - this._logService.info(`[WSLAgentHost] WSL reconnect for ${entry.distro}: auto-resuming after ${Math.round(pausedMs / 1000)}s pause`); - state.resetForResume(); - } - if (!autoConnect) { - this._logService.trace(`[WSLAgentHost] WSL reconnect for ${entry.distro}: auto-connect disabled, skipping`); - continue; - } - void this._attemptWSLReconnect(entry.distro, entry.name, entry.address); - } - - // Drop retry state for distros that are no longer cached. - for (const distro of [...this._reconnectStates.keys()]) { - if (!stillCached.has(distro)) { - this._reconnectStates.deleteAndDispose(distro); - } + this._logService.info(`[RemoteAgentHost] WSL distro '${entry.connection.distro}' is running again; reconnecting`); + this._remoteAgentHostService.reconnect(address, false); } } - private async _attemptWSLReconnect(distro: string, name: string, address: string, options: { userInitiated?: boolean } = {}): Promise { - await this._attemptManagedReconnect({ - kind: 'WSL', - key: distro, - address, - userInitiated: !!options.userInitiated, - reconnectPolicy: getEntryTypeConfig(RemoteAgentHostEntryType.WSL).reconnect, - shouldPause: shouldPauseWSLReconnectAfterFailure, - // WSL-specific gate: never auto-boot a stopped distro. The gate is - // skipped on user-initiated attempts (the user explicitly clicked - // Reconnect — `wsl.exe -d ` will boot it). When the gate - // triggers we return WITHOUT incrementing `attempts` so a long stop - // doesn't burn the retry budget. - preCheck: async userInitiated => { - if (userInitiated) { - return undefined; - } - const stillCached = this._wslService.getCachedDistros().some(d => d.distro === distro); - if (!stillCached) { - this._reconnectStates.deleteAndDispose(distro); - return { skip: true }; - } - const running = new Set(await this._wslService.listRunningDistros().catch(() => [])); - this._lastKnownRunningDistros = running; - if (!running.has(distro)) { - return { skip: true, reason: `distro ${distro} not running` }; - } - return undefined; + protected override _getProviderEntries(): readonly IRemoteAgentHostEntry[] { + if (!this._enabled) { + return []; + } + return this._wslService.getCachedDistros().map(({ distro, name }) => ({ + name, + connection: { + type: RemoteAgentHostEntryType.WSL, + address: `${WSL_ADDRESS_PREFIX}${distro}`, + distro, }, - doConnect: () => this._wslService.reconnect(distro, name).then(() => undefined), - schedule: state => this._scheduleWSLReconnect(distro, name, address, state), - }); + })); } - private _scheduleWSLReconnect(distro: string, name: string, address: string, state: ManagedReconnectState): void { - const reconnectPolicy = getEntryTypeConfig(RemoteAgentHostEntryType.WSL).reconnect; - const delay = computeReconnectDelay(reconnectPolicy, state.attempts); - this._logService.info(`[WSLAgentHost] Scheduling WSL reconnect for ${distro} in ${delay}ms (attempt ${state.attempts + 1}/${reconnectPolicy.maxAttempts})`); - state.scheduleRetry(delay, () => { - if (!this._enabled) { - this._reconnectStates.deleteAndDispose(distro); - return; - } - if (!this._configurationService.getValue(RemoteAgentHostAutoConnectSettingId)) { - return; - } - const live = this._remoteAgentHostService.connections.find(c => c.address === address); - if (live && RemoteAgentHostConnectionStatus.isConnected(live.status)) { - this._reconnectStates.deleteAndDispose(distro); - return; - } - if (live && RemoteAgentHostConnectionStatus.isConnecting(live.status)) { - return; - } - if (live && RemoteAgentHostConnectionStatus.isReconnecting(live.status)) { - // The protocol client is preserving its state while it reconnects; don't replace it. - return; - } - if (this._pendingReconnects.has(distro)) { - return; - } - void this._attemptWSLReconnect(distro, name, address); - }); + protected override _getProviderOptions(entry: IRemoteAgentHostEntry) { + if (entry.connection.type !== RemoteAgentHostEntryType.WSL) { + return {}; + } + const { distro, address } = entry.connection; + return { + connectOnDemand: () => this._connectWSLOnDemand(distro, entry.name, address), + disconnectOnDemand: () => this._disconnectWSLOnDemand(distro, address), + onDidReportConnectProgress: this._wslService.onDidReportConnectProgress, + }; } - // -- On-demand connection -- - private async _connectWSLOnDemand(distro: string, name: string, address: string): Promise { while (true) { const inFlight = this._pendingReconnects.get(distro); @@ -328,26 +127,35 @@ export class WSLAgentHostContribution extends ManagedReconnectAgentHostContribut break; } await inFlight.catch(() => undefined); - const live = this._remoteAgentHostService.connections.find(c => c.address === address); + const live = this._remoteAgentHostService.connections.find(connection => connection.address === address); if (live && RemoteAgentHostConnectionStatus.isConnected(live.status)) { return; } } this._reconnectStates.get(distro)?.resetForResume(); - await this._attemptWSLReconnect(distro, name, address, { userInitiated: true }); + await this._attemptWSLReconnect(distro, name, address, true); + } + + private async _attemptWSLReconnect(distro: string, name: string, address: string, userInitiated: boolean): Promise { + await this._attemptManagedReconnect({ + kind: 'WSL', + key: distro, + address, + userInitiated, + reconnectPolicy: getEntryTypeConfig(RemoteAgentHostEntryType.WSL).reconnect, + shouldPause: shouldPauseWSLReconnectAfterFailure, + doConnect: () => this._wslService.reconnect(distro, name, userInitiated).then(() => undefined), + }); } - /** - * Tear down the active WSL connection for {@link distro} and cancel any - * pending auto-reconnect. Removes the cached distro so it won't auto-reconnect. - * - * Order matters: `removeRemoteAgentHost` MUST run before the WSL service - * teardown so the subsequent close event can't trip auto-reconnect. - */ private async _disconnectWSLOnDemand(distro: string, address: string): Promise { this._reconnectStates.deleteAndDispose(distro); - await this._remoteAgentHostService.removeRemoteAgentHost(address); + // Drop the cached distro before tearing the connection down: the cached + // entry is what makes this address desired, so removing the connection + // first would let reconciliation re-dial it right back. await this._wslService.disconnect(distro); + await this._remoteAgentHostService.removeRemoteAgentHost(address); + this._reconcile(); } } diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/electron-browser/devContainerAgentHostConnector.contribution.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/electron-browser/devContainerAgentHostConnector.contribution.ts index 5bbde6d625fa81..696ae0db248173 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/electron-browser/devContainerAgentHostConnector.contribution.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/electron-browser/devContainerAgentHostConnector.contribution.ts @@ -14,15 +14,13 @@ import { Schemas } from '../../../../../base/common/network.js'; import { ProxyChannel } from '../../../../../base/parts/ipc/common/ipc.js'; import { localize } from '../../../../../nls.js'; import { AGENT_HOST_SCHEME, agentHostAuthority } from '../../../../../platform/agentHost/common/agentHostUri.js'; -import { agentsWindowAgentHostClientInfo } from '../../../../../platform/agentHost/common/agentHostClientInfo.js'; import { AgentHostClientConnectionKind } from '../../../../../platform/agentHost/common/agentHostTelemetry.js'; import { AgentHostAhpJsonlLoggingSettingId } from '../../../../../platform/agentHost/common/agentService.js'; import { AhpJsonlLogger } from '../../../../../platform/agentHost/common/ahpJsonlLogger.js'; import { DEV_CONTAINER_AGENT_HOST_CHANNEL, IDevContainerAgentHostMainService } from '../../../../../platform/agentHost/common/devContainerAgentHost.js'; import { ReconnectingRelayTransport, type IRelayConnectionHandle } from '../../../../../platform/agentHost/common/relayTransport.js'; -import { getEntryTypeConfig, RemoteAgentHostEntryType, RemoteAgentHostsEnabledSettingId } from '../../../../../platform/agentHost/common/remoteAgentHostService.js'; +import { RemoteAgentHostsEnabledSettingId } from '../../../../../platform/agentHost/common/remoteAgentHostService.js'; import { NonReconnectableTransportError } from '../../../../../platform/agentHost/common/state/sessionTransport.js'; -import { AgentHostProtocolClient } from '../../../../../platform/agentHost/browser/agentHostProtocolClient.js'; import { IInstantiationService } from '../../../../../platform/instantiation/common/instantiation.js'; import { ISharedProcessService } from '../../../../../platform/ipc/electron-browser/services.js'; import { ILogService } from '../../../../../platform/log/common/log.js'; @@ -134,7 +132,7 @@ class DevContainerAgentHostConnector implements IDevContainerAgentHostConnector return isDevContainerWorkspaceAvailable(workspaceUri, this._fileService, this._mainService, this._configurationService); } - async connect(workspaceUri: URI, token: CancellationToken): Promise { + async createConnection(workspaceUri: URI, address: string, token: CancellationToken): Promise { ensureDevContainerAgentHostsEnabled(this._configurationService); if (workspaceUri.scheme !== Schemas.file) { throw new Error(localize('devContainerAgentHost.localWorkspaceRequired', "Dev Container Agent Hosts require a local file workspace.")); @@ -148,7 +146,6 @@ class DevContainerAgentHostConnector implements IDevContainerAgentHostConnector this._logService.warn('[DevContainerAgentHostConnector] Failed to cancel connection', error); }); }); - let protocolClient: AgentHostProtocolClient | undefined; try { const result = await this._mainService.connect({ connectionId, @@ -217,24 +214,10 @@ class DevContainerAgentHostConnector implements IDevContainerAgentHostConnector AgentHostClientConnectionKind.DevContainer, ); }; - protocolClient = this._instantiationService.createInstance( - AgentHostProtocolClient, - result.address, - transportFactory, - { - clientInfo: agentsWindowAgentHostClientInfo, - reconnectPolicy: getEntryTypeConfig(RemoteAgentHostEntryType.DevContainer).reconnect, - }, - ); - await protocolClient.connect(); - if (token.isCancellationRequested) { - throw new CancellationError(); - } - return { - address: result.address, + address, name: result.name, - connection: protocolClient, + transportFactory, transportDisposable: combinedDisposable( outputWriter, toDisposable(() => { @@ -245,14 +228,13 @@ class DevContainerAgentHostConnector implements IDevContainerAgentHostConnector ), workspaceUri: workspaceUri.with({ scheme: AGENT_HOST_SCHEME, - authority: agentHostAuthority(result.address), + authority: agentHostAuthority(address), path: result.remoteWorkspaceFolder, }), defaultDirectory: result.remoteWorkspaceFolder, }; } catch (error) { outputWriter.dispose(); - protocolClient?.dispose(); await this._mainService.disconnect(connectionId); throw error; } finally { diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/electron-browser/tunnelAgentHostServiceImpl.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/electron-browser/tunnelAgentHostServiceImpl.ts index cc4092cf54fc73..251ee5945e950c 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/electron-browser/tunnelAgentHostServiceImpl.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/electron-browser/tunnelAgentHostServiceImpl.ts @@ -4,7 +4,8 @@ *--------------------------------------------------------------------------------------------*/ import { Emitter, Event } from '../../../../../base/common/event.js'; -import { Disposable, toDisposable } from '../../../../../base/common/lifecycle.js'; +import { Disposable, IDisposable, toDisposable } from '../../../../../base/common/lifecycle.js'; +import { derived, IObservable, observableSignalFromEvent } from '../../../../../base/common/observable.js'; import { hasKey } from '../../../../../base/common/types.js'; import { ProxyChannel } from '../../../../../base/parts/ipc/common/ipc.js'; import { localize } from '../../../../../nls.js'; @@ -17,10 +18,9 @@ import { ISharedProcessService } from '../../../../../platform/ipc/electron-brow import { ILogService } from '../../../../../platform/log/common/log.js'; import { INotificationService, Severity } from '../../../../../platform/notification/common/notification.js'; import { IProductService } from '../../../../../platform/product/common/productService.js'; -import { IStorageService, StorageScope, StorageTarget } from '../../../../../platform/storage/common/storage.js'; -import { IRemoteAgentHostService, RemoteAgentHostConnectionStatus, RemoteAgentHostEntryType, RemoteAgentHostsEnabledSettingId } from '../../../../../platform/agentHost/common/remoteAgentHostService.js'; +import { IStorageService } from '../../../../../platform/storage/common/storage.js'; +import { IRemoteAgentHostService, RemoteAgentHostConnectionStatus, RemoteAgentHostEntryType, RemoteAgentHostsEnabledSettingId, getEntryAddress, type IRemoteAgentHostConnectOptions, type IRemoteAgentHostConnectionFactory, type IRemoteAgentHostCreatedConnection, type IRemoteAgentHostEntry } from '../../../../../platform/agentHost/common/remoteAgentHostService.js'; import { IRemoteAgentHostLocationPreferenceService } from '../../../../../platform/agentHost/common/remoteAgentHostLocationPreference.js'; -import { PROTOCOL_VERSION } from '../../../../../platform/agentHost/common/state/protocol/version/registry.js'; import { isTunnelGatewaySelectionRejectedError, isTunnelNotFoundError, @@ -28,6 +28,7 @@ import { TUNNEL_ADDRESS_PREFIX, TUNNEL_AGENT_HOST_CHANNEL, TUNNEL_GATEWAY_MIN_PROTOCOL_VERSION, + TUNNEL_MIN_PROTOCOL_VERSION, TunnelAgentHostsSettingId, type ICachedTunnel, type ITunnelAgentHostMainService, @@ -50,6 +51,7 @@ import { AgentHostProtocolClient } from '../../../../../platform/agentHost/brows import { agentsWindowAgentHostClientInfo } from '../../../../../platform/agentHost/common/agentHostClientInfo.js'; import { ReconnectingRelayTransport, type IRelayConnectionHandle } from '../../../../../platform/agentHost/common/relayTransport.js'; import { NonReconnectableTransportError } from '../../../../../platform/agentHost/common/state/sessionTransport.js'; +import { TunnelAgentHostStorage } from '../browser/tunnelAgentHostStorage.js'; export { type IGatewaySelectionRequest, @@ -63,32 +65,86 @@ export { const LOG_PREFIX = '[TunnelAgentHost]'; -/** Storage key for recently used tunnel cache. */ -const CACHED_TUNNELS_KEY = 'tunnelAgentHost.recentTunnels'; -/** Storage key for tunnels the user explicitly disconnected. */ -const AUTO_CONNECT_SUPPRESSED_TUNNELS_KEY = 'tunnelAgentHost.autoConnectSuppressedTunnels'; - /** Whether `selection` picked a live `editor` endpoint out of `inventory`. */ function isEditorGatewaySelection(selection: ITunnelGatewaySelection, inventory: ITunnelGatewayInventory): boolean { return hasKey(selection, { instanceId: true }) && inventory.endpoints.some(endpoint => endpoint.instanceId === selection.instanceId && endpoint.type === 'editor'); } -/** - * Whether the tunnel-failover tracker/notification step should run at all - * for a completed `connect()` attempt. Must be `false` whenever the - * attempt is ultimately a failure — including a registered-for-upgrade - * incompatible handshake (`connectError` set) — even though - * `addManagedConnection` already succeeded and the endpoint is registered. - * A failed reconnect must never update {@link TunnelFailoverTracker} or - * notify: the tracker would otherwise record an endpoint the caller never - * actually got a working connection to, and a subsequent real reconnect - * could then silently skip a notification it should have shown (or vice - * versa). Exported so this ordering guard can be unit tested without - * constructing the full service. - */ -export function shouldTrackTunnelConnection(connectError: unknown): boolean { - return !connectError; +class TunnelConnectionFactory extends Disposable implements IRemoteAgentHostConnectionFactory { + readonly kind = RemoteAgentHostEntryType.Tunnel; + readonly entries: IObservable; + + private readonly _onDidStageTunnel = this._register(new Emitter()); + private readonly _stagedAuthProviders = new Map(); + /** + * Initiation mode for a staged tunnel, consumed by the first + * {@link createConnection} for that address. Staging publishes the entry + * synchronously, so the service's reconciliation can begin dialing before + * the caller's explicit `reconnect` runs — and that dial would otherwise be + * treated as background, suppressing interactive auth and gateway + * selection for the user's own first connect. + */ + private readonly _stagedUserInitiated = new Map(); + private readonly _onDidStageTunnelSignal = observableSignalFromEvent(this, this._onDidStageTunnel.event); + + constructor( + private readonly _storage: TunnelAgentHostStorage, + private readonly _createConnection: (entry: IRemoteAgentHostEntry, authProvider: 'github' | 'microsoft' | undefined, options: IRemoteAgentHostConnectOptions) => Promise, + ) { + super(); + this.entries = derived(this, reader => { + this._onDidStageTunnelSignal.read(reader); + const autoConnectSuppressedTunnels = this._storage.autoConnectSuppressedTunnels.read(reader); + return this._storage.cachedTunnels.read(reader) + .filter(tunnel => !autoConnectSuppressedTunnels.includes(tunnel.tunnelId)) + .map(tunnel => this._entryForTunnel(tunnel, tunnel.authProvider)); + }); + } + + stageTunnel(tunnel: ITunnelInfo, authProvider?: 'github' | 'microsoft', userInitiated = true): IRemoteAgentHostEntry { + const address = `${TUNNEL_ADDRESS_PREFIX}${tunnel.tunnelId}`; + this._stagedAuthProviders.set(address, authProvider); + this._stagedUserInitiated.set(address, userInitiated); + this._storage.cacheTunnel({ tunnelId: tunnel.tunnelId, clusterId: tunnel.clusterId, name: tunnel.name, protocolVersion: tunnel.protocolVersion, authProvider }); + this._onDidStageTunnel.fire(); + return this._entryForTunnel(tunnel, authProvider); + } + + unstageTunnel(address: string): void { + this._stagedUserInitiated.delete(address); + if (this._stagedAuthProviders.delete(address)) { + this._onDidStageTunnel.fire(); + } + } + + createConnection(entry: IRemoteAgentHostEntry, options: IRemoteAgentHostConnectOptions): Promise { + if (entry.connection.type !== RemoteAgentHostEntryType.Tunnel) { + throw new Error(`Tunnel factory cannot create a ${entry.connection.type} connection.`); + } + const address = getEntryAddress(entry); + const stagedUserInitiated = this._stagedUserInitiated.get(address); + // Consume it: only the connect this staging was for is user-initiated, + // and a later automatic reconnect must not prompt. + this._stagedUserInitiated.delete(address); + const connectOptions = stagedUserInitiated === undefined + ? options + : { ...options, userInitiated: stagedUserInitiated }; + return this._createConnection(entry, this._stagedAuthProviders.has(address) ? this._stagedAuthProviders.get(address) : entry.connection.authProvider, connectOptions); + } + + private _entryForTunnel(tunnel: Pick, authProvider?: 'github' | 'microsoft'): IRemoteAgentHostEntry { + return { + name: tunnel.name, + connection: { + type: RemoteAgentHostEntryType.Tunnel, + tunnelId: tunnel.tunnelId, + clusterId: tunnel.clusterId, + label: tunnel.name, + authProvider, + }, + }; + } } /** @@ -100,9 +156,10 @@ export class TunnelAgentHostService extends Disposable implements ITunnelAgentHo declare readonly _serviceBrand: undefined; private readonly _mainService: ITunnelAgentHostMainService; + private readonly _storage: TunnelAgentHostStorage; + private readonly _connectionFactory: TunnelConnectionFactory; - private readonly _onDidChangeTunnels = this._register(new Emitter()); - readonly onDidChangeTunnels: Event = this._onDidChangeTunnels.event; + readonly onDidChangeTunnels: Event; /** Tracks which auth provider was last used successfully. */ private _lastAuthProvider: 'github' | 'microsoft' | undefined; @@ -129,6 +186,13 @@ export class TunnelAgentHostService extends Disposable implements ITunnelAgentHo this._mainService = ProxyChannel.toService( sharedProcessService.getChannel(TUNNEL_AGENT_HOST_CHANNEL), ); + this._storage = this._register(new TunnelAgentHostStorage(this._storageService)); + this.onDidChangeTunnels = this._storage.onDidChangeTunnels; + this._connectionFactory = this._register(new TunnelConnectionFactory( + this._storage, + (entry, authProvider, options) => this._createConnection(entry, authProvider, options), + )); + this._register(this._remoteAgentHostService.registerConnectionFactory(this._connectionFactory)); } async listTunnels(options?: { silent?: boolean }): Promise { @@ -163,132 +227,119 @@ export class TunnelAgentHostService extends Disposable implements ITunnelAgentHo throw new Error('Remote agent host connections are not enabled.'); } + const entry = this._connectionFactory.stageTunnel(tunnel, authProvider, options?.userInitiated ?? true); + const address = getEntryAddress(entry); + this._remoteAgentHostService.reconnect(address, options?.userInitiated ?? true); + await this._remoteAgentHostService.waitForConnection(address); + } + + private async _createConnection(entry: IRemoteAgentHostEntry, authProvider: 'github' | 'microsoft' | undefined, options: IRemoteAgentHostConnectOptions): Promise { + if (entry.connection.type !== RemoteAgentHostEntryType.Tunnel) { + throw new Error(`Tunnel factory cannot create a ${entry.connection.type} connection.`); + } + + // Bind the narrowed connection before the closure: TypeScript does not + // carry the discriminant narrowing into the `find` callback below. + const connection = entry.connection; + const cachedTunnel = this._storage.getCachedTunnels().find(cached => cached.tunnelId === connection.tunnelId); + const tunnel: ITunnelInfo = { + tunnelId: connection.tunnelId, + clusterId: connection.clusterId, + name: connection.label ?? entry.name, + tags: [], + // Legacy cache fallback, not a real capability claim. + protocolVersion: cachedTunnel?.protocolVersion ?? TUNNEL_MIN_PROTOCOL_VERSION, + hostConnectionCount: 0, + }; + const connectOptions = this.getAutoConnectMode(tunnel) === 'prompt' + ? { ...options, userInitiated: true } + : options; const auth = authProvider - ? await this._getTokenForProvider(authProvider, false) - : await this._getToken(false); + ? await this._getTokenForProvider(authProvider, !connectOptions.userInitiated) + : await this._getToken(!connectOptions.userInitiated); if (!auth) { - throw new Error('No authentication available'); + throw new NonReconnectableTransportError('No cached authentication available to connect the tunnel.'); } - this._logService.info(`${LOG_PREFIX} Connecting to tunnel '${tunnel.name}' (${tunnel.tunnelId})`); - - // Protocol-v6 tunnels expose a registry-based endpoint selection - // gateway: prepare it first and resolve a target by the user's saved - // location preference before completing the connection. Protocol-v5 - // tunnels have no gateway — `prepareSelection` returns `undefined` - // and we fall back to the legacy direct-connect path with no prompt. - const session = await this._mainService.prepareSelection(auth.token, auth.provider, tunnel.tunnelId, tunnel.clusterId); let result: ITunnelConnectResult; let editorFallback = false; - if (session) { - const selection = await resolveGatewaySelection(this._locationPreferenceService, this._dialogService, { - hostKey: `${TUNNEL_ADDRESS_PREFIX}${tunnel.tunnelId}`, - hostLabel: tunnel.name, - productName: this._productService.nameShort, - inventory: session.inventory, - userInitiated: options?.userInitiated ?? true, - }); - if (!selection) { - this._logService.info(options?.userInitiated === false - ? `${LOG_PREFIX} Deferring tunnel '${tunnel.name}' until the user chooses an agent host location` - : `${LOG_PREFIX} Agent host selection cancelled for tunnel '${tunnel.name}'`); - await this._mainService.cancelSelection(session.selectionId); - return; + try { + const session = await this._mainService.prepareSelection(auth.token, auth.provider, tunnel.tunnelId, tunnel.clusterId); + if (session) { + const selection = await resolveGatewaySelection(this._locationPreferenceService, this._dialogService, { + hostKey: getEntryAddress(entry), + hostLabel: tunnel.name, + productName: this._productService.nameShort, + inventory: session.inventory, + userInitiated: connectOptions.userInitiated, + }); + if (!selection) { + await this._mainService.cancelSelection(session.selectionId); + throw new NonReconnectableTransportError('Tunnel agent host selection requires user interaction.'); + } + const completed = await this._completeSelectionWithFallback(auth, tunnel, session, selection); + result = completed.result; + editorFallback = completed.editorFallback; + } else { + result = await this._mainService.connect(auth.token, auth.provider, tunnel.tunnelId, tunnel.clusterId); } - const completed = await this._completeSelectionWithFallback(auth, tunnel, session, selection); - result = completed.result; - editorFallback = completed.editorFallback; - } else { - result = await this._mainService.connect(auth.token, auth.provider, tunnel.tunnelId, tunnel.clusterId); + } catch (err) { + if (isTunnelNotFoundError(err)) { + throw new NonReconnectableTransportError(err.message); + } + throw err; } - this._logService.info(`${LOG_PREFIX} Tunnel relay connected, connectionId=${result.connectionId}`); - // Build relay transport + protocol client. If construction itself - // fails (rare — would mean the AHP logger or transport ctor threw) - // tear the just-opened main-side relay down before propagating. - let protocolClient: AgentHostProtocolClient; try { const ahpLoggingEnabled = !!this._configurationService.getValue(AgentHostAhpJsonlLoggingSettingId); let useSeedConnection = true; const establish = async (): Promise => { if (useSeedConnection) { useSeedConnection = false; - // The initial relay belongs to the managed connection's transport disposable. return { connectionId: result.connectionId }; } return this._establishBackgroundRelay(tunnel, auth.provider); }; - const transportFactory = () => new ReconnectingRelayTransport( - establish, - this._mainService, - () => ahpLoggingEnabled ? this._instantiationService.createInstance( - AhpJsonlLogger, - { logsHome: this._environmentService.logsHome, connectionId: result.connectionId, transport: 'tunnel' }, - ) : undefined, - this._logService, - LOG_PREFIX, - AgentHostClientConnectionKind.DevTunnel, - ); - protocolClient = this._instantiationService.createInstance( - AgentHostProtocolClient, result.address, transportFactory, { clientInfo: agentsWindowAgentHostClientInfo }, + const connection = this._instantiationService.createInstance( + AgentHostProtocolClient, + result.address, + () => new ReconnectingRelayTransport( + establish, + this._mainService, + () => ahpLoggingEnabled ? this._instantiationService.createInstance( + AhpJsonlLogger, + { logsHome: this._environmentService.logsHome, connectionId: result.connectionId, transport: 'tunnel' }, + ) : undefined, + this._logService, + LOG_PREFIX, + AgentHostClientConnectionKind.DevTunnel, + ), + { clientInfo: agentsWindowAgentHostClientInfo }, ); + return { + connection, + transportDisposable: this._createTransportDisposable(result, connectOptions.userInitiated, editorFallback), + }; } catch (err) { - this._logService.error(`${LOG_PREFIX} Connection setup failed`, err); this._mainService.disconnect(result.connectionId).catch(() => { /* best effort */ }); throw err; } + } - // Keep an incompatible handshake from tearing down the relay: the - // protocol client must remain registered with IRemoteAgentHostService - // so `triggerServerUpgrade` can locate it and send `_vscodeUpgrade` - // over the still-open transport. - let status: RemoteAgentHostConnectionStatus = RemoteAgentHostConnectionStatus.connected; - let connectError: unknown; - try { - await protocolClient.connect(); - this._logService.info(`${LOG_PREFIX} Protocol handshake completed with ${result.address}`); - } catch (err) { - const incompatible = RemoteAgentHostConnectionStatus.fromConnectError(err, [PROTOCOL_VERSION]); - if (!RemoteAgentHostConnectionStatus.isIncompatible(incompatible)) { - this._logService.error(`${LOG_PREFIX} Connection setup failed`, err); - protocolClient.dispose(); - this._mainService.disconnect(result.connectionId).catch(() => { /* best effort */ }); - throw err; + private _createTransportDisposable(result: ITunnelConnectResult, userInitiated: boolean, editorFallback: boolean): IDisposable { + const listener = this._remoteAgentHostService.onDidChangeConnections(() => { + const status = this._remoteAgentHostService.connections.find(connection => connection.address === result.address)?.status; + if (RemoteAgentHostConnectionStatus.isConnected(status)) { + listener.dispose(); + this._notifyIfTunnelFailover(result, { userInitiated }, editorFallback); + } else if (!status || RemoteAgentHostConnectionStatus.isIncompatible(status)) { + listener.dispose(); } - this._logService.warn(`${LOG_PREFIX} Incompatible with ${result.address}: ${incompatible.message}`); - status = incompatible; - connectError = err; - } - - this.cacheTunnel(tunnel, auth.provider); - - const transportDisposable = toDisposable(() => { + }); + return toDisposable(() => { + listener.dispose(); this._mainService.disconnect(result.connectionId).catch(() => { /* best effort */ }); }); - try { - await this._remoteAgentHostService.addManagedConnection({ - name: result.name, - connectionToken: result.connectionToken, - connection: { - type: RemoteAgentHostEntryType.Tunnel, - tunnelId: tunnel.tunnelId, - clusterId: tunnel.clusterId, - label: tunnel.name, - authProvider: auth.provider, - }, - }, protocolClient, transportDisposable, status); - } catch (err) { - this._logService.error(`${LOG_PREFIX} addManagedConnection failed`, err); - protocolClient.dispose(); - transportDisposable.dispose(); - throw err; - } - - if (!shouldTrackTunnelConnection(connectError)) { - throw connectError; - } - - this._notifyIfTunnelFailover(result, options, editorFallback); } private async _establishBackgroundRelay(tunnel: ITunnelInfo, authProvider: 'github' | 'microsoft'): Promise { @@ -385,7 +436,7 @@ export class TunnelAgentHostService extends Disposable implements ITunnelAgentHo } /** - * After a successful {@link addManagedConnection} registration, compare + * After the service reports a successful connection, compare * the newly selected endpoint's server type against the last one * successfully registered for this tunnel's stable address and, if this * was a silent editor → standalone failover, show a single informational @@ -432,8 +483,8 @@ export class TunnelAgentHostService extends Disposable implements ITunnelAgentHo } async disconnect(address: string): Promise { + this._connectionFactory.unstageTunnel(address); await this._remoteAgentHostService.removeRemoteAgentHost(address); - this._onDidChangeTunnels.fire(); } /** @@ -542,85 +593,39 @@ export class TunnelAgentHostService extends Disposable implements ITunnelAgentHo } getCachedTunnels(): ICachedTunnel[] { - const raw = this._storageService.get(CACHED_TUNNELS_KEY, StorageScope.APPLICATION); - if (!raw) { - return []; - } - try { - return JSON.parse(raw); - } catch { - return []; - } + return this._storage.getCachedTunnels(); } cacheTunnel(tunnel: ITunnelInfo, authProvider?: 'github' | 'microsoft'): void { - const cached = this.getCachedTunnels(); - const filtered = cached.filter(t => t.tunnelId !== tunnel.tunnelId); - filtered.unshift({ - tunnelId: tunnel.tunnelId, - clusterId: tunnel.clusterId, - name: tunnel.name, - authProvider, - }); - this.clearAutoConnectSuppression(tunnel.tunnelId); - this._storeCachedTunnels(filtered); - this._onDidChangeTunnels.fire(); + this._storage.cacheTunnel({ tunnelId: tunnel.tunnelId, clusterId: tunnel.clusterId, name: tunnel.name, protocolVersion: tunnel.protocolVersion, authProvider }); } removeCachedTunnel(tunnelId: string): void { - const cached = this.getCachedTunnels(); - this._storeCachedTunnels(cached.filter(t => t.tunnelId !== tunnelId)); - this.clearAutoConnectSuppression(tunnelId); - this._onDidChangeTunnels.fire(); + this._connectionFactory.unstageTunnel(`${TUNNEL_ADDRESS_PREFIX}${tunnelId}`); + this._storage.removeCachedTunnel(tunnelId); } - isAutoConnectSuppressed(tunnelId: string): boolean { - return this._getAutoConnectSuppressedTunnels().has(tunnelId); + isTunnelDismissed(tunnelId: string): boolean { + return this._storage.isTunnelDismissed(tunnelId); } - suppressAutoConnect(tunnelId: string): void { - const suppressed = this._getAutoConnectSuppressedTunnels(); - suppressed.add(tunnelId); - this._storeAutoConnectSuppressedTunnels(suppressed); + dismissTunnel(tunnelId: string): void { + this._storage.dismissTunnel(tunnelId); } - clearAutoConnectSuppression(tunnelId: string): void { - const suppressed = this._getAutoConnectSuppressedTunnels(); - if (!suppressed.delete(tunnelId)) { - return; - } - this._storeAutoConnectSuppressedTunnels(suppressed); + clearTunnelDismissal(tunnelId: string): void { + this._storage.clearTunnelDismissal(tunnelId); } - private _storeCachedTunnels(tunnels: ICachedTunnel[]): void { - if (tunnels.length === 0) { - this._storageService.remove(CACHED_TUNNELS_KEY, StorageScope.APPLICATION); - } else { - this._storageService.store(CACHED_TUNNELS_KEY, JSON.stringify(tunnels), StorageScope.APPLICATION, StorageTarget.USER); - } + isAutoConnectSuppressed(tunnelId: string): boolean { + return this._storage.isAutoConnectSuppressed(tunnelId); } - private _getAutoConnectSuppressedTunnels(): Set { - const raw = this._storageService.get(AUTO_CONNECT_SUPPRESSED_TUNNELS_KEY, StorageScope.APPLICATION); - if (!raw) { - return new Set(); - } - try { - const parsed: unknown = JSON.parse(raw); - if (!Array.isArray(parsed)) { - return new Set(); - } - return new Set(parsed.filter(item => typeof item === 'string')); - } catch { - return new Set(); - } + suppressAutoConnect(tunnelId: string): void { + this._storage.suppressAutoConnect(tunnelId); } - private _storeAutoConnectSuppressedTunnels(tunnelIds: Set): void { - if (tunnelIds.size === 0) { - this._storageService.remove(AUTO_CONNECT_SUPPRESSED_TUNNELS_KEY, StorageScope.APPLICATION); - } else { - this._storageService.store(AUTO_CONNECT_SUPPRESSED_TUNNELS_KEY, JSON.stringify([...tunnelIds]), StorageScope.APPLICATION, StorageTarget.USER); - } + clearAutoConnectSuppression(tunnelId: string): void { + this._storage.clearAutoConnectSuppression(tunnelId); } } diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/cloudSandboxAgentHostService.test.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/cloudSandboxAgentHostService.test.ts index 04d223ea19bbd5..5f421a4609f1ff 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/cloudSandboxAgentHostService.test.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/cloudSandboxAgentHostService.test.ts @@ -16,7 +16,7 @@ import { type CloudSandboxConnectResult, type ICloudSandboxClientToken, } from '../../../../../../platform/agentHost/common/cloudSandboxAgentHost.js'; -import { IRemoteAgentHostService, RemoteAgentHostsEnabledSettingId } from '../../../../../../platform/agentHost/common/remoteAgentHostService.js'; +import { IRemoteAgentHostConnectionFactory, IRemoteAgentHostService, RemoteAgentHostsEnabledSettingId } from '../../../../../../platform/agentHost/common/remoteAgentHostService.js'; import { IConfigurationService } from '../../../../../../platform/configuration/common/configuration.js'; import { TestConfigurationService } from '../../../../../../platform/configuration/test/common/testConfigurationService.js'; import { IEnvironmentService } from '../../../../../../platform/environment/common/environment.js'; @@ -76,6 +76,7 @@ function createService(store: Pick<{ add(t: T): T override readonly onDidChangeConnections = Event.None; override readonly connections = []; override getConnection() { return undefined; } + override registerConnectionFactory(_factory: IRemoteAgentHostConnectionFactory) { return { dispose() { } }; } }()); instantiationService.stub(IEnvironmentService, new class extends mock() { override readonly logsHome = URI.file('/logs'); diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/devContainerAgentHostService.test.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/devContainerAgentHostService.test.ts index eeec6866cf1319..1185b3fbf78a84 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/devContainerAgentHostService.test.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/devContainerAgentHostService.test.ts @@ -8,13 +8,16 @@ import { DeferredPromise } from '../../../../../../base/common/async.js'; import { CancellationToken, CancellationTokenSource } from '../../../../../../base/common/cancellation.js'; import { CancellationError } from '../../../../../../base/common/errors.js'; import { Emitter, Event } from '../../../../../../base/common/event.js'; +import { StringSHA1 } from '../../../../../../base/common/hash.js'; import { Disposable, IDisposable, toDisposable } from '../../../../../../base/common/lifecycle.js'; +import { getComparisonKey } from '../../../../../../base/common/resources.js'; import { URI } from '../../../../../../base/common/uri.js'; import { mock } from '../../../../../../base/test/common/mock.js'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/test/common/utils.js'; import { IAgentConnection } from '../../../../../../platform/agentHost/common/agentService.js'; +import { AgentHostProtocolClient } from '../../../../../../platform/agentHost/browser/agentHostProtocolClient.js'; import { AGENT_HOST_SCHEME, agentHostAuthority } from '../../../../../../platform/agentHost/common/agentHostUri.js'; -import { getEntryAddress, IRemoteAgentHostConnectionInfo, IRemoteAgentHostEntry, IRemoteAgentHostService, RemoteAgentHostConnectionStatus, RemoteAgentHostEntryType } from '../../../../../../platform/agentHost/common/remoteAgentHostService.js'; +import { getEntryAddress, IRemoteAgentHostConnectionFactory, IRemoteAgentHostConnectionInfo, IRemoteAgentHostEntry, IRemoteAgentHostService, RemoteAgentHostConnectionStatus, RemoteAgentHostEntryType } from '../../../../../../platform/agentHost/common/remoteAgentHostService.js'; import { TestInstantiationService } from '../../../../../../platform/instantiation/test/common/instantiationServiceMock.js'; import { ISessionsProvidersService } from '../../../../../services/sessions/browser/sessionsProvidersService.js'; import { ISessionsProvider } from '../../../../../services/sessions/common/sessionsProvider.js'; @@ -22,21 +25,30 @@ import { IDevContainerAgentHostConnector } from '../../../../../common/devContai import { DevContainerAgentHostService } from '../../browser/devContainerAgentHostService.js'; import { IRemoteAgentHostSessionsProviderConfig, RemoteAgentHostSessionsProvider } from '../../browser/remoteAgentHostSessionsProvider.js'; -class TestAgentConnection extends mock() implements IDisposable { - override readonly clientId = 'dev-container-client'; +/** Stands in for the protocol client the factory hands back to the service. */ +class TestAgentConnection extends mock() implements IDisposable { + override get clientId(): string { return 'dev-container-client'; } disposed = false; - dispose(): void { + override dispose(): void { this.disposed = true; } } +function devContainerAddress(workspaceUri: URI): string { + const sha = new StringSHA1(); + sha.update(getComparisonKey(workspaceUri)); + return `devcontainer:${sha.digest()}`; +} + class TestRemoteAgentHostService extends mock() implements IDisposable { private readonly _onDidChangeConnections = new Emitter(); override readonly onDidChangeConnections = this._onDidChangeConnections.event; private _connections: IRemoteAgentHostConnectionInfo[] = []; + private _factory: IRemoteAgentHostConnectionFactory | undefined; + private _pendingConnect: Promise | undefined; - addedEntry: IRemoteAgentHostEntry | undefined; + stagedEntry: IRemoteAgentHostEntry | undefined; removedAddress: string | undefined; connection: (IAgentConnection & IDisposable) | undefined; transportDisposable: IDisposable | undefined; @@ -51,20 +63,42 @@ class TestRemoteAgentHostService extends mock() impleme : undefined; } - override async addManagedConnection(entry: IRemoteAgentHostEntry, connection: IAgentConnection, transportDisposable?: IDisposable): Promise { - this.addedEntry = entry; - this.connection = connection as IAgentConnection & IDisposable; - this.transportDisposable = transportDisposable; - const connectionInfo = { - address: getEntryAddress(entry), - name: entry.name, - clientId: connection.clientId, - defaultDirectory: '/workspace', - status: RemoteAgentHostConnectionStatus.connected, - }; - this._connections = [connectionInfo]; - this._onDidChangeConnections.fire(); - return connectionInfo; + override registerConnectionFactory(factory: IRemoteAgentHostConnectionFactory): IDisposable { + this._factory = factory; + return toDisposable(() => { + if (this._factory === factory) { + this._factory = undefined; + } + }); + } + + override reconnect(address: string, userInitiated = true): void { + const entry = this._factory?.entries.get().find(entry => getEntryAddress(entry) === address); + if (!entry || !this._factory) { + return; + } + this.stagedEntry = entry; + this._pendingConnect = this._factory.createConnection(entry, { userInitiated }).then(createdConnection => { + this.connection = createdConnection.connection; + this.transportDisposable = createdConnection.transportDisposable; + this._connections = [{ + address, + name: entry.name, + clientId: createdConnection.connection.clientId, + defaultDirectory: '/workspace', + status: RemoteAgentHostConnectionStatus.connected, + }]; + this._onDidChangeConnections.fire(); + }); + } + + override async waitForConnection(address: string): Promise { + await this._pendingConnect; + const connection = this._connections.find(candidate => candidate.address === address); + if (!connection) { + throw new Error(`No connection for ${address}`); + } + return connection; } override async removeRemoteAgentHost(address: string): Promise { @@ -143,7 +177,7 @@ class TestDevContainerAgentHostService extends DevContainerAgentHostService { suite('Dev Container Agent Host Service', () => { const store = ensureNoDisposablesAreLeakedInTestSuite(); - test('registers a runtime provider around a connector-owned Agent Host connection', async () => { + test('registers a runtime provider around a factory-owned Agent Host connection', async () => { const instantiationService = store.add(new TestInstantiationService()); const remoteAgentHostService = store.add(new TestRemoteAgentHostService()); const sessionsProvidersService = store.add(new TestSessionsProvidersService()); @@ -154,7 +188,7 @@ suite('Dev Container Agent Host Service', () => { )); const sourceWorkspace = URI.file('/source'); - const address = 'devcontainer:source'; + const address = devContainerAddress(sourceWorkspace); const remoteWorkspace = URI.from({ scheme: AGENT_HOST_SCHEME, authority: agentHostAuthority(address), @@ -165,18 +199,19 @@ suite('Dev Container Agent Host Service', () => { let connectorCalls = 0; const connector: IDevContainerAgentHostConnector = { isAvailable: async () => true, - connect: async () => { + createConnection: async (_workspaceUri, address) => { connectorCalls++; return { address, name: 'Source Dev Container', - connection, + transportFactory: () => undefined as never, transportDisposable: toDisposable(() => transportDisposed = true), workspaceUri: remoteWorkspace, }; }, }; store.add(service.registerConnector(connector)); + instantiationService.stubInstance(AgentHostProtocolClient, connection); const first = await service.connect(sourceWorkspace, CancellationToken.None); const second = await service.connect(sourceWorkspace, CancellationToken.None); @@ -194,7 +229,7 @@ suite('Dev Container Agent Host Service', () => { reusedConnection: second.providerId === first.providerId && second.workspaceUri.toString() === first.workspaceUri.toString(), afterFirstRelease, connectorCalls, - entry: remoteAgentHostService.addedEntry, + entry: remoteAgentHostService.stagedEntry, provider: service.provider && { config: service.provider.config, connected: service.provider.wiredConnection === connection, @@ -255,15 +290,15 @@ suite('Dev Container Agent Host Service', () => { )); const sourceWorkspace = URI.file('/source'); - const address = 'devcontainer:source'; + const address = devContainerAddress(sourceWorkspace); const connection = new TestAgentConnection(); let transportDisposed = false; store.add(service.registerConnector({ isAvailable: async () => true, - connect: async () => ({ - address, + createConnection: async (_workspaceUri, stagedAddress) => ({ + address: stagedAddress, name: 'Source Dev Container', - connection, + transportFactory: () => undefined as never, transportDisposable: toDisposable(() => transportDisposed = true), workspaceUri: URI.from({ scheme: AGENT_HOST_SCHEME, @@ -272,6 +307,7 @@ suite('Dev Container Agent Host Service', () => { }), }), })); + instantiationService.stubInstance(AgentHostProtocolClient, connection); const target = await service.connect(sourceWorkspace, CancellationToken.None); await service.disconnect(sourceWorkspace); @@ -303,19 +339,19 @@ suite('Dev Container Agent Host Service', () => { )); const sourceWorkspace = URI.file('/source'); - const address = 'devcontainer:source'; + const address = devContainerAddress(sourceWorkspace); const connection = new TestAgentConnection(); let connectorCalls = 0; let connectorToken = CancellationToken.None; const result = new DeferredPromise<{ address: string; name: string; - connection: TestAgentConnection; + transportFactory: () => never; workspaceUri: URI; }>(); store.add(service.registerConnector({ isAvailable: async () => true, - connect: async (_workspaceUri, token) => { + createConnection: async (_workspaceUri, _address, token) => { connectorCalls++; connectorToken = token; return result.p; @@ -330,13 +366,14 @@ suite('Dev Container Agent Host Service', () => { result.complete({ address, name: 'Source Dev Container', - connection, + transportFactory: () => undefined as never, workspaceUri: URI.from({ scheme: AGENT_HOST_SCHEME, authority: agentHostAuthority(address), path: '/workspaces/source', }), }); + instantiationService.stubInstance(AgentHostProtocolClient, connection); const target = await first; await target.release(); @@ -364,20 +401,19 @@ suite('Dev Container Agent Host Service', () => { )); const sourceWorkspace = URI.file('/source'); - const address = 'devcontainer:source'; - const connection = new TestAgentConnection(); + const address = devContainerAddress(sourceWorkspace); let transportDisposed = false; let connectorToken = CancellationToken.None; const result = new DeferredPromise<{ address: string; name: string; - connection: TestAgentConnection; + transportFactory: () => never; transportDisposable: IDisposable; workspaceUri: URI; }>(); store.add(service.registerConnector({ isAvailable: async () => true, - connect: async (_workspaceUri, token) => { + createConnection: async (_workspaceUri, _address, token) => { connectorToken = token; return result.p; }, @@ -389,7 +425,7 @@ suite('Dev Container Agent Host Service', () => { result.complete({ address, name: 'Source Dev Container', - connection, + transportFactory: () => undefined as never, transportDisposable: toDisposable(() => transportDisposed = true), workspaceUri: URI.from({ scheme: AGENT_HOST_SCHEME, @@ -397,20 +433,17 @@ suite('Dev Container Agent Host Service', () => { path: '/workspaces/source', }), }); - await assert.rejects(connect); await disconnect; assert.deepStrictEqual({ - addedEntry: remoteAgentHostService.addedEntry, + stagedEntry: remoteAgentHostService.stagedEntry, provider: service.provider, registeredProviders: sessionsProvidersService.getProviders(), - connectionDisposed: connection.disposed, transportDisposed, }, { - addedEntry: undefined, + stagedEntry: undefined, provider: undefined, registeredProviders: [], - connectionDisposed: true, transportDisposed: true, }); }); @@ -426,15 +459,15 @@ suite('Dev Container Agent Host Service', () => { )); const sourceWorkspace = URI.file('/source'); - const address = 'devcontainer:source'; + const address = devContainerAddress(sourceWorkspace); const connection = new TestAgentConnection(); let transportDisposed = false; store.add(service.registerConnector({ isAvailable: async () => true, - connect: async () => ({ - address, + createConnection: async (_workspaceUri, stagedAddress) => ({ + address: stagedAddress, name: 'Source Dev Container', - connection, + transportFactory: () => undefined as never, transportDisposable: toDisposable(() => transportDisposed = true), workspaceUri: URI.from({ scheme: AGENT_HOST_SCHEME, @@ -443,6 +476,7 @@ suite('Dev Container Agent Host Service', () => { }), }), })); + instantiationService.stubInstance(AgentHostProtocolClient, connection); await service.connect(sourceWorkspace, CancellationToken.None); const provider = service.provider; diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/managedReconnectAgentHostContribution.test.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/managedReconnectAgentHostContribution.test.ts index 88d492c66bb031..7e89e5457df30c 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/managedReconnectAgentHostContribution.test.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/managedReconnectAgentHostContribution.test.ts @@ -99,4 +99,20 @@ suite('ManagedReconnectState', () => { assert.strictEqual(fired, 0, 'pending retry must be cancelled by resetForResume'); }); }); + + test('automatically resumes states that do not require a user action', () => { + const state = store.add(new ManagedReconnectState()); + state.attempts = 1; + state.paused = true; + + assert.deepStrictEqual({ + resumed: state.resumeAutomatically(), + attempts: state.attempts, + paused: state.paused, + }, { + resumed: true, + attempts: 0, + paused: false, + }); + }); }); diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/remoteAgentHost.contribution.test.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/remoteAgentHost.contribution.test.ts index 3d41b7dafb3a50..7125b48dadfd2a 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/remoteAgentHost.contribution.test.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/remoteAgentHost.contribution.test.ts @@ -4,22 +4,20 @@ *--------------------------------------------------------------------------------------------*/ import assert from 'assert'; -import { DeferredPromise, timeout } from '../../../../../../base/common/async.js'; -import { CancellationError } from '../../../../../../base/common/errors.js'; +import { timeout } from '../../../../../../base/common/async.js'; import { AgentHostAuthenticationRecovery, AgentHostAuthTokenCache } from '../../../../../../workbench/contrib/chat/browser/agentSessions/agentHost/agentHostAuth.js'; -import { runWithFakedTimers } from '../../../../../../base/test/common/timeTravelScheduler.js'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/test/common/utils.js'; import { type IAgentConnection } from '../../../../../../platform/agentHost/common/agentService.js'; import { ICommandService } from '../../../../../../platform/commands/common/commands.js'; -import { IRemoteAgentHostSSHConnection, RemoteAgentHostEntryType } from '../../../../../../platform/agentHost/common/remoteAgentHostService.js'; -import { SSHHostKeyDeniedError } from '../../../../../../platform/agentHost/common/sshRemoteAgentHost.js'; +import { type IRemoteAgentHostEntry, getEntryAddress, RemoteAgentHostEntryType } from '../../../../../../platform/agentHost/common/remoteAgentHostService.js'; import { AuthRequiredReason, NotificationType, type INotification } from '../../../../../../platform/agentHost/common/state/sessionActions.js'; import { type ProtectedResourceMetadata } from '../../../../../../platform/agentHost/common/state/protocol/state.js'; import { TestInstantiationService } from '../../../../../../platform/instantiation/test/common/instantiationServiceMock.js'; import { ILogService, NullLogService } from '../../../../../../platform/log/common/log.js'; import { IAuthenticationService } from '../../../../../../workbench/services/authentication/common/authentication.js'; -import { categorizeSSHConnectError } from '../../../../../common/sessionsTelemetry.js'; -import { disconnectSSHEntry, RemoteAgentHostContribution, shouldPauseSSHReconnectAfterFailure, sshConnectionKey, SSHReconnectState } from '../../browser/remoteAgentHost.contribution.js'; +import { RemoteAgentHostContribution } from '../../browser/remoteAgentHost.contribution.js'; +import { SSHAgentHostContribution } from '../../browser/sshAgentHost.contribution.js'; +import { WebSocketAgentHostContribution } from '../../browser/webSocketAgentHost.contribution.js'; interface IRemoteAuthNotificationHarness { _connections: Map; @@ -160,284 +158,61 @@ suite('RemoteAgentHost auth notifications', () => { }); }); -suite('SSHReconnectState', () => { - const store = ensureNoDisposablesAreLeakedInTestSuite(); - - test('scheduleRetry fires the handler after the requested delay', async () => { - return runWithFakedTimers({ useFakeTimers: true, maxTaskCount: 10_000 }, async () => { - const state = store.add(new SSHReconnectState()); - let fired = 0; - state.scheduleRetry(1000, () => fired++); - - assert.strictEqual(state.hasPendingTimer, true); - await timeout(500); - assert.strictEqual(fired, 0); - await timeout(600); - assert.strictEqual(fired, 1); - }); - }); - - test('hasPendingTimer becomes false once the handler has run', async () => { - // Regression guard for the PR-feedback fix: the timer disposable must - // be cleared inside scheduleRetry's tick so that observers that check - // hasPendingTimer after the handler runs see the right value. - return runWithFakedTimers({ useFakeTimers: true, maxTaskCount: 10_000 }, async () => { - const state = store.add(new SSHReconnectState()); - state.scheduleRetry(1000, () => { /* no follow-up */ }); - await timeout(1100); - assert.strictEqual(state.hasPendingTimer, false, 'timer should be cleared after firing'); - }); - }); - - test('cancelTimer prevents the handler from firing', async () => { - return runWithFakedTimers({ useFakeTimers: true, maxTaskCount: 10_000 }, async () => { - const state = store.add(new SSHReconnectState()); - let fired = 0; - state.scheduleRetry(1000, () => fired++); - state.cancelTimer(); - assert.strictEqual(state.hasPendingTimer, false); - await timeout(2000); - assert.strictEqual(fired, 0); - }); - }); - - test('scheduling a second retry replaces the first', async () => { - // MutableDisposable contract: assigning a new value disposes the old. - // If two retries were scheduled simultaneously the contribution would - // double-fire reconnect attempts and inflate the attempt counter. - return runWithFakedTimers({ useFakeTimers: true, maxTaskCount: 10_000 }, async () => { - const state = store.add(new SSHReconnectState()); - let firstFired = 0; - let secondFired = 0; - state.scheduleRetry(5000, () => firstFired++); - state.scheduleRetry(1000, () => secondFired++); - await timeout(6000); - assert.strictEqual(firstFired, 0, 'replaced timer must not fire'); - assert.strictEqual(secondFired, 1); - }); - }); - - test('disposing the state cancels a pending retry timer', async () => { - // This is the safety net for the DisposableMap that owns these states: - // when the contribution is disposed (or a host is removed) the entry's - // pending timer must be cancelled so we don't fire reconnect attempts - // against torn-down services. - return runWithFakedTimers({ useFakeTimers: true, maxTaskCount: 10_000 }, async () => { - const state = new SSHReconnectState(); - let fired = 0; - state.scheduleRetry(1000, () => fired++); - state.dispose(); - await timeout(2000); - assert.strictEqual(fired, 0); - }); - }); - - test('resetForResume clears the timer and zeros attempts/paused state', async () => { - return runWithFakedTimers({ useFakeTimers: true, maxTaskCount: 10_000 }, async () => { - const state = store.add(new SSHReconnectState()); - let fired = 0; - state.attempts = 7; - state.paused = true; - state.requiresUserInitiatedResume = true; - state.scheduleRetry(1000, () => fired++); - - state.resetForResume(); - assert.deepStrictEqual({ - attempts: state.attempts, - paused: state.paused, - requiresUserInitiatedResume: state.requiresUserInitiatedResume, - hasPendingTimer: state.hasPendingTimer, - }, { - attempts: 0, - paused: false, - requiresUserInitiatedResume: false, - hasPendingTimer: false, - }); - - await timeout(2000); - assert.strictEqual(fired, 0, 'pending retry must be cancelled by resetForResume'); - }); - }); - - test('host key denial requires an explicit resume', () => { - const state = store.add(new SSHReconnectState()); - state.attempts = 1; - state.paused = true; - state.requiresUserInitiatedResume = true; - - const automaticResume = state.resumeAutomatically(); - const afterAutomaticResume = { - attempts: state.attempts, - paused: state.paused, - requiresUserInitiatedResume: state.requiresUserInitiatedResume, - }; - state.resetForResume(); - - assert.deepStrictEqual({ - automaticResume, - afterAutomaticResume, - afterExplicitResume: { - attempts: state.attempts, - paused: state.paused, - requiresUserInitiatedResume: state.requiresUserInitiatedResume, - }, - }, { - automaticResume: false, - afterAutomaticResume: { - attempts: 1, - paused: true, - requiresUserInitiatedResume: true, - }, - afterExplicitResume: { - attempts: 0, - paused: false, - requiresUserInitiatedResume: false, - }, - }); - }); -}); - -suite('shouldPauseSSHReconnectAfterFailure', () => { - ensureNoDisposablesAreLeakedInTestSuite(); - - test('pauses reconnect after cancellation or host key denial but not after regular failures', () => { - assert.deepStrictEqual({ - cancellation: shouldPauseSSHReconnectAfterFailure(new CancellationError()), - hostKeyDenial: shouldPauseSSHReconnectAfterFailure(new SSHHostKeyDeniedError('test-host')), - regularError: shouldPauseSSHReconnectAfterFailure(new Error('boom')), - }, { - cancellation: true, - hostKeyDenial: true, - regularError: false, - }); - }); -}); - -suite('categorizeSSHConnectError', () => { - ensureNoDisposablesAreLeakedInTestSuite(); - - test('returns bounded categories without logging error messages', () => { - assert.deepStrictEqual({ - cancellation: categorizeSSHConnectError(new CancellationError()), - hostKeyDenial: categorizeSSHConnectError(new SSHHostKeyDeniedError('test-host')), - authentication: categorizeSSHConnectError(new Error('All configured authentication methods failed')), - network: categorizeSSHConnectError(new Error('connect ETIMEDOUT')), - other: categorizeSSHConnectError(new Error('remote setup failed')), - }, { - cancellation: 'cancelled', - hostKeyDenial: 'hostKeyDenied', - authentication: 'authentication', - network: 'network', - other: 'other', - }); - }); -}); +interface IProviderOwnerHarness { + _configurationService: { getValue(key: string): boolean }; + _remoteAgentHostService: { readonly configuredEntries: readonly IRemoteAgentHostEntry[] }; + _entryType: RemoteAgentHostEntryType; + _providerStores: Map & { deleteAndDispose(address: string): void }; + _providerInstances: Map; + _createProvider(address: string): void; + _getProviderOptions(entry: IRemoteAgentHostEntry): object; + _reconcileProviders(): void; +} -suite('disconnectSSHEntry', () => { +suite('Remote agent host provider ownership', () => { ensureNoDisposablesAreLeakedInTestSuite(); - function makeSSHConfigConnection(overrides: Partial = {}): IRemoteAgentHostSSHConnection { - return { - type: RemoteAgentHostEntryType.SSH, - address: 'localhost:4321', - sshConfigHost: 'myserver', - hostName: 'myserver.example.com', - ...overrides, - }; - } - - test('removes the entry from configured storage BEFORE tearing down the SSH tunnel', async () => { - // Regression guard for the X-button picker fix. `_sshService.disconnect` - // fires `onDidChangeConnections` synchronously, which the contribution - // translates into `_reconcile` → `_reconnectSSHEntries`. If the entry - // is still in configured storage at that point, the auto-reconnect - // path immediately reconnects the host we just told it to disconnect - // (and on the next window reload, the persisted entry reconnects too). - const calls: string[] = []; - const connection = makeSSHConfigConnection(); - - // Block removeRemoteAgentHost so we can prove disconnect waits for it. - const removed = new DeferredPromise(); - - const remoteAgentHostService = { - removeRemoteAgentHost: async (address: string) => { - calls.push(`remove:${address}`); - await removed.p; - }, - }; - const sshService = { - disconnect: async (key: string) => { - calls.push(`ssh:${key}`); - }, + test('gives WebSocket and SSH entries distinct owners while the shared contribution registers none', () => { + const entries: IRemoteAgentHostEntry[] = [ + { name: 'Tunnel', connection: { type: RemoteAgentHostEntryType.Tunnel, tunnelId: 'my-tunnel', clusterId: 'usw2' } }, + { name: 'WSL', connection: { type: RemoteAgentHostEntryType.WSL, address: 'wsl:Ubuntu-24.04', distro: 'Ubuntu-24.04' } }, + { name: 'Sandbox', connection: { type: RemoteAgentHostEntryType.CloudSandbox, address: 'cloudsandbox:abc', environmentId: 'abc' } }, + { name: 'Dev Container', connection: { type: RemoteAgentHostEntryType.DevContainer, address: 'devcontainer:abc', hostPath: '/repo' } }, + { name: 'Socket', connection: { type: RemoteAgentHostEntryType.WebSocket, address: 'ws://host:8080' } }, + { name: 'Remote', connection: { type: RemoteAgentHostEntryType.SSH, address: 'localhost:4321', sshConfigHost: 'myserver', hostName: 'myserver' } }, + ]; + const createHarness = (prototype: object, entryType: RemoteAgentHostEntryType): IProviderOwnerHarness => { + const contribution = Object.create(prototype) as IProviderOwnerHarness; + contribution._configurationService = { getValue: () => true }; + contribution._remoteAgentHostService = { configuredEntries: entries }; + contribution._entryType = entryType; + const providerStores = new Map(); + contribution._providerStores = Object.assign(providerStores, { + deleteAndDispose: (address: string) => { providerStores.delete(address); }, + }); + contribution._providerInstances = new Map(); + return contribution; }; + const sshCreated: string[] = []; + const sshContribution = createHarness(SSHAgentHostContribution.prototype, RemoteAgentHostEntryType.SSH); + sshContribution._getProviderOptions = entry => { sshCreated.push(getEntryAddress(entry)); return {}; }; + sshContribution._createProvider = () => { }; + sshContribution._reconcileProviders(); + const webSocketCreated: string[] = []; + const webSocketContribution = createHarness(WebSocketAgentHostContribution.prototype, RemoteAgentHostEntryType.WebSocket); + webSocketContribution._getProviderOptions = entry => { webSocketCreated.push(getEntryAddress(entry)); return {}; }; + webSocketContribution._createProvider = () => { }; + webSocketContribution._reconcileProviders(); - const pending = disconnectSSHEntry(connection, remoteAgentHostService, sshService); - - // Give microtasks a chance to drain. ssh disconnect must NOT have run yet - // because removeRemoteAgentHost is still pending. - await timeout(0); - assert.deepStrictEqual(calls, ['remove:localhost:4321']); - - removed.complete(); - await pending; - - assert.deepStrictEqual(calls, ['remove:localhost:4321', 'ssh:ssh:myserver']); - }); - - test('uses sshConfigHost-based key when sshConfigHost is set', async () => { - const calls: string[] = []; - await disconnectSSHEntry( - makeSSHConfigConnection({ sshConfigHost: 'myserver' }), - { removeRemoteAgentHost: async () => { /* noop */ } }, - { disconnect: async (key: string) => { calls.push(key); } }, - ); - assert.deepStrictEqual(calls, ['ssh:myserver']); - }); - - test('uses user@host:port key when sshConfigHost is not set', async () => { - const calls: string[] = []; - await disconnectSSHEntry( - { - type: RemoteAgentHostEntryType.SSH, - address: 'localhost:4321', - hostName: 'myserver.example.com', - user: 'me', - port: 2222, - }, - { removeRemoteAgentHost: async () => { /* noop */ } }, - { disconnect: async (key: string) => { calls.push(key); } }, - ); - assert.deepStrictEqual(calls, ['me@myserver.example.com:2222']); - }); -}); - -suite('sshConnectionKey', () => { - ensureNoDisposablesAreLeakedInTestSuite(); - - test('matches the keys the SSH service stores connections under', () => { assert.deepStrictEqual({ - configHost: sshConnectionKey({ - type: RemoteAgentHostEntryType.SSH, - address: 'localhost:4321', - sshConfigHost: 'myserver', - hostName: 'ignored', - }), - userHostPort: sshConnectionKey({ - type: RemoteAgentHostEntryType.SSH, - address: 'localhost:4321', - hostName: 'myserver.example.com', - user: 'me', - port: 2222, - }), - hostOnly: sshConnectionKey({ - type: RemoteAgentHostEntryType.SSH, - address: 'localhost:4321', - hostName: 'myserver.example.com', - }), + sharedProviderMethods: Object.getOwnPropertyNames(RemoteAgentHostContribution.prototype) + .filter(member => member === '_createProvider' || member === '_reconcileProviders'), + sshCreated, + webSocketCreated, }, { - configHost: 'ssh:myserver', - userHostPort: 'me@myserver.example.com:2222', - hostOnly: 'myserver.example.com@myserver.example.com:22', + sharedProviderMethods: [], + sshCreated: ['localhost:4321'], + webSocketCreated: ['ws://host:8080'], }); }); }); diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/sshAgentHost.contribution.test.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/sshAgentHost.contribution.test.ts new file mode 100644 index 00000000000000..330608c557a761 --- /dev/null +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/sshAgentHost.contribution.test.ts @@ -0,0 +1,231 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import assert from 'assert'; +import { DeferredPromise, timeout } from '../../../../../../base/common/async.js'; +import { CancellationError } from '../../../../../../base/common/errors.js'; +import { runWithFakedTimers } from '../../../../../../base/test/common/timeTravelScheduler.js'; +import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/test/common/utils.js'; +import { IRemoteAgentHostSSHConnection, RemoteAgentHostEntryType } from '../../../../../../platform/agentHost/common/remoteAgentHostService.js'; +import { SSHHostKeyDeniedError } from '../../../../../../platform/agentHost/common/sshRemoteAgentHost.js'; +import { categorizeSSHConnectError } from '../../../../../common/sessionsTelemetry.js'; +import { ManagedReconnectState } from '../../browser/managedReconnectAgentHostContribution.js'; +import { disconnectSSHEntry, shouldPauseSSHReconnectAfterFailure, sshConnectionKey } from '../../browser/sshAgentHost.contribution.js'; + +suite('SSH reconnect state', () => { + const store = ensureNoDisposablesAreLeakedInTestSuite(); + + test('manages retry timers and resets state', async () => { + return runWithFakedTimers({ useFakeTimers: true, maxTaskCount: 10_000 }, async () => { + const state = store.add(new ManagedReconnectState()); + let firstFired = 0; + let secondFired = 0; + state.attempts = 7; + state.paused = true; + state.requiresUserInitiatedResume = true; + state.scheduleRetry(5000, () => firstFired++); + state.scheduleRetry(1000, () => secondFired++); + + state.resetForResume(); + assert.deepStrictEqual({ + attempts: state.attempts, + paused: state.paused, + requiresUserInitiatedResume: state.requiresUserInitiatedResume, + hasPendingTimer: state.hasPendingTimer, + }, { + attempts: 0, + paused: false, + requiresUserInitiatedResume: false, + hasPendingTimer: false, + }); + + await timeout(6000); + assert.deepStrictEqual({ firstFired, secondFired }, { firstFired: 0, secondFired: 0 }); + }); + }); + + test('clears a timer once it fires and on disposal', async () => { + return runWithFakedTimers({ useFakeTimers: true, maxTaskCount: 10_000 }, async () => { + const state = store.add(new ManagedReconnectState()); + let fired = 0; + state.scheduleRetry(1000, () => fired++); + await timeout(1100); + + assert.deepStrictEqual({ fired, hasPendingTimer: state.hasPendingTimer }, { fired: 1, hasPendingTimer: false }); + + state.scheduleRetry(1000, () => fired++); + state.dispose(); + await timeout(2000); + assert.strictEqual(fired, 1); + }); + }); + + test('requires explicit resume after host key denial', () => { + const state = store.add(new ManagedReconnectState()); + state.attempts = 1; + state.paused = true; + state.requiresUserInitiatedResume = true; + + const automaticResume = state.resumeAutomatically(); + const afterAutomaticResume = { + attempts: state.attempts, + paused: state.paused, + requiresUserInitiatedResume: state.requiresUserInitiatedResume, + }; + state.resetForResume(); + + assert.deepStrictEqual({ + automaticResume, + afterAutomaticResume, + afterExplicitResume: { + attempts: state.attempts, + paused: state.paused, + requiresUserInitiatedResume: state.requiresUserInitiatedResume, + }, + }, { + automaticResume: false, + afterAutomaticResume: { + attempts: 1, + paused: true, + requiresUserInitiatedResume: true, + }, + afterExplicitResume: { + attempts: 0, + paused: false, + requiresUserInitiatedResume: false, + }, + }); + }); +}); + +suite('shouldPauseSSHReconnectAfterFailure', () => { + ensureNoDisposablesAreLeakedInTestSuite(); + + test('pauses reconnect after cancellation or host key denial but not after regular failures', () => { + assert.deepStrictEqual({ + cancellation: shouldPauseSSHReconnectAfterFailure(new CancellationError()), + hostKeyDenial: shouldPauseSSHReconnectAfterFailure(new SSHHostKeyDeniedError('test-host')), + regularError: shouldPauseSSHReconnectAfterFailure(new Error('boom')), + }, { + cancellation: true, + hostKeyDenial: true, + regularError: false, + }); + }); +}); + +suite('categorizeSSHConnectError', () => { + ensureNoDisposablesAreLeakedInTestSuite(); + + test('returns bounded categories without logging error messages', () => { + assert.deepStrictEqual({ + cancellation: categorizeSSHConnectError(new CancellationError()), + hostKeyDenial: categorizeSSHConnectError(new SSHHostKeyDeniedError('test-host')), + authentication: categorizeSSHConnectError(new Error('All configured authentication methods failed')), + network: categorizeSSHConnectError(new Error('connect ETIMEDOUT')), + other: categorizeSSHConnectError(new Error('remote setup failed')), + }, { + cancellation: 'cancelled', + hostKeyDenial: 'hostKeyDenied', + authentication: 'authentication', + network: 'network', + other: 'other', + }); + }); +}); + +suite('disconnectSSHEntry', () => { + ensureNoDisposablesAreLeakedInTestSuite(); + + function makeSSHConfigConnection(overrides: Partial = {}): IRemoteAgentHostSSHConnection { + return { + type: RemoteAgentHostEntryType.SSH, + address: 'localhost:4321', + sshConfigHost: 'myserver', + hostName: 'myserver.example.com', + ...overrides, + }; + } + + test('drops the persisted entry before tearing down the SSH tunnel', async () => { + const calls: string[] = []; + const connection = makeSSHConfigConnection(); + const disconnected = new DeferredPromise(); + const remoteAgentHostService = { + removeRemoteAgentHost: async (address: string) => { + calls.push(`remove:${address}`); + }, + }; + const sshService = { + disconnect: async (key: string) => { + calls.push(`ssh:${key}`); + await disconnected.p; + }, + }; + + // `sshService.disconnect` is what removes the persisted entry. It has to + // land first, or the teardown's own reconcile still sees the host as + // desired and re-dials it. + const pending = disconnectSSHEntry(connection, remoteAgentHostService, sshService); + await timeout(0); + assert.deepStrictEqual(calls, ['ssh:ssh:myserver']); + + disconnected.complete(); + await pending; + assert.deepStrictEqual(calls, ['ssh:ssh:myserver', 'remove:localhost:4321']); + }); + + test('uses the SSH config host or host connection key on disconnect', async () => { + const calls: string[] = []; + await disconnectSSHEntry( + makeSSHConfigConnection({ sshConfigHost: 'myserver' }), + { removeRemoteAgentHost: async () => { } }, + { disconnect: async key => { calls.push(key); } }, + ); + await disconnectSSHEntry( + { + type: RemoteAgentHostEntryType.SSH, + address: 'localhost:4321', + hostName: 'myserver.example.com', + user: 'me', + port: 2222, + }, + { removeRemoteAgentHost: async () => { } }, + { disconnect: async key => { calls.push(key); } }, + ); + assert.deepStrictEqual(calls, ['ssh:myserver', 'me@myserver.example.com:2222']); + }); +}); + +suite('sshConnectionKey', () => { + ensureNoDisposablesAreLeakedInTestSuite(); + + test('matches the keys the SSH service stores connections under', () => { + assert.deepStrictEqual({ + configHost: sshConnectionKey({ + type: RemoteAgentHostEntryType.SSH, + address: 'localhost:4321', + sshConfigHost: 'myserver', + hostName: 'ignored', + }), + userHostPort: sshConnectionKey({ + type: RemoteAgentHostEntryType.SSH, + address: 'localhost:4321', + hostName: 'myserver.example.com', + user: 'me', + port: 2222, + }), + hostOnly: sshConnectionKey({ + type: RemoteAgentHostEntryType.SSH, + address: 'localhost:4321', + hostName: 'myserver.example.com', + }), + }, { + configHost: 'ssh:myserver', + userHostPort: 'me@myserver.example.com:2222', + hostOnly: 'myserver.example.com@myserver.example.com:22', + }); + }); +}); diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/tunnelAgentHost.contribution.test.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/tunnelAgentHost.contribution.test.ts index f2857d7405a5d2..e84e4eb25e8c7b 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/tunnelAgentHost.contribution.test.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/tunnelAgentHost.contribution.test.ts @@ -13,7 +13,6 @@ import { IAgentConnection } from '../../../../../../platform/agentHost/common/ag import { IRemoteAgentHostConnectionInfo, IRemoteAgentHostService, - RemoteAgentHostAutoConnectSettingId, RemoteAgentHostConnectionStatus, RemoteAgentHostsEnabledSettingId, } from '../../../../../../platform/agentHost/common/remoteAgentHostService.js'; @@ -23,6 +22,7 @@ import { TUNNEL_ADDRESS_PREFIX, type ITunnelHostInfo, type ITunnelInfo, + type TunnelAutoConnectMode, } from '../../../../../../platform/agentHost/common/tunnelAgentHost.js'; import { ConfigurationTarget, IConfigurationService } from '../../../../../../platform/configuration/common/configuration.js'; import { TestConfigurationService } from '../../../../../../platform/configuration/test/common/testConfigurationService.js'; @@ -33,7 +33,6 @@ import { ITelemetryService } from '../../../../../../platform/telemetry/common/t import { IAuthenticationService } from '../../../../../../workbench/services/authentication/common/authentication.js'; import { IHostService } from '../../../../../../workbench/services/host/browser/host.js'; import { ITunnelHostService } from '../../../../../../workbench/contrib/chat/common/tunnelHost.js'; -import type { TunnelConnectFailureReason } from '../../../../../common/sessionsTelemetry.js'; import { ISessionsProvider } from '../../../../../services/sessions/common/sessionsProvider.js'; import { ISessionsProvidersChangeEvent, ISessionsProvidersService } from '../../../../../services/sessions/browser/sessionsProvidersService.js'; import { IAgentHostFilterService } from '../../../../../services/agentHostFilter/common/agentHostFilter.js'; @@ -80,11 +79,13 @@ class StubTunnelService extends Disposable implements ITunnelAgentHostService { private _cached: ICachedTunnel[] = []; private _listed: ITunnelInfo[] | undefined; + private readonly _dismissed = new Set(); private readonly _suppressed = new Set(); - autoConnectMode: 'background' | 'prompt' = 'background'; + autoConnectMode: TunnelAutoConnectMode = 'background'; /** Records every `connect()` call for assertions on the `userInitiated` threading. */ readonly connectCalls: Array<{ tunnel: ITunnelInfo; authProvider: string | undefined; options: { readonly userInitiated?: boolean } | undefined }> = []; + readonly disconnectCalls: string[] = []; setCached(tunnels: ICachedTunnel[]): void { this._cached = tunnels; @@ -93,8 +94,8 @@ class StubTunnelService extends Disposable implements ITunnelAgentHostService { getCachedTunnels(): ICachedTunnel[] { return this._cached; } setListed(tunnels: ITunnelInfo[] | undefined): void { this._listed = tunnels; } - async listTunnels(): Promise { return this._listed ?? []; } - getAutoConnectMode(): 'background' | 'prompt' { return this.autoConnectMode; } + async listTunnels(_options?: { silent?: boolean }): Promise { return this._listed ?? []; } + getAutoConnectMode(_tunnel: ITunnelInfo): TunnelAutoConnectMode { return this.autoConnectMode; } readonly canDeleteTunnels = true; async deleteTunnel(tunnel: ITunnelInfo): Promise { this.removeCachedTunnel(tunnel.tunnelId); } cacheTunnel(tunnel: ITunnelInfo, authProvider?: 'github' | 'microsoft'): void { @@ -105,16 +106,26 @@ class StubTunnelService extends Disposable implements ITunnelAgentHostService { this._cached = this._cached.filter(tunnel => tunnel.tunnelId !== tunnelId); this._onDidChangeTunnels.fire(); } + isTunnelDismissed(id: string): boolean { return this._dismissed.has(id); } + dismissTunnel(id: string): void { + this._dismissed.add(id); + this._onDidChangeTunnels.fire(); + } + clearTunnelDismissal(id: string): void { + if (this._dismissed.delete(id)) { + this._onDidChangeTunnels.fire(); + } + } isAutoConnectSuppressed(id: string): boolean { return this._suppressed.has(id); } suppressAutoConnect(id: string): void { this._suppressed.add(id); } clearAutoConnectSuppression(id: string): void { this._suppressed.delete(id); } - async getAuthProvider(): Promise<'github' | 'microsoft' | undefined> { return undefined; } + async getAuthProvider(_options?: { silent?: boolean }): Promise<'github' | 'microsoft' | undefined> { return undefined; } async connect(tunnel: ITunnelInfo, authProvider?: 'github' | 'microsoft', options?: { readonly userInitiated?: boolean }): Promise { this.connectCalls.push({ tunnel, authProvider, options }); } - async disconnect(_address: string): Promise { /* noop */ } + async disconnect(address: string): Promise { this.disconnectCalls.push(address); } } class StubRemoteAgentHostService extends Disposable { @@ -234,12 +245,10 @@ suite('TunnelAgentHostContribution', () => { const store = ensureNoDisposablesAreLeakedInTestSuite(); test('newly-cached tunnel binds to subsequent live connection', async () => { - // Regression guard for the picker flow: `tunnelService.connect()` is - // contractually obligated to cache the tunnel BEFORE announcing the - // live connection via `addManagedConnection`. That ordering lets the - // `onDidChangeTunnels` handler create the provider first, so the - // `onDidChangeConnections` handler can wire it. Both halves are - // exercised here. + // Tunnel connection staging caches the tunnel before the remote service + // announces its live connection. That ordering lets the cache-change + // handler create the provider first, so the connection-change handler + // can wire it. const tunnelService = store.add(new StubTunnelService()); const remoteService = store.add(new StubRemoteAgentHostService()); const providersService = store.add(new StubSessionsProvidersService()); @@ -292,11 +301,7 @@ suite('TunnelAgentHostContribution', () => { assert.deepStrictEqual(providersService.getProviders(), []); }); - test('background auto-connect threads userInitiated: false through to tunnelService.connect, while explicit connects thread userInitiated: true', async () => { - // Focused regression test for the userInitiated/silent policy: - // background/auto-connect must never be treated as user-initiated - // (so it can reuse, but never prompt for, a saved location), while an - // explicit connect must retain userInitiated: true. + test('on-demand connect threads userInitiated to tunnelService.connect', async () => { const tunnelService = store.add(new StubTunnelService()); const remoteService = store.add(new StubRemoteAgentHostService()); const providersService = store.add(new StubSessionsProvidersService()); @@ -322,67 +327,29 @@ suite('TunnelAgentHostContribution', () => { const address = `${TUNNEL_ADDRESS_PREFIX}${tunnelId}`; tunnelService.setCached([{ tunnelId, clusterId: 'use', name: 'Background Tunnel' }]); - // Access the private connect-orchestration method via a typed seam — - // it's the only place `tunnelService.connect()` is invoked, so this - // exercises the exact threading the fix introduces without needing - // to drive the full `connectOnDemand`/reconnect-timer machinery. + // Access the private on-demand orchestration method via a typed seam. const testable = contribution as unknown as { _connectTunnel(address: string, options: { readonly userInitiated: boolean }): Promise; }; - await testable._connectTunnel(address, { userInitiated: false }); - assert.strictEqual(tunnelService.connectCalls.length, 1); - assert.strictEqual(tunnelService.connectCalls[0].options?.userInitiated, false, 'background connect must pass userInitiated: false'); - + tunnelService.dismissTunnel(tunnelId); await testable._connectTunnel(address, { userInitiated: true }); - assert.strictEqual(tunnelService.connectCalls.length, 2); - assert.strictEqual(tunnelService.connectCalls[1].options?.userInitiated, true, 'explicit/user-initiated connect must pass userInitiated: true'); - }); - - test('auto-connect prompts once for an initial location, then reconnects silently', async () => { - const tunnelService = store.add(new StubTunnelService()); - tunnelService.autoConnectMode = 'prompt'; - const remoteService = store.add(new StubRemoteAgentHostService()); - const providersService = store.add(new StubSessionsProvidersService()); - const configurationService = new TestConfigurationService({ - [RemoteAgentHostsEnabledSettingId]: true, - [RemoteAgentHostAutoConnectSettingId]: true, + assert.deepStrictEqual({ + dismissed: tunnelService.isTunnelDismissed(tunnelId), + connectCalls: tunnelService.connectCalls.map(call => call.options?.userInitiated), + providers: providersService.getProviders().map(provider => provider.id), + }, { + dismissed: false, + connectCalls: [true], + providers: [`agenthost-${address}`], }); - const instantiationService = store.add(new TestInstantiationService()); - instantiationService.stub(ITunnelAgentHostService, tunnelService); - instantiationService.stub(IRemoteAgentHostService, remoteService as unknown as IRemoteAgentHostService); - instantiationService.stub(ISessionsProvidersService, providersService as unknown as ISessionsProvidersService); - instantiationService.stub(IConfigurationService, configurationService); - instantiationService.stub(INotificationService, { notify: () => ({ close() { } }) } as unknown as INotificationService); - instantiationService.stub(ILogService, new NullLogService()); - instantiationService.stub(IAuthenticationService, { onDidChangeSessions: Event.None } as unknown as IAuthenticationService); - instantiationService.stub(ITelemetryService, { publicLog2: () => { } } as unknown as ITelemetryService); - instantiationService.stub(IHostService, new StubHostService()); - instantiationService.stub(ITunnelHostService, store.add(new StubTunnelHostService())); - instantiationService.stub(IAgentHostFilterService, new StubFilterService() as unknown as IAgentHostFilterService); - - const contribution = store.add(instantiationService.createInstance(TestTunnelContribution)); - const tunnel: ITunnelInfo = { tunnelId: 'tunnel-needs-choice', clusterId: 'use', name: 'Needs Choice', tags: ['protocolv6'], protocolVersion: 6, hostConnectionCount: 1 }; - tunnelService.setListed([tunnel]); - const testable = contribution as unknown as { _silentStatusCheck(): Promise }; - - await testable._silentStatusCheck(); - assert.deepStrictEqual(tunnelService.connectCalls.map(call => call.options?.userInitiated), [true]); - - tunnelService.autoConnectMode = 'background'; - await testable._silentStatusCheck(); - assert.deepStrictEqual(tunnelService.connectCalls.map(call => call.options?.userInitiated), [true, false]); }); - test('does not auto-connect the locally hosted tunnel and reconnects it after sharing stops', async () => { + test('suppresses a locally hosted tunnel without removing its provider', () => { const tunnelService = store.add(new StubTunnelService()); const remoteService = store.add(new StubRemoteAgentHostService()); const providersService = store.add(new StubSessionsProvidersService()); - const configurationService = new TestConfigurationService({ - [RemoteAgentHostsEnabledSettingId]: true, - [RemoteAgentHostAutoConnectSettingId]: true, - }); - const hostService = new StubHostService(); + const configurationService = new TestConfigurationService({ [RemoteAgentHostsEnabledSettingId]: true }); const tunnelHostService = store.add(new StubTunnelHostService()); const instantiationService = store.add(new TestInstantiationService()); instantiationService.stub(ITunnelAgentHostService, tunnelService); @@ -393,45 +360,37 @@ suite('TunnelAgentHostContribution', () => { instantiationService.stub(ILogService, new NullLogService()); instantiationService.stub(IAuthenticationService, { onDidChangeSessions: Event.None } as unknown as IAuthenticationService); instantiationService.stub(ITelemetryService, { publicLog2: () => { } } as unknown as ITelemetryService); - instantiationService.stub(IHostService, hostService); + instantiationService.stub(IHostService, new StubHostService()); instantiationService.stub(ITunnelHostService, tunnelHostService); instantiationService.stub(IAgentHostFilterService, new StubFilterService() as unknown as IAgentHostFilterService); - - const locallyHostedTunnel: ITunnelInfo = { tunnelId: 'tunnel-local', clusterId: 'use', name: 'This Machine', tags: [], protocolVersion: 6, hostConnectionCount: 1 }; - const remoteTunnel: ITunnelInfo = { tunnelId: 'tunnel-remote', clusterId: 'use', name: 'Remote Machine', tags: [], protocolVersion: 6, hostConnectionCount: 1 }; - tunnelHostService.setSharingInfo(locallyHostedTunnel.name); - const contribution = store.add(instantiationService.createInstance(TestTunnelContribution)); - tunnelService.setCached([ - { tunnelId: locallyHostedTunnel.tunnelId, clusterId: locallyHostedTunnel.clusterId, name: locallyHostedTunnel.name }, - { tunnelId: remoteTunnel.tunnelId, clusterId: remoteTunnel.clusterId, name: remoteTunnel.name }, - ]); - tunnelService.setListed([locallyHostedTunnel, remoteTunnel]); - const testable = contribution as unknown as { _silentStatusCheck(): Promise }; - await testable._silentStatusCheck(); - const initialConnects = tunnelService.connectCalls.map(call => call.tunnel.tunnelId); + const tunnelId = 'tunnel-hosted'; + const address = `${TUNNEL_ADDRESS_PREFIX}${tunnelId}`; - tunnelHostService.setSharingInfo(undefined); - await Promise.resolve(); - const connectsAfterSharingStopped = tunnelService.connectCalls.map(call => call.tunnel.tunnelId); + tunnelHostService.setSharingInfo('Hosted Tunnel'); + tunnelService.setCached([{ tunnelId, clusterId: 'use', name: 'Hosted Tunnel' }]); + + assert.deepStrictEqual({ + isSuppressed: tunnelService.isAutoConnectSuppressed(tunnelId), + isDismissed: tunnelService.isTunnelDismissed(tunnelId), + hasProvider: contribution.stubProviders.has(address), + }, { + isSuppressed: true, + isDismissed: false, + hasProvider: true, + }); - assert.deepStrictEqual( - { initialConnects, connectsAfterSharingStopped }, - { - initialConnects: [remoteTunnel.tunnelId], - connectsAfterSharingStopped: [remoteTunnel.tunnelId, locallyHostedTunnel.tunnelId, remoteTunnel.tunnelId], - }, - ); + tunnelHostService.setSharingInfo(undefined); + assert.strictEqual(tunnelService.isAutoConnectSuppressed(tunnelId), false); }); - test('recovery signals resume only compatible pause reasons', () => { + test('dismissed tunnel stays removed through discovery until explicitly restored', async () => { const tunnelService = store.add(new StubTunnelService()); const remoteService = store.add(new StubRemoteAgentHostService()); const providersService = store.add(new StubSessionsProvidersService()); const configurationService = new TestConfigurationService({ [RemoteAgentHostsEnabledSettingId]: true }); - const hostService = new StubHostService(); const instantiationService = store.add(new TestInstantiationService()); - instantiationService.stub(ITunnelAgentHostService, tunnelService as unknown as ITunnelAgentHostService); + instantiationService.stub(ITunnelAgentHostService, tunnelService); instantiationService.stub(IRemoteAgentHostService, remoteService as unknown as IRemoteAgentHostService); instantiationService.stub(ISessionsProvidersService, providersService as unknown as ISessionsProvidersService); instantiationService.stub(IConfigurationService, configurationService); @@ -439,125 +398,68 @@ suite('TunnelAgentHostContribution', () => { instantiationService.stub(ILogService, new NullLogService()); instantiationService.stub(IAuthenticationService, { onDidChangeSessions: Event.None } as unknown as IAuthenticationService); instantiationService.stub(ITelemetryService, { publicLog2: () => { } } as unknown as ITelemetryService); - instantiationService.stub(IHostService, hostService); + instantiationService.stub(IHostService, new StubHostService()); instantiationService.stub(ITunnelHostService, store.add(new StubTunnelHostService())); instantiationService.stub(IAgentHostFilterService, new StubFilterService() as unknown as IAgentHostFilterService); const contribution = store.add(instantiationService.createInstance(TestTunnelContribution)); - const maxAttemptsAddress = `${TUNNEL_ADDRESS_PREFIX}tunnel-max-attempts`; - const offlineAddress = `${TUNNEL_ADDRESS_PREFIX}tunnel-offline`; - const authAddress = `${TUNNEL_ADDRESS_PREFIX}tunnel-auth`; - tunnelService.setCached([ - { tunnelId: 'tunnel-max-attempts', clusterId: 'use', name: 'Max Attempts Tunnel' }, - { tunnelId: 'tunnel-offline', clusterId: 'use', name: 'Offline Tunnel' }, - { tunnelId: 'tunnel-auth', clusterId: 'use', name: 'Auth Tunnel' }, - { tunnelId: 'tunnel-idle', clusterId: 'use', name: 'Idle Tunnel' }, - ]); - const testable = contribution as unknown as { - _reconnectPauseReasons: Map; - _reconnectTimeouts: Map>; - _resumeReconnects(trigger: 'sessionAdded'): void; - }; - - testable._reconnectPauseReasons.set(maxAttemptsAddress, 'maxAttemptsReached'); - testable._reconnectPauseReasons.set(offlineAddress, 'hostOffline'); - testable._reconnectPauseReasons.set(authAddress, 'authExpired'); - hostService.fireFocus(true); - const firstResume = { - paused: [...testable._reconnectPauseReasons], - timers: [...testable._reconnectTimeouts.keys()], - }; - - testable._reconnectPauseReasons.set(maxAttemptsAddress, 'maxAttemptsReached'); - hostService.fireFocus(true); - const rateLimitedResume = { - paused: [...testable._reconnectPauseReasons], - timers: [...testable._reconnectTimeouts.keys()], - }; - - testable._resumeReconnects('sessionAdded'); - const sessionResume = { - paused: [...testable._reconnectPauseReasons], - timers: [...testable._reconnectTimeouts.keys()], + const tunnel: ITunnelInfo = { + tunnelId: 'tunnel-dismissed', + clusterId: 'use', + name: 'Dismissed Tunnel', + tags: [], + protocolVersion: 5, + hostConnectionCount: 1, }; - - assert.deepStrictEqual( - { firstResume, rateLimitedResume, sessionResume }, - { - firstResume: { - paused: [[offlineAddress, 'hostOffline'], [authAddress, 'authExpired']], - timers: [maxAttemptsAddress], - }, - rateLimitedResume: { - paused: [[offlineAddress, 'hostOffline'], [authAddress, 'authExpired'], [maxAttemptsAddress, 'maxAttemptsReached']], - timers: [maxAttemptsAddress], - }, - sessionResume: { - paused: [[offlineAddress, 'hostOffline'], [maxAttemptsAddress, 'maxAttemptsReached']], - timers: [maxAttemptsAddress, authAddress], - }, - }, - ); - }); - - test('status checks resume only host-offline pauses and auto-connect preserves other pauses', async () => { - const tunnelService = store.add(new StubTunnelService()); - const remoteService = store.add(new StubRemoteAgentHostService()); - const providersService = store.add(new StubSessionsProvidersService()); - const configurationService = new TestConfigurationService({ - [RemoteAgentHostsEnabledSettingId]: true, - [RemoteAgentHostAutoConnectSettingId]: true, - }); - const hostService = new StubHostService(); - const instantiationService = store.add(new TestInstantiationService()); - instantiationService.stub(ITunnelAgentHostService, tunnelService as unknown as ITunnelAgentHostService); - instantiationService.stub(IRemoteAgentHostService, remoteService as unknown as IRemoteAgentHostService); - instantiationService.stub(ISessionsProvidersService, providersService as unknown as ISessionsProvidersService); - instantiationService.stub(IConfigurationService, configurationService); - instantiationService.stub(INotificationService, { notify: () => ({ close() { } }) } as unknown as INotificationService); - instantiationService.stub(ILogService, new NullLogService()); - instantiationService.stub(IAuthenticationService, { onDidChangeSessions: Event.None } as unknown as IAuthenticationService); - instantiationService.stub(ITelemetryService, { publicLog2: () => { } } as unknown as ITelemetryService); - instantiationService.stub(IHostService, hostService); - instantiationService.stub(ITunnelHostService, store.add(new StubTunnelHostService())); - instantiationService.stub(IAgentHostFilterService, new StubFilterService() as unknown as IAgentHostFilterService); - const contribution = store.add(instantiationService.createInstance(TestTunnelContribution)); - const offlineAddress = `${TUNNEL_ADDRESS_PREFIX}tunnel-offline`; - const authAddress = `${TUNNEL_ADDRESS_PREFIX}tunnel-auth`; - const maxAttemptsAddress = `${TUNNEL_ADDRESS_PREFIX}tunnel-max-attempts`; - tunnelService.setCached([ - { tunnelId: 'tunnel-offline', clusterId: 'use', name: 'Offline Tunnel' }, - { tunnelId: 'tunnel-auth', clusterId: 'use', name: 'Auth Tunnel' }, - { tunnelId: 'tunnel-max-attempts', clusterId: 'use', name: 'Max Attempts Tunnel' }, - ]); - tunnelService.setListed([ - { tunnelId: 'tunnel-offline', clusterId: 'use', name: 'Offline Tunnel', tags: [], protocolVersion: 5, hostConnectionCount: 1 }, - { tunnelId: 'tunnel-auth', clusterId: 'use', name: 'Auth Tunnel', tags: [], protocolVersion: 5, hostConnectionCount: 1 }, - { tunnelId: 'tunnel-max-attempts', clusterId: 'use', name: 'Max Attempts Tunnel', tags: [], protocolVersion: 5, hostConnectionCount: 1 }, - ]); + const address = `${TUNNEL_ADDRESS_PREFIX}${tunnel.tunnelId}`; + tunnelService.setCached([{ tunnelId: tunnel.tunnelId, clusterId: tunnel.clusterId, name: tunnel.name }]); + tunnelService.setListed([tunnel]); const testable = contribution as unknown as { - _reconnectPauseReasons: Map; - _reconnectTimeouts: Map>; + _disconnectTunnel(address: string): Promise; _silentStatusCheck(): Promise; }; - testable._reconnectPauseReasons.set(offlineAddress, 'hostOffline'); - testable._reconnectPauseReasons.set(authAddress, 'authExpired'); - testable._reconnectPauseReasons.set(maxAttemptsAddress, 'maxAttemptsReached'); + await testable._disconnectTunnel(address); + const afterRemove = { + cached: tunnelService.getCachedTunnels().map(cached => cached.tunnelId), + dismissed: tunnelService.isTunnelDismissed(tunnel.tunnelId), + disconnectCalls: tunnelService.disconnectCalls, + providers: providersService.getProviders().map(provider => provider.id), + }; await testable._silentStatusCheck(); - await Promise.resolve(); + const afterDiscovery = { + cached: tunnelService.getCachedTunnels().map(cached => cached.tunnelId), + dismissed: tunnelService.isTunnelDismissed(tunnel.tunnelId), + providers: providersService.getProviders().map(provider => provider.id), + }; - assert.deepStrictEqual( - { - paused: [...testable._reconnectPauseReasons], - connects: tunnelService.connectCalls.map(call => call.tunnel.tunnelId), - timers: [...testable._reconnectTimeouts.keys()], + tunnelService.clearTunnelDismissal(tunnel.tunnelId); + tunnelService.cacheTunnel(tunnel, 'github'); + assert.deepStrictEqual({ + afterRemove, + afterDiscovery, + afterExplicitRestore: { + cached: tunnelService.getCachedTunnels().map(cached => cached.tunnelId), + dismissed: tunnelService.isTunnelDismissed(tunnel.tunnelId), + providers: providersService.getProviders().map(provider => provider.id), }, - { - paused: [[authAddress, 'authExpired'], [maxAttemptsAddress, 'maxAttemptsReached']], - connects: ['tunnel-offline'], - timers: [], + }, { + afterRemove: { + cached: [], + dismissed: true, + disconnectCalls: [address], + providers: [], }, - ); + afterDiscovery: { + cached: [], + dismissed: true, + providers: [], + }, + afterExplicitRestore: { + cached: [tunnel.tunnelId], + dismissed: false, + providers: [`agenthost-${address}`], + }, + }); }); test('clears the provider connection only after a connected transport disconnects', () => { diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/wslAgentHost.contribution.test.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/wslAgentHost.contribution.test.ts index 8f767c3636ecb6..3a07e1d8a6edd0 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/wslAgentHost.contribution.test.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/wslAgentHost.contribution.test.ts @@ -4,9 +4,10 @@ *--------------------------------------------------------------------------------------------*/ import assert from 'assert'; +import { DeferredPromise, timeout } from '../../../../../../base/common/async.js'; import { CancellationError } from '../../../../../../base/common/errors.js'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/test/common/utils.js'; -import { shouldPauseWSLReconnectAfterFailure } from '../../browser/wslAgentHost.contribution.js'; +import { shouldPauseWSLReconnectAfterFailure, WSLAgentHostContribution } from '../../browser/wslAgentHost.contribution.js'; suite('shouldPauseWSLReconnectAfterFailure', () => { ensureNoDisposablesAreLeakedInTestSuite(); @@ -21,3 +22,43 @@ suite('shouldPauseWSLReconnectAfterFailure', () => { }); }); }); + +interface IWSLDisconnectHarness { + _reconnectStates: { deleteAndDispose(key: string): void }; + _wslService: { disconnect(distro: string): Promise }; + _remoteAgentHostService: { removeRemoteAgentHost(address: string): Promise }; + _reconcile(): void; + _disconnectWSLOnDemand(distro: string, address: string): Promise; +} + +suite('WSLAgentHostContribution disconnect', () => { + ensureNoDisposablesAreLeakedInTestSuite(); + + test('drops the cached distro before tearing down the connection', async () => { + const calls: string[] = []; + const disconnected = new DeferredPromise(); + const contribution = Object.create(WSLAgentHostContribution.prototype) as IWSLDisconnectHarness; + contribution._reconnectStates = { deleteAndDispose: key => { calls.push(`state:${key}`); } }; + // `disconnect` is what removes the cached distro. It has to land before + // the connection is torn down, or reconciliation still sees the host as + // desired and re-dials it. + contribution._wslService = { + disconnect: async distro => { + calls.push(`wsl:${distro}`); + await disconnected.p; + }, + }; + contribution._remoteAgentHostService = { + removeRemoteAgentHost: async address => { calls.push(`remove:${address}`); }, + }; + contribution._reconcile = () => { calls.push('reconcile'); }; + + const pending = contribution._disconnectWSLOnDemand('Ubuntu', 'wsl:Ubuntu'); + await timeout(0); + assert.deepStrictEqual(calls, ['state:Ubuntu', 'wsl:Ubuntu']); + + disconnected.complete(); + await pending; + assert.deepStrictEqual(calls, ['state:Ubuntu', 'wsl:Ubuntu', 'remove:wsl:Ubuntu', 'reconcile']); + }); +}); diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/test/electron-browser/tunnelAgentHostServiceImpl.test.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/test/electron-browser/tunnelAgentHostServiceImpl.test.ts index 03e0210d3b839c..51c77ad0221c7d 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/test/electron-browser/tunnelAgentHostServiceImpl.test.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/test/electron-browser/tunnelAgentHostServiceImpl.test.ts @@ -11,7 +11,6 @@ import { selectEditorGatewayEndpoint, selectGatewayFallbackAfterRejection, shouldNotifyTunnelFailover, - shouldTrackTunnelConnection, TunnelFailoverTracker, } from '../../electron-browser/tunnelAgentHostServiceImpl.js'; @@ -194,38 +193,4 @@ suite('tunnelAgentHostServiceImpl - gateway selection', () => { }); }); - suite('shouldTrackTunnelConnection', () => { - test('tracks (and may notify) when the connect attempt has no error', () => { - assert.strictEqual(shouldTrackTunnelConnection(undefined), true); - }); - - test('does not track when the attempt ended in a connectError (e.g. incompatible handshake)', () => { - assert.strictEqual(shouldTrackTunnelConnection(new Error('Unsupported protocol version')), false); - }); - }); - - suite('ordering: connectError must gate the tracker/notification step', () => { - test('an editor -> standalone automatic reconnect that ends in connectError must not update the tracker or notify', () => { - // Models `connect()`'s post-addManagedConnection guard exactly: - // `shouldTrackTunnelConnection(connectError)` must be checked (and - // found false) BEFORE `TunnelFailoverTracker.recordAndShouldNotify` - // is ever called, even though addManagedConnection already - // succeeded and registered the endpoint for a possible upgrade. - const tracker = new TunnelFailoverTracker(); - tracker.recordAndShouldNotify('tunnel:abc', 'editor', true); // initial user-initiated connect - - const connectError: unknown = new Error('Unsupported protocol version'); - let notified: boolean | undefined; - if (shouldTrackTunnelConnection(connectError)) { - notified = tracker.recordAndShouldNotify('tunnel:abc', 'standalone', false); - } - assert.strictEqual(notified, undefined, 'the tracker must never be invoked for a failed (incompatible) reconnect'); - - // A later, fully successful editor -> standalone reconnect must - // still notify: the failed attempt above must not have poisoned - // (or prematurely advanced) the retained state. - assert.strictEqual(shouldTrackTunnelConnection(undefined), true); - assert.strictEqual(tracker.recordAndShouldNotify('tunnel:abc', 'standalone', false), true, 'the retained state must still be "editor" since the failed attempt was never tracked'); - }); - }); }); diff --git a/src/vs/sessions/contrib/sessions/browser/views/sessionsList.ts b/src/vs/sessions/contrib/sessions/browser/views/sessionsList.ts index 249d8182a18e89..a6f225e7f88c2c 100644 --- a/src/vs/sessions/contrib/sessions/browser/views/sessionsList.ts +++ b/src/vs/sessions/contrib/sessions/browser/views/sessionsList.ts @@ -4380,23 +4380,26 @@ export function groupByWorkspace(sessions: ISession[]): ISessionSection[] { /** Maximum number of sessions shown in the "Recent" date section. */ const RECENT_SESSIONS_LIMIT = 10; +const RECENT_SESSIONS_LIMIT_WITH_UPDATES = 15; +const RECENTLY_UPDATED_SESSION_THRESHOLD_MS = 24 * 60 * 60 * 1000; export function groupByDate(sessions: ISession[], sorting: SessionsSorting, getSortKey?: (session: ISession, sorting: SessionsSorting) => number): ISessionSection[] { const key = getSortKey ?? defaultSortKey; const now = new Date(); const startOfToday = new Date(now.getFullYear(), now.getMonth(), now.getDate()).getTime(); const startOfWeek = startOfToday - 7 * 86_400_000; + const recentlyUpdatedThreshold = now.getTime() - RECENTLY_UPDATED_SESSION_THRESHOLD_MS; const recent: ISession[] = []; const older: ISession[] = []; - // `sessions` arrive sorted most-recent-first, so the first sessions within - // the last 7 days (capped at RECENT_SESSIONS_LIMIT) form the "Recent" - // section; everything else falls into "Older". for (const session of sessions) { const time = key(session, sorting); + const wasRecentlyUpdated = sorting === SessionsSorting.Created && session.updatedAt.get().getTime() >= recentlyUpdatedThreshold; + const isWithinRecentLimit = recent.length < RECENT_SESSIONS_LIMIT && time >= startOfWeek; + const isWithinUpdatedRecentLimit = recent.length < RECENT_SESSIONS_LIMIT_WITH_UPDATES && wasRecentlyUpdated; - if (time >= startOfWeek && recent.length < RECENT_SESSIONS_LIMIT) { + if (isWithinRecentLimit || isWithinUpdatedRecentLimit) { recent.push(session); } else { older.push(session); diff --git a/src/vs/sessions/contrib/sessions/test/browser/sessionsList.test.ts b/src/vs/sessions/contrib/sessions/test/browser/sessionsList.test.ts index b2c6f9d1313e04..9cb6c81f6af7e6 100644 --- a/src/vs/sessions/contrib/sessions/test/browser/sessionsList.test.ts +++ b/src/vs/sessions/contrib/sessions/test/browser/sessionsList.test.ts @@ -424,9 +424,27 @@ suite('Sessions - SessionsList', () => { ]); }); - test('"Recent" is capped at 10 sessions; the overflow within 7 days falls into "Older"', () => { - const sessions = Array.from({ length: 13 }, (_, i) => - createSession(`s${i}`, { createdAt: minutesAgo(i + 1) })); + test('sessions updated within the last 24 hours stay in "Recent" when sorting by creation time', () => { + const sessions = [ + createSession('recently-created', { createdAt: daysAgo(3) }), + createSession('recently-updated', { createdAt: daysAgo(10), updatedAt: minutesAgo(30) }), + createSession('old', { createdAt: daysAgo(11), updatedAt: daysAgo(2) }), + ]; + + const sections = groupByDate(sessions, SessionsSorting.Created); + + assert.deepStrictEqual(sections.map(s => ({ id: s.id, sessions: s.sessions.map(session => session.sessionId) })), [ + { id: 'recent', sessions: ['recently-created', 'recently-updated'] }, + { id: 'older', sessions: ['old'] }, + ]); + }); + + test('"Recent" is capped at 10 sessions that were not updated within the last 24 hours', () => { + const twoDaysAgo = daysAgo(2).getTime(); + const sessions = Array.from({ length: 13 }, (_, i) => { + const createdAt = new Date(twoDaysAgo - i * 60_000); + return createSession(`s${i}`, { createdAt }); + }); const sections = groupByDate(sessions, SessionsSorting.Created); @@ -436,6 +454,34 @@ suite('Sessions - SessionsList', () => { ]); }); + test('"Recent" expands from 10 to 15 only for additional recently updated sessions', () => { + const twoDaysAgo = daysAgo(2).getTime(); + const recentlyCreated = Array.from({ length: 10 }, (_, i) => { + const createdAt = new Date(twoDaysAgo - i * 60_000); + return createSession(`created-${i}`, { createdAt }); + }); + const sessions = [ + ...recentlyCreated, + createSession('not-recently-updated', { createdAt: daysAgo(10), updatedAt: daysAgo(2) }), + ...Array.from({ length: 6 }, (_, i) => + createSession(`updated-${i}`, { createdAt: daysAgo(11 + i), updatedAt: minutesAgo(i + 1) })), + ]; + + const sections = groupByDate(sessions, SessionsSorting.Created); + + assert.deepStrictEqual(sections.map(s => ({ id: s.id, sessions: s.sessions.map(session => session.sessionId) })), [ + { + id: 'recent', + sessions: [ + 'created-0', 'created-1', 'created-2', 'created-3', 'created-4', + 'created-5', 'created-6', 'created-7', 'created-8', 'created-9', + 'updated-0', 'updated-1', 'updated-2', 'updated-3', 'updated-4', + ], + }, + { id: 'older', sessions: ['not-recently-updated', 'updated-5'] }, + ]); + }); + test('empty sections are omitted', () => { const sessions = [ createSession('only-old', { createdAt: daysAgo(20) }), diff --git a/src/vs/sessions/sessions.desktop.main.ts b/src/vs/sessions/sessions.desktop.main.ts index 16be5da38d42a4..50d12f94cf4f08 100644 --- a/src/vs/sessions/sessions.desktop.main.ts +++ b/src/vs/sessions/sessions.desktop.main.ts @@ -234,6 +234,8 @@ import './contrib/providers/remoteAgentHost/browser/remoteAgentHost.contribution import './contrib/providers/remoteAgentHost/browser/remoteAgentHostTerminal.contribution.js'; import './contrib/providers/remoteAgentHost/browser/tunnelAgentHost.contribution.js'; import './contrib/providers/remoteAgentHost/browser/wslAgentHost.contribution.js'; +import './contrib/providers/remoteAgentHost/browser/sshAgentHost.contribution.js'; +import './contrib/providers/remoteAgentHost/browser/webSocketAgentHost.contribution.js'; import './contrib/providers/remoteAgentHost/browser/devContainerAgentHostService.js'; import './contrib/providers/remoteAgentHost/electron-browser/devContainerAgentHostConnector.contribution.js'; // Change Preferred Remote Agent Location (Chat: ... command) diff --git a/src/vs/sessions/sessions.web.main.ts b/src/vs/sessions/sessions.web.main.ts index 6789f9dc0ce920..7a757d6ae4f28b 100644 --- a/src/vs/sessions/sessions.web.main.ts +++ b/src/vs/sessions/sessions.web.main.ts @@ -170,6 +170,9 @@ import './contrib/providers/remoteAgentHost/browser/tunnelAgentHost.contribution // WSL agent host — reconciles cached WSL distros into session providers import './contrib/providers/remoteAgentHost/browser/wslAgentHost.contribution.js'; +// WebSocket agent host — reconciles configured WebSocket hosts into session providers +import './contrib/providers/remoteAgentHost/browser/webSocketAgentHost.contribution.js'; + // Remote agent host terminal profiles — registers terminal profiles for connected agent hosts import './contrib/providers/remoteAgentHost/browser/remoteAgentHostTerminal.contribution.js'; diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostResponseFileChanges.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostResponseFileChanges.ts index 36dbb5d2748c26..45f723b5be3fd1 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostResponseFileChanges.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostResponseFileChanges.ts @@ -10,13 +10,14 @@ import { getComparisonKey, isEqual, isEqualOrParent } from '../../../../../../ba import { isDefined } from '../../../../../../base/common/types.js'; import { URI } from '../../../../../../base/common/uri.js'; import { IAgentConnection } from '../../../../../../platform/agentHost/common/agentService.js'; -import { buildTurnChangesetUri, ChangesetKind } from '../../../../../../platform/agentHost/common/changesetUri.js'; +import { buildBranchChangesetUri, buildTurnChangesetUri, ChangesetKind } from '../../../../../../platform/agentHost/common/changesetUri.js'; import { normalizeFileEdit } from '../../../../../../platform/agentHost/common/fileEditDiff.js'; import { toAgentHostContentUri, toAgentHostUri } from '../../../../../../platform/agentHost/common/agentHostUri.js'; import { buildDefaultChatUri, ChangesetStatus, FileEditKind, + readSessionEhcliLastMigratedTurn, ResponsePartKind, StateComponents, ToolCallStatus, @@ -40,7 +41,7 @@ const REQUEST_CACHE_CAPACITY = 1000; * Where a turn's diffs came from, for tracing. `retained` means every source * was momentarily empty and the previous result was kept instead. */ -type TurnDiffSource = 'unsupported' | 'changeset' | 'authoritativeEmpty' | 'response' | 'retained'; +type TurnDiffSource = 'unsupported' | 'changeset' | 'authoritativeEmpty' | 'response' | 'branchFallback' | 'retained'; function uriArrayEquals(a: readonly URI[], b: readonly URI[]): boolean { return a.length === b.length && a.every((uri, index) => isEqual(uri, b[index])); @@ -148,6 +149,13 @@ export class AgentHostResponseFileChangesProvider extends Disposable implements const changesetStateObs = this._subscribe(StateComponents.Changeset, turnChangesetUriObs); const responseFileEditsObs = this._createFileEditDiffsObservable(backendSession, backendChat, requestId); + // Migrated legacy Copilot CLI sessions have no per-turn checkpoints, so + // their turn changeset is always empty even when the session committed + // real work on its branch. Fall back to the session-wide branch changeset + // (the same source the Agents window shows) so those changes surface in + // the chat editor too. Strictly scoped to adopted sessions' latest turn, + // so native sessions and earlier turns are completely unaffected (#333642). + const branchFallbackObs = this._createBranchFallbackDiffsObservable(backendSession, requestId); let lastSource: TurnDiffSource | undefined; const select = (source: TurnDiffSource, diffs: readonly IEditSessionEntryDiff[], status?: ChangesetStatus): readonly IEditSessionEntryDiff[] => { @@ -163,18 +171,33 @@ export class AgentHostResponseFileChangesProvider extends Disposable implements // before anything has been shown. return derivedObservableWithCache(this, (reader, lastValue) => { const retained = lastValue ?? []; - if (!turnChangesetUriObs.read(reader)) { - return select('unsupported', retained); - } - const changesetState = changesetStateObs.read(reader).read(reader); + const turnUri = turnChangesetUriObs.read(reader); + const changesetState = turnUri ? changesetStateObs.read(reader).read(reader) : undefined; const changeset = changesetState instanceof Error ? undefined : changesetState; const changesetDiffs = changeset?.files .map(file => this._changesetFileToEntryDiff(file)) .filter(isDefined); + // A non-empty per-turn changeset is always authoritative (e.g. a turn + // added after migration, which does have checkpoints), so it takes + // precedence over the branch fallback. if (changesetDiffs?.length) { return select('changeset', changesetDiffs, changeset?.status); } + + // The per-turn sources produced nothing. For a migrated session's + // latest turn the session-wide branch changeset carries the committed + // work; `branchFallbackObs` is empty for every non-adopted case, so the + // remaining branches below stay byte-for-byte identical for native + // sessions. + const branchDiffs = branchFallbackObs.read(reader); + if (branchDiffs.length) { + return select('branchFallback', branchDiffs, changeset?.status); + } + + if (!turnUri) { + return select('unsupported', retained); + } if (changeset?.status === ChangesetStatus.Ready && retained.length === 0) { return select('authoritativeEmpty', [], changeset.status); } @@ -186,6 +209,46 @@ export class AgentHostResponseFileChangesProvider extends Disposable implements }); } + /** + * The session-wide branch changeset, exposed as a per-turn fallback but only + * for the specific turn recorded as an adopted legacy Copilot CLI session's + * final migrated turn. That turn has no per-turn checkpoint, so without this + * its committed-on-branch work never appears in the chat editor (#333642). + * Every other case — native sessions, earlier turns, and any turn added after + * adoption (which has its own real per-turn changeset) — yields an empty list, + * so this never alters the changes shown for those turns. + */ + private _createBranchFallbackDiffsObservable(backendSession: URI, requestId: string): IObservable { + const sessionStateObs = this._subscribe(StateComponents.Session, constObservable(backendSession)); + + const branchChangesetUriObs = derivedOpts({ equalsFn: isEqual }, reader => { + const sessionState = sessionStateObs.read(reader).read(reader); + if (!sessionState || sessionState instanceof Error) { + return undefined; + } + // Gate on the durable migration boundary rather than "latest turn": a + // post-adoption no-op turn is also an authoritatively-empty latest turn, + // and must show its own (empty) changes, not the historical aggregate. + if (readSessionEhcliLastMigratedTurn(sessionState._meta) !== requestId) { + return undefined; + } + return URI.parse(buildBranchChangesetUri(backendSession.toString())); + }); + + const branchChangesetStateObs = this._subscribe(StateComponents.Changeset, branchChangesetUriObs); + + return derived(reader => { + if (!branchChangesetUriObs.read(reader)) { + return []; + } + const state = branchChangesetStateObs.read(reader).read(reader); + const changeset = state instanceof Error ? undefined : state; + return changeset?.files + .map(file => this._changesetFileToEntryDiff(file)) + .filter(isDefined) ?? []; + }); + } + private _createFileEditDiffsObservable(backendSession: URI, backendChat: URI | undefined, requestId: string): IObservable { const sessionStateObs = this._subscribe(StateComponents.Session, constObservable(backendSession)); const defaultChatUri = URI.parse(buildDefaultChatUri(backendSession.toString())); diff --git a/src/vs/workbench/contrib/chat/browser/chat.shared.contribution.ts b/src/vs/workbench/contrib/chat/browser/chat.shared.contribution.ts index c725f689f41b38..5aa9aa64424d89 100644 --- a/src/vs/workbench/contrib/chat/browser/chat.shared.contribution.ts +++ b/src/vs/workbench/contrib/chat/browser/chat.shared.contribution.ts @@ -448,7 +448,7 @@ configurationRegistry.registerConfiguration({ enum: [AgentHostExternalSessionsMode.None, AgentHostExternalSessionsMode.Recent, AgentHostExternalSessionsMode.Last24Hours, AgentHostExternalSessionsMode.Last7Days, AgentHostExternalSessionsMode.Last30Days], enumDescriptions: [ nls.localize('chat.agentSessions.showExternal.none', "Do not show external sessions."), - nls.localize('chat.agentSessions.showExternal.recent', "Show up to the 2 most recent external sessions updated in the last 7 days. Once at least 2 local sessions exist, external sessions older than the second-newest local session are hidden."), + nls.localize('chat.agentSessions.showExternal.recent', "Show up to the 2 most recent external sessions updated in the last 7 days. At startup, external sessions older than the second-most-recently updated local session are hidden."), nls.localize('chat.agentSessions.showExternal.last24Hours', "Show external sessions updated in the last 24 hours."), nls.localize('chat.agentSessions.showExternal.last7Days', "Show external sessions updated in the last 7 days."), nls.localize('chat.agentSessions.showExternal.last30Days', "Show external sessions updated in the last 30 days."), diff --git a/src/vs/workbench/contrib/chat/test/browser/agentHost/agentHostResponseFileChanges.test.ts b/src/vs/workbench/contrib/chat/test/browser/agentHost/agentHostResponseFileChanges.test.ts index f9e07342283ddb..f02a26a8a7a14c 100644 --- a/src/vs/workbench/contrib/chat/test/browser/agentHost/agentHostResponseFileChanges.test.ts +++ b/src/vs/workbench/contrib/chat/test/browser/agentHost/agentHostResponseFileChanges.test.ts @@ -12,7 +12,7 @@ import { mock } from '../../../../../../base/test/common/mock.js'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/test/common/utils.js'; import { IAgentConnection } from '../../../../../../platform/agentHost/common/agentService.js'; import { NullLogService } from '../../../../../../platform/log/common/log.js'; -import { buildTurnChangesetUri } from '../../../../../../platform/agentHost/common/changesetUri.js'; +import { buildBranchChangesetUri, buildTurnChangesetUri } from '../../../../../../platform/agentHost/common/changesetUri.js'; import { fromAgentHostUri } from '../../../../../../platform/agentHost/common/agentHostUri.js'; import { IAgentSubscription } from '../../../../../../platform/agentHost/common/state/agentSubscription.js'; import { @@ -87,6 +87,22 @@ suite('AgentHostResponseFileChangesProvider', () => { } as unknown as SessionState; } + /** As {@link sessionStateWithTurnSupport} but flagged as an adopted legacy Copilot CLI session whose final migrated turn is `lastMigratedTurnId`. */ + function adoptedSessionStateWithTurnSupport(lastMigratedTurnId: string): SessionState { + return { + changesets: [{ label: 'This Turn', uriTemplate: buildTurnChangesetUri(backendSession.toString(), '{turnId}'), changeKind: 'turn' }], + _meta: { ehcliAdopted: true, ehcliLastMigratedTurn: lastMigratedTurnId }, + } as unknown as SessionState; + } + + function branchChangesetUri(): string { + return URI.parse(buildBranchChangesetUri(backendSession.toString())).toString(); + } + + function branchFile(path: string, added: number, removed: number): unknown { + return { id: path, edit: { after: { uri: URI.file(path).toString(), content: { uri: `git-blob:/${path}` } }, diff: { added, removed } } }; + } + function createProvider( conn: IAgentConnection, resolveBackendSession: () => URI | undefined = () => backendSession, @@ -231,6 +247,62 @@ suite('AgentHostResponseFileChangesProvider', () => { assert.deepStrictEqual(latest(), []); }); + test('the recorded migrated turn falls back to the branch changeset when its turn changeset is empty', () => { + // #333642: migrated legacy Copilot CLI sessions have no per-turn + // checkpoints, so the committed-on-branch work only lives in the + // session-wide branch changeset. Surface it under the recorded migration + // boundary turn so the chat editor shows the same changes as the Agents window. + const ds = store.add(new DisposableStore()); + const conn = new FakeAgentConnection(); + const defaultChatUri = URI.parse(buildDefaultChatUri(backendSession.toString())); + const provider = ds.add(createProvider(conn, () => backendSession, () => defaultChatUri)); + + conn.setState(backendSession.toString(), adoptedSessionStateWithTurnSupport('t1')); + conn.setState(turnChangesetUri('t1'), { status: ChangesetStatus.Ready, files: [] } satisfies ChangesetState); + conn.setState(branchChangesetUri(), { status: ChangesetStatus.Ready, files: [branchFile('/repo/committed.ts', 4, 2)] } as unknown as ChangesetState); + + const { latest } = observe(provider, ds); + assert.deepStrictEqual(latest().map(d => ({ modified: d.modifiedURI.path, added: d.added, removed: d.removed })), [ + { modified: '/repo/committed.ts', added: 4, removed: 2 }, + ]); + }); + + test('a post-adoption turn with an empty changeset never shows the historical branch aggregate', () => { + // A no-op turn added after migration is authoritatively empty; it must show + // its own (empty) changes, not the migrated session's committed history. + // The recorded boundary turn is 't1'; the requested turn 't2' is later. + const ds = store.add(new DisposableStore()); + const conn = new FakeAgentConnection(); + const defaultChatUri = URI.parse(buildDefaultChatUri(backendSession.toString())); + const provider = ds.add(createProvider(conn, () => backendSession, () => defaultChatUri)); + + conn.setState(backendSession.toString(), adoptedSessionStateWithTurnSupport('t1')); + conn.setState(turnChangesetUri('t2'), { status: ChangesetStatus.Ready, files: [] } satisfies ChangesetState); + conn.setState(branchChangesetUri(), { status: ChangesetStatus.Ready, files: [branchFile('/repo/committed.ts', 4, 2)] } as unknown as ChangesetState); + + const obs = provider.getChangesForRequest(chatResource, 't2')!; + let latest: readonly IEditSessionEntryDiff[] = []; + ds.add(autorun(r => { latest = obs.read(r); })); + assert.deepStrictEqual(latest, []); + }); + + test('a native session never shows the branch changeset in place of an empty turn changeset', () => { + // The fallback is gated on the durable migration boundary, so a normal + // session with an authoritative empty turn changeset stays empty even if a + // branch changeset exists. + const ds = store.add(new DisposableStore()); + const conn = new FakeAgentConnection(); + const defaultChatUri = URI.parse(buildDefaultChatUri(backendSession.toString())); + const provider = ds.add(createProvider(conn, () => backendSession, () => defaultChatUri)); + + conn.setState(backendSession.toString(), sessionStateWithTurnSupport()); + conn.setState(turnChangesetUri('t1'), { status: ChangesetStatus.Ready, files: [] } satisfies ChangesetState); + conn.setState(branchChangesetUri(), { status: ChangesetStatus.Ready, files: [branchFile('/repo/committed.ts', 4, 2)] } as unknown as ChangesetState); + + const { latest } = observe(provider, ds); + assert.deepStrictEqual(latest(), []); + }); + test('keeps a turn visible across changeset recomputes and losses', () => { const ds = store.add(new DisposableStore()); const conn = new FakeAgentConnection(); diff --git a/src/vs/workbench/contrib/terminal/browser/xterm/xtermTerminal.ts b/src/vs/workbench/contrib/terminal/browser/xterm/xtermTerminal.ts index 4fa25000012be2..2bf7f8bbec1535 100644 --- a/src/vs/workbench/contrib/terminal/browser/xterm/xtermTerminal.ts +++ b/src/vs/workbench/contrib/terminal/browser/xterm/xtermTerminal.ts @@ -122,7 +122,6 @@ export class XtermTerminal extends Disposable implements IXtermTerminal, IDetach private readonly _xtermColorProvider: IXtermColorProvider; private readonly _capabilities: ITerminalCapabilityStore; private readonly _disableOverviewRuler: boolean; - private readonly _mainDocument: Document; private static _suggestedRendererType: 'dom' | undefined = undefined; private _attached?: { container: HTMLElement; options: IXtermAttachToElementOptions }; @@ -234,7 +233,6 @@ export class XtermTerminal extends Disposable implements IXtermTerminal, IDetach this._xtermColorProvider = options.xtermColorProvider; this._capabilities = options.capabilities; this._disableOverviewRuler = options.disableOverviewRuler ?? false; - this._mainDocument = layoutService.mainContainer.ownerDocument; const font = this._terminalConfigurationService.getFont(dom.getActiveWindow(), undefined, true); const config = this._terminalConfigurationService.config; @@ -244,7 +242,7 @@ export class XtermTerminal extends Disposable implements IXtermTerminal, IDetach allowProposedApi: true, cols: options.cols, rows: options.rows, - documentOverride: this._mainDocument, + documentOverride: layoutService.mainContainer.ownerDocument, altClickMovesCursor: config.altClickMovesCursor && editorOptions.multiCursorModifier === 'alt', scrollback: config.scrollback, theme: this.getXtermTheme(), @@ -895,7 +893,7 @@ export class XtermTerminal extends Disposable implements IXtermTerminal, IDetach if (!this.raw.element) { return; } - const customGlyphs = this._getWebglCustomGlyphs(); + const customGlyphs = this._terminalConfigurationService.config.customGlyphs; if ((this._webglAddon || this._webglAddonLoading) && this._webglAddonCustomGlyphs === customGlyphs) { return; } @@ -927,7 +925,7 @@ export class XtermTerminal extends Disposable implements IXtermTerminal, IDetach return; } - const currentCustomGlyphs = this._getWebglCustomGlyphs(); + const currentCustomGlyphs = this._terminalConfigurationService.config.customGlyphs; if (customGlyphs !== currentCustomGlyphs) { this._webglAddonCustomGlyphs = undefined; await this._enableWebglRenderer(); @@ -961,11 +959,6 @@ export class XtermTerminal extends Disposable implements IXtermTerminal, IDetach } } - private _getWebglCustomGlyphs(): boolean { - // The custom glyph rasterizer creates a canvas through the rendering document, which is blocked in auxiliary windows. - return this._terminalConfigurationService.config.customGlyphs && this.raw.element?.ownerDocument === this._mainDocument; - } - @debounce(100) private async _refreshLigaturesAddon(): Promise { if (!this.raw.element) { @@ -1150,9 +1143,6 @@ export class XtermTerminal extends Disposable implements IXtermTerminal, IDetach refresh() { this._updateTheme(); this._decorationAddon.refreshLayouts(); - if (this._webglAddon || this._webglAddonLoading) { - this._enableWebglRenderer(); - } } private async _updateUnicodeVersion(): Promise { diff --git a/src/vs/workbench/contrib/terminal/test/browser/xterm/xtermTerminal.test.ts b/src/vs/workbench/contrib/terminal/test/browser/xterm/xtermTerminal.test.ts index 68b48ce1a31e73..66e976ccf9a5b9 100644 --- a/src/vs/workbench/contrib/terminal/test/browser/xterm/xtermTerminal.test.ts +++ b/src/vs/workbench/contrib/terminal/test/browser/xterm/xtermTerminal.test.ts @@ -149,7 +149,7 @@ suite('XtermTerminal', () => { }); }); - test('disables custom glyphs when moved into an auxiliary window', async () => { + test('keeps custom glyphs enabled when moved out of an auxiliary window', async () => { await configurationService.setUserConfiguration('terminal.integrated', { ...defaultTerminalConfig, gpuAcceleration: 'on', @@ -161,12 +161,6 @@ suite('XtermTerminal', () => { } }); - const mainContainer = document.createElement('div'); - document.body.appendChild(mainContainer); - store.add(toDisposable(() => mainContainer.remove())); - xterm.attachToElement(mainContainer); - await timeout(0); - const iframe = document.createElement('iframe'); document.body.appendChild(iframe); store.add(toDisposable(() => iframe.remove())); @@ -179,20 +173,21 @@ suite('XtermTerminal', () => { }; store.add(toDisposable(() => auxiliaryDocument.createElement = createElement)); - auxiliaryContainer.appendChild(xterm.raw.element!); - xterm.raw.open(xterm.raw.element!); - xterm.refresh(); + xterm.attachToElement(auxiliaryContainer); await timeout(0); + const mainContainer = document.createElement('div'); + document.body.appendChild(mainContainer); + store.add(toDisposable(() => mainContainer.remove())); mainContainer.appendChild(xterm.raw.element!); xterm.raw.open(xterm.raw.element!); xterm.refresh(); await timeout(0); - deepStrictEqual(TestWebglAddon.customGlyphOptions, [true, false, true]); + deepStrictEqual(TestWebglAddon.customGlyphOptions, [true]); }); - test('does not load stale custom glyph settings when moved during addon import', async () => { + test('keeps custom glyphs enabled when moved during addon import', async () => { await configurationService.setUserConfiguration('terminal.integrated', { ...defaultTerminalConfig, gpuAcceleration: 'on', @@ -226,7 +221,7 @@ suite('XtermTerminal', () => { xterm.refresh(); await timeout(0); - deepStrictEqual(TestWebglAddon.customGlyphOptions, [false]); + deepStrictEqual(TestWebglAddon.customGlyphOptions, [true]); }); suite('getContentsAsText', () => {