Git - wire-up git repository state (#296563)

* Initial implementation

* Get the initial state working

* Pull request feedback

src/vs/workbench/api/browser/mainThreadGitExtensionService.ts

@@ -3,12 +3,14 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { CancellationToken } from '../../../base/common/cancellation.js';
+import { Sequencer } from '../../../base/common/async.js'; import { CancellationToken } from '../../../base/common/cancellation.js';
 import { Disposable } from '../../../base/common/lifecycle.js';
+import { ResourceMap } from '../../../base/common/map.js';
 import { URI } from '../../../base/common/uri.js';
-import { IGitExtensionDelegate, IGitService, GitRef, GitRefQuery, GitRefType } from '../../contrib/git/common/gitService.js';
+import { GitRepository } from '../../contrib/git/browser/gitService.js';
+import { IGitExtensionDelegate, IGitService, GitRef, GitRefQuery, GitRefType, GitRepositoryState, GitBranch, IGitRepository } from '../../contrib/git/common/gitService.js';
 import { extHostNamedCustomer, IExtHostContext } from '../../services/extensions/common/extHostCustomers.js';
-import { ExtHostContext, ExtHostGitExtensionShape, GitRefTypeDto, MainContext, MainThreadGitExtensionShape } from '../common/extHost.protocol.js';
+import { ExtHostContext, ExtHostGitExtensionShape, GitRefTypeDto, GitRepositoryStateDto, MainContext, MainThreadGitExtensionShape } from '../common/extHost.protocol.js';
 
 function toGitRefType(type: GitRefTypeDto): GitRefType {
 	switch (type) {
@@ -19,13 +21,35 @@ function toGitRefType(type: GitRefTypeDto): GitRefType {
 	}
 }
 
+function toGitRepositoryState(dto: GitRepositoryStateDto | undefined): GitRepositoryState {
+	return {
+		HEAD: dto?.HEAD ? {
+			type: toGitRefType(dto.HEAD.type),
+			name: dto.HEAD.name,
+			commit: dto.HEAD.commit,
+			remote: dto.HEAD.remote,
+			upstream: dto.HEAD.upstream,
+			ahead: dto.HEAD.ahead,
+			behind: dto.HEAD.behind,
+		} satisfies GitBranch : undefined,
+	};
+}
+
 @extHostNamedCustomer(MainContext.MainThreadGitExtension)
 export class MainThreadGitExtensionService extends Disposable implements MainThreadGitExtensionShape, IGitExtensionDelegate {
 	private readonly _proxy: ExtHostGitExtensionShape;
+	private readonly _openRepositorySequencer = new Sequencer();
+
+	private _repositoryHandles = new ResourceMap<number>();
+	private _repositories = new Map<number, IGitRepository>();
+
+	get repositories(): Iterable<IGitRepository> {
+		return this._repositories.values();
+	}
 
 	constructor(
 		extHostContext: IExtHostContext,
-		@IGitService private readonly gitService: IGitService,
+		@IGitService private readonly gitService: IGitService
 	) {
 		super();
 
@@ -44,13 +68,51 @@ export class MainThreadGitExtensionService extends Disposable implements MainThr
 		}
 	}
 
-	async openRepository(uri: URI): Promise<URI | undefined> {
-		const result = await this._proxy.$openRepository(uri);
-		return result ? URI.revive(result) : undefined;
+	private _getRepositoryByUri(uri: URI): IGitRepository | undefined {
+		const handle = this._repositoryHandles.get(uri);
+		return handle !== undefined ? this._repositories.get(handle) : undefined;
+	}
+
+	async openRepository(uri: URI): Promise<IGitRepository | undefined> {
+		return this._openRepositorySequencer.queue(async () => {
+			// Check if we already have a repository for the given URI
+			const existingRepository = this._getRepositoryByUri(uri);
+			if (existingRepository) {
+				return existingRepository;
+			}
+
+			// Open the repository
+			const result = await this._proxy.$openRepository(uri);
+			if (!result) {
+				return undefined;
+			}
+
+			const repositoryRootUri = URI.revive(result.rootUri);
+
+			// Check if we already have a repository for the given root
+			const existingRepositoryForRoot = this._getRepositoryByUri(repositoryRootUri);
+			if (existingRepositoryForRoot) {
+				return existingRepositoryForRoot;
+			}
+
+			// Create a new repository and store it in the maps
+			const state = toGitRepositoryState(result.state);
+			const repository = new GitRepository(this, repositoryRootUri, state);
+
+			this._repositories.set(result.handle, repository);
+			this._repositoryHandles.set(repositoryRootUri, result.handle);
+
+			return repository;
+		});
 	}
 
 	async getRefs(root: URI, query: GitRefQuery, token?: CancellationToken): Promise<GitRef[]> {
-		const result = await this._proxy.$getRefs(root, query, token);
+		const handle = this._repositoryHandles.get(root);
+		if (handle === undefined) {
+			return [];
+		}
+
+		const result = await this._proxy.$getRefs(handle, query, token);
 
 		if (token?.isCancellationRequested) {
 			return [];
@@ -61,4 +123,18 @@ export class MainThreadGitExtensionService extends Disposable implements MainThr
 			type: toGitRefType(ref.type)
 		} satisfies GitRef));
 	}
+
+	async $onDidChangeRepository(handle: number): Promise<void> {
+		const repository = this._repositories.get(handle);
+		if (!repository) {
+			return;
+		}
+
+		const state = await this._proxy.$getRepositoryState(handle);
+		if (!state) {
+			return;
+		}
+
+		repository.setState(toGitRepositoryState(state));
+	}
 }

