Skip to content
Closed
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
79 changes: 79 additions & 0 deletions apps/server/src/orchestration-v2/Adapters/ClaudeAdapterV2.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11137,6 +11137,85 @@ describe("ClaudeAdapterV2 background wake turns", () => {
),
);

it.effect("lets a released continuation offer a later wake without dropping its output", () =>

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This repeats the existing buffers wake output and requests a single continuation run setup and endpoint. Could you fold the release and stale-offer assertions into that case instead of adding another test?

Posted via Macroscope — Test Signal

Effect.scoped(
Effect.gen(function* () {
const harness = yield* makeWakeHarness;
const now = yield* DateTime.now;
const SECOND_WAKE_TEXT = "SECOND_WAKE_DONE";
const nextWake = (index: number, text: string) => [
claudeSdkFrame({
type: "system",
subtype: "init",
uuid: `00000000-0000-4000-8000-00000000012${index}`,
session_id: WAKE_NATIVE_SESSION,
}),
makeAssistantTextFrame({ uuid: `00000000-0000-4000-8000-00000000013${index}`, text }),
makeResultFrame({ uuid: `00000000-0000-4000-8000-00000000014${index}`, result: text }),
];

yield* harness.runtime.startTurn(
makeClaudeTestTurnInput({
threadId: harness.threadId,
providerThread: harness.providerThread,
now,
attemptId: RunAttemptId.make("attempt-claude-wake-release-1"),
text: "Run the build in the background.",
attachments: [],
}),
);
yield* Queue.offer(harness.sdkMessages, wakeTaskStarted);
yield* Queue.offer(harness.sdkMessages, turnOneResult);
yield* harness.awaitTerminal(1);
yield* Queue.offer(harness.sdkMessages, wakeNotification);
yield* Queue.offer(harness.sdkMessages, wakeAssistant);
yield* harness.offerAndWait(wakeResult);
assert.lengthOf(harness.continuationRequests, 1);
const cancelled = harness.continuationRequests[0]?.clearIfCurrent;
assert.isDefined(cancelled, "Claude must let a continuation that never starts release it");

// The first continuation run ended without starting, e.g. cancelled while queued.
yield* cancelled!();
for (const frame of nextWake(1, SECOND_WAKE_TEXT)) yield* harness.offerAndWait(frame);
assert.lengthOf(harness.continuationRequests, 2);

// A late release for the old offer must not reopen the latch the newer offer holds.
yield* cancelled!();
for (const frame of nextWake(2, "THIRD_WAKE_DONE")) yield* harness.offerAndWait(frame);
assert.lengthOf(harness.continuationRequests, 2);

yield* harness.runtime.startTurn(
makeClaudeTestTurnInput({
threadId: harness.threadId,
providerThread: harness.providerThread,
now,
attemptId: RunAttemptId.make("attempt-claude-wake-release-2"),
text: "Background task completed.",
attachments: [],
providerTurnOrdinal: 2,
messageCreatedBy: "agent",
messageCreationSource: "provider",
}),
);
yield* harness.awaitTerminal(2);
const texts = harness.events.flatMap((event) =>
event.type === "message.updated" ? [event.message.text] : [],
);
assert.include(texts, WAKE_ASSISTANT_TEXT);
assert.include(texts, SECOND_WAKE_TEXT);
}).pipe(
Effect.provide(
Layer.mergeAll(
IdAllocator.layer,
McpProviderSessions.layer,
NodeServices.layer,
NodeCrypto.layer,
),
),
),
),
);

