refactor(agent): consolidate harness session abstraction

packages/agent/src/harness/compaction/branch-summarization.ts

@@ -5,7 +5,7 @@
  * a summary of the branch being left so context isn't lost.
  */
 
-import type { Model } from "@mariozechner/pi-ai";
+import type { ImageContent, Model, TextContent } from "@mariozechner/pi-ai";
 import { completeSimple } from "@mariozechner/pi-ai";
 import type { AgentMessage } from "../../types.js";
 import {
@@ -14,7 +14,7 @@ import {
 	createCompactionSummaryMessage,
 	createCustomMessage,
 } from "../messages.js";
-import type { SessionTree, SessionTreeEntry } from "../types.js";
+import type { Session, SessionTreeEntry } from "../types.js";
 import { estimateTokens } from "./compaction.js";
 import {
 	computeFileLists,
@@ -96,7 +96,7 @@ export interface GenerateBranchSummaryOptions {
  * @returns Entries to summarize and the common ancestor
  */
 export async function collectEntriesForBranchSummary(
-	session: SessionTree,
+	session: Session,
 	oldLeafId: string | null,
 	targetId: string,
 ): Promise<CollectEntriesResult> {
@@ -125,7 +125,7 @@ export async function collectEntriesForBranchSummary(
 	while (current && current !== commonAncestorId) {
 		const entry = await session.getEntry(current);
 		if (!entry) break;
-		entries.push(entry);
+		entries.push(entry as SessionTreeEntry);
 		current = entry.parentId;
 	}
 
@@ -148,10 +148,16 @@ function getMessageFromEntry(entry: SessionTreeEntry): AgentMessage | undefined
 		case "message":
 			// Skip tool results - context is in assistant's tool call
 			if (entry.message.role === "toolResult") return undefined;
-			return entry.message;
+			return entry.message as AgentMessage;
 
 		case "custom_message":
-			return createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp);
+			return createCustomMessage(
+				entry.customType,
+				entry.content as string | (TextContent | ImageContent)[],
+				entry.display,
+				entry.details,
+				entry.timestamp,
+			);
 
 		case "branch_summary":
 			return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp);

packages/agent/src/harness/compaction/compaction.ts

@@ -5,7 +5,7 @@
  * and after compaction the session is reloaded.
  */
 
-import type { AssistantMessage, Model, Usage } from "@mariozechner/pi-ai";
+import type { AssistantMessage, ImageContent, Model, TextContent, Usage } from "@mariozechner/pi-ai";
 import { completeSimple } from "@mariozechner/pi-ai";
 import type { AgentMessage, ThinkingLevel } from "../../types.js";
 import {
@@ -79,10 +79,16 @@ function extractFileOperations(
  */
 function getMessageFromEntry(entry: SessionTreeEntry): AgentMessage | undefined {
 	if (entry.type === "message") {
-		return entry.message;
+		return entry.message as AgentMessage;
 	}
 	if (entry.type === "custom_message") {
-		return createCustomMessage(entry.customType, entry.content, entry.display, entry.details, entry.timestamp);
+		return createCustomMessage(
+			entry.customType,
+			entry.content as string | (TextContent | ImageContent)[],
+			entry.display,
+			entry.details,
+			entry.timestamp,
+		);
 	}
 	if (entry.type === "branch_summary") {
 		return createBranchSummaryMessage(entry.summary, entry.fromId, entry.timestamp);
@@ -158,7 +164,7 @@ export function getLastAssistantUsage(entries: SessionTreeEntry[]): Usage | unde
 	for (let i = entries.length - 1; i >= 0; i--) {
 		const entry = entries[i];
 		if (entry.type === "message") {
-			const usage = getAssistantUsage(entry.message);
+			const usage = getAssistantUsage(entry.message as AgentMessage);
 			if (usage) return usage;
 		}
 	}
@@ -405,7 +411,7 @@ export function findCutPoint(
 		if (entry.type !== "message") continue;
 
 		// Estimate this message's size
-		const messageTokens = estimateTokens(entry.message);
+		const messageTokens = estimateTokens(entry.message as AgentMessage);
 		accumulatedTokens += messageTokens;
 
 		// Check if we've exceeded the budget