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
121 changes: 121 additions & 0 deletions packages/agent/src/server/agent-server.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ import {
it,
vi,
} from "vitest";
import { POSTHOG_NOTIFICATIONS } from "../acp-extensions";
import { getSessionJsonlPath } from "../adapters/claude/session/jsonl-hydration";
import type { PermissionMode } from "../execution-mode";
import type { PostHogAPIClient } from "../posthog-api";
Expand Down Expand Up @@ -1433,6 +1434,126 @@ describe("AgentServer HTTP Mode", () => {
expect(body.error).toBe("No active session for this run");
}, 20000);

it("continues a cloud task after a manual compact command", async () => {
const s = createServer();
await s.start();
const broadcastEvent = vi.fn();
let serverInternals!: {
session: { clientConnection: { prompt: typeof prompt } };
broadcastEvent: typeof broadcastEvent;
handleAcpTransportMessage(message: unknown): void;
};
const prompt = vi.fn(async (_params: { prompt: ContentBlock[] }) => {
serverInternals.handleAcpTransportMessage({
jsonrpc: "2.0",
method: POSTHOG_NOTIFICATIONS.TURN_COMPLETE,
params: { sessionId: "session-1", stopReason: "end_turn" },
});
return { stopReason: "end_turn" };
});
serverInternals = s as unknown as typeof serverInternals;
serverInternals.session.clientConnection.prompt = prompt;
serverInternals.broadcastEvent = broadcastEvent;

const token = createToken();
const response = await fetch(`http://localhost:${port}/command`, {
method: "POST",
headers: {
Authorization: `Bearer ${token}`,
"Content-Type": "application/json",
},
body: JSON.stringify({
jsonrpc: "2.0",
id: "compact-and-continue",
method: "user_message",
params: {
content:
"/compact Continue with the task using the question tool and plan.",
},
}),
});

expect(response.status).toBe(200);
const body = (await response.json()) as {
result?: { stopReason?: string };
};
expect(body.result?.stopReason).toBe("end_turn");
expect(prompt).toHaveBeenCalledTimes(2);
expect(prompt.mock.calls[0]?.[0].prompt).toEqual([
{
type: "text",
text: "/compact Continue with the task using the question tool and plan.",
},
]);
expect(prompt.mock.calls[1]?.[0].prompt).toEqual([
{
type: "text",
text: expect.stringContaining("Continue working on the task"),
_meta: { ui: { hidden: true } },
},
]);
const turnCompleteEvents = broadcastEvent.mock.calls.filter(
([event]) =>
(event as { notification?: { method?: string } }).notification
?.method === POSTHOG_NOTIFICATIONS.TURN_COMPLETE,
);
expect(turnCompleteEvents).toHaveLength(1);
}, 20000);

it("retries only the continuation after compact follow-up failure", async () => {
const s = createServer();
await s.start();
const prompt = vi
.fn(async (_params: { prompt: ContentBlock[] }) => ({
stopReason: "end_turn",
}))
.mockResolvedValueOnce({ stopReason: "end_turn" })
.mockRejectedValueOnce(new Error("sdk connection lost"));
const serverInternals = s as unknown as {
session: { clientConnection: { prompt: typeof prompt } };
};
serverInternals.session.clientConnection.prompt = prompt;

const token = createToken();
const send = async () =>
fetch(`http://localhost:${port}/command`, {
method: "POST",
headers: {
Authorization: `Bearer ${token}`,
"Content-Type": "application/json",
},
body: JSON.stringify({
jsonrpc: "2.0",
id: "compact-retry",
method: "user_message",
params: {
content: "/compact Continue the task.",
messageId: "compact-retry",
},
}),
});

const first = await send();
expect(first.status).toBe(200);
expect(prompt).toHaveBeenCalledTimes(2);

const retry = await send();
expect(retry.status).toBe(200);
expect(prompt).toHaveBeenCalledTimes(3);
expect(prompt.mock.calls[0]?.[0].prompt[0]).toMatchObject({
type: "text",
text: "/compact Continue the task.",
});
expect(prompt.mock.calls[1]?.[0].prompt[0]).toMatchObject({
type: "text",
text: expect.stringContaining("Continue working on the task"),
});
expect(prompt.mock.calls[2]?.[0].prompt[0]).toMatchObject({
type: "text",
text: expect.stringContaining("Continue working on the task"),
});
}, 20000);