it.effect("starts the wake run when Claude opens the wake turn, before its output", () =>
Effect.scoped(
Effect.gen(function* () {
Expand Down
22 changes: 17 additions & 5 deletions apps/server/src/orchestration-v2/Adapters/ClaudeAdapterV2.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5090,7 +5090,8 @@ export const makeClaudeAdapterV2 = Effect.fn("makeClaudeAdapterV2")(function* (
return { kind: "command", label, outcome };
}
});
const requestedContinuations = yield* Ref.make(new Set<string>());
// Each offer owns a token so a late release for an old offer cannot clear a newer one.
const requestedContinuations = yield* Ref.make(new Map<string, symbol>());
const absorbedWakes = yield* Ref.make(new Set<string>());
const heldProposedPlansByToolUseId = new Map<string, string>();
const runtimeContext = yield* Effect.context<never>();
Expand Down Expand Up @@ -5551,7 +5552,7 @@ export const makeClaudeAdapterV2 = Effect.fn("makeClaudeAdapterV2")(function* (
if (!current.has(nativeThreadId)) {
return current;
}
const updated = new Set(current);
const updated = new Map(current);
updated.delete(nativeThreadId);
return updated;
});
Expand Down Expand Up @@ -8194,12 +8195,13 @@ export const makeClaudeAdapterV2 = Effect.fn("makeClaudeAdapterV2")(function* (
});
return;
}
const continuationToken = Symbol(wakeInput.nativeThreadId);
const shouldOffer = yield* Ref.modify(requestedContinuations, (current) => {
if (current.has(wakeInput.nativeThreadId)) {
return [false, current] as const;
}
const updated = new Set(current);
updated.add(wakeInput.nativeThreadId);
const updated = new Map(current);
updated.set(wakeInput.nativeThreadId, continuationToken);
return [true, updated] as const;
});
if (!shouldOffer) {
Expand All @@ -8222,6 +8224,16 @@ export const makeClaudeAdapterV2 = Effect.fn("makeClaudeAdapterV2")(function* (
driver: CLAUDE_PROVIDER,
detail,
...(notification === null ? {} : { notification }),
// Buffered wake output stays for the next continuation to replay.
clearIfCurrent: () =>
Ref.update(requestedContinuations, (current) => {
if (current.get(wakeInput.nativeThreadId) !== continuationToken) {
return current;
}
const updated = new Map(current);
updated.delete(wakeInput.nativeThreadId);
return updated;
}),
});
});

Expand Down Expand Up @@ -10896,7 +10908,7 @@ export const makeClaudeAdapterV2 = Effect.fn("makeClaudeAdapterV2")(function* (
yield* attachTurn;
const drained = (yield* Ref.get(wakeBuffers)).get(nativeThreadId)?.messages ?? [];
yield* Ref.update(requestedContinuations, (current) => {
const updated = new Set(current);
const updated = new Map(current);
updated.delete(nativeThreadId);
return updated;
});
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import {
ProviderThreadId,
RunId,
ThreadId,
type OrchestrationV2StoredEvent,
type OrchestrationV2ThreadProjection,
type OrchestrationV2Notification,
} from "@t3tools/contracts";
Expand All @@ -18,6 +19,7 @@ import * as Option from "effect/Option";
import * as Queue from "effect/Queue";
import * as Ref from "effect/Ref";
import * as Semaphore from "effect/Semaphore";
import * as Stream from "effect/Stream";
import * as TestClock from "effect/testing/TestClock";

import * as IdAllocator from "@t3tools/provider-core/server/IdAllocator";
Expand Down Expand Up @@ -420,6 +422,96 @@ describe("ProviderContinuationService", () => {
},
);

it.effect(
"releases an accepted continuation once its exact run ends, replaying a cancel committed before the watch",
() => {
return Effect.gen(function* () {
const continuationRunId = RunId.make("run-provider-continuation-queued");
const otherRunId = RunId.make("run-provider-continuation-other");
const runEvent = (sequence: number, id: RunId, status: string) =>
({
sequence,
commandId: null,
event: { type: "run.updated", threadId, runId: id, payload: { id, status } },
}) as unknown as OrchestrationV2StoredEvent;
const log = yield* Queue.unbounded<OrchestrationV2StoredEvent>();
const continuationEnded = yield* Ref.make(false);
const watchedFrom = yield* Deferred.make<number | undefined>();
const released = yield* Deferred.make<boolean>();
const layerThreads = Layer.mock(ThreadManagementService.ThreadManagementService)({
getThreadRecords: () => Effect.succeed(projection),
dispatch: (command) =>
Effect.gen(function* () {
const { messageId } = command as { readonly messageId: MessageId };
// The run ends before the service can subscribe, so only replay from the dispatch sequence finds it.
yield* Queue.offerAll(log, [
runEvent(11, otherRunId, "cancelled"),
runEvent(12, continuationRunId, "running"),
runEvent(13, continuationRunId, "cancelled"),
]);
return {
sequence: 10,
storedEvents: [
{
sequence: 10,
commandId: null,
event: {
type: "run.created",
threadId,
runId: continuationRunId,
payload: {
id: continuationRunId,
userMessageId: messageId,
status: "queued",
},
},
},
],
} as never;
}),
streamStoredEventsFrom: (input) =>
Stream.unwrap(
Deferred.succeed(watchedFrom, input?.afterSequence).pipe(
Effect.as(
Stream.fromEffectRepeat(
Queue.take(log).pipe(
Effect.tap((stored) =>
stored.sequence === 13 ? Ref.set(continuationEnded, true) : Effect.void,
),
),
),
),
),
),
});
const layerWorker = ProviderContinuationService.layer.pipe(
Layer.provide(
Layer.mergeAll(IdAllocator.layer, ProviderContinuationRequests.layer, layerThreads),
),
);

yield* Effect.gen(function* () {
const requests = yield* ProviderContinuationRequests.ProviderContinuationRequests;
yield* requests.offer({
...request(),
clearIfCurrent: () =>
Ref.get(continuationEnded).pipe(
Effect.flatMap((ended) => Deferred.succeed(released, ended)),
),
});
assert.equal(yield* Deferred.await(watchedFrom), 10);
assert.isTrue(
yield* Deferred.await(released),
"only the continuation run's own terminal event may release the adapter offer",
);
}).pipe(
Effect.provide(Layer.merge(ProviderContinuationRequests.layer, layerWorker)),
Effect.scoped,
);
});
},
);

it.effect(
"drops stopped and disposed delegated completions instead of reviving them after recovery",
() => {
Expand Down
85 changes: 72 additions & 13 deletions apps/server/src/orchestration-v2/ProviderContinuationService.ts
Original file line number Diff line number Diff line change
@@ -1,10 +1,17 @@
import { CommandId, type OrchestrationV2ThreadProjection } from "@t3tools/contracts";
import {
CommandId,
type MessageId,
type OrchestrationV2ThreadProjection,
} from "@t3tools/contracts";
import * as Effect from "effect/Effect";
import * as Layer from "effect/Layer";
import * as Option from "effect/Option";
import * as Ref from "effect/Ref";
import * as Stream from "effect/Stream";

import * as IdAllocator from "@t3tools/provider-core/server/IdAllocator";
import * as ProviderContinuationRequests from "@t3tools/provider-core/server/ProviderContinuationRequests";
import type * as Orchestrator from "./Orchestrator.ts";
import * as ThreadManagementService from "./ThreadManagementService.ts";
import { isUndeliveredMailboxSteer } from "./NotificationMailbox.ts";

Expand Down Expand Up @@ -56,6 +63,7 @@ export const layer = Layer.effectDiscard(
const requests = yield* ProviderContinuationRequests.ProviderContinuationRequests;
const threads = yield* ThreadManagementService.ThreadManagementService;
const retryAttempts = yield* Ref.make(new Map<string, number>());
const layerScope = yield* Effect.scope;

const clearRetryAttempt = (key: string) =>
Ref.update(retryAttempts, (current) => {
Expand All @@ -73,6 +81,50 @@ export const layer = Layer.effectDiscard(
return [Math.min(100 * 2 ** Math.min(attempt, 6), 5_000), updated] as const;
});

// A continuation run can end without starting (a queued cancel), so the adapter never sees it.
const releaseWhenRunEnds = Effect.fn("ProviderContinuationService.releaseWhenRunEnds")(
function* (input: {
readonly request: ProviderContinuationRequests.ProviderContinuationRequest;
readonly clearIfCurrent: () => Effect.Effect<void>;
readonly messageId: MessageId;
readonly dispatched: Orchestrator.OrchestratorV2DispatchResult;
}) {
const runId = input.dispatched.storedEvents.flatMap(({ event }) =>
event.type === "run.created" && event.payload.userMessageId === input.messageId
? [event.payload.id]
: [],
)[0];
if (runId === undefined) return;
// Replays from the dispatch so a run that ended before this subscription still counts.
const ended = yield* threads
.streamStoredEventsFrom({
threadId: input.request.threadId,
afterSequence: input.dispatched.sequence,
eventType: "run.updated",
})
.pipe(
Stream.filter(
({ event }) =>
event.type === "run.updated" &&
event.payload.id === runId &&
ThreadManagementService.isTerminalRunStatus(event.payload.status),
),
Stream.runHead,
);
if (Option.isSome(ended)) yield* input.clearIfCurrent();
},
(effect, input) =>
effect.pipe(
Effect.catchCause((cause) =>
Effect.logWarning("orchestration-v2.provider-continuation.release-watch-failed", {
threadId: input.request.threadId,
providerThreadId: input.request.providerThreadId,
cause,
}),
),
),
);

const dispatchContinuation = Effect.fn("ProviderContinuationService.dispatchContinuation")(
function* (request: ProviderContinuationRequests.ProviderContinuationRequest) {
const projection = yield* threads.getThreadRecords(
Expand Down Expand Up @@ -218,19 +270,26 @@ export const layer = Layer.effectDiscard(
request.clearIfCurrent !== undefined ? request.clearIfCurrent() : Effect.void,
),
);
if (request.dispatchIfCurrent === undefined) {
yield* dispatch.pipe(
Effect.catchTags({
ThreadManagementProviderInstanceDisabledError: logInstanceDisabled,
}),
);
return;
const refused = (
error: ThreadManagementService.ThreadManagementProviderInstanceDisabledError,
) => logInstanceDisabled(error).pipe(Effect.as(Option.none()));
const dispatched =
request.dispatchIfCurrent === undefined
? yield* dispatch.pipe(
Effect.asSome,
Effect.catchTags({ ThreadManagementProviderInstanceDisabledError: refused }),
)
: yield* request
.dispatchIfCurrent(dispatch)
.pipe(Effect.catchTags({ ThreadManagementProviderInstanceDisabledError: refused }));
if (Option.isSome(dispatched) && request.clearIfCurrent !== undefined) {
yield* releaseWhenRunEnds({
request,
clearIfCurrent: request.clearIfCurrent,
messageId,
dispatched: dispatched.value,
}).pipe(Effect.forkIn(layerScope));
}
yield* request.dispatchIfCurrent(dispatch).pipe(
Effect.catchTags({
ThreadManagementProviderInstanceDisabledError: logInstanceDisabled,
}),
);
},
);

Expand Down
Loading