Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions packages/agent/src/adapters/error-classification.ts
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import { getErrorMessage } from "@posthog/shared";

export type AgentErrorClassification =
| "upstream_stream_terminated"
| "upstream_connection_error"
Expand Down Expand Up @@ -32,3 +34,8 @@ export function classifyAgentError(
}
return "agent_error";
}

/** Hard API rejection: the assembled prompt exceeds the model's context window. */
export function isPromptTooLongError(error: unknown): boolean {
return /prompt is too long/i.test(getErrorMessage(error));
}
92 changes: 92 additions & 0 deletions packages/agent/src/server/agent-server.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1808,6 +1808,98 @@ describe("AgentServer HTTP Mode", () => {
});

describe("native resume", () => {
it.each([
{ retryOutcome: "succeeds", retryFails: false },
{ retryOutcome: "fails", retryFails: true },
])(
"clears resume state when the fresh-session retry $retryOutcome",
async ({ retryFails }) => {
const s = createServer();
await s.start();

const prompts: ContentBlock[][] = [];
const prompt = vi.fn(async (params: { prompt: ContentBlock[] }) => {
prompts.push(params.prompt);
if (prompts.length === 1) {
throw new Error("Internal error: Prompt is too long");
}
if (retryFails) {
throw new Error("Fresh-session retry failed");
}
return { stopReason: "end_turn" };
});
const newSession = vi.fn(async () => ({ sessionId: "fresh-session" }));

const internals = s as unknown as {
session: {
acpSessionId: string;
clientConnection: {
prompt: typeof prompt;
newSession: typeof newSession;
};
};
resumeState: ResumeState | null;
nativeResume: { sessionId: string; warm: boolean } | null;
loadResumeState(
taskId: string,
resumeRunId: string,
runId: string,
): Promise<void>;
sendResumeContinuation(
payload: JwtPayload,
taskRun: TaskRun | null,
): Promise<void>;
};
internals.session.clientConnection.prompt = prompt;
internals.session.clientConnection.newSession = newSession;
internals.nativeResume = { sessionId: "prior-session", warm: true };
internals.loadResumeState = vi.fn(async () => {
internals.resumeState = {
conversation: [
{
role: "user",
content: [{ type: "text", text: "original task" }],
},
{
role: "assistant",
content: [{ type: "text", text: "progress so far" }],
},
],
latestGitCheckpoint: null,
interrupted: false,
logEntryCount: 2,
sessionId: "prior-session",
};
});

await internals.sendResumeContinuation(
{
task_id: "test-task-id",
run_id: "test-run-id",
team_id: 1,
user_id: 1,
distinct_id: "test-distinct-id",
mode: "interactive",
},
createTaskRun({
id: "test-run-id",
state: { resume_from_run_id: "previous-run" },
}),
);

expect(newSession).toHaveBeenCalledOnce();
expect(internals.session.acpSessionId).toBe("fresh-session");
expect(internals.resumeState).toBeNull();
expect(internals.nativeResume).toBeNull();
expect(prompts).toHaveLength(2);
const retryText = prompts[1]
.map((block) => ("text" in block ? block.text : ""))
.join("\n");
expect(retryText).toContain("progress so far");
},
20000,
);

it("hydrates cold sessions from S3 logs instead of cached resume conversation", async () => {
const originalConfigDir = process.env.CLAUDE_CONFIG_DIR;
process.env.CLAUDE_CONFIG_DIR = join(repo.path, ".claude-test");
Expand Down
83 changes: 80 additions & 3 deletions packages/agent/src/server/agent-server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ import { getCurrentBranch } from "@posthog/git/queries";
import {
type Adapter,
buildPrOutput,
getErrorMessage,
mergePrUrls,
readPrUrls,
} from "@posthog/shared";
Expand All @@ -40,6 +41,7 @@ import { hasCodexThreadState } from "../adapters/codex-app-server/thread-state";
import {
type AgentErrorClassification,
classifyAgentError,
isPromptTooLongError,
} from "../adapters/error-classification";
import {
SIGNED_COMMIT_QUALIFIED_TOOL_NAME,
Expand Down Expand Up @@ -249,6 +251,8 @@ interface ActiveSession {
/** Whether a desktop client has ever connected via SSE during this session */
hasDesktopConnected: boolean;
pendingHandoffGitState?: HandoffLocalGitState;
/** Meta the session was created with, reused when a retry needs a fresh session */
sessionMeta: Record<string, unknown>;
}

interface InstalledSkillBundle {
Expand Down Expand Up @@ -328,6 +332,7 @@ export class AgentServer {
private lastReportedBranch: string | null = null;
private resumeState: ResumeState | null = null;
private nativeResume: { sessionId: string; warm: boolean } | null = null;
private oversizedResumeRetried = false;
// Prewarmed runs boot before the user's first message exists, so the boot-time
// --autoPublish flag can't carry the user's choice; it is resolved from run
// state when the first message arrives (see resolveWarmAutoPublishUpgrade).
Expand Down Expand Up @@ -1409,6 +1414,7 @@ export class AgentServer {
permissionMode: initialPermissionMode,
hasDesktopConnected: sseController !== null,
pendingHandoffGitState: undefined,
sessionMeta,
};

this.logger = new Logger({
Expand Down Expand Up @@ -1745,7 +1751,6 @@ export class AgentServer {
gitCheckpointBranch: resumeState.latestGitCheckpoint?.branch ?? null,
});

this.resumeState = null;
return {
prompt: resumePromptBlocks,
...(resumePromptMeta ? { meta: resumePromptMeta } : {}),
Expand Down Expand Up @@ -1786,21 +1791,82 @@ export class AgentServer {
hasPendingUserMessage: !!pendingUserPrompt?.prompt.length,
});

this.resumeState = null;
this.nativeResume = null;
return {
prompt,
...(pendingUserPrompt?.meta ? { meta: pendingUserPrompt.meta } : {}),
};
},
{ retryOnOversizedPrompt: true },
);
}

/**
* A native resume replays the prior transcript verbatim; when that
* transcript no longer fits the context window, every request (including
* auto-compaction) is rejected, so the only way forward is a fresh session
* seeded with the summarized history the non-native resume path uses.
*/
private async retryOversizedResumeOnFreshSession(
payload: JwtPayload,
taskRun: TaskRun | null,
): Promise<boolean> {
if (this.oversizedResumeRetried || !this.session) {
return false;
}
this.oversizedResumeRetried = true;

const resumeRunId = this.getResumeRunId(taskRun);
if (!resumeRunId) return false;
if (!this.resumeState) {
try {
await this.loadResumeState(
payload.task_id,
resumeRunId,
payload.run_id,
);
} catch (error) {
this.logger.warn("Failed to reload resume state for retry", {
error: getErrorMessage(error),
});
return false;
}
}
if (!this.resumeState?.conversation.length) return false;

this.logger.warn(
"Resume prompt exceeded the context window; retrying on a fresh session with summarized history",
{ taskId: payload.task_id, runId: payload.run_id },
);

try {
const response = await this.session.clientConnection.newSession({
cwd: this.config.repositoryPath ?? "/tmp/workspace",
mcpServers: this.config.mcpServers ?? [],
_meta: this.session.sessionMeta,
});
this.session.acpSessionId = response.sessionId;
} catch (error) {
this.logger.warn("Failed to start fresh session for oversized resume", {
error: getErrorMessage(error),
});
return false;
}

try {
await this.sendResumeMessage(payload, taskRun);
return true;
} finally {
this.resumeState = null;
this.nativeResume = null;
}
}

private async runResumeTurn(
payload: JwtPayload,
taskRun: TaskRun | null,
logLabel: string,
buildPrompt: () => Promise<BuiltPrompt>,
opts: { retryOnOversizedPrompt?: boolean } = {},
): Promise<void> {
if (!this.session) return;

Expand All @@ -1819,6 +1885,10 @@ export class AgentServer {
stopReason: result.stopReason,
});

// Kept until the turn succeeds so a prompt-too-long retry can reuse it.
this.resumeState = null;
this.nativeResume = null;

await this.clearPendingInitialPromptState(payload, taskRun);

if (result.stopReason === "end_turn") {
Expand All @@ -1836,6 +1906,13 @@ export class AgentServer {
if (this.session) {
await this.session.logWriter.flushAll();
}
if (
opts.retryOnOversizedPrompt &&
isPromptTooLongError(error) &&
(await this.retryOversizedResumeOnFreshSession(payload, taskRun))
) {
return;
}
await this.handleTurnFailure(payload, "resume", error);
}
}
Expand Down
Loading