Skip to content
Merged
2 changes: 2 additions & 0 deletions packages/agent/src/agent-loop.ts
Original file line number Diff line number Diff line change
Expand Up @@ -737,6 +737,7 @@ async function finalizeExecutedToolCall(
...result,
content: afterResult.content ?? result.content,
details: afterResult.details ?? result.details,
usage: afterResult.usage ?? result.usage,
terminate: afterResult.terminate ?? result.terminate,
};
isError = afterResult.isError ?? isError;
Expand Down Expand Up @@ -780,6 +781,7 @@ function createToolResultMessage(finalized: FinalizedToolCallOutcome): ToolResul
// so the null never enters session history or provider payloads.
content: finalized.result.content ?? [],
details: finalized.result.details,
usage: finalized.result.usage,
...(finalized.result.addedToolNames?.length ? { addedToolNames: finalized.result.addedToolNames } : {}),
isError: finalized.isError,
timestamp: Date.now(),
Expand Down
24 changes: 19 additions & 5 deletions packages/agent/src/harness/agent-harness.ts
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ import type {
AgentHarnessResources,
AgentHarnessStreamOptions,
AgentHarnessStreamOptionsPatch,
CompactResult,
ExecutionEnv,
NavigateTreeResult,
PendingSessionWrite,
Expand Down Expand Up @@ -434,9 +435,16 @@ export class AgentHarness<
content: result.content,
details: result.details,
isError,
usage: result.usage,
});
return patch
? { content: patch.content, details: patch.details, isError: patch.isError, terminate: patch.terminate }
? {
content: patch.content,
details: patch.details,
isError: patch.isError,
usage: patch.usage,
terminate: patch.terminate,
}
: undefined;
},
prepareNextTurn: async () => {
Expand Down Expand Up @@ -690,9 +698,7 @@ export class AgentHarness<
}
}

