Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor execution results #14697

Merged
merged 3 commits into from
Nov 10, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions src/api.proposed.kernelApi.d.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,11 @@ declare module './api' {
>;
}
export interface Kernels {
/**
* Whether the access to the Kernels has been revoked.
* This happens when the user has not provided consent to the API being used by the requesting extension.
*/
isRevoked: boolean;
/**
* Finds a kernel for a given resource.
* For instance if the resource is a notebook, then look for a kernel associated with the given Notebook document.
Expand Down
11 changes: 5 additions & 6 deletions src/api.proposed.kernelApiAccess.d.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,12 +5,11 @@
import type { CancellationToken } from 'vscode';

declare module './api' {
/**
* Provides access to the Jupyter Kernel API.
* As Kernels can be used to execute code on local or remote machines, this poses a threat to security.
* As a result users will be prompted to allow access to the Kernel API.
*/
export interface Jupyter {
getKernelApi(): Promise<Kernels | undefined>;
/**
* Request access to Kernels.
* As Kernels can be used to execute code on local or remote machines, user consent will be required.
*/
requestKernelAccess(): Thenable<Kernels>;
}
}
22 changes: 5 additions & 17 deletions src/api.proposed.kernelCodeExecution.d.ts
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

import type { CancellationToken, Event } from 'vscode';
import type { CancellationToken } from 'vscode';

declare module './api' {
interface OutputItem {
Expand All @@ -14,28 +14,16 @@ declare module './api' {
*/
data: Uint8Array;
}

interface ExecutionResult {
/**
* Resolves when the execution has completed.
*/
done: Promise<void>;

/**
* Event fired with the output items emitted as a result of the execution.
*/
onDidEmitOutput: Event<OutputItem[]>;
}
export interface Kernel {
/**
* Executes code in the kernel.
* The code executed will not result in changes to the execution count
* & will not show up in the Kernel execution history.
*
* @param {string} code Code to be executed.
* @param {CancellationToken} token Triggers the cancellation of the execution.
* @return {*} {ExecutionResult}
* @param code Code to be executed.
* @param token Triggers the cancellation of the execution.
* @returns Async iterable of output items, that completes when the execution is complete.
*/
executeCode(code: string, token: CancellationToken): ExecutionResult;
executeCode(code: string, token: CancellationToken): AsyncIterable<OutputItem[]>;
}
}
2 changes: 1 addition & 1 deletion src/extension.node.ts
Original file line number Diff line number Diff line change
Expand Up @@ -140,7 +140,7 @@ export async function activate(context: IExtensionContext): Promise<IExtensionAp
createJupyterServerCollection: () => {
throw new Error('Not Implemented');
},
getKernelApi: () => Promise.resolve(undefined)
requestKernelAccess: () => Promise.reject(new Error('Not Implemented'))
};
}
}
Expand Down
2 changes: 1 addition & 1 deletion src/extension.web.ts
Original file line number Diff line number Diff line change
Expand Up @@ -144,7 +144,7 @@ export async function activate(context: IExtensionContext): Promise<IExtensionAp
createJupyterServerCollection: () => {
throw new Error('Not Implemented');
},
getKernelApi: () => Promise.resolve(undefined)
requestKernelAccess: () => Promise.reject(new Error('Not Implemented'))
};
}
}
Expand Down
1 change: 1 addition & 0 deletions src/kernels/api/api.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ export function getKernelsApi(extensionId: string): Kernels {
}

