@vscode/chat-lib 0.0.1

package/src/_internal/platform/tokenizer/node/tikTokenizerImpl.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 { createTokenizer, getRegexByEncoder, getSpecialTokensByEncoder, TikTokenizer } from '@microsoft/tiktokenizer';
+import { MovingAverage } from '../../../util/vs/base/common/numbers';
+import { StopWatch } from '../../../util/vs/base/common/stopwatch';
+import { parseTikTokenBinary } from './parseTikTokens';
+
+export type TokenDictionaryParser = (file: string) => string | Map<Uint8Array, number>;
+
+export class TikTokenImpl {
+
+	private static _instance: TikTokenImpl | undefined;
+
+	private _values: (TikTokenizer | undefined)[] = [];
+	private _stats = {
+		encodeDuration: new MovingAverage(),
+		textLength: new MovingAverage(),
+		callCount: 0,
+	};
+
+	private constructor() { }
+
+	static get instance(): TikTokenImpl {
+		if (!this._instance) {
+			this._instance = new TikTokenImpl();
+		}
+		return this._instance;
+	}
+
+	init(tokenFilePath: string, encoderName: string, useBinaryTokens: boolean): number {
+		const handle = this._values.length;
+		const parser: TokenDictionaryParser = useBinaryTokens ? parseTikTokenBinary : f => f;
+
+		this._values.push(createTokenizer(
+			parser(tokenFilePath),
+			getSpecialTokensByEncoder(encoderName),
+			getRegexByEncoder(encoderName),
+			64000
+		));
+
+		return handle;
+	}
+
+	encode(handle: number, text: string, allowedSpecial?: readonly string[]): number[] {
+		const sw = StopWatch.create(true);
+		const result = this._values[handle]!.encode(text, allowedSpecial);
+
+		this._stats.callCount += 1;
+		this._stats.encodeDuration.update(sw.elapsed());
+		this._stats.textLength.update(text.length);
+
+		return result;
+	}
+
+	destroy(handle: number) {
+		this._values[handle] = undefined;
+	}
+
+	resetStats() {
+		const oldValue = this._stats;
+		const result = {
+			callCount: oldValue.callCount,
+			encodeDuration: oldValue.encodeDuration.value,
+			textLength: oldValue.textLength.value
+		};
+		this._stats.encodeDuration = new MovingAverage();
+		this._stats.textLength = new MovingAverage();
+		this._stats.callCount = 0;
+		return result;
+	}
+}

package/src/_internal/platform/tokenizer/node/tikTokenizerWorker.ts

@@ -0,0 +1,24 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { parentPort } from 'worker_threads';
+import { TikTokenImpl } from './tikTokenizerImpl';
+
+function main() {
+	const port = parentPort;
+	if (!port) {
+		throw new Error(`This module should only be used in a worker thread.`);
+	}
+	port.on('message', async (message: { id: number; fn: string; args: any[] }) => {
+		try {
+			const res = await (<any>TikTokenImpl.instance)[message.fn](...message.args);
+			port.postMessage({ id: message.id, res });
+		} catch (err) {
+			port.postMessage({ id: message.id, err });
+		}
+	});
+}
+
+main();

package/src/_internal/platform/tokenizer/node/tokenizer.ts

