Improve navigateTree API

This commit is contained in:
Armin Ronacher 2026-01-16 19:52:54 +01:00 committed by Mario Zechner
parent ffdc8d686b
commit 6b6707f30c
5 changed files with 68 additions and 10 deletions

View file

@ -1972,11 +1972,13 @@ export class AgentSession {
* @param targetId The entry ID to navigate to
* @param options.summarize Whether user wants to summarize abandoned branch
* @param options.customInstructions Custom instructions for summarizer
* @param options.replaceInstructions If true, customInstructions replaces the default prompt
* @param options.label Label to attach to the branch summary entry
* @returns Result with editorText (if user message) and cancelled status
*/
async navigateTree(
targetId: string,
options: { summarize?: boolean; customInstructions?: string } = {},
options: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string } = {},
): Promise<{ editorText?: string; cancelled: boolean; aborted?: boolean; summaryEntry?: BranchSummaryEntry }> {
const oldLeafId = this.sessionManager.getLeafId();
@ -2002,13 +2004,20 @@ export class AgentSession {
targetId,
);
// Prepare event data
// Prepare event data - mutable so extensions can override
let customInstructions = options.customInstructions;
let replaceInstructions = options.replaceInstructions;
let label = options.label;
const preparation: TreePreparation = {
targetId,
oldLeafId,
commonAncestorId,
entriesToSummarize,
userWantsSummary: options.summarize ?? false,
customInstructions,
replaceInstructions,
label,
};
// Set up abort controller for summarization
@ -2032,6 +2041,17 @@ export class AgentSession {
extensionSummary = result.summary;
fromExtension = true;
}
// Allow extensions to override instructions and label
if (result?.customInstructions !== undefined) {
customInstructions = result.customInstructions;
}
if (result?.replaceInstructions !== undefined) {
replaceInstructions = result.replaceInstructions;
}
if (result?.label !== undefined) {
label = result.label;
}
}
// Run default summarizer if needed
@ -2048,7 +2068,8 @@ export class AgentSession {
model,
apiKey,
signal: this._branchSummaryAbortController.signal,
customInstructions: options.customInstructions,
customInstructions,
replaceInstructions,
reserveTokens: branchSummarySettings.reserveTokens,
});
this._branchSummaryAbortController = undefined;
@ -2098,6 +2119,11 @@ export class AgentSession {
// Create summary at target position (can be null for root)
const summaryId = this.sessionManager.branchWithSummary(newLeafId, summaryText, summaryDetails, fromExtension);
summaryEntry = this.sessionManager.getEntry(summaryId) as BranchSummaryEntry;
// Attach label to the summary entry
if (label) {
this.sessionManager.appendLabelChange(summaryId, label);
}
} else if (newLeafId === null) {
// No summary, navigating to root - reset leaf
this.sessionManager.resetLeaf();
@ -2106,6 +2132,11 @@ export class AgentSession {
this.sessionManager.branch(newLeafId);
}
// Attach label to target entry when not summarizing (no summary entry to label)
if (label && !summaryText) {
this.sessionManager.appendLabelChange(targetId, label);
}
// Update agent state
const sessionContext = this.sessionManager.buildSessionContext();
this.agent.replaceMessages(sessionContext.messages);

View file

@ -71,6 +71,8 @@ export interface GenerateBranchSummaryOptions {
signal: AbortSignal;
/** Optional custom instructions for summarization */
customInstructions?: string;
/** If true, customInstructions replaces the default prompt instead of being appended */
replaceInstructions?: boolean;
/** Tokens reserved for prompt + LLM response (default 16384) */
reserveTokens?: number;
}
@ -279,7 +281,7 @@ export async function generateBranchSummary(
entries: SessionEntry[],
options: GenerateBranchSummaryOptions,
): Promise<BranchSummaryResult> {
const { model, apiKey, signal, customInstructions, reserveTokens = 16384 } = options;
const { model, apiKey, signal, customInstructions, replaceInstructions, reserveTokens = 16384 } = options;
// Token budget = context window minus reserved space for prompt + response
const contextWindow = model.contextWindow || 128000;
@ -297,9 +299,14 @@ export async function generateBranchSummary(
const conversationText = serializeConversation(llmMessages);
// Build prompt
const instructions = customInstructions
? `${BRANCH_SUMMARY_PROMPT}\n\nAdditional focus: ${customInstructions}`
: BRANCH_SUMMARY_PROMPT;
let instructions: string;
if (replaceInstructions && customInstructions) {
instructions = customInstructions;
} else if (customInstructions) {
instructions = `${BRANCH_SUMMARY_PROMPT}\n\nAdditional focus: ${customInstructions}`;
} else {
instructions = BRANCH_SUMMARY_PROMPT;
}
const promptText = `<conversation>\n${conversationText}\n</conversation>\n\n${instructions}`;
const summarizationMessages = [

View file

@ -57,7 +57,7 @@ export type ForkHandler = (entryId: string) => Promise<{ cancelled: boolean }>;
export type NavigateTreeHandler = (
targetId: string,
options?: { summarize?: boolean },
options?: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string },
) => Promise<{ cancelled: boolean }>;
export type ShutdownHandler = () => void;

View file

@ -344,6 +344,12 @@ export interface TreePreparation {
commonAncestorId: string | null;
entriesToSummarize: SessionEntry[];
userWantsSummary: boolean;
/** Custom instructions for summarization */
customInstructions?: string;
/** If true, customInstructions replaces the default prompt instead of being appended */
replaceInstructions?: boolean;
/** Label to attach to the branch summary entry */
label?: string;
}
/** Fired before navigating in the session tree (can be cancelled) */
@ -633,6 +639,12 @@ export interface SessionBeforeTreeResult {
summary: string;
details?: unknown;
};
/** Override custom instructions for summarization */
customInstructions?: string;
/** Override whether customInstructions replaces the default prompt */
replaceInstructions?: boolean;
/** Override label to attach to the branch summary entry */
label?: string;
}
// ============================================================================
@ -917,7 +929,10 @@ export interface ExtensionCommandContextActions {
setup?: (sessionManager: SessionManager) => Promise<void>;
}) => Promise<{ cancelled: boolean }>;
fork: (entryId: string) => Promise<{ cancelled: boolean }>;
navigateTree: (targetId: string, options?: { summarize?: boolean }) => Promise<{ cancelled: boolean }>;
navigateTree: (
targetId: string,
options?: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string },
) => Promise<{ cancelled: boolean }>;
}
/**

View file

@ -717,7 +717,12 @@ export class InteractiveMode {
return { cancelled: false };
},
navigateTree: async (targetId, options) => {
const result = await this.session.navigateTree(targetId, { summarize: options?.summarize });
const result = await this.session.navigateTree(targetId, {
summarize: options?.summarize,
customInstructions: options?.customInstructions,
replaceInstructions: options?.replaceInstructions,
label: options?.label,
});
if (result.cancelled) {
return { cancelled: true };
}