ton/pi_harness
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge main into model-registry

Browse Source
This commit is contained in:
Mario Zechner
2026-06-22 14:00:18 +02:00
parent 019b022068 717a8f958a
commit abbd911693
220 changed files with 10488 additions and 4354 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/coding-agent/src/core/agent-session.ts
+68 -46
View File
@@ -33,7 +33,7 @@ import {
resetApiProviders,
streamSimple,
} from "@earendil-works/pi-ai/compat";
import { theme } from "../modes/interactive/theme/theme.ts";
import { getThemeByName, theme } from "../modes/interactive/theme/theme.ts";
import { stripFrontmatter } from "../utils/frontmatter.ts";
import { resolvePath } from "../utils/paths.ts";
import { sleep } from "../utils/sleep.ts";
@@ -45,6 +45,7 @@ import {
collectEntriesForBranchSummary,
compact,
estimateContextTokens,
estimateTokens,
generateBranchSummary,
prepareCompaction,
shouldCompact,
@@ -242,6 +243,14 @@ interface ToolDefinitionEntry {
sourceInfo: SourceInfo;
}
function estimateMessagesTokens(messages: AgentMessage[]): number {
let tokens = 0;
for (const message of messages) {
tokens += estimateTokens(message);
}
return tokens;
}
// ============================================================================
// Constants
// ============================================================================
@@ -357,6 +366,7 @@ export class AgentSession {
private async _getRequiredRequestAuth(model: Model<any>): Promise<{
apiKey: string;
headers?: Record<string, string>;
env?: Record<string, string>;
}> {
const result = await this._modelRegistry.getApiKeyAndHeaders(model);
if (!result.ok) {
@@ -366,7 +376,7 @@ export class AgentSession {
throw new Error(result.error);
}
if (result.apiKey) {
return { apiKey: result.apiKey, headers: result.headers };
return { apiKey: result.apiKey, headers: result.headers, env: result.env };
}
const isOAuth = this._modelRegistry.isUsingOAuth(model);
@@ -383,13 +393,14 @@ export class AgentSession {
private async _getCompactionRequestAuth(model: Model<any>): Promise<{
apiKey?: string;
headers?: Record<string, string>;
env?: Record<string, string>;
}> {
if (this.agent.streamFn === streamSimple) {
return this._getRequiredRequestAuth(model);
}
const result = await this._modelRegistry.getApiKeyAndHeaders(model);
return result.ok ? { apiKey: result.apiKey, headers: result.headers } : {};
return result.ok ? { apiKey: result.apiKey, headers: result.headers, env: result.env } : {};
}
/**
@@ -1649,7 +1660,7 @@ export class AgentSession {
throw new Error(formatNoModelSelectedMessage());
}
const { apiKey, headers } = await this._getCompactionRequestAuth(this.model);
const { apiKey, headers, env } = await this._getCompactionRequestAuth(this.model);
const pathEntries = this.sessionManager.getBranch();
const settings = this.settingsManager.getCompactionSettings();
@@ -1673,6 +1684,8 @@ export class AgentSession {
preparation,
branchEntries: pathEntries,
customInstructions,
reason: "manual",
willRetry: false,
signal: this._compactionAbortController.signal,
})) as SessionBeforeCompactResult | undefined;
@@ -1708,6 +1721,7 @@ export class AgentSession {
this._compactionAbortController.signal,
this.thinkingLevel,
this.agent.streamFn,
env,
);
summary = result.summary;
firstKeptEntryId = result.firstKeptEntryId;
@@ -1723,6 +1737,7 @@ export class AgentSession {
const newEntries = this.sessionManager.getEntries();
const sessionContext = this.sessionManager.buildSessionContext();
this.agent.state.messages = sessionContext.messages;
const estimatedTokensAfter = estimateMessagesTokens(sessionContext.messages);
// Get the saved compaction entry for the extension event
const savedCompactionEntry = newEntries.find((e) => e.type === "compaction" && e.summary === summary) as
@@ -1734,13 +1749,16 @@ export class AgentSession {
type: "session_compact",
compactionEntry: savedCompactionEntry,
fromExtension,
reason: "manual",
willRetry: false,
});
}
const compactionResult = {
const compactionResult: CompactionResult = {
summary,
firstKeptEntryId,
tokensBefore,
estimatedTokensAfter,
details,
};
this._emit({
@@ -1821,8 +1839,17 @@ export class AgentSession {
return false;
}
// Case 1: Overflow - LLM returned context overflow error
// Case 1: Overflow - LLM returned context overflow error, or reported usage exceeded
// the configured window. A successful response over the configured window should compact
// but must not retry: the assistant answer already completed and agent.continue() cannot
// continue from an assistant message.
if (sameModel && isContextOverflow(assistantMessage, contextWindow)) {
const willRetry = assistantMessage.stopReason !== "stop";
if (!willRetry) {
return await this._runAutoCompaction("overflow", false);
}
if (this._overflowRecoveryAttempted) {
this._emit({
type: "compaction_end",
@@ -1843,7 +1870,7 @@ export class AgentSession {
if (messages.length > 0 && messages[messages.length - 1].role === "assistant") {
this.agent.state.messages = messages.slice(0, -1);
}
return await this._runAutoCompaction("overflow", true);
return await this._runAutoCompaction("overflow", willRetry);
}
// Case 2: Threshold - context is getting large
@@ -1880,56 +1907,39 @@ export class AgentSession {
*/
private async _runAutoCompaction(reason: "overflow" | "threshold", willRetry: boolean): Promise<boolean> {
const settings = this.settingsManager.getCompactionSettings();
this._emit({ type: "compaction_start", reason });
this._autoCompactionAbortController = new AbortController();
let started = false;
try {
if (!this.model) {
this._emit({
type: "compaction_end",
reason,
result: undefined,
aborted: false,
willRetry: false,
});
return false;
}
let apiKey: string | undefined;
let headers: Record<string, string> | undefined;
let env: Record<string, string> | undefined;
if (this.agent.streamFn === streamSimple) {
const authResult = await this._modelRegistry.getApiKeyAndHeaders(this.model);
if (!authResult.ok || !authResult.apiKey) {
this._emit({
type: "compaction_end",
reason,
result: undefined,
aborted: false,
willRetry: false,
});
return false;
}
apiKey = authResult.apiKey;
headers = authResult.headers;
env = authResult.env;
} else {
({ apiKey, headers } = await this._getCompactionRequestAuth(this.model));
({ apiKey, headers, env } = await this._getCompactionRequestAuth(this.model));
}
const pathEntries = this.sessionManager.getBranch();
const preparation = prepareCompaction(pathEntries, settings);
if (!preparation) {
this._emit({
type: "compaction_end",
reason,
result: undefined,
aborted: false,
willRetry: false,
});
return false;
}
this._emit({ type: "compaction_start", reason });
this._autoCompactionAbortController = new AbortController();
started = true;
let extensionCompaction: CompactionResult | undefined;
let fromExtension = false;
@@ -1939,6 +1949,8 @@ export class AgentSession {
preparation,
branchEntries: pathEntries,
customInstructions: undefined,
reason,
willRetry,
signal: this._autoCompactionAbortController.signal,
})) as SessionBeforeCompactResult | undefined;
@@ -1981,6 +1993,7 @@ export class AgentSession {
this._autoCompactionAbortController.signal,
this.thinkingLevel,
this.agent.streamFn,
env,
);
summary = compactResult.summary;
firstKeptEntryId = compactResult.firstKeptEntryId;
@@ -2003,6 +2016,7 @@ export class AgentSession {
const newEntries = this.sessionManager.getEntries();
const sessionContext = this.sessionManager.buildSessionContext();
this.agent.state.messages = sessionContext.messages;
const estimatedTokensAfter = estimateMessagesTokens(sessionContext.messages);
// Get the saved compaction entry for the extension event
const savedCompactionEntry = newEntries.find((e) => e.type === "compaction" && e.summary === summary) as
@@ -2014,6 +2028,8 @@ export class AgentSession {
type: "session_compact",
compactionEntry: savedCompactionEntry,
fromExtension,
reason,
willRetry,
});
}
@@ -2021,6 +2037,7 @@ export class AgentSession {
summary,
firstKeptEntryId,
tokensBefore,
estimatedTokensAfter,
details,
};
this._emit({ type: "compaction_end", reason, result, aborted: false, willRetry });
@@ -2039,17 +2056,19 @@ export class AgentSession {
return this.agent.hasQueuedMessages();
} catch (error) {
const errorMessage = error instanceof Error ? error.message : "compaction failed";
this._emit({
type: "compaction_end",
reason,
result: undefined,
aborted: false,
willRetry: false,
errorMessage:
reason === "overflow"
? `Context overflow recovery failed: ${errorMessage}`
: `Auto-compaction failed: ${errorMessage}`,
});
if (started) {
this._emit({
type: "compaction_end",
reason,
result: undefined,
aborted: false,
willRetry: false,
errorMessage:
reason === "overflow"
? `Context overflow recovery failed: ${errorMessage}`
: `Auto-compaction failed: ${errorMessage}`,
});
}
return false;
} finally {
this._autoCompactionAbortController = undefined;
@@ -2432,7 +2451,7 @@ export class AgentSession {
});
}
async reload(): Promise<void> {
async reload(options?: { beforeSessionStart?: () => void | Promise<void> }): Promise<void> {
const previousFlagValues = this._extensionRunner.getFlagValues();
await emitSessionShutdownEvent(this._extensionRunner, { type: "session_shutdown", reason: "reload" });
await this.settingsManager.reload();
@@ -2451,6 +2470,7 @@ export class AgentSession {
this._extensionShutdownHandler ||
this._extensionErrorListener;
if (hasBindings) {
await options?.beforeSessionStart?.();
await this._extensionRunner.emit({ type: "session_start", reason: "reload" });
await this.extendResourcesFromExtensions("reload");
}
@@ -2784,12 +2804,13 @@ export class AgentSession {
let summaryDetails: unknown;
if (options.summarize && entriesToSummarize.length > 0 && !extensionSummary) {
const model = this.model!;
const { apiKey, headers } = await this._getRequiredRequestAuth(model);
const { apiKey, headers, env } = await this._getRequiredRequestAuth(model);
const branchSummarySettings = this.settingsManager.getBranchSummarySettings();
const result = await generateBranchSummary(entriesToSummarize, {
model,
apiKey,
headers,
env,
signal: this._branchSummaryAbortController.signal,
customInstructions,
replaceInstructions,
@@ -3017,7 +3038,8 @@ export class AgentSession {
* @returns Path to exported file
*/
async exportToHtml(outputPath?: string): Promise<string> {
const themeName = this.settingsManager.getTheme();
const configuredThemeName = this.settingsManager.getTheme();
const themeName = configuredThemeName && getThemeByName(configuredThemeName) ? configuredThemeName : undefined;
// Create tool renderer if we have an extension runner (for custom tool HTML rendering)
const toolRenderer: ToolHtmlRenderer = createToolHtmlRenderer({