@@ -0,0 +1,357 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { OutputMode, Raw, toMode } from '@vscode/prompt-tsx';
+import type { LanguageModelChatTool } from 'vscode';
+import { LRUCache } from '../../../util/common/cache';
+import { getImageDimensions } from '../../../util/common/imageUtils';
+import { createServiceIdentifier } from '../../../util/common/services';
+import { ITokenizer, TokenizerType } from '../../../util/common/tokenizer';
+import { WorkerWithRpcProxy } from '../../../util/node/worker';
+import { assertNever } from '../../../util/vs/base/common/assert';
+import { Disposable, toDisposable } from '../../../util/vs/base/common/lifecycle';
+import { basename, join } from '../../../util/vs/base/common/path';
+import { ITelemetryService } from '../../telemetry/common/telemetry';
+import { TikTokenImpl } from './tikTokenizerImpl';
+
+export const ITokenizerProvider = createServiceIdentifier<ITokenizerProvider>('ITokenizerProvider');
+
+export interface TokenizationEndpoint {
+	readonly tokenizer: TokenizerType;
+}
+
+export interface ITokenizerProvider {
+	readonly _serviceBrand: undefined;
+
+	acquireTokenizer(endpoint: TokenizationEndpoint): ITokenizer;
+}
+
+/**
+ * BaseTokensPerCompletion is the minimum tokens for a completion request.
+ * Replies are primed with <|im_start|>assistant<|message|>, so these tokens represent the
+ * special token and the role name.
+ */
+export const BaseTokensPerCompletion = 3;
+/*
+ * Each GPT 3.5 / GPT 4 message comes with 3 tokens per message due to special characters
+ */
+export const BaseTokensPerMessage = 3;
+/*
+ * Since gpt-3.5-turbo-0613 each name costs 1 token
+ */
+export const BaseTokensPerName = 1;
+
+
+export class TokenizerProvider implements ITokenizerProvider {
+	declare readonly _serviceBrand: undefined;
+
+	// These files are copied to `dist` via the `build/postinstall.ts` script
+	private readonly _cl100kTokenizer: BPETokenizer;
+	private readonly _o200kTokenizer: BPETokenizer;
+
+	constructor(
+		useWorker: boolean,
+		@ITelemetryService telmetryService: ITelemetryService
+	) {
+		// if we're running from dist, the dictionary is compressed, but if we're  running
+		// in e.g. a `spec` file we should load the dictionary using default behavior.
+		// todo: cleanup a bit, have an IS_BUILT constant?
+		this._cl100kTokenizer = new BPETokenizer(useWorker, join(__dirname, './cl100k_base.tiktoken'), 'cl100k_base', telmetryService);
+		this._o200kTokenizer = new BPETokenizer(useWorker, join(__dirname, './o200k_base.tiktoken'), 'o200k_base', telmetryService);
+	}
+
+	dispose() {
+		this._cl100kTokenizer.dispose();
+		this._o200kTokenizer.dispose();
+	}
+
+	/**
+	 * Gets a tokenizer for a given model family
+	 * @param endpoint The endpoint you want to acquire a tokenizer for
+	 */
+	public acquireTokenizer(endpoint: TokenizationEndpoint): ITokenizer {
+		switch (endpoint.tokenizer) {
+			case TokenizerType.CL100K:
+				return this._cl100kTokenizer;
+			case TokenizerType.O200K:
+				return this._o200kTokenizer;
+			default:
+				throw new Error(`Unknown tokenizer: ${endpoint.tokenizer}`);
+		}
+	}
+}
+
+type TikTokenWorker = {
+	encode(text: string, allowedSpecial?: ReadonlyArray<string>): Promise<number[]>;
+};
+
+class BPETokenizer extends Disposable implements ITokenizer {
+
+	private _tokenizer?: Promise<TikTokenWorker>;
+
+	/**
+	 * TikToken has its own cache, but it still does some processing
+	 * until a cache hit. We can have a much more efficient cache that
+	 * directly looks up string -> token length
+	 */
+	private readonly _cache = new LRUCache<number>(5000);
+
+	protected readonly baseTokensPerMessage = BaseTokensPerMessage;
+	protected readonly baseTokensPerName = BaseTokensPerName;
+
+	public readonly mode = OutputMode.Raw;
+
+	constructor(
+		private readonly _useWorker: boolean,
+		private readonly _tokenFilePath: string,
+		private readonly _encoderName: string,
+		@ITelemetryService private readonly _telemetryService: ITelemetryService
+	) {
+		super();
+	}
+
+	async countMessagesTokens(messages: Raw.ChatMessage[]): Promise<number> {
+		let numTokens = BaseTokensPerMessage;
+		for (const message of messages) {
+			numTokens += await this.countMessageTokens(message);
+		}
+		return numTokens;
+	}
+
+	/**
+	 * Tokenizes the given text.
+	 * @param text The text to tokenize.
+	 * @returns The tokenized text.
+	 */
+	private async tokenize(text: string): Promise<number[]> {
+		return (await this.ensureTokenizer()).encode(text);
+	}
+
+	/**
+	 * Calculates the token length of the given text.
+	 * @param text The text to calculate the token length for.
+	 * @returns The number of tokens in the text.
+	 */
+	async tokenLength(text: string | Raw.ChatCompletionContentPart): Promise<number> {
+		if (typeof text === 'string') {
+			return this._textTokenLength(text);
+		}
+
+		switch (text.type) {
+			case Raw.ChatCompletionContentPartKind.Text:
+				return this._textTokenLength(text.text);
+			case Raw.ChatCompletionContentPartKind.Opaque:
+				return text.tokenUsage || 0;
+			case Raw.ChatCompletionContentPartKind.Image:
+				if (text.imageUrl.url.startsWith('data:image/')) {
+					try {
+						return calculateImageTokenCost(text.imageUrl.url, text.imageUrl.detail);
+					} catch {
+						return this._textTokenLength(text.imageUrl.url);
+					}
+				}
+				return this._textTokenLength(text.imageUrl.url);
+			case Raw.ChatCompletionContentPartKind.CacheBreakpoint:
+				return 0;
+			default:
+				assertNever(text, `unknown content part (${JSON.stringify(text)})`);
+		}
+	}
+
+	private async _textTokenLength(text: string) {
+		if (!text) {
+			return 0;
+		}
+		let cacheValue = this._cache.get(text);
+		if (!cacheValue) {
+			cacheValue = (await this.tokenize(text)).length;
+			this._cache.put(text, cacheValue);
+		}
+		return cacheValue;
+	}
+
+	/**
+	 * Counts tokens for a single chat message within a completion request.
+	 *
+	 * Follows https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb for GPT 3.5/4 models.
+	 *
+	 * **Note**: The result does not include base tokens for the completion itself.
+	 */
+	async countMessageTokens(message: Raw.ChatMessage): Promise<number> {
+		return this.baseTokensPerMessage + (await this.countMessageObjectTokens(toMode(OutputMode.OpenAI, message)));
+	}
+
+	async countToolTokens(tools: LanguageModelChatTool[]): Promise<number> {
+		const baseToolTokens = 16;
+		let numTokens = 0;
+		if (tools.length) {
+			numTokens += baseToolTokens;
+		}
+
+		const baseTokensPerTool = 8;
+		for (const tool of tools) {
+			numTokens += baseTokensPerTool;
+			numTokens += await this.countObjectTokens({ name: tool.name, description: tool.description, parameters: tool.inputSchema });
+		}
+
+		// This is an estimate, so give a little safety margin
+		return Math.floor(numTokens * 1.1);
+	}
+
+	private async countMessageObjectTokens(obj: any): Promise<number> {
+		let numTokens = 0;
+		for (const [key, value] of Object.entries(obj)) {
+			if (!value) {
+				continue;
+			}
+
+			if (typeof value === 'string') {
+				numTokens += await this.tokenLength(value);
+			} else if (value) {
+				const casted = value as any;
+				if (casted.type === 'text') {
+					numTokens += await this.tokenLength(casted.text);
+				} else if (casted.type === 'image_url' && casted.image_url) {
+					if (casted.image_url.url.startsWith('data:image/')) {
+						try {
+							numTokens += calculateImageTokenCost(casted.image_url.url, casted.image_url.detail);
+						} catch {
+							numTokens += await this.tokenLength(casted.image_url.url);
+						}
+					} else {
+						numTokens += await this.tokenLength(casted.image_url.url);
+					}
+				} else {
+					let newTokens = await this.countMessageObjectTokens(value);
+					if (key === 'tool_calls') {
+						// This is an estimate, not including all of the overhead, so give a little safety margin
+						newTokens = Math.floor(newTokens * 1.5);
+					}
+
+					numTokens += newTokens;
+				}
+			}
+
+			if (key === 'name' && value !== undefined) {
+				numTokens += this.baseTokensPerName;
+			}
+		}
+
+		return numTokens;
+	}
+
+	private async countObjectTokens(obj: any): Promise<number> {
+		let numTokens = 0;
+		for (const [key, value] of Object.entries(obj)) {
+			if (!value) {
+				continue;
+			}
+
+			numTokens += await this.tokenLength(key);
+			if (typeof value === 'string') {
+				numTokens += await this.tokenLength(value);
+			} else if (value) {
+				numTokens += await this.countMessageObjectTokens(value);
+			}
+		}
+
+		return numTokens;
+	}
+
+	private ensureTokenizer(): Promise<TikTokenWorker> {
+		this._tokenizer ??= this.doInitTokenizer();
+		return this._tokenizer;
+	}
+
+	private async doInitTokenizer(): Promise<TikTokenWorker> {
+
+		const useBinaryTokens = basename(__dirname) === 'dist';
+
+		if (!this._useWorker) {
+			const handle = TikTokenImpl.instance.init(this._tokenFilePath, this._encoderName, useBinaryTokens);
+
+			const cleanup = toDisposable(() => {
+				TikTokenImpl.instance.destroy(handle);
+				this._store.deleteAndLeak(cleanup);
+				this._tokenizer = undefined;
+			});
+			this._store.add(cleanup);
+
+			return {
+				encode: async (text, allowedSpecial) => {
+					return TikTokenImpl.instance.encode(handle, text, allowedSpecial);
+				}
+			};
+		} else {
+
+			const workerPath = join(__dirname, 'tikTokenizerWorker.js');
+			const worker = new WorkerWithRpcProxy<TikTokenImpl>(workerPath, { name: `TikToken worker (${this._encoderName})` });
+			const handle = await worker.proxy.init(this._tokenFilePath, this._encoderName, useBinaryTokens);
+
+			const cleanup = toDisposable(() => {
+				worker.terminate();
+				this._store.deleteAndLeak(cleanup);
+				this._tokenizer = undefined;
+			});
+
+			let timeout: TimeoutHandle;
+
+			return {
+				encode: (text, allowedSpecial) => {
+					const result = worker.proxy.encode(handle, text, allowedSpecial);
+
+					clearTimeout(timeout);
+					timeout = setTimeout(() => cleanup.dispose(), 15000);
+
+					if (Math.random() < 1 / 1000) {
+						worker.proxy.resetStats().then(stats => {
+							/* __GDPR__
+								"tokenizer.stats" : {
+									"owner": "jrieken",
+									"comment": "Perf stats about tokenizers",
+									"callCount": { "classification": "SystemMetaData", "purpose": "FeatureInsight", "isMeasurement": true, "comment": "How often tokenize was called" },
+									"encodeDuration": { "classification": "SystemMetaData", "purpose": "FeatureInsight", "isMeasurement": true, "comment": "Average time encode took" },
+									"textLength": { "classification": "SystemMetaData", "purpose": "FeatureInsight", "isMeasurement": true, "comment": "Average length of text that got encoded" }
+								}
+							*/
+							this._telemetryService.sendMSFTTelemetryEvent('tokenizer.stats', undefined, stats);
+						});
+					}
+
+					return result;
+				}
+			};
+		}
+	}
+}
+
+
+//#region Image tokenizer helpers
+
+// https://platform.openai.com/docs/guides/vision#calculating-costs
+export function calculateImageTokenCost(imageUrl: string, detail: 'low' | 'high' | undefined): number {
+	let { width, height } = getImageDimensions(imageUrl);
+
+	if (detail === 'low') {
+		return 85;
+	}
+
+	// Scale image to fit within a 2048 x 2048 square if necessary.
+	if (width > 2048 || height > 2048) {
+		const scaleFactor = 2048 / Math.max(width, height);
+		width = Math.round(width * scaleFactor);
+		height = Math.round(height * scaleFactor);
+	}
+
+	const scaleFactor = 768 / Math.min(width, height);
+	width = Math.round(width * scaleFactor);
+	height = Math.round(height * scaleFactor);
+
+	const tiles = Math.ceil(width / 512) * Math.ceil(height / 512);
+
+	return tiles * 170 + 85;
+}
+
+//#endregion

