feat: offer inline-script env setup as an unresolved-import quick fix

src/common/localize.ts

@@ -26,6 +26,19 @@ export namespace WorkbenchStrings {
 export namespace InlineScriptStrings {
     export const updateExtension = l10n.t('Update Extension');
 
+    /**
+     * Quick fix title offered on an unresolved import in a PEP 723 script.
+     *
+     * Deliberately describes what the action *does* (set the environment up) rather than promising
+     * to fix the import: setup installs the block's declared `dependencies` verbatim, which may not
+     * include the module that is actually unresolved.
+     */
+    export const setUpScriptEnvironment = l10n.t("Set up this script's Python environment");
+
+    export const saveFailedBeforeSetup = l10n.t(
+        'Could not save this script, so its environment was not set up. Save the file and try again.',
+    );
+
     export const updatePythonExtension = l10n.t(
         'The environment for this script was created. Update the Python extension for the full inline script experience.',
     );

src/common/telemetry/constants.ts

@@ -253,8 +253,29 @@ export enum EventNames {
      * - duration: number (ms between the detection and the first edit)
      */
     INLINE_SCRIPT_EDITED = 'inlineScript.edited',
+    /**
+     * Telemetry event fired once per user-initiated inline-script environment
+     * setup, recording which surface started it. This is the adoption measure
+     * for the unresolved-import quick fix: the CodeLens is hidden while a
+     * document is dirty, so `codeaction` counts setups that the CodeLens alone
+     * could not have produced.
+     * Properties:
+     * - trigger: 'codelens' | 'codeaction' | 'bulk' (which surface invoked setup)
+     * - outcome: 'created' | 'notCreated' | 'error'
+     *
+     * `notCreated` covers every benign or reported non-creation (cancelled,
+     * skipped, no compatible Python, ...); the failure taxonomy itself already
+     * ships on `inlineScript.envError` and is not duplicated here.
+     */
+    INLINE_SCRIPT_SETUP_INVOKED = 'inlineScript.setupInvoked',
 }
 
+/** Surface that started an inline-script environment setup. */
+export type InlineScriptSetupTrigger = 'codelens' | 'codeaction' | 'bulk';
+
+/** Result of one inline-script environment setup attempt, as seen by the invoking surface. */
+export type InlineScriptSetupOutcomeKind = 'created' | 'notCreated' | 'error';
+
 export type InlineScriptEnvErrorCategory =
     | 'compatible-python-declined'
     | 'discovery-failure'
@@ -778,4 +799,15 @@ export interface IEventNamePropertyMapping {
         }
     */
     [EventNames.INLINE_SCRIPT_EDITED]: never | undefined;
+
+    /* __GDPR__
+        "inlineScript.setupInvoked": {
+            "trigger": { "classification": "SystemMetaData", "purpose": "FeatureInsight", "owner": "StellaHuang95" },
+            "outcome": { "classification": "SystemMetaData", "purpose": "FeatureInsight", "owner": "StellaHuang95" }
+        }
+    */
+    [EventNames.INLINE_SCRIPT_SETUP_INVOKED]: {
+        trigger: InlineScriptSetupTrigger;
+        outcome: InlineScriptSetupOutcomeKind;
+    };
 }

src/features/inlineScript/codeLens.ts

@@ -13,6 +13,7 @@ import {
     TextDocument,
 } from 'vscode';
 import { InlineScriptRoutingRegistry } from '../../common/inlineScript/routingRegistry';
+import { InlineScriptSetupTrigger } from '../../common/telemetry/constants';
 
 /**
  * Shows a single "Set up environment for this script" CodeLens above a `.py` file's PEP 723
@@ -67,11 +68,12 @@ export class InlineScriptCodeLensProvider implements CodeLensProvider, Disposabl
         const offset = metadata.sourceRange?.start ?? metadata.range.start;
         const position = document.positionAt(offset);
         const range = new Range(position, position);
+        const trigger: InlineScriptSetupTrigger = 'codelens';
         return [
             new CodeLens(range, {
                 title: l10n.t('Set up environment for this script'),
                 command: this.setupCommand,
-                arguments: [uri],
+                arguments: [uri, trigger],
             }),
         ];
     }

src/features/inlineScript/setupCodeAction.ts

@@ -0,0 +1,170 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import {
+    CancellationToken,
+    CodeAction,
+    CodeActionContext,
+    CodeActionKind,
+    CodeActionProvider,
+    Diagnostic,
+    Disposable,
+    languages,
+    Range,
+    TextDocument,
+} from 'vscode';
+import { MAX_HEADER_BYTES, readInlineScriptMetadata } from '../../common/inlineScript/metadata';
+import { getInlineScriptRoutingKey, InlineScriptRoutingRegistry } from '../../common/inlineScript/routingRegistry';
+import { InlineScriptStrings } from '../../common/localize';
+import { InlineScriptSetupTrigger } from '../../common/telemetry/constants';
+import { isInlineScriptsFeatureEnabled } from '../../helpers';
+
+/**
+ * Diagnostic codes that mean "this import did not resolve", across every type checker a user of this
+ * extension is likely to have enabled. Stored lowercased; compare with {@link normalizeDiagnosticCode}.
+ *
+ * Matching is on `code` only, never on `source`: Pyrefly-backed Pylance reports its source as the
+ * literal string `pylance + pyrefly`, so any source allow-list would be wrong somewhere.
+ *
+ * `reportMissingModuleSource` is included deliberately, diverging from Pylance's own
+ * `isMissingImportDiagnostic`, which excludes it because a stub-resolved module is correctly spelled
+ * and so has nothing for a "change spelling" fix to suggest. For us the meaning is the opposite kind
+ * of useful: a stub was found but the source was not, i.e. the package is not installed — which is
+ * exactly what setting the script's environment up addresses.
+ */
+const UNRESOLVED_IMPORT_DIAGNOSTIC_CODES: ReadonlySet<string> = new Set([
+    // Pyright / Pylance / basedpyright.
+    'reportmissingimports',
+    'reportmissingmodulesource',
+    // Ty (kebab-case), mapped to the two rules above by Pylance's TyDiagnosticCodeMapper.
+    'unresolved-import',
+    'possibly-missing-import',
+    // Pyrefly (kebab-case `ErrorKind` names), mapped by Pylance's PyreflyDiagnosticCodeMapper.
+    'missing-import',
+    'missing-source',
+    'missing-source-for-stubs',
+    // mypy, via the separate ms-python.mypy-type-checker extension.
+    'import-not-found',
+    'import-untyped',
+]);
+
+/**
+ * Reduce `Diagnostic.code` — which is `string | number | { value: string | number; target: Uri }` —
+ * to a lowercased string, or `undefined` when the diagnostic carries no code.
+ */
+function normalizeDiagnosticCode(code: Diagnostic['code']): string | undefined {
+    if (code === undefined || code === null) {
+        return undefined;
+    }
+    const value = typeof code === 'object' ? code.value : code;
+    return typeof value === 'string' || typeof value === 'number' ? String(value).toLowerCase() : undefined;
+}
+
+/**
+ * Whether `diagnostic` reports an import that could not be resolved. See
+ * {@link UNRESOLVED_IMPORT_DIAGNOSTIC_CODES} for the dialects covered and for why `source` is ignored.
+ */
+export function isUnresolvedImportDiagnostic(diagnostic: Diagnostic): boolean {
+    const code = normalizeDiagnosticCode(diagnostic.code);
+    return code !== undefined && UNRESOLVED_IMPORT_DIAGNOSTIC_CODES.has(code);
+}
+
+/**
+ * The head of `document`'s in-memory text, bounded to the same byte budget that
+ * `readInlineScriptMetadataFromFile` reads from disk.
+ *
+ * Bounding it matters twice over: it keeps this provider's work constant regardless of file size,
+ * and it keeps what the quick fix can see identical to what setup will later parse off disk, so the
+ * action is never offered for a block that setup would not find.
+ */
+function getInlineScriptHeaderText(document: TextDocument): string {
+    const text = document.getText();
+    if (Buffer.byteLength(text, 'utf-8') <= MAX_HEADER_BYTES) {
+        return text;
+    }
+    // Truncating on a byte boundary can split a multi-byte character; the disk reader's bounded
+    // `read` has exactly the same behaviour, so the two stay in agreement.
+    return Buffer.from(text, 'utf-8').subarray(0, MAX_HEADER_BYTES).toString('utf-8');
+}
+
+/**
+ * Offers "Set up this script's Python environment" as a quick fix on an unresolved import in a `.py`
+ * file that declares a PEP 723 `# /// script` block and has no inline-script environment yet.
+ *
+ * This exists because the CodeLens is the feature's only other entry point and `provideCodeLenses`
+ * returns nothing while `document.isDirty` — so it is absent at the one moment a user most needs it,
+ * right after typing `import requests` and seeing the squiggle. This provider parses the in-memory
+ * buffer instead, so it works on an unsaved edit.
+ *
+ * Deliberate non-promises, both in wording and in mechanics:
+ *  - the title says what the action does, not that the squiggle will clear — the unresolved module
+ *    may be undeclared in the block, or declared under a different distribution name (`PIL` vs
+ *    `pillow`);
+ *  - `diagnostics` is left unset, because populating it would tell VS Code this action *resolves*
+ *    those diagnostics and would opt it into fix-all affordances;
+ *  - `isPreferred` is left unset, so it never pre-empts a real import fix such as "add import".
+ *
+ * The action disappears once the script is set up (`shouldRoute`) and comes back by itself if the
+ * user later edits the metadata block, because a changed metadata identity resets the registry's
+ * validated association.
+ */
+export class InlineScriptSetupCodeActionProvider implements CodeActionProvider {
+    constructor(
+        private readonly routing: InlineScriptRoutingRegistry,
+        private readonly setupCommand: string,
+    ) {}
+
+    /**
+     * Gates are ordered cheapest-first because VS Code may call this on every cursor move. The
+     * `context.diagnostics` test comes before any parsing: the common case is a file with no
+     * unresolved import at the cursor, and that case must cost nothing but a short array scan.
+     */
+    public provideCodeActions(
+        document: TextDocument,
+        _range: Range,
+        context: CodeActionContext,
+        _token: CancellationToken,
+    ): CodeAction[] {
+        if (!isInlineScriptsFeatureEnabled()) {
+            return [];
+        }
+        if (!context.diagnostics.some(isUnresolvedImportDiagnostic)) {
+            return [];
+        }
+        const uri = document.uri;
+        if (!getInlineScriptRoutingKey(uri)) {
+            // Not a local `.py` file, so it can never carry an inline-script environment.
+            return [];
+        }
+        if (this.routing.shouldRoute(uri)) {
+            // A validated inline-script environment matching the current metadata already exists.
+            return [];
+        }
+        if (!readInlineScriptMetadata(getInlineScriptHeaderText(document), uri.fsPath)) {
+            return [];
+        }
+        const action = new CodeAction(InlineScriptStrings.setUpScriptEnvironment, CodeActionKind.QuickFix);
+        const trigger: InlineScriptSetupTrigger = 'codeaction';
+        action.command = {
+            title: InlineScriptStrings.setUpScriptEnvironment,
+            command: this.setupCommand,
+            arguments: [uri, trigger],
+        };
+        return [action];
+    }
+}
+
+/**
+ * Register the inline-script quick fix for local `.py` files. Only called when the PEP 723
+ * inline-script feature flag is enabled, so it is a no-op for everyone else.
+ */
+export function registerInlineScriptSetupCodeAction(
+    routing: InlineScriptRoutingRegistry,
+    setupCommand: string,
+): Disposable {
+    return languages.registerCodeActionsProvider(
+        { scheme: 'file', language: 'python' },
+        new InlineScriptSetupCodeActionProvider(routing, setupCommand),
+        { providedCodeActionKinds: [CodeActionKind.QuickFix] },
+    );
+}