Skip to content

Commit f9a8396

Browse files
adam-birddogadam-birddog
committed
fix: prevent sessions_yield completion wake freeze
1 parent 9159013 commit f9a8396

6 files changed

Lines changed: 118 additions & 73 deletions

File tree

‎src/agents/pi-embedded-runner/run-state.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ export type EmbeddedPiQueueHandle = {
1111
queueMessage: (text: string, options?: EmbeddedPiQueueMessageOptions) => Promise<void>;
1212
isStreaming: () => boolean;
1313
isCompacting: () => boolean;
14+
isAcceptingMessages?: () => boolean;
1415
supportsTranscriptCommitWait?: boolean;
1516
cancel?: (reason?: "user_abort" | "restart" | "superseded") => void;
1617
abort: () => void;

‎src/agents/pi-embedded-runner/run/attempt.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1792,6 +1792,7 @@ export async function runEmbeddedAttempt(
17921792
onYield: (message) => {
17931793
yieldDetected = true;
17941794
yieldMessage = message;
1795+
yieldAbortInProgress = true;
17951796
queueYieldInterruptForSession?.();
17961797
runAbortController.abort("sessions_yield");
17971798
abortSessionForYield?.();
@@ -1945,6 +1946,7 @@ export async function runEmbeddedAttempt(
19451946
// Track sessions_yield tool invocation (callback pattern, like clientToolCallDetected)
19461947
let yieldDetected = false;
19471948
let yieldMessage: string | null = null;
1949+
let yieldAbortInProgress = false;
19481950
// Late-binding reference so onYield can abort the session (declared after tool creation)
19491951
let abortSessionForYield: (() => void) | null = null;
19501952
let queueYieldInterruptForSession: (() => void) | null = null;
@@ -3746,6 +3748,7 @@ export async function runEmbeddedAttempt(
37463748
},
37473749
isStreaming: () => activeSession.isStreaming,
37483750
isCompacting: () => subscription.isCompacting(),
3751+
isAcceptingMessages: () => !yieldAbortInProgress,
37493752
supportsTranscriptCommitWait: true,
37503753
sourceReplyDeliveryMode: params.sourceReplyDeliveryMode,
37513754
cancel: () => {

‎src/agents/pi-embedded-runner/runs.test.ts‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,13 +41,17 @@ function createRunHandle(
4141
isCompacting?: boolean;
4242
isStreaming?: boolean;
4343
supportsTranscriptCommitWait?: boolean;
44+
isAcceptingMessages?: boolean;
4445
} = {},
4546
): RunHandle {
4647
const abort = overrides.abort ?? (() => {});
4748
return {
4849
queueMessage: async () => {},
4950
isStreaming: () => overrides.isStreaming ?? true,
5051
isCompacting: () => overrides.isCompacting ?? false,
52+
...(overrides.isAcceptingMessages !== undefined
53+
? { isAcceptingMessages: () => overrides.isAcceptingMessages ?? true }
54+
: {}),
5155
supportsTranscriptCommitWait: overrides.supportsTranscriptCommitWait,
5256
abort,
5357
};
@@ -228,6 +232,31 @@ describe("pi-embedded runner run registry", () => {
228232
});
229233
});
230234

235+
it("rejects queueing into a sessions_yield-aborting run before the stream fully clears", () => {
236+
const queueMessage = vi.fn(async () => {});
237+
setActiveEmbeddedRun(
238+
"session-yield-aborting",
239+
{
240+
...createRunHandle({ isStreaming: true, isAcceptingMessages: false }),
241+
queueMessage,
242+
},
243+
"agent:main:yield-parent",
244+
);
245+
246+
const outcome = queueEmbeddedPiMessageWithOutcome(
247+
"session-yield-aborting",
248+
"subagent completion",
249+
);
250+
251+
expect(outcome).toEqual({
252+
queued: false,
253+
sessionId: "session-yield-aborting",
254+
reason: "not_streaming",
255+
gatewayHealth: "live",
256+
});
257+
expect(queueMessage).not.toHaveBeenCalled();
258+
});
259+
231260
it("returns runtime rejection details when async queue delivery fails", async () => {
232261
setActiveEmbeddedRun("session-rejected", {
233262
...createRunHandle(),

‎src/agents/pi-embedded-runner/runs.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -261,7 +261,7 @@ function prepareEmbeddedPiQueueMessage(
261261
diag.debug(`queue message failed: sessionId=${sessionId} reason=no_active_run`);
262262
return { kind: "complete", outcome: createQueueFailureOutcome(sessionId, "no_active_run") };
263263
}
264-
if (!handle.isStreaming()) {
264+
if (handle.isAcceptingMessages?.() === false || !handle.isStreaming()) {
265265
diag.debug(`queue message failed: sessionId=${sessionId} reason=not_streaming`);
266266
return { kind: "complete", outcome: createQueueFailureOutcome(sessionId, "not_streaming") };
267267
}

‎src/agents/subagent-announce-delivery.test.ts‎

Lines changed: 49 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -884,7 +884,7 @@ describe("deliverSubagentAnnouncement active requester steering", () => {
884884
});
885885

886886
describe("deliverSubagentAnnouncement completion delivery", () => {
887-
it("uses an active requester queue as the completion handoff when message-tool delivery is not required", async () => {
887+
it("uses direct requester-agent handoff before steering active completion requesters", async () => {
888888
const callGateway = createGatewayMock();
889889
const queueEmbeddedPiMessageWithOutcome = createQueueOutcomeMock(true);
890890
const result = await deliverSlackThreadAnnouncement({
@@ -898,24 +898,21 @@ describe("deliverSubagentAnnouncement completion delivery", () => {
898898

899899
expectRecordFields(result, {
900900
delivered: true,
901-
path: "steered",
902-
enqueuedAt: 4_100,
903-
deliveredAt: 4_200,
901+
path: "direct",
902+
phases: [{ phase: "direct-primary", delivered: true, path: "direct", error: undefined }],
904903
});
905-
expect(queueEmbeddedPiMessageWithOutcome).toHaveBeenCalledWith(
906-
"requester-session-1",
907-
"child done",
908-
{
909-
steeringMode: "all",
910-
debounceMs: 500,
911-
waitForTranscriptCommit: true,
912-
deliveryTimeoutMs: 120_000,
913-
},
914-
);
915-
expect(callGateway).not.toHaveBeenCalled();
904+
expectGatewayAgentParams(callGateway, {
905+
deliver: true,
906+
channel: "slack",
907+
accountId: "acct-1",
908+
to: "channel:C123",
909+
threadId: "171.222",
910+
bestEffortDeliver: true,
911+
});
912+
expect(queueEmbeddedPiMessageWithOutcome).not.toHaveBeenCalled();
916913
});
917914

918-
it("does not also direct-run a queued active completion", async () => {
915+
it("falls back to queued steering only after direct active completion handoff fails", async () => {
919916
const callGateway = createGatewayMock();
920917
const queueEmbeddedPiMessageWithOutcome = createQueueOutcomeMock(true);
921918
const result = await deliverSlackThreadAnnouncement({
@@ -931,11 +928,27 @@ describe("deliverSubagentAnnouncement completion delivery", () => {
931928
expectRecordFields(result, {
932929
delivered: true,
933930
path: "steered",
934-
enqueuedAt: 4_100,
935-
deliveredAt: 4_200,
931+
phases: [
932+
{
933+
phase: "direct-primary",
934+
delivered: false,
935+
path: "direct",
936+
deliveredAt: undefined,
937+
enqueuedAt: undefined,
938+
error: "completion agent did not produce a visible reply",
939+
},
940+
{
941+
phase: "steer-fallback",
942+
delivered: true,
943+
path: "steered",
944+
deliveredAt: 4_200,
945+
enqueuedAt: 4_100,
946+
error: undefined,
947+
},
948+
],
936949
});
950+
expect(callGateway).toHaveBeenCalledTimes(1);
937951
expect(queueEmbeddedPiMessageWithOutcome).toHaveBeenCalledTimes(1);
938-
expect(callGateway).not.toHaveBeenCalled();
939952
});
940953

941954
it("keeps direct external delivery for dormant completion requesters", async () => {
@@ -1525,28 +1538,7 @@ describe("deliverSubagentAnnouncement completion delivery", () => {
15251538
to: "channel:C123",
15261539
threadId: "171.222",
15271540
});
1528-
expect(queueEmbeddedPiMessageWithOutcome).toHaveBeenCalledTimes(2);
1529-
expect(queueEmbeddedPiMessageWithOutcome).toHaveBeenNthCalledWith(
1530-
1,
1531-
"requester-session-4",
1532-
"child done",
1533-
{
1534-
debounceMs: 500,
1535-
deliveryTimeoutMs: 120_000,
1536-
steeringMode: "all",
1537-
waitForTranscriptCommit: true,
1538-
},
1539-
);
1540-
expect(queueEmbeddedPiMessageWithOutcome).toHaveBeenNthCalledWith(
1541-
2,
1542-
"requester-session-4",
1543-
"child done",
1544-
{
1545-
debounceMs: 500,
1546-
deliveryTimeoutMs: 120_000,
1547-
steeringMode: "all",
1548-
},
1549-
);
1541+
expect(queueEmbeddedPiMessageWithOutcome).not.toHaveBeenCalled();
15501542
expect(sendMessage).not.toHaveBeenCalled();
15511543
});
15521544

@@ -1765,17 +1757,7 @@ describe("deliverSubagentAnnouncement completion delivery", () => {
17651757
},
17661758
],
17671759
});
1768-
expect(queueEmbeddedPiMessageWithOutcome).toHaveBeenCalledTimes(1);
1769-
expect(queueEmbeddedPiMessageWithOutcome).toHaveBeenCalledWith(
1770-
"requester-session-telegram",
1771-
"child done",
1772-
{
1773-
steeringMode: "all",
1774-
debounceMs: 500,
1775-
waitForTranscriptCommit: true,
1776-
deliveryTimeoutMs: 10,
1777-
},
1778-
);
1760+
expect(queueEmbeddedPiMessageWithOutcome).not.toHaveBeenCalled();
17791761
expect(callGateway).toHaveBeenCalledTimes(1);
17801762
expect(sendMessage).not.toHaveBeenCalled();
17811763
});
@@ -1835,7 +1817,7 @@ describe("deliverSubagentAnnouncement completion delivery", () => {
18351817
},
18361818
],
18371819
});
1838-
expect(queueEmbeddedPiMessageWithOutcome).toHaveBeenCalledTimes(1);
1820+
expect(queueEmbeddedPiMessageWithOutcome).not.toHaveBeenCalled();
18391821
expect(callGateway).toHaveBeenCalledTimes(1);
18401822
});
18411823

@@ -2507,7 +2489,7 @@ describe("deliverSubagentAnnouncement completion delivery", () => {
25072489
);
25082490
});
25092491

2510-
it("keeps generated media completions on the active requester session path", async () => {
2492+
it("keeps generated media completions on the direct requester handoff path", async () => {
25112493
const callGateway = createGatewayMock();
25122494
const queueEmbeddedPiMessageWithOutcome = createQueueOutcomeMock(true);
25132495
const sendMessage = createSendMessageMock();
@@ -2540,23 +2522,20 @@ describe("deliverSubagentAnnouncement completion delivery", () => {
25402522

25412523
expectRecordFields(result, {
25422524
delivered: true,
2543-
path: "steered",
2544-
enqueuedAt: 4_100,
2545-
deliveredAt: 4_200,
2525+
path: "direct",
25462526
});
2547-
expect(queueEmbeddedPiMessageWithOutcome).toHaveBeenCalledWith(
2548-
"requester-session-channel",
2549-
"child done",
2550-
{
2551-
steeringMode: "all",
2552-
sourceReplyDeliveryMode: "message_tool_only",
2553-
debounceMs: 500,
2554-
waitForTranscriptCommit: true,
2555-
deliveryTimeoutMs: 120_000,
2556-
},
2527+
expect(callGateway).toHaveBeenCalledTimes(1);
2528+
expect(queueEmbeddedPiMessageWithOutcome).not.toHaveBeenCalled();
2529+
expect(sendMessage).toHaveBeenCalledWith(
2530+
expect.objectContaining({
2531+
channel: "slack",
2532+
accountId: "acct-1",
2533+
to: "channel:C123",
2534+
content: "The generated video is ready.",
2535+
mediaUrls: ["/tmp/generated-corgi.mp4"],
2536+
idempotencyKey: "announce-channel-media-active-direct:generated-media-direct",
2537+
}),
25572538
);
2558-
expect(callGateway).not.toHaveBeenCalled();
2559-
expect(sendMessage).not.toHaveBeenCalled();
25602539
});
25612540

25622541
it("directly delivers missing generated media after active requester wake failure", async () => {
@@ -2612,7 +2591,7 @@ describe("deliverSubagentAnnouncement completion delivery", () => {
26122591
delivered: true,
26132592
path: "direct",
26142593
});
2615-
expect(queueEmbeddedPiMessageWithOutcome).toHaveBeenCalledTimes(2);
2594+
expect(queueEmbeddedPiMessageWithOutcome).not.toHaveBeenCalled();
26162595
expect(callGateway).toHaveBeenCalledTimes(1);
26172596
expect(sendMessage).toHaveBeenCalledWith(
26182597
expect.objectContaining({

‎src/agents/subagent-announce-delivery.ts‎

Lines changed: 35 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -459,6 +459,7 @@ async function maybeSteerSubagentAnnounce(params: {
459459
deliveryTimeoutMs?: number;
460460
requesterSessionKey: string;
461461
steerMessage: string;
462+
sourceReplyDeliveryMode?: "message_tool_only";
462463
signal?: AbortSignal;
463464
}): Promise<
464465
{ status: "steered"; deliveredAt?: number; enqueuedAt?: number } | { status: "none" | "dropped" }
@@ -484,6 +485,9 @@ async function maybeSteerSubagentAnnounce(params: {
484485
const queueOptions: EmbeddedPiQueueMessageOptions = {
485486
deliveryTimeoutMs: params.deliveryTimeoutMs,
486487
steeringMode: "all",
488+
...(params.sourceReplyDeliveryMode
489+
? { sourceReplyDeliveryMode: params.sourceReplyDeliveryMode }
490+
: {}),
487491
...(queueSettings.debounceMs !== undefined ? { debounceMs: queueSettings.debounceMs } : {}),
488492
waitForTranscriptCommit: true,
489493
};
@@ -566,7 +570,7 @@ function requiresAgentMediatedCompletionDelivery(params: {
566570
}
567571

568572
function collectExpectedMediaFromInternalEvents(
569-
events: AgentInternalEvent[] | undefined,
573+
events: readonly AgentInternalEvent[] | undefined,
570574
): string[] {
571575
if (!events?.length) {
572576
return [];
@@ -880,6 +884,7 @@ async function sendSubagentAnnounceDirectly(params: {
880884
triggerMessage: string;
881885
internalEvents?: AgentInternalEvent[];
882886
expectsCompletionMessage: boolean;
887+
allowActiveRequesterWake?: boolean;
883888
bestEffortDeliver?: boolean;
884889
directIdempotencyKey: string;
885890
completionDirectOrigin?: DeliveryContext;
@@ -965,7 +970,11 @@ async function sendSubagentAnnounceDirectly(params: {
965970
const requesterActivity = resolveRequesterSessionActivity(canonicalRequesterSessionKey);
966971
let activeRequesterWakeFailed = false;
967972
const tryGeneratedMediaDirectDelivery = async (announceResponse?: unknown) => {
968-
if (requesterActivity.isActive && !activeRequesterWakeFailed) {
973+
if (
974+
params.allowActiveRequesterWake === true &&
975+
requesterActivity.isActive &&
976+
!activeRequesterWakeFailed
977+
) {
969978
return undefined;
970979
}
971980
const missingMediaUrls = resolveGeneratedMediaDirectFallbackUrls({
@@ -998,6 +1007,7 @@ async function sendSubagentAnnounceDirectly(params: {
9981007
});
9991008
if (
10001009
params.expectsCompletionMessage &&
1010+
params.allowActiveRequesterWake === true &&
10011011
requesterActivity.sessionId &&
10021012
requesterActivity.isActive
10031013
) {
@@ -1252,6 +1262,23 @@ async function sendSubagentAnnounceDirectly(params: {
12521262
}
12531263
}
12541264

1265+
function resolveFallbackSteerSourceReplyDeliveryMode(params: {
1266+
expectsCompletionMessage: boolean;
1267+
sourceTool?: string;
1268+
internalEvents?: readonly AgentInternalEvent[];
1269+
}): "message_tool_only" | undefined {
1270+
if (
1271+
requiresAgentMediatedCompletionDelivery({
1272+
expectsCompletionMessage: params.expectsCompletionMessage,
1273+
sourceTool: params.sourceTool,
1274+
}) &&
1275+
collectExpectedMediaFromInternalEvents(params.internalEvents).length > 0
1276+
) {
1277+
return "message_tool_only";
1278+
}
1279+
return undefined;
1280+
}
1281+
12551282
export async function deliverSubagentAnnouncement(params: {
12561283
requesterSessionKey: string;
12571284
announceId?: string;
@@ -1283,6 +1310,11 @@ export async function deliverSubagentAnnouncement(params: {
12831310
),
12841311
requesterSessionKey: params.requesterSessionKey,
12851312
steerMessage: params.steerMessage,
1313+
sourceReplyDeliveryMode: resolveFallbackSteerSourceReplyDeliveryMode({
1314+
expectsCompletionMessage: params.expectsCompletionMessage,
1315+
sourceTool: params.sourceTool,
1316+
internalEvents: params.internalEvents,
1317+
}),
12861318
signal: params.signal,
12871319
}),
12881320
direct: async () =>
@@ -1300,6 +1332,7 @@ export async function deliverSubagentAnnouncement(params: {
13001332
sourceTool: params.sourceTool,
13011333
requesterIsSubagent: params.requesterIsSubagent,
13021334
expectsCompletionMessage: params.expectsCompletionMessage,
1335+
allowActiveRequesterWake: !params.expectsCompletionMessage,
13031336
signal: params.signal,
13041337
bestEffortDeliver: params.bestEffortDeliver,
13051338
}),

0 commit comments

Comments
 (0)