package/src/_internal/platform/workspace/common/workspaceService.ts

@@ -0,0 +1,226 @@
+/*---------------------------------------------------------------------------------------------
+ *  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, FileSystem, NotebookData, NotebookDocument, NotebookDocumentChangeEvent, TextDocument, TextDocumentChangeEvent, TextEditorSelectionChangeEvent, Uri, WorkspaceEdit, WorkspaceFolder, WorkspaceFoldersChangeEvent } from 'vscode';
+import { findNotebook } from '../../../util/common/notebooks';
+import { createServiceIdentifier } from '../../../util/common/services';
+import * as path from '../../../util/vs/base/common/path';
+import { extUriBiasedIgnorePathCase, relativePath } from '../../../util/vs/base/common/resources';
+import { URI } from '../../../util/vs/base/common/uri';
+import { NotebookDocumentSnapshot } from '../../editing/common/notebookDocumentSnapshot';
+import { TextDocumentSnapshot } from '../../editing/common/textDocumentSnapshot';
+import { DisposableStore, IDisposable } from '../../../util/vs/base/common/lifecycle';
+import { Emitter } from '../../../util/vs/base/common/event';
+
+export const IWorkspaceService = createServiceIdentifier<IWorkspaceService>('IWorkspaceService');
+
+export interface IWorkspaceService {
+	readonly _serviceBrand: undefined;
+	textDocuments: readonly TextDocument[];
+	notebookDocuments: readonly NotebookDocument[];
+	readonly onDidOpenTextDocument: Event<TextDocument>;
+	readonly onDidCloseTextDocument: Event<TextDocument>;
+	readonly onDidOpenNotebookDocument: Event<NotebookDocument>;
+	readonly onDidCloseNotebookDocument: Event<NotebookDocument>;
+	readonly onDidChangeTextDocument: Event<TextDocumentChangeEvent>;
+	readonly onDidChangeNotebookDocument: Event<NotebookDocumentChangeEvent>;
+	readonly onDidChangeWorkspaceFolders: Event<WorkspaceFoldersChangeEvent>;
+	readonly onDidChangeTextEditorSelection: Event<TextEditorSelectionChangeEvent>;
+	openTextDocument(uri: Uri): Promise<TextDocument>;
+	fs: FileSystem;
+	showTextDocument(document: TextDocument): Promise<void>;
+	openTextDocumentAndSnapshot(uri: Uri): Promise<TextDocumentSnapshot>;
+	openNotebookDocumentAndSnapshot(uri: Uri, format: 'xml' | 'json' | 'text'): Promise<NotebookDocumentSnapshot>;
+	openNotebookDocument(uri: Uri): Promise<NotebookDocument>;
+	openNotebookDocument(notebookType: string, content?: NotebookData): Promise<NotebookDocument>;
+	getWorkspaceFolders(): URI[];
+	getWorkspaceFolder(resource: URI): URI | undefined;
+	getWorkspaceFolderName(workspaceFolderUri: URI): string;
+	showWorkspaceFolderPicker(): Promise<WorkspaceFolder | undefined>;
+
+	asRelativePath(pathOrUri: string | Uri, includeWorkspaceFolder?: boolean): string;
+	applyEdit(edit: WorkspaceEdit): Thenable<boolean>;
+
+	/**
+	 * Ensures that the workspace has fully loaded before returning. This is useful for
+	 * virtual workspaces where we need to ensure that the contents of the workspace
+	 * has been downloaded before we can use them.
+	 */
+	ensureWorkspaceIsFullyLoaded(): Promise<void>;
+}
+
+export abstract class AbstractWorkspaceService implements IWorkspaceService {
+	declare readonly _serviceBrand: undefined;
+	abstract textDocuments: readonly TextDocument[];
+	abstract notebookDocuments: readonly NotebookDocument[];
+	abstract readonly onDidOpenTextDocument: Event<TextDocument>;
+	abstract readonly onDidCloseTextDocument: Event<TextDocument>;
+	abstract readonly onDidOpenNotebookDocument: Event<NotebookDocument>;
+	abstract readonly onDidCloseNotebookDocument: Event<NotebookDocument>;
+	abstract readonly onDidChangeTextDocument: Event<TextDocumentChangeEvent>;
+	abstract readonly onDidChangeWorkspaceFolders: Event<WorkspaceFoldersChangeEvent>;
+	abstract readonly onDidChangeNotebookDocument: Event<NotebookDocumentChangeEvent>;
+	abstract readonly onDidChangeTextEditorSelection: Event<TextEditorSelectionChangeEvent>;
+	abstract openTextDocument(uri: Uri): Promise<TextDocument>;
+	abstract fs: FileSystem;
+	abstract showTextDocument(document: TextDocument): Promise<void>;
+	abstract openNotebookDocument(uri: Uri): Promise<NotebookDocument>;
+	abstract openNotebookDocument(notebookType: string, content?: NotebookData): Promise<NotebookDocument>;
+	abstract getWorkspaceFolders(): URI[];
+	abstract ensureWorkspaceIsFullyLoaded(): Promise<void>;
+	abstract showWorkspaceFolderPicker(): Promise<WorkspaceFolder | undefined>;
+	abstract getWorkspaceFolderName(workspaceFolderUri: URI): string;
+	abstract applyEdit(edit: WorkspaceEdit): Thenable<boolean>;
+	asRelativePath(pathOrUri: string | Uri, includeWorkspaceFolder?: boolean): string {
+		// Copied from the implementation in vscode/extHostWorkspace.ts
+		let resource: URI | undefined;
+		let path: string = '';
+		if (typeof pathOrUri === 'string') {
+			resource = URI.file(pathOrUri);
+			path = pathOrUri;
+		} else if (typeof pathOrUri !== 'undefined') {
+			resource = pathOrUri;
+			path = pathOrUri.fsPath;
+		}
+
+		if (!resource) {
+			return path;
+		}
+
+		const folder = this.getWorkspaceFolder(resource);
+
+		if (!folder) {
+			return path;
+		}
+
+		if (typeof includeWorkspaceFolder === 'undefined') {
+			includeWorkspaceFolder = this.getWorkspaceFolders().length > 1;
+		}
+
+		let result = relativePath(folder, resource);
+		if (includeWorkspaceFolder) {
+			const name = this.getWorkspaceFolderName(folder);
+			result = `${name}/${result}`;
+		}
+		return result!;
+	}
+
+	async openTextDocumentAndSnapshot(uri: Uri): Promise<TextDocumentSnapshot> {
+		const doc = await this.openTextDocument(uri);
+		return TextDocumentSnapshot.create(doc);
+	}
+
+	async openNotebookDocumentAndSnapshot(uri: Uri, format: 'xml' | 'json' | 'text'): Promise<NotebookDocumentSnapshot> {
+		// Possible we have an untitled file opened as a notebook.
+		const doc = findNotebook(uri, this.notebookDocuments) || await this.openNotebookDocument(uri);
+
+		return NotebookDocumentSnapshot.create(doc, format);
+	}
+
+	getWorkspaceFolder(resource: URI): URI | undefined {
+		return this.getWorkspaceFolders().find(folder => extUriBiasedIgnorePathCase.isEqualOrParent(resource, folder));
+	}
+}
+
+export function getWorkspaceFileDisplayPath(workspaceService: IWorkspaceService, file: URI): string {
+	const workspaceUri = workspaceService.getWorkspaceFolder(file);
+	return workspaceUri ? path.posix.relative(workspaceUri.path, file.path) : file.path;
+}
+
+export class NullWorkspaceService extends AbstractWorkspaceService implements IDisposable {
+	override fs!: FileSystem;
+	private readonly disposables = new DisposableStore();
+
+	public readonly didOpenTextDocumentEmitter = this.disposables.add(new Emitter<TextDocument>());
+	public readonly didCloseTextDocumentEmitter = this.disposables.add(new Emitter<TextDocument>());
+	public readonly didOpenNotebookDocumentEmitter = this.disposables.add(new Emitter<NotebookDocument>());
+	public readonly didCloseNotebookDocumentEmitter = this.disposables.add(new Emitter<NotebookDocument>());
+	public readonly didChangeTextDocumentEmitter = this.disposables.add(new Emitter<TextDocumentChangeEvent>());
+	public readonly didChangeWorkspaceFoldersEmitter = this.disposables.add(new Emitter<WorkspaceFoldersChangeEvent>());
+	public readonly didChangeNotebookDocumentEmitter = this.disposables.add(new Emitter<NotebookDocumentChangeEvent>());
+	public readonly didChangeTextEditorSelectionEmitter = this.disposables.add(new Emitter<TextEditorSelectionChangeEvent>());
+
+	public override readonly onDidChangeTextDocument = this.didChangeTextDocumentEmitter.event;
+	public override readonly onDidCloseTextDocument = this.didCloseTextDocumentEmitter.event;
+	public override readonly onDidOpenNotebookDocument = this.didOpenNotebookDocumentEmitter.event;
+	public override readonly onDidCloseNotebookDocument = this.didCloseNotebookDocumentEmitter.event;
+	public override readonly onDidOpenTextDocument = this.didOpenTextDocumentEmitter.event;
+	public override readonly onDidChangeWorkspaceFolders = this.didChangeWorkspaceFoldersEmitter.event;
+	public override readonly onDidChangeNotebookDocument = this.didChangeNotebookDocumentEmitter.event;
+	public override readonly onDidChangeTextEditorSelection = this.didChangeTextEditorSelectionEmitter.event;
+
+	private readonly workspaceFolder: URI[];
+	private readonly _textDocuments: TextDocument[] = [];
+	private readonly _notebookDocuments: NotebookDocument[] = [];
+
+	constructor(workspaceFolders: URI[] = [], textDocuments: TextDocument[] = [], notebookDocuments: NotebookDocument[] = []) {
+		super();
+		this.workspaceFolder = workspaceFolders;
+		this._textDocuments = textDocuments;
+		this._notebookDocuments = notebookDocuments;
+	}
+
+	get textDocuments(): TextDocument[] {
+		return this._textDocuments;
+	}
+
+	override showTextDocument(document: TextDocument): Promise<void> {
+		return Promise.resolve();
+	}
+
+	override async openTextDocument(uri: Uri): Promise<TextDocument> {
+		const doc = this.textDocuments.find(d => d.uri.toString() === uri.toString());
+		if (doc) {
+			return doc;
+		}
+
+		throw new Error(`Unknown document: ${uri}`);
+	}
+
+	override async openNotebookDocument(uri: Uri): Promise<NotebookDocument>;
+	override async openNotebookDocument(notebookType: string, content?: NotebookData): Promise<NotebookDocument>;
+	override async openNotebookDocument(arg1: Uri | string, arg2?: NotebookData): Promise<NotebookDocument> {
+		if (typeof arg1 === 'string') {
+			// Handle the overload for notebookType and content
+			throw new Error('Not implemented');
+		} else {
+			const notebook = this.notebookDocuments.find(d => d.uri.toString() === arg1.toString());
+			if (notebook) {
+				return notebook;
+			}
+
+			throw new Error(`Unknown notebook: ${arg1}`);
+		}
+	}
+
+	get notebookDocuments(): readonly NotebookDocument[] {
+		return this._notebookDocuments;
+	}
+
+	getWorkspaceFolders(): URI[] {
+		return this.workspaceFolder;
+	}
+
+	override getWorkspaceFolderName(workspaceFolderUri: URI): string {
+		return 'default';
+	}
+
+	override ensureWorkspaceIsFullyLoaded(): Promise<void> {
+		// We aren't using virtual workspaces here, so we can just return
+		return Promise.resolve();
+	}
+
+	showWorkspaceFolderPicker(): Promise<undefined> {
+		return Promise.resolve(undefined);
+	}
+
+	override applyEdit(): Promise<boolean> {
+		return Promise.resolve(true);
+	}
+
+	public dispose() {
+		this.disposables.dispose();
+	}
+}