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
This commit is contained in:
Vladimir Krivosheev
2026-07-09 18:59:26 +00:00
committed by intellij-monorepo-bot
parent 15bcfedc2a
commit 09ff8a47d5
11 changed files with 629 additions and 221 deletions
@@ -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",
@@ -26,6 +26,7 @@
<orderEntry type="module" module-name="intellij.platform.ai.agent.core" />
<orderEntry type="module" module-name="intellij.platform.ai.agent.json" />
<orderEntry type="module" module-name="intellij.agent.workbench.prompt.core" />
<orderEntry type="module" module-name="intellij.agent.workbench.sessions" />
<orderEntry type="module" module-name="intellij.platform.ai.agent.sessions.core" />
<orderEntry type="module" module-name="intellij.libraries.jackson3" />
<orderEntry type="module" module-name="intellij.libraries.netty.buffer" />
@@ -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;
@@ -13,12 +13,7 @@ const STATUS_TOKEN = process.env[STATUS_TOKEN_ENV];
type AgentWorkbenchControlBridge = {
setContext: (ctx: ExtensionContext) => void;
getCurrentTaskFolder: () => Promise<AgentWorkbenchTaskFolder | undefined>;
listTaskFolderThreads: (folderId?: string) => Promise<AgentWorkbenchTaskFolderThread[]>;
createAndAssignTaskFolder: (name: string, metadata?: Record<string, string>) => Promise<AgentWorkbenchCreatedTaskFolder>;
getTaskFolderMetadata: (folderId?: string) => Promise<Record<string, string> | undefined>;
setTaskFolderMetadata: (key: string, value: string, folderId?: string) => Promise<boolean>;
deleteTaskFolderMetadata: (key: string, folderId?: string) => Promise<boolean>;
requestTaskFolder: <T = unknown>(operation: string, args?: Record<string, unknown>) => Promise<T>;
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<string, string>;
assigned?: boolean;
operation?: string;
arguments?: Record<string, unknown>;
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> | 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 <T = unknown>(operation: string, args: Record<string, unknown> = {}): Promise<T> => {
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<typeof Type.Object>;
operation: string;
arguments?: (params: Record<string, unknown>) => Record<string, unknown>;
confirm?: (params: Record<string, unknown>, ctx: ExtensionContext) => Promise<boolean>;
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<string, unknown>;
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<boolean> {
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<string, unknown> {
return typeof value === "object" && value !== null;
}
function resolveCapabilities(ctx: ExtensionContext): { navigateTree: boolean; fork: boolean } {
@@ -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) },
)
}
@@ -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<String, PiControlMessageType> = 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<String, String>? = 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<String, String>? = 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<String, String>? = 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,
)
}
@@ -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<AgentTaskFolderService>() },
private val archiveServiceProvider: () -> AgentSessionArchiveService = { service<AgentSessionArchiveService>() },
) {
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<AgentTaskFolder>): 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<AgentTaskFolderThreadAssignment>): 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, String>): 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"
@@ -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<String, String>? = 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<String>()
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(",")}}"
}
@@ -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",
)
}
@@ -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()
}
@@ -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 <launchToken>`, 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