Skip to content

Commit b97e463

Browse files
maria-rcksmaria-rcks
committed
fix(server): refresh cached sessions after provider routing changes
1 parent 35014f5 commit b97e463

3 files changed

Lines changed: 105 additions & 2 deletions

File tree

‎apps/server/src/orchestration-v2/ProviderAdapterRegistry.ts‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,9 @@ export interface ProviderAdapterRegistryV2Shape {
4242
readonly get: (
4343
instanceId: ProviderInstanceId,
4444
) => Effect.Effect<ProviderAdapter.ProviderAdapterV2["Service"], ProviderAdapterRegistryV2Error>;
45+
readonly getRuntimeIdentity?: (
46+
instanceId: ProviderInstanceId,
47+
) => Effect.Effect<object | undefined>;
4548
readonly list: () => Effect.Effect<ReadonlyArray<ProviderInstanceId>>;
4649
readonly getMetadata?: (instanceId: ProviderInstanceId) => Effect.Effect<
4750
{
@@ -68,6 +71,11 @@ export const layerFromProviderInstanceRegistry: Layer.Layer<
6871
Effect.gen(function* () {
6972
const instances = yield* ProviderInstanceRegistry.ProviderInstanceRegistry;
7073
return ProviderAdapterRegistryV2.of({
74+
getRuntimeIdentity: Effect.fn("ProviderAdapterRegistryV2.getRuntimeIdentity")((instanceId) =>
75+
instances
76+
.getInstance(instanceId)
77+
.pipe(Effect.map((instance) => instance?.orchestrationAdapter)),
78+
),
7179
get: Effect.fn("ProviderAdapterRegistryV2.get")(function (instanceId) {
7280
return instances.getInstance(instanceId).pipe(
7381
Effect.flatMap((instance) => {

‎apps/server/src/orchestration-v2/ProviderSessionManager.test.ts‎

Lines changed: 57 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -464,6 +464,7 @@ function layerTest(input: {
464464
readonly state: Ref.Ref<TestProviderRuntimeState>;
465465
readonly idleTimeoutMs: number;
466466
readonly maxIdlePinMs?: number;
467+
readonly runtimeIdentity?: Ref.Ref<object>;
467468
readonly failEventStream?: boolean;
468469
readonly capabilities?: OrchestrationV2ProviderCapabilities;
469470
readonly mcpConfigs?: Ref.Ref<
@@ -528,7 +529,17 @@ function layerTest(input: {
528529
...(input.scopeCloseReached === undefined
529530
? {}
530531
: { scopeCloseReached: input.scopeCloseReached }),
531-
}).pipe(Effect.map(ProviderAdapterRegistry.layerSingle)),
532+
}).pipe(
533+
Effect.map((adapter) =>
534+
input.runtimeIdentity === undefined
535+
? ProviderAdapterRegistry.layerSingle(adapter)
536+
: Layer.succeed(ProviderAdapterRegistry.ProviderAdapterRegistryV2, {
537+
get: () => Effect.succeed(adapter),
538+
list: () => Effect.succeed([adapter.instanceId]),
539+
getRuntimeIdentity: () => Ref.get(input.runtimeIdentity!),
540+
}),
541+
),
542+
),
532543
);
533544
const layerConfiguredMcpRegistry =
534545
input.pauseResolve === undefined
@@ -3411,6 +3422,51 @@ it.effect(
34113422
}),
34123423
);
34133424

3425+
it.effect(
3426+
"ProviderSessionManagerV2 replaces a cached session when its provider route changes",
3427+
() =>
3428+
Effect.gen(function* () {
3429+
const state = yield* Ref.make(emptyState);
3430+
const runtimeIdentity = yield* Ref.make<object>({});
3431+
const pendingWork = yield* Ref.make(false);
3432+
yield* Effect.gen(function* () {
3433+
const manager = yield* ProviderSessionManager.ProviderSessionManagerV2;
3434+
const eventSink = yield* EventSink.EventSinkV2;
3435+
const idAllocator = yield* IdAllocator.IdAllocatorV2;
3436+
const threadId = ThreadId.make("proxy-route-change");
3437+
const now = yield* DateTime.now;
3438+
yield* eventSink.write({
3439+
events: [yield* makeThreadCreatedEvent({ idAllocator, threadId, now })],
3440+
});
3441+
const providerSessionId = idAllocator.derive.providerSession({
3442+
providerInstanceId: modelSelection.instanceId,
3443+
});
3444+
const input = { threadId, providerSessionId, modelSelection, runtimePolicy };
3445+
const direct = yield* manager.open(input);
3446+
assert.strictEqual(yield* manager.open(input), direct);
3447+
yield* Ref.set(runtimeIdentity, {});
3448+
yield* Ref.set(pendingWork, true);
3449+
assert.isTrue(Exit.isFailure(yield* Effect.exit(manager.open(input))));
3450+
assert.equal((yield* Ref.get(state)).closeCount, 0);
3451+
yield* Ref.set(pendingWork, false);
3452+
const proxied = yield* manager.open(input);
3453+
assert.notStrictEqual(proxied, direct);
3454+
assert.equal((yield* Ref.get(state)).openCount, 2);
3455+
assert.equal((yield* Ref.get(state)).closeCount, 1);
3456+
assert.strictEqual(yield* manager.open(input), proxied);
3457+
}).pipe(
3458+
Effect.provide(
3459+
layerTest({
3460+
state,
3461+
runtimeIdentity,
3462+
idleTimeoutMs: 60000,
3463+
hasPendingBackgroundWork: Ref.get(pendingWork),
3464+
}),
3465+
),
3466+
);
3467+
}),
3468+
);
3469+
34143470
it.effect(
34153471
"ProviderSessionManagerV2 invalidates only errored threads and keeps reusable session errors cached",
34163472
() =>

‎apps/server/src/orchestration-v2/ProviderSessionManager.ts‎

Lines changed: 40 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -216,6 +216,7 @@ export class ProviderSessionManagerV2 extends Context.Service<
216216
>()("t3/orchestration-v2/ProviderSessionManager/ProviderSessionManagerV2") {}
217217

218218
interface LiveSessionEntry {
219+
readonly runtimeIdentity: object | undefined;
219220
readonly attachedThreadIds: ReadonlySet<ThreadId>;
220221
readonly loadedProviderThreadKeyByThread: ReadonlyMap<ThreadId, string>;
221222
/** Scope revocation to recorded credential ids so superseded sessions cannot revoke replacements opened before detach. */
@@ -2135,7 +2136,44 @@ export const layerWithOptions = (
21352136
}
21362137
}
21372138
const key = sessionKey(input.providerSessionId);
2138-
const existing = (yield* Ref.get(sessions)).get(key);
2139+
const runtimeIdentity =
2140+
registry.getRuntimeIdentity === undefined
2141+
? undefined
2142+
: yield* registry.getRuntimeIdentity(input.modelSelection.instanceId);
2143+
let existing = (yield* Ref.get(sessions)).get(key);
2144+
if (existing !== undefined && existing.runtimeIdentity !== runtimeIdentity) {
2145+
const changedSettings = () =>
2146+
new ProviderSessionOpenError({
2147+
instanceId: input.modelSelection.instanceId,
2148+
providerSessionId: input.providerSessionId,
2149+
cause:
2150+
"Provider settings changed. Wait for its active work to finish before starting another turn.",
2151+
});
2152+
const pending =
2153+
existing.runtime.hasPendingBackgroundWork === undefined
2154+
? false
2155+
: yield* existing.runtime.hasPendingBackgroundWork;
2156+
if (existing.busyTurns.size > 0 || pending) {
2157+
return yield* changedSettings();
2158+
}
2159+
yield* releaseEntry({
2160+
providerSessionId: input.providerSessionId,
2161+
reason: "manual_shutdown",
2162+
detail: "Provider settings changed.",
2163+
onlyIfIdleGeneration: existing.idleGeneration,
2164+
}).pipe(
2165+
Effect.mapError(
2166+
(cause) =>
2167+
new ProviderSessionOpenError({
2168+
instanceId: input.modelSelection.instanceId,
2169+
providerSessionId: input.providerSessionId,
2170+
cause,
2171+
}),
2172+
),
2173+
);
2174+
if ((yield* Ref.get(sessions)).has(key)) return yield* changedSettings();
2175+
existing = undefined;
2176+
}
21392177
if (existing !== undefined) {
21402178
const attached = existing.attachedThreadIds.has(input.threadId);
21412179
if (!attached && !existing.supportsMultipleProviderThreads) {
@@ -2234,6 +2272,7 @@ export const layerWithOptions = (
22342272
>(new Map());
22352273
const exposedRuntime = decorateRuntime(runtime, eventSubscribers);
22362274
const entry: LiveSessionEntry = {
2275+
runtimeIdentity,
22372276
attachedThreadIds: new Set([input.threadId]),
22382277
loadedProviderThreadKeyByThread: new Map(),
22392278
mcpCredentialIdByThread:

0 commit comments

Comments
 (0)