src/vs/workbench/api/common/extHost.protocol.ts

@@ -119,6 +119,7 @@ export interface IMainContext extends IRPCProtocol {
 // --- main thread
 
 export interface MainThreadGitExtensionShape extends IDisposable {
+	$onDidChangeRepository(handle: number): Promise<void>;
 }
 
 export interface MainThreadClipboardShape extends IDisposable {
@@ -3476,10 +3477,31 @@ export interface GitRefDto {
 	readonly revision: string;
 }
 
+export interface GitRepositoryStateDto {
+	readonly HEAD?: GitBranchDto;
+}
+
+export interface GitBranchDto {
+	readonly name?: string;
+	readonly commit?: string;
+	readonly type: GitRefTypeDto;
+	readonly remote?: string;
+	readonly upstream?: GitUpstreamRefDto;
+	readonly ahead?: number;
+	readonly behind?: number;
+}
+
+export interface GitUpstreamRefDto {
+	readonly remote: string;
+	readonly name: string;
+	readonly commit?: string;
+}
+
 export interface ExtHostGitExtensionShape {
 	$isGitExtensionAvailable(): Promise<boolean>;
-	$openRepository(root: UriComponents): Promise<UriComponents | undefined>;
-	$getRefs(root: UriComponents, query: GitRefQueryDto, token?: CancellationToken): Promise<GitRefDto[]>;
+	$openRepository(root: UriComponents): Promise<{ handle: number; rootUri: UriComponents; state: GitRepositoryStateDto } | undefined>;
+	$getRefs(handle: number, query: GitRefQueryDto, token?: CancellationToken): Promise<GitRefDto[]>;
+	$getRepositoryState(handle: number): Promise<GitRepositoryStateDto | undefined>;
 }
 
 // --- proxy identifiers

src/vs/workbench/api/common/extHostGitExtensionService.ts

@@ -4,13 +4,15 @@
  *--------------------------------------------------------------------------------------------*/
 
 import type * as vscode from 'vscode';
+import { Event } from '../../../base/common/event.js';
 import { Disposable, DisposableStore } from '../../../base/common/lifecycle.js';
 import { URI, UriComponents } from '../../../base/common/uri.js';
 import { ExtensionIdentifier } from '../../../platform/extensions/common/extensions.js';
 import { createDecorator } from '../../../platform/instantiation/common/instantiation.js';
 import { IExtHostExtensionService } from './extHostExtensionService.js';
 import { IExtHostRpcService } from './extHostRpcService.js';
-import { ExtHostGitExtensionShape, GitRefDto, GitRefQueryDto, GitRefTypeDto } from './extHost.protocol.js';
+import { ExtHostGitExtensionShape, GitBranchDto, GitRefDto, GitRefQueryDto, GitRefTypeDto, GitRepositoryStateDto, GitUpstreamRefDto, MainContext, MainThreadGitExtensionShape } from './extHost.protocol.js';
+import { ResourceMap } from '../../../base/common/map.js';
 
 const GIT_EXTENSION_ID = 'vscode.git';
 
@@ -23,11 +25,50 @@ function toGitRefTypeDto(type: GitRefType): GitRefTypeDto {
 	}
 }
 
+function toGitBranchDto(branch: Branch): GitBranchDto {
+	return {
+		name: branch.name,
+		commit: branch.commit,
+		type: toGitRefTypeDto(branch.type),
+		remote: branch.remote,
+		upstream: branch.upstream ? toGitUpstreamRefDto(branch.upstream) : undefined,
+		ahead: branch.ahead,
+		behind: branch.behind,
+	};
+}
+
+function toGitUpstreamRefDto(upstream: UpstreamRef): GitUpstreamRefDto {
+	return {
+		remote: upstream.remote,
+		name: upstream.name,
+		commit: upstream.commit,
+	};
+}
+
 interface Repository {
 	readonly rootUri: vscode.Uri;
+	readonly state: RepositoryState;
+
 	getRefs(query: GitRefQuery, token?: vscode.CancellationToken): Promise<GitRef[]>;
 }
 
+interface RepositoryState {
+	readonly HEAD: Branch | undefined;
+	readonly onDidChange: Event<void>;
+}
+
+interface Branch extends GitRef {
+	readonly upstream?: UpstreamRef;
+	readonly ahead?: number;
+	readonly behind?: number;
+}
+
+interface UpstreamRef {
+	readonly remote: string;
+	readonly name: string;
+	readonly commit?: string;
+}
+
 interface GitRef {
 	type: GitRefType;
 	name?: string;
@@ -65,14 +106,24 @@ export const IExtHostGitExtensionService = createDecorator<IExtHostGitExtensionS
 export class ExtHostGitExtensionService extends Disposable implements IExtHostGitExtensionService {
 	declare readonly _serviceBrand: undefined;
 
+	private static _handlePool: number = 0;
+
 	private _gitApi: GitExtensionAPI | undefined;
+
+	private readonly _proxy: MainThreadGitExtensionShape;
+
+	private readonly _repositories = new Map<number, Repository>();
+	private readonly _repositoryByUri = new ResourceMap<number>();
+
 	private readonly _disposables = this._register(new DisposableStore());
 
 	constructor(
-		@IExtHostRpcService _extHostRpc: IExtHostRpcService,
+		@IExtHostRpcService extHostRpc: IExtHostRpcService,
 		@IExtHostExtensionService private readonly _extHostExtensionService: IExtHostExtensionService,
 	) {
 		super();
+
+		this._proxy = extHostRpc.getProxy(MainContext.MainThreadGitExtension);
 	}
 
 	async $isGitExtensionAvailable(): Promise<boolean> {
@@ -80,23 +131,55 @@ export class ExtHostGitExtensionService extends Disposable implements IExtHostGi
 		return !!registry.getExtensionDescription(GIT_EXTENSION_ID);
 	}
 
-	async $openRepository(uri: UriComponents): Promise<UriComponents | undefined> {
+	async $openRepository(uri: UriComponents): Promise<{ handle: number; rootUri: UriComponents; state: GitRepositoryStateDto } | undefined> {
 		const api = await this._ensureGitApi();
 		if (!api) {
 			return undefined;
 		}
 
 		const repository = await api.openRepository(URI.revive(uri));
-		return repository?.rootUri;
-	}
-
-	async $getRefs(uri: UriComponents, query: GitRefQueryDto, token?: vscode.CancellationToken): Promise<GitRefDto[]> {
-		const api = await this._ensureGitApi();
-		if (!api) {
-			return [];
+		if (!repository) {
+			return undefined;
 		}
 
-		const repository = await api.openRepository(URI.revive(uri));
+		const existingHandle = this._repositoryByUri.get(repository.rootUri);
+		if (existingHandle !== undefined) {
+			return {
+				handle: existingHandle,
+				rootUri: repository.rootUri,
+				state: {
+					HEAD: repository.state.HEAD ? toGitBranchDto(repository.state.HEAD) : undefined
+				}
+			};
+		}
+
+		let repositoryState = repository.state;
+		if (repository.state.HEAD === undefined) {
+			// If the repository is not initialized, wait for it
+			await Event.toPromise(repository.state.onDidChange, this._disposables);
+			repositoryState = repository.state;
+		}
+
+		const handle = ExtHostGitExtensionService._handlePool++;
+
+		this._repositories.set(handle, repository);
+		this._repositoryByUri.set(repository.rootUri, handle);
+
+		this._disposables.add(repository.state.onDidChange(() => {
+			this._proxy.$onDidChangeRepository(handle);
+		}));
+
+		return {
+			handle,
+			rootUri: repository.rootUri,
+			state: {
+				HEAD: repositoryState.HEAD ? toGitBranchDto(repositoryState.HEAD) : undefined
+			}
+		};
+	}
+
+	async $getRefs(handle: number, query: GitRefQueryDto, token?: vscode.CancellationToken): Promise<GitRefDto[]> {
+		const repository = this._repositories.get(handle);
 		if (!repository) {
 			return [];
 		}
@@ -128,6 +211,16 @@ export class ExtHostGitExtensionService extends Disposable implements IExtHostGi
 		}
 	}
 
+	async $getRepositoryState(handle: number): Promise<GitRepositoryStateDto | undefined> {
+		const repository = this._repositories.get(handle);
+		if (!repository) {
+			return undefined;
+		}
+
+		const state = repository.state;
+		return { HEAD: state.HEAD ? toGitBranchDto(state.HEAD) : undefined };
+	}
+
 	private async _ensureGitApi(): Promise<GitExtensionAPI | undefined> {
 		if (this._gitApi) {
 			return this._gitApi;

src/vs/workbench/contrib/git/browser/gitService.ts

@@ -3,23 +3,20 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { Sequencer } from '../../../../base/common/async.js';
 import { CancellationToken } from '../../../../base/common/cancellation.js';
+import { Emitter, Event } from '../../../../base/common/event.js';
 import { BugIndicatingError } from '../../../../base/common/errors.js';
 import { Disposable, IDisposable, toDisposable } from '../../../../base/common/lifecycle.js';
-import { ResourceMap } from '../../../../base/common/map.js';
 import { URI } from '../../../../base/common/uri.js';
-import { IGitService, IGitExtensionDelegate, GitRef, GitRefQuery, IGitRepository } from '../common/gitService.js';
+import { IGitService, IGitExtensionDelegate, GitRef, GitRefQuery, IGitRepository, GitRepositoryState } from '../common/gitService.js';
 
 export class GitService extends Disposable implements IGitService {
 	declare readonly _serviceBrand: undefined;
 
 	private _delegate: IGitExtensionDelegate | undefined;
-	private readonly _openRepositorySequencer = new Sequencer();
 
-	private readonly _repositories = new ResourceMap<IGitRepository>();
 	get repositories(): Iterable<IGitRepository> {
-		return this._repositories.values();
+		return this._delegate?.repositories ?? [];
 	}
 
 	setDelegate(delegate: IGitExtensionDelegate): IDisposable {
@@ -33,48 +30,40 @@ export class GitService extends Disposable implements IGitService {
 		this._delegate = delegate;
 
 		return toDisposable(() => {
-			this._repositories.clear();
 			this._delegate = undefined;
 		});
 	}
 
 	async openRepository(uri: URI): Promise<IGitRepository | undefined> {
-		return this._openRepositorySequencer.queue(async () => {
-			if (!this._delegate) {
-				return undefined;
-			}
+		if (!this._delegate) {
+			return undefined;
+		}
 
-			// Check whether we have an opened repository for the uri
-			let repository = this._repositories.get(uri);
-			if (repository) {
-				return repository;
-			}
-
-			// Open the repository to get the repository root
-			const root = await this._delegate.openRepository(uri);
-			if (!root) {
-				return undefined;
-			}
-
-			const rootUri = URI.revive(root);
-
-			// Check whether we have an opened repository for the root
-			repository = this._repositories.get(rootUri);
-			if (repository) {
-				return repository;
-			}
-
-			// Create a new repository
-			repository = new GitRepository(this._delegate, rootUri);
-			this._repositories.set(rootUri, repository);
-
-			return repository;
-		});
+		return this._delegate.openRepository(uri);
 	}
 }
 
-export class GitRepository implements IGitRepository {
-	constructor(private readonly delegate: IGitExtensionDelegate, readonly rootUri: URI) { }
+export class GitRepository extends Disposable implements IGitRepository {
+	private readonly _onDidChangeState = this._register(new Emitter<void>());
+	readonly onDidChangeState: Event<void> = this._onDidChangeState.event;
+
+	private _state: GitRepositoryState;
+	get state(): GitRepositoryState { return this._state; }
+
+	setState(state: GitRepositoryState): void {
+		this._state = state;
+		this._onDidChangeState.fire();
+	}
+
+	constructor(
+		private readonly delegate: IGitExtensionDelegate,
+		readonly rootUri: URI,
+		state: GitRepositoryState
+	) {
+		super();
+
+		this._state = state;
+	}
 
 	async getRefs(query: GitRefQuery, token?: CancellationToken): Promise<GitRef[]> {
 		return this.delegate.getRefs(this.rootUri, query, token);

src/vs/workbench/contrib/git/common/gitService.ts

@@ -4,8 +4,9 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { CancellationToken } from '../../../../base/common/cancellation.js';
+import { Event } from '../../../../base/common/event.js';
 import { IDisposable } from '../../../../base/common/lifecycle.js';
-import { URI, UriComponents } from '../../../../base/common/uri.js';
+import { URI } from '../../../../base/common/uri.js';
 import { createDecorator } from '../../../../platform/instantiation/common/instantiation.js';
 
 export enum GitRefType {
@@ -28,14 +29,38 @@ export interface GitRefQuery {
 	readonly sort?: 'alphabetically' | 'committerdate' | 'creatordate';
 }
 
+export interface GitRepositoryState {
+	readonly HEAD?: GitBranch;
+}
+
+export interface GitBranch extends GitRef {
+	readonly upstream?: GitUpstreamRef;
+	readonly ahead?: number;
+	readonly behind?: number;
+}
+
+export interface GitUpstreamRef {
+	readonly remote: string;
+	readonly name: string;
+	readonly commit?: string;
+}
+
 export interface IGitRepository {
 	readonly rootUri: URI;
+
+	readonly state: GitRepositoryState;
+	setState(state: GitRepositoryState): void;
+
+	readonly onDidChangeState: Event<void>;
+
 	getRefs(query: GitRefQuery, token?: CancellationToken): Promise<GitRef[]>;
 }
 
 export interface IGitExtensionDelegate {
-	getRefs(uri: UriComponents, query?: GitRefQuery, token?: CancellationToken): Promise<GitRef[]>;
-	openRepository(uri: UriComponents): Promise<UriComponents | undefined>;
+	readonly repositories: Iterable<IGitRepository>;
+	openRepository(uri: URI): Promise<IGitRepository | undefined>;
+
+	getRefs(root: URI, query?: GitRefQuery, token?: CancellationToken): Promise<GitRef[]>;
 }
 
 export const IGitService = createDecorator<IGitService>('gitService');