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<void> 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
This commit is contained in:
Mario Zechner 2026-01-02 00:31:23 +01:00
parent 0d9fddec1e
commit 03159d2f4b
10 changed files with 68 additions and 9 deletions

View file

@ -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) {

View file

@ -47,6 +47,12 @@ export interface CustomToolContext {
modelRegistry: ModelRegistry;
/** Current model (may be undefined if no model is selected yet) */
model: Model<any> | 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 */

View file

@ -71,7 +71,7 @@ export class HookRunner {
private getModel: () => Model<any> | undefined = () => undefined;
private isIdleFn: () => boolean = () => true;
private waitForIdleFn: () => Promise<void> = async () => {};
private abortFn: () => Promise<void> = 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<void>;
/** Function to abort current operation */
abort?: () => Promise<void>;
/** 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) {

View file

@ -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<void>;
abort(): void;
/** Whether there are queued messages waiting to be processed */
hasQueuedMessages(): boolean;
}

View file

@ -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,

View file

@ -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));

View file

@ -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) {

View file

@ -196,6 +196,11 @@ export async function runRpcMode(session: AgentSession): Promise<never> {
sessionManager: session.sessionManager,
modelRegistry: session.modelRegistry,
model: session.model,
isIdle: () => !session.isStreaming,
hasQueuedMessages: () => session.queuedMessageCount > 0,
abort: () => {
session.abort();
},
},
);
} catch (_err) {