mirror of
https://github.com/getcompanion-ai/co-mono.git
synced 2026-04-21 05:02:14 +00:00
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:
parent
0d9fddec1e
commit
03159d2f4b
10 changed files with 68 additions and 9 deletions
|
|
@ -76,7 +76,10 @@ execute(toolCallId, params, signal, onUpdate)
|
||||||
execute(toolCallId, params, onUpdate, ctx, signal?)
|
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:**
|
**Session event changes:**
|
||||||
- `CustomToolSessionEvent` now only has `reason` and `previousSessionFile`
|
- `CustomToolSessionEvent` now only has `reason` and `previousSessionFile`
|
||||||
|
|
|
||||||
|
|
@ -230,11 +230,33 @@ interface CustomToolContext {
|
||||||
sessionManager: ReadonlySessionManager; // Read-only access to session
|
sessionManager: ReadonlySessionManager; // Read-only access to session
|
||||||
modelRegistry: ModelRegistry; // For API key resolution
|
modelRegistry: ModelRegistry; // For API key resolution
|
||||||
model: Model | undefined; // Current model (may be undefined)
|
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.
|
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
|
## Session Lifecycle
|
||||||
|
|
||||||
Tools can implement `onSession` to react to session changes:
|
Tools can implement `onSession` to react to session changes:
|
||||||
|
|
|
||||||
|
|
@ -1878,6 +1878,11 @@ export class AgentSession {
|
||||||
sessionManager: this.sessionManager,
|
sessionManager: this.sessionManager,
|
||||||
modelRegistry: this._modelRegistry,
|
modelRegistry: this._modelRegistry,
|
||||||
model: this.agent.state.model,
|
model: this.agent.state.model,
|
||||||
|
isIdle: () => !this.isStreaming,
|
||||||
|
hasQueuedMessages: () => this.queuedMessageCount > 0,
|
||||||
|
abort: () => {
|
||||||
|
this.abort();
|
||||||
|
},
|
||||||
};
|
};
|
||||||
|
|
||||||
for (const { tool } of this._customTools) {
|
for (const { tool } of this._customTools) {
|
||||||
|
|
|
||||||
|
|
@ -47,6 +47,12 @@ export interface CustomToolContext {
|
||||||
modelRegistry: ModelRegistry;
|
modelRegistry: ModelRegistry;
|
||||||
/** Current model (may be undefined if no model is selected yet) */
|
/** Current model (may be undefined if no model is selected yet) */
|
||||||
model: Model<any> | undefined;
|
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 */
|
/** Session event passed to onSession callback */
|
||||||
|
|
|
||||||
|
|
@ -71,7 +71,7 @@ export class HookRunner {
|
||||||
private getModel: () => Model<any> | undefined = () => undefined;
|
private getModel: () => Model<any> | undefined = () => undefined;
|
||||||
private isIdleFn: () => boolean = () => true;
|
private isIdleFn: () => boolean = () => true;
|
||||||
private waitForIdleFn: () => Promise<void> = async () => {};
|
private waitForIdleFn: () => Promise<void> = async () => {};
|
||||||
private abortFn: () => Promise<void> = async () => {};
|
private abortFn: () => void = () => {};
|
||||||
private hasQueuedMessagesFn: () => boolean = () => false;
|
private hasQueuedMessagesFn: () => boolean = () => false;
|
||||||
private newSessionHandler: NewSessionHandler = async () => ({ cancelled: false });
|
private newSessionHandler: NewSessionHandler = async () => ({ cancelled: false });
|
||||||
private branchHandler: BranchHandler = async () => ({ cancelled: false });
|
private branchHandler: BranchHandler = async () => ({ cancelled: false });
|
||||||
|
|
@ -107,8 +107,8 @@ export class HookRunner {
|
||||||
isIdle?: () => boolean;
|
isIdle?: () => boolean;
|
||||||
/** Function to wait for agent to be idle */
|
/** Function to wait for agent to be idle */
|
||||||
waitForIdle?: () => Promise<void>;
|
waitForIdle?: () => Promise<void>;
|
||||||
/** Function to abort current operation */
|
/** Function to abort current operation (fire-and-forget) */
|
||||||
abort?: () => Promise<void>;
|
abort?: () => void;
|
||||||
/** Function to check if there are queued messages */
|
/** Function to check if there are queued messages */
|
||||||
hasQueuedMessages?: () => boolean;
|
hasQueuedMessages?: () => boolean;
|
||||||
/** UI context for interactive prompts */
|
/** UI context for interactive prompts */
|
||||||
|
|
@ -119,7 +119,7 @@ export class HookRunner {
|
||||||
this.getModel = options.getModel;
|
this.getModel = options.getModel;
|
||||||
this.isIdleFn = options.isIdle ?? (() => true);
|
this.isIdleFn = options.isIdle ?? (() => true);
|
||||||
this.waitForIdleFn = options.waitForIdle ?? (async () => {});
|
this.waitForIdleFn = options.waitForIdle ?? (async () => {});
|
||||||
this.abortFn = options.abort ?? (async () => {});
|
this.abortFn = options.abort ?? (() => {});
|
||||||
this.hasQueuedMessagesFn = options.hasQueuedMessages ?? (() => false);
|
this.hasQueuedMessagesFn = options.hasQueuedMessages ?? (() => false);
|
||||||
// Store session handlers for HookCommandContext
|
// Store session handlers for HookCommandContext
|
||||||
if (options.newSessionHandler) {
|
if (options.newSessionHandler) {
|
||||||
|
|
|
||||||
|
|
@ -150,7 +150,7 @@ export interface HookContext {
|
||||||
/** Whether the agent is idle (not streaming) */
|
/** Whether the agent is idle (not streaming) */
|
||||||
isIdle(): boolean;
|
isIdle(): boolean;
|
||||||
/** Abort the current agent operation (fire-and-forget, does not wait) */
|
/** 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 */
|
/** Whether there are queued messages waiting to be processed */
|
||||||
hasQueuedMessages(): boolean;
|
hasQueuedMessages(): boolean;
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -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 agent: Agent;
|
||||||
|
let session: AgentSession;
|
||||||
const wrappedCustomTools = wrapCustomTools(customToolsResult.tools, () => ({
|
const wrappedCustomTools = wrapCustomTools(customToolsResult.tools, () => ({
|
||||||
sessionManager,
|
sessionManager,
|
||||||
modelRegistry,
|
modelRegistry,
|
||||||
model: agent.state.model,
|
model: agent.state.model,
|
||||||
|
isIdle: () => !session.isStreaming,
|
||||||
|
hasQueuedMessages: () => session.queuedMessageCount > 0,
|
||||||
|
abort: () => {
|
||||||
|
session.abort();
|
||||||
|
},
|
||||||
}));
|
}));
|
||||||
|
|
||||||
let allToolsArray: Tool[] = [...builtInTools, ...wrappedCustomTools];
|
let allToolsArray: Tool[] = [...builtInTools, ...wrappedCustomTools];
|
||||||
|
|
@ -646,7 +652,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
|
||||||
sessionManager.appendThinkingLevelChange(thinkingLevel);
|
sessionManager.appendThinkingLevelChange(thinkingLevel);
|
||||||
}
|
}
|
||||||
|
|
||||||
const session = new AgentSession({
|
session = new AgentSession({
|
||||||
agent,
|
agent,
|
||||||
sessionManager,
|
sessionManager,
|
||||||
settingsManager,
|
settingsManager,
|
||||||
|
|
|
||||||
|
|
@ -476,7 +476,9 @@ export class InteractiveMode {
|
||||||
},
|
},
|
||||||
isIdle: () => !this.session.isStreaming,
|
isIdle: () => !this.session.isStreaming,
|
||||||
waitForIdle: () => this.session.agent.waitForIdle(),
|
waitForIdle: () => this.session.agent.waitForIdle(),
|
||||||
abort: () => this.session.abort(),
|
abort: () => {
|
||||||
|
this.session.abort();
|
||||||
|
},
|
||||||
hasQueuedMessages: () => this.session.queuedMessageCount > 0,
|
hasQueuedMessages: () => this.session.queuedMessageCount > 0,
|
||||||
uiContext,
|
uiContext,
|
||||||
hasUI: true,
|
hasUI: true,
|
||||||
|
|
@ -512,6 +514,11 @@ export class InteractiveMode {
|
||||||
sessionManager: this.session.sessionManager,
|
sessionManager: this.session.sessionManager,
|
||||||
modelRegistry: this.session.modelRegistry,
|
modelRegistry: this.session.modelRegistry,
|
||||||
model: this.session.model,
|
model: this.session.model,
|
||||||
|
isIdle: () => !this.session.isStreaming,
|
||||||
|
hasQueuedMessages: () => this.session.queuedMessageCount > 0,
|
||||||
|
abort: () => {
|
||||||
|
this.session.abort();
|
||||||
|
},
|
||||||
});
|
});
|
||||||
} catch (err) {
|
} catch (err) {
|
||||||
this.showToolError(tool.name, err instanceof Error ? err.message : String(err));
|
this.showToolError(tool.name, err instanceof Error ? err.message : String(err));
|
||||||
|
|
|
||||||
|
|
@ -63,6 +63,11 @@ export async function runPrintMode(
|
||||||
sessionManager: session.sessionManager,
|
sessionManager: session.sessionManager,
|
||||||
modelRegistry: session.modelRegistry,
|
modelRegistry: session.modelRegistry,
|
||||||
model: session.model,
|
model: session.model,
|
||||||
|
isIdle: () => !session.isStreaming,
|
||||||
|
hasQueuedMessages: () => session.queuedMessageCount > 0,
|
||||||
|
abort: () => {
|
||||||
|
session.abort();
|
||||||
|
},
|
||||||
},
|
},
|
||||||
);
|
);
|
||||||
} catch (_err) {
|
} catch (_err) {
|
||||||
|
|
|
||||||
|
|
@ -196,6 +196,11 @@ export async function runRpcMode(session: AgentSession): Promise<never> {
|
||||||
sessionManager: session.sessionManager,
|
sessionManager: session.sessionManager,
|
||||||
modelRegistry: session.modelRegistry,
|
modelRegistry: session.modelRegistry,
|
||||||
model: session.model,
|
model: session.model,
|
||||||
|
isIdle: () => !session.isStreaming,
|
||||||
|
hasQueuedMessages: () => session.queuedMessageCount > 0,
|
||||||
|
abort: () => {
|
||||||
|
session.abort();
|
||||||
|
},
|
||||||
},
|
},
|
||||||
);
|
);
|
||||||
} catch (_err) {
|
} catch (_err) {
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue