Skip to content

fix(server): release Codex sessions after transport termination #20992

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
matthias-scale wants to merge 1 commit into from
Closed

fix(server): release Codex sessions after transport termination #20992

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
22 changes: 20 additions & 2 deletions apps/server/src/orchestration-v2/Adapters/CodexAdapterV2.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1791,7 +1791,25 @@ export const makeCodexAdapterV2 = Effect.fn("makeCodexAdapterV2")(function* (
model: input.modelSelection.model,
now,
});
const events = yield* Queue.unbounded<ProviderAdapter.ProviderAdapterV2Event>();
const events = yield* Queue.unbounded<
ProviderAdapter.ProviderAdapterV2Event,
ProviderAdapter.ProviderAdapterEventStreamError
>();
// Drain decoded events before reporting transport death. The manager
// owns eviction and cleanup; never replay a possibly accepted turn.
yield* client.termination.pipe(
Effect.catch((cause) =>
Queue.fail(
events,
new ProviderAdapter.ProviderAdapterEventStreamError({
driver: CODEX_PROVIDER,
providerSessionId: input.providerSessionId,
cause,
}),
),
),
Effect.forkScoped,
);
const rateLimitSnapshot = yield* Ref.make<CodexRateLimitSnapshot | undefined>(undefined);
const limitedTurnItems = yield* Ref.make(
new Map<ProviderThreadId, Extract<OrchestrationV2TurnItem, { type: "error" }>>(),
Expand Down Expand Up @@ -6405,7 +6423,7 @@ export const makeCodexAdapterV2 = Effect.fn("makeCodexAdapterV2")(function* (
driver: CODEX_PROVIDER,
providerSessionId: input.providerSessionId,
providerSession: session,
events: Stream.fromEffectRepeat(Queue.take(events)),
events: Stream.fromQueue(events),
canReuseContextUsage: canReuseCodexContextUsage,
// Known gap: a subagent that Codex resumes later reads as completed
// (not pending) between turns, so idle release can win the race
Expand Down
280 changes: 255 additions & 25 deletions apps/server/src/orchestration-v2/ProviderSessionManager.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ import * as NodeServices from "@effect/platform-node/NodeServices";
import { assert, it } from "@effect/vitest";
import { AgentScope } from "@t3tools/shared/AgentScope";
import {
CodexSettings,
EnvironmentId,
type ModelSelection,
type OrchestrationV2AppThread,
Expand All @@ -17,6 +18,15 @@ import {
type ProviderSessionId,
ThreadId,
} from "@t3tools/contracts";
import * as CodexClient from "effect-codex-app-server/client";
import * as CodexError from "effect-codex-app-server/errors";
import * as Crypto from "effect/Crypto";
import * as Schema from "effect/Schema";
import * as Path from "effect/Path";
import * as Sink from "effect/Sink";
import * as Stdio from "effect/Stdio";
import * as CodexAdapter from "./Adapters/CodexAdapterV2.ts";
import { makeReplayServerConfig } from "./Adapters/CodexAdapterV2.testkit.ts";
import * as Cause from "effect/Cause";
import * as DateTime from "effect/DateTime";
import * as Deferred from "effect/Deferred";
Expand Down Expand Up @@ -160,6 +170,8 @@ const layerBeforeUnchangedWriteEventSink = (before: Ref.Ref<Effect.Effect<void>>
}),
).pipe(Layer.provide(layerTestEventSink));

const decodeCodexSettings = Schema.decodeEffect(CodexSettings);

const CodexCapabilities: OrchestrationV2ProviderCapabilities = CodexProviderCapabilitiesV2;
const ExclusiveCapabilities: OrchestrationV2ProviderCapabilities = {
...CodexCapabilities,
Expand Down Expand Up @@ -458,6 +470,15 @@ const makeProviderAdapter = Effect.fnUntraced(function* (
});

function layerTest(input: {
readonly adapters?: Effect.Effect<
ReadonlyArray<ProviderAdapter.ProviderAdapterV2["Service"]>,
never,
| McpProviderSessions.McpProviderSessions
| FileSystem.FileSystem
| Path.Path
| Crypto.Crypto
| IdAllocator.IdAllocatorV2
>;
readonly state: Ref.Ref<TestProviderRuntimeState>;
readonly idleTimeoutMs: number;
readonly maxIdlePinMs?: number;
Expand Down Expand Up @@ -500,30 +521,33 @@ function layerTest(input: {
? layerFailingReleaseEventSink
: layerTestEventSink;
const layerRegistry = Layer.unwrap(
makeProviderAdapter(input.state, {
failEventStream: input.failEventStream ?? false,
...(input.capabilities === undefined ? {} : { capabilities: input.capabilities }),
...(input.mcpConfigs === undefined ? {} : { mcpConfigs: input.mcpConfigs }),
...(input.beforeOpen === undefined ? {} : { beforeOpen: input.beforeOpen }),
...(input.hasPendingBackgroundWork === undefined
? {}
: { hasPendingBackgroundWork: input.hasPendingBackgroundWork }),
...(input.hasPendingBackgroundWorkForThread === undefined
? {}
: { hasPendingBackgroundWorkForThread: input.hasPendingBackgroundWorkForThread }),
...(input.backgroundTasksEndWithProcess === undefined
? {}
: { backgroundTasksEndWithProcess: input.backgroundTasksEndWithProcess }),
...(input.hangSessionScopeClose === undefined
? {}
: { hangSessionScopeClose: input.hangSessionScopeClose }),
...(input.startTurn === undefined ? {} : { startTurn: input.startTurn }),
...(input.beforeUnload === undefined ? {} : { beforeUnload: input.beforeUnload }),
...(input.spawnBeforeOpen === undefined ? {} : { spawnBeforeOpen: input.spawnBeforeOpen }),
...(input.scopeCloseReached === undefined
? {}
: { scopeCloseReached: input.scopeCloseReached }),
}).pipe(Effect.map(ProviderAdapterRegistry.layerSingle)),
(
input.adapters ??
makeProviderAdapter(input.state, {
failEventStream: input.failEventStream ?? false,
...(input.capabilities === undefined ? {} : { capabilities: input.capabilities }),
...(input.mcpConfigs === undefined ? {} : { mcpConfigs: input.mcpConfigs }),
...(input.beforeOpen === undefined ? {} : { beforeOpen: input.beforeOpen }),
...(input.hasPendingBackgroundWork === undefined
? {}
: { hasPendingBackgroundWork: input.hasPendingBackgroundWork }),
...(input.hasPendingBackgroundWorkForThread === undefined
? {}
: { hasPendingBackgroundWorkForThread: input.hasPendingBackgroundWorkForThread }),
...(input.backgroundTasksEndWithProcess === undefined
? {}
: { backgroundTasksEndWithProcess: input.backgroundTasksEndWithProcess }),
...(input.hangSessionScopeClose === undefined
? {}
: { hangSessionScopeClose: input.hangSessionScopeClose }),
...(input.startTurn === undefined ? {} : { startTurn: input.startTurn }),
...(input.beforeUnload === undefined ? {} : { beforeUnload: input.beforeUnload }),
...(input.spawnBeforeOpen === undefined ? {} : { spawnBeforeOpen: input.spawnBeforeOpen }),
...(input.scopeCloseReached === undefined
? {}
: { scopeCloseReached: input.scopeCloseReached }),
}).pipe(Effect.map((adapter) => [adapter]))
).pipe(Effect.map(ProviderAdapterRegistry.layerFromAdapters)),
);
const layerConfiguredMcpRegistry =
input.pauseResolve === undefined
Expand Down Expand Up @@ -554,7 +578,10 @@ function layerTest(input: {
}).pipe(
Layer.provide(
Layer.mergeAll(
layerRegistry.pipe(Layer.provideMerge(McpProviderSessions.layer)),
layerRegistry.pipe(
Layer.provide(IdAllocator.layer),
Layer.provideMerge(McpProviderSessions.layer),
),
CommandOutputHub.layer.pipe(
Layer.provide(Layer.merge(layerConfiguredEventSink, layerTestStores)),
),
Expand Down Expand Up @@ -4908,3 +4935,206 @@ it.effect(
assert.isFalse(denied?.capabilities?.has("device"));
}),
);

it.effect.each([false, true])(
"Codex transport death reopens the saved native thread (in flight: %s)",
(inFlight) =>
Effect.gen(function* () {
const state = yield* Ref.make(emptyState);
const frames: Array<{ process: number; method: string; params?: { threadId?: string } }> = [];
const processes: Array<{
input: Queue.Queue<Uint8Array, Cause.Done>;
closed: Deferred.Deferred<void>;
submitted: Deferred.Deferred<void>;
}> = [];
const healthyInstanceId = ProviderInstanceId.make("codex-healthy");
const settings = yield* decodeCodexSettings({
homePath: "/custom-codex-home",
}).pipe(Effect.orDie);
const openedSettings: Array<CodexSettings> = [];
const adapters = Effect.forEach(
[modelSelection.instanceId, healthyInstanceId],
(instanceId) =>
Effect.gen(function* () {
return yield* CodexAdapter.makeCodexAdapterV2({
instanceId,
settings,
environment: {},
crypto: yield* Crypto.Crypto,
fileSystem: yield* FileSystem.FileSystem,
idAllocator: yield* IdAllocator.IdAllocatorV2,
serverConfig: yield* makeReplayServerConfig("dead-session").pipe(Effect.orDie),
clientFactory: {
open: (openInput) =>
Effect.gen(function* () {
openedSettings.push(openInput.settings);
const processIndex = processes.length;
const input = yield* Queue.unbounded<Uint8Array, Cause.Done>();
const closed = yield* Deferred.make<void>();
const submitted = yield* Deferred.make<void>();
processes.push({ input, closed, submitted });
yield* Effect.addFinalizer(() => Deferred.succeed(closed, undefined));
const encoder = new TextEncoder();
return yield* CodexClient.make(
Stdio.make({
args: Effect.succeed([]),
stdin: Stream.fromQueue(input),
stdout: () =>
Sink.forEach((chunk: string | Uint8Array) =>
Effect.gen(function* () {
const frame = JSON.parse(
typeof chunk === "string" ? chunk : new TextDecoder().decode(chunk),
);
frames.push({
process: processIndex,
method: frame.method,
params: frame.params,
});
if (frame.id === undefined) return;
if (frame.method === "turn/start") {
yield* Deferred.succeed(submitted, undefined);
return; // Accepted, but the reply is lost when the process dies.
}
const result =
frame.method === "initialize"
? {
userAgent: "test",
codexHome: process.cwd(),
platformFamily: "unix",
platformOs: "linux",
}
: { thread: { id: frame.params.threadId, updatedAt: 1 } };
yield* Queue.offer(
input,
encoder.encode(JSON.stringify({ id: frame.id, result }) + "\n"),
);
}),
),
stderr: () => Sink.drain,
}),
{},
Effect.succeed(
new CodexError.CodexAppServerTransportError({
operation: "read-process-exit-status",
cause: "SIGKILL",
}),
),
);
}),
},
});
}),
);
yield* Effect.gen(function* () {
const manager = yield* ProviderSessionManager.ProviderSessionManagerV2;
const eventSink = yield* EventSink.EventSinkV2;
const allocator = yield* IdAllocator.IdAllocatorV2;
const projections = yield* ProjectionStore.ProjectionStoreV2;
const now = yield* DateTime.now;
const threadId = ThreadId.make("dead-codex");
const healthyThreadId = ThreadId.make("healthy-codex");
for (const id of [threadId, healthyThreadId]) {
yield* eventSink.write({
events: [yield* makeThreadCreatedEvent({ idAllocator: allocator, threadId: id, now })],
});
}
const providerSessionId = yield* allocator.allocate.providerSession({
providerInstanceId: modelSelection.instanceId,
threadId,
});
const healthySessionId = yield* allocator.allocate.providerSession({
providerInstanceId: healthyInstanceId,
threadId: healthyThreadId,
});
const openInput = { threadId, providerSessionId, modelSelection, runtimePolicy };
const first = yield* manager.open(openInput);
const saved = makeProviderThread({
idAllocator: allocator,
threadId,
providerSessionId,
now,
nativeThreadId: "saved-native-id",
});
const resumeInput = { threadId, providerThread: saved, modelSelection, runtimePolicy };
yield* first.resumeThread(resumeInput);
const healthy = yield* manager.open({
...openInput,
threadId: healthyThreadId,
providerSessionId: healthySessionId,
modelSelection: { ...modelSelection, instanceId: healthyInstanceId },
});
const subscription = yield* first.subscribeEvents!;
const streamExit = yield* subscription.events.pipe(
Stream.runDrain,
Effect.exit,
Effect.forkChild,
);
const runId = allocator.derive.run({ threadId, ordinal: 1 });
const starting = inFlight
? yield* first
.startTurn({
appThread: (yield* projections.getThreadProjection(threadId)).thread,
threadId,
runId,
runOrdinal: 1,
providerTurnOrdinal: 1,
attemptId: allocator.derive.runAttempt({ runId, attemptOrdinal: 1 }),
rootNodeId: allocator.derive.rootNode({ runId }),
providerThread: saved,
message: {
createdBy: "user",
creationSource: "web",
messageId: yield* allocator.allocate.message({ threadId, ordinal: 1 }),
text: "do work once",
attachments: [],
},
modelSelection,
runtimePolicy,
})
.pipe(Effect.exit, Effect.forkChild)
: undefined;
if (starting) yield* Deferred.await(processes[0]!.submitted);
yield* Queue.end(processes[0]!.input);
assert.isTrue(Exit.isFailure(yield* Fiber.join(streamExit)));
if (starting) assert.isTrue(Exit.isFailure(yield* Fiber.join(starting)));
yield* Deferred.await(processes[0]!.closed);
assert.isTrue(Option.isNone(yield* manager.get(providerSessionId)));
assert.strictEqual(Option.getOrThrow(yield* manager.get(healthySessionId)), healthy);
assert.isFalse(yield* Deferred.isDone(processes[1]!.closed));
yield* healthy.resumeThread({
threadId: healthyThreadId,
providerThread: makeProviderThread({
idAllocator: allocator,
threadId: healthyThreadId,
providerSessionId: healthySessionId,
now,
nativeThreadId: "healthy-native-id",
}),
modelSelection: { ...modelSelection, instanceId: healthyInstanceId },
runtimePolicy,
});
const replacement = yield* manager.open(openInput);
const resumed = yield* replacement.resumeThread(resumeInput);
assert.notStrictEqual(replacement, first);
assert.equal(resumed.nativeThreadRef?.nativeId, "saved-native-id");
assert.equal(processes.length, 3);
assert.deepEqual(openedSettings, [settings, settings, settings]);
assert.deepEqual(
frames
.filter((f) => f.method === "thread/resume")
.map((f) => [f.process, f.params?.threadId]),
[
[0, "saved-native-id"],
[1, "healthy-native-id"],
[2, "saved-native-id"],
],
);
assert.equal(frames.filter((f) => f.method === "thread/start").length, 0);
assert.equal(frames.filter((f) => f.method === "turn/start").length, inFlight ? 1 : 0);
assert.deepEqual(
frames.filter((f) => f.method === "initialize").map((f) => f.process),
[0, 1, 2],
);
}).pipe(Effect.provide(layerTest({ state, adapters, idleTimeoutMs: 60_000 })));
}),
);
21 changes: 21 additions & 0 deletions packages/effect-codex-app-server/src/client.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import * as Fiber from "effect/Fiber";
import * as Exit from "effect/Exit";
import * as Layer from "effect/Layer";
import * as Path from "effect/Path";
Expand Down Expand Up @@ -124,6 +125,26 @@ it.layer(NodeServices.layer)("effect-codex-app-server client", (it) => {
]);
}),
);
it.effect("reports SIGKILL between requests to every termination observer", () =>
Effect.gen(function* () {
const handle = yield* makeHandle();
const context = yield* Layer.build(CodexClient.layerChildProcess(handle));
const client = yield* Effect.service(CodexClient.CodexAppServerClient).pipe(
Effect.provide(context),
);
yield* client.request("initialize", {
clientInfo: { name: "termination-test", version: "0.0.0" },
capabilities: { experimentalApi: true, optOutNotificationMethods: null },
});
const observer = yield* client.termination.pipe(Effect.exit, Effect.forkChild);
yield* handle.kill({ killSignal: "SIGKILL" });
const exit = yield* Fiber.join(observer);
assert.isTrue(Exit.isFailure(exit));
// The signal is durable even for observers attached after the death.
assert.deepEqual(yield* Effect.exit(client.termination), exit);
assert.isTrue(Exit.isFailure(yield* Effect.exit(client.raw.request("account/read", {}))));
}),
);
it.effect("drains child stderr so large diagnostics cannot block protocol responses", () =>
Effect.gen(function* () {
const handle = yield* makeHandle({
Expand Down
Loading
Loading