From 09ff8a47d5412a2fe43dde2ee0583a687b080285 Mon Sep 17 00:00:00 2001 From: Vladimir Krivosheev Date: Mon, 29 Jun 2026 11:43:59 +0200 Subject: [PATCH] IJPL-248623 expose Pi task folder workflow tools Route Pi task-folder operations through one taskFolderRequest control message with an operation name and structured arguments. This keeps the backend protocol extensible and lets the Pi extension register tools from a descriptor table instead of duplicating bridge code for each operation. Task folder metadata is now represented uniformly as string key/value data. The extension, prompt, and spec document conventional issue and review keys instead of adding dedicated shortcut parameters that can conflict with metadata. Marking an empty folder done now avoids initializing archive infrastructure; non-empty folders still archive assigned threads before changing folder status. GitOrigin-RevId: fc4a71386fab37c2ce4fb35ef35179525d338e3c --- .../providers/pi/sessions/BUILD.bazel | 2 + ...intellij.platform.ai.agent.pi.sessions.iml | 1 + .../pi-extension/agent-workbench-extension.ts | 4 +- .../resources/pi-extension/control.ts | 318 +++++++++++++----- .../sessions/src/PiExtensionControlBridge.kt | 20 +- .../src/PiExtensionControlProtocol.kt | 66 ++-- .../src/PiTaskFolderControlHandler.kt | 309 +++++++++++++---- .../PiExtensionControlWebSocketHandlerTest.kt | 101 ++++-- .../pi/sessions/testSrc/PiThemeSupportTest.kt | 15 +- ...SessionsTreePopupTaskFolderAgentActions.kt | 2 +- .../spec/sessions/agent-sessions-pi.spec.md | 12 +- 11 files changed, 629 insertions(+), 221 deletions(-) diff --git a/plugins/agent-workbench/lib-agent/providers/pi/sessions/BUILD.bazel b/plugins/agent-workbench/lib-agent/providers/pi/sessions/BUILD.bazel index d96dd9fa0be2..8c88bfd96d0d 100644 --- a/plugins/agent-workbench/lib-agent/providers/pi/sessions/BUILD.bazel +++ b/plugins/agent-workbench/lib-agent/providers/pi/sessions/BUILD.bazel @@ -39,6 +39,7 @@ jvm_library( "//plugins/agent-workbench/lib-agent/core", "//plugins/agent-workbench/lib-agent/json", "//plugins/agent-workbench/prompt/core", + "//plugins/agent-workbench/sessions", "//plugins/agent-workbench/lib-agent/sessions-core:ai-agent-sessions-core", "//libraries/jackson3/jackson:libraries-jackson3", "//libraries/netty/buffer", @@ -77,6 +78,7 @@ jvm_library( "//plugins/agent-workbench/lib-agent/core:core_test_lib", "//plugins/agent-workbench/lib-agent/json:json_test_lib", "//plugins/agent-workbench/prompt/core:core_test_lib", + "//plugins/agent-workbench/sessions:sessions_test_lib", "//plugins/agent-workbench/lib-agent/sessions-core:ai-agent-sessions-core_test_lib", "//libraries/jackson3/jackson:libraries-jackson3_test_lib", "//libraries/netty/buffer:buffer_test_lib", diff --git a/plugins/agent-workbench/lib-agent/providers/pi/sessions/intellij.platform.ai.agent.pi.sessions.iml b/plugins/agent-workbench/lib-agent/providers/pi/sessions/intellij.platform.ai.agent.pi.sessions.iml index 99c165a23ba6..4ed17cbebb4d 100644 --- a/plugins/agent-workbench/lib-agent/providers/pi/sessions/intellij.platform.ai.agent.pi.sessions.iml +++ b/plugins/agent-workbench/lib-agent/providers/pi/sessions/intellij.platform.ai.agent.pi.sessions.iml @@ -26,6 +26,7 @@ + diff --git a/plugins/agent-workbench/lib-agent/providers/pi/sessions/resources/pi-extension/agent-workbench-extension.ts b/plugins/agent-workbench/lib-agent/providers/pi/sessions/resources/pi-extension/agent-workbench-extension.ts index 57f42ad772c5..6b20cb7d62cd 100644 --- a/plugins/agent-workbench/lib-agent/providers/pi/sessions/resources/pi-extension/agent-workbench-extension.ts +++ b/plugins/agent-workbench/lib-agent/providers/pi/sessions/resources/pi-extension/agent-workbench-extension.ts @@ -1,5 +1,5 @@ import {type ExtensionAPI, type ExtensionContext} from "@earendil-works/pi-coding-agent"; -import {registerTaskFolderTool, startControlBridge} from "./control.ts"; +import {registerTaskFolderTools, startControlBridge} from "./control.ts"; import {registerJbCentralProvider} from "./jbcentral.ts"; import {parseModelCatalogMetadata} from "./modelCatalog.ts"; import {registerOmlxProviders} from "./omlx.ts"; @@ -33,7 +33,7 @@ export default async function agentWorkbenchTheme(pi: ExtensionAPI) { let lastSessionInfoSignature: string | undefined; let lastSessionInfoLeafId: string | null | undefined; - registerTaskFolderTool(pi, () => controlBridge); + registerTaskFolderTools(pi, () => controlBridge); const updateLastStatusSignature = (signature: string) => { lastStatusSignature = signature; diff --git a/plugins/agent-workbench/lib-agent/providers/pi/sessions/resources/pi-extension/control.ts b/plugins/agent-workbench/lib-agent/providers/pi/sessions/resources/pi-extension/control.ts index b9c2545fa021..20f86faea915 100644 --- a/plugins/agent-workbench/lib-agent/providers/pi/sessions/resources/pi-extension/control.ts +++ b/plugins/agent-workbench/lib-agent/providers/pi/sessions/resources/pi-extension/control.ts @@ -13,12 +13,7 @@ const STATUS_TOKEN = process.env[STATUS_TOKEN_ENV]; type AgentWorkbenchControlBridge = { setContext: (ctx: ExtensionContext) => void; - getCurrentTaskFolder: () => Promise; - listTaskFolderThreads: (folderId?: string) => Promise; - createAndAssignTaskFolder: (name: string, metadata?: Record) => Promise; - getTaskFolderMetadata: (folderId?: string) => Promise | undefined>; - setTaskFolderMetadata: (key: string, value: string, folderId?: string) => Promise; - deleteTaskFolderMetadata: (key: string, folderId?: string) => Promise; + requestTaskFolder: (operation: string, args?: Record) => Promise; close: () => void; }; @@ -28,21 +23,10 @@ type AgentWorkbenchControlMessageType = | "response" | "navigateTree" | "forkFromEntry" - | "getCurrentTaskFolder" - | "listTaskFolderThreads" - | "createAndAssignTaskFolder" - | "getTaskFolderMetadata" - | "setTaskFolderMetadata" - | "deleteTaskFolderMetadata"; + | "taskFolderRequest"; type AgentWorkbenchControlSessionMessageType = "hello" | "sessionState"; -type AgentWorkbenchControlRequestType = - | "getCurrentTaskFolder" - | "listTaskFolderThreads" - | "createAndAssignTaskFolder" - | "getTaskFolderMetadata" - | "setTaskFolderMetadata" - | "deleteTaskFolderMetadata"; +type AgentWorkbenchControlRequestType = "taskFolderRequest"; type AgentWorkbenchControlCommand = { type?: AgentWorkbenchControlMessageType; @@ -54,15 +38,9 @@ type AgentWorkbenchControlCommand = { ok?: boolean; cancelled?: boolean; error?: string; - folderId?: string; - name?: string; - key?: string; - value?: string; - changed?: boolean; - folder?: AgentWorkbenchTaskFolder | null; - threads?: AgentWorkbenchTaskFolderThread[]; - metadata?: Record; - assigned?: boolean; + operation?: string; + arguments?: Record; + result?: unknown; }; type AgentWorkbenchControlThread = { @@ -90,11 +68,6 @@ type AgentWorkbenchTaskFolderThread = { assignedAt?: number; }; -type AgentWorkbenchCreatedTaskFolder = { - folder: AgentWorkbenchTaskFolder; - assigned: boolean; -}; - type AgentWorkbenchControlContext = ExtensionContext & { navigateTree?: (entryId: string) => Promise | void; fork?: ( @@ -296,32 +269,9 @@ export function startControlBridge(ctx: ExtensionContext): AgentWorkbenchControl currentCtx = nextCtx; sendSessionMessage("sessionState"); }, - getCurrentTaskFolder: async () => { - const response = requireOk(await sendRequest("getCurrentTaskFolder")); - return response.folder ?? undefined; - }, - listTaskFolderThreads: async (folderId) => { - const response = requireOk(await sendRequest("listTaskFolderThreads", {folderId})); - return response.threads ?? []; - }, - createAndAssignTaskFolder: async (name, metadata) => { - const response = requireOk(await sendRequest("createAndAssignTaskFolder", {name, metadata})); - if (response.folder === undefined || response.folder === null) { - throw new Error("Agent Workbench did not return the created task folder"); - } - return {folder: response.folder, assigned: response.assigned === true}; - }, - getTaskFolderMetadata: async (folderId) => { - const response = requireOk(await sendRequest("getTaskFolderMetadata", {folderId})); - return response.metadata; - }, - setTaskFolderMetadata: async (key, value, folderId) => { - const response = requireOk(await sendRequest("setTaskFolderMetadata", {folderId, key, value})); - return response.changed === true; - }, - deleteTaskFolderMetadata: async (key, folderId) => { - const response = requireOk(await sendRequest("deleteTaskFolderMetadata", {folderId, key})); - return response.changed === true; + requestTaskFolder: async (operation: string, args: Record = {}): Promise => { + const response = requireOk(await sendRequest("taskFolderRequest", {operation, arguments: args})); + return response.result as T; }, close: () => { for (const [requestId, pending] of pendingRequests) { @@ -334,37 +284,253 @@ export function startControlBridge(ctx: ExtensionContext): AgentWorkbenchControl }; } -export function registerTaskFolderTool( +type AgentWorkbenchTaskFolderToolDefinition = { + name: string; + label: string; + description: string; + promptSnippet?: string; + promptGuidelines?: string[]; + parameters: ReturnType; + operation: string; + arguments?: (params: Record) => Record; + confirm?: (params: Record, ctx: ExtensionContext) => Promise; + resultText: (result: unknown) => string; +}; + +export function registerTaskFolderTools( pi: ExtensionAPI, bridgeProvider: () => AgentWorkbenchControlBridge | undefined, ): void { - pi.registerTool(defineTool({ + for (const definition of TASK_FOLDER_TOOL_DEFINITIONS) { + pi.registerTool(defineTool({ + name: definition.name, + label: definition.label, + description: definition.description, + promptSnippet: definition.promptSnippet, + promptGuidelines: definition.promptGuidelines, + parameters: definition.parameters, + async execute(_toolCallId, params, _signal, _onUpdate, ctx) { + const normalizedParams = params as Record; + const bridge = bridgeProvider(); + if (bridge === undefined) { + throw new Error("Agent Workbench control bridge is unavailable"); + } + if (definition.confirm !== undefined && !(await definition.confirm(normalizedParams, ctx))) { + return { + content: [{type: "text", text: "Cancelled task folder operation."}], + details: {cancelled: true}, + }; + } + const result = await bridge.requestTaskFolder(definition.operation, definition.arguments?.(normalizedParams) ?? normalizedParams); + return { + content: [{type: "text", text: definition.resultText(result)}], + details: result, + }; + }, + })); + } +} + +const TASK_FOLDER_METADATA_DESCRIPTION = "String metadata. Conventional keys are 'issue' and 'review'; custom keys are allowed."; + +const TASK_FOLDER_TOOL_DEFINITIONS: AgentWorkbenchTaskFolderToolDefinition[] = [ + { + name: "agent_workbench_get_current_task_folder", + label: "Get Current Task Folder", + description: "Get the Agent Workbench task folder assigned to the current Pi session, if any.", + promptSnippet: "Inspect the task folder assigned to the current Pi thread", + promptGuidelines: [ + "Use agent_workbench_get_current_task_folder before updating a task folder when you are unsure whether this Pi thread is already assigned.", + ], + parameters: Type.Object({}), + operation: "getCurrent", + resultText: (result) => { + const folder = resultFolder(result); + return folder === undefined ? "No task folder is assigned to this thread." : `Current task folder: '${folder.name}'.`; + }, + }, + { + name: "agent_workbench_list_task_folders", + label: "List Task Folders", + description: "List Agent Workbench task folders for the current project.", + promptSnippet: "List Agent Workbench task folders for the current project", + promptGuidelines: [ + "Use agent_workbench_list_task_folders when you need a task folder id or need to choose an existing task folder.", + ], + parameters: Type.Object({ + includeDone: Type.Optional(Type.Boolean({description: "Whether to include done task folders"})), + }), + operation: "listFolders", + resultText: (result) => `Found ${resultFolders(result).length} task folder(s).`, + }, + { + name: "agent_workbench_list_task_folder_threads", + label: "List Task Folder Threads", + description: "List threads assigned to an Agent Workbench task folder. Defaults to the current task folder.", + promptSnippet: "List threads assigned to an Agent Workbench task folder", + parameters: Type.Object({ + folderId: Type.Optional(Type.String({description: "Task folder id; defaults to the current task folder"})), + }), + operation: "listThreads", + resultText: (result) => `Found ${resultThreads(result).length} assigned thread(s).`, + }, + { name: "agent_workbench_create_task_folder", label: "Create Task Folder", description: "Create an Agent Workbench task folder for the current project and assign the current Pi session to it.", promptSnippet: "Create an Agent Workbench task folder and assign this Pi thread to it", promptGuidelines: [ - "Use agent_workbench_create_task_folder when the user asks to create a task folder or start work in a task folder.", - "Pass the issue id as issue when the user mentions an issue tracker id.", + "Use agent_workbench_create_task_folder when the user asks to create a task folder or start work in a new task folder.", + "Use metadata key 'issue' for issue tracker ids and 'review' for review ids; do not use separate issue parameters.", + "If the current Pi thread already has a task folder, update it with metadata or rename tools instead of creating another folder.", ], parameters: Type.Object({ name: Type.String({description: "Task folder name"}), - issue: Type.Optional(Type.String({description: "Issue tracker id to store as task folder metadata key 'issue'"})), + metadata: Type.Optional(Type.Record(Type.String(), Type.String({description: TASK_FOLDER_METADATA_DESCRIPTION}))), }), - async execute(_toolCallId, params, _signal, _onUpdate, _ctx) { - const bridge = bridgeProvider(); - if (bridge === undefined) { - throw new Error("Agent Workbench control bridge is unavailable"); + operation: "createAndAssign", + resultText: (result) => { + const folder = resultFolder(result); + if (folder === undefined) { + return "Task folder request completed."; } - const issue = params.issue?.trim(); - const metadata = issue === undefined || issue.length === 0 ? undefined : {issue}; - const result = await bridge.createAndAssignTaskFolder(params.name, metadata); - return { - content: [{type: "text", text: `Created task folder '${result.folder.name}' and assigned this thread.`}], - details: result, - }; + return resultBoolean(result, "created") === false + ? `This thread is already assigned to task folder '${folder.name}'. Use metadata or rename tools for updates.` + : `Created task folder '${folder.name}' and assigned this thread.`; }, - })); + }, + { + name: "agent_workbench_assign_current_thread_to_task_folder", + label: "Assign Current Thread", + description: "Assign the current Pi session to an existing Agent Workbench task folder in the current project.", + promptSnippet: "Assign this Pi thread to an existing task folder", + parameters: Type.Object({ + folderId: Type.String({description: "Task folder id"}), + }), + operation: "assignCurrentThread", + resultText: mutationResultText("Assigned current thread to task folder."), + }, + { + name: "agent_workbench_remove_current_thread_from_task_folder", + label: "Remove Current Thread", + description: "Remove the current Pi session from its Agent Workbench task folder.", + promptSnippet: "Remove this Pi thread from its current task folder", + parameters: Type.Object({}), + operation: "unassignCurrentThread", + resultText: mutationResultText("Removed current thread from task folder."), + }, + { + name: "agent_workbench_rename_task_folder", + label: "Rename Task Folder", + description: "Rename an Agent Workbench task folder. Defaults to the current task folder.", + promptSnippet: "Rename an Agent Workbench task folder", + parameters: Type.Object({ + folderId: Type.Optional(Type.String({description: "Task folder id; defaults to the current task folder"})), + name: Type.String({description: "New task folder name"}), + }), + operation: "rename", + resultText: (result) => { + const folder = resultFolder(result); + return folder === undefined ? "Task folder rename completed." : `Renamed task folder to '${folder.name}'.`; + }, + }, + { + name: "agent_workbench_set_task_folder_metadata", + label: "Set Task Folder Metadata", + description: "Set a string metadata key on an Agent Workbench task folder. Defaults to the current task folder.", + promptSnippet: "Set task folder metadata such as issue or review", + promptGuidelines: [ + "Use agent_workbench_set_task_folder_metadata to associate an issue id with an existing task folder using key 'issue'.", + "Use agent_workbench_set_task_folder_metadata to associate a review id with an existing task folder using key 'review'.", + ], + parameters: Type.Object({ + folderId: Type.Optional(Type.String({description: "Task folder id; defaults to the current task folder"})), + key: Type.String({description: "Metadata key. Conventional keys are 'issue' and 'review'."}), + value: Type.String({description: "Metadata value"}), + }), + operation: "setMetadata", + resultText: mutationResultText("Updated task folder metadata."), + }, + { + name: "agent_workbench_delete_task_folder_metadata", + label: "Delete Task Folder Metadata", + description: "Delete a metadata key from an Agent Workbench task folder. Defaults to the current task folder.", + promptSnippet: "Delete task folder metadata", + parameters: Type.Object({ + folderId: Type.Optional(Type.String({description: "Task folder id; defaults to the current task folder"})), + key: Type.String({description: "Metadata key to delete"}), + }), + operation: "deleteMetadata", + resultText: mutationResultText("Deleted task folder metadata."), + }, + { + name: "agent_workbench_mark_task_folder_done", + label: "Mark Task Folder Done", + description: "Mark an Agent Workbench task folder done after archiving assigned threads. Defaults to the current task folder.", + promptSnippet: "Mark a task folder done and archive assigned threads", + parameters: Type.Object({ + folderId: Type.Optional(Type.String({description: "Task folder id; defaults to the current task folder"})), + }), + operation: "markDone", + confirm: (_params, ctx) => confirmTaskFolderOperation(ctx, "Mark Task Folder Done", "Archive assigned threads and mark this task folder done?"), + resultText: (result) => { + const archived = resultNumber(result, "archivedCount"); + const requested = resultNumber(result, "requestedCount"); + return `Marked task folder done. Archived ${archived} of ${requested} assigned thread(s).`; + }, + }, + { + name: "agent_workbench_delete_task_folder", + label: "Delete Task Folder", + description: "Delete an Agent Workbench task folder and remove its assignments. Threads are not archived.", + parameters: Type.Object({ + folderId: Type.Optional(Type.String({description: "Task folder id; defaults to the current task folder"})), + }), + operation: "delete", + confirm: (_params, ctx) => confirmTaskFolderOperation(ctx, "Delete Task Folder", "Delete this task folder and remove its thread assignments? Threads will not be archived."), + resultText: mutationResultText("Deleted task folder."), + }, +]; + +async function confirmTaskFolderOperation(ctx: ExtensionContext, title: string, message: string): Promise { + if (!ctx.hasUI) { + throw new Error("Task folder operation requires interactive confirmation"); + } + return ctx.ui.confirm(title, message); +} + +function mutationResultText(changedText: string): (result: unknown) => string { + return (result) => resultBoolean(result, "changed") ? changedText : "Task folder was already up to date."; +} + +function resultFolder(result: unknown): AgentWorkbenchTaskFolder | undefined { + if (!isRecord(result)) return undefined; + const folder = result.folder; + return isRecord(folder) && typeof folder.name === "string" ? folder as AgentWorkbenchTaskFolder : undefined; +} + +function resultFolders(result: unknown): AgentWorkbenchTaskFolder[] { + if (!isRecord(result) || !Array.isArray(result.folders)) return []; + return result.folders.filter(isRecord) as AgentWorkbenchTaskFolder[]; +} + +function resultThreads(result: unknown): AgentWorkbenchTaskFolderThread[] { + if (!isRecord(result) || !Array.isArray(result.threads)) return []; + return result.threads.filter(isRecord) as AgentWorkbenchTaskFolderThread[]; +} + +function resultBoolean(result: unknown, key: string): boolean | undefined { + if (!isRecord(result)) return undefined; + return typeof result[key] === "boolean" ? result[key] : undefined; +} + +function resultNumber(result: unknown, key: string): number { + if (!isRecord(result)) return 0; + return typeof result[key] === "number" ? result[key] : 0; +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null; } function resolveCapabilities(ctx: ExtensionContext): { navigateTree: boolean; fork: boolean } { diff --git a/plugins/agent-workbench/lib-agent/providers/pi/sessions/src/PiExtensionControlBridge.kt b/plugins/agent-workbench/lib-agent/providers/pi/sessions/src/PiExtensionControlBridge.kt index ce7c971c5418..efbdc9cea5a6 100644 --- a/plugins/agent-workbench/lib-agent/providers/pi/sessions/src/PiExtensionControlBridge.kt +++ b/plugins/agent-workbench/lib-agent/providers/pi/sessions/src/PiExtensionControlBridge.kt @@ -87,13 +87,7 @@ internal object PiExtensionControlBridge { PiControlMessageType.HELLO -> handleHello(webSocketClient, payload) PiControlMessageType.SESSION_STATE -> handleSessionState(webSocketClient, payload) PiControlMessageType.RESPONSE -> handleResponse(webSocketClient, payload) - PiControlMessageType.GET_CURRENT_TASK_FOLDER, - PiControlMessageType.LIST_TASK_FOLDER_THREADS, - PiControlMessageType.CREATE_AND_ASSIGN_TASK_FOLDER, - PiControlMessageType.GET_TASK_FOLDER_METADATA, - PiControlMessageType.SET_TASK_FOLDER_METADATA, - PiControlMessageType.DELETE_TASK_FOLDER_METADATA, - -> handleTaskFolderRequest(webSocketClient, payload) + PiControlMessageType.TASK_FOLDER_REQUEST -> handleTaskFolderRequest(webSocketClient, payload) PiControlMessageType.NAVIGATE_TREE, PiControlMessageType.FORK_FROM_ENTRY, null, @@ -185,13 +179,11 @@ internal object PiExtensionControlBridge { return } - sendControlText( - client = client, - text = taskFolderControlHandler.handle( - context = PiControlSessionContext(projectPath = connection.projectPath, sessionId = connection.sessionId), - payload = payload, - requestId = requestId, - ), + taskFolderControlHandler.handle( + context = PiControlSessionContext(projectPath = connection.projectPath, sessionId = connection.sessionId), + payload = payload, + requestId = requestId, + sendResponse = { response -> sendControlText(client = client, text = response) }, ) } diff --git a/plugins/agent-workbench/lib-agent/providers/pi/sessions/src/PiExtensionControlProtocol.kt b/plugins/agent-workbench/lib-agent/providers/pi/sessions/src/PiExtensionControlProtocol.kt index 9f150619048e..dd4ee1796923 100644 --- a/plugins/agent-workbench/lib-agent/providers/pi/sessions/src/PiExtensionControlProtocol.kt +++ b/plugins/agent-workbench/lib-agent/providers/pi/sessions/src/PiExtensionControlProtocol.kt @@ -19,12 +19,7 @@ internal enum class PiControlMessageType(@JvmField val wireName: String) { RESPONSE("response"), NAVIGATE_TREE("navigateTree"), FORK_FROM_ENTRY("forkFromEntry"), - GET_CURRENT_TASK_FOLDER("getCurrentTaskFolder"), - LIST_TASK_FOLDER_THREADS("listTaskFolderThreads"), - CREATE_AND_ASSIGN_TASK_FOLDER("createAndAssignTaskFolder"), - GET_TASK_FOLDER_METADATA("getTaskFolderMetadata"), - SET_TASK_FOLDER_METADATA("setTaskFolderMetadata"), - DELETE_TASK_FOLDER_METADATA("deleteTaskFolderMetadata"); + TASK_FOLDER_REQUEST("taskFolderRequest"); companion object { private val byWireName: Map = entries.associateBy(PiControlMessageType::wireName) @@ -69,10 +64,16 @@ internal data class PiControlPayload( @JvmField val error: String? = null, @JvmField val thread: PiControlThreadPayload? = null, @JvmField val capabilities: PiControlCapabilities? = null, + @JvmField val operation: String? = null, + @JvmField val arguments: PiTaskFolderControlArguments? = null, +) + +internal data class PiTaskFolderControlArguments( @JvmField val folderId: String? = null, @JvmField val name: String? = null, @JvmField val key: String? = null, @JvmField val value: String? = null, + @JvmField val includeDone: Boolean? = null, @JvmField val metadata: Map? = null, ) @@ -128,15 +129,6 @@ internal fun buildPiControlHelloAcknowledgement(requestId: String?, sessionId: S } } -internal fun buildPiControlMutationResponse(requestId: String, changed: Boolean): String { - return buildPiControlJsonObject { generator -> - generator.writeStringProperty("type", PiControlMessageType.RESPONSE.wireName) - generator.writeStringProperty("requestId", requestId) - generator.writeBooleanProperty("ok", true) - generator.writeBooleanProperty("changed", changed) - } -} - internal fun buildPiControlErrorResponse(requestId: String?, error: String): String { return buildPiControlJsonObject { generator -> generator.writeStringProperty("type", PiControlMessageType.RESPONSE.wireName) @@ -167,11 +159,8 @@ private fun readControlPayload(parser: JsonParser): PiControlPayload { var error: String? = null var thread: PiControlThreadPayload? = null var capabilities: PiControlCapabilities? = null - var folderId: String? = null - var name: String? = null - var key: String? = null - var value: String? = null - var metadata: Map? = null + var operation: String? = null + var arguments: PiTaskFolderControlArguments? = null forEachJsonObjectField(parser) { fieldName -> when (fieldName) { "type" -> type = PiControlMessageType.fromWireName(readJsonStringOrNull(parser)) @@ -184,11 +173,8 @@ private fun readControlPayload(parser: JsonParser): PiControlPayload { "error" -> error = readJsonStringOrNull(parser) "thread" -> thread = readControlThreadPayload(parser) "capabilities" -> capabilities = readControlCapabilities(parser) - "folderId" -> folderId = readJsonStringOrNull(parser) - "name" -> name = readJsonStringOrNull(parser) - "key" -> key = readJsonStringOrNull(parser) - "value" -> value = readJsonStringOrNull(parser) - "metadata" -> metadata = readControlStringMap(parser) + "operation" -> operation = readJsonStringOrNull(parser) + "arguments" -> arguments = readTaskFolderControlArguments(parser) else -> parser.skipChildren() } true @@ -204,10 +190,40 @@ private fun readControlPayload(parser: JsonParser): PiControlPayload { error = error, thread = thread, capabilities = capabilities, + operation = operation, + arguments = arguments, + ) +} + +private fun readTaskFolderControlArguments(parser: JsonParser): PiTaskFolderControlArguments? { + if (parser.currentToken() != JsonToken.START_OBJECT) { + parser.skipChildren() + return null + } + var folderId: String? = null + var name: String? = null + var key: String? = null + var value: String? = null + var includeDone: Boolean? = null + var metadata: Map? = null + forEachJsonObjectField(parser) { fieldName -> + when (fieldName) { + "folderId" -> folderId = readJsonStringOrNull(parser) + "name" -> name = readJsonStringOrNull(parser) + "key" -> key = readJsonStringOrNull(parser) + "value" -> value = readJsonStringOrNull(parser) + "includeDone" -> includeDone = readJsonBooleanOrNull(parser) + "metadata" -> metadata = readControlStringMap(parser) + else -> parser.skipChildren() + } + true + } + return PiTaskFolderControlArguments( folderId = folderId, name = name, key = key, value = value, + includeDone = includeDone, metadata = metadata, ) } diff --git a/plugins/agent-workbench/lib-agent/providers/pi/sessions/src/PiTaskFolderControlHandler.kt b/plugins/agent-workbench/lib-agent/providers/pi/sessions/src/PiTaskFolderControlHandler.kt index fd15c98356b4..4652e6fd0b83 100644 --- a/plugins/agent-workbench/lib-agent/providers/pi/sessions/src/PiTaskFolderControlHandler.kt +++ b/plugins/agent-workbench/lib-agent/providers/pi/sessions/src/PiTaskFolderControlHandler.kt @@ -1,110 +1,235 @@ // Copyright 2000-2026 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. package com.intellij.platform.ai.agent.pi.sessions +import com.intellij.agent.workbench.sessions.model.ArchiveThreadTarget +import com.intellij.agent.workbench.sessions.service.AgentSessionArchiveService +import com.intellij.agent.workbench.sessions.statistics.AgentWorkbenchEntryPoint import com.intellij.openapi.components.service import com.intellij.platform.ai.agent.sessions.core.folders.AgentTaskFolder import com.intellij.platform.ai.agent.sessions.core.folders.AgentTaskFolderService +import com.intellij.platform.ai.agent.sessions.core.folders.AgentTaskFolderStatus import com.intellij.platform.ai.agent.sessions.core.folders.AgentTaskFolderThreadAssignment import tools.jackson.core.JsonGenerator internal class PiTaskFolderControlHandler( private val taskFolderServiceProvider: () -> AgentTaskFolderService = { service() }, + private val archiveServiceProvider: () -> AgentSessionArchiveService = { service() }, ) { - fun handle(context: PiControlSessionContext, payload: PiControlPayload, requestId: String): String { + fun handle( + context: PiControlSessionContext, + payload: PiControlPayload, + requestId: String, + sendResponse: (String) -> Unit, + ) { val service = taskFolderServiceProvider() - return when (payload.type) { - PiControlMessageType.GET_CURRENT_TASK_FOLDER -> { + val arguments = payload.arguments ?: PiTaskFolderControlArguments() + when (payload.operation?.trim()) { + OP_GET_CURRENT -> { val folder = service.getFolderForThread(context.projectPath, PI_AGENT_SESSION_PROVIDER, context.sessionId) - buildTaskFolderResponse(requestId = requestId, folder = folder) + sendResponse(buildTaskFolderResponse(requestId = requestId, folder = folder)) } - PiControlMessageType.LIST_TASK_FOLDER_THREADS -> { - val folderId = resolveTaskFolderId(service, context, payload) - if (folderId == null) { - buildPiControlErrorResponse(requestId, "Task folder is not available") - } - else { - buildTaskFolderAssignmentsResponse(requestId, service.listFolderThreadAssignments(folderId)) - } + OP_LIST_FOLDERS -> { + val folders = service.listFolders(context.projectPath, includeDone = arguments.includeDone == true) + sendResponse(buildTaskFoldersResponse(requestId = requestId, folders = folders)) } - PiControlMessageType.CREATE_AND_ASSIGN_TASK_FOLDER -> { - val name = payload.name?.trim()?.takeIf { it.isNotEmpty() } - if (name == null) { - buildPiControlErrorResponse(requestId, "Task folder name is required") - } - else { - val folder = service.createFolder(context.projectPath, name, payload.metadata.orEmpty()) - if (folder == null) { - buildPiControlErrorResponse(requestId, "Task folder could not be created") - } - else { - val assigned = service.assignThread(context.projectPath, PI_AGENT_SESSION_PROVIDER, context.sessionId, folder.id) - buildTaskFolderCreatedResponse(requestId = requestId, folder = folder, assigned = assigned) - } - } - } - PiControlMessageType.GET_TASK_FOLDER_METADATA -> { - val folder = resolveTaskFolder(service, context, payload) + OP_LIST_THREADS -> { + val folder = resolveTaskFolder(service, context, arguments) if (folder == null) { - buildPiControlErrorResponse(requestId, "Task folder is not available") + sendResponse(buildPiControlErrorResponse(requestId, "Task folder is not available")) } else { - buildTaskFolderMetadataResponse(requestId, folder.metadata) + sendResponse(buildTaskFolderAssignmentsResponse(requestId, service.listFolderThreadAssignments(folder.id))) } } - PiControlMessageType.SET_TASK_FOLDER_METADATA -> { - val folderId = resolveTaskFolderId(service, context, payload) - val key = payload.key?.trim()?.takeIf { it.isNotEmpty() } - val value = payload.value - if (folderId == null || key == null || value == null) { - buildPiControlErrorResponse(requestId, "Task folder metadata request is incomplete") + OP_CREATE_AND_ASSIGN -> { + createAndAssignCurrentThread(service, context, arguments, requestId, sendResponse) + } + OP_ASSIGN_CURRENT_THREAD -> { + val folder = resolveTaskFolder(service, context, arguments) + if (folder == null) { + sendResponse(buildPiControlErrorResponse(requestId, "Task folder is not available")) } else { - val changed = service.setMetadata(folderId, key, value) - buildPiControlMutationResponse(requestId, changed) + val changed = service.assignThread(context.projectPath, PI_AGENT_SESSION_PROVIDER, context.sessionId, folder.id) + sendResponse(buildTaskFolderMutationResponse(requestId = requestId, + changed = changed, + folder = refreshedFolder(service, context, folder.id))) } } - PiControlMessageType.DELETE_TASK_FOLDER_METADATA -> { - val folderId = resolveTaskFolderId(service, context, payload) - val key = payload.key?.trim()?.takeIf { it.isNotEmpty() } - if (folderId == null || key == null) { - buildPiControlErrorResponse(requestId, "Task folder metadata request is incomplete") + OP_UNASSIGN_CURRENT_THREAD -> { + val folder = service.getFolderForThread(context.projectPath, PI_AGENT_SESSION_PROVIDER, context.sessionId) + val changed = service.unassignThread(context.projectPath, PI_AGENT_SESSION_PROVIDER, context.sessionId) + sendResponse(buildTaskFolderMutationResponse(requestId = requestId, changed = changed, folder = folder)) + } + OP_RENAME -> { + val folder = resolveTaskFolder(service, context, arguments) + val name = arguments.name?.trim()?.takeIf { it.isNotEmpty() } + if (folder == null || name == null) { + sendResponse(buildPiControlErrorResponse(requestId, "Task folder rename request is incomplete")) } else { - val changed = service.deleteMetadata(folderId, key) - buildPiControlMutationResponse(requestId, changed) + val changed = service.renameFolder(folder.id, name) + sendResponse(buildTaskFolderMutationResponse(requestId = requestId, + changed = changed, + folder = refreshedFolder(service, context, folder.id))) } } - else -> buildPiControlErrorResponse(requestId, "Unsupported task folder request") + OP_SET_METADATA -> { + val folder = resolveTaskFolder(service, context, arguments) + val key = arguments.key?.trim()?.takeIf { it.isNotEmpty() } + val value = arguments.value + if (folder == null || key == null || value == null) { + sendResponse(buildPiControlErrorResponse(requestId, "Task folder metadata request is incomplete")) + } + else { + val changed = service.setMetadata(folder.id, key, value) + sendResponse(buildTaskFolderMutationResponse(requestId = requestId, + changed = changed, + folder = refreshedFolder(service, context, folder.id))) + } + } + OP_DELETE_METADATA -> { + val folder = resolveTaskFolder(service, context, arguments) + val key = arguments.key?.trim()?.takeIf { it.isNotEmpty() } + if (folder == null || key == null) { + sendResponse(buildPiControlErrorResponse(requestId, "Task folder metadata request is incomplete")) + } + else { + val changed = service.deleteMetadata(folder.id, key) + sendResponse(buildTaskFolderMutationResponse(requestId = requestId, + changed = changed, + folder = refreshedFolder(service, context, folder.id))) + } + } + OP_MARK_DONE -> { + val folder = resolveTaskFolder(service, context, arguments) + if (folder == null) { + sendResponse(buildPiControlErrorResponse(requestId, "Task folder is not available")) + } + else { + markDone(service, archiveServiceProvider, context, folder, requestId, sendResponse) + } + } + OP_DELETE -> { + val folder = resolveTaskFolder(service, context, arguments) + if (folder == null) { + sendResponse(buildPiControlErrorResponse(requestId, "Task folder is not available")) + } + else { + val changed = service.deleteFolder(folder.id) + sendResponse(buildTaskFolderMutationResponse(requestId = requestId, changed = changed, folder = folder)) + } + } + else -> sendResponse(buildPiControlErrorResponse(requestId, "Unsupported task folder operation")) } } } -private fun buildTaskFolderCreatedResponse(requestId: String, folder: AgentTaskFolder, assigned: Boolean): String { - return buildPiControlJsonObject { generator -> - generator.writeStringProperty("type", PiControlMessageType.RESPONSE.wireName) - generator.writeStringProperty("requestId", requestId) - generator.writeBooleanProperty("ok", true) +private fun createAndAssignCurrentThread( + service: AgentTaskFolderService, + context: PiControlSessionContext, + arguments: PiTaskFolderControlArguments, + requestId: String, + sendResponse: (String) -> Unit, +) { + val existing = service.getFolderForThread(context.projectPath, PI_AGENT_SESSION_PROVIDER, context.sessionId) + if (existing != null) { + sendResponse(buildTaskFolderCreatedResponse(requestId = requestId, folder = existing, created = false, assigned = true)) + return + } + + val name = arguments.name?.trim()?.takeIf { it.isNotEmpty() } + if (name == null) { + sendResponse(buildPiControlErrorResponse(requestId, "Task folder name is required")) + return + } + val folder = service.createFolder(context.projectPath, name, arguments.metadata.orEmpty()) + if (folder == null) { + sendResponse(buildPiControlErrorResponse(requestId, "Task folder could not be created")) + return + } + val assigned = service.assignThread(context.projectPath, PI_AGENT_SESSION_PROVIDER, context.sessionId, folder.id) + sendResponse(buildTaskFolderCreatedResponse(requestId = requestId, folder = folder, created = true, assigned = assigned)) +} + +private fun markDone( + service: AgentTaskFolderService, + archiveServiceProvider: () -> AgentSessionArchiveService, + context: PiControlSessionContext, + folder: AgentTaskFolder, + requestId: String, + sendResponse: (String) -> Unit, +) { + if (folder.status == AgentTaskFolderStatus.DONE) { + sendResponse(buildTaskFolderDoneResponse(requestId = requestId, + changed = false, + folder = folder, + requestedCount = 0, + archivedCount = 0)) + return + } + val targets = service.listFolderThreadAssignments(folder.id) + .map { assignment -> ArchiveThreadTarget.Thread(assignment.path, assignment.provider, assignment.threadId) } + .distinctBy { target -> ArchiveTargetKey(target.path, target.provider.value, target.threadId) } + if (targets.isEmpty()) { + val changed = service.setFolderStatus(folder.id, AgentTaskFolderStatus.DONE) + sendResponse( + buildTaskFolderDoneResponse( + requestId = requestId, + changed = changed, + folder = refreshedFolder(service, context, folder.id), + requestedCount = 0, + archivedCount = 0, + ) + ) + return + } + val archiveService = archiveServiceProvider() + if (!targets.all { target -> archiveService.canArchiveProvider(target.provider) }) { + sendResponse(buildPiControlErrorResponse(requestId, "Task folder contains threads that cannot be archived")) + return + } + archiveService.archiveThreads(targets, AgentWorkbenchEntryPoint.TREE_POPUP) { result -> + val changed = result.allRequestedArchived && service.setFolderStatus(folder.id, AgentTaskFolderStatus.DONE) + sendResponse( + buildTaskFolderDoneResponse( + requestId = requestId, + changed = changed, + folder = refreshedFolder(service, context, folder.id) ?: folder, + requestedCount = result.requestedCount, + archivedCount = result.archivedCount, + ) + ) + } +} + +private fun buildTaskFolderCreatedResponse(requestId: String, folder: AgentTaskFolder, created: Boolean, assigned: Boolean): String { + return buildTaskFolderResultResponse(requestId) { generator -> generator.writeName("folder") writeTaskFolder(generator, folder) + generator.writeBooleanProperty("created", created) generator.writeBooleanProperty("assigned", assigned) } } private fun buildTaskFolderResponse(requestId: String, folder: AgentTaskFolder?): String { - return buildPiControlJsonObject { generator -> - generator.writeStringProperty("type", PiControlMessageType.RESPONSE.wireName) - generator.writeStringProperty("requestId", requestId) - generator.writeBooleanProperty("ok", true) + return buildTaskFolderResultResponse(requestId) { generator -> generator.writeName("folder") writeTaskFolder(generator, folder) } } +private fun buildTaskFoldersResponse(requestId: String, folders: List): String { + return buildTaskFolderResultResponse(requestId) { generator -> + generator.writeName("folders") + generator.writeStartArray() + folders.forEach { folder -> writeTaskFolder(generator, folder) } + generator.writeEndArray() + } +} + private fun buildTaskFolderAssignmentsResponse(requestId: String, assignments: List): String { - return buildPiControlJsonObject { generator -> - generator.writeStringProperty("type", PiControlMessageType.RESPONSE.wireName) - generator.writeStringProperty("requestId", requestId) - generator.writeBooleanProperty("ok", true) + return buildTaskFolderResultResponse(requestId) { generator -> generator.writeName("threads") generator.writeStartArray() assignments.forEach { assignment -> writeTaskFolderAssignment(generator, assignment) } @@ -112,14 +237,38 @@ private fun buildTaskFolderAssignmentsResponse(requestId: String, assignments: L } } -private fun buildTaskFolderMetadataResponse(requestId: String, metadata: Map): String { +private fun buildTaskFolderMutationResponse(requestId: String, changed: Boolean, folder: AgentTaskFolder?): String { + return buildTaskFolderResultResponse(requestId) { generator -> + generator.writeBooleanProperty("changed", changed) + generator.writeName("folder") + writeTaskFolder(generator, folder) + } +} + +private fun buildTaskFolderDoneResponse( + requestId: String, + changed: Boolean, + folder: AgentTaskFolder?, + requestedCount: Int, + archivedCount: Int, +): String { + return buildTaskFolderResultResponse(requestId) { generator -> + generator.writeBooleanProperty("changed", changed) + generator.writeNumberProperty("requestedCount", requestedCount) + generator.writeNumberProperty("archivedCount", archivedCount) + generator.writeName("folder") + writeTaskFolder(generator, folder) + } +} + +private fun buildTaskFolderResultResponse(requestId: String, writeResult: (JsonGenerator) -> Unit): String { return buildPiControlJsonObject { generator -> generator.writeStringProperty("type", PiControlMessageType.RESPONSE.wireName) generator.writeStringProperty("requestId", requestId) generator.writeBooleanProperty("ok", true) - generator.writeName("metadata") + generator.writeName("result") generator.writeStartObject() - metadata.forEach { (key, value) -> generator.writeStringProperty(key, value) } + writeResult(generator) generator.writeEndObject() } } @@ -156,13 +305,33 @@ private fun writeTaskFolderAssignment(generator: JsonGenerator, assignment: Agen private fun resolveTaskFolder( service: AgentTaskFolderService, context: PiControlSessionContext, - payload: PiControlPayload, + arguments: PiTaskFolderControlArguments, ): AgentTaskFolder? { - val folderId = resolveTaskFolderId(service, context, payload) ?: return null - return service.getFolder(folderId) + val explicitFolderId = arguments.folderId?.trim()?.takeIf { it.isNotEmpty() } + if (explicitFolderId != null) { + return service.snapshot(includeDone = true).folder(context.projectPath, explicitFolderId) + } + return service.getFolderForThread(context.projectPath, PI_AGENT_SESSION_PROVIDER, context.sessionId) } -private fun resolveTaskFolderId(service: AgentTaskFolderService, context: PiControlSessionContext, payload: PiControlPayload): String? { - return payload.folderId?.trim()?.takeIf { it.isNotEmpty() } - ?: service.getFolderForThread(context.projectPath, PI_AGENT_SESSION_PROVIDER, context.sessionId)?.id +private fun refreshedFolder(service: AgentTaskFolderService, context: PiControlSessionContext, folderId: String): AgentTaskFolder? { + return service.snapshot(includeDone = true).folder(context.projectPath, folderId) } + +private data class ArchiveTargetKey( + @JvmField val path: String, + @JvmField val providerId: String, + @JvmField val threadId: String, +) + +private const val OP_GET_CURRENT: String = "getCurrent" +private const val OP_LIST_FOLDERS: String = "listFolders" +private const val OP_LIST_THREADS: String = "listThreads" +private const val OP_CREATE_AND_ASSIGN: String = "createAndAssign" +private const val OP_ASSIGN_CURRENT_THREAD: String = "assignCurrentThread" +private const val OP_UNASSIGN_CURRENT_THREAD: String = "unassignCurrentThread" +private const val OP_RENAME: String = "rename" +private const val OP_SET_METADATA: String = "setMetadata" +private const val OP_DELETE_METADATA: String = "deleteMetadata" +private const val OP_MARK_DONE: String = "markDone" +private const val OP_DELETE: String = "delete" diff --git a/plugins/agent-workbench/lib-agent/providers/pi/sessions/testSrc/PiExtensionControlWebSocketHandlerTest.kt b/plugins/agent-workbench/lib-agent/providers/pi/sessions/testSrc/PiExtensionControlWebSocketHandlerTest.kt index fd831ceebe92..fd93308eaf43 100644 --- a/plugins/agent-workbench/lib-agent/providers/pi/sessions/testSrc/PiExtensionControlWebSocketHandlerTest.kt +++ b/plugins/agent-workbench/lib-agent/providers/pi/sessions/testSrc/PiExtensionControlWebSocketHandlerTest.kt @@ -6,6 +6,7 @@ import com.intellij.platform.ai.agent.json.createJsonParser import com.intellij.platform.ai.agent.json.forEachJsonObjectField import com.intellij.platform.ai.agent.json.readJsonStringOrNull import com.intellij.platform.ai.agent.sessions.core.folders.AgentTaskFolderService +import com.intellij.platform.ai.agent.sessions.core.folders.AgentTaskFolderStatus import com.intellij.platform.ai.agent.sessions.core.providers.AgentSessionSourceUpdate import com.intellij.openapi.Disposable import com.intellij.openapi.components.service @@ -210,42 +211,60 @@ class PiExtensionControlWebSocketHandlerTest { ).join() listener.nextMessage() - webSocket.sendText(controlTaskFolderRequest("request-current", "getCurrentTaskFolder"), true).join() + webSocket.sendText(controlTaskFolderRequest("request-current", "getCurrent"), true).join() val currentFolderResponse = listener.nextMessage() assertThat(currentFolderResponse).contains("\"requestId\":\"request-current\"") assertThat(currentFolderResponse).contains("\"ok\":true") + assertThat(currentFolderResponse).contains("\"result\":") assertThat(currentFolderResponse).contains("\"id\":${folder.id.jsonString()}") assertThat(currentFolderResponse).contains("\"issue\":\"IJPL-248623\"") - webSocket.sendText(controlTaskFolderRequest("request-list", "listTaskFolderThreads", folderId = folder.id), true).join() + webSocket.sendText(controlTaskFolderRequest("request-folders", "listFolders"), true).join() + val foldersResponse = listener.nextMessage() + assertThat(foldersResponse).contains("\"folders\":") + assertThat(foldersResponse).contains("\"name\":\"Research\"") + + webSocket.sendText(controlTaskFolderRequest("request-list", "listThreads", folderId = folder.id), true).join() val listResponse = listener.nextMessage() assertThat(listResponse).contains("\"threadId\":${sessionId.jsonString()}") assertThat(listResponse).contains("\"folderId\":${folder.id.jsonString()}") webSocket.sendText( - controlTaskFolderRequest("request-set", "setTaskFolderMetadata", folderId = folder.id, key = "review", value = "backend"), + controlTaskFolderRequest("request-set", "setMetadata", folderId = folder.id, key = "review", value = "backend"), true, ).join() val setResponse = listener.nextMessage() assertThat(setResponse).contains("\"changed\":true") assertThat(folderService.getFolder(folder.id)?.metadata).containsEntry("review", "backend") - webSocket.sendText(controlTaskFolderRequest("request-get-metadata", "getTaskFolderMetadata", folderId = folder.id), true).join() - val metadataResponse = listener.nextMessage() - assertThat(metadataResponse).contains("\"requestId\":\"request-get-metadata\"") - assertThat(metadataResponse).contains("\"metadata\":") - assertThat(metadataResponse).contains("\"review\":\"backend\"") - - webSocket.sendText(controlTaskFolderRequest("request-delete", "deleteTaskFolderMetadata", folderId = folder.id, key = "review"), + webSocket.sendText(controlTaskFolderRequest("request-delete-metadata", "deleteMetadata", folderId = folder.id, key = "review"), true).join() - val deleteResponse = listener.nextMessage() - assertThat(deleteResponse).contains("\"changed\":true") + val deleteMetadataResponse = listener.nextMessage() + assertThat(deleteMetadataResponse).contains("\"changed\":true") assertThat(folderService.getFolder(folder.id)?.metadata).doesNotContainKey("review") + webSocket.sendText( + controlTaskFolderRequest( + requestId = "request-duplicate-create", + operation = "createAndAssign", + name = "Should not be created", + metadata = mapOf("issue" to "IJPL-000000"), + ), + true, + ).join() + val duplicateCreateResponse = listener.nextMessage() + assertThat(duplicateCreateResponse).contains("\"created\":false") + assertThat(folderService.getFolderForThread(projectDir.toString(), PI_AGENT_SESSION_PROVIDER, sessionId)?.id).isEqualTo(folder.id) + + webSocket.sendText(controlTaskFolderRequest("request-unassign", "unassignCurrentThread"), true).join() + val unassignResponse = listener.nextMessage() + assertThat(unassignResponse).contains("\"changed\":true") + assertThat(folderService.getFolderForThread(projectDir.toString(), PI_AGENT_SESSION_PROVIDER, sessionId)).isNull() + webSocket.sendText( controlTaskFolderRequest( requestId = "request-create", - type = "createAndAssignTaskFolder", + operation = "createAndAssign", name = "Follow-up task", metadata = mapOf("issue" to "IJPL-999999"), ), @@ -255,10 +274,39 @@ class PiExtensionControlWebSocketHandlerTest { val createdFolder = requireNotNull(folderService.getFolderForThread(projectDir.toString(), PI_AGENT_SESSION_PROVIDER, sessionId)) assertThat(createResponse).contains("\"requestId\":\"request-create\"") assertThat(createResponse).contains("\"ok\":true") + assertThat(createResponse).contains("\"created\":true") assertThat(createResponse).contains("\"assigned\":true") assertThat(createResponse).contains("\"id\":${createdFolder.id.jsonString()}") assertThat(createdFolder.name).isEqualTo("Follow-up task") assertThat(createdFolder.metadata).containsEntry("issue", "IJPL-999999") + + webSocket.sendText(controlTaskFolderRequest("request-rename", "rename", name = "Renamed task"), true).join() + val renameResponse = listener.nextMessage() + assertThat(renameResponse).contains("\"changed\":true") + assertThat(folderService.getFolder(createdFolder.id)?.name).isEqualTo("Renamed task") + + webSocket.sendText(controlTaskFolderRequest("request-assign-existing", "assignCurrentThread", folderId = folder.id), true).join() + val assignResponse = listener.nextMessage() + assertThat(assignResponse).contains("\"changed\":true") + assertThat(folderService.getFolderForThread(projectDir.toString(), PI_AGENT_SESSION_PROVIDER, sessionId)?.id).isEqualTo(folder.id) + + val otherProjectFolder = requireNotNull(folderService.createFolder(tempDir.resolve("other-project").toString(), "Other")) + webSocket.sendText(controlTaskFolderRequest("request-cross-path", "listThreads", folderId = otherProjectFolder.id), true).join() + val crossPathResponse = listener.nextMessage() + assertThat(crossPathResponse).contains("\"ok\":false") + assertThat(crossPathResponse).contains("Task folder is not available") + + val emptyFolder = requireNotNull(folderService.createFolder(projectDir.toString(), "Empty")) + webSocket.sendText(controlTaskFolderRequest("request-mark-done", "markDone", folderId = emptyFolder.id), true).join() + val markDoneResponse = listener.nextMessage() + assertThat(markDoneResponse).contains("\"changed\":true") + assertThat(markDoneResponse).contains("\"requestedCount\":0") + assertThat(folderService.getFolder(emptyFolder.id)?.status).isEqualTo(AgentTaskFolderStatus.DONE) + + webSocket.sendText(controlTaskFolderRequest("request-delete-folder", "delete", folderId = emptyFolder.id), true).join() + val deleteFolderResponse = listener.nextMessage() + assertThat(deleteFolderResponse).contains("\"changed\":true") + assertThat(folderService.getFolder(emptyFolder.id)).isNull() } finally { webSocket.sendClose(WebSocket.NORMAL_CLOSURE, "done").join() @@ -285,7 +333,7 @@ class PiExtensionControlWebSocketHandlerTest { ).join() listener.nextMessage() - webSocket.sendText(controlTaskFolderRequest("request-unsupported", "unknownControl"), true).join() + webSocket.sendText("{\"type\":\"unknownControl\",\"requestId\":\"request-unsupported\"}", true).join() val unsupportedResponse = listener.nextMessage() assertThat(unsupportedResponse).contains("\"type\":\"response\"") assertThat(unsupportedResponse).contains("\"requestId\":\"request-unsupported\"") @@ -440,27 +488,32 @@ private fun controlForkResponsePayload(requestId: String, forkedSessionId: Strin private fun controlTaskFolderRequest( requestId: String, - type: String, + operation: String, folderId: String? = null, name: String? = null, key: String? = null, value: String? = null, + includeDone: Boolean? = null, metadata: Map? = null, ): String { - val fields = mutableListOf( - "\"type\":${type.jsonString()}", - "\"requestId\":${requestId.jsonString()}", - ) - folderId?.let { fields += "\"folderId\":${it.jsonString()}" } - name?.let { fields += "\"name\":${it.jsonString()}" } - key?.let { fields += "\"key\":${it.jsonString()}" } - value?.let { fields += "\"value\":${it.jsonString()}" } + val arguments = mutableListOf() + folderId?.let { arguments += "\"folderId\":${it.jsonString()}" } + name?.let { arguments += "\"name\":${it.jsonString()}" } + key?.let { arguments += "\"key\":${it.jsonString()}" } + value?.let { arguments += "\"value\":${it.jsonString()}" } + includeDone?.let { arguments += "\"includeDone\":$it" } metadata?.let { values -> val metadataFields = values.entries.joinToString(",") { (metadataKey, metadataValue) -> "${metadataKey.jsonString()}:${metadataValue.jsonString()}" } - fields += "\"metadata\":{$metadataFields}" + arguments += "\"metadata\":{$metadataFields}" } + val fields = mutableListOf( + "\"type\":\"taskFolderRequest\"", + "\"requestId\":${requestId.jsonString()}", + "\"operation\":${operation.jsonString()}", + "\"arguments\":{${arguments.joinToString(",")}}", + ) return "{${fields.joinToString(",")}}" } diff --git a/plugins/agent-workbench/lib-agent/providers/pi/sessions/testSrc/PiThemeSupportTest.kt b/plugins/agent-workbench/lib-agent/providers/pi/sessions/testSrc/PiThemeSupportTest.kt index e4eb2ebb7374..4b914ffbe018 100644 --- a/plugins/agent-workbench/lib-agent/providers/pi/sessions/testSrc/PiThemeSupportTest.kt +++ b/plugins/agent-workbench/lib-agent/providers/pi/sessions/testSrc/PiThemeSupportTest.kt @@ -211,16 +211,23 @@ class PiThemeSupportTest { val control = readBundledPiExtensionText("control.ts") assertThat(entrypoint).contains( - "import {registerTaskFolderTool, startControlBridge} from \"./control.ts\";", - "registerTaskFolderTool(pi, () => controlBridge);", + "import {registerTaskFolderTools, startControlBridge} from \"./control.ts\";", + "registerTaskFolderTools(pi, () => controlBridge);", ) assertThat(control).contains( + "const TASK_FOLDER_TOOL_DEFINITIONS", "pi.registerTool(defineTool({", + "name: \"agent_workbench_get_current_task_folder\"", + "name: \"agent_workbench_list_task_folders\"", "name: \"agent_workbench_create_task_folder\"", + "name: \"agent_workbench_set_task_folder_metadata\"", + "name: \"agent_workbench_mark_task_folder_done\"", + "name: \"agent_workbench_delete_task_folder\"", "parameters: Type.Object({", "name: Type.String({description: \"Task folder name\"})", - "issue: Type.Optional(Type.String({description: \"Issue tracker id to store as task folder metadata key 'issue'\"}))", - "const result = await bridge.createAndAssignTaskFolder(params.name, metadata);", + "metadata: Type.Optional(Type.Record(Type.String(), Type.String({description: TASK_FOLDER_METADATA_DESCRIPTION})))", + "Use metadata key 'issue' for issue tracker ids and 'review' for review ids; do not use separate issue parameters.", + "const result = await bridge.requestTaskFolder(definition.operation, definition.arguments?.(normalizedParams) ?? normalizedParams);", "details: result", ) } diff --git a/plugins/agent-workbench/sessions-toolwindow/src/actions/AgentSessionsTreePopupTaskFolderAgentActions.kt b/plugins/agent-workbench/sessions-toolwindow/src/actions/AgentSessionsTreePopupTaskFolderAgentActions.kt index 7428673401e6..a3c8d921b5f4 100644 --- a/plugins/agent-workbench/sessions-toolwindow/src/actions/AgentSessionsTreePopupTaskFolderAgentActions.kt +++ b/plugins/agent-workbench/sessions-toolwindow/src/actions/AgentSessionsTreePopupTaskFolderAgentActions.kt @@ -199,7 +199,7 @@ internal fun buildTaskFolderAgentPrompt( Project path: $path - When you know the task folder name, call the agent_workbench_create_task_folder tool. If the user mentions an issue tracker id, pass it as the issue parameter. + When you know the task folder name, call the agent_workbench_create_task_folder tool. If the user mentions an issue tracker id, store it in metadata key "issue". If the task details are missing, ask one concise follow-up before creating the folder. """.trimIndent() } diff --git a/plugins/agent-workbench/spec/sessions/agent-sessions-pi.spec.md b/plugins/agent-workbench/spec/sessions/agent-sessions-pi.spec.md index de1ce730bf97..100fabdbbb18 100644 --- a/plugins/agent-workbench/spec/sessions/agent-sessions-pi.spec.md +++ b/plugins/agent-workbench/spec/sessions/agent-sessions-pi.spec.md @@ -90,11 +90,13 @@ Agent Workbench treats Pi as a first-class terminal-backed provider. Pi sessions - The bundled Pi extension must open one private IDE-local WebSocket control connection for live Structure View actions. The WebSocket handshake must be authenticated with `Authorization: Bearer `, the first `hello` frame must bind the connection to the same launch-scoped token/session id/cwd, and every command must target the currently bound session id. The IDE may send `navigateTree` and `forkFromEntry` commands only when the live connection advertises those capabilities. Fork commands must use Pi `fork(entryId, { position: "at", withSession })` and return the fresh replacement session state captured from `withSession`; there is no long-polling, polling, slash-command, or unauthenticated fallback. [@test] ../../lib-agent/providers/pi/sessions/testSrc/PiExtensionControlWebSocketHandlerTest.kt -- The same authenticated control WebSocket may carry explicit task-folder capability requests from the bundled Pi extension to the IDE: - `getCurrentTaskFolder`, `listTaskFolderThreads`, `createAndAssignTaskFolder`, `getTaskFolderMetadata`, - `setTaskFolderMetadata`, and `deleteTaskFolderMetadata`. These requests must use the bound cwd/session id to resolve current-session context, use global folder id - for explicit folder mutations, must not inject prompt context automatically, and must return normal `response` frames with `ok`, - `requestId`, and the requested folder, assignment, metadata, or mutation result fields. +- The same authenticated control WebSocket may carry explicit task-folder capability requests from the bundled Pi extension to the IDE using + `taskFolderRequest` frames with an `operation` string and nested `arguments` object. Supported operations cover current-folder lookup, + folder listing, thread assignment listing, create-and-assign, current-thread assign/unassign, rename, metadata set/delete, mark done, + and delete. These requests must use the bound cwd/session id to resolve current-session context, use global folder id for explicit + folder mutations, must not inject prompt context automatically, and must return normal `response` frames with `ok`, `requestId`, and a + `result` object containing the requested folder, assignment, metadata, or mutation fields. Metadata is ordinary string key/value data; + `issue` and `review` are conventional keys, not separate protocol fields. [@test] ../../lib-agent/providers/pi/sessions/testSrc/PiExtensionControlWebSocketHandlerTest.kt [@test] ../../lib-agent/providers/pi/sessions/testSrc/PiThemeSupportTest.kt