feat: negotiate client-independent prompt suggestions

README.md

@@ -111,3 +111,5 @@ See [AIR extensions](docs/air-extensions.md#async-tasks) for the capability, lif
 ## License
 
 By contributing, you agree that your contributions will be licensed under the Apache 2.0 License.
+
+See [Prompt suggestions](docs/prompt-suggestions.md) for the optional, independently negotiated next-Prompt prediction extension.

docs/prompt-suggestions.md

@@ -0,0 +1,47 @@
+# Prompt suggestions
+
+This experimental extension lets any ACP client request a predicted next user Prompt without declaring AIR identity. The adapter uses the current Session's authentication, model and conversation. A prediction is optional and never executes tools.
+
+## Negotiation
+
+For ACP v1, declare `initialize.clientCapabilities._meta.promptSuggestions: {version: 1}`. The adapter advertises the same marker in `agentCapabilities._meta.promptSuggestions`.
+
+For ACP v2, declare `initialize.capabilities._meta.promptSuggestions: {version: 1}`. The adapter advertises the marker in `capabilities._meta.promptSuggestions`.
+
+Only version 1 is supported. Existing AIR clients retain their previous behavior when they do not declare this independent extension. A client declaring both uses this extension, with one prediction request at most.
+
+## Per-Prompt consent and correlation
+
+An opted-in client adds this metadata to `session/prompt`:
+
+```json
+{"_meta":{"promptSuggestions":{"enabled":true,"promptId":"opaque-client-prompt-id"}}}
+```
+
+`promptId` is an opaque, nonblank string, at most 1024 characters. The client must use a unique identifier for each Prompt in a live Session. The adapter echoes it without adding it to the model's conversation. It is independent of Codex's own Turn identifier.
+
+Missing or malformed optional metadata, or `enabled: false`, disables generation for that Prompt. Changing the preference affects subsequent Prompts. A client hiding suggestions immediately must also discard already authorized predictions for the current Prompt.
+
+## Updates
+
+A prediction follows successful completion through `session/update`:
+
+```json
+{
+  "sessionId":"agent-session",
+  "update":{
+    "sessionUpdate":"session_info_update",
+    "_meta":{"promptSuggestions":{
+      "version":1,
+      "promptId":"opaque-client-prompt-id",
+      "suggestions":[{"text":"Run the tests"}]
+    }}
+  }
+}
+```
+
+An empty `suggestions` array clears the prediction for that same identifier. Starting another Prompt invalidates a sample in flight and clears the old prediction with its original identifier. Clients must reject updates for older identifiers or replaced Sessions.
+
+No prediction is guaranteed. Plan mode, unsuccessful Turns, disabled generation, blank model output and inference errors can all produce none. Prediction work does not delay the user's Turn result. `CODEX_ACP_PROMPT_SUGGESTIONS=false` remains an authoritative opt-out.
+
+Clients keep ghost text separate from their drafts. Acceptance fills an editable draft and does not send it automatically.

src/CodexAcpServer.ts

@@ -200,7 +200,7 @@ import {
 } from "./AuthStatusMeta";
 import {randomUUID} from "node:crypto";
 import {TitleGenerator} from "./TitleGenerator";
-import {PromptSuggestionGenerator, promptSuggestionsEnabled, promptSuggestionsMeta} from "./PromptSuggestionGenerator";
+import {PromptSuggestionGenerator, promptSuggestionsEnabled, promptSuggestionsMeta, clientSupportsPromptSuggestions, promptSuggestionId, PROMPT_SUGGESTIONS_CAPABILITY} from "./PromptSuggestionGenerator";
 import {once} from "node:events";
 import {
     AIR_AGENT_FILE_CHANGE_REPORT_KEY,
@@ -293,7 +293,7 @@ export interface SessionState extends SessionIndexTitleState {
     sessionTitleSource: "unset" | "fallback" | "explicit" | "unknown";
     sessionFailure?: SessionFailure;
     titleGen?: TitleGenerator;
-    /** Predicts the follow-up prompt of a finished turn; set for AIR clients only. */
+    /** Predicts the follow-up prompt for independently opted-in clients and legacy AIR clients. */
     promptSuggestions?: PromptSuggestionGenerator | undefined;
     subagents: SubagentEventRouter;
     subagentTransition?: Promise<void>;
@@ -861,6 +861,7 @@ export class CodexAcpServer {
                     // Presence means "this agent pushes `_auth/status_update`". It
                     // never carries a payload, and the client never asks for one.
                     [AUTH_STATUS_META_KEY]: authStatusCapability(),
+                    promptSuggestions: PROMPT_SUGGESTIONS_CAPABILITY,
                 },
             },
             authMethods: getCodexAuthMethods(_params.clientCapabilities),
@@ -895,6 +896,7 @@ export class CodexAcpServer {
                 version: packageJson.version,
             },
             capabilities: {
+                _meta: {promptSuggestions: PROMPT_SUGGESTIONS_CAPABILITY},
                 auth: {
                     _meta: {
                         // Presence means "this agent pushes `_auth/status_update`".
@@ -4037,15 +4039,17 @@ export class CodexAcpServer {
     }
 
     private createPromptSuggestions(sessionState: SessionState): PromptSuggestionGenerator | undefined {
-        if (!sessionState.clientCapabilities.airClient || !promptSuggestionsEnabled()) return undefined;
+        const shared = clientSupportsPromptSuggestions(this.clientCapabilities);
+        if ((!sessionState.clientCapabilities.airClient && !shared) || !promptSuggestionsEnabled()) return undefined;
         return new PromptSuggestionGenerator(
             this.codexAcpClient.appServerClient,
-            async text => {
+            async (text, promptId) => {
                 await new ACPSessionConnection(this.connection, sessionState.sessionId).update({
                     sessionUpdate: "session_info_update",
-                    _meta: promptSuggestionsMeta(text),
+                    _meta: promptSuggestionsMeta(text, promptId),
                 });
             },
+            shared,
         );
     }
 
@@ -4906,8 +4910,12 @@ export class CodexAcpServer {
         });
         const sessionState = this.getSessionState(params.sessionId);
         // Awaited only when a clear is sent, so a prompt with no held suggestion starts in the same tick.
-        const clearing = sessionState.promptSuggestions?.clear();
-        if (clearing) await clearing;
+        if (clientSupportsPromptSuggestions(this.clientCapabilities)) {
+            await sessionState.promptSuggestions?.beginPrompt(promptSuggestionId(params._meta));
+        } else {
+            const clearing = sessionState.promptSuggestions?.clear();
+            if (clearing) await clearing;
+        }
         const agentFileChangeReportRequest = clientSupportsAgentFileChangeReports(this.clientCapabilities)
             ? parseAgentFileChangeReportRequest(params._meta)
             : null;

src/PromptSuggestionGenerator.ts

@@ -3,6 +3,22 @@ import type {CodexAppServerClient} from "./CodexAppServerClient";
 /** The `_meta` key the suggestion travels under. It names no vendor and sits outside `jetbrains.air`. */
 export const PROMPT_SUGGESTIONS_META_KEY = "promptSuggestions";
 
+export const PROMPT_SUGGESTIONS_CAPABILITY = {version: 1} as const;
+
+/** The independent extension, not the client's identity or other capabilities. */
+export function clientSupportsPromptSuggestions(capabilities: {_meta?: Record<string, unknown> | null} | null | undefined): boolean {
+    const value = capabilities?._meta?.[PROMPT_SUGGESTIONS_META_KEY];
+    return typeof value === "object" && value !== null && "version" in value && value.version === 1;
+}
+
+/** Each prompt opts in separately. Malformed optional metadata leaves prediction off. */
+export function promptSuggestionId(meta: Record<string, unknown> | null | undefined): string | undefined {
+    const value = meta?.[PROMPT_SUGGESTIONS_META_KEY];
+    if (typeof value !== "object" || value === null || !("enabled" in value) || value.enabled !== true) return undefined;
+    if (!("promptId" in value) || typeof value.promptId !== "string" || !value.promptId.trim() || value.promptId.length > 1024) return undefined;
+    return value.promptId;
+}
+
 // Picking the grounded next move is worth the reasoning; the sample reuses the cached conversation either way.
 const SUGGESTION_EFFORT = "high";
 
@@ -27,12 +43,15 @@ export function promptSuggestionsEnabled(env: NodeJS.ProcessEnv = process.env):
 }
 
 /** The `_meta` of the session info update that publishes `text`. `null` sends an empty list, which clears it. */
-export function promptSuggestionsMeta(text: string | null): Record<string, unknown> {
-    return {[PROMPT_SUGGESTIONS_META_KEY]: {suggestions: text === null ? [] : [{text}]}};
+export function promptSuggestionsMeta(text: string | null, promptId?: string): Record<string, unknown> {
+    return {[PROMPT_SUGGESTIONS_META_KEY]: {
+        ...(promptId === undefined ? {} : {version: 1, promptId}),
+        suggestions: text === null ? [] : [{text}],
+    }};
 }
 
 /** Publishes `text` to the client; `null` clears the suggestion the client holds. */
-export type PublishPromptSuggestion = (text: string | null) => Promise<void>;
+export type PublishPromptSuggestion = (text: string | null, promptId?: string) => Promise<void>;
 
 /**
  * Predicts the follow-up prompt of a finished turn with `thread/sample`: one model response on an ephemeral fork
@@ -41,16 +60,27 @@ export type PublishPromptSuggestion = (text: string | null) => Promise<void>;
 export class PromptSuggestionGenerator {
     private generation = 0;
     private lastPublished: string | null = null;
+    private promptId: string | undefined;
+    private publishedPromptId: string | undefined;
 
     constructor(
         private client: CodexAppServerClient,
         private readonly publish: PublishPromptSuggestion,
+        private readonly requiresPromptOptIn = false,
     ) {}
 
+    /** Clears the previous prediction before associating work with a new client prompt. */
+    async beginPrompt(promptId?: string): Promise<void> {
+        const clearing = this.clear();
+        this.promptId = promptId;
+        if (clearing) await clearing;
+    }
+
     /** Fire-and-forget hook, called after a turn ends normally. A newer turn or {@link invalidate} drops its result. */
     onTurnCompleted(threadId: string, turnId: string): void {
+        if (this.requiresPromptOptIn && this.promptId === undefined) return;
         const generation = ++this.generation;
-        this.generateAndPublish(generation, threadId, turnId).catch(() => {
+        this.generateAndPublish(generation, threadId, turnId, this.promptId).catch(() => {
             // best-effort; never surface errors to the user
         });
     }
@@ -67,7 +97,7 @@ export class PromptSuggestionGenerator {
     clear(): Promise<void> | null {
         this.invalidate();
         if (this.lastPublished === null) return null;
-        return this.publishIfChanged(null).catch(() => {
+        return this.publishIfChanged(null, this.publishedPromptId).catch(() => {
             // best-effort
         });
     }
@@ -77,7 +107,7 @@ export class PromptSuggestionGenerator {
         this.client = client;
     }
 
-    private async generateAndPublish(generation: number, threadId: string, turnId: string): Promise<void> {
+    private async generateAndPublish(generation: number, threadId: string, turnId: string, promptId: string | undefined): Promise<void> {
         // The app-server a generation starts on ends it, see `rebindAppServer`.
         const client = this.client;
         // The app-server rejects the sample when `turnId` is no longer the thread's latest turn.
@@ -89,13 +119,15 @@ export class PromptSuggestionGenerator {
         });
 
         if (generation !== this.generation) return;
-        await this.publishIfChanged(parseSuggestion(response.text));
+        await this.publishIfChanged(parseSuggestion(response.text), promptId);
     }
 
-    private async publishIfChanged(text: string | null): Promise<void> {
-        if (text === this.lastPublished) return;
+    private async publishIfChanged(text: string | null, promptId?: string): Promise<void> {
+        if (text === this.lastPublished && promptId === this.publishedPromptId) return;
         this.lastPublished = text;
-        await this.publish(text);
+        this.publishedPromptId = promptId;
+        if (promptId === undefined) await this.publish(text);
+        else await this.publish(text, promptId);
     }
 }
 

src/__tests__/CodexACPAgent/data/initialize-v1-routed-response.json

@@ -31,7 +31,10 @@
       "sse": false
     },
     "_meta": {
-      "authStatus": {}
+      "authStatus": {},
+      "promptSuggestions": {
+        "version": 1
+      }
     }
   },
   "authMethods": [

src/__tests__/CodexACPAgent/data/initialize-v2-response.json

@@ -26,7 +26,12 @@
         "authStatus": {}
       }
     },
-    "providers": {}
+    "providers": {},
+    "_meta": {
+      "promptSuggestions": {
+        "version": 1
+      }
+    }
   },
   "authMethods": [
     {

src/__tests__/CodexACPAgent/initialize.test.ts

@@ -64,6 +64,7 @@ describe('CodexACPAgent - initialize', () => {
                 },
                 _meta: {
                     authStatus: {},
+                    promptSuggestions: {version: 1},
                 },
             },
             authMethods: getCodexAuthMethods(),

src/__tests__/CodexACPAgent/prompt-suggestions.test.ts

@@ -29,9 +29,9 @@ function suggestionUpdates(client: PromptSession): unknown[] {
 }
 
 /** Runs turn number `index` (from 0) to completion; its id is `turn-${index + 1}`. */
-async function runTurn(client: PromptSession, protocolVersion: 1 | 2, index: number, text: string, answer: string) {
+async function runTurn(client: PromptSession, protocolVersion: 1 | 2, index: number, text: string, answer: string, meta?: Record<string, unknown>) {
     const turnId = `turn-${index + 1}`;
-    const response = client.sendPrompt([{type: "text", text}]);
+    const response = meta ? client.request("session/prompt", {sessionId, prompt: [{type: "text", text}], _meta: meta}) : client.sendPrompt([{type: "text", text}]);
     await vi.waitFor(() => expect(client.turnStartParams).toHaveLength(index + 1));
     const params = client.turnStartParams[index]!;
     client.emit(turnStarted(turnId));
@@ -56,10 +56,13 @@ describe.each([1, 2] as const)("predicted follow-up prompt over ACP v%i", (proto
 
     let sampleParams: Array<Record<string, unknown>> = [];
 
-    function connect(air: boolean) {
+    function connect(air: boolean | "shared") {
         sampleParams = [];
         return connectSession(protocolVersion, {
-            ...(air ? {clientCapabilities: airCapabilities, v1ClientCapabilities: airCapabilities} : {}),
+            ...(air ? {
+                clientCapabilities: air === "shared" ? {_meta: {promptSuggestions: {version: 1}}} : airCapabilities,
+                v1ClientCapabilities: air === "shared" ? {_meta: {promptSuggestions: {version: 1}}} : airCapabilities,
+            } : {}),
             codexResponses: {
                 "thread/sample": async params => {
                     sampleParams.push(params as Record<string, unknown>);
@@ -92,6 +95,22 @@ describe.each([1, 2] as const)("predicted follow-up prompt over ACP v%i", (proto
         ]);
     });
 
+    it("serves a generic opted-in client with correlation, without AIR metadata", async () => {
+        const client = await connect("shared");
+        closeClient = () => client.connection.close();
+        await runTurn(client, protocolVersion, 0, "Fix the bug", "Fixed.", {promptSuggestions: {enabled: true, promptId: "client-1"}});
+        await vi.waitFor(() => expect(suggestionUpdates(client)).toEqual([
+            {version: 1, promptId: "client-1", suggestions: [{text: "Run the tests"}]},
+        ]));
+        await runTurn(client, protocolVersion, 1, "Thanks", "Done.", {promptSuggestions: {enabled: false, promptId: "client-2"}});
+        await settle();
+        expect(suggestionUpdates(client)).toEqual([
+            {version: 1, promptId: "client-1", suggestions: [{text: "Run the tests"}]},
+            {version: 1, promptId: "client-1", suggestions: []},
+        ]);
+        expect(sampleParams).toHaveLength(1);
+    });
+
     /** The prediction samples the session's own thread, anchored at the turn that just completed. */
     it("samples the session thread with the completed turn id", async () => {
         const client = await connect(true);

src/__tests__/PromptSuggestionGenerator.test.ts

@@ -1,5 +1,5 @@
 import {describe, expect, it, vi} from "vitest";
-import {PromptSuggestionGenerator, promptSuggestionsEnabled} from "../PromptSuggestionGenerator";
+import {PromptSuggestionGenerator, promptSuggestionsEnabled, promptSuggestionId, clientSupportsPromptSuggestions, promptSuggestionsMeta} from "../PromptSuggestionGenerator";
 import type {CodexAppServerClient} from "../CodexAppServerClient";
 import {deferred} from "./acp-test-utils";
 
@@ -170,3 +170,66 @@ describe("promptSuggestionsEnabled", () => {
         expect(promptSuggestionsEnabled({CODEX_ACP_PROMPT_SUGGESTIONS: " False "})).toBe(false);
     });
 });
+
+
+describe("independent prompt suggestions", () => {
+    it.each([
+        [undefined, false],
+        [null, false],
+        [{}, false],
+        [{_meta: {promptSuggestions: null}}, false],
+        [{_meta: {promptSuggestions: {version: 2}}}, false],
+        [{_meta: {promptSuggestions: {version: 1}}}, true],
+    ])("negotiates only the supported version: %j", (capabilities, expected) => {
+        expect(clientSupportsPromptSuggestions(capabilities)).toBe(expected);
+    });
+
+    it.each([
+        [undefined, undefined],
+        [{promptSuggestions: {enabled: false, promptId: "client-1"}}, undefined],
+        [{promptSuggestions: {enabled: true, promptId: " "}}, undefined],
+        [{promptSuggestions: {enabled: true, promptId: 1}}, undefined],
+        [{promptSuggestions: {enabled: true, promptId: "client-1"}}, "client-1"],
+    ])("reads explicit per-prompt consent and correlation: %j", (meta, expected) => {
+        expect(promptSuggestionId(meta)).toBe(expected);
+    });
+
+    it("does not sample without per-prompt opt-in", async () => {
+        const threadSample = vi.fn().mockResolvedValue(suggestionSample("Run the tests"));
+        const publish = vi.fn().mockResolvedValue(undefined);
+        const generator = new PromptSuggestionGenerator({threadSample} as unknown as CodexAppServerClient, publish, true);
+        await generator.beginPrompt();
+        generator.onTurnCompleted("thread", "turn");
+        await flush();
+        expect(threadSample).not.toHaveBeenCalled();
+        expect(publish).not.toHaveBeenCalled();
+    });
+
+    it("echoes the captured client id for predictions and clears, independently of vendor turn ids", async () => {
+        const threadSample = vi.fn().mockResolvedValue(suggestionSample("Run the tests"));
+        const publish = vi.fn().mockResolvedValue(undefined);
+        const generator = new PromptSuggestionGenerator({threadSample} as unknown as CodexAppServerClient, publish, true);
+        await generator.beginPrompt("client-1");
+        generator.onTurnCompleted("thread", "vendor-turn-1");
+        await vi.waitFor(() => expect(publish).toHaveBeenCalledWith("Run the tests", "client-1"));
+        await generator.beginPrompt("client-2");
+        expect(publish).toHaveBeenLastCalledWith(null, "client-1");
+        expect(promptSuggestionsMeta(null, "client-1")).toEqual({promptSuggestions: {version: 1, promptId: "client-1", suggestions: []}});
+    });
+
+    it("discards an older sample after the next client prompt opts out", async () => {
+        const result = deferred<SampleResult>();
+        const threadSample = vi.fn().mockReturnValue(result.promise);
+        const publish = vi.fn().mockResolvedValue(undefined);
+        const generator = new PromptSuggestionGenerator({threadSample} as unknown as CodexAppServerClient, publish, true);
+        await generator.beginPrompt("client-1");
+        generator.onTurnCompleted("thread", "vendor-turn-1");
+        await generator.beginPrompt();
+        result.resolve(suggestionSample("Too late"));
+        await flush();
+        generator.onTurnCompleted("thread", "vendor-turn-2");
+        await flush();
+        expect(publish).not.toHaveBeenCalled();
+        expect(threadSample).toHaveBeenCalledTimes(1);
+    });
+});