api = {
isRevoked: false,
findKernel(query: { uri: Uri }) {
const kernelProvider = ServiceContainer.instance.get<IKernelProvider>(IKernelProvider);
const notebooks = ServiceContainer.instance.get<IVSCodeNotebook>(IVSCodeNotebook);
Expand Down
39 changes: 19 additions & 20 deletions src/kernels/api/api.vscode.common.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -24,10 +24,10 @@ import {
} from '../../test/datascience/notebook/helper';
import { getKernelsApi } from './api';
import { raceTimeoutError } from '../../platform/common/utils/async';
import { ExecutionResult } from '../../api';
import { dispose } from '../../platform/common/utils/lifecycle';
import { IKernel, IKernelProvider } from '../types';
import { IControllerRegistration, IVSCodeNotebookController } from '../../notebooks/controllers/types';
import { OutputItem } from '../../api';

suiteMandatory('Kernel API Tests @python', function () {
const disposables: IDisposable[] = [];
Expand Down Expand Up @@ -127,29 +127,28 @@ suiteMandatory('Kernel API Tests @python', function () {
);
});

async function waitForOutput(executionResult: ExecutionResult, expectedOutput: string, expectedMimetype: string) {
async function waitForOutput(
executionResult: AsyncIterable<OutputItem[]>,
expectedOutput: string,
expectedMimetype: string
) {
const disposables: IDisposable[] = [];
const outputsReceived: string[] = [];
const outputPromise = new Promise<void>((resolve, reject) => {
executionResult.onDidEmitOutput(
(e) => {
traceInfo(`Output received ${e.length} & mime types are ${e.map((item) => item.mime).join(', ')}}`);
e.forEach((item) => {
if (item.mime === expectedMimetype) {
const output = new TextDecoder().decode(item.data).trim();
if (output === expectedOutput.trim()) {
resolve();
} else {
reject(new Error(`Unexpected output ${output}`));
}
const outputPromise = new Promise<void>(async (resolve, reject) => {
for await (const items of executionResult) {
items.forEach((item) => {
if (item.mime === expectedMimetype) {
const output = new TextDecoder().decode(item.data).trim();
if (output === expectedOutput.trim()) {
resolve();
} else {
outputsReceived.push(`${item.mime} ${new TextDecoder().decode(item.data).trim()}`);
reject(new Error(`Unexpected output ${output}`));
}
});
},
undefined,
disposables
);
} else {
outputsReceived.push(`${item.mime} ${new TextDecoder().decode(item.data).trim()}`);
}
});
}
});

await raceTimeoutError(
Expand Down
30 changes: 19 additions & 11 deletions src/kernels/api/kernel.ts
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

import { l10n, CancellationToken, Event, EventEmitter, ProgressLocation, extensions, window, Disposable } from 'vscode';
import { ExecutionResult, Kernel } from '../../api';
import { l10n, CancellationToken, ProgressLocation, extensions, window, Disposable, Event } from 'vscode';
import { Kernel, OutputItem } from '../../api';
import { ServiceContainer } from '../../platform/ioc/container';
import { IKernel, IKernelProvider, INotebookKernelExecution } from '../types';
import { getDisplayNameOrNameOfKernelConnection } from '../helpers';
Expand Down Expand Up @@ -123,7 +123,7 @@ class WrappedKernelPerExtension implements Kernel {
once(kernel.onDisposed)(() => this.progress.dispose());
}

executeCode(code: string, token: CancellationToken): ExecutionResult {
async *executeCode(code: string, token: CancellationToken): AsyncGenerator<OutputItem[], void, unknown> {
this.previousProgress?.dispose();
let completed = false;
const measures = {
Expand Down Expand Up @@ -156,16 +156,14 @@ class WrappedKernelPerExtension implements Kernel {
if (!this.kernel.session?.kernel) {
properties.failed = true;
sendApiExecTelemetry(this.kernel, measures, properties).catch(noop);
if (this.status === 'dead' || this.status === 'terminating') {
if (this.kernel.status === 'dead' || this.kernel.status === 'terminating') {
throw new Error('Kernel is dead or terminating');
}
throw new Error('Kernel connection not available to execute 3rd party code');
}

const disposables: IDisposable[] = [];
const done = createDeferred<void>();
const onDidEmitOutput = new EventEmitter<{ mime: string; data: Uint8Array }[]>();
disposables.push(onDidEmitOutput);
disposables.push({
dispose: () => {
measures.duration = stopwatch.elapsedTime;
Expand All @@ -179,6 +177,8 @@ class WrappedKernelPerExtension implements Kernel {
const kernelExecution = ServiceContainer.instance
.get<IKernelProvider>(IKernelProvider)
.getKernelExecution(this.kernel);
const outputs: OutputItem[][] = [];
let outputsReceieved = createDeferred<void>();
kernelExecution
.executeCode(code, this.extensionId, token)
.then((codeExecution) => {
Expand Down Expand Up @@ -206,7 +206,7 @@ class WrappedKernelPerExtension implements Kernel {
codeExecution.onDidEmitOutput(
(e) => {
e.forEach((item) => mimeTypes.add(item.mime));
onDidEmitOutput.fire(e);
outputs.push(e);
},
this,
disposables
Expand All @@ -232,10 +232,18 @@ class WrappedKernelPerExtension implements Kernel {
this,
disposables
);
return {
done: done.promise,
onDidEmitOutput: onDidEmitOutput.event
};
while (true) {
await Promise.race([outputsReceieved.promise, done.promise]);
if (outputsReceieved.completed) {
outputsReceieved = createDeferred<void>();
}
while (outputs.length) {
yield outputs.shift()!;
}
if (done.completed) {
break;
}
}
}
}

Expand Down
4 changes: 2 additions & 2 deletions src/standalone/api/api.ts
Original file line number Diff line number Diff line change
Expand Up @@ -213,8 +213,8 @@ export function buildApi(
extensions.determineExtensionFromCallStack().extensionId
);
},
getKernelApi() {
return Promise.resolve(undefined);
requestKernelAccess() {
return Promise.reject(new Error('Not Implemeted'));
}
};

Expand Down
Loading