it("rewrites a bundled local skill slash command before sending the prompt", async () => {
const skillDefinition = [
"---",
Expand Down
106 changes: 83 additions & 23 deletions packages/agent/src/server/agent-server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -289,6 +289,10 @@ function hiddenTextBlock(text: string): ContentBlock {
} as ContentBlock;
}

function isManualCompactPrompt(prompt: ContentBlock[]): boolean {
return /^\/compact(?:\s|$)/.test(promptBlocksToText(prompt).trimStart());
}

interface LocalSkillPromptContext {
/** Set when the message is a bare `/skill` invocation the adapter should strip. */
skillName?: string;
Expand Down Expand Up @@ -339,6 +343,7 @@ export class AgentServer {
private rtkSavingsAttempted = false;
private questionRelayedToSlack = false;
private adapterEmittedTurnComplete = false;
private suppressAdapterTurnComplete = false;
private runUsage = new RunUsageAccumulator();
private detectedPrUrl: string | null = null;
// Reset per session. `evaluatedPrUrls` dedupes per URL; `prAttributionChain` serializes
Expand All @@ -364,6 +369,7 @@ export class AgentServer {
private initializationPromise: Promise<void> | null = null;
private pendingEvents: Record<string, unknown>[] = [];
private deliveredMessageIds = new Set<string>();
private pendingCompactContinuationMessageIds = new Set<string>();
private pendingPermissions = new Map<
string,
{
Expand Down Expand Up @@ -909,18 +915,28 @@ export class AgentServer {
typeof params.messageId === "string" && params.messageId
? params.messageId
: undefined;
let retryCompactContinuation = false;
if (messageId) {
if (this.deliveredMessageIds.has(messageId)) {
this.logger.info("Duplicate user_message delivery ignored", {
messageId,
});
return { stopReason: "duplicate_delivery", duplicate: true };
if (this.pendingCompactContinuationMessageIds.has(messageId)) {
retryCompactContinuation = true;
this.logger.info("Retrying pending compact continuation", {
messageId,
});
} else {
this.logger.info("Duplicate user_message delivery ignored", {
messageId,
});
return { stopReason: "duplicate_delivery", duplicate: true };
}
} else {
this.deliveredMessageIds.add(messageId);
}
this.deliveredMessageIds.add(messageId);
if (this.deliveredMessageIds.size > 500) {
const oldest = this.deliveredMessageIds.values().next().value;
if (oldest !== undefined) {
this.deliveredMessageIds.delete(oldest);
this.pendingCompactContinuationMessageIds.delete(oldest);
}
}
}
Expand Down Expand Up @@ -951,17 +967,53 @@ export class AgentServer {
: {}),
};

const manualCompactPrompt = isManualCompactPrompt(prompt);
const acpSessionId = this.session.acpSessionId;
const continueAfterCompaction = (): Promise<PromptResponse> =>
this.promptWithUpstreamRetry({
sessionId: acpSessionId,
prompt: [
hiddenTextBlock(
"Compaction is complete. Continue working on the task from the compacted context, following the user's instructions from the /compact command.",
),
],
});

let compactCommandCompleted = retryCompactContinuation;
let result: PromptResponse;
this.suppressAdapterTurnComplete =
manualCompactPrompt || retryCompactContinuation;
try {
result = await this.session.clientConnection.prompt({
sessionId: this.session.acpSessionId,
prompt,
...(Object.keys(promptMeta).length > 0
? { _meta: promptMeta }
: {}),
});
if (retryCompactContinuation) {
result = await continueAfterCompaction();
if (messageId) {
this.pendingCompactContinuationMessageIds.delete(messageId);
}
} else {
result = await this.session.clientConnection.prompt({
sessionId: this.session.acpSessionId,
prompt,
...(Object.keys(promptMeta).length > 0
? { _meta: promptMeta }
: {}),
});

if (result.stopReason === "end_turn" && manualCompactPrompt) {
compactCommandCompleted = true;
if (messageId) {
this.pendingCompactContinuationMessageIds.add(messageId);
}
// `/compact` is an SDK-local command, so without a follow-up the
// cloud run reports completion before the model resumes the task.
this.recordTurnUsage(result.usage);
result = await continueAfterCompaction();
if (messageId) {
this.pendingCompactContinuationMessageIds.delete(messageId);
}
}
}
} catch (error) {
if (messageId) {
if (messageId && !compactCommandCompleted) {
this.deliveredMessageIds.delete(messageId);
}
await this.session.logWriter.flushAll();
Expand All @@ -974,6 +1026,8 @@ export class AgentServer {
throw error;
}
return { stopReason: "error_recoverable" };
} finally {
this.suppressAdapterTurnComplete = false;
}

this.logger.debug("User message completed", {
Expand Down Expand Up @@ -1299,16 +1353,8 @@ export class AgentServer {

// Tap both streams to broadcast all ACP messages via SSE (mimics local transport)
this.adapterEmittedTurnComplete = false;
const onAcpMessage = (message: unknown) => {
if (isTurnCompleteNotification(message)) {
this.adapterEmittedTurnComplete = true;
}
this.broadcastEvent({
type: "notification",
timestamp: new Date().toISOString(),
notification: message,
});
};
const onAcpMessage = (message: unknown) =>
this.handleAcpTransportMessage(message);

const tappedReadable = createTappedReadableStream(
acpConnection.clientStreams.readable as ReadableStream<Uint8Array>,
Expand Down Expand Up @@ -4045,6 +4091,20 @@ ${signedCommitInstructions}${prLinkInstructions}${shellEfficiencyInstructions}
});
}

private handleAcpTransportMessage(message: unknown): void {
if (isTurnCompleteNotification(message)) {
if (this.suppressAdapterTurnComplete) {
return;
}
this.adapterEmittedTurnComplete = true;
}
this.broadcastEvent({
type: "notification",
timestamp: new Date().toISOString(),
notification: message,
});
}

private broadcastTurnComplete(stopReason: string): void {
if (!this.session) return;
if (this.adapterEmittedTurnComplete) {
Expand Down
Loading