From 03159d2f4b37a8dde0f80ea844468cf070617ca8 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Fri, 2 Jan 2026 00:31:23 +0100 Subject: [PATCH] Add agent state methods to CustomToolContext and fix abort signature CustomToolContext now has: - isIdle() - check if agent is streaming - hasQueuedMessages() - check if user has queued messages - abort() - abort current operation (fire-and-forget) Changed abort() signature from Promise to void in both HookContext and CustomToolContext. The abort is fire-and-forget: it calls session.abort() without awaiting, so the abort signal is set immediately while waitForIdle() runs in the background. Fixes #388 --- packages/coding-agent/CHANGELOG.md | 5 ++++- packages/coding-agent/docs/custom-tools.md | 22 +++++++++++++++++++ .../coding-agent/src/core/agent-session.ts | 5 +++++ .../src/core/custom-tools/types.ts | 6 +++++ .../coding-agent/src/core/hooks/runner.ts | 8 +++---- packages/coding-agent/src/core/hooks/types.ts | 2 +- packages/coding-agent/src/core/sdk.ts | 10 +++++++-- .../src/modes/interactive/interactive-mode.ts | 9 +++++++- packages/coding-agent/src/modes/print-mode.ts | 5 +++++ .../coding-agent/src/modes/rpc/rpc-mode.ts | 5 +++++ 10 files changed, 68 insertions(+), 9 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 58b128e0..01d2711d 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -76,7 +76,10 @@ execute(toolCallId, params, signal, onUpdate) execute(toolCallId, params, onUpdate, ctx, signal?) ``` -The new `ctx: CustomToolContext` provides `sessionManager`, `modelRegistry`, and `model`. +The new `ctx: CustomToolContext` provides `sessionManager`, `modelRegistry`, `model`, and agent state methods: +- `ctx.isIdle()` - check if agent is streaming +- `ctx.hasQueuedMessages()` - check if user has queued messages (skip interactive prompts) +- `ctx.abort()` - abort current operation (fire-and-forget) **Session event changes:** - `CustomToolSessionEvent` now only has `reason` and `previousSessionFile` diff --git a/packages/coding-agent/docs/custom-tools.md b/packages/coding-agent/docs/custom-tools.md index 3940b43e..2e546618 100644 --- a/packages/coding-agent/docs/custom-tools.md +++ b/packages/coding-agent/docs/custom-tools.md @@ -230,11 +230,33 @@ interface CustomToolContext { sessionManager: ReadonlySessionManager; // Read-only access to session modelRegistry: ModelRegistry; // For API key resolution model: Model | undefined; // Current model (may be undefined) + isIdle(): boolean; // Whether agent is streaming + hasQueuedMessages(): boolean; // Whether user has queued messages + abort(): void; // Abort current operation (fire-and-forget) } ``` Use `ctx.sessionManager.getBranch()` to get entries on the current branch for state reconstruction. +### Checking Queue State + +Interactive tools can skip prompts when the user has already queued a message: + +```typescript +async execute(toolCallId, params, onUpdate, ctx, signal) { + // If user already queued a message, skip the interactive prompt + if (ctx.hasQueuedMessages()) { + return { + content: [{ type: "text", text: "Skipped - user has queued input" }], + }; + } + + // Otherwise, prompt for input + const answer = await pi.ui.input("What would you like to do?"); + // ... +} +``` + ## Session Lifecycle Tools can implement `onSession` to react to session changes: diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index aaec0ee0..9a4d752c 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -1878,6 +1878,11 @@ export class AgentSession { sessionManager: this.sessionManager, modelRegistry: this._modelRegistry, model: this.agent.state.model, + isIdle: () => !this.isStreaming, + hasQueuedMessages: () => this.queuedMessageCount > 0, + abort: () => { + this.abort(); + }, }; for (const { tool } of this._customTools) { diff --git a/packages/coding-agent/src/core/custom-tools/types.ts b/packages/coding-agent/src/core/custom-tools/types.ts index 69a95ad4..59d800f9 100644 --- a/packages/coding-agent/src/core/custom-tools/types.ts +++ b/packages/coding-agent/src/core/custom-tools/types.ts @@ -47,6 +47,12 @@ export interface CustomToolContext { modelRegistry: ModelRegistry; /** Current model (may be undefined if no model is selected yet) */ model: Model | undefined; + /** Whether the agent is idle (not streaming) */ + isIdle(): boolean; + /** Whether there are queued messages waiting to be processed */ + hasQueuedMessages(): boolean; + /** Abort the current agent operation (fire-and-forget, does not wait) */ + abort(): void; } /** Session event passed to onSession callback */ diff --git a/packages/coding-agent/src/core/hooks/runner.ts b/packages/coding-agent/src/core/hooks/runner.ts index 2c029285..d20bb170 100644 --- a/packages/coding-agent/src/core/hooks/runner.ts +++ b/packages/coding-agent/src/core/hooks/runner.ts @@ -71,7 +71,7 @@ export class HookRunner { private getModel: () => Model | undefined = () => undefined; private isIdleFn: () => boolean = () => true; private waitForIdleFn: () => Promise = async () => {}; - private abortFn: () => Promise = async () => {}; + private abortFn: () => void = () => {}; private hasQueuedMessagesFn: () => boolean = () => false; private newSessionHandler: NewSessionHandler = async () => ({ cancelled: false }); private branchHandler: BranchHandler = async () => ({ cancelled: false }); @@ -107,8 +107,8 @@ export class HookRunner { isIdle?: () => boolean; /** Function to wait for agent to be idle */ waitForIdle?: () => Promise; - /** Function to abort current operation */ - abort?: () => Promise; + /** Function to abort current operation (fire-and-forget) */ + abort?: () => void; /** Function to check if there are queued messages */ hasQueuedMessages?: () => boolean; /** UI context for interactive prompts */ @@ -119,7 +119,7 @@ export class HookRunner { this.getModel = options.getModel; this.isIdleFn = options.isIdle ?? (() => true); this.waitForIdleFn = options.waitForIdle ?? (async () => {}); - this.abortFn = options.abort ?? (async () => {}); + this.abortFn = options.abort ?? (() => {}); this.hasQueuedMessagesFn = options.hasQueuedMessages ?? (() => false); // Store session handlers for HookCommandContext if (options.newSessionHandler) { diff --git a/packages/coding-agent/src/core/hooks/types.ts b/packages/coding-agent/src/core/hooks/types.ts index 4910fe6f..a15f2e7d 100644 --- a/packages/coding-agent/src/core/hooks/types.ts +++ b/packages/coding-agent/src/core/hooks/types.ts @@ -150,7 +150,7 @@ export interface HookContext { /** Whether the agent is idle (not streaming) */ isIdle(): boolean; /** Abort the current agent operation (fire-and-forget, does not wait) */ - abort(): Promise; + abort(): void; /** Whether there are queued messages waiting to be processed */ hasQueuedMessages(): boolean; } diff --git a/packages/coding-agent/src/core/sdk.ts b/packages/coding-agent/src/core/sdk.ts index 5ab85ef6..7f945bcc 100644 --- a/packages/coding-agent/src/core/sdk.ts +++ b/packages/coding-agent/src/core/sdk.ts @@ -567,12 +567,18 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} } } - // Wrap custom tools with context getter (agent is assigned below, accessed at execute time) + // Wrap custom tools with context getter (agent/session assigned below, accessed at execute time) let agent: Agent; + let session: AgentSession; const wrappedCustomTools = wrapCustomTools(customToolsResult.tools, () => ({ sessionManager, modelRegistry, model: agent.state.model, + isIdle: () => !session.isStreaming, + hasQueuedMessages: () => session.queuedMessageCount > 0, + abort: () => { + session.abort(); + }, })); let allToolsArray: Tool[] = [...builtInTools, ...wrappedCustomTools]; @@ -646,7 +652,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} sessionManager.appendThinkingLevelChange(thinkingLevel); } - const session = new AgentSession({ + session = new AgentSession({ agent, sessionManager, settingsManager, diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index d6fe474b..663cc09a 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -476,7 +476,9 @@ export class InteractiveMode { }, isIdle: () => !this.session.isStreaming, waitForIdle: () => this.session.agent.waitForIdle(), - abort: () => this.session.abort(), + abort: () => { + this.session.abort(); + }, hasQueuedMessages: () => this.session.queuedMessageCount > 0, uiContext, hasUI: true, @@ -512,6 +514,11 @@ export class InteractiveMode { sessionManager: this.session.sessionManager, modelRegistry: this.session.modelRegistry, model: this.session.model, + isIdle: () => !this.session.isStreaming, + hasQueuedMessages: () => this.session.queuedMessageCount > 0, + abort: () => { + this.session.abort(); + }, }); } catch (err) { this.showToolError(tool.name, err instanceof Error ? err.message : String(err)); diff --git a/packages/coding-agent/src/modes/print-mode.ts b/packages/coding-agent/src/modes/print-mode.ts index b0aec7fa..04ff4f36 100644 --- a/packages/coding-agent/src/modes/print-mode.ts +++ b/packages/coding-agent/src/modes/print-mode.ts @@ -63,6 +63,11 @@ export async function runPrintMode( sessionManager: session.sessionManager, modelRegistry: session.modelRegistry, model: session.model, + isIdle: () => !session.isStreaming, + hasQueuedMessages: () => session.queuedMessageCount > 0, + abort: () => { + session.abort(); + }, }, ); } catch (_err) { diff --git a/packages/coding-agent/src/modes/rpc/rpc-mode.ts b/packages/coding-agent/src/modes/rpc/rpc-mode.ts index a5770767..6d863230 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-mode.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-mode.ts @@ -196,6 +196,11 @@ export async function runRpcMode(session: AgentSession): Promise { sessionManager: session.sessionManager, modelRegistry: session.modelRegistry, model: session.model, + isIdle: () => !session.isStreaming, + hasQueuedMessages: () => session.queuedMessageCount > 0, + abort: () => { + session.abort(); + }, }, ); } catch (_err) {