async compact(
customInstructions?: string,
): Promise<{ summary: string; firstKeptEntryId: string; tokensBefore: number; details?: unknown }> {
async compact(customInstructions?: string): Promise<CompactResult> {
if (this.phase !== "idle") throw new AgentHarnessError("busy", "compact() requires idle harness");
this.phase = "compaction";
try {
Expand Down Expand Up @@ -723,6 +729,7 @@ export class AgentHarness<
result.tokensBefore,
result.details,
provided !== undefined,
result.usage,
);
const entry = await this.session.getEntry(entryId);
if (entry?.type === "compaction") {
Expand Down Expand Up @@ -764,6 +771,7 @@ export class AgentHarness<
let summaryEntry: NavigateTreeResult["summaryEntry"];
let summaryText: string | undefined = hookResult?.summary?.summary;
let summaryDetails: unknown = hookResult?.summary?.details;
let summaryUsage = hookResult?.summary?.usage;
if (!summaryText && options?.summarize && entries.length > 0) {
const model = this.model;
if (!model) throw new AgentHarnessError("invalid_state", "No model set for branch summary");
Expand All @@ -779,6 +787,7 @@ export class AgentHarness<
throw new AgentHarnessError("branch_summary", branchSummary.error.message, branchSummary.error);
}
summaryText = branchSummary.value.summary;
summaryUsage = branchSummary.value.usage;
summaryDetails = {
readFiles: branchSummary.value.readFiles,
modifiedFiles: branchSummary.value.modifiedFiles,
Expand All @@ -798,7 +807,12 @@ export class AgentHarness<
const summaryId = await this.session.moveTo(
newLeafId,
summaryText
? { summary: summaryText, details: summaryDetails, fromHook: hookResult?.summary !== undefined }
? {
summary: summaryText,
details: summaryDetails,
usage: summaryUsage,
fromHook: hookResult?.summary !== undefined,
}
: undefined,
);
if (summaryId) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -252,6 +252,7 @@ export async function generateBranchSummary(

return ok({
summary: summary || "No summary generated",
usage: response.usage,
readFiles,
modifiedFiles,
});
Expand Down
77 changes: 57 additions & 20 deletions packages/agent/src/harness/compaction/compaction.ts
Original file line number Diff line number Diff line change
Expand Up @@ -101,10 +101,35 @@ export interface CompactionResult<T = unknown> {
firstKeptEntryId: string;
/** Estimated context tokens before compaction. */
tokensBefore: number;
/** Usage from the LLM call(s) that generated this summary, if available. */
usage?: Usage;
/** Optional implementation-specific details stored with the compaction entry. */
details?: T;
}

function combineUsage(first: Usage, second: Usage): Usage {
return {
input: first.input + second.input,
output: first.output + second.output,
cacheRead: first.cacheRead + second.cacheRead,
cacheWrite: first.cacheWrite + second.cacheWrite,
...(first.cacheWrite1h !== undefined || second.cacheWrite1h !== undefined
? { cacheWrite1h: (first.cacheWrite1h ?? 0) + (second.cacheWrite1h ?? 0) }
: {}),
...(first.reasoning !== undefined || second.reasoning !== undefined
? { reasoning: (first.reasoning ?? 0) + (second.reasoning ?? 0) }
: {}),
totalTokens: first.totalTokens + second.totalTokens,
cost: {
input: first.cost.input + second.cost.input,
output: first.cost.output + second.cost.output,
cacheRead: first.cost.cacheRead + second.cost.cacheRead,
cacheWrite: first.cost.cacheWrite + second.cost.cacheWrite,
total: first.cost.total + second.cost.total,
},
};
}

/** Compaction thresholds and retention settings. */
export interface CompactionSettings {
/** Enable automatic compaction decisions. */
Expand Down Expand Up @@ -474,7 +499,7 @@ export async function generateSummary(
customInstructions?: string,
previousSummary?: string,
thinkingLevel?: ThinkingLevel,
): Promise<Result<string, CompactionError>> {
): Promise<Result<{ text: string; usage: Usage }, CompactionError>> {
const maxTokens = Math.min(
Math.floor(0.8 * reserveTokens),
model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY,
Expand Down Expand Up @@ -523,7 +548,7 @@ export async function generateSummary(

const textContent = contentText(response.content);

return ok(textContent);
return ok({ text: textContent, usage: response.usage });
}

/** Prepared inputs for a compaction run. */
Expand Down Expand Up @@ -656,22 +681,26 @@ export async function compact(
}

let summary: string;
let summaryUsage: Usage;

if (isSplitTurn && turnPrefixMessages.length > 0) {
const historyResult =
messagesToSummarize.length > 0
? await generateSummary(
messagesToSummarize,
models,
model,
settings.reserveTokens,
signal,
customInstructions,
previousSummary,
thinkingLevel,
)
: ok<string, CompactionError>("No prior history.");
if (!historyResult.ok) return err(historyResult.error);
let historyText = "No prior history.";
let historyUsage: Usage | undefined;
if (messagesToSummarize.length > 0) {
const historyResult = await generateSummary(
messagesToSummarize,
models,
model,
settings.reserveTokens,
signal,
customInstructions,
previousSummary,
thinkingLevel,
);
if (!historyResult.ok) return err(historyResult.error);
historyText = historyResult.value.text;
historyUsage = historyResult.value.usage;
}
const turnPrefixResult = await generateTurnPrefixSummary(
turnPrefixMessages,
models,
Expand All @@ -681,7 +710,10 @@ export async function compact(
thinkingLevel,
);
if (!turnPrefixResult.ok) return err(turnPrefixResult.error);
summary = `${historyResult.value}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult.value}`;
summary = `${historyText}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult.value.text}`;
summaryUsage = historyUsage
? combineUsage(historyUsage, turnPrefixResult.value.usage)
: turnPrefixResult.value.usage;
} else {
const summaryResult = await generateSummary(
messagesToSummarize,
Expand All @@ -694,7 +726,8 @@ export async function compact(
thinkingLevel,
);
if (!summaryResult.ok) return err(summaryResult.error);
summary = summaryResult.value;
summary = summaryResult.value.text;
summaryUsage = summaryResult.value.usage;
}

const { readFiles, modifiedFiles } = computeFileLists(fileOps);
Expand All @@ -704,6 +737,7 @@ export async function compact(
summary,
firstKeptEntryId,
tokensBefore,
usage: summaryUsage,
details: { readFiles, modifiedFiles } as CompactionDetails,
});
}
Expand All @@ -714,7 +748,7 @@ async function generateTurnPrefixSummary(
reserveTokens: number,
signal?: AbortSignal,
thinkingLevel?: ThinkingLevel,
): Promise<Result<string, CompactionError>> {
): Promise<Result<{ text: string; usage: Usage }, CompactionError>> {
const maxTokens = Math.min(
Math.floor(0.5 * reserveTokens),
model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY,
Expand Down Expand Up @@ -749,5 +783,8 @@ async function generateTurnPrefixSummary(
);
}

return ok(contentText(response.content));
return ok({
text: contentText(response.content),
usage: response.usage,
});
}
7 changes: 5 additions & 2 deletions packages/agent/src/harness/session/session.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import type { ImageContent, TextContent } from "@earendil-works/pi-ai";
import type { ImageContent, TextContent, Usage } from "@earendil-works/pi-ai";
import type { AgentMessage } from "../../types.ts";
import { createBranchSummaryMessage, createCompactionSummaryMessage, createCustomMessage } from "../messages.ts";
import type {
Expand Down Expand Up @@ -247,6 +247,7 @@ export class Session<TMetadata extends SessionMetadata = SessionMetadata> {
tokensBefore: number,
details?: T,
fromHook?: boolean,
usage?: Usage,
): Promise<string> {
return this.appendTypedEntry({
type: "compaction",
Expand All @@ -257,6 +258,7 @@ export class Session<TMetadata extends SessionMetadata = SessionMetadata> {
firstKeptEntryId,
tokensBefore,
details,
usage,
fromHook,
} satisfies CompactionEntry<T>);
}
Expand Down Expand Up @@ -317,7 +319,7 @@ export class Session<TMetadata extends SessionMetadata = SessionMetadata> {

async moveTo(
entryId: string | null,
summary?: { summary: string; details?: unknown; fromHook?: boolean },
summary?: { summary: string; details?: unknown; usage?: Usage; fromHook?: boolean },
): Promise<string | undefined> {
if (entryId !== null && !(await this.storage.getEntry(entryId))) {
throw new SessionError("not_found", `Entry ${entryId} not found`);
Expand All @@ -332,6 +334,7 @@ export class Session<TMetadata extends SessionMetadata = SessionMetadata> {
fromId: entryId ?? "root",
summary: summary.summary,
details: summary.details,
usage: summary.usage,
fromHook: summary.fromHook,
} satisfies BranchSummaryEntry);
}
Expand Down
24 changes: 22 additions & 2 deletions packages/agent/src/harness/types.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,12 @@
import type { ImageContent, Model, Models, SimpleStreamOptions, TextContent, Transport } from "@earendil-works/pi-ai";
import type {
ImageContent,
Model,
Models,
SimpleStreamOptions,
TextContent,
Transport,
Usage,
} from "@earendil-works/pi-ai";
import type { AgentEvent, AgentMessage, AgentTool, QueueMode, ThinkingLevel } from "../index.ts";
import type { Session } from "./session/session.ts";

Expand Down Expand Up @@ -365,6 +373,7 @@ export interface CompactionEntry<T = unknown> extends SessionTreeEntryBase {
firstKeptEntryId: string;
tokensBefore: number;
details?: T;
usage?: Usage;
fromHook?: boolean;
}

Expand All @@ -373,6 +382,7 @@ export interface BranchSummaryEntry<T = unknown> extends SessionTreeEntryBase {
fromId: string;
summary: string;
details?: T;
usage?: Usage;
fromHook?: boolean;
}

Expand Down Expand Up @@ -572,6 +582,7 @@ export interface ToolResultEvent {
content: Array<TextContent | ImageContent>;
details: unknown;
isError: boolean;
usage?: Usage;
}

export interface SessionBeforeCompactEvent {
Expand Down Expand Up @@ -687,6 +698,7 @@ export interface ToolResultPatch {
content?: Array<TextContent | ImageContent>;
details?: unknown;
isError?: boolean;
usage?: Usage;
terminate?: boolean;
}

Expand All @@ -697,7 +709,12 @@ export interface SessionBeforeCompactResult {

export interface SessionBeforeTreeResult {
cancel?: boolean;
summary?: { summary: string; details?: unknown };
summary?: {
summary: string;
details?: unknown;
/** Usage from the LLM call that generated this summary, if available. */
usage?: Usage;
};
customInstructions?: string;
replaceInstructions?: boolean;
label?: string;
Expand Down Expand Up @@ -738,6 +755,8 @@ export interface CompactResult {
summary: string;
firstKeptEntryId: string;
tokensBefore: number;
/** Usage from the LLM call(s) that generated this summary, if available. */
usage?: Usage;
details?: unknown;
}

Expand Down Expand Up @@ -793,6 +812,7 @@ export interface GenerateBranchSummaryOptions {

export interface BranchSummaryResult {
summary: string;
usage?: Usage;
readFiles: string[];
modifiedFiles: string[];
}
Expand Down
Loading
Loading